refactor(agent): 按需加载 Agent 运行时 (#6336)

tests/test_agent_protocol_lifecycle.py

@@ -0,0 +1,216 @@
+"""兼容协议请求的 AgentManager ownership 与关闭竞态合同。"""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from fastapi.security import HTTPAuthorizationCredentials
+
+from app import schemas
+from app.api.endpoints import anthropic, openai
+from app.runtime.config import settings
+
+_API_TOKEN = "test-agent-protocol-token"
+
+
+class _ManagerClosedError(RuntimeError):
+    """模拟 enqueue 时 manager 已关闭的 acceptance gate 错误。"""
+
+    code = "agent_manager_unavailable"
+
+
+class _ClosingManager:
+    """拒绝新任务并记录请求级清理的 manager 替身。"""
+
+    def __init__(self) -> None:
+        self.process_calls = []
+        self.clear_calls = []
+
+    async def process_message(self, **kwargs):
+        self.process_calls.append(kwargs)
+        raise _ManagerClosedError("AgentManager 已关闭")
+
+    async def clear_session(self, **kwargs):
+        self.clear_calls.append(kwargs)
+
+    async def stop_current_task(self, _session_id):
+        return False
+
+
+async def _collect(response) -> str:
+    """收集 StreamingResponse 的全部文本块。"""
+    chunks = []
+    async for chunk in response.body_iterator:
+        chunks.append(chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk)
+    return "".join(chunks)
+
+
+def test_streaming_protocols_reject_config_disable_before_manager_lookup() -> None:
+    """配置关闭后流式请求保持 503，且不得接触运行态 manager。"""
+    credentials = HTTPAuthorizationCredentials(
+        scheme="Bearer",
+        credentials=_API_TOKEN,
+    )
+    openai_payload = schemas.OpenAIChatCompletionsRequest(
+        messages=[schemas.OpenAIChatMessage(role="user", content="hello")],
+        stream=True,
+    )
+    anthropic_payload = schemas.AnthropicMessagesRequest(
+        messages=[schemas.AnthropicMessage(role="user", content="hello")],
+        stream=True,
+    )
+
+    async def scenario():
+        return (
+            await openai.chat_completions(
+                openai_payload,
+                SimpleNamespace(headers={}),
+                credentials,
+            ),
+            await anthropic.messages(
+                anthropic_payload,
+                x_api_key=_API_TOKEN,
+            ),
+        )
+
+    with patch.object(settings, "AI_AGENT_ENABLE", False), patch.object(
+        settings,
+        "API_TOKEN",
+        _API_TOKEN,
+    ), patch.object(
+        openai,
+        "get_running_agent_manager",
+    ) as openai_manager, patch.object(
+        anthropic,
+        "get_running_agent_manager",
+    ) as anthropic_manager:
+        responses = asyncio.run(scenario())
+
+    assert [response.status_code for response in responses] == [503, 503]
+    openai_manager.assert_not_called()
+    anthropic_manager.assert_not_called()
+
+
+def test_openai_stream_rejects_shutdown_race_and_cleans_request_session() -> None:
+    """随机 OpenAI 流在 enqueue 竞态失败时返回协议错误并清理临时会话。"""
+    manager = _ClosingManager()
+    credentials = HTTPAuthorizationCredentials(
+        scheme="Bearer",
+        credentials=_API_TOKEN,
+    )
+    payload = schemas.OpenAIChatCompletionsRequest(
+        messages=[schemas.OpenAIChatMessage(role="user", content="hello")],
+        stream=True,
+    )
+
+    async def scenario() -> str:
+        response = await openai.chat_completions(
+            payload,
+            SimpleNamespace(headers={}),
+            credentials,
+        )
+        return await _collect(response)
+
+    with patch.object(settings, "AI_AGENT_ENABLE", True), patch.object(
+        settings,
+        "API_TOKEN",
+        _API_TOKEN,
+    ), patch.object(
+        openai,
+        "get_running_agent_manager",
+        return_value=manager,
+    ):
+        body = asyncio.run(scenario())
+
+    assert '"type": "server_error"' in body
+    assert "data: [DONE]" in body
+    assert len(manager.process_calls) == 1
+    assert manager.process_calls[0]["wait_for_completion"] is True
+    assert callable(manager.process_calls[0]["agent_setup"])
+    assert len(manager.clear_calls) == 1
+
+
+def test_anthropic_stream_rejects_shutdown_race_and_cleans_request_session() -> None:
+    """Anthropic 流在 enqueue 竞态失败时返回 error 终态并清理临时会话。"""
+    manager = _ClosingManager()
+    payload = schemas.AnthropicMessagesRequest(
+        messages=[schemas.AnthropicMessage(role="user", content="hello")],
+        stream=True,
+    )
+
+    async def scenario() -> str:
+        response = await anthropic.messages(
+            payload,
+                x_api_key=_API_TOKEN,
+        )
+        return await _collect(response)
+
+    with patch.object(settings, "AI_AGENT_ENABLE", True), patch.object(
+        settings,
+        "API_TOKEN",
+        _API_TOKEN,
+    ), patch.object(
+        anthropic,
+        "get_running_agent_manager",
+        return_value=manager,
+    ):
+        body = asyncio.run(scenario())
+
+    assert "event: error" in body
+    assert "event: message_stop" in body
+    assert len(manager.process_calls) == 1
+    assert manager.process_calls[0]["wait_for_completion"] is True
+    assert callable(manager.process_calls[0]["agent_setup"])
+    assert len(manager.clear_calls) == 1
+
+
+def test_managed_protocol_request_releases_its_stream_queue() -> None:
+    """协议请求完成后不应由持久会话 Agent 继续强引用请求队列。"""
+    event_queue = asyncio.Queue()
+    created_agents = []
+
+    class ProtocolAgent:
+        """记录请求绑定与释放的最小协议 Agent。"""
+
+        def __init__(self, **_kwargs):
+            self.collected_messages = ["done"]
+            self.bound_queue = None
+            created_agents.append(self)
+
+        def configure_protocol_request(self, *, stream_mode, event_queue):
+            assert stream_mode is True
+            self.bound_queue = event_queue
+
+        def release_protocol_request(self, queue):
+            if self.bound_queue is queue:
+                self.bound_queue = None
+
+    class RunningManager:
+        """在 worker 边界执行 agent_setup 的 manager 替身。"""
+
+        async def process_message(self, **kwargs):
+            agent = kwargs["agent_factory"]()
+            kwargs["agent_setup"](agent)
+            return "done"
+
+    async def scenario():
+        with patch.object(
+            openai,
+            "_get_collecting_agent_type",
+            return_value=ProtocolAgent,
+        ):
+            return await openai._run_managed_agent(
+                manager=RunningManager(),
+                session_id="persistent",
+                user_id="1",
+                username="api",
+                source="openai",
+                prompt="hello",
+                images=[],
+                stream_mode=True,
+                event_queue=event_queue,
+            )
+
+    assert asyncio.run(scenario()) == ("done", ["done"])
+    assert len(created_agents) == 1
+    assert created_agents[0].bound_queue is None