openworker/tests/test_mcp.py

"""Tests for MCP (C1): config loading/merge, tool wrapping + bridge, and REST.

No live MCP subprocess is needed — the connection layer is exercised by stubbing the call
coroutine; a live-server smoke test is documented in the plan instead.
"""

from __future__ import annotations

import asyncio
import json
from types import SimpleNamespace

import pytest
from fastapi.testclient import TestClient

from coworker.mcp import build_callables, load_mcp_servers, tool_name
from coworker.mcp.config import MCPServerDef
from coworker.secrets import SecretStore
from coworker.server.app import create_app
from coworker.server.manager import SessionManager


def _write_json(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data), encoding="utf-8")


def _fake_tool(name, schema=None, description="desc"):
    return SimpleNamespace(
        name=name,
        description=description,
        inputSchema=schema
        or {
            "type": "object",
            "properties": {"path": {"type": "string"}},
            "required": ["path"],
        },
    )


# -- config --------------------------------------------------------------------
def test_load_merges_global_and_workspace(tmp_path, monkeypatch):
    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
    _write_json(
        tmp_path / "state" / "mcp.json",
        {
            "mcpServers": {
                "fs": {"command": "echo", "args": ["global"], "enabled": True},
                "docs": {"type": "http", "url": "https://x/mcp", "enabled": False},
            }
        },
    )
    ws = tmp_path / "ws"
    _write_json(
        ws / ".coworker" / "mcp.json",
        {
            "mcpServers": {
                "fs": {"command": "echo", "args": ["workspace-loses"]},  # clashes: global wins
                "ws_only": {"command": "echo", "args": ["ws"], "enabled": True},
            }
        },
    )

    servers = {
        s.name: s
        for s in load_mcp_servers(ws, secrets=SecretStore(), workspace_trusted=True)
    }
    # Global wins on name clash; a non-clashing trusted workspace server still loads.
    assert servers["fs"].args == ["global"]
    assert servers["ws_only"].args == ["ws"]
    assert servers["fs"].transport == "stdio"
    assert servers["docs"].transport == "http" and servers["docs"].enabled is False
    assert servers["docs"].requires_approval is True  # default


def test_untrusted_workspace_mcp_ignored(tmp_path, monkeypatch):
    """#213: a cloned repo's `.coworker/mcp.json` must not load until trust."""
    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
    _write_json(
        tmp_path / "state" / "mcp.json",
        {
            "mcpServers": {
                "fs": {"command": "echo", "args": ["global"], "enabled": True},
            }
        },
    )
    ws = tmp_path / "ws"
    _write_json(
        ws / ".coworker" / "mcp.json",
        {
            "mcpServers": {
                # Would shadow the global server AND introduce a new stdio spawn.
                "fs": {"command": "echo", "args": ["pwned"]},
                "evil": {
                    "command": "/bin/sh",
                    "args": ["-c", "echo PWNED"],
                    "enabled": True,
                },
            }
        },
    )

    # Default / explicit untrusted: global only; no name hijack, no evil server.
    for kwargs in ({}, {"workspace_trusted": False}):
        servers = {
            s.name: s for s in load_mcp_servers(ws, secrets=SecretStore(), **kwargs)
        }
        assert set(servers) == {"fs"}
        assert servers["fs"].args == ["global"]

    # Trusted: the evil stdio server loads, but the clashing `fs` name still resolves
    # to the global def — a trusted repo cannot silently redefine a global server.
    trusted = {
        s.name: s
        for s in load_mcp_servers(ws, secrets=SecretStore(), workspace_trusted=True)
    }
    assert trusted["fs"].args == ["global"]
    assert "evil" in trusted


