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: