import asyncio import contextlib import logging from collections.abc import Awaitable, Callable from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any, Protocol, cast from uuid import uuid4 from sqlalchemy import select, update from app.core.config import Settings, get_settings from app.core.contracts import AgentRequest, AgentResult, RequestContext from app.core.errors import AgentError, RecoverableAgentError, RunLeaseLostError from app.infrastructure.db import SessionFactory from app.model.audit import InteractionAudit from app.model.conversation import ConversationMessage from app.model.platform import AgentRun, DomainEventOutbox, RequestIdempotency from app.repository.agent_run_repository import AgentRunRepository from app.service.agent.bootstrap import ( get_agent_factory, get_memory_cache_adapter, get_model_service, ) from app.service.agent.executor import AgentExecutor from app.service.agent.factory import AgentFactory from app.service.agent_persistence_service import AgentPersistenceService from app.service.identity_service import IdentityService from app.service.memory_extraction_service import ( ExtractionEndpointResolver, MemoryExtractionService, ) from app.service.memory_lifecycle_service import MemoryLifecycleService, Mode from app.service.memory_recall_service import MemoryRecallService from app.service.memory_service import CacheDeleteAdapter, MemoryService from app.service.memory_taxonomy import BUSINESS_EVENT_TYPES from app.service.model_gateway import DatabaseModelEndpointResolver, ModelGenerationService from app.worker.episode_worker import ( EpisodeConsumptionResult, EpisodeExtractionConsumer, EpisodeWorker, ) from app.worker.memory_extraction_worker import MemoryExtractionWorker from app.worker.outbox_worker import OutboxWorker logger = logging.getLogger(__name__) # episode 聚合按轮次节流：每 N 轮顺带处理一次已静默的会话片段， # 避免每轮轮询都做一次客户维度的 distinct 查询。 EPISODE_INTERVAL_ROUNDS = 30 # 一轮内最多消费的事件条数。`OutboxWorker.publish_one` 每次只领一条（单条事务让幂等与 # skip_locked 锁语义保持简单），所以清空速度原先等于轮询速度：262 条积压要 262 轮。 OUTBOX_DISPATCH_LIMIT = 10 # 单轮消费的片段条数上限：批处理必须可中断，不能一次吃完整库。 EPISODE_CONSUME_LIMIT = 20 PROJECTION_AUDIT_ACTION = "memory.projection_cleanup" @dataclass(frozen=True) class ProjectionCleanupOutcome: """投影清理结果；`cleaned=False` 时 `detail` 必须说明真实缺口。""" cleaned: bool detail: str = "" class ProjectionCleaner(Protocol): """Milvus/Neo4j 投影清理边界：由组装层注入，未注入即显式降级（绝不伪造成功）。""" async def cleanup(self, *, memory_uuid: str, operation: str) -> ProjectionCleanupOutcome: ... class WorkerRuntime: def __init__( self, factory: AgentFactory | None = None, settings: Settings | None = None, resolve_identity: Callable[[RequestContext], Awaitable[RequestContext]] | None = None, model_service: ModelGenerationService | None = None, endpoint_resolver: ExtractionEndpointResolver | None = None, memory_cache: CacheDeleteAdapter | None = None, projection_cleaner: ProjectionCleaner | None = None, relationships: Any = None, ) -> None: self.factory = factory if factory is not None else get_agent_factory() self.settings = settings or get_settings() self.resolve_identity = resolve_identity or IdentityService().resolve # 记忆抽取必须走与业务 Agent 相同的模型路由入口。默认取生产装配 # （app/service/agent/bootstrap.py），允许构造参数注入替身模型与端点解析器， # 使验收探针不依赖库中真实模型端点配置。 self.model_service = model_service if model_service is not None else get_model_service() self.memory_extraction = MemoryExtractionService( self.model_service, endpoint_resolver if endpoint_resolver is not None else DatabaseModelEndpointResolver(), ) # 召回热缓存适配器：任何记忆写入（事件消费、片段消费）都必须让该客户的 # 召回热缓存失效，否则新记忆在 TTL 内召回不到。默认取生产装配，可注入替身。 self.memory_cache: CacheDeleteAdapter | None = ( memory_cache if memory_cache is not None else get_memory_cache_adapter() ) # Milvus/Neo4j 删除客户端：由**组装层**（app/worker/__main__.py）注入生产实现， # 这里不兜底。组件内部给默认实现会把"尚未装配"这一事实悄悄盖住——而"未注入即显式 # 降级并留痕"是本模块刻意保留的语义（有单测守着），因此默认值保持 None。 self.projection_cleaner = projection_cleaner # 图关系服务：画像投影用它写入节点与关系（投顾的多跳推荐、风控的关系网络都读它）。 # 默认取生产装配；图库不可用时该值为 None，投影如实降级而不是失败。 if relationships is not None: self.relationships = relationships else: # 延迟导入：bootstrap 会间接导入本模块，模块级导入会形成循环依赖 from app.service.agent.bootstrap import get_relationship_service self.relationships = get_relationship_service() # episode 聚合是低频批处理，按轮次节流而不是每轮都查。 self._episode_rounds = 0 async def dispatch_one(self, *, run_id: str | None = None) -> bool: # Outbox acknowledges a durable SQL queue entry, not an in-memory task. async with SessionFactory() as session: async def dispatch(payload: dict[str, Any]) -> None: run = await AgentRunRepository(session).get(str(payload["run_id"])) if run is None: raise ValueError("run not found") async def dispatch_memory_extraction(payload: dict[str, Any]) -> None: if "message_id" not in payload or "customer_id" not in payload: raise ValueError("memory extraction payload is incomplete") # 幂等键只认事件 id，由 worker 自己按 payload 回查，避免调用方漏传。 # 注入召回缓存适配器：写入生效后立即失效该客户的热缓存。 await MemoryExtractionWorker( session, extractor=self.memory_extraction, cache=self.memory_cache ).handle(payload) async def dispatch_run_completed(payload: dict[str, Any]) -> None: # 结果消息与审计已由 complete_run 同事务落库，此事件只承担 # "运行已完成"的对外通知职责。当前没有独立外部消费者， # 这里显式消费以免事件永久滞留；接入推送链路时在此处扩展。 if not str(payload.get("run_id", "")): raise ValueError("agent.run_completed payload is incomplete") async def dispatch_cache_invalidate(payload: dict[str, Any]) -> None: await self._invalidate_config_cache(payload) async def dispatch_profile_rebuild(payload: dict[str, Any]) -> None: """画像重建 + 图投影，由记忆写入后发出的事件驱动。 为什么绕一层事件而不在记忆抽取处直接调用：抽取时那条记忆还在**未提交**的 事务里，另开 session 去重建画像看不到它（实测：快照加了、事实没进、图里 也没多出关系）。事件只可能在本事务提交之后被消费，届时数据一定可见。 """ customer_id = payload.get("customer_id") if not customer_id: raise ValueError("profile.rebuild_requested payload is incomplete") # 延迟导入：bootstrap 会间接导入本模块，模块级导入会形成循环依赖 from app.service.profile_assembly_service import ProfileAssemblyService from app.service.profile_graph_projection_service import ( ProfileGraphProjectionService, ) outcome = await ProfileAssemblyService(session).rebuild(int(customer_id)) if outcome.get("profile") is None: # 客户尚未开户：画像行由开户流程创建（trade_account 等字段属注册侧所有）， # 这里不代建。事实已提升进 user_facts，开户后再重建即可。 logger.info("profile rebuild skipped (not opened) customer_id=%s", customer_id) return projection = ProfileGraphProjectionService(session, self.relationships) result = await projection.project_customer(int(customer_id)) if result.degraded: logger.warning("graph projection degraded customer_id=%s reason=%s", customer_id, result.reason) async def dispatch_projection_cleanup(payload: dict[str, Any]) -> None: # memory.invalidated / memory.deleted 由 MemoryLifecycleService 按 # memory_uuid 写入，这里做幂等的投影清理（Milvus 向量、Neo4j 关系）。 await self._cleanup_projection(payload, session=session) async def dispatch_memory_deletion(payload: dict[str, Any]) -> None: # 客户级联失效/删除的公共入口：业务侧（用户注销、合规删除令） # 只写 memory.deletion_requested 事件，不直接操作记忆表。 customer_id = payload.get("customer_id") if not customer_id: raise ValueError("memory.deletion_requested payload is incomplete") mode = str(payload.get("mode", "invalidate")) if mode not in {"invalidate", "delete"}: raise ValueError("memory.deletion_requested mode is invalid") await MemoryLifecycleService(session).run( int(customer_id), mode=cast("Mode", mode), reason=str(payload.get("reason", "customer_lifecycle")), trace_id=str(payload.get("trace_id", "")), ) handlers: dict[str, Callable[[dict[str, Any]], Awaitable[None]]] = { "agent.run_requested": dispatch, "memory.extraction_requested": dispatch_memory_extraction, "agent.run_completed": dispatch_run_completed, "config.cache_invalidate_requested": dispatch_cache_invalidate, "memory.deletion_requested": dispatch_memory_deletion, # 投影清理：这两类事件此前没有消费者，永久 pending。 "memory.invalidated": dispatch_projection_cleanup, "memory.deleted": dispatch_projection_cleanup, # 画像重建：记忆写入后自动触发，使「记忆 → 画像 → 图」全链路无需手工介入 "profile.rebuild_requested": dispatch_profile_rebuild, } return await OutboxWorker(session, handlers).publish_one(aggregate_id=run_id) async def dispatch_batch(self, *, limit: int = OUTBOX_DISPATCH_LIMIT) -> int: """一轮内尽量多消费事件，返回实际消费条数。 单条领取的语义与幂等边界都不变，只是不再让"每轮一条"限制清空速度： 队列空时提前退出，因此稳态下与原来的一次调用开销相同，积压时才提速。 """ consumed = 0 for _ in range(max(1, limit)): if not await self.dispatch_one(): break consumed += 1 return consumed async def _invalidate_config_cache(self, payload: dict[str, Any]) -> None: """删除发布配置与记忆召回热缓存键；Redis 不可用时只记录告警，不阻塞事件消费。 记忆召回热缓存的键必须由 `MemoryRecallService.cache_keys` 枚举：此前这里手写的 `mem:mid:hot:{customer_id}` 并不是召回缓存的真实前缀（真实前缀是 `mem:recall`）， 失效动作一直打在并不存在的键上。事件 payload 里没有查询摘要与 limit，客户级 枚举是唯一可靠的失效方式（召回热缓存本身也只服务常见参数组合）。 """ release_id = payload.get("release_id") raw_customer_id = payload.get("customer_id") keys: list[str] = [] if raw_customer_id: keys.extend(MemoryRecallService.cache_keys(int(raw_customer_id))) if release_id is not None: keys.append(f"config:release:{release_id}") if not keys: return client = await self._redis_client() if client is None: return try: await client.delete(*keys) except Exception: logger.warning("config cache invalidation degraded keys=%s", ",".join(keys)) finally: await client.aclose() @staticmethod async def _conversation_history( session: Any, *, session_id: str, user_id: int, before_message_id: int | None, limit: int = 10, ) -> tuple[Any, ...]: """取该会话最近若干轮对话，按时间正序（旧 → 新）返回。 以 MySQL 的会话消息为**唯一来源**，不引入 Redis 双写：消息在受理时已经落库， 再同步一份到 Redis 只会带来不一致与 TTL 管理成本，换来的仅是一次索引查询的节省。 方案 §2.2 设想的是 Redis 列表，这里取等价语义（同样"最近若干轮、超出即截断"） 而不复制存储。 `before_message_id` 排除本轮请求消息本身：它刚写入库，若也算进历史， 模型会在上下文里看到自己的问题被重复一遍。 截断按**条数**而非 token：这里没有与模型一致的分词器，按 token 截断只能靠估算、 边界会随实现漂移；按条数是确定性的，宁可少给几轮，也不给一个不稳定的边界。 """ from app.core.contracts import ConversationTurn from app.repository.conversation_repository import ConversationRepository rows = await ConversationRepository(session).messages( session_id, user_id, limit + 1, before=before_message_id ) # repository 按 id DESC 返回（最新在前），这里翻正为旧 → 新 ordered = list(reversed(rows))[-limit:] turns: list[Any] = [] for row in ordered: content = str(row.content or "").strip() if not content: continue turns.append(ConversationTurn( role="assistant" if str(row.role) == "assistant" else "user", content=content, )) return tuple(turns) async def _cleanup_projection(self, payload: dict[str, Any], *, session: Any) -> None: """幂等清理一条记忆的派生投影（Milvus 向量、Neo4j 关系）。 为什么选择"显式降级 + 留痕"而不是假装成功：当前组装层没有提供任何删除客户端 （`bootstrap` 只装配了召回用的 Milvus 读适配器，Neo4j 连读适配器都没有）， 伪造"已删除"会让合规删除令在投影侧静默失效。因此无客户端时只做两件事： 记录告警，并写一条 `interaction_audit`（`status=skipped_no_client`），事件本身 照常标记为已消费——权威库（MySQL）状态已经正确，投影是可重建的派生数据， 让事件永久 pending 只会阻塞队列里其它事件。 注入 `projection_cleaner` 后同一入口执行真实删除，审计只记录适配器返回的真实 结论（`cleaned` 为假即写 `skipped`），不做任何"假定成功"的兜底。 """ memory_uuid = str(payload.get("memory_uuid", "")).strip() if not memory_uuid: raise ValueError("memory projection cleanup payload is incomplete") operation = str(payload.get("operation", "invalidate")) raw_customer_id = payload.get("customer_id") customer_id = int(raw_customer_id) if raw_customer_id else None cleaner = self.projection_cleaner if cleaner is None: logger.warning( "projection cleanup degraded: milvus/neo4j delete client not configured " "memory_uuid=%s operation=%s", memory_uuid, operation, ) await self._audit_projection( session, memory_uuid, customer_id, operation, status="skipped_no_client", reason="projection delete client not configured", ) return outcome = await cleaner.cleanup(memory_uuid=memory_uuid, operation=operation) if outcome.cleaned: logger.info("projection cleanup done memory_uuid=%s operation=%s", memory_uuid, operation) else: logger.warning("projection cleanup degraded memory_uuid=%s reason=%s", memory_uuid, outcome.detail) await self._audit_projection( session, memory_uuid, customer_id, operation, status="cleaned" if outcome.cleaned else "skipped", reason=outcome.detail, ) async def _audit_projection( self, session: Any, memory_uuid: str, customer_id: int | None, operation: str, *, status: str, reason: str, ) -> None: """投影清理留痕：审计行与事件同事务提交，降级状态在库里可查。""" session.add(InteractionAudit( actor_type="system", actor_id=None, target_customer_id=customer_id, session_id=None, portal=None, action_type=PROJECTION_AUDIT_ACTION, detail={ "memory_uuid": memory_uuid, "operation": operation, "status": status, "reason": reason, }, created_at=datetime.now(UTC).replace(tzinfo=None), )) await session.flush() async def _redis_client(self) -> Any: try: from redis.asyncio import Redis except ImportError: logger.warning("redis client unavailable; cache invalidation skipped") return None client: Any = Redis.from_url( self.settings.redis_url, socket_connect_timeout=self.settings.redis_connect_timeout_seconds, socket_timeout=self.settings.redis_connect_timeout_seconds, ) return client async def run_once(self) -> bool: dispatched = await self.dispatch_batch() > 0 self._episode_rounds += 1 if self._episode_rounds % EPISODE_INTERVAL_ROUNDS == 0: # 会话片段聚合：内部幂等（content_hash 唯一键），失败只告警， # 不得影响 run 的处理与轮询节奏。 try: await self.aggregate_episodes() except Exception: logger.warning("episode aggregation failed", exc_info=True) # 聚合之后立刻消费：片段只有被消费才会变成记忆，否则"待提取"永久滞留。 try: await self.consume_episodes() except Exception: logger.warning("episode extraction failed", exc_info=True) now = datetime.now(UTC).replace(tzinfo=None) async with SessionFactory() as session: run_id = await session.scalar(select(AgentRun.run_id).where( AgentRun.status.in_(("queued", "running", "cancel_requested")), (AgentRun.locked_until.is_(None) | (AgentRun.locked_until < now)), ).order_by(AgentRun.created_at).limit(1)) if run_id is None: return dispatched return await self.execute(run_id) or dispatched async def aggregate_episodes(self, *, customer_limit: int = 50) -> int: """把已静默的会话片段聚合为 episode，返回新写入的片段数。 可重复调用：同一片段的 `content_hash` 命中唯一键时直接跳过， 不需要额外的幂等表。 """ # 必须由调用方开启并提交事务：EpisodeWorker 只做 add/flush（savepoint 幂等）， # 自己不提交。漏掉提交会同时出现"返回已插入"与"表里没有行"。 async with SessionFactory() as session, session.begin(): # customer_id 列可空：SQL 的 is_not(None) 不足以让类型检查器收窄， # 这里显式再过滤一次。 raw_customers = list(await session.scalars( select(ConversationMessage.customer_id) .where(ConversationMessage.customer_id.is_not(None)) .distinct() .limit(max(1, customer_limit)) )) customers = [int(item) for item in raw_customers if item is not None] worker = EpisodeWorker(session) inserted = 0 for customer_id in customers: result = await worker.aggregate(customer_id) inserted += len(result.inserted) return inserted async def consume_episodes( self, *, limit: int = EPISODE_CONSUME_LIMIT ) -> EpisodeConsumptionResult: """消费一批待提取片段（低频批处理），把片段提升为长期记忆。 必须由调用方开启并提交事务：消费者只做 flush 与状态更新，自己不提交， 否则会出现"返回已处理但表里状态没变"。单个片段失败不冒泡（在消费者内部 转换为 `失败` + 重试计数），因此本方法只在数据库层面失败时抛错。 """ async with SessionFactory() as session, session.begin(): return await EpisodeExtractionConsumer( session, extractor=self.memory_extraction, cache=self.memory_cache, limit=limit, ).consume_pending() async def execute(self, run_id: str) -> bool: # A new fencing token for each claim also fences restarts of the same process. worker_id = str(uuid4()) async with SessionFactory() as session, session.begin(): run = await session.scalar(select(AgentRun).where( AgentRun.run_id == run_id).with_for_update()) if run is None: return False if run.status == "cancel_requested": # 文档 §6.4：运行真正落到 cancelled 时，原请求在 request_idempotency 中 # 以 failed + RUN_CANCELLED 结束（HTTP 层取消受理时可能已写过一次， # 这里幂等覆盖，保证 worker 抢先落终态的场景也不漏）。 run.status = "cancelled" run.completed_at = datetime.now(UTC).replace(tzinfo=None) run.locked_until = None await session.execute(update(RequestIdempotency).where( RequestIdempotency.id == run.idempotency_id ).values(status="failed", error_code="RUN_CANCELLED", updated_at=run.completed_at)) return True claimed = await AgentRunRepository(session).claim( run_id, worker_id, self.settings.worker_lease_seconds) if not claimed: return False task = asyncio.create_task(self._execute_claimed(run_id, worker_id)) heartbeat = asyncio.create_task(self._heartbeat(run_id, worker_id, task)) try: await task except asyncio.CancelledError: await self._failure(run_id, worker_id, "RUN_INTERRUPTED", retryable=True) # Runtime cancellation is shutdown; a lost lease only cancels the child. current = asyncio.current_task() if current is not None and current.cancelling(): raise except RunLeaseLostError: await self._failure(run_id, worker_id, "RUN_LEASE_LOST", retryable=True) except Exception as exc: code = exc.code if isinstance(exc, AgentError) else "AGENT_INTERNAL_ERROR" await self._failure(run_id, worker_id, code, retryable=isinstance(exc, RecoverableAgentError)) # 记录堆栈：只记录异常类型会让线上排障无从下手（run 的 error_code 只有 # AGENT_INTERNAL_ERROR，看不到真实原因）。 logger.warning("run failed run_id=%s error_type=%s", run_id, type(exc).__name__, exc_info=True) finally: heartbeat.cancel() with contextlib.suppress(asyncio.CancelledError): await heartbeat return True async def _heartbeat( self, run_id: str, worker_id: str, task: asyncio.Task[None] ) -> None: try: while True: await asyncio.sleep(self.settings.worker_lease_seconds / 3) async with SessionFactory() as session, session.begin(): renewed = await AgentRunRepository(session).renew( run_id, worker_id, self.settings.worker_lease_seconds) if not renewed: task.cancel() return except Exception: task.cancel() logger.warning("lease renewal failed run_id=%s", run_id) async def _execute_claimed(self, run_id: str, worker_id: str) -> None: async with SessionFactory() as session: run = await AgentRunRepository(session).get(run_id) if run is None: raise ValueError("run not found") message = await session.get(ConversationMessage, run.request_message_id) idem = await session.get(RequestIdempotency, run.idempotency_id) event = await session.scalar(select(DomainEventOutbox).where( DomainEventOutbox.aggregate_id == run_id, DomainEventOutbox.event_type == "agent.run_requested").limit(1)) if message is None or idem is None: raise ValueError("run input missing") # 短期会话记忆：加载本次之前的对话，模型靠它解析指代。 history = await self._conversation_history( session, session_id=run.session_id, user_id=int(run.user_id), before_message_id=run.request_message_id, ) request = AgentRequest( agent_type=run.agent_type, message=message.content, session_id=run.session_id, idempotency_key=idem.idempotency_key, metadata=event.payload.get("metadata", {}) if event else {}, history=history, ) identity = RequestContext(user_id=str(run.user_id), trace_id=run.trace_id) # Re-check account and permissions at execution time, including delayed jobs. context = await self.resolve_identity(identity) result: AgentResult | None = None async for event_data in AgentExecutor(self.factory).execute( request.agent_type, request, context, run_id ): if event_data.event_type == "done": result = AgentResult.model_validate(event_data.payload["result"]) if result is None: raise ValueError("Agent produced no terminal result") # 业务事件查询必须用独立 Session：在 complete_run 的 Session 上先跑 SELECT 会触发 # SQLAlchemy 的 autobegin，使 complete_run 内部的 session.begin() 抛 # "A transaction is already begun on this Session"，运行直接失败。 async with SessionFactory() as events_session: business_events = list(await events_session.scalars( select(DomainEventOutbox.event_type).where( DomainEventOutbox.aggregate_id == run_id, DomainEventOutbox.event_type.in_(tuple(BUSINESS_EVENT_TYPES)), ) )) async with SessionFactory() as session: await AgentPersistenceService(session).complete_run( run_id, result, worker_id=worker_id, memory_extraction_requested=MemoryService.should_extract_memory( conversation_content=request.message, role="user", # 工具产出的权威事实同样构成持久记忆（工具调用记录来自终态结果）。 tool_result=any( call.status == "succeeded" for call in result.result.tool_calls ), # 本 run 落库的业务事件（风险评估完成、交易完成等）。 event_type=business_events[0] if business_events else None, # 用户明确陈述的偏好/约束/身份/目标，命中才触发抽取。 signals=MemoryService.detect_memory_signals(request.message), ), ) async def _failure( self, run_id: str, worker_id: str, error_code: str, *, retryable: bool ) -> None: now = datetime.now(UTC).replace(tzinfo=None) async with SessionFactory() as session, session.begin(): run = await session.scalar(select(AgentRun).where( AgentRun.run_id == run_id).with_for_update()) if run is None or run.worker_id != worker_id: return if run.status not in {"running", "cancel_requested"}: return if run.status == "cancel_requested": run.status = "cancelled" elif retryable and run.attempt_count < self.settings.worker_retry_limit: run.status = "queued" else: run.status = "failed" run.error_code, run.updated_at = error_code, now run.locked_until = ( now + timedelta(seconds=min(60, 2**run.attempt_count)) if run.status == "queued" else None ) if run.status != "queued": run.completed_at = now # 取消导致的终态属于“原请求已终止”，按文档 §6.4 在 request_idempotency # 中写 failed + RUN_CANCELLED；`agent_run.error_code` 仍保留真实触发原因 #（如 RUN_INTERRUPTED），两者层级不同。 await session.execute(update(RequestIdempotency).where( RequestIdempotency.id == run.idempotency_id ).values( status="failed", error_code="RUN_CANCELLED" if run.status == "cancelled" else error_code, updated_at=now, )) session.add(InteractionAudit( actor_type="agent", actor_id=run.user_id, session_id=run.session_id, portal="api", action_type=f"agent.run_{run.status}", detail={"run_id": run_id, "error_code": error_code}, created_at=now, ))