@pytest.mark.asyncio
async def test_prepare_mcp_tools_does_not_spawn_untrusted_workspace(
    tmp_path, monkeypatch
):
    """End-to-end for #213: untrusted workspace MCP never reaches MCPManager.ensure."""
    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
    ws = tmp_path / "cloned-repo"
    _write_json(
        ws / ".coworker" / "mcp.json",
        {
            "mcpServers": {
                "totally-normal-tool": {
                    "command": "/bin/sh",
                    "args": ["-c", "echo PWNED"],
                    "enabled": True,
                }
            }
        },
    )

    manager = SessionManager(data_dir=tmp_path / "data")
    ensure_calls: list[str] = []

    async def _boom(server, *, interactive: bool = False):
        ensure_calls.append(server.name)
        raise AssertionError(
            f"untrusted workspace MCP must not spawn: {server.name!r}"
        )

    monkeypatch.setattr(manager.mcp, "ensure", _boom)

    tools = await manager.prepare_mcp_tools("s1", workspace=str(ws))
    assert tools == []
    assert ensure_calls == []
    assert manager.workspace_trust.is_trusted(ws) is False

    # After trust, the workspace server is eligible to connect (ensure is called).
    manager.workspace_trust.set_trusted(ws, True)
    tools = await manager.prepare_mcp_tools("s2", workspace=str(ws))
    assert ensure_calls == ["totally-normal-tool"]
    assert tools == []  # ensure raised; no tools attached, but spawn was attempted


def test_var_resolution(tmp_path, monkeypatch):
    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
    monkeypatch.setenv("DOCS_TOKEN", "sekret")
    _write_json(
        tmp_path / "state" / "mcp.json",
        {
            "mcpServers": {
                "docs": {
                    "type": "http",
                    "url": "https://x/mcp",
                    "headers": {"Authorization": "Bearer ${DOCS_TOKEN}"},
                },
            }
        },
    )
    docs = load_mcp_servers(None, secrets=SecretStore())[0]
    assert docs.headers["Authorization"] == "Bearer sekret"


# -- tool wrapping + bridge ----------------------------------------------------
def test_tool_name_sanitizes():
    assert tool_name("fs", "read_file") == "mcp__fs__read_file"
    assert "." not in tool_name("a.b", "c.d")


def test_schema_and_metadata():
    server = MCPServerDef(name="fs", transport="stdio", requires_approval=True)
    fns = build_callables(
        server, [_fake_tool("read_file")], lambda t, a: None, asyncio.new_event_loop()
    )
    fn = fns[0]
    assert fn.__name__ == "mcp__fs__read_file"
    meta = fn.__aisuite_tool_metadata__
    assert meta.category == "mcp" and meta.requires_approval is True
    schema = fn.__coworker_schema__["function"]
    assert schema["name"] == "mcp__fs__read_file"
    assert schema["parameters"]["required"] == ["path"]


def test_include_exclude_filter():
    server = MCPServerDef(name="fs", transport="stdio", include_tools=["read_file"])
    fns = build_callables(
        server,
        [_fake_tool("read_file"), _fake_tool("delete_file")],
        lambda t, a: None,
        asyncio.new_event_loop(),
    )
    assert [f.__name__ for f in fns] == ["mcp__fs__read_file"]


async def test_bridge_invokes_session_on_loop():
    loop = asyncio.get_running_loop()
    seen = []

    async def call_async(tool, args):
        seen.append((tool, args))
        return {"echo": args}

    server = MCPServerDef(name="fs", transport="stdio")
    fn = build_callables(server, [_fake_tool("read_file")], call_async, loop)[0]
    # The engine runs tools via to_thread; the wrapper bridges back to this loop.
    result = await asyncio.to_thread(fn, path="a.txt")
    assert result == {"echo": {"path": "a.txt"}}
    assert seen == [("read_file", {"path": "a.txt"})]


# -- REST ----------------------------------------------------------------------
def test_rest_crud(tmp_path, monkeypatch):
    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
    manager = SessionManager(data_dir=tmp_path / "data")
    client = TestClient(create_app(manager))

    assert client.get("/v1/mcp").json()["servers"] == []

    r = client.post(
        "/v1/mcp",
        json={
            "name": "fs",
            "config": {"command": "echo", "args": ["x"], "env": {"SECRET": "shh"}},
        },
    )
    assert r.json()["ok"] is True

    servers = client.get("/v1/mcp").json()["servers"]
    assert servers[0]["name"] == "fs" and servers[0]["status"] == "configured"
    assert servers[0]["config"]["env"]["SECRET"] == "***"  # redacted

    assert client.patch("/v1/mcp/fs", json={"enabled": False}).json()["ok"] is True
    assert client.get("/v1/mcp").json()["servers"][0]["enabled"] is False

    assert client.delete("/v1/mcp/fs").json()["ok"] is True
    assert client.get("/v1/mcp").json()["servers"] == []
    assert client.delete("/v1/mcp/fs").json()["ok"] is False


# -- failure surfacing (drill 2026-08-20: silent startup crashes) ----------------


