diff --git a/app/agent/orchestrator.py b/app/agent/orchestrator.py index 9fd8481c3..fff264b9c 100644 --- a/app/agent/orchestrator.py +++ b/app/agent/orchestrator.py @@ -3034,6 +3034,10 @@ class AgentManager: # 等待消息,超时后自动退出worker task = await asyncio.wait_for(queue.get(), timeout=60.0) except asyncio.TimeoutError: + # 超时回调与入队可能在同一轮事件循环就绪;已有消息时继续消费, + # 避免旧 worker 退出后留下没有消费者的非空队列。 + if not queue.empty(): + continue # 队列空闲超时,退出worker logger.debug(f"会话 {session_id} 的消息队列空闲,worker退出") break @@ -3306,22 +3310,17 @@ class AgentManager: self._session_cleanup_pending.add(session_id) self._session_cancel_requested.add(session_id) worker.cancel() - stopped_cleanly = await self._wait_for_worker_shutdown( - session_id, - worker, - reason="clear_session", - ) - if not stopped_cleanly: - queue = self._session_queues.pop(session_id, None) - if queue: - self._discard_queued_messages(queue) - worker.add_done_callback( - lambda done: self._schedule_deferred_session_cleanup( - session_id, - user_id, - worker, - ) + try: + stopped_cleanly = await self._wait_for_worker_shutdown( + session_id, + worker, + reason="clear_session", ) + except asyncio.CancelledError: + self._defer_session_cleanup(session_id, user_id, worker) + raise + if not stopped_cleanly: + self._defer_session_cleanup(session_id, user_id, worker) return if self._session_workers.get(session_id) is worker: self._session_workers.pop(session_id, None) # noqa @@ -3346,14 +3345,36 @@ class AgentManager: memory_manager.clear_memory(session_id, user_id) logger.info(f"会话 {session_id} 的记忆已清空") + def _defer_session_cleanup( + self, + session_id: str, + user_id: str, + worker: asyncio.Task, + ) -> None: + """把中断或超时的清理转交给 worker 终态回调。""" + queue = self._session_queues.pop(session_id, None) + if queue: + self._discard_queued_messages(queue) + self._session_shutdown_pending[session_id] = worker + worker.add_done_callback( + lambda done: self._schedule_deferred_session_cleanup( + session_id, + user_id, + done, + ) + ) + def _schedule_deferred_session_cleanup( self, session_id: str, user_id: str, worker: asyncio.Task, ) -> None: - """worker 超时后延迟释放会话资源,避免与仍在运行的 Agent 竞态。""" - if session_id in self._session_shutdown_pending: + """worker 取得终态后释放会话资源,避免与仍在运行的 Agent 竞态。""" + if self._session_shutdown_pending.get(session_id) is worker: + existing = self._session_deferred_cleanup_tasks.get(session_id) + if existing is not None and not existing.done(): + return cleanup_task = asyncio.create_task( self._finish_deferred_session_cleanup( session_id=session_id, @@ -3438,6 +3459,9 @@ class AgentManager: ) return True except asyncio.CancelledError: + current = asyncio.current_task() + if worker.cancelled() and (current is None or not current.cancelling()): + return True raise except asyncio.TimeoutError: self._session_shutdown_pending[session_id] = worker diff --git a/app/application/agenttask.py b/app/application/agenttask.py index 497a6827d..eeb5f4f1a 100644 --- a/app/application/agenttask.py +++ b/app/application/agenttask.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import threading from collections.abc import Callable from dataclasses import dataclass from typing import Protocol, TypeVar @@ -155,6 +156,7 @@ class AgentTaskExecutionService: """认领一次执行;取消发生在提交后时先补偿收口再传播取消。""" run_id = uuid4().hex + run_created = threading.Event() def transaction(session: object) -> AgentTaskClaim: repository = self._repository(session) @@ -173,6 +175,7 @@ class AgentTaskExecutionService: else "Agent 定时任务当前不可执行" ), ) + run_created.set() return AgentTaskClaim(run=self._snapshot(run)) async def claim_to_terminal() -> AgentTaskClaim: @@ -189,6 +192,9 @@ class AgentTaskExecutionService: try: return await claim_task except asyncio.CancelledError as cancellation: + # 纯容量拒绝尚未进入事务,不存在需要等待数据库容量的补偿对象。 + if not run_created.is_set(): + raise cancellation finalize_task = asyncio.create_task(self._finalize( run_id=run_id, task_id=task_id, diff --git a/tests/test_agent_lifecycle.py b/tests/test_agent_lifecycle.py index a41e1b208..bf2d5cdef 100644 --- a/tests/test_agent_lifecycle.py +++ b/tests/test_agent_lifecycle.py @@ -697,3 +697,209 @@ async def test_background_prompt_is_owned_and_cancelled_by_manager_close( assert manager._session_queues == {} assert manager._session_workers == {} assert manager.active_agents == {} + + +@pytest.mark.anyio +async def test_stop_current_task_handles_worker_cancelled_before_first_run() -> None: + """worker 尚未首次运行时停止也必须完成收口。""" + manager = AgentManager() + session_id = "stop-before-worker-start" + worker = None + await manager.initialize() + + try: + await manager.process_message(session_id, "1", "message") + worker = manager._session_workers[session_id] + assert worker.done() is False + + assert await manager.stop_current_task(session_id) is True + assert worker.done() is True + assert session_id not in manager._session_workers + assert session_id not in manager._session_queues + finally: + if worker is not None: + if not worker.done(): + worker.cancel() + try: + await worker + except BaseException: + pass + if manager._session_workers.get(session_id) is worker: + manager._session_workers.pop(session_id, None) + if manager._accepting_tasks: + await manager.close() + + +@pytest.mark.anyio +async def test_clear_session_cancellation_does_not_stick_cleanup_pending( + monkeypatch, +) -> None: + """clear_session 被调用方取消后必须能重试或已转交延迟清理。""" + manager = AgentManager() + memory_manager = MemoryManager() + memory_manager.clear_memory = MagicMock() + monkeypatch.setattr(agent_module, "memory_manager", memory_manager) + session_id = "clear-caller-cancelled" + started = asyncio.Event() + cancellation_seen = asyncio.Event() + release = asyncio.Event() + execution = None + clear_request = None + + class BlockingAgent: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + async def process(self, _message, **_kwargs): + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancellation_seen.set() + await release.wait() + raise + + async def cleanup(self): + return True + + def get_session_status(self): + return {} + + await manager.initialize() + try: + execution = asyncio.create_task( + manager.process_message( + session_id, + "1", + "message", + agent_factory=BlockingAgent, + wait_for_completion=True, + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + clear_request = asyncio.create_task( + manager.clear_session(session_id, "1") + ) + await asyncio.wait_for(cancellation_seen.wait(), timeout=1) + assert clear_request.done() is False + clear_request.cancel() + with pytest.raises(asyncio.CancelledError): + await clear_request + + release.set() + with pytest.raises(asyncio.CancelledError): + await execution + + # 取消安全的延迟清理可能异步完成;若未发生转交,重试仍必须完成清理。 + await asyncio.wait_for( + manager.clear_session(session_id, "1"), + timeout=1, + ) + for _ in range(100): + if ( + session_id not in manager._session_cleanup_pending + and session_id not in manager.active_agents + ): + break + await asyncio.sleep(0) + + assert session_id not in manager._session_cleanup_pending + assert session_id not in manager._session_shutdown_pending + assert session_id not in manager._session_deferred_cleanup_tasks + assert session_id not in manager._session_workers + assert session_id not in manager._session_queues + assert session_id not in manager.active_agents + assert memory_manager.clear_memory.call_count == 1 + finally: + release.set() + if clear_request is not None and not clear_request.done(): + clear_request.cancel() + try: + await clear_request + except BaseException: + pass + if execution is not None and not execution.done(): + execution.cancel() + try: + await execution + except BaseException: + pass + for deferred in list(manager._session_deferred_cleanup_tasks.values()): + if not deferred.done(): + deferred.cancel() + try: + await deferred + except BaseException: + pass + if manager._accepting_tasks: + await manager.close() + + +@pytest.mark.anyio +async def test_session_worker_restarts_after_idle_timeout_races_with_full_enqueue( + monkeypatch, +) -> None: + """空闲退出与满队列入队交错时必须保留会话消费者。""" + manager = AgentManager() + memory_manager = MemoryManager() + monkeypatch.setattr(agent_module, "memory_manager", memory_manager) + session_id = "idle-timeout-full-queue" + processed = [] + first_processed = asyncio.Event() + idle_waiting = asyncio.Event() + release_timeout = asyncio.Event() + timeout_intercepted = False + real_wait_for = asyncio.wait_for + + async def process(task): + processed.append(task.message) + if task.message == "initial": + first_processed.set() + return task.message + + async def controlled_wait_for(awaitable, timeout): + nonlocal timeout_intercepted + if timeout == 60.0 and processed and not timeout_intercepted: + timeout_intercepted = True + idle_waiting.set() + await release_timeout.wait() + awaitable.close() + raise asyncio.TimeoutError + return await real_wait_for(awaitable, timeout) + + monkeypatch.setattr(agent_module.asyncio, "wait_for", controlled_wait_for) + await manager.initialize() + try: + manager._process_message_internal = process + await manager.process_message(session_id, "1", "initial") + await real_wait_for(first_processed.wait(), timeout=1) + await real_wait_for(idle_waiting.wait(), timeout=1) + + for index in range(AGENT_SESSION_QUEUE_MAX_SIZE): + await manager.process_message( + session_id, + "1", + f"queued-{index}", + ) + + assert manager._session_queues[session_id].full() + release_timeout.set() + + async def all_messages_processed() -> None: + while len(processed) < AGENT_SESSION_QUEUE_MAX_SIZE + 1: + await asyncio.sleep(0) + + await real_wait_for(all_messages_processed(), timeout=1) + assert processed == [ + "initial", + *[f"queued-{index}" for index in range(AGENT_SESSION_QUEUE_MAX_SIZE)], + ] + worker = manager._session_workers.get(session_id) + assert worker is not None + assert worker.done() is False + finally: + release_timeout.set() + if manager._accepting_tasks: + await manager.clear_session(session_id, "1") + await manager.close() diff --git a/tests/test_agent_task_execution_service.py b/tests/test_agent_task_execution_service.py index 0fd72ed49..3a9a61684 100644 --- a/tests/test_agent_task_execution_service.py +++ b/tests/test_agent_task_execution_service.py @@ -8,11 +8,15 @@ from uuid import uuid4 import pytest from app.application.agenttask import AgentTaskExecutionService +from app.application.database import AsyncDatabaseExecutor from app.db.adapters.transaction import TransactionalWriteRunner from app.db.oper.agenttask import AgentTaskOper from app.db.session import SessionFactory, async_session_scope from app.db.worker import DatabaseWorker -from app.schemas.exception import DatabaseWorkerClosedError +from app.schemas.exception import ( + DatabaseWorkerClosedError, + DatabaseWorkerOverloadedError, +) def _add_task(prefix: str, *, trigger_type: str = "cron"): @@ -34,7 +38,7 @@ def _add_task(prefix: str, *, trigger_type: str = "cron"): def _build_service( - worker: DatabaseWorker, + worker: AsyncDatabaseExecutor, repository: Callable[[object], object] | None = None, ) -> AgentTaskExecutionService: """按生产事务和 worker 边界构造独立服务。""" @@ -91,6 +95,41 @@ async def test_cancelled_queued_claim_does_not_create_run() -> None: await worker.shutdown() +@pytest.mark.anyio +async def test_cancelled_overloaded_claim_does_not_start_compensation() -> None: + """认领尚未获得 admission 时取消,不得启动不存在运行的补偿收口。""" + + class OverloadedExecutor: + """在认领取得 worker admission 前稳定制造取消窗口。""" + + def __init__(self) -> None: + self.calls = 0 + self.first_call = asyncio.Event() + + async def run(self, _operation): + self.calls += 1 + self.first_call.set() + if self.calls > 1: + raise AssertionError("未获 admission 的认领不应启动终态补偿") + raise DatabaseWorkerOverloadedError("worker full") + + executor = OverloadedExecutor() + task = _add_task("overload-cancel") + service = _build_service(executor) + claim = asyncio.create_task(service.claim(task.id)) + await executor.first_call.wait() + claim.cancel() + + with pytest.raises(asyncio.CancelledError): + await claim + + assert executor.calls == 1 + current = AgentTaskOper().get(task.id) + assert current.last_status == "waiting" + assert current.last_run_id is None + assert AgentTaskOper().list_runs(task.id) == [] + + @pytest.mark.anyio async def test_cancelled_started_claim_is_compensated_before_return() -> None: """认领事务已开始时取消,返回前必须把已提交运行收口为失败。"""