mulit-agent/mulit_agent/tests/test_agent_service_mcp.py

import json
import time
from threading import Lock

from mulit_agent_app.application.agent_service import AgentService
from mulit_agent_app.config import AgentConfig
from mulit_agent_app.infrastructure.database import AgentDatabase


class _Session:
    def __init__(self) -> None:
        self.acks: list[tuple[str, int, str, str]] = []

    def send_mcp_ack(self, tool_id: str, version: int, status: str, error: str = "") -> None:
        self.acks.append((tool_id, version, status, error))


class _Executor:
    def __init__(self) -> None:
        self.desktop_enabled = False
        self.endpoints: tuple = ()

    def set_desktop_filesystem_enabled(self, enabled: bool) -> None:
        self.desktop_enabled = enabled

    def set_mcp_endpoints(self, endpoints: tuple) -> None:
        for endpoint in self.endpoints:
            if endpoint not in endpoints:
                endpoint.close()
        self.endpoints = tuple(endpoints)

    def mcp_endpoints(self) -> tuple:
        return self.endpoints


def _service(tmp_path) -> AgentService:
    config = AgentConfig(tmp_path, tmp_path / "settings.json", tmp_path / "agent.sqlite3", tmp_path / "agent.log")
    service = AgentService.__new__(AgentService)
    service._config = config
    service._database = AgentDatabase(config.database_path)
    service._database.initialize()
    service.session = _Session()
    service.task_executor = _Executor()
    service._mcp_lock = Lock()
    service._mcp_sequence = 0
    return service


def _desktop_snapshot(version: int) -> dict:
    return {
        "tools": [
            {
                "id": "desktop-filesystem",
                "name": "桌面文件工具",
                "kind": "builtin",
                "version": "1.0.0",
                "deployment_version": version,
                "enabled": True,
                "scope": "desktop-directories",
            }
        ]
    }


def test_mcp_snapshot_persists_a_builtin_tool_and_reacknowledges_after_reconnect(tmp_path):
    service = _service(tmp_path)
    snapshot = _desktop_snapshot(3)

    service.apply_mcp_snapshot(snapshot)
    service.apply_mcp_snapshot(snapshot)

    stored = json.loads(service._database.get_meta("mcp.tool.desktop-filesystem"))
    assert stored["deployment_version"] == 3 and stored["kind"] == "builtin"
    assert service.task_executor.desktop_enabled is True
    assert service.session.acks == [
        ("desktop-filesystem", 3, "applied", ""),
        ("desktop-filesystem", 3, "applied", ""),
    ]


def test_mcp_snapshot_rejects_a_lower_version_without_changing_state(tmp_path):
    service = _service(tmp_path)
    service.apply_mcp_snapshot(_desktop_snapshot(4))
    service.session.acks.clear()

    service.apply_mcp_snapshot(_desktop_snapshot(3))

    assert json.loads(service._database.get_meta("mcp.tool.desktop-filesystem"))["deployment_version"] == 4
    # The rejected snapshot is answered, and the already-deployed tool is re-acknowledged.
    assert service.session.acks == [
        ("desktop-filesystem", 3, "failed", "拒绝低版本工具部署。"),
        ("desktop-filesystem", 4, "applied", ""),
    ]


def test_mcp_snapshot_rejects_an_unsupported_builtin_and_removes_unselected_tools(tmp_path):
    service = _service(tmp_path)
    service.apply_mcp_snapshot(_desktop_snapshot(1))
    service.session.acks.clear()

    service.apply_mcp_snapshot(
        {
            "tools": [
                {"id": "docker", "kind": "builtin", "deployment_version": 1, "enabled": True},
            ]
        }
    )
    assert service.session.acks == [("docker", 1, "failed", "Agent 不支持该内置工具。")]

    service.session.acks.clear()
    service.apply_mcp_snapshot({"tools": []})
    assert service._database.get_meta("mcp.tool.desktop-filesystem") == ""
    assert service.task_executor.desktop_enabled is False


def test_http_mcp_tool_is_reported_failed_when_its_endpoint_is_unreachable(tmp_path):
    service = _service(tmp_path)

    service.apply_mcp_snapshot(
        {
            "tools": [
                {
                    "id": "imageforge",
                    "name": "ImageForge",
                    "kind": "http",
                    "endpoint": "http://127.0.0.1:6/mcp",
                    "version": "1.0.0",
                    "deployment_version": 2,
                    "enabled": True,
                }
            ]
        }
    )
    deadline = time.time() + 5
    while not service.session.acks and time.time() < deadline:
        time.sleep(0.05)

    assert len(service.session.acks) == 1
    tool_id, version, status, error = service.session.acks[0]
    assert (tool_id, version, status) == ("imageforge", 2, "failed") and error
    # A failed endpoint must not be published as a usable tool.
    assert service.task_executor.endpoints == ()


def _stdio_snapshot(command: str, args: list[str], version: int = 1) -> dict:
    return {
        "tools": [
            {
                "id": "echo",
                "name": "回显工具",
                "kind": "stdio",
                "command": command,
                "args": args,
                "version": "1.0.0",
                "deployment_version": version,
                "enabled": True,
            }
        ]
    }


def test_stdio_tool_is_spawned_discovered_and_reused_across_reconnects(tmp_path, python_executable, mcp_server_script):
    service = _service(tmp_path)
    snapshot = _stdio_snapshot(python_executable, [mcp_server_script])

    try:
        service.apply_mcp_snapshot(snapshot)
        _wait_for_acks(service, 1)
        assert service.session.acks == [("echo", 1, "applied", "")]
        endpoints = service.task_executor.endpoints
        assert [endpoint.tool_id for endpoint in endpoints] == ["echo"]
        process = endpoints[0]._client._transport._process

        # A reconnect replays the same definition: it must keep the live process
        # instead of spawning a second server.
        service.session.acks.clear()
        service.apply_mcp_snapshot({**_stdio_snapshot(python_executable, [mcp_server_script], version=2)})
        _wait_for_acks(service, 1)

        assert service.session.acks == [("echo", 2, "applied", "")]
        assert service.task_executor.endpoints[0] is endpoints[0]
        assert process.poll() is None, "未变更的定义不应重启 MCP 进程"
    finally:
        service.task_executor.set_mcp_endpoints(())
    assert process.poll() is not None, "Agent 必须回收自己拉起的 MCP 进程"


def test_stdio_tool_is_reported_failed_when_its_command_is_missing(tmp_path):
    service = _service(tmp_path)

    service.apply_mcp_snapshot(_stdio_snapshot("definitely-not-installed-command", []))
    _wait_for_acks(service, 1)

    tool_id, version, status, error = service.session.acks[0]
    assert (tool_id, version, status) == ("echo", 1, "failed") and "找不到要执行的命令" in error
    assert service.task_executor.endpoints == ()


def test_stdio_tool_rejects_a_definition_with_invalid_arguments(tmp_path):
    service = _service(tmp_path)

    service.apply_mcp_snapshot(_stdio_snapshot("python", ["bad\nargument"]))

    assert service.session.acks == [("echo", 1, "failed", "本机命令参数无效。")]


def _wait_for_acks(service: AgentService, count: int, timeout: float = 10.0) -> None:
    deadline = time.time() + timeout
    while len(service.session.acks) < count and time.time() < deadline:
        time.sleep(0.05)