_server.py

"""MCP JSON-RPC server core. Transport-agnostic — see ``_mount`` for HTTP."""

from __future__ import annotations

import asyncio
import json
import re
import time
from typing import Any

from hawkapi_mcp._dispatch import dispatch_tool
from hawkapi_mcp._errors import (
    INTERNAL_ERROR,
    INVALID_PARAMS,
    INVALID_REQUEST,
    METHOD_NOT_FOUND,
    RpcError,
)
from hawkapi_mcp._tools import openapi_to_tools

MCP_PROTOCOL_VERSION = "2025-06-18"

# Response headers never surfaced to the agent via ``structuredContent``.
# Anything matching these would leak credentials/session material set by the
# inner route back to the (untrusted) MCP client.
_SENSITIVE_HEADERS = frozenset({"set-cookie", "authorization"})
_SENSITIVE_HEADER_PATTERN = re.compile(r"^x-.*-(token|secret)$", re.IGNORECASE)

# Re-exported for backward compatibility with anyone importing the constants
# from this module.
_PARSE_ERROR = -32700
_INVALID_REQUEST = INVALID_REQUEST
_METHOD_NOT_FOUND = METHOD_NOT_FOUND
_INVALID_PARAMS = INVALID_PARAMS
_INTERNAL_ERROR = INTERNAL_ERROR