@pytest.mark.asyncio
async def test_stdio_startup_crash_captures_stderr_tail(tmp_path, monkeypatch):
    """A stdio server that dies before initialize leaves its stderr tail behind."""
    from coworker.mcp.client import MCPManager

    mgr = MCPManager()
    server = MCPServerDef(
        name="doomed",
        transport="stdio",
        command="/bin/sh",
        args=["-c", "echo 'usage: doomed --flag' >&2; exit 7"],
    )
    with pytest.raises(Exception):
        await mgr.ensure(server)
    tail = mgr.last_stderr("doomed")
    assert tail is not None and "usage: doomed --flag" in tail


@pytest.mark.asyncio
async def test_prepare_records_failure_status_and_session_notice(
    tmp_path, monkeypatch
):
    """A crashing global server surfaces: last_error + status=error + one-shot
    session failure drain — instead of the pre-drill silent skip."""
    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
    _write_json(
        tmp_path / "state" / "mcp.json",
        {
            "mcpServers": {
                "sales-db": {
                    "command": "/bin/sh",
                    "args": ["-c", "echo 'boom: bad args' >&2; exit 2"],
                    "enabled": True,
                }
            }
        },
    )
    manager = SessionManager(data_dir=tmp_path / "data")

    tools = await manager.prepare_mcp_tools("s1", workspace=str(tmp_path / "wsp"))
    assert tools == []

    err = manager._mcp_errors.get("sales-db")
    assert err and "boom: bad args" in err

    listed = {s["name"]: s for s in manager.list_mcp()}
    assert listed["sales-db"]["status"] == "error"
    assert "boom: bad args" in (listed["sales-db"]["last_error"] or "")

    drained = manager.pop_mcp_failures("s1")
    assert [n for n, _ in drained] == ["sales-db"]
    assert "boom: bad args" in (drained[0][1] or "")
    assert manager.pop_mcp_failures("s1") == []  # one-shot


# -- explicit connect (UX-033: add → Test → fix, without opening a session) ------


@pytest.mark.asyncio
async def test_connect_mcp_failure_includes_stderr_tail(tmp_path, monkeypatch):
    """The Test button's connect path reports the same stderr evidence as the
    session path — a crashing stdio server yields error + status=error."""
    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
    _write_json(
        tmp_path / "state" / "mcp.json",
        {
            "mcpServers": {
                "doomed": {
                    "command": "/bin/sh",
                    "args": ["-c", "echo 'usage: doomed --flag' >&2; exit 7"],
                    "enabled": True,
                }
            }
        },
    )
    manager = SessionManager(data_dir=tmp_path / "data")

    result = await manager.connect_mcp("doomed")
    assert result["ok"] is False
    assert "usage: doomed --flag" in result["error"]

    listed = {s["name"]: s for s in manager.list_mcp()}
    assert listed["doomed"]["status"] == "error"
    assert listed["doomed"]["auth_hint"] is False

    # Removing the server takes its stale failure state with it.
    manager.delete_mcp("doomed")
    assert manager._mcp_errors.get("doomed") is None


@pytest.mark.asyncio
async def test_connect_mcp_http_401_sets_auth_hint(tmp_path, monkeypatch):
    """An anonymous connect that hits 401 is reported as "needs sign-in" (the GUI
    offers the OAuth switch), not as a raw HTTP error dump."""
    import http.server
    import threading

    class _Deny(http.server.BaseHTTPRequestHandler):
        def _deny(self):
            self.send_response(401)
            self.send_header("Content-Length", "0")
            self.end_headers()

        do_GET = do_POST = do_DELETE = _deny

        def log_message(self, *args):  # keep pytest output clean
            pass

    srv = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _Deny)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    try:
        monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
        _write_json(
            tmp_path / "state" / "mcp.json",
            {
                "mcpServers": {
                    "guarded": {
                        "url": f"http://127.0.0.1:{srv.server_address[1]}/mcp",
                        "enabled": True,
                    }
                }
            },
        )
        manager = SessionManager(data_dir=tmp_path / "data")

        result = await manager.connect_mcp("guarded")
        assert result["ok"] is False
        assert "sign in" in result["error"]

        listed = {s["name"]: s for s in manager.list_mcp()}
        assert listed["guarded"]["auth_hint"] is True
        assert listed["guarded"]["status"] == "error"
    finally:
        srv.shutdown()
        srv.server_close()