protocol.py

"""MCP core: JSON-RPC validation, lifecycle, and method dispatch.

`MCPServer.dispatch()` is transport-agnostic and is the single seam to subclass
when crafting adversarial / non-spec-compliant behaviour.
"""
from __future__ import annotations

import json

from .constants import (
    INSTRUCTIONS,
    INTERNAL_ERROR,
    INVALID_PARAMS,
    INVALID_REQUEST,
    METHOD_NOT_FOUND,
    PROTOCOL_VERSION,
    SERVER_INFO,
    SERVER_NOT_INITIALIZED,
    SUPPORTED_VERSIONS,
)
from .registry import PROMPTS, RESOURCES, TOOLS
from .tracing import log


class JsonRpcError(Exception):
    """Raise inside a handler to produce a JSON-RPC error response."""

    def __init__(self, code: int, message: str, data: object | None = None):
        super().__init__(message)
        self.code = code
        self.message = message
        self.data = data


class MCPServer:
    """Transport-agnostic MCP core. One instance holds the state for one logical
    session (stdio: the process; HTTP: one Mcp-Session-Id)."""

    def __init__(self, *, strict: bool = False):
        self.strict = strict          # enforce lifecycle ordering
        self.initialized = False      # received notifications/initialized
        self.client_info: dict | None = None
        self.negotiated_version = PROTOCOL_VERSION

    # ---- entry point -------------------------------------------------------
    def dispatch(self, message: object) -> dict | None:
        """Route one parsed JSON-RPC message. Returns a response dict, or None
        for notifications (which never get a reply)."""
        if not isinstance(message, dict) or message.get("jsonrpc") != "2.0":
            mid = message.get("id") if isinstance(message, dict) else None
            return self._error(mid, INVALID_REQUEST, "Invalid Request: expected a JSON-RPC 2.0 object")

        is_notification = "id" not in message
        msg_id = message.get("id")
        method = message.get("method")
        params = message.get("params") or {}

        if not isinstance(method, str):
            if is_notification:
                return None
            return self._error(msg_id, INVALID_REQUEST, "Invalid Request: 'method' must be a string")

        try:
            if (self.strict and not self.initialized
                    and method not in ("initialize", "ping", "notifications/initialized")):
                raise JsonRpcError(SERVER_NOT_INITIALIZED, "Server not initialized")

            handler = self._route(method)
            if handler is None:
                if is_notification:
                    return None  # unknown notifications are silently ignored
                raise JsonRpcError(METHOD_NOT_FOUND, f"Method not found: {method}")

            result = handler(params)
            if is_notification:
                return None
            return {"jsonrpc": "2.0", "id": msg_id, "result": result}

        except JsonRpcError as exc:
            if is_notification:
                log(f"error while handling notification {method!r}: {exc.message}")
                return None
            return self._error(msg_id, exc.code, exc.message, exc.data)
        except Exception as exc:  # defensive: never crash the loop on a handler bug
            log(f"internal error in {method!r}: {exc!r}")
            if is_notification:
                return None
            return self._error(msg_id, INTERNAL_ERROR, f"Internal error: {exc}")

    # ---- routing -----------------------------------------------------------
    def _route(self, method: str):
        return {
            "initialize": self._initialize,
            "notifications/initialized": self._initialized,
            "ping": self._ping,
            "tools/list": self._tools_list,
            "tools/call": self._tools_call,
            "resources/list": self._resources_list,
            "resources/read": self._resources_read,
            "resources/templates/list": self._resource_templates_list,
            "prompts/list": self._prompts_list,
            "prompts/get": self._prompts_get,
        }.get(method)

    # ---- lifecycle ---------------------------------------------------------
    def _initialize(self, params: dict) -> dict:
        requested = params.get("protocolVersion")
        self.client_info = params.get("clientInfo")
        self.negotiated_version = requested if requested in SUPPORTED_VERSIONS else PROTOCOL_VERSION
        return {
            "protocolVersion": self.negotiated_version,
            "capabilities": {
                "tools": {"listChanged": True},
                "resources": {"listChanged": True},
                "prompts": {"listChanged": True},
            },
            "serverInfo": SERVER_INFO,
            "instructions": INSTRUCTIONS,
        }

    def _initialized(self, params: dict) -> None:
        self.initialized = True
        return None

    def _ping(self, params: dict) -> dict:
        return {}

    # ---- tools -------------------------------------------------------------
    def _tools_list(self, params: dict) -> dict:
        return {"tools": [t["meta"] for t in TOOLS.values()]}

    def _tools_call(self, params: dict) -> dict:
        name = params.get("name")
        arguments = params.get("arguments") or {}
        entry = TOOLS.get(name)
        if entry is None:
            # Unknown tool is a protocol error per spec (not an execution error).
            raise JsonRpcError(INVALID_PARAMS, f"Unknown tool: {name}")
        try:
            result = entry["handler"](arguments)
        except Exception as exc:
            # Execution failures are reported in-band so a model can self-correct.
            return {"content": [{"type": "text", "text": f"Tool execution error: {exc}"}],
                    "isError": True}
        return self._normalize_tool_result(result)

    @staticmethod
    def _normalize_tool_result(result: object) -> dict:
        if isinstance(result, str):
            return {"content": [{"type": "text", "text": result}]}
        if isinstance(result, list):
            return {"content": result}
        if isinstance(result, dict):
            if "content" in result:
                return result  # handler supplied a complete CallToolResult
            # structured output + serialized-JSON text block (spec backward-compat)
            return {"content": [{"type": "text", "text": json.dumps(result)}],
                    "structuredContent": result}
        return {"content": [{"type": "text", "text": str(result)}]}

    # ---- resources ---------------------------------------------------------
    def _resources_list(self, params: dict) -> dict:
        return {"resources": [r["meta"] for r in RESOURCES.values()]}

    def _resources_read(self, params: dict) -> dict:
        uri = params.get("uri")
        entry = RESOURCES.get(uri)
        if entry is None:
            raise JsonRpcError(INVALID_PARAMS, f"Unknown resource: {uri}")
        content = entry["handler"](uri)
        if isinstance(content, str):
            mime = entry["meta"].get("mimeType", "text/plain")
            return {"contents": [{"uri": uri, "mimeType": mime, "text": content}]}
        if isinstance(content, dict):
            return {"contents": [content]}
        if isinstance(content, list):
            return {"contents": content}
        return {"contents": [{"uri": uri, "text": str(content)}]}

    def _resource_templates_list(self, params: dict) -> dict:
        return {"resourceTemplates": []}

    # ---- prompts -----------------------------------------------------------
    def _prompts_list(self, params: dict) -> dict:
        return {"prompts": [p["meta"] for p in PROMPTS.values()]}

    def _prompts_get(self, params: dict) -> dict:
        name = params.get("name")
        arguments = params.get("arguments") or {}
        entry = PROMPTS.get(name)
        if entry is None:
            raise JsonRpcError(INVALID_PARAMS, f"Unknown prompt: {name}")
        result = entry["handler"](arguments)
        if isinstance(result, dict):
            return result  # complete GetPromptResult
        if isinstance(result, list):
            return {"messages": result}
        return {"messages": [{"role": "user", "content": {"type": "text", "text": str(result)}}]}

    # ---- helpers -----------------------------------------------------------
    @staticmethod
    def _error(msg_id: object, code: int, message: str, data: object | None = None) -> dict:
        err = {"code": code, "message": message}
        if data is not None:
            err["data"] = data
        return {"jsonrpc": "2.0", "id": msg_id, "error": err}