mcp-server/common/mcp_service.py

# coding=utf-8
"""
@Time   :2026/6/30
@Author :tao.chen

The MCP service registry: a single place that knows which FastAPI sub-apps
to mount under the root app and how to gate them by the MCP_SERVICES env
var. Each service package declares itself via a `service.py` exporting a
module-level `SERVICE: McpService` instance.
"""
from __future__ import annotations

import importlib
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING

from common.logging import logger

if TYPE_CHECKING:
    from fastapi import FastAPI


# The codebase standardizes on loguru (see common.logging), but pytest's
# `caplog` fixture captures stdlib `logging` records, not loguru. We mirror
# the single skip-warning through a stdlib logger named `common.logging`
# (matching how `caplog.at_level(..., logger="common.logging")` is invoked
# in the tests) so the warning is assertable. A NullHandler keeps the
# stdlib logger silent in production; `caplog` attaches its own handler.
import logging as _logging
_stdlib_logger = _logging.getLogger("common.logging")
if not _stdlib_logger.handlers:
    _stdlib_logger.addHandler(_logging.NullHandler())
_stdlib_logger.propagate = False


@dataclass(frozen=True)
class McpService:
    """One MCP-exposed FastAPI sub-app to be mounted under the root."""

    name: str          # short id used in startup logs and the MCP_SERVICES env whitelist
    app: "FastAPI"     # the sub-app whose routes fastapi-mcp exposes as tools
    mount_path: str    # HTTP path under the root app, e.g. "/spark-executor-mcp"


# Every known MCP service's import path. Add a new service by appending one
# line here and creating <package>/service.py that exports SERVICE.
BUILTIN_SERVICES: list[str] = [
    "spark_executor.service",
    "files_mcp.service",
]


def _filter_by_env(services: list[McpService]) -> list[McpService]:
    """Whitelist by `MCP_SERVICES` env. Unset = no filter (backward compat).
    Empty string = zero services (explicit operator choice, distinct from unset)."""
    if "MCP_SERVICES" not in os.environ:
        return services
    raw = os.environ["MCP_SERVICES"]
    if not raw:
        return []
    allowed = {n.strip() for n in raw.split(",") if n.strip()}
    return [s for s in services if s.name in allowed]


def discover_and_filter() -> list[McpService]:
    """Import each BUILTIN_SERVICES module, read its SERVICE attribute,
    then apply the env whitelist.

    Raises:
        RuntimeError: see `Error Handling` in the spec.
    """
    services: list[McpService] = []
    raw = os.environ.get("MCP_SERVICES")
    requested = {n.strip() for n in raw.split(",") if n.strip()} if raw else set()
    origins: dict[str, str] = {}  # service.name -> import_path (for diagnostics)

    for import_path in BUILTIN_SERVICES:
        try:
            module = importlib.import_module(import_path)
        except ImportError as e:
            pkg_name = import_path.rsplit(".", 1)[0]
            if pkg_name in requested or import_path in requested:
                raise RuntimeError(
                    f"MCP service {import_path!r} is requested via MCP_SERVICES "
                    f"but cannot be imported: {e}"
                ) from e
            logger.warning(f"Skipping missing MCP service {import_path!r}: {e}")
            _stdlib_logger.warning(f"Skipping missing MCP service {import_path!r}: {e}")
            continue
        svc = getattr(module, "SERVICE", None)
        if svc is None:
            raise RuntimeError(
                f"{import_path!r} does not export a SERVICE attribute"
            )
        if not isinstance(svc, McpService):
            raise RuntimeError(
                f"{import_path!r}.SERVICE is not a McpService "
                f"(got {type(svc).__name__})"
            )
        services.append(svc)
        origins[svc.name] = import_path

    # Duplicate detection - surface typos at startup.
    seen_names: dict[str, str] = {}
    seen_mounts: dict[str, str] = {}
    for s in services:
        if s.name in seen_names:
            raise RuntimeError(
                f"duplicate MCP service name {s.name!r}: "
                f"{seen_names[s.name]!r} and {origins[s.name]!r}"
            )
        seen_names[s.name] = origins[s.name]
        if s.mount_path in seen_mounts:
            raise RuntimeError(
                f"duplicate MCP mount_path {s.mount_path!r}: "
                f"{seen_mounts[s.mount_path]!r} and {origins[s.name]!r}"
            )
        seen_mounts[s.mount_path] = origins[s.name]

    filtered = _filter_by_env(services)

    found_names = {s.name for s in services}
    missing = requested - found_names
    if missing:
        raise RuntimeError(
            f"MCP_SERVICES env references unknown service names: "
            f"{sorted(missing)}. Available: {sorted(found_names)}"
        )

    return filtered