Files
AIRegulation-DocAnalysis/docs/superpowers/plans/2026-07-29-mcp-search-regulations.md
T
wangweiandCopilot 49ee50c104 fix: harden MCP endpoint after code review
Critical: the MCP SDK auto-enables DNS-rebinding protection when its host
parameter is left at the 127.0.0.1 default, hard-coding a loopback-only Host
allow-list. Every remote client (the only deployment this feature targets) was
refused with HTTP 421 before auth or the tool ran. Now driven by a new
MCP_ALLOWED_HOSTS setting, with '*' as an explicit, logged opt-out.

Also bounds query/top_k to match AskRequest (top_k is amplified 4x downstream,
so an unbounded value was a resource-exhaustion vector), decodes the
Authorization header as latin-1 per the ASGI spec instead of raising a 500 on
malformed bytes, and returns WWW-Authenticate on 401 per RFC 7235.

Moves the psycopg2 import guard into backend/tests/conftest.py: duplicated
across four test modules, it only worked because of alphabetical collection
order, and any earlier-sorting package would have reintroduced a live
connection attempt against the production database.

Registers the mcp module in the authoritative backend architecture doc.

84 backend tests pass. Verified against a live server: allowed remote Host
returns a valid initialize result, unknown Host returns 421, missing token
returns 401 with WWW-Authenticate.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-29 17:11:54 +08:00

23 KiB
Raw Blame History

MCP Regulation Search Server — Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [x]) syntax for tracking.

Goal: Expose the existing compliance knowledge base as an MCP tool (search_regulations) via a standalone backend/app/mcp/ module, mounted into the existing FastAPI backend over Streamable HTTP, reusing existing JWT auth and the existing AgentConversationService.

Architecture: A new top-level backend/app/mcp/server.py builds a FastMCP instance with one @mcp.tool()-decorated search_regulations function that calls the existing get_agent_conversation_service().ask(...) (zero new business logic). A small MCPAuthMiddleware ASGI middleware validates the existing JWT bearer scheme in front of the mounted MCP app. backend/app/api/main.py mounts the resulting ASGI app at /mcp and wires its lifespan into the existing lifespan() function via AsyncExitStack (required — app.mount() does not propagate nested ASGI lifespans, so without this the MCP session manager never starts and every tool call fails).

Tech Stack: Python 3.12, FastAPI, Starlette, official mcp SDK (mcp.server.fastmcp.FastMCP), pytest, unittest.mock, Starlette TestClient.

