MoviePilot/tests/test_agent_protocol_lifecycle.py

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