quality-inspection-platform/app/mcp_native.py

"""Native MCP Server (Streamable HTTP) mounted on the FastAPI app."""

from __future__ import annotations

import json
import logging
import os
from contextlib import asynccontextmanager
from typing import Any

from fastapi import FastAPI, HTTPException, Request
from mcp import types
from mcp.server.lowlevel.server import Server
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.server.transport_security import TransportSecuritySettings
from sqlalchemy.orm import Session
from starlette.requests import Request as StarletteRequest

from .database import SessionLocal

logger = logging.getLogger(__name__)

SERVER_NAME = "quality-inspection-platform"
SERVER_VERSION = "0.2.0"
MCP_INSTRUCTIONS = (
    "Quality inspection platform (质量检测平台). Before complex tasks, read project file "
    "docs/mcp_tools_for_ai.md (tool picker, auth, recipes). "
    "Typical start: catalog_snapshot. "
    "Batch runs: workflow_batch_create → workflow_batch_update → workflow_batch_run. "
    "Writes and execution require sto- API Key (Authorization: Bearer or X-API-Key)."
)

_mcp_server = Server(SERVER_NAME, version=SERVER_VERSION, instructions=MCP_INSTRUCTIONS)
_session_manager: StreamableHTTPSessionManager | None = None


def _mcp_transport_security() -> TransportSecuritySettings | None:
    if os.getenv("QIP_MCP_DISABLE_DNS_PROTECTION", "").strip().lower() in {"1", "true", "yes", "on"}:
        return TransportSecuritySettings(enable_dns_rebinding_protection=False)
    hosts = [item.strip() for item in os.getenv("QIP_MCP_ALLOWED_HOSTS", "").split(",") if item.strip()]
    origins = [item.strip() for item in os.getenv("QIP_MCP_ALLOWED_ORIGINS", "").split(",") if item.strip()]
    if not hosts and not origins:
        return None
    return TransportSecuritySettings(
        enable_dns_rebinding_protection=True,
        allowed_hosts=hosts,
        allowed_origins=origins,
    )


def _tool_specs() -> list[dict[str, Any]]:
    from .main import MCP_TOOL_SPECS

    return MCP_TOOL_SPECS


def _specs_to_mcp_tools() -> list[types.Tool]:
    tools: list[types.Tool] = []
    for spec in _tool_specs():
        schema = spec.get("input_schema") or {"type": "object"}
        tools.append(
            types.Tool(
                name=str(spec.get("name", "")),
                description=str(spec.get("description", "")),
                inputSchema=schema,
            )
        )
    return tools


def _starlette_request() -> StarletteRequest | None:
    ctx = _mcp_server.request_context
    request = ctx.request
    return request if isinstance(request, StarletteRequest) else None


def _resolve_actor_for_tool(tool: str, request: StarletteRequest, db: Session):
    from .main import _mcp_global_require_api_key, _mcp_tool_requires_credential, _resolve_mcp_actor

    require = _mcp_global_require_api_key() or _mcp_tool_requires_credential(tool)
    return _resolve_mcp_actor(request, db, require_credential=require)


@_mcp_server.list_tools()
async def _handle_list_tools() -> list[types.Tool]:
    request = _starlette_request()
    if request is None:
        return _specs_to_mcp_tools()
    db = SessionLocal()
    try:
        from .main import _mcp_global_require_api_key

        if _mcp_global_require_api_key():
            _resolve_actor_for_tool("catalog_snapshot", request, db)
        return _specs_to_mcp_tools()
    finally:
        db.close()


@_mcp_server.call_tool()
async def _handle_call_tool(
    name: str,
    arguments: dict[str, Any] | None,
) -> types.CallToolResult:
    request = _starlette_request()
    if request is None:
        raise ValueError("MCP request context missing HTTP request")

    db = SessionLocal()
    try:
        try:
            actor = _resolve_actor_for_tool(name, request, db)
            from .main import _dispatch_mcp_tool

            result = await _dispatch_mcp_tool(name, arguments or {}, db, actor)
        except HTTPException as exc:
            payload = {"ok": False, "tool": name, "data": {}, "error": str(exc.detail)}
            return types.CallToolResult(
                content=[types.TextContent(type="text", text=json.dumps(payload, ensure_ascii=False, indent=2))],
                isError=True,
            )
        text = json.dumps(result.model_dump(), ensure_ascii=False, indent=2)
        return types.CallToolResult(
            content=[types.TextContent(type="text", text=text)],
            isError=not result.ok,
        )
    finally:
        db.close()


def _get_session_manager() -> StreamableHTTPSessionManager:
    global _session_manager
    if _session_manager is None:
        _session_manager = StreamableHTTPSessionManager(
            app=_mcp_server,
            json_response=False,
            stateless=True,
            security_settings=_mcp_transport_security(),
        )
    return _session_manager


@asynccontextmanager
async def mcp_lifespan(_app: FastAPI):
    manager = _get_session_manager()
    async with manager.run():
        yield


def mount_mcp_streamable_http(fastapi_app: FastAPI, *, path: str = "/mcp") -> None:
    """Register Streamable HTTP MCP endpoint on the FastAPI app."""
    manager = _get_session_manager()

    @fastapi_app.api_route(path, methods=["GET", "POST", "DELETE"], include_in_schema=False)
    async def mcp_streamable_http_endpoint(request: Request) -> None:
        await manager.handle_request(request.scope, request.receive, request._send)  # type: ignore[attr-defined]

    logger.info("MCP Streamable HTTP mounted at %s (stateless)", path)