""" 兼容协议请求的 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 _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): """ 模拟没有可停止任务的会话,返回 False。 """ 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, ), ) runtime_config = SimpleNamespace( ai_agent_enable=False, api_token=_API_TOKEN, ) with patch.object( openai, "get_api_runtime_config_snapshot", return_value=runtime_config, ), patch.object( anthropic, "get_api_runtime_config_snapshot", return_value=runtime_config, ), patch.object( openai, "validate_api_credential_identity", ), patch.object( anthropic, "validate_api_credential_identity", ), 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( openai, "get_api_runtime_config_snapshot", return_value=SimpleNamespace( ai_agent_enable=True, api_token=_API_TOKEN, ), ), patch.object( openai, "validate_api_credential_identity", ), 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( anthropic, "get_api_runtime_config_snapshot", return_value=SimpleNamespace( ai_agent_enable=True, api_token=_API_TOKEN, ), ), patch.object( anthropic, "validate_api_credential_identity", ), 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,返回预设处理结果。 """ 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