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()