AIRegulation-DocAnalysis/docs/superpowers/plans/2026-07-29-mcp-search-regulations.md

# 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 -q` → `69 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).

- [x] **Step 1: Write the failing test**

Create `backend/tests/mcp/__init__.py`:

```python
"""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`:

```python
"""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):

```powershell
C:\software\Python312\python.exe -m pytest backend/tests/mcp/test_search_regulations_tool.py -v
```

- [x] **Step 2: Implement the tool**

Create `backend/app/mcp/__init__.py`:

```python
"""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`:

```python
"""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],
    }
```

- [x] **Step 3: Run the test — confirm it passes**

```powershell
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.

- [x] **Step 1: Write the failing test**

Create `backend/tests/mcp/test_mcp_auth_middleware.py`:

```python
"""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):

```powershell
C:\software\Python312\python.exe -m pytest backend/tests/mcp/test_mcp_auth_middleware.py -v
```

- [x] **Step 2: Implement the middleware and ASGI app builder**

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

```python
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
```

- [x] **Step 3: Run the test — confirm it passes**

```powershell
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.

- [x] **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):

```powershell
C:\software\Python312\python.exe -m pip install -r backend/requirements.txt
```

- [x] **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):

```python
from contextlib import AsyncExitStack

from app.mcp.server import build_mcp_asgi_app
```

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

```python
# 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:

```python
@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:

```python
app.include_router(api_router, prefix="/api/v1")
app.mount("/mcp", mcp_app)
```

- [x] **Step 3: Run the full backend test suite**

```powershell
C:\software\Python312\python.exe -m pytest backend/tests -q
```

Expected: `76 passed` (69 existing + 3 + 4 new from Tasks 1–2). No regressions.

- [x] **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:

```powershell
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:

- [x] **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.
- [x] **Important — `top_k` unbounded on the MCP path.** `AskRequest` constrains the same parameter to 1–20, 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 1–2000 chars), which also publishes the bounds in the advertised JSON schema.
- [x] **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.
- [x] **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.
- [x] **Minor — 401 missing `WWW-Authenticate`.** Added `WWW-Authenticate: Bearer`, matching `get_current_user` and RFC 7235.
- [x] **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`.