MoviePilot/tests/test_agent_steering.py

"""
运行中 Agent 消息排队、真实 HumanMessage 注入及身份隔离测试。
"""

import asyncio
import json
from unittest.mock import AsyncMock, Mock

import pytest
from langchain.agents import create_agent
from langchain.agents.middleware.types import ModelRequest, ModelResponse
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver

from app.agent.middleware.steering import SteeringMiddleware
from app.agent.orchestrator import MoviePilotAgent
from app.agent.session import AgentSessionOwner, _MessageTask
from app.agent.steering import STEERING_QUEUE_MAX_SIZE, SteeringInbox, bind_steering_inbox, reset_steering_inbox
from app.schemas.types import NotificationChannel, ReplyMode


@pytest.mark.asyncio
async def test_steering_message_is_injected_as_real_human_message() -> None:
    """
    下一次模型调用应在图状态中看到真实 HumanMessage，而不是提示词拼接。
    """
    inbox = SteeringInbox("session", "user")
    await inbox.begin_run()
    queued = await inbox.enqueue(user_id="user", text="补充要求")
    assert queued is not None

    graph = create_agent(
        model=GenericFakeChatModel(messages=iter([AIMessage(content="完成")])),
        tools=[],
        middleware=[SteeringMiddleware()],
        checkpointer=InMemorySaver(),
    )
    token = bind_steering_inbox(inbox)
    try:
        result = await graph.ainvoke(
            {"messages": [HumanMessage(content="开始任务")]},
            config={"configurable": {"thread_id": "session"}},
        )
    finally:
        reset_steering_inbox(token)

    messages = result["messages"]
    assert [message.type for message in messages] == ["human", "human", "ai"]
    assert "补充要求" in messages[1].content[0]["text"]
    assert messages[1].additional_kwargs["moviepilot_steering_message_id"] == queued.message_id
    steering_payload = json.loads(messages[1].content[0]["text"])
    assert "保留此前的范围" in steering_payload["continuation_context"]


@pytest.mark.asyncio
async def test_steering_does_not_break_unresolved_tool_call_pairing() -> None:
    """
    未交付工具调用所在的状态边界不应插入用户消息。
    """
    inbox = SteeringInbox("session", "user")
    await inbox.begin_run()
    assert await inbox.enqueue(user_id="user", text="稍等") is not None
    middleware = SteeringMiddleware()
    token = bind_steering_inbox(inbox)
    seen_messages = []

    async def handler(request):
        """
        记录未注入的模型请求，模拟下游模型处理。
        """
        seen_messages.extend(request.messages)
        return ModelResponse(result=[AIMessage(content="继续执行工具")])

    try:
        pending_tool_message = AIMessage(
            content="",
            tool_calls=[{"name": "tool", "args": {}, "id": "call"}],
        )
        await middleware.awrap_model_call(
            request=ModelRequest(
                model=GenericFakeChatModel(messages=iter([AIMessage(content="unused")])),
                tools=[],
                system_message=None,
                response_format=None,
                messages=[pending_tool_message],
                tool_choice=None,
                state={"messages": [pending_tool_message]},
                runtime=None,
            ),
            handler=handler,
        )
    finally:
        reset_steering_inbox(token)
    assert seen_messages[-1].tool_calls
    assert inbox.pending_count == 1


