feat(error-shim): degraded-mode MCP server for setup/bootstrap failures

Author: Broccolito

src/codegraphagent/error_shim.py

@@ -0,0 +1,135 @@
+"""Degraded-mode MCP server.
+
+When bootstrap or layout setup fails, the main shim hands control here. This
+server still satisfies the MCP `initialize` + `tools/list` + `tools/call`
+handshake — but the only tool it offers returns the underlying error so the
+agent (and the user reading the log) can see exactly what went wrong.
+
+Implemented in pure stdlib to keep it functional even if fastmcp's import
+itself is what failed.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from typing import BinaryIO
+
+from codegraphagent.errors import (
+    BootstrapError,
+    CodeGraphAgentError,
+    LayoutConflictError,
+)
+
+
+def _tool_for(exc: CodeGraphAgentError) -> dict:
+    """Return the synthetic tool descriptor matching the given error type."""
+    if isinstance(exc, BootstrapError):
+        return {
+            "name": "codegraphagent_bootstrap_error",
+            "description": (
+                "CodeGraphAgent failed to download or verify the engine bundle. "
+                "Call this tool to see the underlying error."
+            ),
+            "inputSchema": {"type": "object", "properties": {}, "required": []},
+        }
+    if isinstance(exc, LayoutConflictError):
+        return {
+            "name": "codegraphagent_setup_error",
+            "description": (
+                "CodeGraphAgent could not set up the .biorouter/codegraph "
+                "symlink. Call this tool to see the remediation steps."
+            ),
+            "inputSchema": {"type": "object", "properties": {}, "required": []},
+        }
+    return {
+        "name": "codegraphagent_error",
+        "description": "CodeGraphAgent encountered an internal error.",
+        "inputSchema": {"type": "object", "properties": {}, "required": []},
+    }
+
+
+def _error_text(exc: CodeGraphAgentError) -> str:
+    parts = [f"CodeGraphAgent error: {exc}"]
+    if isinstance(exc, BootstrapError):
+        if exc.url:
+            parts.append(f"URL: {exc.url}")
+        if exc.expected_sha and exc.observed_sha:
+            parts.append(f"Expected SHA-256: {exc.expected_sha}")
+            parts.append(f"Observed SHA-256: {exc.observed_sha}")
+        parts.append(
+            "Recovery: retry; or set CODEGRAPH_ENGINE_PATH to a pre-downloaded "
+            "bundle; or pin CODEGRAPH_ENGINE_VERSION to a different release."
+        )
+    elif isinstance(exc, LayoutConflictError):
+        parts.append(f"Path: {exc.path}")
+        parts.append(
+            "Recovery: rename or remove the conflicting .codegraph directory, "
+            "then restart the CodeGraphAgent extension."
+        )
+    return "\n".join(parts)
+
+
+def serve(
+    exc: CodeGraphAgentError,
+    *,
+    stdin: BinaryIO | None = None,
+    stdout: BinaryIO | None = None,
+) -> None:
+    """Serve MCP requests in degraded mode until `shutdown` is received."""
+    inp = stdin or sys.stdin.buffer
+    out = stdout or sys.stdout.buffer
+    tool = _tool_for(exc)
+
+    while True:
+        line = inp.readline()
+        if not line:
+            break
+        try:
+            req = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+
+        method = req.get("method")
+        req_id = req.get("id")
+
+        if method == "initialize":
+            resp = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": {
+                    "protocolVersion": req.get("params", {}).get(
+                        "protocolVersion", "2024-11-05"
+                    ),
+                    "capabilities": {"tools": {}},
+                    "serverInfo": {"name": "codegraphagent (degraded)", "version": "0.1.0"},
+                },
+            }
+        elif method == "tools/list":
+            resp = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": {"tools": [tool]},
+            }
+        elif method == "tools/call":
+            resp = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": {
+                    "content": [{"type": "text", "text": _error_text(exc)}],
+                    "isError": True,
+                },
+            }
+        elif method == "shutdown":
+            resp = {"jsonrpc": "2.0", "id": req_id, "result": None}
+            out.write((json.dumps(resp) + "\n").encode())
+            out.flush()
+            break
+        else:
+            resp = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "error": {"code": -32601, "message": f"Method not found: {method}"},
+            }
+        out.write((json.dumps(resp) + "\n").encode())
+        out.flush()

tests/test_error_shim.py

@@ -0,0 +1,78 @@
+"""error_shim — minimal MCP server that surfaces a single error tool when the
+shim couldn't reach the engine."""
+
+from __future__ import annotations
+
+import io
+import json
+import threading
+
+from codegraphagent import error_shim
+from codegraphagent.errors import BootstrapError, LayoutConflictError
+
+
+def _send(stream: io.BytesIO, frame: dict) -> None:
+    stream.write((json.dumps(frame) + "\n").encode())
+
+
+def test_error_shim_lists_bootstrap_error_tool():
+    parent_in = io.BytesIO()
+    _send(parent_in, {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}})
+    _send(parent_in, {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}})
+    _send(parent_in, {"jsonrpc": "2.0", "id": 3, "method": "shutdown"})
+    parent_in.seek(0)
+    parent_out = io.BytesIO()
+
+    exc = BootstrapError(
+        "boom",
+        url="https://example.com/x.tar.gz",
+        expected_sha="abc",
+        observed_sha="def",
+    )
+    error_shim.serve(exc, stdin=parent_in, stdout=parent_out)
+
+    frames = [json.loads(l) for l in parent_out.getvalue().decode().splitlines() if l]
+    # frames: initialize response, tools/list response, shutdown response
+    assert len(frames) == 3
+    tools_list = frames[1]["result"]["tools"]
+    assert any(t["name"] == "codegraphagent_bootstrap_error" for t in tools_list)
+
+
+def test_error_shim_call_returns_error_details():
+    parent_in = io.BytesIO()
+    _send(parent_in, {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}})
+    _send(parent_in, {
+        "jsonrpc": "2.0", "id": 2, "method": "tools/call",
+        "params": {"name": "codegraphagent_bootstrap_error", "arguments": {}},
+    })
+    _send(parent_in, {"jsonrpc": "2.0", "id": 3, "method": "shutdown"})
+    parent_in.seek(0)
+    parent_out = io.BytesIO()
+
+    exc = BootstrapError(
+        "download failed",
+        url="https://example.com/x.tar.gz",
+    )
+    error_shim.serve(exc, stdin=parent_in, stdout=parent_out)
+
+    frames = [json.loads(l) for l in parent_out.getvalue().decode().splitlines() if l]
+    call_result = frames[1]["result"]
+    text = call_result["content"][0]["text"]
+    assert "download failed" in text
+    assert "https://example.com/x.tar.gz" in text
+
+
+def test_error_shim_layout_conflict():
+    parent_in = io.BytesIO()
+    _send(parent_in, {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}})
+    _send(parent_in, {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}})
+    _send(parent_in, {"jsonrpc": "2.0", "id": 3, "method": "shutdown"})
+    parent_in.seek(0)
+    parent_out = io.BytesIO()
+
+    exc = LayoutConflictError("path exists", path="/tmp/foo/.codegraph")
+    error_shim.serve(exc, stdin=parent_in, stdout=parent_out)
+
+    frames = [json.loads(l) for l in parent_out.getvalue().decode().splitlines() if l]
+    tools_list = frames[1]["result"]["tools"]
+    assert any(t["name"] == "codegraphagent_setup_error" for t in tools_list)