AIRegulation-DocAnalysis/backend/app/mcp/server.py

"""MCPServer instance exposing the compliance knowledge base as an MCP tool.

This module is a pure protocol adapter: search_regulations() below calls the
existing AgentConversationService.ask() (the same application service backing
the /api/v1/agent/ask REST endpoint) and reshapes its result into a plain
dict. No new retrieval, ranking, or LLM orchestration logic lives here.
"""

from __future__ import annotations

import logging
import time
from typing import Annotated

from mcp.server import MCPServer
from mcp.server.transport_security import TransportSecuritySettings
from pydantic import Field
from starlette.responses import PlainTextResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from app.config.settings import settings
from app.mcp.stats import get_mcp_stats_tracker
from app.shared.bootstrap import get_agent_conversation_service, get_jwt_handler

logger = logging.getLogger(__name__)

# Single shared MCPServer instance — analogous to the single shared FastAPI
# `app` instance in app/api/main.py. Tools registered via @mcp.tool() below.
# Note: the installed mcp SDK (2.0.0) renamed the older "FastMCP" class to
# "MCPServer" (mcp.server.mcpserver.MCPServer); the .tool()/.streamable_http_app()
# API surface used here is unchanged across that rename.
mcp = MCPServer("ai-regulations")


@mcp.tool()
def search_regulations(
    query: Annotated[str, Field(min_length=1, max_length=2000)],
    top_k: Annotated[int, Field(ge=1, le=20)] = 5,
) -> dict:
    """Search the compliance knowledge base and return a grounded answer.

    query: Natural-language search question, e.g. "国六排放标准最新要求".
    top_k: Maximum number of cited sources to return (1-20, default 5).
    """
    # Bounds mirror AskRequest in app/api/models/agent.py so the MCP path cannot
    # be used to bypass the REST endpoint's limits. They matter more here than
    # there: KnowledgeRetrievalService amplifies top_k (candidate_k = top_k * 4)
    # when reranking, and an LLM client can easily hallucinate a huge value.
    # Declaring them via Annotated puts them in the advertised JSON schema too,
    # so well-behaved clients never send an out-of-range value in the first place.
    #
    # No session_id is passed: this keeps each call stateless (no
    # ConversationStore reads/writes), matching "search" semantics rather
    # than multi-turn chat semantics.
    started = time.perf_counter()
    try:
        _, result = get_agent_conversation_service().ask(query=query, top_k=top_k)
    except Exception:
        # Record the failure, then re-raise unchanged so the MCP SDK still
        # converts it into a protocol-level error for the client. Swallowing
        # it here would report success to the caller.
        get_mcp_stats_tracker().record(
            tool="search_regulations",
            duration_ms=(time.perf_counter() - started) * 1000,
            success=False,
        )
        raise
    get_mcp_stats_tracker().record(
        tool="search_regulations",
        duration_ms=(time.perf_counter() - started) * 1000,
        success=True,
    )
    return {
        "answer": result.answer,
        "sources": [source.__dict__ for source in result.sources],
    }