@pytest.mark.asyncio
async def test_owner_routes_matching_user_message_to_active_steering_inbox() -> None:
    """
    活动会话只接受同一用户的补充消息，其他用户回退到普通入口。
    """
    started = asyncio.Event()
    release = asyncio.Event()
    processed_messages: list[str] = []

    class BlockingAgent:
        """
        模拟持久会话 Agent，暴露运行边界供 owner 测试。
        """

        def __init__(self, **kwargs):
            """
            保存宿主注入字段。
            """
            self.__dict__.update(kwargs)
            self.user_id = str(kwargs["user_id"])

        def set_output_callback(self, callback):
            """
            兼容会话 owner 的输出回调更新。
            """
            self.output_callback = callback

        def set_protected_output_callback(self, callback):
            """
            兼容会话 owner 的敏感输出回调更新。
            """
            self.protected_output_callback = callback

        def get_session_status(self):
            """
            返回 owner 状态合并所需的最小快照。
            """
            return {"model": "fake", "is_processing": True}

        async def process(self, message, **kwargs):
            """
            阻塞到测试释放，模拟长时间模型运行。
            """
            processed_messages.append(message)
            started.set()
            await release.wait()
            return message

        async def cleanup(self):
            """
            报告测试 Agent 已收敛。
            """
            return True

    owner = AgentSessionOwner()
    owner._accepting_tasks = True
    running = asyncio.create_task(
        owner.process_message(
            session_id="session",
            user_id="user",
            message="开始",
            agent_factory=BlockingAgent,
            wait_for_completion=True,
        )
    )
    await asyncio.wait_for(started.wait(), 1)
    accepted = await owner.submit_steering_message(
        session_id="session",
        user_id="user",
        message="继续检查",
    )
    rejected = await owner.submit_steering_message(
        session_id="session",
        user_id="other",
        message="越权",
    )
    assert accepted is not None
    assert rejected is None
    assert owner.get_session_status("session")["steering_pending"] == 1

    release.set()
    assert await asyncio.wait_for(running, 1) == "继续检查"
    assert processed_messages == ["开始", "继续检查"]
    worker = owner._session_workers.get("session")
    if worker:
        worker.cancel()
        await asyncio.gather(worker, return_exceptions=True)


@pytest.mark.asyncio
async def test_steering_finish_boundary_returns_pending_messages_once() -> None:
    """
    运行收尾与入队共享锁时，未注入消息只能被交给下一轮一次。
    """
    inbox = SteeringInbox("session", "user")
    await inbox.begin_run()
    message = await inbox.enqueue(user_id="user", text="收尾前到达")
    assert message is not None
    pending = await inbox.finish_run()
    assert pending == (message,)
    assert await inbox.finish_run() == ()
    assert await inbox.enqueue(user_id="user", text="已关闭运行") is None


@pytest.mark.asyncio
@pytest.mark.parametrize("channel", [NotificationChannel.Telegram, NotificationChannel.Feishu, NotificationChannel.Wechat])
async def test_channel_followup_enters_running_graph_without_starting_another_task(monkeypatch, channel) -> None:
    """
    真实会话入口应把连续补充消息注入当前图，并仅在消费时切分渠道回复。
    """
    started = asyncio.Event()
    release = asyncio.Event()
    results = []
    processed = []
    monkeypatch.setattr("app.agent.session._async_start_processing_status", AsyncMock(return_value=None))
    monkeypatch.setattr("app.agent.session._async_finish_processing_status", AsyncMock())

    class ChannelAgent:
        """
        使用真实 steering 中间件，隔离模型与消息渠道的测试 Agent。
        """

        def __init__(self, **kwargs):
            """
            接收会话上下文并记录消息切分调用。
            """
            self.__dict__.update(kwargs)
            self.stream_handler = Mock()
            self._tool_context = {"user_reply_sent": True}

        def start_steering_reply(self, original_message_id=None):
            """
            复用生产回复切分，验证补充输入解除前段主动回复的抑制状态。
            """
            MoviePilotAgent.start_steering_reply(self, original_message_id)

        async def process(self, message, **_kwargs):
            """
            暂停当前任务，让补充输入在下一次真实模型调用前进入 inbox。
            """
            processed.append(message)
            started.set()
            await release.wait()
            graph = create_agent(
                model=GenericFakeChatModel(messages=iter([AIMessage(content="已按补充要求完成")])),
                tools=[],
                middleware=[SteeringMiddleware()],
                checkpointer=InMemorySaver(),
            )
            token = bind_steering_inbox(self.steering_inbox)
            try:
                result = await graph.ainvoke(
                    {"messages": [HumanMessage(content=message), AIMessage(content="正在检查")]},
                    config={"configurable": {"thread_id": "channel-session"}},
                )
                results.extend(result["messages"])
            finally:
                reset_steering_inbox(token)
            return "完成"

    owner = AgentSessionOwner()
    owner._accepting_tasks = True
    context = {
        "session_id": "channel-session", "user_id": "user", "channel": channel.value,
        "source": "configured-channel", "original_chat_id": "chat", "is_channel_admin": True,
    }
    try:
        await owner.process_message(**context, message="检查服务", agent_factory=ChannelAgent)
        await asyncio.wait_for(started.wait(), 1)
        agent = owner.active_agents["channel-session"]
        await owner.process_message(
            **context, message="补充检查目录", original_message_id="user-2",
            images=["data:image/png;base64,test"], files=[{"ref": "log.txt"}],
        )
        await owner.process_message(**context, message="仅报告异常", original_message_id="user-3")
        assert owner._session_queues["channel-session"].empty()
        assert agent.steering_inbox.pending_count == 2
        agent.stream_handler.start_new_message.assert_not_called()
        release.set()
        await asyncio.wait_for(owner._session_queues["channel-session"].join(), 2)

        assert processed == ["检查服务"]
        assert [message.type for message in results] == ["human", "ai", "human", "human", "ai"]
        payload = json.loads(results[2].content[0]["text"])
        assert payload["message"] == "补充检查目录"
        assert payload["files"] == [{"ref": "log.txt"}]
        assert results[2].content[1]["image_url"]["url"] == "data:image/png;base64,test"
        assert json.loads(results[3].content[0]["text"])["message"] == "仅报告异常"
        assert [call.args for call in agent.stream_handler.start_new_message.call_args_list] == [("user-2",), ("user-3",)]
        assert agent.original_message_id == "user-3"
        assert agent._tool_context["user_reply_sent"] is False
    finally:
        release.set()
        worker = owner._session_workers.get("channel-session")
        if worker:
            worker.cancel()
            await asyncio.gather(worker, return_exceptions=True)


