test_mcp_server.py

from __future__ import annotations

import importlib
import sys
import types
from pathlib import Path


def _install_fake_mcp_module() -> list[tuple[str, callable]]:
    """Install a stub for mcp.server.fastmcp.FastMCP so we can exercise the
    server construction path without installing the real SDK."""
    registrations: list[tuple[str, callable]] = []

    class FakeFastMCP:
        def __init__(self, name: str) -> None:
            self.name = name

        def tool(self, *args, **kwargs):
            def decorator(fn):
                registrations.append((fn.__name__, fn))
                return fn

            return decorator

        def run(self) -> None:  # pragma: no cover - not invoked in tests
            raise NotImplementedError

    fastmcp_module = types.ModuleType("mcp.server.fastmcp")
    fastmcp_module.FastMCP = FakeFastMCP
    server_module = types.ModuleType("mcp.server")
    server_module.fastmcp = fastmcp_module
    mcp_module = types.ModuleType("mcp")
    mcp_module.server = server_module

    sys.modules["mcp"] = mcp_module
    sys.modules["mcp.server"] = server_module
    sys.modules["mcp.server.fastmcp"] = fastmcp_module
    return registrations


def test_mcp_server_registers_core_tools(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setenv("STATE_TRACE_STORAGE_PATH", str(tmp_path / "mcp.db"))
    monkeypatch.setenv("STATE_TRACE_NAMESPACE", "test-repo")
    monkeypatch.setenv("STATE_TRACE_CAPACITY_LIMIT", "48")

    registrations = _install_fake_mcp_module()

    mcp_server = importlib.import_module("state_trace.mcp_server")
    importlib.reload(mcp_server)
    mcp_server.build_server()

    names = {name for name, _ in registrations}
    assert {
        "store",
        "retrieve",
        "retrieve_brief",
        "record_action",
        "record_observation",
        "record_test_result",
        "ingest_agent_log_file",
        "list_namespaces",
        "graph_snapshot",
        "current_state",
        "failed_hypotheses",
    }.issubset(names)

    tools = dict(registrations)
    task_payload = tools["store"](
        text="Fix login by tracing the refresh token path",
        context={"type": "task", "session": "auth", "file": "auth.ts"},
    )
    assert task_payload["type"] == "task"
    assert task_payload["metadata"].get("namespace") == "test-repo"

    retrieval = tools["retrieve"](
        query="why is login broken?",
        context={"session": "auth"},
        explain=True,
    )
    assert retrieval["namespace"] == "test-repo"
    assert "explanations" in retrieval

    namespaces = tools["list_namespaces"]()
    assert "test-repo" in namespaces

    state = tools["current_state"](session="auth")
    assert state["session"] == "auth"
    assert state["active_task"] is not None
    assert state["active_task"]["type"] == "task"

    failures = tools["failed_hypotheses"](session="auth")
    assert isinstance(failures, list)

    for mod in ("mcp.server.fastmcp", "mcp.server", "mcp"):
        sys.modules.pop(mod, None)