fix(agent): propagate task cancellation (#6284)

tests/test_agent_cancellation.py

@@ -0,0 +1,153 @@
+import asyncio
+from unittest.mock import AsyncMock
+
+import pytest
+
+from app.agent import AgentManager, MoviePilotAgent
+
+
+def test_execute_agent_propagates_task_cancellation():
+    """取消 Agent 执行时应终止调用方任务，不能转换成普通完成结果。"""
+    started = asyncio.Event()
+
+    class _BlockingAgent:
+        """等待取消的最小 LangGraph 替身。"""
+
+        async def ainvoke(self, _payload, config=None):  # noqa: ARG002
+            """阻塞到外层任务取消。"""
+            started.set()
+            await asyncio.Event().wait()
+
+    async def _run_scenario():
+        agent = MoviePilotAgent(session_id="session-1", user_id="10001")
+        agent._should_stream = lambda: False
+        agent._create_agent = AsyncMock(return_value=_BlockingAgent())
+        agent.stream_handler.stop_streaming = AsyncMock(return_value=(False, ""))
+
+        execution = asyncio.create_task(agent._execute_agent([]))
+        await asyncio.wait_for(started.wait(), timeout=1)
+        execution.cancel()
+
+        with pytest.raises(asyncio.CancelledError):
+            await execution
+        agent.stream_handler.stop_streaming.assert_awaited_once()
+
+    asyncio.run(_run_scenario())
+
+
+def test_stop_current_task_cancels_waiters_and_allows_next_message():
+    """停止会话应结束当前及排队请求，并允许同一会话继续处理消息。"""
+
+    async def _run_scenario():
+        manager = AgentManager()
+        started = asyncio.Event()
+
+        async def _block_current_task(_task):
+            started.set()
+            await asyncio.Event().wait()
+
+        manager._process_message_internal = _block_current_task
+        first_waiter = asyncio.create_task(
+            manager.process_message(
+                session_id="session-1",
+                user_id="10001",
+                message="first",
+                wait_for_completion=True,
+            )
+        )
+        await asyncio.wait_for(started.wait(), timeout=1)
+        second_waiter = asyncio.create_task(
+            manager.process_message(
+                session_id="session-1",
+                user_id="10001",
+                message="second",
+                wait_for_completion=True,
+            )
+        )
+        await asyncio.sleep(0)
+
+        try:
+            assert await asyncio.wait_for(
+                manager.stop_current_task("session-1"), timeout=1
+            ) is True
+            with pytest.raises(asyncio.CancelledError):
+                await first_waiter
+            with pytest.raises(asyncio.CancelledError):
+                await asyncio.wait_for(second_waiter, timeout=1)
+
+            manager._process_message_internal = AsyncMock(return_value="resumed")
+            result = await asyncio.wait_for(
+                manager.process_message(
+                    session_id="session-1",
+                    user_id="10001",
+                    message="next",
+                    wait_for_completion=True,
+                ),
+                timeout=1,
+            )
+            assert result == "resumed"
+        finally:
+            await manager.stop_current_task("session-1")
+            for waiter in (first_waiter, second_waiter):
+                if not waiter.done():
+                    waiter.cancel()
+            await asyncio.gather(
+                first_waiter,
+                second_waiter,
+                return_exceptions=True,
+            )
+
+    asyncio.run(_run_scenario())
+
+
+def test_stop_queues_new_message_until_cancellation_cleanup_finishes():
+    """旧 worker 清理期间到达的新消息应保留，并在清理完成后执行。"""
+
+    async def _run_scenario():
+        manager = AgentManager()
+        current_started = asyncio.Event()
+        cancellation_cleanup_started = asyncio.Event()
+        release_cleanup = asyncio.Event()
+
+        async def _process(task):
+            if task.message == "current":
+                current_started.set()
+                await asyncio.Event().wait()
+            return "next-completed"
+
+        async def _finish_status(_task):
+            cancellation_cleanup_started.set()
+            await release_cleanup.wait()
+
+        manager._process_message_internal = _process
+        manager._finish_task_processing_status = _finish_status
+        current_waiter = asyncio.create_task(
+            manager.process_message(
+                session_id="session-1",
+                user_id="10001",
+                message="current",
+                wait_for_completion=True,
+            )
+        )
+        await asyncio.wait_for(current_started.wait(), timeout=1)
+        stop_task = asyncio.create_task(manager.stop_current_task("session-1"))
+        await asyncio.wait_for(cancellation_cleanup_started.wait(), timeout=1)
+
+        next_waiter = asyncio.create_task(
+            manager.process_message(
+                session_id="session-1",
+                user_id="10001",
+                message="next",
+                wait_for_completion=True,
+            )
+        )
+        await asyncio.sleep(0)
+        assert not next_waiter.done()
+
+        release_cleanup.set()
+        assert await asyncio.wait_for(stop_task, timeout=1) is True
+        assert await asyncio.wait_for(next_waiter, timeout=1) == "next-completed"
+        with pytest.raises(asyncio.CancelledError):
+            await current_waiter
+
+    asyncio.run(_run_scenario())