class MCPAuthMiddleware:
    """Reject unauthenticated requests before they reach the MCP protocol handler.

    Mirrors the existing get_current_user dependency's behavior (auth.py) but
    implemented as raw ASGI middleware, since the mounted MCP app is a plain
    ASGI app, not a FastAPI/APIRouter instance that supports Depends().
    """

    def __init__(self, app: ASGIApp) -> None:
        """Store the wrapped ASGI app to delegate to once auth passes."""
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Validate the bearer token for HTTP requests; pass non-HTTP scopes through."""
        # Only HTTP requests carry an Authorization header to check; lifespan
        # and other scope types must always pass through untouched.
        if scope["type"] != "http" or not settings.auth_enabled:
            await self.app(scope, receive, send)
            return

        headers = dict(scope["headers"])
        # ASGI header values are raw bytes specified as latin-1, not UTF-8;
        # decoding strictly as UTF-8 would raise on a malformed byte and turn a
        # bad request into an unhandled 500.
        auth_header = headers.get(b"authorization", b"").decode("latin-1")
        token = auth_header.removeprefix("Bearer ").strip()
        try:
            get_jwt_handler().decode_token(token)
        except ValueError as exc:
            # Reject before the MCP session/protocol layer ever sees the request.
            # WWW-Authenticate matches the get_current_user dependency (auth.py)
            # and is required by RFC 7235 so clients can tell "needs credentials"
            # apart from a generic failure.
            response = PlainTextResponse(
                str(exc), status_code=401, headers={"WWW-Authenticate": "Bearer"}
            )
            await response(scope, receive, send)
            return

        await self.app(scope, receive, send)


def _parse_allowed_hosts() -> list[str]:
    """Split the configured MCP host allow-list into individual entries."""
    # Shared by the transport-security builder and the status endpoint so the
    # panel can never display an allow-list different from the enforced one.
    return [h.strip() for h in settings.mcp_allowed_hosts.split(",") if h.strip()]


def _build_transport_security() -> TransportSecuritySettings:
    """Translate the configured MCP host allow-list into SDK transport settings.

    Without this the SDK infers its own allow-list from the bind host, which
    defaults to 127.0.0.1 and therefore rejects every remote client with HTTP
    421 — fatal for a remotely deployed backend.
    """
    allowed = _parse_allowed_hosts()
    if "*" in allowed:
        # Explicit, logged opt-out. Kept as an escape hatch for environments
        # behind a proxy that rewrites Host unpredictably, but never the default.
        logger.warning(
            "MCP DNS-rebinding protection is disabled (mcp_allowed_hosts='*'). "
            "Set MCP_ALLOWED_HOSTS to the real deployment host(s) instead."
        )
        return TransportSecuritySettings(enable_dns_rebinding_protection=False)
    return TransportSecuritySettings(
        enable_dns_rebinding_protection=True,
        allowed_hosts=allowed,
        # Browser clients send Origin; reuse the already-maintained CORS list so
        # there is one place to declare trusted web origins. Non-browser MCP
        # clients send no Origin at all, which the SDK treats as allowed.
        allowed_origins=[o.strip() for o in settings.cors_allow_origins.split(",") if o.strip()],
    )


def build_mcp_asgi_app() -> ASGIApp:
    """Return the Streamable HTTP ASGI app for the MCP server, auth-guarded.

    streamable_http_path="/" is required here: MCPServer.streamable_http_app()
    registers its own internal route at "/mcp" by default, and this app is
    itself mounted at "/mcp" in api/main.py — without overriding the internal
    path to "/", the effective external path would be the confusing "/mcp/mcp"
    instead of "/mcp".
    """
    asgi_app = mcp.streamable_http_app(
        streamable_http_path="/",
        transport_security=_build_transport_security(),
    )
    asgi_app.add_middleware(MCPAuthMiddleware)
    return asgi_app


async def get_mcp_status(public_url: str) -> dict:
    """Assemble the MCP status payload shown on the System Status page.

    Owned by this module rather than the status route so that MCP internals
    (the tool registry, the allow-list format, the stats tracker) stay behind
    one boundary; the route only supplies public_url, which is the one value
    only the HTTP layer can know.
    """
    stats = get_mcp_stats_tracker().snapshot()
    # list_tools() reads the in-memory registry populated by @mcp.tool() at
    # import time, so the panel always reflects what is actually advertised
    # rather than a hand-maintained duplicate list.
    tools = await mcp.list_tools()
    return {
        "endpoint_url": public_url,
        "auth_required": settings.auth_enabled,
        "allowed_hosts": _parse_allowed_hosts(),
        "tools": [
            {
                "name": tool.name,
                "description": (tool.description or "").strip().split("\n")[0],
                "calls": entry.calls if entry else 0,
                "errors": entry.errors if entry else 0,
                "avg_duration_ms": entry.avg_duration_ms if entry else None,
                "last_called_at": (
                    entry.last_called_at.isoformat() if entry and entry.last_called_at else None
                ),
            }
            for tool, entry in ((tool, stats.get(tool.name)) for tool in tools)
        ],
    }