doc-forge/docforge/servers/mcp_server.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Literal

from mcp.server.fastmcp import FastMCP


class MCPServer:
    """
    MCP server for serving a pre-generated documentation bundle.

    The server exposes documentation resources and diagnostic tools through
    MCP endpoints backed by JSON files generated by the MCP renderer.
    """

    def __init__(self, mcp_root: Path, name: str) -> None:
        """
        Initialize the MCP server.

        Args:
            mcp_root: Directory containing the generated MCP documentation
                bundle (for example ``index.json``, ``nav.json``, and
                ``modules/``).
            name: Identifier used for the MCP server instance.
        """
        self.mcp_root = mcp_root
        self.app = FastMCP(name)

        self._register_resources()
        self._register_tools()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _read_json(self, path: Path) -> Any:
        """
        Load and parse a JSON file.

        If the file does not exist, a structured error response is returned
        instead of raising an exception.

        Args:
            path: Path to the JSON file to read.

        Returns:
            Parsed JSON data if the file exists, otherwise an error dictionary
            describing the missing resource.
        """
        if not path.exists():
            return {
                "error": "not_found",
                "path": str(path),
            }

        return json.loads(path.read_text(encoding="utf-8"))

    # ------------------------------------------------------------------
    # MCP resources
    # ------------------------------------------------------------------

    def _register_resources(self) -> None:
        """
        Register MCP resource endpoints.

        The server exposes documentation resources through the following
        endpoints:

        - ``docs://index`` – Project metadata
        - ``docs://nav`` – Navigation structure
        - ``docs://modules/{module}`` – Individual module documentation
        """

        @self.app.resource("docs://index")
        def index():
            return self._read_json(self.mcp_root / "index.json")

        @self.app.resource("docs://nav")
        def nav():
            return self._read_json(self.mcp_root / "nav.json")

        @self.app.resource("docs://modules/{module}")
        def module(module: str):
            return self._read_json(
                self.mcp_root / "modules" / f"{module}.json"
            )

    # ------------------------------------------------------------------
    # MCP tools
    # ------------------------------------------------------------------

    def _register_tools(self) -> None:
        """
        Register optional MCP diagnostic tools.

        These tools provide lightweight endpoints useful for verifying that
        the MCP server is operational.
        """

        @self.app.tool()
        def ping() -> str:
            """Return a simple health check response."""
            return "pong"

    # ------------------------------------------------------------------
    # Server lifecycle
    # ------------------------------------------------------------------

    def run(
        self,
        transport: Literal["stdio", "sse", "streamable-http"] = "streamable-http",
    ) -> None:
        """
        Start the MCP server.

        Args:
            transport: Transport mechanism used by the MCP server. Supported
                options include ``stdio``, ``sse``, and ``streamable-http``.
        """
        self.app.run(transport=transport)