From 21ce70fbfca19cc4370e9e4ea90f56a89ac2ede2 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 02:56:07 +0800 Subject: [PATCH 01/14] refactor(agent): route chat persistence through database worker --- app/agent/memory/__init__.py | 58 +++++++ app/agent/orchestrator.py | 40 ++--- app/api/endpoints/agent.py | 62 ++++++-- app/application/messaging/chat.py | 220 +++++++++++++++++++++++++- app/startup/modules_initializer.py | 13 +- tests/conftest.py | 10 ++ tests/test_agent_background_output.py | 16 +- tests/test_agent_chat_persistence.py | 154 ++++++++++++++++++ tests/test_web_agent_stream.py | 69 ++++++-- 9 files changed, 586 insertions(+), 56 deletions(-) create mode 100644 tests/test_agent_chat_persistence.py diff --git a/app/agent/memory/__init__.py b/app/agent/memory/__init__.py index 32b23fbee..d65241926 100644 --- a/app/agent/memory/__init__.py +++ b/app/agent/memory/__init__.py @@ -10,6 +10,7 @@ from app.runtime.settings import RuntimeSettingsCompat settings = RuntimeSettingsCompat() from app.application.agentdata import AgentChatPort as AgentChatOper +from app.application.messaging.chat import get_configured_agent_chat_persistence from app.runtime.log import logger from app.schemas.agent import ConversationMemory @@ -105,6 +106,42 @@ class MemoryManager: self.save_memory(memory) return memory.messages + async def async_get_agent_messages( + self, session_id: str, user_id: str + ) -> List[BaseMessage]: + """异步恢复 Agent 消息，持久化读取经有界数据库 worker 承接。""" + memory = self.get_memory(session_id, user_id) + if memory: + return memory.messages + + try: + persistence = get_configured_agent_chat_persistence() + chat = await persistence.async_get( + session_id=session_id, + user_id=user_id, + ) + if not chat: + chat = await persistence.async_get(session_id=session_id) + except Exception as e: + logger.debug(f"读取持久化Agent会话失败: {e}") + return [] + if not chat or not chat.agent_messages: + return [] + + try: + messages = messages_from_dict(chat.agent_messages) + except Exception as e: + logger.debug(f"恢复持久化Agent消息失败: {e}") + return [] + + memory = ConversationMemory( + session_id=session_id, + user_id=user_id, + messages=messages, + ) + self.save_memory(memory) + return memory.messages + def save_agent_messages( self, session_id: str, user_id: str, messages: List[BaseMessage] ): @@ -129,6 +166,27 @@ class MemoryManager: except Exception as e: logger.debug(f"持久化Agent消息失败: {e}") + async def async_save_agent_messages( + self, session_id: str, user_id: str, messages: List[BaseMessage] + ) -> None: + """异步保存 Agent 消息，持久化写入经有界数据库 worker 承接。""" + memory = self.get_memory(session_id, user_id) + if not memory: + memory = ConversationMemory(session_id=session_id, user_id=user_id) + + memory.messages = messages + memory.updated_at = datetime.now() + self.save_memory(memory) + try: + persistence = get_configured_agent_chat_persistence() + await persistence.async_save_agent_messages( + session_id=session_id, + user_id=user_id, + messages=messages_to_dict(messages), + ) + except Exception as e: + logger.debug(f"持久化Agent消息失败: {e}") + def save_memory(self, memory: ConversationMemory): """ 保存记忆到内存缓存 diff --git a/app/agent/orchestrator.py b/app/agent/orchestrator.py index 8687cf80a..c12995e16 100644 --- a/app/agent/orchestrator.py +++ b/app/agent/orchestrator.py @@ -79,9 +79,12 @@ from app.application.plugin.runtime import get_plugin_manager def _get_plugin_tools_revision() -> int: """读取插件工具目录修订号，避免 Agent 编排依赖具体管理器类型。""" return get_plugin_manager().get_plugin_agent_tools_revision() -from app.application.agentdata import AgentChatPort as AgentChatOper from app.application.agentdata import AgentTaskPort as AgentTaskOper from app.application.agentdata import UserPort as UserOper +from app.application.messaging.chat import ( + get_configured_agent_chat_persistence, + has_custom_agent_chat_title, +) from app.runtime.log import logger from app.schemas.event import AgentLLMProviderEventData from app.schemas.event import AgentTokensUsageEventData @@ -471,14 +474,14 @@ class MoviePilotAgent: """ return bool(self.channel and self.source) - def _save_display_history_messages(self, messages: List[dict]) -> None: + async def _save_display_history_messages(self, messages: List[dict]) -> None: """ 将一组可见消息追加到 Agent 会话历史表。 """ if not messages or not self._should_save_display_history(): return try: - AgentChatOper().append_display_messages( + await get_configured_agent_chat_persistence().async_append_display_messages( session_id=self.session_id, user_id=self.user_id, username=self.username, @@ -490,13 +493,13 @@ class MoviePilotAgent: except Exception as e: logger.debug(f"写入Agent展示历史失败: {e}") - def _save_assistant_display_message_once(self, message: str) -> None: + async def _save_assistant_display_message_once(self, message: str) -> None: """ 保存一条助手回复展示记录，并标记本轮已写入。 """ if not message or self._tool_context.get("assistant_display_saved"): return - self._save_display_history_messages( + await self._save_display_history_messages( [self.build_display_message(role="assistant", content=message)] ) self._tool_context["assistant_display_saved"] = True @@ -576,18 +579,16 @@ class MoviePilotAgent: return self._tool_context["chat_title_prepared"] = True try: - chat = await run_in_threadpool( - AgentChatOper().get, + chat = await get_configured_agent_chat_persistence().async_get( session_id=self.session_id, user_id=self.user_id, ) - if chat and AgentChatOper.has_custom_title(chat.title): + if chat and has_custom_agent_chat_title(chat.title): return title = await self._generate_chat_title(message) if not title: return - await run_in_threadpool( - AgentChatOper().update_title_if_empty, + await get_configured_agent_chat_persistence().async_update_title_if_empty( session_id=self.session_id, user_id=self.user_id, title=title, @@ -2167,9 +2168,12 @@ class MoviePilotAgent: return confirmation_result # 获取历史消息 - messages = list(memory_manager.get_agent_messages( - session_id=self.session_id, user_id=self.user_id - )) + messages = list( + await memory_manager.async_get_agent_messages( + session_id=self.session_id, + user_id=self.user_id, + ) + ) # 构建结构化用户消息内容 request_payload = { @@ -2194,7 +2198,7 @@ class MoviePilotAgent: content.append({"type": "image_url", "image_url": {"url": img}}) messages.append(HumanMessage(content=content)) await self.prepare_chat_title(message) - self._save_display_history_messages( + await self._save_display_history_messages( [ self.build_display_message( role="user", @@ -2219,7 +2223,7 @@ class MoviePilotAgent: error_message = f"处理消息时发生错误: {str(e)}" logger.error(error_message) if not user_display_saved: - self._save_display_history_messages( + await self._save_display_history_messages( [self.build_display_message(role="user", content=message)] ) if not self.should_dispatch_reply: @@ -2460,10 +2464,10 @@ class MoviePilotAgent: if hasattr(msg, "type") and msg.type == "ai" and msg.content: display_text = LLMHelper.extract_text_content(msg.content).strip() break - self._save_assistant_display_message_once(display_text) + await self._save_assistant_display_message_once(display_text) if self._should_persist_agent_chat(): - memory_manager.save_agent_messages( + await memory_manager.async_save_agent_messages( session_id=self.session_id, user_id=self.user_id, messages=agent.get_state(agent_config).values.get("messages", []), @@ -2517,7 +2521,7 @@ class MoviePilotAgent: and self.channel == NotificationChannel.Telegram.value else None ) - self._save_assistant_display_message_once(message) + await self._save_assistant_display_message_once(message) await AgentChain().async_post_message( Message( channel=None if broadcast else self.channel, diff --git a/app/api/endpoints/agent.py b/app/api/endpoints/agent.py index eaa5bf218..eef533bb7 100644 --- a/app/api/endpoints/agent.py +++ b/app/api/endpoints/agent.py @@ -53,6 +53,7 @@ from app.application.messaging.chat import ( AgentChatRecord, AgentChatService, get_configured_agent_chat_service, + get_configured_agent_chat_persistence, ) from app.application.security.user import get_configured_user_id_lookup from app.application.configuration import get_api_runtime_config_snapshot @@ -513,6 +514,25 @@ def _build_web_agent_session_id(user: ApiPrincipal, session_id: Optional[str]) - return f"{WEB_AGENT_SESSION_PREFIX}{digest[:32]}" +async def _build_web_agent_session_id_async( + user: ApiPrincipal, + session_id: Optional[str], +) -> str: + """异步解析 Web Agent 会话 ID，历史查询经有界数据库 worker 承接。""" + seed = str(session_id or "").strip() or uuid.uuid4().hex + if seed.startswith(WEB_AGENT_SESSION_PREFIX): + return seed + try: + existing_chat = await get_configured_agent_chat_persistence().async_get(seed) + if existing_chat and AgentChatService.can_access(existing_chat, user): + return seed + except Exception as e: + logger.debug(f"读取WebAgent历史会话失败: {e}") + user_part = user.name or str(user.id) + digest = hashlib.sha256(f"{user_part}:{seed}".encode("utf-8")).hexdigest() + return f"{WEB_AGENT_SESSION_PREFIX}{digest[:32]}" + + def _can_access_agent_chat(chat: Any, user: ApiPrincipal) -> bool: """ 判断当前登录用户是否可以访问指定 Agent 会话。 @@ -632,7 +652,7 @@ def _apply_web_agent_display_event(event: dict, assistant_message: dict) -> None tool["status"] = "done" -def _save_web_agent_display_snapshot( +async def _save_web_agent_display_snapshot( *, session_id: str, current_user: ApiPrincipal, @@ -643,9 +663,10 @@ def _save_web_agent_display_snapshot( 保存 WebAgent 当前展示消息快照。 """ try: - service = get_configured_agent_chat_service() - existing_chat = service.get_sync(session_id) - service.save_display_sync( + existing_chat = await get_configured_agent_chat_persistence().async_get( + session_id + ) + await get_configured_agent_chat_persistence().async_save_display_messages( session_id=session_id, user_id=(existing_chat.user_id if existing_chat else str(current_user.id)), username=(existing_chat.username if existing_chat else current_user.name), @@ -735,7 +756,10 @@ def _sanitize_web_agent_upload_name( return safe_name -def _get_web_agent_upload_dir(user: ApiPrincipal, session_id: Optional[str]) -> Path: +async def _get_web_agent_upload_dir( + user: ApiPrincipal, + session_id: Optional[str], +) -> Path: """ 计算当前 Web Agent 会话的临时附件目录。 @@ -743,7 +767,7 @@ def _get_web_agent_upload_dir(user: ApiPrincipal, session_id: Optional[str]) -> :param session_id: 前端会话标识 :return: 已创建的临时附件目录 """ - server_session_id = _build_web_agent_session_id(user, session_id) + server_session_id = await _build_web_agent_session_id_async(user, session_id) safe_session_id = server_session_id.replace(":", "_") upload_dir = get_api_runtime_config_snapshot().temp_path / "agent_uploads" / safe_session_id upload_dir.mkdir(parents=True, exist_ok=True) @@ -1690,7 +1714,7 @@ async def upload_web_agent_file( """ mime_type = file.content_type or mimetypes.guess_type(file.filename or "")[0] safe_name = _sanitize_web_agent_upload_name(file.filename, mime_type) - upload_dir = _get_web_agent_upload_dir(current_user, session_id) + upload_dir = await _get_web_agent_upload_dir(current_user, session_id) target_path = upload_dir / f"{uuid.uuid4().hex[:8]}_{safe_name}" size = await _save_web_agent_upload(file, target_path) attachment = _register_web_agent_file( @@ -1821,7 +1845,10 @@ async def get_agent_chat_session( chat = await _get_accessible_agent_chat(service, session_id, current_user) server_session_id = session_id if not chat: - server_session_id = _build_web_agent_session_id(current_user, session_id) + server_session_id = await _build_web_agent_session_id_async( + current_user, + session_id, + ) if server_session_id != session_id: chat = await _get_accessible_agent_chat( service, @@ -1881,8 +1908,7 @@ async def save_agent_chat_display( message.model_dump(exclude_none=True) for message in payload.messages ] - await run_in_threadpool( - _save_web_agent_display_snapshot, + await _save_web_agent_display_snapshot( session_id=session_id, current_user=current_user, messages=messages, @@ -1937,7 +1963,10 @@ async def stop_web_agent_session_task( :param service: Agent 会话应用服务 :return: 停止结果 """ - server_session_id = _build_web_agent_session_id(current_user, session_id) + server_session_id = await _build_web_agent_session_id_async( + current_user, + session_id, + ) chat = await _get_accessible_agent_chat( service, server_session_id, @@ -1973,7 +2002,10 @@ async def _web_agent_stream_impl( prompt = payload.text.strip() locale = LocaleHelper.get_locale_from_request(request) display_prompt = (payload.display_text or payload.text).strip() - session_id = _build_web_agent_session_id(current_user, payload.session_id) + session_id = await _build_web_agent_session_id_async( + current_user, + payload.session_id, + ) is_secret_confirmation_candidate = ( prompt in {"确认", "取消"} and not payload.images @@ -2075,8 +2107,7 @@ async def _web_agent_stream_impl( async def save_display_snapshot() -> None: """后台保存传统消息展示快照，不阻塞 SSE 终态。""" try: - await run_in_threadpool( - _save_web_agent_display_snapshot, + await _save_web_agent_display_snapshot( session_id=session_id, current_user=current_user, messages=display_messages, @@ -2239,8 +2270,7 @@ async def _web_agent_stream_impl( # 终态先进入 SSE 队列，避免展示快照落库延迟前端结束动画。 event_publisher.publish(done_event) if not is_secret_confirmation_control: - await run_in_threadpool( - _save_web_agent_display_snapshot, + await _save_web_agent_display_snapshot( session_id=session_id, current_user=current_user, messages=display_messages, diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index 71c4545ad..f1e40b576 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -3,11 +3,21 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Optional, Protocol +from collections.abc import Callable +from typing import Any, Optional, Protocol, TypeVar +from app.application.database import AsyncDatabaseExecutor from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary +T = TypeVar("T") + + +def has_custom_agent_chat_title(value: Optional[str]) -> bool: + """判断会话标题是否已经脱离默认占位标题。""" + return bool(value and value.strip() and value.strip() != "未命名会话") + + class AgentChatPrincipal(Protocol): """会话访问控制所需的最小用户身份。""" @@ -72,6 +82,72 @@ class AsyncAgentChatRepository(Protocol): ... +class SyncAgentChatRepository(Protocol): + """仅包含 Agent 编排所需同步持久化方法的适配器端口。""" + + def get( + self, + session_id: str, + user_id: Optional[str] = None, + ) -> Optional[Any]: + """读取服务端会话。""" + ... + + def append_display_messages( + self, + session_id: str, + user_id: Optional[str] = None, + messages: Optional[list[dict]] = None, + username: Optional[str] = None, + channel: Optional[Any] = None, + source: Optional[str] = None, + original_chat_id: Optional[str] = None, + client_session_id: Optional[str] = None, + ) -> Optional[Any]: + """追加用户可见消息。""" + ... + + def save_display_messages( + self, + session_id: str, + user_id: Optional[str] = None, + messages: Optional[list[dict]] = None, + username: Optional[str] = None, + channel: Optional[Any] = None, + source: Optional[str] = None, + original_chat_id: Optional[str] = None, + client_session_id: Optional[str] = None, + ) -> Optional[Any]: + """保存用户可见消息快照。""" + ... + + def save_agent_messages( + self, + session_id: str, + user_id: Optional[str], + messages: list[dict], + ) -> None: + """保存可恢复的原始 Agent 消息。""" + ... + + def update_title_if_empty( + self, + session_id: str, + user_id: Optional[str], + title: Optional[str], + username: Optional[str] = None, + channel: Optional[Any] = None, + source: Optional[str] = None, + original_chat_id: Optional[str] = None, + client_session_id: Optional[str] = None, + ) -> None: + """在会话尚无标题时写入标题。""" + ... + + +SyncAgentChatRepositoryFactory = Callable[[], SyncAgentChatRepository] + + @dataclass(frozen=True, slots=True) class AgentChatRecord: """脱离 ORM 会话的 Agent 会话持久化投影。""" @@ -265,7 +341,134 @@ class AgentChatService: ) +class AgentChatPersistenceService: + """把 Agent 编排所需的同步短事务委托给有界数据库 worker。""" + + def __init__( + self, + repository: SyncAgentChatRepositoryFactory, + async_executor: AsyncDatabaseExecutor, + ) -> None: + """保存同步仓储工厂和异步执行端口。""" + self._repository = repository + self._async_executor = async_executor + + async def _run(self, operation: Callable[[SyncAgentChatRepository], T]) -> T: + """在线程 worker 中执行一个完整的同步 AgentChat 短事务。""" + return await self._async_executor.run( + lambda: operation(self._repository()) + ) + + async def async_get( + self, + session_id: str, + user_id: Optional[str] = None, + ) -> Optional[Any]: + """异步读取会话，实际查询由有界 worker 承接。""" + return await self._run( + lambda repository: repository.get( + session_id=session_id, + user_id=user_id, + ) + ) + + async def async_append_display_messages( + self, + *, + session_id: str, + user_id: Optional[str] = None, + messages: Optional[list[dict]] = None, + username: Optional[str] = None, + channel: Optional[Any] = None, + source: Optional[str] = None, + original_chat_id: Optional[str] = None, + client_session_id: Optional[str] = None, + ) -> Optional[Any]: + """异步追加展示消息，等待同步事务取得确定终态。""" + return await self._run( + lambda repository: repository.append_display_messages( + session_id=session_id, + user_id=user_id, + messages=messages, + username=username, + channel=channel, + source=source, + original_chat_id=original_chat_id, + client_session_id=client_session_id, + ) + ) + + async def async_save_display_messages( + self, + *, + session_id: str, + user_id: Optional[str] = None, + messages: Optional[list[dict]] = None, + username: Optional[str] = None, + channel: Optional[Any] = None, + source: Optional[str] = None, + original_chat_id: Optional[str] = None, + client_session_id: Optional[str] = None, + ) -> Optional[Any]: + """异步保存展示消息快照，实际写入由有界 worker 承接。""" + return await self._run( + lambda repository: repository.save_display_messages( + session_id=session_id, + user_id=user_id, + messages=messages, + username=username, + channel=channel, + source=source, + original_chat_id=original_chat_id, + client_session_id=client_session_id, + ) + ) + + async def async_save_agent_messages( + self, + *, + session_id: str, + user_id: str, + messages: list[dict], + ) -> None: + """异步保存可恢复的原始消息。""" + await self._run( + lambda repository: repository.save_agent_messages( + session_id=session_id, + user_id=user_id, + messages=messages, + ) + ) + + async def async_update_title_if_empty( + self, + *, + session_id: str, + user_id: Optional[str], + title: Optional[str], + username: Optional[str] = None, + channel: Optional[Any] = None, + source: Optional[str] = None, + original_chat_id: Optional[str] = None, + client_session_id: Optional[str] = None, + ) -> None: + """异步写入首次生成的会话标题。""" + await self._run( + lambda repository: repository.update_title_if_empty( + session_id=session_id, + user_id=user_id, + title=title, + username=username, + channel=channel, + source=source, + original_chat_id=original_chat_id, + client_session_id=client_session_id, + ) + ) + + _configured_agent_chat_service: AgentChatService | None = None +_configured_agent_chat_persistence: AgentChatPersistenceService | None = None def configure_agent_chat_service(service: AgentChatService) -> None: @@ -279,3 +482,18 @@ def get_configured_agent_chat_service() -> AgentChatService: if _configured_agent_chat_service is None: raise RuntimeError("Agent 会话服务尚未配置") return _configured_agent_chat_service + + +def configure_agent_chat_persistence( + service: AgentChatPersistenceService, +) -> None: + """由启动组合根登记 Agent 编排所需的同步持久化端口。""" + global _configured_agent_chat_persistence + _configured_agent_chat_persistence = service + + +def get_configured_agent_chat_persistence() -> AgentChatPersistenceService: + """返回由启动组合根登记的 AgentChat worker 端口。""" + if _configured_agent_chat_persistence is None: + raise RuntimeError("Agent 会话持久化服务尚未配置") + return _configured_agent_chat_persistence diff --git a/app/startup/modules_initializer.py b/app/startup/modules_initializer.py index 4cd1f2286..f14c415a8 100644 --- a/app/startup/modules_initializer.py +++ b/app/startup/modules_initializer.py @@ -62,7 +62,12 @@ from app.application.database import configure_database_governance from app.application.service import configure_service_directory from app.application.plugin.runtime import configure_plugin_runtime from app.application.module import configure_module_runtime -from app.application.messaging.chat import AgentChatService, configure_agent_chat_service +from app.application.messaging.chat import ( + AgentChatPersistenceService, + AgentChatService, + configure_agent_chat_persistence, + configure_agent_chat_service, +) from app.application.security.user import configure_user_lookups from app.application.security.auth import AuthService, configure_auth_service from app.application.security.passkeys import PasskeyService, configure_passkey_service @@ -709,6 +714,12 @@ async def init_modules() -> HostRuntime: ) configure_database_governance(build_database_governance()) configure_agent_chat_service(AgentChatService(repository=AgentChatOper())) + configure_agent_chat_persistence( + AgentChatPersistenceService( + repository=AgentChatOper, + async_executor=database_worker, + ) + ) configure_user_lookups( by_id=lambda user_id: UserOper().get_by_id(user_id), by_name=lambda username: UserOper().get_by_name(username), diff --git a/tests/conftest.py b/tests/conftest.py index 378b14565..e53f522d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -118,6 +118,10 @@ def configure_plugin_system_services(): configure_chain_runtime_context_provider, ) from app.application.messaging.message import MessageHelper, MessageQueueManager + from app.application.messaging.chat import ( + AgentChatPersistenceService, + configure_agent_chat_persistence, + ) from app.runtime.cache import AsyncFileCache, FileCache from app.runtime.events import EventManager from app.runtime.extensions.module_manager import ModuleManager @@ -255,6 +259,12 @@ def configure_plugin_system_services(): workflow=lambda: WorkflowOper(), plugin_data=lambda: PluginDataOper(), ) + configure_agent_chat_persistence( + AgentChatPersistenceService( + repository=AgentChatOper, + async_executor=database_executor, + ) + ) from app.adapters.external.market import ( PluginHelper, VERSION_BACKWARD_COMPATIBLE_FLAGS, diff --git a/tests/test_agent_background_output.py b/tests/test_agent_background_output.py index bed906a57..af6700fa0 100644 --- a/tests/test_agent_background_output.py +++ b/tests/test_agent_background_output.py @@ -249,7 +249,11 @@ class AgentBackgroundOutputTest(unittest.IsolatedAsyncioTestCase): ) agent.send_agent_message = AsyncMock() - with patch.object(memory_manager, "save_agent_messages") as save_messages: + with patch.object( + memory_manager, + "async_save_agent_messages", + new=AsyncMock(), + ) as save_messages: await agent._execute_agent([HumanMessage(content="测试")]) save_messages.assert_called_once() @@ -277,7 +281,9 @@ class AgentBackgroundOutputTest(unittest.IsolatedAsyncioTestCase): with ( patch.object( - memory_manager, "get_agent_messages", return_value=cached_messages + memory_manager, + "async_get_agent_messages", + new=AsyncMock(return_value=cached_messages), ), patch.object(agent, "prepare_chat_title", new=AsyncMock()), patch.object(agent, "_save_display_history_messages"), @@ -305,7 +311,11 @@ class AgentBackgroundOutputTest(unittest.IsolatedAsyncioTestCase): ) agent.send_agent_message = AsyncMock() - with patch.object(memory_manager, "save_agent_messages") as save_messages: + with patch.object( + memory_manager, + "async_save_agent_messages", + new=AsyncMock(), + ) as save_messages: await agent._execute_agent([]) agent.send_agent_message.assert_awaited_once_with( diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py new file mode 100644 index 000000000..8b2cd6875 --- /dev/null +++ b/tests/test_agent_chat_persistence.py @@ -0,0 +1,154 @@ +"""AgentChat 同步短事务经有界 worker 委托的应用端口测试。""" + +from __future__ import annotations + +import asyncio +import threading +from uuid import uuid4 +from types import SimpleNamespace + +import pytest + +from app.application.messaging.chat import AgentChatPersistenceService +from app.db.oper.agentchat import AgentChatOper +from app.db.models.agentchat import AgentChat +from app.db.worker import DatabaseWorker + + +class _Executor: + """用独立线程模拟 G2B worker，验证调用方不会直接执行同步仓储。""" + + def __init__(self) -> None: + self.calls = 0 + self.worker_thread_id: int | None = None + + async def run(self, operation): + """在线程中执行一个完整的同步操作。""" + self.calls += 1 + + def invoke(): + self.worker_thread_id = threading.get_ident() + return operation() + + return await asyncio.to_thread(invoke) + + +class _Repository: + """记录 AgentChat 端口调用的同步仓储替身。""" + + def __init__(self) -> None: + self.calls: list[tuple[str, dict]] = [] + + def get(self, **kwargs): + self.calls.append(("get", kwargs)) + return SimpleNamespace(agent_messages=[]) + + def append_display_messages(self, **kwargs): + self.calls.append(("append_display_messages", kwargs)) + return None + + def save_display_messages(self, **kwargs): + self.calls.append(("save_display_messages", kwargs)) + return None + + def save_agent_messages(self, **kwargs): + self.calls.append(("save_agent_messages", kwargs)) + + def update_title_if_empty(self, **kwargs): + self.calls.append(("update_title_if_empty", kwargs)) + + +@pytest.mark.asyncio +async def test_agent_chat_persistence_runs_sync_repository_inside_worker() -> None: + """同步查询和写入都必须经过一次 worker admission。""" + executor = _Executor() + repository = _Repository() + service = AgentChatPersistenceService( + repository=lambda: repository, + async_executor=executor, + ) + caller_thread_id = threading.get_ident() + + await service.async_get("session-1", user_id="1") + await service.async_append_display_messages( + session_id="session-1", + user_id="1", + messages=[{"role": "user", "content": "hello"}], + ) + await service.async_save_display_messages( + session_id="session-1", + user_id="1", + messages=[], + ) + await service.async_save_agent_messages( + session_id="session-1", + user_id="1", + messages=[], + ) + await service.async_update_title_if_empty( + session_id="session-1", + user_id="1", + title="标题", + ) + + assert executor.calls == 5 + assert executor.worker_thread_id != caller_thread_id + assert [name for name, _kwargs in repository.calls] == [ + "get", + "append_display_messages", + "save_display_messages", + "save_agent_messages", + "update_title_if_empty", + ] + + +@pytest.mark.asyncio +async def test_agent_chat_persistence_propagates_worker_failure() -> None: + """worker admission 或事务异常必须原样返回给 async 应用调用方。""" + + class FailingExecutor: + async def run(self, _operation): + raise RuntimeError("worker failed") + + service = AgentChatPersistenceService( + repository=_Repository, + async_executor=FailingExecutor(), + ) + + with pytest.raises(RuntimeError, match="worker failed"): + await service.async_save_agent_messages( + session_id="session-1", + user_id="1", + messages=[], + ) + + +@pytest.mark.asyncio +async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() -> None: + """真实 AgentChat Oper 经 worker 写入后可被后续 worker 查询恢复。""" + worker = DatabaseWorker(max_workers=1, capacity=4) + await worker.start() + session_id = f"worker-{uuid4().hex}" + service = AgentChatPersistenceService( + repository=AgentChatOper, + async_executor=worker, + ) + + try: + await service.async_save_display_messages( + session_id=session_id, + user_id="worker-user", + username="worker-user", + channel="WebAgent", + source="worker-test", + messages=[{"role": "user", "content": "worker"}], + ) + chat = await service.async_get(session_id, user_id="worker-user") + assert chat is not None + assert chat.message_count == 1 + assert chat.display_messages[0]["content"] == "worker" + finally: + chat = AgentChatOper().get(session_id=session_id, user_id="worker-user") + if chat is not None: + AgentChat.delete(rid=chat.id) + await worker.shutdown() diff --git a/tests/test_web_agent_stream.py b/tests/test_web_agent_stream.py index ea0faeed0..9b9efcbcb 100644 --- a/tests/test_web_agent_stream.py +++ b/tests/test_web_agent_stream.py @@ -19,6 +19,7 @@ from app.api.endpoints.agent import ( _build_web_agent_message_events, _build_web_agent_command_items, _build_web_agent_session_id, + _build_web_agent_session_id_async, _build_web_agent_traditional_callback_payload, _build_web_agent_display_message_from_events, _collect_web_agent_traditional_events, @@ -174,6 +175,30 @@ def test_build_web_agent_session_id_reuses_accessible_history(): assert _build_web_agent_session_id(user, "telegram-session") == "telegram-session" +def test_build_web_agent_session_id_async_uses_worker_persistence(): + """异步 Web 会话解析应通过 AgentChat worker 端口读取历史。""" + user = SimpleNamespace(id=1, name="admin", is_superuser=True) + persistence = SimpleNamespace( + async_get=AsyncMock( + return_value=SimpleNamespace( + user_id="telegram-user", + username="tester", + ) + ) + ) + + with patch( + "app.api.endpoints.agent.get_configured_agent_chat_persistence", + return_value=persistence, + ): + session_id = asyncio.run( + _build_web_agent_session_id_async(user, "telegram-session") + ) + + assert session_id == "telegram-session" + persistence.async_get.assert_awaited_once_with("telegram-session") + + def test_apply_web_agent_display_event_updates_snapshot(): """WebAgent SSE 事件应按到达顺序聚合为服务端展示快照。""" message = { @@ -904,6 +929,7 @@ def test_web_agent_stream_emits_secret_result_only_as_protected_event(): return_value=FakeProtectedAgent, ), patch( "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, ) as save_snapshot: body = asyncio.run(scenario()) @@ -970,6 +996,7 @@ def test_web_agent_cancel_keeps_existing_display_history(): new=AsyncMock(return_value="已取消敏感设置读取。"), ), patch( "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, ) as save_snapshot: body = asyncio.run(scenario()) @@ -1113,6 +1140,7 @@ def test_web_agent_stream_drops_secret_result_after_disconnect(): new=AsyncMock(side_effect=finish_after_disconnect), ) as process, patch( "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, ) as save_snapshot: body = asyncio.run(scenario()) @@ -1156,7 +1184,7 @@ def test_web_agent_stream_emits_heartbeat_during_idle_tool_wait(): "app.api.endpoints.agent._has_web_agent_traditional_interaction", return_value=False, ), patch( - "app.api.endpoints.agent._build_web_agent_session_id", + "app.api.endpoints.agent._build_web_agent_session_id_async", return_value="web-agent:heartbeat", ), patch.object( MessageChain, @@ -1167,6 +1195,7 @@ def test_web_agent_stream_emits_heartbeat_during_idle_tool_wait(): side_effect=slow_process_message, ), patch( "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, ): body = asyncio.run(scenario()) @@ -1228,7 +1257,7 @@ def test_web_agent_stop_finishes_stream_without_error(): "app.api.endpoints.agent._has_web_agent_traditional_interaction", return_value=False, ), patch( - "app.api.endpoints.agent._build_web_agent_session_id", + "app.api.endpoints.agent._build_web_agent_session_id_async", return_value=session_id, ), patch.object( MessageChain, @@ -1238,6 +1267,7 @@ def test_web_agent_stop_finishes_stream_without_error(): return_value=BlockingWebAgent, ), patch( "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, ): body = asyncio.run(scenario()) finally: @@ -1275,9 +1305,10 @@ def test_web_agent_stream_rechecks_running_service_before_enqueue(): ), patch( "app.api.endpoints.agent.get_running_agent_manager", side_effect=[stale_manager, None], - ), patch( - "app.api.endpoints.agent._save_web_agent_display_snapshot", - ): + ), patch( + "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, + ): body = asyncio.run(scenario()) assert '"type": "error"' in body @@ -1299,11 +1330,11 @@ def test_web_agent_traditional_stream_keeps_alive_and_saves_after_done(): await asyncio.sleep(0.035) return [{"type": "delta", "content": "状态正常"}] - def slow_snapshot(**_kwargs): + async def slow_snapshot(**_kwargs): """阻塞快照写入，便于断言 done 不等待落库。""" - snapshot_started.set() - snapshot_release.wait(timeout=2) - snapshot_finished.set() + await asyncio.to_thread(snapshot_started.set) + await asyncio.to_thread(snapshot_release.wait, 2) + await asyncio.to_thread(snapshot_finished.set) async def scenario(): response = await web_agent_stream(payload, request, user) @@ -1323,6 +1354,8 @@ def test_web_agent_traditional_stream_keeps_alive_and_saves_after_done(): await asyncio.sleep(0.001) assert snapshot_started.is_set() assert not snapshot_finished.is_set() + snapshot_release.set() + await asyncio.to_thread(snapshot_finished.wait, 1) await iterator.aclose() return "".join(received) @@ -1340,13 +1373,14 @@ def test_web_agent_traditional_stream_keeps_alive_and_saves_after_done(): "app.api.endpoints.agent._get_web_agent_unknown_command_message", return_value=None, ), patch( - "app.api.endpoints.agent._build_web_agent_session_id", + "app.api.endpoints.agent._build_web_agent_session_id_async", return_value="web-agent:traditional-heartbeat", ), patch( "app.api.endpoints.agent._collect_web_agent_traditional_events", side_effect=slow_collect, ), patch( "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, side_effect=slow_snapshot, ): body = asyncio.run(scenario()) @@ -1354,7 +1388,6 @@ def test_web_agent_traditional_stream_keeps_alive_and_saves_after_done(): assert ": heartbeat\n\n" in body assert '"type": "delta"' in body assert '"type": "done"' in body - assert not snapshot_finished.is_set() finally: snapshot_release.set() @@ -1374,11 +1407,11 @@ def test_web_agent_stream_sends_done_before_snapshot_persistence_finishes(): """立即生成一段文本，随后进入终态。""" kwargs["output_callback"]("检查完成") - def slow_snapshot(**_kwargs): + async def slow_snapshot(**_kwargs): """阻塞快照写入，便于验证 done 的发送时机。""" - snapshot_started.set() - snapshot_release.wait(timeout=2) - snapshot_finished.set() + await asyncio.to_thread(snapshot_started.set) + await asyncio.to_thread(snapshot_release.wait, 2) + await asyncio.to_thread(snapshot_finished.set) async def scenario(): response = await web_agent_stream(payload, request, user) @@ -1398,6 +1431,8 @@ def test_web_agent_stream_sends_done_before_snapshot_persistence_finishes(): assert snapshot_started.is_set() assert not snapshot_finished.is_set() + snapshot_release.set() + await asyncio.to_thread(snapshot_finished.wait, 1) await iterator.aclose() return "".join(received) @@ -1412,7 +1447,7 @@ def test_web_agent_stream_sends_done_before_snapshot_persistence_finishes(): "app.api.endpoints.agent._has_web_agent_traditional_interaction", return_value=False, ), patch( - "app.api.endpoints.agent._build_web_agent_session_id", + "app.api.endpoints.agent._build_web_agent_session_id_async", return_value="web-agent:snapshot", ), patch.object( MessageChain, @@ -1423,12 +1458,12 @@ def test_web_agent_stream_sends_done_before_snapshot_persistence_finishes(): side_effect=immediate_process_message, ), patch( "app.api.endpoints.agent._save_web_agent_display_snapshot", + new_callable=AsyncMock, side_effect=slow_snapshot, ): body = asyncio.run(scenario()) assert '"type": "done"' in body - assert not snapshot_finished.is_set() finally: snapshot_release.set() From f560c19d653f3931e75d40e9277577332fb1d25c Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:04:19 +0800 Subject: [PATCH 02/14] fix(agent): keep chat reads on native async path --- app/agent/memory/__init__.py | 13 +++++--- app/agent/orchestrator.py | 3 +- app/api/endpoints/agent.py | 8 ++--- app/application/messaging/chat.py | 48 +++++++++++----------------- tests/conftest.py | 3 ++ tests/test_agent_chat_history.py | 46 ++++++++++++++++++++++++++ tests/test_agent_chat_persistence.py | 35 ++++++++++---------- tests/test_agent_chat_transaction.py | 1 + tests/test_web_agent_stream.py | 15 +++++---- 9 files changed, 106 insertions(+), 66 deletions(-) diff --git a/app/agent/memory/__init__.py b/app/agent/memory/__init__.py index d65241926..8775e34b8 100644 --- a/app/agent/memory/__init__.py +++ b/app/agent/memory/__init__.py @@ -10,7 +10,10 @@ from app.runtime.settings import RuntimeSettingsCompat settings = RuntimeSettingsCompat() from app.application.agentdata import AgentChatPort as AgentChatOper -from app.application.messaging.chat import get_configured_agent_chat_persistence +from app.application.messaging.chat import ( + get_configured_agent_chat_persistence, + get_configured_agent_chat_service, +) from app.runtime.log import logger from app.schemas.agent import ConversationMemory @@ -109,19 +112,19 @@ class MemoryManager: async def async_get_agent_messages( self, session_id: str, user_id: str ) -> List[BaseMessage]: - """异步恢复 Agent 消息，持久化读取经有界数据库 worker 承接。""" + """异步恢复 Agent 消息，查询与会话应用服务保持同一异步端口。""" memory = self.get_memory(session_id, user_id) if memory: return memory.messages try: - persistence = get_configured_agent_chat_persistence() - chat = await persistence.async_get( + service = get_configured_agent_chat_service() + chat = await service.get( session_id=session_id, user_id=user_id, ) if not chat: - chat = await persistence.async_get(session_id=session_id) + chat = await service.get(session_id=session_id) except Exception as e: logger.debug(f"读取持久化Agent会话失败: {e}") return [] diff --git a/app/agent/orchestrator.py b/app/agent/orchestrator.py index c12995e16..249b3b0cb 100644 --- a/app/agent/orchestrator.py +++ b/app/agent/orchestrator.py @@ -82,6 +82,7 @@ def _get_plugin_tools_revision() -> int: from app.application.agentdata import AgentTaskPort as AgentTaskOper from app.application.agentdata import UserPort as UserOper from app.application.messaging.chat import ( + get_configured_agent_chat_service, get_configured_agent_chat_persistence, has_custom_agent_chat_title, ) @@ -579,7 +580,7 @@ class MoviePilotAgent: return self._tool_context["chat_title_prepared"] = True try: - chat = await get_configured_agent_chat_persistence().async_get( + chat = await get_configured_agent_chat_service().get( session_id=self.session_id, user_id=self.user_id, ) diff --git a/app/api/endpoints/agent.py b/app/api/endpoints/agent.py index eef533bb7..c644da0ca 100644 --- a/app/api/endpoints/agent.py +++ b/app/api/endpoints/agent.py @@ -518,12 +518,12 @@ async def _build_web_agent_session_id_async( user: ApiPrincipal, session_id: Optional[str], ) -> str: - """异步解析 Web Agent 会话 ID，历史查询经有界数据库 worker 承接。""" + """异步解析 Web Agent 会话 ID，并复用异步会话查询端口。""" seed = str(session_id or "").strip() or uuid.uuid4().hex if seed.startswith(WEB_AGENT_SESSION_PREFIX): return seed try: - existing_chat = await get_configured_agent_chat_persistence().async_get(seed) + existing_chat = await get_configured_agent_chat_service().get(seed) if existing_chat and AgentChatService.can_access(existing_chat, user): return seed except Exception as e: @@ -663,9 +663,7 @@ async def _save_web_agent_display_snapshot( 保存 WebAgent 当前展示消息快照。 """ try: - existing_chat = await get_configured_agent_chat_persistence().async_get( - session_id - ) + existing_chat = await get_configured_agent_chat_service().get(session_id) await get_configured_agent_chat_persistence().async_save_display_messages( session_id=session_id, user_id=(existing_chat.user_id if existing_chat else str(current_user.id)), diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index f1e40b576..9bbf18fab 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -85,14 +85,6 @@ class AsyncAgentChatRepository(Protocol): class SyncAgentChatRepository(Protocol): """仅包含 Agent 编排所需同步持久化方法的适配器端口。""" - def get( - self, - session_id: str, - user_id: Optional[str] = None, - ) -> Optional[Any]: - """读取服务端会话。""" - ... - def append_display_messages( self, session_id: str, @@ -165,6 +157,7 @@ class AgentChatRecord: created_at: Any updated_at: Any messages: list[dict] + agent_messages: list[dict] class AsyncUnitOfWork(Protocol): @@ -222,9 +215,16 @@ class AgentChatService: return None return projected - async def get(self, session_id: str) -> Optional[AgentChatRecord]: + async def get( + self, + session_id: str, + user_id: Optional[str] = None, + ) -> Optional[AgentChatRecord]: """读取不附带授权判断的会话投影。""" - record = await self._repository.async_get(session_id=session_id) + record = await self._repository.async_get( + session_id=session_id, + user_id=user_id, + ) if record is None: return None return self._project(record) @@ -338,11 +338,12 @@ class AgentChatService: created_at=record.created_at, updated_at=record.updated_at, messages=list(record.display_messages or []), + agent_messages=list(record.agent_messages or []), ) class AgentChatPersistenceService: - """把 Agent 编排所需的同步短事务委托给有界数据库 worker。""" + """把 Agent 编排所需的同步持久化操作委托给有界数据库 worker。""" def __init__( self, @@ -354,24 +355,11 @@ class AgentChatPersistenceService: self._async_executor = async_executor async def _run(self, operation: Callable[[SyncAgentChatRepository], T]) -> T: - """在线程 worker 中执行一个完整的同步 AgentChat 短事务。""" + """在线程 worker 中执行一个同步 AgentChat 持久化操作。""" return await self._async_executor.run( lambda: operation(self._repository()) ) - async def async_get( - self, - session_id: str, - user_id: Optional[str] = None, - ) -> Optional[Any]: - """异步读取会话，实际查询由有界 worker 承接。""" - return await self._run( - lambda repository: repository.get( - session_id=session_id, - user_id=user_id, - ) - ) - async def async_append_display_messages( self, *, @@ -383,9 +371,9 @@ class AgentChatPersistenceService: source: Optional[str] = None, original_chat_id: Optional[str] = None, client_session_id: Optional[str] = None, - ) -> Optional[Any]: + ) -> None: """异步追加展示消息，等待同步事务取得确定终态。""" - return await self._run( + await self._run( lambda repository: repository.append_display_messages( session_id=session_id, user_id=user_id, @@ -397,6 +385,7 @@ class AgentChatPersistenceService: client_session_id=client_session_id, ) ) + return None async def async_save_display_messages( self, @@ -409,9 +398,9 @@ class AgentChatPersistenceService: source: Optional[str] = None, original_chat_id: Optional[str] = None, client_session_id: Optional[str] = None, - ) -> Optional[Any]: + ) -> None: """异步保存展示消息快照，实际写入由有界 worker 承接。""" - return await self._run( + await self._run( lambda repository: repository.save_display_messages( session_id=session_id, user_id=user_id, @@ -423,6 +412,7 @@ class AgentChatPersistenceService: client_session_id=client_session_id, ) ) + return None async def async_save_agent_messages( self, diff --git a/tests/conftest.py b/tests/conftest.py index e53f522d1..9b3395f0a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -119,7 +119,9 @@ def configure_plugin_system_services(): ) from app.application.messaging.message import MessageHelper, MessageQueueManager from app.application.messaging.chat import ( + AgentChatService, AgentChatPersistenceService, + configure_agent_chat_service, configure_agent_chat_persistence, ) from app.runtime.cache import AsyncFileCache, FileCache @@ -265,6 +267,7 @@ def configure_plugin_system_services(): async_executor=database_executor, ) ) + configure_agent_chat_service(AgentChatService(repository=AgentChatOper())) from app.adapters.external.market import ( PluginHelper, VERSION_BACKWARD_COMPATIBLE_FLAGS, diff --git a/tests/test_agent_chat_history.py b/tests/test_agent_chat_history.py index b87ebe460..25d1207fc 100644 --- a/tests/test_agent_chat_history.py +++ b/tests/test_agent_chat_history.py @@ -1,6 +1,7 @@ import asyncio import json from types import SimpleNamespace +from unittest.mock import AsyncMock from langchain_core.messages import AIMessage, HumanMessage @@ -278,3 +279,48 @@ def test_memory_manager_restores_agent_messages_from_database(): assert len(messages) == 1 assert isinstance(messages[0], HumanMessage) assert messages[0].content == "继续之前的话题" + + +def test_async_memory_manager_restores_through_native_async_service(monkeypatch): + """异步记忆恢复只能通过会话应用服务的异步查询端口。""" + session_id = "session-memory-async" + user_id = "3" + memory_manager.clear_memory(session_id, user_id) + service = SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + agent_messages=[ + { + "type": "human", + "data": { + "content": "异步恢复", + "additional_kwargs": {}, + "response_metadata": {}, + "type": "human", + "name": None, + "id": None, + "example": False, + }, + } + ] + ) + ) + ) + monkeypatch.setattr( + "app.agent.memory.get_configured_agent_chat_service", + lambda: service, + ) + + messages = asyncio.run( + memory_manager.async_get_agent_messages( + session_id=session_id, + user_id=user_id, + ) + ) + + assert len(messages) == 1 + assert messages[0].content == "异步恢复" + service.get.assert_awaited_once_with( + session_id=session_id, + user_id=user_id, + ) diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index 8b2cd6875..6501aaef3 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -5,13 +5,11 @@ from __future__ import annotations import asyncio import threading from uuid import uuid4 -from types import SimpleNamespace import pytest -from app.application.messaging.chat import AgentChatPersistenceService +from app.application.messaging.chat import AgentChatPersistenceService, AgentChatService from app.db.oper.agentchat import AgentChatOper -from app.db.models.agentchat import AgentChat from app.db.worker import DatabaseWorker @@ -39,10 +37,6 @@ class _Repository: def __init__(self) -> None: self.calls: list[tuple[str, dict]] = [] - def get(self, **kwargs): - self.calls.append(("get", kwargs)) - return SimpleNamespace(agent_messages=[]) - def append_display_messages(self, **kwargs): self.calls.append(("append_display_messages", kwargs)) return None @@ -60,7 +54,7 @@ class _Repository: @pytest.mark.asyncio async def test_agent_chat_persistence_runs_sync_repository_inside_worker() -> None: - """同步查询和写入都必须经过一次 worker admission。""" + """同步 AgentChat 写入必须经过一次 worker admission。""" executor = _Executor() repository = _Repository() service = AgentChatPersistenceService( @@ -69,7 +63,6 @@ async def test_agent_chat_persistence_runs_sync_repository_inside_worker() -> No ) caller_thread_id = threading.get_ident() - await service.async_get("session-1", user_id="1") await service.async_append_display_messages( session_id="session-1", user_id="1", @@ -91,10 +84,9 @@ async def test_agent_chat_persistence_runs_sync_repository_inside_worker() -> No title="标题", ) - assert executor.calls == 5 + assert executor.calls == 4 assert executor.worker_thread_id != caller_thread_id assert [name for name, _kwargs in repository.calls] == [ - "get", "append_display_messages", "save_display_messages", "save_agent_messages", @@ -125,17 +117,18 @@ async def test_agent_chat_persistence_propagates_worker_failure() -> None: @pytest.mark.asyncio async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() -> None: - """真实 AgentChat Oper 经 worker 写入后可被后续 worker 查询恢复。""" + """真实 AgentChat Oper 经 worker 写入后可被 native async 查询恢复。""" worker = DatabaseWorker(max_workers=1, capacity=4) await worker.start() session_id = f"worker-{uuid4().hex}" - service = AgentChatPersistenceService( + persistence = AgentChatPersistenceService( repository=AgentChatOper, async_executor=worker, ) + query = AgentChatService(repository=AgentChatOper()) try: - await service.async_save_display_messages( + await persistence.async_save_display_messages( session_id=session_id, user_id="worker-user", username="worker-user", @@ -143,12 +136,16 @@ async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() source="worker-test", messages=[{"role": "user", "content": "worker"}], ) - chat = await service.async_get(session_id, user_id="worker-user") + chat = await query.get( + session_id, + user_id="worker-user", + ) assert chat is not None assert chat.message_count == 1 - assert chat.display_messages[0]["content"] == "worker" + assert chat.messages[0]["content"] == "worker" finally: - chat = AgentChatOper().get(session_id=session_id, user_id="worker-user") - if chat is not None: - AgentChat.delete(rid=chat.id) + await AgentChatOper().async_delete( + session_id=session_id, + user_id="worker-user", + ) await worker.shutdown() diff --git a/tests/test_agent_chat_transaction.py b/tests/test_agent_chat_transaction.py index aa3ee1705..89e55b710 100644 --- a/tests/test_agent_chat_transaction.py +++ b/tests/test_agent_chat_transaction.py @@ -29,6 +29,7 @@ def _chat() -> SimpleNamespace: created_at=None, updated_at=None, display_messages=[], + agent_messages=[], ) diff --git a/tests/test_web_agent_stream.py b/tests/test_web_agent_stream.py index 9b9efcbcb..01c0bd447 100644 --- a/tests/test_web_agent_stream.py +++ b/tests/test_web_agent_stream.py @@ -175,28 +175,29 @@ def test_build_web_agent_session_id_reuses_accessible_history(): assert _build_web_agent_session_id(user, "telegram-session") == "telegram-session" -def test_build_web_agent_session_id_async_uses_worker_persistence(): - """异步 Web 会话解析应通过 AgentChat worker 端口读取历史。""" +def test_build_web_agent_session_id_async_uses_native_async_persistence(): + """异步 Web 会话解析应通过 native async 会话服务读取历史。""" user = SimpleNamespace(id=1, name="admin", is_superuser=True) - persistence = SimpleNamespace( - async_get=AsyncMock( + service = SimpleNamespace( + get=AsyncMock( return_value=SimpleNamespace( user_id="telegram-user", username="tester", + agent_messages=[], ) ) ) with patch( - "app.api.endpoints.agent.get_configured_agent_chat_persistence", - return_value=persistence, + "app.api.endpoints.agent.get_configured_agent_chat_service", + return_value=service, ): session_id = asyncio.run( _build_web_agent_session_id_async(user, "telegram-session") ) assert session_id == "telegram-session" - persistence.async_get.assert_awaited_once_with("telegram-session") + service.get.assert_awaited_once_with("telegram-session") def test_apply_web_agent_display_event_updates_snapshot(): From 9c4bead782918b20db21c955d7915959c96c452e Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:05:06 +0800 Subject: [PATCH 03/14] fix(agent): keep ORM results inside database worker --- app/application/messaging/chat.py | 28 +++++++++++++++------------- tests/test_agent_chat_persistence.py | 6 +++++- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index 9bbf18fab..984a69e8a 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -4,15 +4,12 @@ from __future__ import annotations from dataclasses import dataclass from collections.abc import Callable -from typing import Any, Optional, Protocol, TypeVar +from typing import Any, Optional, Protocol from app.application.database import AsyncDatabaseExecutor from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary -T = TypeVar("T") - - def has_custom_agent_chat_title(value: Optional[str]) -> bool: """判断会话标题是否已经脱离默认占位标题。""" return bool(value and value.strip() and value.strip() != "未命名会话") @@ -354,11 +351,16 @@ class AgentChatPersistenceService: self._repository = repository self._async_executor = async_executor - async def _run(self, operation: Callable[[SyncAgentChatRepository], T]) -> T: - """在线程 worker 中执行一个同步 AgentChat 持久化操作。""" - return await self._async_executor.run( - lambda: operation(self._repository()) - ) + async def _run_write( + self, + operation: Callable[[SyncAgentChatRepository], object], + ) -> None: + """在线程 worker 内完成同步写入并丢弃仓储对象返回值。""" + def execute() -> None: + """执行同步写入，不让 ORM 对象越过 worker 边界。""" + operation(self._repository()) + + await self._async_executor.run(execute) async def async_append_display_messages( self, @@ -373,7 +375,7 @@ class AgentChatPersistenceService: client_session_id: Optional[str] = None, ) -> None: """异步追加展示消息，等待同步事务取得确定终态。""" - await self._run( + await self._run_write( lambda repository: repository.append_display_messages( session_id=session_id, user_id=user_id, @@ -400,7 +402,7 @@ class AgentChatPersistenceService: client_session_id: Optional[str] = None, ) -> None: """异步保存展示消息快照，实际写入由有界 worker 承接。""" - await self._run( + await self._run_write( lambda repository: repository.save_display_messages( session_id=session_id, user_id=user_id, @@ -422,7 +424,7 @@ class AgentChatPersistenceService: messages: list[dict], ) -> None: """异步保存可恢复的原始消息。""" - await self._run( + await self._run_write( lambda repository: repository.save_agent_messages( session_id=session_id, user_id=user_id, @@ -443,7 +445,7 @@ class AgentChatPersistenceService: client_session_id: Optional[str] = None, ) -> None: """异步写入首次生成的会话标题。""" - await self._run( + await self._run_write( lambda repository: repository.update_title_if_empty( session_id=session_id, user_id=user_id, diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index 6501aaef3..ff5f7520f 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -19,6 +19,7 @@ class _Executor: def __init__(self) -> None: self.calls = 0 self.worker_thread_id: int | None = None + self.results: list[object] = [] async def run(self, operation): """在线程中执行一个完整的同步操作。""" @@ -26,7 +27,9 @@ class _Executor: def invoke(): self.worker_thread_id = threading.get_ident() - return operation() + result = operation() + self.results.append(result) + return result return await asyncio.to_thread(invoke) @@ -85,6 +88,7 @@ async def test_agent_chat_persistence_runs_sync_repository_inside_worker() -> No ) assert executor.calls == 4 + assert executor.results == [None, None, None, None] assert executor.worker_thread_id != caller_thread_id assert [name for name, _kwargs in repository.calls] == [ "append_display_messages", From d724dfa3f997d51754d75724e884e626756dc4de Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:07:40 +0800 Subject: [PATCH 04/14] test(agent): verify snapshot survives stream close --- tests/test_web_agent_stream.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_web_agent_stream.py b/tests/test_web_agent_stream.py index 01c0bd447..e13dc0ca9 100644 --- a/tests/test_web_agent_stream.py +++ b/tests/test_web_agent_stream.py @@ -1355,9 +1355,10 @@ def test_web_agent_traditional_stream_keeps_alive_and_saves_after_done(): await asyncio.sleep(0.001) assert snapshot_started.is_set() assert not snapshot_finished.is_set() + await iterator.aclose() + assert not snapshot_finished.is_set() snapshot_release.set() await asyncio.to_thread(snapshot_finished.wait, 1) - await iterator.aclose() return "".join(received) try: @@ -1432,9 +1433,10 @@ def test_web_agent_stream_sends_done_before_snapshot_persistence_finishes(): assert snapshot_started.is_set() assert not snapshot_finished.is_set() + await iterator.aclose() + assert not snapshot_finished.is_set() snapshot_release.set() await asyncio.to_thread(snapshot_finished.wait, 1) - await iterator.aclose() return "".join(received) try: From 8d05c6d8bf499e88ae4bb766bb88889f9de8eb17 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:18:02 +0800 Subject: [PATCH 05/14] fix(agent): serialize session persistence writes --- app/application/messaging/chat.py | 28 ++++++++++-- tests/test_agent_chat_persistence.py | 67 ++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index 984a69e8a..ce0e9e2bc 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -2,9 +2,11 @@ from __future__ import annotations +import asyncio from dataclasses import dataclass from collections.abc import Callable from typing import Any, Optional, Protocol +from weakref import WeakValueDictionary from app.application.database import AsyncDatabaseExecutor from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary @@ -350,17 +352,31 @@ class AgentChatPersistenceService: """保存同步仓储工厂和异步执行端口。""" self._repository = repository self._async_executor = async_executor + # append_display_messages 属于读取旧快照后整列写回的复合操作；按会话串行化， + # 才能在 worker 并发下保持首次建行和既有会话追加的完整性。弱引用避免长期运行 + # 中为一次性会话永久保留锁对象，不限制不同会话之间的 worker 并行度。 + self._session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() + + def _session_lock(self, session_id: str) -> asyncio.Lock: + """返回当前进程内指定会话的写锁。""" + lock = self._session_locks.get(session_id) + if lock is None: + lock = asyncio.Lock() + self._session_locks[session_id] = lock + return lock async def _run_write( self, + session_id: str, operation: Callable[[SyncAgentChatRepository], object], ) -> None: """在线程 worker 内完成同步写入并丢弃仓储对象返回值。""" - def execute() -> None: - """执行同步写入，不让 ORM 对象越过 worker 边界。""" - operation(self._repository()) + async with self._session_lock(session_id): + def execute() -> None: + """执行同步写入，不让 ORM 对象越过 worker 边界。""" + operation(self._repository()) - await self._async_executor.run(execute) + await self._async_executor.run(execute) async def async_append_display_messages( self, @@ -376,6 +392,7 @@ class AgentChatPersistenceService: ) -> None: """异步追加展示消息，等待同步事务取得确定终态。""" await self._run_write( + session_id, lambda repository: repository.append_display_messages( session_id=session_id, user_id=user_id, @@ -403,6 +420,7 @@ class AgentChatPersistenceService: ) -> None: """异步保存展示消息快照，实际写入由有界 worker 承接。""" await self._run_write( + session_id, lambda repository: repository.save_display_messages( session_id=session_id, user_id=user_id, @@ -425,6 +443,7 @@ class AgentChatPersistenceService: ) -> None: """异步保存可恢复的原始消息。""" await self._run_write( + session_id, lambda repository: repository.save_agent_messages( session_id=session_id, user_id=user_id, @@ -446,6 +465,7 @@ class AgentChatPersistenceService: ) -> None: """异步写入首次生成的会话标题。""" await self._run_write( + session_id, lambda repository: repository.update_title_if_empty( session_id=session_id, user_id=user_id, diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index ff5f7520f..225125a96 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -7,9 +7,12 @@ import threading from uuid import uuid4 import pytest +from sqlalchemy import delete, select from app.application.messaging.chat import AgentChatPersistenceService, AgentChatService +from app.db.models.agentchat import AgentChat from app.db.oper.agentchat import AgentChatOper +from app.db.session import SessionFactory, async_session_scope from app.db.worker import DatabaseWorker @@ -153,3 +156,67 @@ async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() user_id="worker-user", ) await worker.shutdown() + + +@pytest.mark.asyncio +async def test_agent_chat_persistence_serializes_same_session_writes() -> None: + """同一会话的首次创建和既有快照追加都必须串行。""" + worker = DatabaseWorker(max_workers=4, capacity=16) + await worker.start() + session_id = f"worker-race-{uuid4().hex}" + existing_session_id = f"worker-race-existing-{uuid4().hex}" + persistence = AgentChatPersistenceService( + repository=AgentChatOper, + async_executor=worker, + ) + + async def append(content: str) -> None: + await persistence.async_append_display_messages( + session_id=session_id, + user_id="worker-race-user", + messages=[{"role": "user", "content": content}], + ) + + async def append_existing(content: str) -> None: + await persistence.async_append_display_messages( + session_id=existing_session_id, + user_id="worker-race-user", + messages=[{"role": "user", "content": content}], + ) + + try: + await asyncio.gather(*(append(f"message-{index}") for index in range(4))) + await persistence.async_save_display_messages( + session_id=existing_session_id, + user_id="worker-race-user", + messages=[{"role": "user", "content": "seed"}], + ) + await asyncio.gather( + *(append_existing(f"existing-{index}") for index in range(4)) + ) + async with async_session_scope() as session: + result = await session.execute( + select(AgentChat).where( + AgentChat.session_id.in_((session_id, existing_session_id)) + ) + ) + rows = list(result.scalars().all()) + assert len(rows) == 2 + row_by_session = {row.session_id: row for row in rows} + assert { + message["content"] + for message in row_by_session[session_id].display_messages + } == {f"message-{index}" for index in range(4)} + assert { + message["content"] + for message in row_by_session[existing_session_id].display_messages + } == {"seed"} | {f"existing-{index}" for index in range(4)} + finally: + with SessionFactory() as session: + session.execute( + delete(AgentChat).where( + AgentChat.session_id.in_((session_id, existing_session_id)) + ) + ) + session.commit() + await worker.shutdown() From 598c004efeaba201b9ea093226c71e5c7727a153 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:21:34 +0800 Subject: [PATCH 06/14] fix(agent): bound session persistence admission --- app/application/messaging/chat.py | 36 +++++++++++++++--- app/runtime/observability/__init__.py | 2 + app/startup/modules_initializer.py | 1 + tests/test_agent_chat_persistence.py | 55 +++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 6 deletions(-) diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index ce0e9e2bc..24c012611 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -8,8 +8,15 @@ from collections.abc import Callable from typing import Any, Optional, Protocol from weakref import WeakValueDictionary -from app.application.database import AsyncDatabaseExecutor +from app.application.database import ( + AsyncDatabaseExecutor, + DatabaseWorkerOverloadedError, +) from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary +from app.runtime.observability import record_metric + + +DEFAULT_AGENT_CHAT_WRITE_CAPACITY = 32 def has_custom_agent_chat_title(value: Optional[str]) -> bool: @@ -348,10 +355,15 @@ class AgentChatPersistenceService: self, repository: SyncAgentChatRepositoryFactory, async_executor: AsyncDatabaseExecutor, + capacity: int = DEFAULT_AGENT_CHAT_WRITE_CAPACITY, ) -> None: """保存同步仓储工厂和异步执行端口。""" + if capacity < 1: + raise ValueError("AgentChat 写入容量必须大于 0") self._repository = repository self._async_executor = async_executor + self._capacity = capacity + self._pending_writes = 0 # append_display_messages 属于读取旧快照后整列写回的复合操作；按会话串行化， # 才能在 worker 并发下保持首次建行和既有会话追加的完整性。弱引用避免长期运行 # 中为一次性会话永久保留锁对象，不限制不同会话之间的 worker 并行度。 @@ -371,12 +383,24 @@ class AgentChatPersistenceService: operation: Callable[[SyncAgentChatRepository], object], ) -> None: """在线程 worker 内完成同步写入并丢弃仓储对象返回值。""" - async with self._session_lock(session_id): - def execute() -> None: - """执行同步写入，不让 ORM 对象越过 worker 边界。""" - operation(self._repository()) + # 会话锁前的等待也纳入固定总量，避免公开展示保存入口形成无界应用层队列。 + if self._pending_writes >= self._capacity: + record_metric("agent.chat.persistence.rejected") + raise DatabaseWorkerOverloadedError( + f"AgentChat 写入容量已用尽（上限 {self._capacity}）" + ) + self._pending_writes += 1 + record_metric("agent.chat.persistence.pending", self._pending_writes) + try: + async with self._session_lock(session_id): + def execute() -> None: + """执行同步写入，不让 ORM 对象越过 worker 边界。""" + operation(self._repository()) - await self._async_executor.run(execute) + await self._async_executor.run(execute) + finally: + self._pending_writes -= 1 + record_metric("agent.chat.persistence.pending", self._pending_writes) async def async_append_display_messages( self, diff --git a/app/runtime/observability/__init__.py b/app/runtime/observability/__init__.py index 1da6e1d10..bbebbca15 100644 --- a/app/runtime/observability/__init__.py +++ b/app/runtime/observability/__init__.py @@ -60,6 +60,8 @@ METRIC_SPECS = { MetricSpec("scheduler.job.dead_letter", MetricKind.COUNTER, frozenset({"owner"})), MetricSpec("plugin.lifecycle.duration", MetricKind.HISTOGRAM, frozenset({"operation", "outcome"})), MetricSpec("agent.active_tasks", MetricKind.GAUGE, frozenset({"task_type"})), + MetricSpec("agent.chat.persistence.pending", MetricKind.GAUGE, frozenset()), + MetricSpec("agent.chat.persistence.rejected", MetricKind.COUNTER, frozenset()), MetricSpec("agent.cancel", MetricKind.COUNTER, frozenset({"task_type", "outcome"})), MetricSpec("agent.provider.duration", MetricKind.HISTOGRAM, frozenset({"provider_type", "outcome"})), MetricSpec("agent.token_usage", MetricKind.COUNTER, frozenset({"provider_type", "direction"})), diff --git a/app/startup/modules_initializer.py b/app/startup/modules_initializer.py index f14c415a8..b889cc9e1 100644 --- a/app/startup/modules_initializer.py +++ b/app/startup/modules_initializer.py @@ -718,6 +718,7 @@ async def init_modules() -> HostRuntime: AgentChatPersistenceService( repository=AgentChatOper, async_executor=database_worker, + capacity=database_worker.snapshot().capacity, ) ) configure_user_lookups( diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index 225125a96..b894d0c0b 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -9,6 +9,7 @@ from uuid import uuid4 import pytest from sqlalchemy import delete, select +from app.application.database import DatabaseWorkerOverloadedError from app.application.messaging.chat import AgentChatPersistenceService, AgentChatService from app.db.models.agentchat import AgentChat from app.db.oper.agentchat import AgentChatOper @@ -122,6 +123,60 @@ async def test_agent_chat_persistence_propagates_worker_failure() -> None: ) +@pytest.mark.asyncio +async def test_agent_chat_persistence_bounds_session_waiters_and_releases_cancelled() -> None: + """同会话锁等待受总量限制，取消等待不会遗留 admission。""" + + class BlockingExecutor: + def __init__(self) -> None: + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def run(self, operation): + self.started.set() + await self.release.wait() + return operation() + + executor = BlockingExecutor() + service = AgentChatPersistenceService( + repository=_Repository, + async_executor=executor, + capacity=2, + ) + first = asyncio.create_task( + service.async_save_agent_messages( + session_id="session-admission", + user_id="1", + messages=[], + ) + ) + await executor.started.wait() + second = asyncio.create_task( + service.async_save_agent_messages( + session_id="session-admission", + user_id="1", + messages=[], + ) + ) + await asyncio.sleep(0) + third = asyncio.create_task( + service.async_save_agent_messages( + session_id="session-admission", + user_id="1", + messages=[], + ) + ) + with pytest.raises(DatabaseWorkerOverloadedError): + await third + second.cancel() + with pytest.raises(asyncio.CancelledError): + await second + assert service._pending_writes == 1 + executor.release.set() + await first + assert service._pending_writes == 0 + + @pytest.mark.asyncio async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() -> None: """真实 AgentChat Oper 经 worker 写入后可被 native async 查询恢复。""" From 6f5ee96152aaebe8ca93737d171eb14c8965b404 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:27:41 +0800 Subject: [PATCH 07/14] test: refresh dependency baseline for agent chat --- tests/fixtures/architecture/dependency-baseline.json | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/fixtures/architecture/dependency-baseline.json b/tests/fixtures/architecture/dependency-baseline.json index 0d992facf..d4d4e8fbd 100644 --- a/tests/fixtures/architecture/dependency-baseline.json +++ b/tests/fixtures/architecture/dependency-baseline.json @@ -13,8 +13,8 @@ "runtime_to_db": [], "workflow_to_db": [] }, - "edge_count": 6435, - "edge_sha256": "ed079837faf943d744ba7c6dee9172449bc204a2068a89db75e030e0be12fbb2", + "edge_count": 6443, + "edge_sha256": "da29d186d8ca672b117535abf9976a46123563a8a12532689496ce7dd56008f5", "edges": [ "app -> app.runtime", "app -> app.runtime.compat", @@ -237,6 +237,8 @@ "app.agent.mcp -> app.schemas.types", "app.agent.memory -> app.application", "app.agent.memory -> app.application.agentdata", + "app.agent.memory -> app.application.messaging", + "app.agent.memory -> app.application.messaging.chat", "app.agent.memory -> app.runtime", "app.agent.memory -> app.runtime.log", "app.agent.memory -> app.runtime.settings", @@ -345,6 +347,8 @@ "app.agent.orchestrator -> app.agent.tools.impl.query_system_settings", "app.agent.orchestrator -> app.application", "app.agent.orchestrator -> app.application.agentdata", + "app.agent.orchestrator -> app.application.messaging", + "app.agent.orchestrator -> app.application.messaging.chat", "app.agent.orchestrator -> app.application.plugin", "app.agent.orchestrator -> app.application.plugin.runtime", "app.agent.orchestrator -> app.chain", @@ -2553,6 +2557,10 @@ "app.application.mediaserver -> app.schemas.types", "app.application.messaging.agent -> app.schemas", "app.application.messaging.agent -> app.schemas.types", + "app.application.messaging.chat -> app.application", + "app.application.messaging.chat -> app.application.database", + "app.application.messaging.chat -> app.runtime", + "app.application.messaging.chat -> app.runtime.observability", "app.application.messaging.chat -> app.schemas", "app.application.messaging.chat -> app.schemas.agent", "app.application.messaging.interaction -> app.schemas", From 0ba4a7e5e3f5c71561f8108e22c9d715c507dc84 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:55:10 +0800 Subject: [PATCH 08/14] fix(agent): close async chat persistence lifecycle --- app/api/dependencies/agent.py | 9 + app/api/deps.py | 2 + app/api/endpoints/agent.py | 136 +++++++++---- app/application/messaging/agent.py | 23 ++- app/application/messaging/chat.py | 58 +++++- app/db/oper/agentchat.py | 3 +- app/startup/context.py | 2 + app/startup/modules_initializer.py | 22 +- tests/conftest.py | 3 +- .../architecture/dependency-baseline.json | 1 + tests/test_agent_chat_persistence.py | 192 +++++++++++++++++- tests/test_agent_lifecycle.py | 26 +++ tests/test_host_runtime_context.py | 14 +- tests/test_lifecycle_shutdown.py | 29 +++ 14 files changed, 448 insertions(+), 72 deletions(-) diff --git a/app/api/dependencies/agent.py b/app/api/dependencies/agent.py index fcac8e71f..cb283f95d 100644 --- a/app/api/dependencies/agent.py +++ b/app/api/dependencies/agent.py @@ -4,6 +4,7 @@ from fastapi import Depends from sqlalchemy.ext.asyncio import AsyncSession from app.api.context import ( + get_agent_chat_runtime, get_agent_chat_repository, get_agent_chat_transaction, get_async_session, @@ -11,6 +12,7 @@ from app.api.context import ( ) from app.application.messaging.chat import ( AgentChatService, + AgentChatPersistenceService, AsyncAgentChatRepository, AsyncUnitOfWork, ) @@ -26,6 +28,13 @@ def get_agent_chat_service( return AgentChatService(chat_repository, unit_of_work) +def get_agent_chat_persistence( + runtime: HostRuntime = Depends(get_agent_chat_runtime), +) -> AgentChatPersistenceService: + """从类型化 Agent 运行时获取有界会话写入端口。""" + return runtime.persistence + + def get_message_query_service( db: AsyncSession = Depends(get_async_session), runtime: HostRuntime = Depends(get_host_runtime), diff --git a/app/api/deps.py b/app/api/deps.py index 9007667af..6073df004 100644 --- a/app/api/deps.py +++ b/app/api/deps.py @@ -5,6 +5,7 @@ """ from app.api.dependencies.agent import ( + get_agent_chat_persistence, get_agent_chat_service, get_message_query_service, ) @@ -52,6 +53,7 @@ from app.api.dependencies.workflow import ( # 兼容聚合入口只显式列出既有 FastAPI 依赖，不向插件制造新的动态导出规则。 __all__ = [ + "get_agent_chat_persistence", "get_agent_chat_service", "get_auth_service", "get_current_active_manage_user", diff --git a/app/api/endpoints/agent.py b/app/api/endpoints/agent.py index c644da0ca..522af2a44 100644 --- a/app/api/endpoints/agent.py +++ b/app/api/endpoints/agent.py @@ -47,17 +47,25 @@ from app.command import Command from app.runtime.config import global_vars from app.runtime.events import Event, EventManager from app.api.principal import ApiPrincipal -from app.api.dependencies.agent import get_agent_chat_service +from app.api.dependencies.agent import ( + get_agent_chat_persistence, + get_agent_chat_service, +) from app.api.dependencies.auth import get_current_active_user from app.application.messaging.chat import ( AgentChatRecord, + AgentChatPersistenceService, AgentChatService, get_configured_agent_chat_service, get_configured_agent_chat_persistence, ) from app.application.security.user import get_configured_user_id_lookup from app.application.configuration import get_api_runtime_config_snapshot -from app.application.messaging.agent import attach_web_agent_edit_queue, detach_web_agent_edit_queue +from app.application.messaging.agent import ( + attach_web_agent_edit_queue, + create_web_agent_background_task, + detach_web_agent_edit_queue, +) from app.application.messaging.agent import agent_interaction_manager from app.application.messaging.agent import ( build_agent_choice_button_rows, @@ -88,7 +96,6 @@ _WEB_AGENT_FILE_REGISTRY: dict[str, dict[str, Any]] = {} _WEB_AGENT_MESSAGE_QUEUES: dict[str, list[Queue[_SchemaMessage]]] = {} _WEB_AGENT_MESSAGE_LOCK = Lock() _WEB_AGENT_MESSAGE_LISTENER_REGISTERED = False -_WEB_AGENT_BACKGROUND_TASKS: set[asyncio.Task] = set() class _WebAgentEventPublisher: @@ -517,13 +524,16 @@ def _build_web_agent_session_id(user: ApiPrincipal, session_id: Optional[str]) - async def _build_web_agent_session_id_async( user: ApiPrincipal, session_id: Optional[str], + service: Optional[AgentChatService] = None, ) -> str: """异步解析 Web Agent 会话 ID，并复用异步会话查询端口。""" seed = str(session_id or "").strip() or uuid.uuid4().hex if seed.startswith(WEB_AGENT_SESSION_PREFIX): return seed try: - existing_chat = await get_configured_agent_chat_service().get(seed) + if service is None: + service = get_configured_agent_chat_service() + existing_chat = await service.get(seed) if existing_chat and AgentChatService.can_access(existing_chat, user): return seed except Exception as e: @@ -658,36 +668,41 @@ async def _save_web_agent_display_snapshot( current_user: ApiPrincipal, messages: list[dict], client_session_id: Optional[str] = None, + service: Optional[AgentChatService] = None, + persistence: Optional[AgentChatPersistenceService] = None, ) -> None: """ 保存 WebAgent 当前展示消息快照。 """ - try: - existing_chat = await get_configured_agent_chat_service().get(session_id) - await get_configured_agent_chat_persistence().async_save_display_messages( - session_id=session_id, - user_id=(existing_chat.user_id if existing_chat else str(current_user.id)), - username=(existing_chat.username if existing_chat else current_user.name), - channel=( - existing_chat.channel - if existing_chat and existing_chat.channel - else NotificationChannel.WebAgent - ), - source=( - existing_chat.source - if existing_chat and existing_chat.source - else WEB_AGENT_SOURCE - ), - original_chat_id=existing_chat.original_chat_id if existing_chat else None, - client_session_id=( - existing_chat.client_session_id - if existing_chat and existing_chat.client_session_id - else client_session_id - ), - messages=messages, - ) - except Exception as e: - logger.debug(f"保存WebAgent展示历史失败: {e}") + if service is None: + # 直接调用该内部 helper 时没有 FastAPI 依赖注入上下文。 + service = get_configured_agent_chat_service() + existing_chat = await service.get(session_id) + if persistence is None: + # 直接调用该内部 helper 时没有 FastAPI 依赖注入上下文。 + persistence = get_configured_agent_chat_persistence() + await persistence.async_save_display_messages( + session_id=session_id, + user_id=(existing_chat.user_id if existing_chat else str(current_user.id)), + username=(existing_chat.username if existing_chat else current_user.name), + channel=( + existing_chat.channel + if existing_chat and existing_chat.channel + else NotificationChannel.WebAgent + ), + source=( + existing_chat.source + if existing_chat and existing_chat.source + else WEB_AGENT_SOURCE + ), + original_chat_id=existing_chat.original_chat_id if existing_chat else None, + client_session_id=( + existing_chat.client_session_id + if existing_chat and existing_chat.client_session_id + else client_session_id + ), + messages=messages, + ) def _build_web_agent_sse( @@ -757,6 +772,7 @@ def _sanitize_web_agent_upload_name( async def _get_web_agent_upload_dir( user: ApiPrincipal, session_id: Optional[str], + service: Optional[AgentChatService] = None, ) -> Path: """ 计算当前 Web Agent 会话的临时附件目录。 @@ -765,7 +781,11 @@ async def _get_web_agent_upload_dir( :param session_id: 前端会话标识 :return: 已创建的临时附件目录 """ - server_session_id = await _build_web_agent_session_id_async(user, session_id) + server_session_id = await _build_web_agent_session_id_async( + user, + session_id, + service, + ) safe_session_id = server_session_id.replace(":", "_") upload_dir = get_api_runtime_config_snapshot().temp_path / "agent_uploads" / safe_session_id upload_dir.mkdir(parents=True, exist_ok=True) @@ -1701,6 +1721,7 @@ async def upload_web_agent_file( file: UploadFile = File(...), session_id: Optional[str] = Form(None), current_user: ApiPrincipal = Depends(get_current_active_user), + service: AgentChatService = Depends(get_agent_chat_service), ) -> _SchemaResponse: """ 上传 Web 智能助手对话附件。 @@ -1712,7 +1733,7 @@ async def upload_web_agent_file( """ mime_type = file.content_type or mimetypes.guess_type(file.filename or "")[0] safe_name = _sanitize_web_agent_upload_name(file.filename, mime_type) - upload_dir = await _get_web_agent_upload_dir(current_user, session_id) + upload_dir = await _get_web_agent_upload_dir(current_user, session_id, service) target_path = upload_dir / f"{uuid.uuid4().hex[:8]}_{safe_name}" size = await _save_web_agent_upload(file, target_path) attachment = _register_web_agent_file( @@ -1846,6 +1867,7 @@ async def get_agent_chat_session( server_session_id = await _build_web_agent_session_id_async( current_user, session_id, + service, ) if server_session_id != session_id: chat = await _get_accessible_agent_chat( @@ -1884,6 +1906,7 @@ async def save_agent_chat_display( payload: _SchemaAgentChatDisplaySaveRequest, current_user: ApiPrincipal = Depends(get_current_active_user), service: AgentChatService = Depends(get_agent_chat_service), + persistence: AgentChatPersistenceService = Depends(get_agent_chat_persistence), ) -> _SchemaResponse: """ 保存前端聚合后的 Agent 展示消息。 @@ -1911,6 +1934,8 @@ async def save_agent_chat_display( current_user=current_user, messages=messages, client_session_id=existing_chat.client_session_id if existing_chat else session_id, + service=service, + persistence=persistence, ) chat = await service.get_accessible(session_id, current_user) if not chat: @@ -1964,6 +1989,7 @@ async def stop_web_agent_session_task( server_session_id = await _build_web_agent_session_id_async( current_user, session_id, + service, ) chat = await _get_accessible_agent_chat( service, @@ -1988,6 +2014,8 @@ async def _web_agent_stream_impl( payload: _SchemaAgentWebChatRequest, request: Request, current_user: ApiPrincipal = Depends(get_current_active_user), + service: Optional[AgentChatService] = None, + persistence: Optional[AgentChatPersistenceService] = None, ) -> StreamingResponse: """ Web 智能助手流式对话。 @@ -1998,11 +2026,18 @@ async def _web_agent_stream_impl( :return: SSE 流式响应 """ prompt = payload.text.strip() + if not isinstance(service, AgentChatService): + # 直接调用公开函数时不经过 FastAPI 依赖解析；生产路由总是传入运行时服务。 + service = get_configured_agent_chat_service() + if not isinstance(persistence, AgentChatPersistenceService): + # 直接调用公开函数时不经过 FastAPI 依赖解析；生产路由总是传入运行时端口。 + persistence = get_configured_agent_chat_persistence() locale = LocaleHelper.get_locale_from_request(request) display_prompt = (payload.display_text or payload.text).strip() session_id = await _build_web_agent_session_id_async( current_user, payload.session_id, + service, ) is_secret_confirmation_candidate = ( prompt in {"确认", "取消"} @@ -2110,13 +2145,13 @@ async def _web_agent_stream_impl( current_user=current_user, messages=display_messages, client_session_id=payload.session_id or session_id, + service=service, + persistence=persistence, ) except Exception as err: logger.error(f"保存WebAgent传统消息快照失败: {str(err)}") - snapshot_task = asyncio.create_task(save_display_snapshot()) - _WEB_AGENT_BACKGROUND_TASKS.add(snapshot_task) - snapshot_task.add_done_callback(_WEB_AGENT_BACKGROUND_TASKS.discard) + snapshot_task = create_web_agent_background_task(save_display_snapshot()) await asyncio.sleep(0) for event in events: event_payload = copy.deepcopy(event) @@ -2268,16 +2303,19 @@ async def _web_agent_stream_impl( # 终态先进入 SSE 队列，避免展示快照落库延迟前端结束动画。 event_publisher.publish(done_event) if not is_secret_confirmation_control: - await _save_web_agent_display_snapshot( - session_id=session_id, - current_user=current_user, - messages=display_messages, - client_session_id=payload.session_id or session_id, - ) + try: + await _save_web_agent_display_snapshot( + session_id=session_id, + current_user=current_user, + messages=display_messages, + client_session_id=payload.session_id or session_id, + service=service, + persistence=persistence, + ) + except Exception as err: + logger.error(f"保存WebAgent展示历史失败：{err}") - task = asyncio.create_task(run_agent()) - _WEB_AGENT_BACKGROUND_TASKS.add(task) - task.add_done_callback(_WEB_AGENT_BACKGROUND_TASKS.discard) + task = create_web_agent_background_task(run_agent()) disconnected = False terminal_sent = False try: @@ -2349,6 +2387,14 @@ async def web_agent_stream( payload: _SchemaAgentWebChatRequest, request: Request, current_user: ApiPrincipal = Depends(get_current_active_user), + service: AgentChatService = Depends(get_agent_chat_service), + persistence: AgentChatPersistenceService = Depends(get_agent_chat_persistence), ) -> StreamingResponse: """Web 智能助手流式对话的稳定公开路由入口。""" - return await _web_agent_stream_impl(payload, request, current_user) + return await _web_agent_stream_impl( + payload, + request, + current_user, + service, + persistence, + ) diff --git a/app/application/messaging/agent.py b/app/application/messaging/agent.py index 61887983e..8378c9770 100644 --- a/app/application/messaging/agent.py +++ b/app/application/messaging/agent.py @@ -1,9 +1,10 @@ +import asyncio import uuid from dataclasses import dataclass, field from datetime import datetime, timedelta from queue import Queue from threading import Lock -from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union +from typing import Awaitable, Callable, Dict, Iterable, List, Optional, Tuple, Union from app.schemas.types import NotificationChannel @@ -177,6 +178,26 @@ _WEB_AGENT_EDIT_QUEUES: dict[str, list[Queue[dict]]] = {} _WEB_AGENT_EDIT_LOCK = Lock() _ChannelAdminResolver = Callable[[Optional[dict]], Iterable[Union[str, int]]] _CHANNEL_ADMIN_RESOLVERS: dict[str, _ChannelAdminResolver] = {} +_WEB_AGENT_BACKGROUND_TASKS: set[asyncio.Task[object]] = set() + + +def create_web_agent_background_task( + coroutine: Awaitable[object], +) -> asyncio.Task[object]: + """登记 Web Agent 后台任务，使应用关闭时可以统一收口。""" + task = asyncio.create_task(coroutine) + _WEB_AGENT_BACKGROUND_TASKS.add(task) + task.add_done_callback(_WEB_AGENT_BACKGROUND_TASKS.discard) + return task + + +async def shutdown_web_agent_background_tasks() -> None: + """取消并等待 Web Agent 后台任务，避免关闭数据库后仍提交快照。""" + tasks = tuple(_WEB_AGENT_BACKGROUND_TASKS) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) def register_channel_admin_resolver( diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index 24c012611..9f509e295 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -10,6 +10,7 @@ from weakref import WeakValueDictionary from app.application.database import ( AsyncDatabaseExecutor, + DatabaseWorkerClosedError, DatabaseWorkerOverloadedError, ) from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary @@ -17,6 +18,7 @@ from app.runtime.observability import record_metric DEFAULT_AGENT_CHAT_WRITE_CAPACITY = 32 +DEFAULT_AGENT_CHAT_SESSION_CAPACITY = 4 def has_custom_agent_chat_title(value: Optional[str]) -> bool: @@ -143,7 +145,8 @@ class SyncAgentChatRepository(Protocol): ... -SyncAgentChatRepositoryFactory = Callable[[], SyncAgentChatRepository] +SyncAgentChatRepositoryFactory = Callable[[object], SyncAgentChatRepository] +SyncAgentChatTransaction = Callable[[Callable[[object], object]], object] @dataclass(frozen=True, slots=True) @@ -355,15 +358,24 @@ class AgentChatPersistenceService: self, repository: SyncAgentChatRepositoryFactory, async_executor: AsyncDatabaseExecutor, + sync_transaction: SyncAgentChatTransaction, capacity: int = DEFAULT_AGENT_CHAT_WRITE_CAPACITY, + session_capacity: int = DEFAULT_AGENT_CHAT_SESSION_CAPACITY, ) -> None: - """保存同步仓储工厂和异步执行端口。""" + """保存同步仓储工厂、事务端口和两级写入容量。""" if capacity < 1: raise ValueError("AgentChat 写入容量必须大于 0") + if session_capacity < 1: + raise ValueError("AgentChat 单会话写入容量必须大于 0") self._repository = repository self._async_executor = async_executor + self._sync_transaction = sync_transaction self._capacity = capacity + self._session_capacity = session_capacity self._pending_writes = 0 + self._pending_by_session: dict[str, int] = {} + self._active_tasks: set[asyncio.Task[object]] = set() + self._closing = False # append_display_messages 属于读取旧快照后整列写回的复合操作；按会话串行化， # 才能在 worker 并发下保持首次建行和既有会话追加的完整性。弱引用避免长期运行 # 中为一次性会话永久保留锁对象，不限制不同会话之间的 worker 并行度。 @@ -383,24 +395,52 @@ class AgentChatPersistenceService: operation: Callable[[SyncAgentChatRepository], object], ) -> None: """在线程 worker 内完成同步写入并丢弃仓储对象返回值。""" - # 会话锁前的等待也纳入固定总量，避免公开展示保存入口形成无界应用层队列。 - if self._pending_writes >= self._capacity: + # 同时限制全局和单会话等待量，避免一个热点会话占满总 admission 后饿死其他会话。 + if self._closing: + raise DatabaseWorkerClosedError("AgentChat 持久化服务当前不可接收任务") + session_pending = self._pending_by_session.get(session_id, 0) + if ( + self._pending_writes >= self._capacity + or session_pending >= self._session_capacity + ): record_metric("agent.chat.persistence.rejected") raise DatabaseWorkerOverloadedError( - f"AgentChat 写入容量已用尽（上限 {self._capacity}）" + f"AgentChat 写入容量已用尽（全局上限 {self._capacity}，" + f"单会话上限 {self._session_capacity}）" ) self._pending_writes += 1 - record_metric("agent.chat.persistence.pending", self._pending_writes) + self._pending_by_session[session_id] = session_pending + 1 + current = asyncio.current_task() + if current is not None: + self._active_tasks.add(current) + record_metric("agent.chat.persistence.pending", 1) try: async with self._session_lock(session_id): def execute() -> None: - """执行同步写入，不让 ORM 对象越过 worker 边界。""" - operation(self._repository()) + """在单一同步事务中执行写入，不让 ORM 对象越过 worker 边界。""" + self._sync_transaction( + lambda session: operation(self._repository(session)) + ) await self._async_executor.run(execute) finally: self._pending_writes -= 1 - record_metric("agent.chat.persistence.pending", self._pending_writes) + remaining = self._pending_by_session.get(session_id, 1) - 1 + if remaining: + self._pending_by_session[session_id] = remaining + else: + self._pending_by_session.pop(session_id, None) + if current is not None: + self._active_tasks.discard(current) + record_metric("agent.chat.persistence.pending", -1) + + async def shutdown(self) -> None: + """拒绝新写入并等待当前会话锁和 worker 操作取得终态。""" + self._closing = True + current = asyncio.current_task() + tasks = tuple(task for task in self._active_tasks if task is not current) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) async def async_append_display_messages( self, diff --git a/app/db/oper/agentchat.py b/app/db/oper/agentchat.py index acc3bccce..b81d7b612 100644 --- a/app/db/oper/agentchat.py +++ b/app/db/oper/agentchat.py @@ -269,7 +269,8 @@ class AgentChatOper(DbOper): ) if not chat: return None - display_messages = self._normalize_messages(chat.display_messages) + # JSON 列不是 MutableList；必须复制旧列表，原地 extend 会让 SQLAlchemy 误认为字段未变化。 + display_messages = list(self._normalize_messages(chat.display_messages)) display_messages.extend(self._normalize_messages(messages)) title = chat.title if self.has_custom_title(chat.title) else None return self.save_display_messages( diff --git a/app/startup/context.py b/app/startup/context.py index 2ba54b23c..a00690c55 100644 --- a/app/startup/context.py +++ b/app/startup/context.py @@ -6,6 +6,7 @@ from typing import Protocol from app.application.messaging.chat import ( AsyncAgentChatRepository, + AgentChatPersistenceService, AsyncUnitOfWork, ) from app.application.outbox import AsyncOutboxTransaction @@ -113,6 +114,7 @@ class AgentChatRuntime: async_session: AsyncSessionProvider repository: AgentChatRepositoryFactory transaction: AsyncUnitOfWorkFactory + persistence: AgentChatPersistenceService @dataclass(frozen=True, slots=True) diff --git a/app/startup/modules_initializer.py b/app/startup/modules_initializer.py index b889cc9e1..c222acf04 100644 --- a/app/startup/modules_initializer.py +++ b/app/startup/modules_initializer.py @@ -67,7 +67,9 @@ from app.application.messaging.chat import ( AgentChatService, configure_agent_chat_persistence, configure_agent_chat_service, + get_configured_agent_chat_persistence, ) +from app.application.messaging.agent import shutdown_web_agent_background_tasks from app.application.security.user import configure_user_lookups from app.application.security.auth import AuthService, configure_auth_service from app.application.security.passkeys import PasskeyService, configure_passkey_service @@ -576,6 +578,11 @@ async def stop_modules(): await run_step("消息服务", stop_message) await run_step("Redis缓存连接", lambda: RedisHelper().close()) await run_step("异步Redis缓存连接", lambda: AsyncRedisHelper().close()) + await run_step("Web Agent后台任务", shutdown_web_agent_background_tasks) + await run_step( + "Agent会话持久化", + lambda: get_configured_agent_chat_persistence().shutdown(), + ) await run_step("数据库任务", stop_database_worker) if _database_worker is None: await run_step("数据库连接", close_database) @@ -642,11 +649,18 @@ async def init_modules() -> HostRuntime: chain=lambda: build_chain_runtime_config(settings), ) runtime_settings = _build_runtime_settings_service() + agent_chat_persistence = AgentChatPersistenceService( + repository=lambda session: AgentChatOper(session), + async_executor=database_worker, + sync_transaction=transaction_runner.sync, + capacity=database_worker.snapshot().capacity, + ) host_runtime = HostRuntime( agent_chat=AgentChatRuntime( async_session=get_async_db, repository=AgentChatOper, transaction=SqlAlchemyAsyncUnitOfWork, + persistence=agent_chat_persistence, ), persistence=PersistenceRuntime( sync_session=get_db, @@ -714,13 +728,7 @@ async def init_modules() -> HostRuntime: ) configure_database_governance(build_database_governance()) configure_agent_chat_service(AgentChatService(repository=AgentChatOper())) - configure_agent_chat_persistence( - AgentChatPersistenceService( - repository=AgentChatOper, - async_executor=database_worker, - capacity=database_worker.snapshot().capacity, - ) - ) + configure_agent_chat_persistence(agent_chat_persistence) configure_user_lookups( by_id=lambda user_id: UserOper().get_by_id(user_id), by_name=lambda username: UserOper().get_by_name(username), diff --git a/tests/conftest.py b/tests/conftest.py index 9b3395f0a..9a746c245 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -263,8 +263,9 @@ def configure_plugin_system_services(): ) configure_agent_chat_persistence( AgentChatPersistenceService( - repository=AgentChatOper, + repository=lambda session: AgentChatOper(session), async_executor=database_executor, + sync_transaction=transaction_runner.sync, ) ) configure_agent_chat_service(AgentChatService(repository=AgentChatOper())) diff --git a/tests/fixtures/architecture/dependency-baseline.json b/tests/fixtures/architecture/dependency-baseline.json index d4d4e8fbd..06210fa82 100644 --- a/tests/fixtures/architecture/dependency-baseline.json +++ b/tests/fixtures/architecture/dependency-baseline.json @@ -6097,6 +6097,7 @@ "app.startup.modules_initializer -> app.application.history", "app.startup.modules_initializer -> app.application.image", "app.startup.modules_initializer -> app.application.messaging", + "app.startup.modules_initializer -> app.application.messaging.agent", "app.startup.modules_initializer -> app.application.messaging.chat", "app.startup.modules_initializer -> app.application.messaging.message", "app.startup.modules_initializer -> app.application.module", diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index b894d0c0b..396bb8e23 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -4,16 +4,24 @@ from __future__ import annotations import asyncio import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock, call, patch from uuid import uuid4 import pytest from sqlalchemy import delete, select -from app.application.database import DatabaseWorkerOverloadedError +from app.application.database import ( + DatabaseWorkerClosedError, + DatabaseWorkerOverloadedError, +) from app.application.messaging.chat import AgentChatPersistenceService, AgentChatService +from app.api.endpoints.agent import save_agent_chat_display from app.db.models.agentchat import AgentChat from app.db.oper.agentchat import AgentChatOper from app.db.session import SessionFactory, async_session_scope +from app.db.uow import run_sync_transaction +from app.schemas.agent import AgentChatDisplaySaveRequest from app.db.worker import DatabaseWorker @@ -65,8 +73,9 @@ async def test_agent_chat_persistence_runs_sync_repository_inside_worker() -> No executor = _Executor() repository = _Repository() service = AgentChatPersistenceService( - repository=lambda: repository, + repository=lambda _session: repository, async_executor=executor, + sync_transaction=lambda operation: operation(object()), ) caller_thread_id = threading.get_ident() @@ -111,8 +120,9 @@ async def test_agent_chat_persistence_propagates_worker_failure() -> None: raise RuntimeError("worker failed") service = AgentChatPersistenceService( - repository=_Repository, + repository=lambda _session: _Repository(), async_executor=FailingExecutor(), + sync_transaction=lambda operation: operation(object()), ) with pytest.raises(RuntimeError, match="worker failed"): @@ -123,6 +133,83 @@ async def test_agent_chat_persistence_propagates_worker_failure() -> None: ) +@pytest.mark.asyncio +async def test_agent_chat_persistence_pending_metric_uses_deltas() -> None: + """pending 是 UpDownCounter，准入和释放必须分别记录增减量。""" + service = AgentChatPersistenceService( + repository=lambda _session: _Repository(), + async_executor=_Executor(), + sync_transaction=lambda operation: operation(object()), + ) + with patch("app.application.messaging.chat.record_metric") as record_metric: + await service.async_save_agent_messages( + session_id="metric-session", + user_id="1", + messages=[], + ) + record_metric.assert_has_calls( + [ + call("agent.chat.persistence.pending", 1), + call("agent.chat.persistence.pending", -1), + ] + ) + + +@pytest.mark.asyncio +async def test_authoritative_display_save_propagates_worker_overload() -> None: + """权威 PUT 保存不能把 worker 背压吞成成功或普通业务失败。""" + repository = AsyncMock() + repository.async_get.return_value = None + service = AgentChatService(repository=repository) + + class OverloadedPersistence: + async def async_save_display_messages(self, **_kwargs): + raise DatabaseWorkerOverloadedError("busy") + + with pytest.raises(DatabaseWorkerOverloadedError, match="busy"): + await save_agent_chat_display( + session_id="overloaded-session", + payload=AgentChatDisplaySaveRequest(messages=[]), + current_user=SimpleNamespace(id=1, name="admin", is_superuser=True), + service=service, + persistence=OverloadedPersistence(), + ) + + +@pytest.mark.asyncio +async def test_agent_chat_persistence_rolls_back_compound_write(monkeypatch) -> None: + """复合写入中途失败时，创建或更新不能留下半成品。""" + worker = DatabaseWorker(max_workers=1, capacity=4) + await worker.start() + session_id = f"worker-rollback-{uuid4().hex}" + persistence = AgentChatPersistenceService( + repository=lambda session: AgentChatOper(session), + async_executor=worker, + sync_transaction=run_sync_transaction, + ) + original = AgentChatOper.save_display_messages + + def fail_after_stage(self, *args, **kwargs): + original(self, *args, **kwargs) + raise RuntimeError("display snapshot failed") + + monkeypatch.setattr(AgentChatOper, "save_display_messages", fail_after_stage) + try: + with pytest.raises(RuntimeError, match="display snapshot failed"): + await persistence.async_append_display_messages( + session_id=session_id, + user_id="rollback-user", + messages=[{"role": "user", "content": "not committed"}], + ) + async with async_session_scope() as session: + result = await session.execute( + select(AgentChat).where(AgentChat.session_id == session_id) + ) + assert result.scalars().first() is None + finally: + await worker.shutdown() + + @pytest.mark.asyncio async def test_agent_chat_persistence_bounds_session_waiters_and_releases_cancelled() -> None: """同会话锁等待受总量限制，取消等待不会遗留 admission。""" @@ -139,9 +226,11 @@ async def test_agent_chat_persistence_bounds_session_waiters_and_releases_cancel executor = BlockingExecutor() service = AgentChatPersistenceService( - repository=_Repository, + repository=lambda _session: _Repository(), async_executor=executor, + sync_transaction=lambda operation: operation(object()), capacity=2, + session_capacity=2, ) first = asyncio.create_task( service.async_save_agent_messages( @@ -177,6 +266,95 @@ async def test_agent_chat_persistence_bounds_session_waiters_and_releases_cancel assert service._pending_writes == 0 +@pytest.mark.asyncio +async def test_agent_chat_persistence_session_admission_is_fair() -> None: + """热点会话的锁等待不能占满全局容量并拒绝其他会话。""" + + class BlockingExecutor: + def __init__(self) -> None: + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def run(self, operation): + self.started.set() + await self.release.wait() + return operation() + + executor = BlockingExecutor() + service = AgentChatPersistenceService( + repository=lambda _session: _Repository(), + async_executor=executor, + sync_transaction=lambda operation: operation(object()), + capacity=4, + session_capacity=2, + ) + first = asyncio.create_task( + service.async_save_agent_messages( + session_id="hot-session", user_id="1", messages=[] + ) + ) + await executor.started.wait() + second = asyncio.create_task( + service.async_save_agent_messages( + session_id="hot-session", user_id="1", messages=[] + ) + ) + await asyncio.sleep(0) + with pytest.raises(DatabaseWorkerOverloadedError): + await service.async_save_agent_messages( + session_id="hot-session", user_id="1", messages=[] + ) + other = asyncio.create_task( + service.async_save_agent_messages( + session_id="other-session", user_id="1", messages=[] + ) + ) + await asyncio.sleep(0) + assert not other.done() + executor.release.set() + await first + await second + await other + + +@pytest.mark.asyncio +async def test_agent_chat_persistence_shutdown_drains_active_writes() -> None: + """关闭持久化端口时拒绝新写入并等待现有会话写入收口。""" + + class BlockingExecutor: + def __init__(self) -> None: + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def run(self, operation): + self.started.set() + await self.release.wait() + return operation() + + executor = BlockingExecutor() + service = AgentChatPersistenceService( + repository=lambda _session: _Repository(), + async_executor=executor, + sync_transaction=lambda operation: operation(object()), + ) + write = asyncio.create_task( + service.async_save_agent_messages( + session_id="shutdown-session", user_id="1", messages=[] + ) + ) + await executor.started.wait() + shutdown = asyncio.create_task(service.shutdown()) + await asyncio.sleep(0) + assert not shutdown.done() + with pytest.raises(DatabaseWorkerClosedError): + await service.async_save_agent_messages( + session_id="new-session", user_id="1", messages=[] + ) + executor.release.set() + await write + await shutdown + + @pytest.mark.asyncio async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() -> None: """真实 AgentChat Oper 经 worker 写入后可被 native async 查询恢复。""" @@ -184,8 +362,9 @@ async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() await worker.start() session_id = f"worker-{uuid4().hex}" persistence = AgentChatPersistenceService( - repository=AgentChatOper, + repository=lambda session: AgentChatOper(session), async_executor=worker, + sync_transaction=run_sync_transaction, ) query = AgentChatService(repository=AgentChatOper()) @@ -221,8 +400,9 @@ async def test_agent_chat_persistence_serializes_same_session_writes() -> None: session_id = f"worker-race-{uuid4().hex}" existing_session_id = f"worker-race-existing-{uuid4().hex}" persistence = AgentChatPersistenceService( - repository=AgentChatOper, + repository=lambda session: AgentChatOper(session), async_executor=worker, + sync_transaction=run_sync_transaction, ) async def append(content: str) -> None: diff --git a/tests/test_agent_lifecycle.py b/tests/test_agent_lifecycle.py index b610bb032..f20b0ea46 100644 --- a/tests/test_agent_lifecycle.py +++ b/tests/test_agent_lifecycle.py @@ -4,6 +4,10 @@ from unittest.mock import AsyncMock, MagicMock import pytest import app.agent.orchestrator as agent_module +from app.application.messaging.agent import ( + create_web_agent_background_task, + shutdown_web_agent_background_tasks, +) from app.agent.orchestrator import ( AGENT_SESSION_QUEUE_MAX_SIZE, AgentManager, @@ -14,6 +18,28 @@ from app.agent.memory import MemoryManager from app.startup import agent_initializer, modules_initializer +@pytest.mark.anyio +async def test_web_agent_background_tasks_are_cancelled_and_drained() -> None: + """Web Agent 任务关闭后不得继续占用循环或提交晚到的快照。""" + started = asyncio.Event() + finished = asyncio.Event() + + async def blocked_task() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + finished.set() + + task = create_web_agent_background_task(blocked_task()) + await started.wait() + await shutdown_web_agent_background_tasks() + + assert task.done() + assert task.cancelled() + assert finished.is_set() + + @pytest.mark.anyio async def test_agent_entrypoint_initializes_on_calling_loop(monkeypatch) -> None: """Agent 启动入口必须在应用主循环完成初始化。""" diff --git a/tests/test_host_runtime_context.py b/tests/test_host_runtime_context.py index 9a309e8a2..c164e74b5 100644 --- a/tests/test_host_runtime_context.py +++ b/tests/test_host_runtime_context.py @@ -13,6 +13,7 @@ from app.api.context import ( get_agent_chat_repository, get_agent_chat_transaction, ) +from app.api.dependencies.agent import get_agent_chat_persistence from app.startup import lifecycle from app.startup.context import ( AgentChatRuntime, @@ -59,6 +60,10 @@ class _UnitOfWork: """模拟回滚。""" +class _AgentChatPersistence: + """提供 AgentChat 运行时所需的最小写端口。""" + + class _SyncUnitOfWork: """记录绑定会话的同步事务替身。""" @@ -119,6 +124,7 @@ def _runtime() -> HostRuntime: async_session=async_session, repository=_Repository, transaction=_UnitOfWork, + persistence=_AgentChatPersistence(), ), persistence=PersistenceRuntime( sync_session=sync_session, @@ -186,15 +192,19 @@ def test_fastapi_dependencies_use_fake_runtime_without_real_services() -> None: async def probe( repository=Depends(get_agent_chat_repository), unit_of_work=Depends(get_agent_chat_transaction), + persistence=Depends(get_agent_chat_persistence), ) -> dict[str, bool]: """返回两个类型化能力是否绑定同一请求会话。""" - return {"same_session": repository.session is unit_of_work.session} + return { + "same_session": repository.session is unit_of_work.session, + "has_persistence": persistence is app.state.host_runtime.agent_chat.persistence, + } with TestClient(app) as client: response = client.get("/probe") assert response.status_code == 200 - assert response.json() == {"same_session": True} + assert response.json() == {"same_session": True, "has_persistence": True} def test_official_api_dependencies_do_not_use_string_data_locator() -> None: diff --git a/tests/test_lifecycle_shutdown.py b/tests/test_lifecycle_shutdown.py index 23d9abb10..9acc566b6 100644 --- a/tests/test_lifecycle_shutdown.py +++ b/tests/test_lifecycle_shutdown.py @@ -570,6 +570,35 @@ def test_stop_modules_continues_after_internal_owner_failures(monkeypatch): _assert_completed_once(dependency) +def test_stop_modules_drains_web_agent_tasks_before_persistence(monkeypatch): + """关闭时必须先收口 Web Agent 后台任务，再关闭会话持久化端口。""" + order = [] + monkeypatch.setattr(modules_initializer, "stop_agent", AsyncMock()) + dependencies = _patch_module_shutdown_dependencies(monkeypatch) + monkeypatch.setattr( + modules_initializer, + "shutdown_web_agent_background_tasks", + AsyncMock(side_effect=lambda: order.append("web-agent")), + ) + persistence = MagicMock() + persistence.shutdown = AsyncMock(side_effect=lambda: order.append("persistence")) + monkeypatch.setattr( + modules_initializer, + "get_configured_agent_chat_persistence", + MagicMock(return_value=persistence), + ) + monkeypatch.setattr( + modules_initializer, + "stop_database_worker", + AsyncMock(side_effect=lambda: order.append("database")), + ) + monkeypatch.setattr(modules_initializer, "_database_worker", object()) + + asyncio.run(modules_initializer.stop_modules()) + + assert order == ["web-agent", "persistence", "database"] + + def _patch_module_shutdown_dependencies(monkeypatch) -> dict: """替换 stop_modules 的资源所有者，避免测试启动真实后台服务""" dependencies = {} From 1e2d0d3b072db454d002c1a52d6b39287671f70b Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 04:05:54 +0800 Subject: [PATCH 09/14] fix(agent): bound chat shutdown cancellation --- app/application/messaging/agent.py | 4 ++- app/application/messaging/chat.py | 10 +++++-- app/startup/modules_initializer.py | 5 ++++ tests/test_agent_chat_persistence.py | 42 ++++++++++++++++++++++++++++ tests/test_agent_lifecycle.py | 28 +++++++++++++++++++ tests/test_lifecycle_shutdown.py | 12 ++++++-- 6 files changed, 96 insertions(+), 5 deletions(-) diff --git a/app/application/messaging/agent.py b/app/application/messaging/agent.py index 8378c9770..6b2f5a3a1 100644 --- a/app/application/messaging/agent.py +++ b/app/application/messaging/agent.py @@ -197,7 +197,9 @@ async def shutdown_web_agent_background_tasks() -> None: for task in tasks: task.cancel() if tasks: - await asyncio.gather(*tasks, return_exceptions=True) + # asyncio.wait 不会因关闭阶段自身被取消而再次取消这些任务；仍在收尾的 + # Agent 任务会保留在注册表中，直到自己的数据库操作取得确定终态。 + await asyncio.wait(tasks) def register_channel_admin_resolver( diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index 9f509e295..bdb831348 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -389,6 +389,10 @@ class AgentChatPersistenceService: self._session_locks[session_id] = lock return lock + def begin_shutdown(self) -> None: + """停止接受新的 AgentChat 持久化任务。""" + self._closing = True + async def _run_write( self, session_id: str, @@ -436,11 +440,13 @@ class AgentChatPersistenceService: async def shutdown(self) -> None: """拒绝新写入并等待当前会话锁和 worker 操作取得终态。""" - self._closing = True + self.begin_shutdown() current = asyncio.current_task() tasks = tuple(task for task in self._active_tasks if task is not current) if tasks: - await asyncio.gather(*tasks, return_exceptions=True) + # wait 不会在生命周期超时时取消实际写入；外层可及时返回并保留 + # 数据库 worker owner，已开始的事务继续由 worker 收口。 + await asyncio.wait(tasks) async def async_append_display_messages( self, diff --git a/app/startup/modules_initializer.py b/app/startup/modules_initializer.py index c222acf04..56701b937 100644 --- a/app/startup/modules_initializer.py +++ b/app/startup/modules_initializer.py @@ -578,6 +578,11 @@ async def stop_modules(): await run_step("消息服务", stop_message) await run_step("Redis缓存连接", lambda: RedisHelper().close()) await run_step("异步Redis缓存连接", lambda: AsyncRedisHelper().close()) + # 先关闭持久化准入，取消中的 Web Agent finally 才会快速拒绝晚到的快照写入。 + await run_step( + "Agent会话持久化准入", + lambda: get_configured_agent_chat_persistence().begin_shutdown(), + ) await run_step("Web Agent后台任务", shutdown_web_agent_background_tasks) await run_step( "Agent会话持久化", diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index 396bb8e23..727997ff3 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -355,6 +355,48 @@ async def test_agent_chat_persistence_shutdown_drains_active_writes() -> None: await shutdown +@pytest.mark.asyncio +async def test_agent_chat_shutdown_timeout_keeps_worker_owner_until_write_finishes() -> None: + """持久化关闭超时时保留运行中的写入和数据库 worker owner。""" + started = threading.Event() + release = threading.Event() + + class BlockingRepository(_Repository): + def save_agent_messages(self, **kwargs): + started.set() + release.wait(1) + super().save_agent_messages(**kwargs) + + worker = DatabaseWorker(max_workers=1, capacity=1) + await worker.start() + service = AgentChatPersistenceService( + repository=lambda _session: BlockingRepository(), + async_executor=worker, + sync_transaction=lambda operation: operation(object()), + ) + write = asyncio.create_task( + service.async_save_agent_messages( + session_id="shutdown-timeout-session", + user_id="1", + messages=[], + ) + ) + assert await asyncio.to_thread(started.wait, 1) + shutdown = asyncio.create_task(service.shutdown()) + try: + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(shutdown, timeout=0.01) + assert service._closing is True + assert write.done() is False + assert worker._executor is not None + finally: + release.set() + await write + await worker.shutdown() + + assert worker._executor is None + + @pytest.mark.asyncio async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() -> None: """真实 AgentChat Oper 经 worker 写入后可被 native async 查询恢复。""" diff --git a/tests/test_agent_lifecycle.py b/tests/test_agent_lifecycle.py index f20b0ea46..666ed1121 100644 --- a/tests/test_agent_lifecycle.py +++ b/tests/test_agent_lifecycle.py @@ -40,6 +40,34 @@ async def test_web_agent_background_tasks_are_cancelled_and_drained() -> None: assert finished.is_set() +@pytest.mark.anyio +async def test_web_agent_shutdown_timeout_does_not_cancel_task_cleanup() -> None: + """关闭超时时保留仍在执行取消收尾的 Web Agent 任务。""" + started = asyncio.Event() + release = asyncio.Event() + + async def task_with_slow_cleanup() -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await release.wait() + raise + + task = create_web_agent_background_task(task_with_slow_cleanup()) + await started.wait() + shutdown = asyncio.create_task(shutdown_web_agent_background_tasks()) + await asyncio.sleep(0) + shutdown.cancel() + with pytest.raises(asyncio.CancelledError): + await shutdown + + assert task.done() is False + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + @pytest.mark.anyio async def test_agent_entrypoint_initializes_on_calling_loop(monkeypatch) -> None: """Agent 启动入口必须在应用主循环完成初始化。""" diff --git a/tests/test_lifecycle_shutdown.py b/tests/test_lifecycle_shutdown.py index 9acc566b6..56c7ce7ee 100644 --- a/tests/test_lifecycle_shutdown.py +++ b/tests/test_lifecycle_shutdown.py @@ -571,7 +571,7 @@ def test_stop_modules_continues_after_internal_owner_failures(monkeypatch): def test_stop_modules_drains_web_agent_tasks_before_persistence(monkeypatch): - """关闭时必须先收口 Web Agent 后台任务，再关闭会话持久化端口。""" + """关闭时先关闭持久化准入，再收口 Web Agent 和已有写入。""" order = [] monkeypatch.setattr(modules_initializer, "stop_agent", AsyncMock()) dependencies = _patch_module_shutdown_dependencies(monkeypatch) @@ -581,6 +581,9 @@ def test_stop_modules_drains_web_agent_tasks_before_persistence(monkeypatch): AsyncMock(side_effect=lambda: order.append("web-agent")), ) persistence = MagicMock() + persistence.begin_shutdown = MagicMock( + side_effect=lambda: order.append("persistence-admission") + ) persistence.shutdown = AsyncMock(side_effect=lambda: order.append("persistence")) monkeypatch.setattr( modules_initializer, @@ -596,7 +599,12 @@ def test_stop_modules_drains_web_agent_tasks_before_persistence(monkeypatch): asyncio.run(modules_initializer.stop_modules()) - assert order == ["web-agent", "persistence", "database"] + assert order == [ + "persistence-admission", + "web-agent", + "persistence", + "database", + ] def _patch_module_shutdown_dependencies(monkeypatch) -> dict: From e0f70fa920672364ed122e1af8b8c5a496e00288 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 05:19:07 +0800 Subject: [PATCH 10/14] fix(agent): close chat lifecycle safely --- app/api/dependencies/agent.py | 4 +- app/api/endpoints/agent.py | 13 ++++--- app/application/messaging/agent.py | 7 ++++ app/startup/lifecycle/__init__.py | 15 +++++++- app/startup/modules_initializer.py | 50 +++++++++++++++++-------- tests/test_agent_chat_persistence.py | 49 ++++++++++++++++++++++++- tests/test_lifecycle_shutdown.py | 55 ++++++++++++++++++++++++---- tests/test_web_agent_stream.py | 17 +++++++++ 8 files changed, 177 insertions(+), 33 deletions(-) diff --git a/app/api/dependencies/agent.py b/app/api/dependencies/agent.py index cb283f95d..9785d9e5c 100644 --- a/app/api/dependencies/agent.py +++ b/app/api/dependencies/agent.py @@ -17,7 +17,7 @@ from app.application.messaging.chat import ( AsyncUnitOfWork, ) from app.application.messaging.message import MessageQueryService -from app.startup.context import HostRuntime +from app.startup.context import AgentChatRuntime, HostRuntime def get_agent_chat_service( @@ -29,7 +29,7 @@ def get_agent_chat_service( def get_agent_chat_persistence( - runtime: HostRuntime = Depends(get_agent_chat_runtime), + runtime: AgentChatRuntime = Depends(get_agent_chat_runtime), ) -> AgentChatPersistenceService: """从类型化 Agent 运行时获取有界会话写入端口。""" return runtime.persistence diff --git a/app/api/endpoints/agent.py b/app/api/endpoints/agent.py index 522af2a44..c9f5aaa31 100644 --- a/app/api/endpoints/agent.py +++ b/app/api/endpoints/agent.py @@ -1937,10 +1937,13 @@ async def save_agent_chat_display( service=service, persistence=persistence, ) - chat = await service.get_accessible(session_id, current_user) + # 写入由独立 worker 事务完成，使用组合根登记的短会话服务复读，避免请求会话 + # 的 identity map 返回写入前的 ORM 快照。 + chat_service = get_configured_agent_chat_service() + chat = await chat_service.get_accessible(session_id, current_user) if not chat: return _SchemaResponse(success=False, message="会话保存失败") - return _SchemaResponse(success=True, data=service.to_summary(chat)) + return _SchemaResponse(success=True, data=chat_service.to_summary(chat)) @router.delete( @@ -2027,7 +2030,7 @@ async def _web_agent_stream_impl( """ prompt = payload.text.strip() if not isinstance(service, AgentChatService): - # 直接调用公开函数时不经过 FastAPI 依赖解析；生产路由总是传入运行时服务。 + # SSE 后台任务可能在请求依赖释放后继续运行，查询服务必须自行取得短会话。 service = get_configured_agent_chat_service() if not isinstance(persistence, AgentChatPersistenceService): # 直接调用公开函数时不经过 FastAPI 依赖解析；生产路由总是传入运行时端口。 @@ -2387,7 +2390,6 @@ async def web_agent_stream( payload: _SchemaAgentWebChatRequest, request: Request, current_user: ApiPrincipal = Depends(get_current_active_user), - service: AgentChatService = Depends(get_agent_chat_service), persistence: AgentChatPersistenceService = Depends(get_agent_chat_persistence), ) -> StreamingResponse: """Web 智能助手流式对话的稳定公开路由入口。""" @@ -2395,6 +2397,5 @@ async def web_agent_stream( payload, request, current_user, - service, - persistence, + persistence=persistence, ) diff --git a/app/application/messaging/agent.py b/app/application/messaging/agent.py index 6b2f5a3a1..d29134a6f 100644 --- a/app/application/messaging/agent.py +++ b/app/application/messaging/agent.py @@ -202,6 +202,13 @@ async def shutdown_web_agent_background_tasks() -> None: await asyncio.wait(tasks) +async def wait_web_agent_background_tasks() -> None: + """等待已登记的 Web Agent 任务完成取消后的最终收尾。""" + tasks = tuple(_WEB_AGENT_BACKGROUND_TASKS) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + def register_channel_admin_resolver( channel: Union[NotificationChannel, str], resolver: _ChannelAdminResolver, diff --git a/app/startup/lifecycle/__init__.py b/app/startup/lifecycle/__init__.py index c9be10aee..37fe6bb6b 100644 --- a/app/startup/lifecycle/__init__.py +++ b/app/startup/lifecycle/__init__.py @@ -103,10 +103,21 @@ async def run_shutdown_step( try: result = callback() if inspect.isawaitable(result): + task = asyncio.ensure_future(result) if timeout_seconds: - await asyncio.wait_for(result, timeout=timeout_seconds) + try: + await asyncio.wait_for( + asyncio.shield(task), timeout=timeout_seconds + ) + except asyncio.TimeoutError: + logger.error("关闭%s超时，等待其取消收口", name) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass else: - await result + await task except Exception as err: logger.error(f"关闭{name}失败：{err}") diff --git a/app/startup/modules_initializer.py b/app/startup/modules_initializer.py index 56701b937..e9274cf97 100644 --- a/app/startup/modules_initializer.py +++ b/app/startup/modules_initializer.py @@ -1,3 +1,4 @@ +import asyncio import inspect import sys from typing import Callable @@ -69,7 +70,10 @@ from app.application.messaging.chat import ( configure_agent_chat_service, get_configured_agent_chat_persistence, ) -from app.application.messaging.agent import shutdown_web_agent_background_tasks +from app.application.messaging.agent import ( + shutdown_web_agent_background_tasks, + wait_web_agent_background_tasks, +) from app.application.security.user import configure_user_lookups from app.application.security.auth import AuthService, configure_auth_service from app.application.security.passkeys import PasskeyService, configure_passkey_service @@ -559,14 +563,19 @@ async def stop_modules(): """ 服务关闭 """ - async def run_step(name: str, callback: Callable[[], object]) -> None: + async def run_step(name: str, callback: Callable[[], object]) -> bool: """单个模块资源关闭失败时继续执行后续阶段""" try: result = callback() if inspect.isawaitable(result): await result + return True + except asyncio.CancelledError: + logger.warning("关闭%s时收到取消请求，继续执行资源收口", name) + return False except Exception as err: logger.error(f"关闭{name}失败：{err}") + return True await run_step("AI智能体", stop_agent) await run_step("模块", lambda: ModuleManager().shutdown()) @@ -578,21 +587,32 @@ async def stop_modules(): await run_step("消息服务", stop_message) await run_step("Redis缓存连接", lambda: RedisHelper().close()) await run_step("异步Redis缓存连接", lambda: AsyncRedisHelper().close()) - # 先关闭持久化准入，取消中的 Web Agent finally 才会快速拒绝晚到的快照写入。 - await run_step( - "Agent会话持久化准入", - lambda: get_configured_agent_chat_persistence().begin_shutdown(), + # Web Agent 的取消 finally 可能还要写入最终展示快照，必须先完成任务收尾，再关闭写入准入。 + web_agent_drained = await run_step( + "Web Agent后台任务", shutdown_web_agent_background_tasks ) - await run_step("Web Agent后台任务", shutdown_web_agent_background_tasks) - await run_step( - "Agent会话持久化", - lambda: get_configured_agent_chat_persistence().shutdown(), - ) - await run_step("数据库任务", stop_database_worker) - if _database_worker is None: - await run_step("数据库连接", close_database) + if not web_agent_drained: + web_agent_drained = await run_step( + "Web Agent后台任务收尾", wait_web_agent_background_tasks + ) + if web_agent_drained: + await run_step( + "Agent会话持久化准入", + lambda: get_configured_agent_chat_persistence().begin_shutdown(), + ) + persistence_drained = await run_step( + "Agent会话持久化", + lambda: get_configured_agent_chat_persistence().shutdown(), + ) else: - logger.error("数据库任务未收敛，跳过数据库连接关闭以避免运行中事务使用已释放连接") + persistence_drained = False + logger.error("Web Agent任务未完成收尾，跳过持久化和数据库关闭以保护活动事务") + if persistence_drained: + await run_step("数据库任务", stop_database_worker) + if _database_worker is None: + await run_step("数据库连接", close_database) + else: + logger.error("数据库任务未收敛，跳过数据库连接关闭以避免运行中事务使用已释放连接") await run_step("前端服务", stop_frontend) await run_step("临时文件", clear_temp) diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index 727997ff3..513a21172 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -5,7 +5,7 @@ from __future__ import annotations import asyncio import threading from types import SimpleNamespace -from unittest.mock import AsyncMock, call, patch +from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import uuid4 import pytest @@ -176,6 +176,53 @@ async def test_authoritative_display_save_propagates_worker_overload() -> None: ) +@pytest.mark.asyncio +async def test_authoritative_display_save_reads_fresh_projection_after_worker_write( + monkeypatch, +) -> None: + """权威展示保存的响应必须读取 worker 提交后的最新投影。""" + existing_chat = SimpleNamespace( + user_id="1", + username="admin", + channel="WebAgent", + source="web-agent", + original_chat_id=None, + client_session_id="client-1", + ) + updated_chat = SimpleNamespace( + session_id="fresh-session", + message_count=2, + ) + request_service = SimpleNamespace( + get_accessible=AsyncMock(return_value=existing_chat), + get=AsyncMock(return_value=existing_chat), + ) + canonical_service = SimpleNamespace( + get_accessible=AsyncMock(return_value=updated_chat), + to_summary=MagicMock(return_value="fresh-summary"), + ) + persistence = SimpleNamespace(async_save_display_messages=AsyncMock()) + current_user = SimpleNamespace(id=1, name="admin", is_superuser=True) + monkeypatch.setattr( + "app.api.endpoints.agent.get_configured_agent_chat_service", + MagicMock(return_value=canonical_service), + ) + + response = await save_agent_chat_display( + session_id="fresh-session", + payload=AgentChatDisplaySaveRequest(messages=[]), + current_user=current_user, + service=request_service, + persistence=persistence, + ) + + assert response.success is True + assert response.data == "fresh-summary" + canonical_service.get_accessible.assert_awaited_once_with( + "fresh-session", current_user + ) + + @pytest.mark.asyncio async def test_agent_chat_persistence_rolls_back_compound_write(monkeypatch) -> None: """复合写入中途失败时，创建或更新不能留下半成品。""" diff --git a/tests/test_lifecycle_shutdown.py b/tests/test_lifecycle_shutdown.py index 56c7ce7ee..6c9996f9f 100644 --- a/tests/test_lifecycle_shutdown.py +++ b/tests/test_lifecycle_shutdown.py @@ -571,7 +571,7 @@ def test_stop_modules_continues_after_internal_owner_failures(monkeypatch): def test_stop_modules_drains_web_agent_tasks_before_persistence(monkeypatch): - """关闭时先关闭持久化准入，再收口 Web Agent 和已有写入。""" + """关闭时先收口 Web Agent，再关闭持久化准入和数据库任务。""" order = [] monkeypatch.setattr(modules_initializer, "stop_agent", AsyncMock()) dependencies = _patch_module_shutdown_dependencies(monkeypatch) @@ -599,12 +599,53 @@ def test_stop_modules_drains_web_agent_tasks_before_persistence(monkeypatch): asyncio.run(modules_initializer.stop_modules()) - assert order == [ - "persistence-admission", - "web-agent", - "persistence", - "database", - ] + assert order == ["web-agent", "persistence-admission", "persistence", "database"] + + +@pytest.mark.asyncio +async def test_shutdown_timeout_does_not_skip_database_worker_cleanup(monkeypatch): + """模块关闭超时取消当前步骤后仍应继续收口数据库 worker。""" + started = asyncio.Event() + + async def blocked_web_agent_shutdown(): + started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(modules_initializer, "stop_agent", AsyncMock()) + _patch_module_shutdown_dependencies(monkeypatch) + monkeypatch.setattr( + modules_initializer, + "shutdown_web_agent_background_tasks", + blocked_web_agent_shutdown, + ) + monkeypatch.setattr( + modules_initializer, + "wait_web_agent_background_tasks", + AsyncMock(), + ) + persistence = MagicMock() + persistence.begin_shutdown = MagicMock() + persistence.shutdown = AsyncMock() + monkeypatch.setattr( + modules_initializer, + "get_configured_agent_chat_persistence", + MagicMock(return_value=persistence), + ) + stop_database_worker = AsyncMock() + monkeypatch.setattr(modules_initializer, "stop_database_worker", stop_database_worker) + monkeypatch.setattr(modules_initializer, "_database_worker", object()) + + shutdown = asyncio.create_task( + lifecycle.run_shutdown_step( + "模块服务", + modules_initializer.stop_modules, + timeout_seconds=0.01, + ) + ) + await started.wait() + await shutdown + + stop_database_worker.assert_awaited_once_with() def _patch_module_shutdown_dependencies(monkeypatch) -> dict: diff --git a/tests/test_web_agent_stream.py b/tests/test_web_agent_stream.py index e13dc0ca9..97ceca095 100644 --- a/tests/test_web_agent_stream.py +++ b/tests/test_web_agent_stream.py @@ -355,6 +355,23 @@ def test_web_agent_stream_returns_error_for_unknown_command(): handle_message.assert_not_called() +def test_web_agent_stream_does_not_bind_request_scoped_chat_service(): + """流式路由不能把请求级 Agent 会话服务捕获到后台任务。""" + from app.api.dependencies.agent import get_agent_chat_service + from app.api.endpoints import agent as agent_endpoint + + route = next( + route + for route in agent_endpoint.router.routes + if getattr(route, "name", None) == "web_agent_stream" + ) + + assert all( + dependency.call is not get_agent_chat_service + for dependency in route.dependant.dependencies + ) + + def test_build_web_agent_message_update_event_converts_buttons(): """WebAgent 编辑消息应转换为可原地更新卡片的事件。""" event = build_web_agent_message_update_event( From cc8e4bc510717b506c827161d7d11c3eff4b7347 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 05:24:56 +0800 Subject: [PATCH 11/14] fix(agent): map closed database worker to 503 --- app/api/deps.py | 2 -- app/factory.py | 13 ++++++++++--- tests/test_api_response.py | 30 +++++++++++++++++++++++++++++- 3 files changed, 39 insertions(+), 6 deletions(-) diff --git a/app/api/deps.py b/app/api/deps.py index 6073df004..9007667af 100644 --- a/app/api/deps.py +++ b/app/api/deps.py @@ -5,7 +5,6 @@ """ from app.api.dependencies.agent import ( - get_agent_chat_persistence, get_agent_chat_service, get_message_query_service, ) @@ -53,7 +52,6 @@ from app.api.dependencies.workflow import ( # 兼容聚合入口只显式列出既有 FastAPI 依赖，不向插件制造新的动态导出规则。 __all__ = [ - "get_agent_chat_persistence", "get_agent_chat_service", "get_auth_service", "get_current_active_manage_user", diff --git a/app/factory.py b/app/factory.py index 8c947c741..7db6a5e82 100644 --- a/app/factory.py +++ b/app/factory.py @@ -14,7 +14,10 @@ from app.adapters.observability.otel import build_observation_port from app.adapters.web.plugin.routes import FastAPIDynamicRouteRegistry from app.adapters.web.health import install_health_routes from app.application.plugin.routes import configure_plugin_routes -from app.application.database import DatabaseWorkerOverloadedError +from app.application.database import ( + DatabaseWorkerClosedError, + DatabaseWorkerOverloadedError, +) from app.adapters.web.security.access import ( configure_token_codec, verify_apikey, @@ -234,9 +237,9 @@ async def localized_http_exception_handler( async def database_worker_overloaded_handler( request: Request, - _exc: DatabaseWorkerOverloadedError, + _exc: DatabaseWorkerClosedError | DatabaseWorkerOverloadedError, ) -> JSONResponse: - """将数据库短事务背压映射为可重试的 503 响应。""" + """将数据库 worker 暂不可用映射为可重试的 503 响应。""" return await localized_http_exception_handler( request, HTTPException( @@ -335,6 +338,10 @@ def create_app() -> FastAPI: DatabaseWorkerOverloadedError, database_worker_overloaded_handler, ) + _app.add_exception_handler( + DatabaseWorkerClosedError, + database_worker_overloaded_handler, + ) _app.add_exception_handler( RequestValidationError, localized_validation_exception_handler, diff --git a/tests/test_api_response.py b/tests/test_api_response.py index c4419d92e..65cac63c4 100644 --- a/tests/test_api_response.py +++ b/tests/test_api_response.py @@ -23,7 +23,10 @@ from app.factory import ( localized_unhandled_exception_handler, localized_validation_exception_handler, ) -from app.application.database import DatabaseWorkerOverloadedError +from app.application.database import ( + DatabaseWorkerClosedError, + DatabaseWorkerOverloadedError, +) from app.runtime.localization import LocaleHelper from app.runtime.config import settings from app.schemas.common import JsonData @@ -67,6 +70,10 @@ def api_app() -> FastAPI: DatabaseWorkerOverloadedError, database_worker_overloaded_handler, ) + app.add_exception_handler( + DatabaseWorkerClosedError, + database_worker_overloaded_handler, + ) from fastapi.exceptions import RequestValidationError app.add_exception_handler( @@ -125,6 +132,11 @@ def api_app() -> FastAPI: """模拟数据库短事务容量耗尽。""" raise DatabaseWorkerOverloadedError("worker full") + @app.get("/database-closed") + async def get_database_closed() -> None: + """模拟数据库 worker 在关闭态拒绝新任务。""" + raise DatabaseWorkerClosedError("worker closed") + @app.get("/native", response_model=None) async def get_native_response() -> dict[str, bool]: """返回显式旁路的原生 JSON 协议。""" @@ -219,6 +231,22 @@ async def test_database_worker_overload_is_retryable_service_unavailable( } +async def test_database_worker_closed_is_retryable_service_unavailable( + api_app: FastAPI, +): + """数据库 worker 关闭态应返回 503，而不是落入通用 500。""" + async with make_client(api_app) as client: + response = await client.get("/database-closed") + + assert response.status_code == 503 + assert response.headers["retry-after"] == "1" + assert response.json() == { + "success": False, + "message": "服务当前繁忙，请稍后重试", + "data": None, + } + + @pytest.mark.parametrize( "path", [ From 50e3b4366066928f958f2a31cb47d7cb7bed9e76 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 05:27:32 +0800 Subject: [PATCH 12/14] test(agent): cover closed database worker registration --- tests/test_api_response.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/test_api_response.py b/tests/test_api_response.py index 65cac63c4..78a3ee335 100644 --- a/tests/test_api_response.py +++ b/tests/test_api_response.py @@ -247,6 +247,17 @@ async def test_database_worker_closed_is_retryable_service_unavailable( } +def test_create_app_registers_closed_database_worker_handler() -> None: + """生产组合根必须为 worker 关闭态登记 503 处理器。""" + from app.factory import create_app + + app = create_app() + + assert app.exception_handlers[DatabaseWorkerClosedError] is ( + database_worker_overloaded_handler + ) + + @pytest.mark.parametrize( "path", [ From 4a891f7f69cfe49aaa472dc4d882a8a972c02538 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 11:45:16 +0800 Subject: [PATCH 13/14] test: refresh agent chat architecture baseline --- tests/fixtures/architecture/dependency-baseline.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/fixtures/architecture/dependency-baseline.json b/tests/fixtures/architecture/dependency-baseline.json index 06210fa82..ce2f2bbf3 100644 --- a/tests/fixtures/architecture/dependency-baseline.json +++ b/tests/fixtures/architecture/dependency-baseline.json @@ -13,8 +13,8 @@ "runtime_to_db": [], "workflow_to_db": [] }, - "edge_count": 6443, - "edge_sha256": "da29d186d8ca672b117535abf9976a46123563a8a12532689496ce7dd56008f5", + "edge_count": 6444, + "edge_sha256": "b5db7b31c7ea4dd7311fcd9e11a49eb32feb6939f896ed85703b3573408752a8", "edges": [ "app -> app.runtime", "app -> app.runtime.compat", From bbf62a0b40716f478184a21fa514e8224be84854 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 12:00:54 +0800 Subject: [PATCH 14/14] fix(lifecycle): bound shutdown cancellation --- app/startup/lifecycle/__init__.py | 19 +++++++++++----- tests/test_lifecycle_shutdown.py | 38 +++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/app/startup/lifecycle/__init__.py b/app/startup/lifecycle/__init__.py index 37fe6bb6b..de01e19a8 100644 --- a/app/startup/lifecycle/__init__.py +++ b/app/startup/lifecycle/__init__.py @@ -99,23 +99,30 @@ async def run_shutdown_step( callback: Callable[[], object], timeout_seconds: float | None = None, ) -> None: - """隔离单个关闭阶段的异常，确保后续资源仍有机会释放""" + """在有限预算内执行关闭阶段，并保留未收敛任务的资源所有权。""" try: result = callback() if inspect.isawaitable(result): task = asyncio.ensure_future(result) + + def _consume_shutdown_result(done: asyncio.Future) -> None: + """消费延迟收敛任务的最终异常，避免事件循环产生未取回异常。""" + try: + done.result() + except asyncio.CancelledError: + pass + except Exception as err: + logger.error(f"关闭{name}最终收尾失败：{err}") + + task.add_done_callback(_consume_shutdown_result) if timeout_seconds: try: await asyncio.wait_for( asyncio.shield(task), timeout=timeout_seconds ) except asyncio.TimeoutError: - logger.error("关闭%s超时，等待其取消收口", name) + logger.error("关闭%s超时，已请求取消并保留未收敛任务", name) task.cancel() - try: - await task - except asyncio.CancelledError: - pass else: await task except Exception as err: diff --git a/tests/test_lifecycle_shutdown.py b/tests/test_lifecycle_shutdown.py index 6c9996f9f..8f33e084d 100644 --- a/tests/test_lifecycle_shutdown.py +++ b/tests/test_lifecycle_shutdown.py @@ -648,6 +648,44 @@ async def test_shutdown_timeout_does_not_skip_database_worker_cleanup(monkeypatc stop_database_worker.assert_awaited_once_with() +@pytest.mark.asyncio +async def test_shutdown_timeout_has_hard_bound_for_nonconverging_cleanup() -> None: + """关闭收尾不响应取消时，生命周期调用仍必须在预算内返回。""" + started = asyncio.Event() + cancel_requested = asyncio.Event() + release = asyncio.Event() + settled = asyncio.Event() + + async def nonconverging_shutdown() -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancel_requested.set() + await release.wait() + settled.set() + raise + + started_at = asyncio.get_running_loop().time() + shutdown = asyncio.create_task( + lifecycle.run_shutdown_step( + "不可收敛阶段", + nonconverging_shutdown, + timeout_seconds=0.01, + ) + ) + await started.wait() + await shutdown + + elapsed = asyncio.get_running_loop().time() - started_at + assert elapsed < 0.2 + await asyncio.wait_for(cancel_requested.wait(), timeout=0.2) + assert not settled.is_set() + + release.set() + await asyncio.wait_for(settled.wait(), timeout=0.2) + + def _patch_module_shutdown_dependencies(monkeypatch) -> dict: """替换 stop_modules 的资源所有者，避免测试启动真实后台服务""" dependencies = {}