"""兼容协议请求的 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