@pytest.mark.asyncio
@pytest.mark.parametrize("changed", [
    {"user_id": "another-user"}, {"source": "another-bot"}, {"original_chat_id": "another-chat"},
    {"is_channel_admin": False}, {"channel": NotificationChannel.Feishu.value},
    {"channel": NotificationChannel.WebAgent.value}, {"reply_mode": ReplyMode.CAPTURE_ONLY},
    {"allow_message_tools": False}, {"scheduled_run_id": "scheduled"},
])
async def test_channel_steering_preserves_identity_and_independent_task_boundaries(changed) -> None:
    """
    其它用户、聊天、渠道、权限或独立调用不得进入已有渠道任务的上下文。
    """
    context = {
        "session_id": "session", "user_id": "user", "channel": NotificationChannel.Telegram.value,
        "source": "bot", "original_chat_id": "chat", "is_channel_admin": True,
    }
    owner = AgentSessionOwner()
    owner._accepting_tasks = True
    inbox = SteeringInbox("session", "user")
    await inbox.begin_run()
    owner._session_active_tasks["session"] = _MessageTask(**context, message="开始")
    owner._session_steering_inboxes["session"] = inbox
    owner._session_queues["session"] = asyncio.Queue()
    owner._session_workers["session"] = Mock(done=lambda: False)
    await owner.process_message(**{**context, **changed}, message="独立请求")
    assert inbox.pending_count == 0
    assert owner._session_queues["session"].get_nowait().message == "独立请求"


@pytest.mark.asyncio
@pytest.mark.parametrize("running", [True, False])
async def test_channel_steering_falls_back_to_queue_when_full_or_finished(running) -> None:
    """
    inbox 满或运行已收尾时，补充消息仍进入普通队列且只保留一次。
    """
    context = {
        "session_id": "session", "user_id": "user", "channel": NotificationChannel.Telegram.value,
        "source": "bot",
    }
    owner = AgentSessionOwner()
    owner._accepting_tasks = True
    inbox = SteeringInbox("session", "user")
    if running:
        await inbox.begin_run()
        for index in range(STEERING_QUEUE_MAX_SIZE):
            await inbox.enqueue(user_id="user", text=str(index))
    owner._session_active_tasks["session"] = _MessageTask(**context, message="开始")
    owner._session_steering_inboxes["session"] = inbox
    owner._session_queues["session"] = asyncio.Queue()
    owner._session_workers["session"] = Mock(done=lambda: False)
    await owner.process_message(**context, message="不能丢失")
    assert inbox.pending_count == (STEERING_QUEUE_MAX_SIZE if running else 0)
    assert owner._session_queues["session"].get_nowait().message == "不能丢失"