Global Constraints

  • Design source of truth: docs/superpowers/specs/2026-07-29-mcp-search-regulations-design.md.
  • Correction discovered during implementation: the mcp package's latest release is 2.0.0, which renamed FastMCP to MCPServer (mcp.server.MCPServer) and its client helper streamablehttp_client to streamable_http_client. mcp>=2.0.0 is pinned in requirements.txt (not >=1.9.0 as originally estimated below) since the shipped code uses the MCPServer name. Also discovered: MCPServer.streamable_http_app() defaults to registering its route at /mcp, requiring streamable_http_path="/" to avoid a doubled /mcp/mcp when mounted at /mcp; the effective client URL is /mcp/ (trailing slash, due to Starlette's Mount redirect behavior).
  • All comments and docstrings in backend/**/*.py must be in English; every function/method needs a docstring; every file (including __init__.py) needs a module docstring + at least one meaningful # comment (AGENTS.md).
  • No new business orchestration — search_regulations is a thin protocol adapter over the existing AgentConversationService, same tier as app/api/routes/agent.py.
  • Python interpreter for this repo checkout: C:\software\Python312\python.exe (no .venv present in this checkout; this is the interpreter with all project dependencies already installed and is what the previous session's work was verified against).
  • Verified baseline test command (run from repo root, before any change in this plan): C:\software\Python312\python.exe -m pytest backend/tests -q69 passed (9.10s). Re-run this after every task.
  • Confirmed via direct import check: starlette is installed and starlette.testclient.TestClient works; the mcp package is not yet installed (ModuleNotFoundError: No module named 'mcp') — Task 3 installs it.
  • Latest published mcp version on PyPI at plan time: 2.0.0. Pin mcp>=1.9.0 in requirements (first version line with stable streamable_http_app() support) and let pip resolve to latest.

Task 1: search_regulations MCP tool

Files:

  • Create: backend/app/mcp/__init__.py
  • Create: backend/app/mcp/server.py (tool definition only — middleware and ASGI app builder added in Task 2)
  • Create: backend/tests/mcp/__init__.py
  • Create: backend/tests/mcp/test_search_regulations_tool.py

Interfaces:

  • Consumes: app.shared.bootstrap.get_agent_conversation_service() (existing, returns AgentConversationService).

  • Produces: app.mcp.server.search_regulations(query: str, top_k: int = 5) -> dict — a plain function (before the @mcp.tool() decorator is applied, it remains directly callable/testable; the decorator only adds MCP schema metadata, it does not change the function's Python call signature or return value).

  • Step 1: Write the failing test

Create backend/tests/mcp/__init__.py:

"""Test package for the MCP module (backend/app/mcp/)."""
# Empty package marker — no shared fixtures needed yet for this small test suite.

Create backend/tests/mcp/test_search_regulations_tool.py:

"""Unit tests for the search_regulations MCP tool function.

Mocks AgentConversationService so no real retrieval/LLM call happens —
verifies only the protocol-adapter contract: correct call shape in,
correct dict shape out.
"""

from __future__ import annotations

from dataclasses import dataclass
from unittest.mock import MagicMock, patch


@dataclass
class _FakeSource:
    """Minimal stand-in for a real Source dataclass (only __dict__ is used)."""

    doc_id: str
    doc_title: str
    score: float


@dataclass
class _FakeAnswerResult:
    """Minimal stand-in for AnswerResult — only .answer/.sources are read."""

    answer: str
    sources: list


def test_search_regulations_calls_agent_ask_without_session():
    """search_regulations must call ask() with no session_id (stateless search)."""
    from app.mcp.server import search_regulations

    fake_service = MagicMock()
    fake_service.ask.return_value = (
        None,
        _FakeAnswerResult(answer="国六排放标准要求...", sources=[_FakeSource("doc-1", "国六标准", 0.9)]),
    )

    with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
        result = search_regulations(query="国六排放标准最新要求", top_k=3)

    fake_service.ask.assert_called_once_with(query="国六排放标准最新要求", top_k=3)
    assert "session_id" not in fake_service.ask.call_args.kwargs


def test_search_regulations_shapes_response_dict():
    """The returned dict must expose 'answer' and 'sources' (list of plain dicts)."""
    from app.mcp.server import search_regulations

    fake_service = MagicMock()
    fake_service.ask.return_value = (
        None,
        _FakeAnswerResult(answer="答案文本", sources=[_FakeSource("doc-2", "国标GB1589", 0.8)]),
    )

    with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
        result = search_regulations(query="q")

    assert result == {
        "answer": "答案文本",
        "sources": [{"doc_id": "doc-2", "doc_title": "国标GB1589", "score": 0.8}],
    }


def test_search_regulations_default_top_k():
    """top_k defaults to 5 when the caller omits it."""
    from app.mcp.server import search_regulations

    fake_service = MagicMock()
    fake_service.ask.return_value = (None, _FakeAnswerResult(answer="a", sources=[]))

    with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
        search_regulations(query="q")

    assert fake_service.ask.call_args.kwargs["top_k"] == 5

Run it — confirm it fails on import (app.mcp.server does not exist yet):

C:\software\Python312\python.exe -m pytest backend/tests/mcp/test_search_regulations_tool.py -v
  • Step 2: Implement the tool

Create backend/app/mcp/__init__.py:

"""MCP (Model Context Protocol) server module.

Exposes selected read-only platform capabilities — currently only regulation
search — as MCP tools so external MCP clients (Claude Desktop, GitHub Copilot,
Cursor, etc.) can query this platform's compliance knowledge base directly.
"""
# Kept deliberately empty beyond this docstring — see server.py for the
# actual FastMCP instance and tool/middleware definitions.

Create backend/app/mcp/server.py:

"""FastMCP server 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

from mcp.server.fastmcp import FastMCP

from app.shared.bootstrap import get_agent_conversation_service

# Single shared FastMCP instance — analogous to the single shared FastAPI
# `app` instance in app/api/main.py. Tools registered via @mcp.tool() below.
mcp = FastMCP("ai-regulations")


@mcp.tool()
def search_regulations(query: str, top_k: int = 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 (default 5).
    """
    # No session_id is passed: this keeps each call stateless (no
    # ConversationStore reads/writes), matching "search" semantics rather
    # than multi-turn chat semantics.
    _, result = get_agent_conversation_service().ask(query=query, top_k=top_k)
    return {
        "answer": result.answer,
        "sources": [source.__dict__ for source in result.sources],
    }
  • Step 3: Run the test — confirm it passes
C:\software\Python312\python.exe -m pytest backend/tests/mcp/test_search_regulations_tool.py -v

Expected: 3 passed. Note: this step imports mcp.server.fastmcp, which is not yet installed — if it fails with ModuleNotFoundError: No module named 'mcp', that is expected until Task 3 installs the dependency; run C:\software\Python312\python.exe -m pip install "mcp>=1.9.0" locally first so this task's tests can actually execute now (Task 3 formalizes the requirements.txt entry — installing it now is just so this task's own tests are green before moving on).


Task 2: MCPAuthMiddleware — reuse existing JWT auth

Files:

  • Modify: backend/app/mcp/server.py (add middleware + ASGI app builder)
  • Create: backend/tests/mcp/test_mcp_auth_middleware.py

Interfaces:

  • Consumes: app.config.settings.settings.auth_enabled (existing), app.shared.bootstrap.get_jwt_handler() (existing, returns JWTHandler).

  • Produces: app.mcp.server.MCPAuthMiddleware (ASGI middleware class), app.mcp.server.build_mcp_asgi_app() -> ASGIApp (returns mcp.streamable_http_app() with the middleware already attached). Task 3's main.py change consumes build_mcp_asgi_app() directly — it does not need to attach the middleware itself.

  • Step 1: Write the failing test

Create backend/tests/mcp/test_mcp_auth_middleware.py:

"""Unit tests for MCPAuthMiddleware.

Wraps a minimal dummy ASGI app (not the real MCP app) so these tests exercise
only the auth gate, not the MCP protocol itself — keeps the test fast and
independent of FastMCP internals.
"""

from __future__ import annotations

from unittest.mock import patch

from starlette.applications import Starlette
from starlette.responses import PlainTextResponse
from starlette.routing import Route
from starlette.testclient import TestClient

from app.mcp.server import MCPAuthMiddleware


def _dummy_app() -> Starlette:
    """Build a minimal Starlette app that MCPAuthMiddleware can wrap."""
    async def _ok(request):
        """Return a fixed 200 response so tests can assert pass-through."""
        return PlainTextResponse("ok")

    app = Starlette(routes=[Route("/ping", _ok)])
    app.add_middleware(MCPAuthMiddleware)
    return app


def test_missing_token_rejected_when_auth_enabled():
    """No Authorization header + auth_enabled=True -> 401."""
    with patch("app.mcp.server.settings") as fake_settings:
        fake_settings.auth_enabled = True
        client = TestClient(_dummy_app())
        response = client.get("/ping")
    assert response.status_code == 401


def test_invalid_token_rejected_when_auth_enabled():
    """A token that fails decode_token() -> 401, request never reaches the app."""
    fake_handler = type("H", (), {"decode_token": lambda self, t: (_ for _ in ()).throw(ValueError("bad token"))})()
    with patch("app.mcp.server.settings") as fake_settings, \
         patch("app.mcp.server.get_jwt_handler", return_value=fake_handler):
        fake_settings.auth_enabled = True
        client = TestClient(_dummy_app())
        response = client.get("/ping", headers={"Authorization": "Bearer garbage"})
    assert response.status_code == 401


def test_valid_token_passes_through_when_auth_enabled():
    """A token that decodes successfully -> request reaches the wrapped app."""
    fake_handler = type("H", (), {"decode_token": lambda self, t: object()})()
    with patch("app.mcp.server.settings") as fake_settings, \
         patch("app.mcp.server.get_jwt_handler", return_value=fake_handler):
        fake_settings.auth_enabled = True
        client = TestClient(_dummy_app())
        response = client.get("/ping", headers={"Authorization": "Bearer good"})
    assert response.status_code == 200
    assert response.text == "ok"


def test_auth_disabled_always_passes_through():
    """auth_enabled=False (dev mode) -> no token needed, matches get_current_user's dev bypass."""
    with patch("app.mcp.server.settings") as fake_settings:
        fake_settings.auth_enabled = False
        client = TestClient(_dummy_app())
        response = client.get("/ping")
    assert response.status_code == 200

Run it — confirm it fails (MCPAuthMiddleware does not exist yet):

C:\software\Python312\python.exe -m pytest backend/tests/mcp/test_mcp_auth_middleware.py -v
  • Step 2: Implement the middleware and ASGI app builder

Append to backend/app/mcp/server.py:

from starlette.responses import PlainTextResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from app.config.settings import settings
from app.shared.bootstrap import get_jwt_handler


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"])
        auth_header = headers.get(b"authorization", b"").decode()
        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.
            response = PlainTextResponse(str(exc), status_code=401)
            await response(scope, receive, send)
            return

        await self.app(scope, receive, send)


def build_mcp_asgi_app() -> ASGIApp:
    """Return the Streamable HTTP ASGI app for the MCP server, auth-guarded."""
    asgi_app = mcp.streamable_http_app()
    asgi_app.add_middleware(MCPAuthMiddleware)
    return asgi_app
  • Step 3: Run the test — confirm it passes
C:\software\Python312\python.exe -m pytest backend/tests/mcp/test_mcp_auth_middleware.py -v

Expected: 4 passed.


Task 3: Mount into the FastAPI app

Files:

  • Modify: backend/requirements.txt (add mcp dependency)
  • Modify: backend/app/api/main.py (mount /mcp, wire lifespan via AsyncExitStack)

Interfaces:

  • Consumes: app.mcp.server.build_mcp_asgi_app() (from Task 2).

  • Produces: a running /mcp Streamable HTTP endpoint on the existing FastAPI app/port — no new port, process, or deployment step.

  • Step 1: Add the dependency

In backend/requirements.txt, add to the "Web framework" section (or a new small section — either is fine, keep it near fastapi/uvicorn since it is another transport-layer concern):

mcp>=1.9.0

Install it (already done ad hoc in Task 1 to unblock those tests — this step just formalizes the pin in the manifest; re-run install to be certain the pinned version resolves cleanly):

C:\software\Python312\python.exe -m pip install -r backend/requirements.txt
  • Step 2: Mount the MCP app and fix the lifespan gap

In backend/app/api/main.py, add the import and build the ASGI app at module scope (before lifespan() is defined, since lifespan() needs to reference it):

from contextlib import AsyncExitStack

from app.mcp.server import build_mcp_asgi_app

Add right after the existing imports, before @asynccontextmanager def lifespan(...):

# Built once at module scope so both lifespan() and app.mount() below reference
# the same instance — mounting a second, separately-built instance would start
# a second, unrelated MCP session manager.
mcp_app = build_mcp_asgi_app()

Replace the existing lifespan() function body:

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifecycle hooks."""
    # FastMCP's streamable_http_app() owns a session manager that only starts
    # via its own lifespan context. app.mount() does NOT propagate nested ASGI
    # lifespans automatically (confirmed Starlette/ASGI limitation — see
    # https://github.com/modelcontextprotocol/python-sdk/issues/1367) — without
    # this, every search_regulations call fails because the session manager
    # was never started.
    async with AsyncExitStack() as stack:
        await stack.enter_async_context(mcp_app.router.lifespan_context(mcp_app))

        logger.info(f"启动 {settings.app_name} v{settings.app_version}")
        logger.info(f"调试模式: {settings.debug}")
        logger.info("预加载LLM客户端...")
        preload_runtime_dependencies()

        yield

        logger.info("应用关闭,执行清理...")
        cleanup_runtime_dependencies()

Add the mount call right after the existing app.include_router(api_router, prefix="/api/v1") line:

app.include_router(api_router, prefix="/api/v1")
app.mount("/mcp", mcp_app)
  • Step 3: Run the full backend test suite
C:\software\Python312\python.exe -m pytest backend/tests -q

Expected: 76 passed (69 existing + 3 + 4 new from Tasks 12). No regressions.

  • Step 4: Manual end-to-end verification (not automated)

Start the backend the normal way (dev.bat start api --foreground or the documented uvicorn command) and, from a separate shell, run:

C:\software\Python312\python.exe -c "
import asyncio
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

async def main():
    url = 'http://127.0.0.1:8000/mcp'
    headers = {'Authorization': 'Bearer <put a real JWT here if AUTH_ENABLED=true>'}
    async with streamablehttp_client(url, headers=headers) as (read, write, _):
        async with ClientSession(read, write) as session:
            await session.initialize()
            tools = await session.list_tools()
            print([t.name for t in tools.tools])
            result = await session.call_tool('search_regulations', {'query': '国六排放标准'})
            print(result)

asyncio.run(main())
"

Confirm search_regulations appears in the tool list and returns a real answer + sources. If AUTH_ENABLED=false locally, omit the headers argument entirely.


Summary

Task New files Modified files Tests added
1 app/mcp/__init__.py, app/mcp/server.py, tests/mcp/__init__.py, tests/mcp/test_search_regulations_tool.py 3
2 tests/mcp/test_mcp_auth_middleware.py app/mcp/server.py 4
3 requirements.txt, api/main.py 0 (full-suite regression check + manual e2e)

Task 4 (post-review): code-review fixes

Added after a code review of the three implementation commits. Findings and resolutions:

  • Critical — every remote client rejected with HTTP 421. The MCP SDK auto-enables DNS-rebinding protection when its host parameter is left at the 127.0.0.1 default, hard-coding a loopback-only Host allow-list; a client at http://6.86.80.9:8000/mcp/ was refused before auth or the tool ran, making the feature non-functional in the only deployment it targets. Fixed by passing an explicit TransportSecuritySettings built from a new MCP_ALLOWED_HOSTS setting (app/config/settings.py, documented in .env.example), with * as a logged, explicit opt-out. Deliberately not fixed by passing host="0.0.0.0", which would silently disable the protection.
  • Important — top_k unbounded on the MCP path. AskRequest constrains the same parameter to 120, but the tool accepted any integer and KnowledgeRetrievalService amplifies it (top_k * 4), so top_k=100000 would request 400,000 Milvus candidates. Fixed with Annotated[int, Field(ge=1, le=20)] (and query bounded to 12000 chars), which also publishes the bounds in the advertised JSON schema.
  • Important — order-dependent psycopg2 test guard. The guard was duplicated across four test modules and only worked because of pytest's alphabetical collection order; any new test package sorting earlier would have reintroduced a multi-second TCP timeout against the production database. Moved into a single backend/tests/conftest.py (imported before any test module regardless of order) and the four in-file copies deleted. bootstrap.py's eager imports were left alone — restructuring the composition root every route depends on is disproportionate to a test-harness ordering problem.
  • Minor — non-UTF-8 Authorization header caused a 500. ASGI header values are latin-1; strict UTF-8 decoding let any remote client trigger an unhandled UnicodeDecodeError. Now decoded as latin-1, yielding a clean 401.
  • Minor — 401 missing WWW-Authenticate. Added WWW-Authenticate: Bearer, matching get_current_user and RFC 7235.
  • Minor — missing # comment in tests/mcp/test_search_regulations_tool.py (AGENTS.md requires at least one per file). Added.

Reviewer-confirmed as correct, no change needed: the AsyncExitStack lifespan wiring (including its failure path), the absence of auth-bypass vectors, and the statelessness of ask() without a session_id.

New tests: tests/mcp/test_mcp_transport_security.py (5, exercising the real MCP app end-to-end) plus 3 more across the existing two files — 84 backend tests pass. Verified against a live server: Host: 6.86.80.9:8000 → 200 with a valid initialize result, Host: evil.example.com → 421, no token → 401 with WWW-Authenticate: Bearer.