fix(agent): close cancellation lifecycle races (#6440)

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

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,

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

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:
     """认领事务已开始时取消，返回前必须把已提交运行收口为失败。"""