ObsiGate/tests/test_mcp.py

# tests/test_mcp.py — Unit tests for the MCP server (Phase E)
"""Tests for the Streamable HTTP MCP endpoint: handshake, tools, resources,
prompts, per-vault permissions and single-use confirmation tokens.

The endpoint is mounted at ``/mcp``. Auth is disabled in the ``client``
fixture, so the caller is an anonymous admin unless a test monkeypatches
``backend.mcp.server._authenticate``.
"""

import json

import pytest
from fastapi.testclient import TestClient

ACCEPT = "application/json, text/event-stream"
PROTOCOL = "2025-03-26"


def _post(client, payload, session=None, headers=None):
    merged = {"Accept": ACCEPT, "Content-Type": "application/json"}
    if session:
        merged["Mcp-Session-Id"] = session
    if headers:
        merged.update(headers)
    return client.post("/mcp", content=json.dumps(payload), headers=merged)


def _initialize(client):
    resp = _post(
        client,
        {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": PROTOCOL,
                "capabilities": {},
                "clientInfo": {"name": "pytest", "version": "1.0"},
            },
        },
    )
    assert resp.status_code == 200, resp.text
    session = resp.headers.get("mcp-session-id")
    assert session
    # Client acknowledges initialization (notification, no id).
    ack = _post(client, {"jsonrpc": "2.0", "method": "notifications/initialized"}, session=session)
    assert ack.status_code == 202
    return session, resp.json()


def _call(client, session, method, params=None, req_id=2):
    resp = _post(
        client,
        {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params or {}},
        session=session,
    )
    assert resp.status_code == 200, resp.text
    return resp.json()


def _text_json(result):
    return json.loads(result["content"][0]["text"])


@pytest.fixture
def mcp_client(app_with_vault, monkeypatch):
    """Persistent TestClient (lifespan + single event loop) with a fresh MCP manager."""
    import backend.main as main

    async def _noop_build(*args, **kwargs):
        return None

    monkeypatch.setattr(main, "build_index", _noop_build)
    monkeypatch.setattr(main, "init_inverted_index", lambda: None)

    # The manager can only run() once per instance; reset it between tests so
    # each test gets a manager bound to its own event loop.
    main.mcp_app._manager = None
    main.mcp_app._run_task = None
    main.mcp_app._start_lock = None

    with TestClient(main.app) as client:
        yield client


# ═══════════════════════════════════════════════════════════════════
# Handshake & primitives
# ═══════════════════════════════════════════════════════════════════


class TestHandshake:
    def test_initialize(self, mcp_client):
        session, body = _initialize(mcp_client)
        assert session
        result = body["result"]
        assert result["serverInfo"]["name"] == "obsigate"
        assert result["protocolVersion"]
        assert "tools" in result["capabilities"]
        assert "resources" in result["capabilities"]
        assert "prompts" in result["capabilities"]

    def test_tools_list(self, mcp_client):
        session, _ = _initialize(mcp_client)
        body = _call(mcp_client, session, "tools/list")
        names = {t["name"] for t in body["result"]["tools"]}
        # Read tools exposed directly.
        assert "read_file" in names
        assert "search_fulltext" in names
        # Mutations exposed only through propose_/apply_.
        assert "edit_file" not in names
        assert "propose_edit_file" in names
        assert "apply_edit_file" in names
        assert "propose_delete_file" in names
        assert "apply_delete_file" in names

    def test_resources_list(self, mcp_client):
        session, _ = _initialize(mcp_client)
        body = _call(mcp_client, session, "resources/list")
        uris = {r["uri"] for r in body["result"]["resources"]}
        assert "vault://TestVault" in uris

    def test_resource_templates(self, mcp_client):
        session, _ = _initialize(mcp_client)
        body = _call(mcp_client, session, "resources/templates/list")
        templates = [t["uriTemplate"] for t in body["result"]["resourceTemplates"]]
        assert "vault://{vault}/{path}" in templates

    def test_prompts_list_and_get(self, mcp_client):
        session, _ = _initialize(mcp_client)
        listing = _call(mcp_client, session, "prompts/list")
        names = {p["name"] for p in listing["result"]["prompts"]}
        assert {"summarize-directory", "generate-note", "find-related"} <= names

        got = _call(
            mcp_client,
            session,
            "prompts/get",
            {"name": "find-related", "arguments": {"vault": "TestVault", "path": "note1.md"}},
        )
        text = got["result"]["messages"][0]["content"]["text"]
        assert "note1.md" in text


# ═══════════════════════════════════════════════════════════════════
# Tools: read + two-step mutations
# ═══════════════════════════════════════════════════════════════════


