""" 进程内 / PostgreSQL inbox·outbox 账本。 本文件职责：定义任务结构、Memory 实现、按配置选择后端工厂。 禁止：Redis 全局 Worker 单例锁；禁止在回调线程里等 LLM。 有 PG_DSN 且 MESSAGE_STORE_BACKEND=postgres|auto 时用表；否则内存（本机联调）。 """ from __future__ import annotations import logging import threading import time import uuid from collections import OrderedDict from dataclasses import dataclass, field from enum import Enum from typing import Any, Optional, Union from agent.channel.wecom.models import InboundMessage logger = logging.getLogger(__name__) class TaskState(str, Enum): PENDING = "pending" CLAIMED = "claimed" DONE = "done" FAILED = "failed" @dataclass class InboxItem: """inbox 一条入站任务。""" id: str message: InboundMessage state: TaskState = TaskState.PENDING claimed_at: float = 0.0 attempts: int = 0 last_error: str = "" @dataclass class OutboxItem: """ outbox 一条出站任务。 touser：必须是入站 sender_id，禁止另换 ID。 发送成功 ≠ 业务成功；此处仅通道送达。 """ id: str touser: str content: str state: TaskState = TaskState.PENDING claimed_at: float = 0.0 attempts: int = 0 last_error: str = "" created_at: float = field(default_factory=time.time) dedupe_key: str = "" class MemoryMessageStore: """ 线程安全的内存 inbox/outbox。 认领：同一条任务同一时刻只被一个消费者持有；超时可回收（默认 120s）。 去重：相同 message_id / dedupe_key 不重复入队。 """ def __init__(self, *, claim_timeout_sec: float = 120.0, dedupe_maxlen: int = 5000) -> None: self._lock = threading.RLock() self._inbox: OrderedDict[str, InboxItem] = OrderedDict() self._outbox: OrderedDict[str, OutboxItem] = OrderedDict() self._seen_msg: OrderedDict[str, float] = OrderedDict() self._seen_out: OrderedDict[str, float] = OrderedDict() self._claim_timeout = claim_timeout_sec self._dedupe_maxlen = dedupe_maxlen self._inbox_event = threading.Event() self._outbox_event = threading.Event() def enqueue_inbound(self, message: InboundMessage) -> tuple[bool, str]: """ 写入 inbox。 返回：(是否新入队, 说明)。无 sender_id 拒绝；message_id 重复则跳过。 副作用：唤醒 inbox 等待线程。 """ if not message.has_sender(): return False, "missing_sender_id" with self._lock: if message.message_id in self._seen_msg: return False, "duplicate_message_id" self._remember(self._seen_msg, message.message_id) item = InboxItem(id=str(uuid.uuid4()), message=message) self._inbox[item.id] = item self._inbox_event.set() logger.info( "inbox 入队 id=%s msg_id=%s sender=%s type=%s", item.id, message.message_id, message.sender_id, message.msg_type, ) return True, item.id def claim_inbound(self) -> Optional[InboxItem]: """认领一条 pending/超时 inbox；无任务返回 None。""" now = time.time() with self._lock: for item in self._inbox.values(): if item.state == TaskState.PENDING or ( item.state == TaskState.CLAIMED and item.claimed_at and now - item.claimed_at > self._claim_timeout ): item.state = TaskState.CLAIMED item.claimed_at = now item.attempts += 1 return item self._inbox_event.clear() return None def complete_inbound(self, item_id: str, *, error: str = "") -> None: """标记 inbox 完成或失败。""" with self._lock: item = self._inbox.get(item_id) if not item: return if error: item.state = TaskState.FAILED item.last_error = error else: item.state = TaskState.DONE def enqueue_outbound( self, *, touser: str, content: str, dedupe_key: str = "", ) -> tuple[bool, str]: """写入 outbox；dedupe_key 重复则跳过。""" if not touser or not content: return False, "invalid_outbox" with self._lock: if dedupe_key and dedupe_key in self._seen_out: return False, "duplicate_outbox" if dedupe_key: self._remember(self._seen_out, dedupe_key) item = OutboxItem( id=str(uuid.uuid4()), touser=touser, content=content, dedupe_key=dedupe_key, ) self._outbox[item.id] = item self._outbox_event.set() logger.info("outbox 入队 id=%s touser=%s", item.id, touser) return True, item.id def claim_outbound(self) -> Optional[OutboxItem]: """认领一条 pending/超时 outbox。""" now = time.time() with self._lock: for item in self._outbox.values(): if item.state == TaskState.PENDING or ( item.state == TaskState.CLAIMED and item.claimed_at and now - item.claimed_at > self._claim_timeout ): item.state = TaskState.CLAIMED item.claimed_at = now item.attempts += 1 return item self._outbox_event.clear() return None def complete_outbound(self, item_id: str, *, error: str = "", permanent: bool = False) -> None: """ 标记 outbox 完成或失败。 permanent=True：不再重试（如 IP 未加白 60020）。 """ with self._lock: item = self._outbox.get(item_id) if not item: return if error: if permanent or item.attempts >= 5: item.state = TaskState.FAILED else: item.state = TaskState.PENDING item.claimed_at = 0.0 self._outbox_event.set() item.last_error = error else: item.state = TaskState.DONE def wait_inbox(self, timeout: float = 1.0) -> None: self._inbox_event.wait(timeout=timeout) def wait_outbox(self, timeout: float = 1.0) -> None: self._outbox_event.wait(timeout=timeout) def stats(self) -> dict[str, Any]: with self._lock: return { "backend": "memory", "inbox": len(self._inbox), "outbox": len(self._outbox), "seen_msg": len(self._seen_msg), } def _remember(self, store: OrderedDict[str, float], key: str) -> None: store[key] = time.time() while len(store) > self._dedupe_maxlen: store.popitem(last=False) MessageStoreImpl = Union[MemoryMessageStore, Any] _STORE: Optional[MessageStoreImpl] = None _STORE_LOCK = threading.Lock() def get_message_store() -> MessageStoreImpl: """ 获取消息账本单例。 MESSAGE_STORE_BACKEND=postgres 且 PG 可用 → PostgresMessageStore； auto 有 DSN 则尝试 PG，失败回退内存；memory 强制内存。 """ global _STORE with _STORE_LOCK: if _STORE is None: _STORE = _create_store() return _STORE def reset_message_store_for_tests() -> None: """仅自检用：清空单例。""" global _STORE with _STORE_LOCK: _STORE = None def _create_store() -> MessageStoreImpl: from agent.config import get_settings settings = get_settings() backend = (getattr(settings, "message_store_backend", None) or "auto").strip().lower() dsn = settings.resolve_pg_dsn() use_pg = backend == "postgres" or (backend == "auto" and bool(dsn)) if use_pg and dsn: try: from agent.channel.pg_store import PostgresMessageStore from agent.runtime_db import assert_safe_database assert_safe_database(dsn) store = PostgresMessageStore(dsn) store.stats() logger.info("消息账本后端=postgres") return store except Exception as exc: # noqa: BLE001 if backend == "postgres": raise logger.warning("PG 消息账本不可用，回退内存：%s", exc) logger.info("消息账本后端=memory") return MemoryMessageStore()