labyricorn-site/labyricorn_mcp/server.py

"""Dependency-free MCP stdio transport for the local Labyricorn repository."""

from __future__ import annotations

import json
import logging
import sys
from typing import Any, TextIO

from . import __version__
from .errors import LabyricornMcpError
from .safety import RepositoryGuard
from .tools import ToolRegistry


SUPPORTED_PROTOCOLS = ("2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05")
LATEST_PROTOCOL = SUPPORTED_PROTOCOLS[0]


class StdioMcpServer:
    def __init__(self, guard: RepositoryGuard, *, input_stream: TextIO | None = None, output_stream: TextIO | None = None):
        self.guard = guard
        self.registry = ToolRegistry(guard)
        self.input = input_stream or sys.stdin
        self.output = output_stream or sys.stdout
        self.initialized = False
        self.client_initialized = False
        self.protocol_version: str | None = None
        self.log = logging.getLogger("labyricorn_mcp")

    def serve_forever(self) -> int:
        self.log.info("server started for repository %s", self.guard.root)
        for raw_line in self.input:
            line = raw_line.rstrip("\r\n")
            if not line:
                continue
            try:
                message = json.loads(line)
            except json.JSONDecodeError:
                self._write_error(None, -32700, "Parse error")
                continue
            if not isinstance(message, dict) or isinstance(message.get("id"), (dict, list)):
                self._write_error(message.get("id") if isinstance(message, dict) else None, -32600, "Invalid Request")
                continue
            self._handle(message)
        self.log.info("server stopped after stdin closed")
        return 0

    def _handle(self, message: dict[str, Any]) -> None:
        method = message.get("method")
        request_id = message.get("id")
        is_notification = "id" not in message
        if message.get("jsonrpc") != "2.0" or not isinstance(method, str):
            if not is_notification:
                self._write_error(request_id, -32600, "Invalid Request")
            return
        params = message.get("params", {})
        if not isinstance(params, dict):
            if not is_notification:
                self._write_error(request_id, -32602, "Invalid params")
            return

        if is_notification:
            if method == "notifications/initialized" and self.initialized:
                self.client_initialized = True
            return

        if method == "initialize":
            self._initialize(request_id, params)
            return
        if not self.initialized:
            self._write_error(request_id, -32002, "Server is not initialized")
            return
        if method == "ping":
            self._write_result(request_id, {})
        elif method == "tools/list":
            self._write_result(request_id, {"tools": self.registry.tools})
        elif method == "tools/call":
            self._call_tool(request_id, params)
        else:
            self._write_error(request_id, -32601, "Method not found")

    def _initialize(self, request_id: Any, params: dict[str, Any]) -> None:
        if self.initialized:
            self._write_error(request_id, -32600, "Server is already initialized")
            return
        requested = params.get("protocolVersion")
        self.protocol_version = requested if requested in SUPPORTED_PROTOCOLS else LATEST_PROTOCOL
        self.initialized = True
        self._write_result(
            request_id,
            {
                "protocolVersion": self.protocol_version,
                "capabilities": {"tools": {"listChanged": False}},
                "serverInfo": {"name": "labyricorn-mcp", "version": __version__},
                "instructions": (
                    "This server manages only the configured local Labyricorn Git checkout. "
                    "It has no production administration capability. Commits and pushes are separate explicit tools."
                ),
            },
        )

    def _call_tool(self, request_id: Any, params: dict[str, Any]) -> None:
        name = params.get("name")
        if not isinstance(name, str):
            self._write_error(request_id, -32602, "Tool name is required")
            return
        self.log.info("tool invocation: %s", name)
        try:
            result = self.registry.call(name, params.get("arguments"))
        except LabyricornMcpError as exc:
            self.log.warning("tool rejected: %s code=%s", name, exc.code)
            payload = {"error": exc.as_dict()}
            self._write_result(
                request_id,
                {
                    "content": [{"type": "text", "text": json.dumps(payload, ensure_ascii=False)}],
                    "structuredContent": payload,
                    "isError": True,
                },
            )
        except Exception:
            self.log.exception("unexpected tool failure: %s", name)
            payload = {"error": {"code": "INTERNAL_ERROR", "message": "The tool failed unexpectedly; see local stderr diagnostics."}}
            self._write_result(
                request_id,
                {
                    "content": [{"type": "text", "text": json.dumps(payload)}],
                    "structuredContent": payload,
                    "isError": True,
                },
            )
        else:
            self._write_result(
                request_id,
                {
                    "content": [{"type": "text", "text": json.dumps(result, ensure_ascii=False, default=str)}],
                    "structuredContent": result,
                    "isError": False,
                },
            )

    def _write_result(self, request_id: Any, result: dict[str, Any]) -> None:
        self._write({"jsonrpc": "2.0", "id": request_id, "result": result})

    def _write_error(self, request_id: Any, code: int, message: str) -> None:
        self._write({"jsonrpc": "2.0", "id": request_id, "error": {"code": code, "message": message}})

    def _write(self, message: dict[str, Any]) -> None:
        self.output.write(json.dumps(message, ensure_ascii=False, separators=(",", ":")) + "\n")
        self.output.flush()