class MCPServer:
    """Stateless MCP server that exposes a HawkAPI app's routes as tools.

    Methods supported (subset of the spec sufficient for tool execution):

    * ``initialize``
    * ``ping``
    * ``tools/list``
    * ``tools/call``
    """

    def __init__(
        self,
        app: Any,
        *,
        server_name: str | None = None,
        server_version: str | None = None,
        include_only: set[str] | None = None,
        exclude: set[str] | None = None,
        cache_ttl_seconds: float | None = None,
        strip_response_headers: set[str] | None = None,
    ) -> None:
        self._app = app
        self._name = server_name or getattr(app, "title", "hawkapi-mcp")
        self._version = server_version or getattr(app, "version", "0.0.0")
        self._include_only = include_only
        self._exclude = exclude
        self._cache_ttl = cache_ttl_seconds
        self._strip_response_headers = {h.lower() for h in strip_response_headers or set()}
        self._tools_cache: list[dict[str, Any]] | None = None
        self._tools_cached_at: float = 0.0
        self._tools_lock = asyncio.Lock()

    async def _tools(self) -> list[dict[str, Any]]:
        if self._tools_cache is not None and not self._cache_expired():
            return self._tools_cache
        async with self._tools_lock:
            if self._tools_cache is None or self._cache_expired():
                spec = self._app.openapi() if hasattr(self._app, "openapi") else {}
                self._tools_cache = openapi_to_tools(
                    spec, include_only=self._include_only, exclude=self._exclude
                )
                self._tools_cached_at = time.monotonic()
            return self._tools_cache

    def _cache_expired(self) -> bool:
        if self._cache_ttl is None:
            return False
        return (time.monotonic() - self._tools_cached_at) >= self._cache_ttl

    def invalidate_tools(self) -> None:
        self._tools_cache = None

    async def handle(
        self,
        request: dict[str, Any],
        *,
        extra_headers: list[tuple[bytes, bytes]] | None = None,
        client: tuple[str, int] | None = None,
    ) -> dict[str, Any] | None:
        """Process a single JSON-RPC message. ``None`` for valid notifications.

        ``extra_headers`` and ``client`` are forwarded to ``tools/call`` so the
        synthetic inner request carries the real caller's credentials and
        address (see :func:`hawkapi_mcp._dispatch.dispatch_tool`).
        """
        if not isinstance(request, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
            return _error(None, _INVALID_REQUEST, "Request must be a JSON object")
        if request.get("jsonrpc") != "2.0":
            return _error(request.get("id"), _INVALID_REQUEST, "jsonrpc must be '2.0'")
        method = request.get("method")
        msg_id = request.get("id")
        params = request.get("params") or {}
        if not isinstance(method, str):
            return _error(msg_id, _INVALID_REQUEST, "method must be a string")
        if not isinstance(params, dict):
            return _error(msg_id, _INVALID_PARAMS, "params must be an object")

        # Notification — no id, no response expected.
        is_notification = msg_id is None and "id" not in request

        try:
            result = await self._dispatch_method(
                method, params, extra_headers=extra_headers, client=client
            )
        except RpcError as exc:
            if is_notification:
                return None
            return _error(msg_id, exc.code, exc.message)
        except Exception as exc:  # noqa: BLE001
            if is_notification:
                return None
            return _error(msg_id, _INTERNAL_ERROR, f"Internal error: {exc!r}")

        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": msg_id, "result": result}

    async def _dispatch_method(
        self,
        method: str,
        params: dict[str, Any],
        *,
        extra_headers: list[tuple[bytes, bytes]] | None = None,
        client: tuple[str, int] | None = None,
    ) -> Any:
        if method == "initialize":
            return {
                "protocolVersion": MCP_PROTOCOL_VERSION,
                "capabilities": {"tools": {"listChanged": False}},
                "serverInfo": {"name": self._name, "version": self._version},
            }
        if method == "ping":
            return {}
        if method == "notifications/initialized":
            return {}
        if method == "tools/list":
            return {"tools": [_external_tool(t) for t in await self._tools()]}
        if method == "tools/call":
            return await self._call_tool(params, extra_headers=extra_headers, client=client)
        raise RpcError(_METHOD_NOT_FOUND, f"Unknown method: {method}")

    async def _call_tool(
        self,
        params: dict[str, Any],
        *,
        extra_headers: list[tuple[bytes, bytes]] | None = None,
        client: tuple[str, int] | None = None,
    ) -> dict[str, Any]:
        name = params.get("name")
        arguments = params.get("arguments") or {}
        if not isinstance(name, str):
            raise RpcError(_INVALID_PARAMS, "'name' is required")
        if not isinstance(arguments, dict):
            raise RpcError(_INVALID_PARAMS, "'arguments' must be an object")
        tools = await self._tools()
        tool = next((t for t in tools if t["name"] == name), None)
        if tool is None:
            raise RpcError(_METHOD_NOT_FOUND, f"Tool '{name}' not found")
        meta = tool["_meta"]
        status, body, headers = await dispatch_tool(
            self._app,
            meta["method"],
            meta["path"],
            arguments,
            client=client,
            extra_headers=extra_headers,
        )
        is_error = status >= 400
        return {
            "content": [
                {"type": "text", "text": _decode_body(body, headers)},
            ],
            "isError": is_error,
            "structuredContent": {
                "status": status,
                "headers": self._safe_headers(headers),
            },
        }

    def _safe_headers(self, headers: dict[str, str]) -> dict[str, str]:
        """Strip credential-bearing response headers before exposing to the agent."""
        return {
            k: v
            for k, v in headers.items()
            if k.lower() not in _SENSITIVE_HEADERS
            and k.lower() not in self._strip_response_headers
            and not _SENSITIVE_HEADER_PATTERN.match(k)
        }


# Backward-compatible alias — older code may import _RpcError from this module.
_RpcError = RpcError


def _error(msg_id: Any, code: int, message: str) -> dict[str, Any]:
    return {"jsonrpc": "2.0", "id": msg_id, "error": {"code": code, "message": message}}


def _external_tool(tool: dict[str, Any]) -> dict[str, Any]:
    return {
        "name": tool["name"],
        "description": tool["description"],
        "inputSchema": tool["inputSchema"],
    }


def _decode_body(body: bytes, headers: dict[str, str]) -> str:
    ctype = headers.get("content-type", "")
    if "application/json" in ctype:
        try:
            return json.dumps(json.loads(body.decode("utf-8")), ensure_ascii=False)
        except (UnicodeDecodeError, json.JSONDecodeError):
            return body.decode("utf-8", errors="replace")
    return body.decode("utf-8", errors="replace")


__all__ = ["MCPServer", "MCP_PROTOCOL_VERSION"]