class TestTools:
    def test_call_read_tool(self, mcp_client):
        session, _ = _initialize(mcp_client)
        body = _call(
            mcp_client,
            session,
            "tools/call",
            {"name": "read_file", "arguments": {"vault": "TestVault", "path": "note1.md"}},
        )
        payload = _text_json(body["result"])
        assert payload["ok"] is True
        assert "Python" in payload["data"]["content"]

    def test_resource_read(self, mcp_client):
        session, _ = _initialize(mcp_client)
        body = _call(
            mcp_client,
            session,
            "resources/read",
            {"uri": "vault://TestVault/note1.md"},
        )
        text = body["result"]["contents"][0]["text"]
        assert "Python" in text

    def test_propose_then_apply(self, mcp_client):
        session, _ = _initialize(mcp_client)
        proposed = _call(
            mcp_client,
            session,
            "tools/call",
            {
                "name": "propose_edit_file",
                "arguments": {"vault": "TestVault", "path": "note1.md", "content": "# Edited via MCP\n"},
            },
        )
        payload = _text_json(proposed["result"])
        token = payload["confirmation_token"]
        assert payload["tool"] == "edit_file"
        assert "diff" in payload

        applied = _call(
            mcp_client,
            session,
            "tools/call",
            {"name": "apply_edit_file", "arguments": {"confirmation_token": token}},
            req_id=3,
        )
        result = _text_json(applied["result"])
        assert result["ok"] is True

        from backend.indexer import get_vault_data
        from pathlib import Path

        content = (Path(get_vault_data("TestVault")["path"]) / "note1.md").read_text(encoding="utf-8")
        assert content == "# Edited via MCP\n"

    def test_token_replay_rejected(self, mcp_client):
        session, _ = _initialize(mcp_client)
        proposed = _call(
            mcp_client,
            session,
            "tools/call",
            {
                "name": "propose_edit_file",
                "arguments": {"vault": "TestVault", "path": "note1.md", "content": "# Once\n"},
            },
        )
        token = _text_json(proposed["result"])["confirmation_token"]

        first = _call(
            mcp_client,
            session,
            "tools/call",
            {"name": "apply_edit_file", "arguments": {"confirmation_token": token}},
            req_id=3,
        )
        assert _text_json(first["result"])["ok"] is True

        second = _call(
            mcp_client,
            session,
            "tools/call",
            {"name": "apply_edit_file", "arguments": {"confirmation_token": token}},
            req_id=4,
        )
        payload = _text_json(second["result"])
        assert payload["ok"] is False
        assert payload["error"]["code"] == "token_reused"

    def test_apply_without_token_rejected(self, mcp_client):
        session, _ = _initialize(mcp_client)
        body = _call(
            mcp_client,
            session,
            "tools/call",
            {"name": "apply_edit_file", "arguments": {}},
        )
        # The low-level handler converts the ValueError into an MCP error result.
        assert body["result"].get("isError") is True


# ═══════════════════════════════════════════════════════════════════
# Permissions & auth
# ═══════════════════════════════════════════════════════════════════


class TestPermissions:
    def test_denied_vault_tool_call(self, mcp_client, monkeypatch):
        import backend.mcp.server as mcp_server

        monkeypatch.setattr(
            mcp_server,
            "_authenticate",
            lambda request: {"username": "limited", "vaults": ["OtherVault"], "_token_vaults": ["OtherVault"]},
        )
        session, _ = _initialize(mcp_client)
        body = _call(
            mcp_client,
            session,
            "tools/call",
            {"name": "read_file", "arguments": {"vault": "TestVault", "path": "note1.md"}},
        )
        payload = _text_json(body["result"])
        assert payload["ok"] is False
        assert payload["error"]["code"] in ("vault_access_denied", "permission_denied")

    def test_denied_vault_resource_read(self, mcp_client, monkeypatch):
        import backend.mcp.server as mcp_server

        monkeypatch.setattr(
            mcp_server,
            "_authenticate",
            lambda request: {"username": "limited", "vaults": ["OtherVault"], "_token_vaults": ["OtherVault"]},
        )
        session, _ = _initialize(mcp_client)
        body = _call(
            mcp_client,
            session,
            "resources/read",
            {"uri": "vault://TestVault/note1.md"},
        )
        # A denied read surfaces as a JSON-RPC error (or an isError result).
        assert "error" in body or body.get("result", {}).get("isError") is True

    def test_missing_token_returns_401_when_auth_enabled(self, mcp_client, monkeypatch):
        import backend.mcp.server as mcp_server

        monkeypatch.setattr(mcp_server, "is_auth_enabled", lambda: True)
        resp = _post(
            mcp_client,
            {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}},
        )
        assert resp.status_code == 401