feat: add MCP server module exposing search_regulations tool

- New backend/app/mcp/ module: MCPServer instance with a single
  search_regulations tool backed by the existing AgentConversationService.
- MCPAuthMiddleware reuses existing JWT auth (no new auth mechanism).
- Mounted at /mcp/ in api/main.py via Streamable HTTP transport; wired the
  MCP session manager into the existing lifespan() via AsyncExitStack
  (app.mount() does not propagate nested ASGI lifespans automatically).
- Fixed a doubled /mcp/mcp path by setting streamable_http_path to "/"
  (MCPServer.streamable_http_app() defaults to registering its own /mcp route).
- Verified end-to-end with the real mcp Python client: list_tools() returns
  search_regulations, auth correctly 401s without or with an invalid token.
- 7 new tests, 76 total (up from 69), all passing.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-29 13:00:52 +08:00
co-authored by Copilot
parent e78c8a989f
commit bd3dc38d1d
9 changed files with 302 additions and 20 deletions
+23 -8
View File
@@ -1,6 +1,6 @@
"""FastAPI application entrypoint."""
from contextlib import asynccontextmanager
from contextlib import AsyncExitStack, asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.encoders import jsonable_encoder
@@ -13,6 +13,7 @@ from app.api.models import ErrorResponse
from app.api.routes import api_router
from app.config.logging import setup_logging
from app.config.settings import settings
from app.mcp.server import build_mcp_asgi_app
from app.shared.bootstrap import cleanup_runtime_dependencies, preload_runtime_dependencies
from app.shared.errors import VectorStoreSchemaError
# Keep module behavior explicit so the backend flow stays easy to audit.
@@ -20,19 +21,32 @@ from app.shared.errors import VectorStoreSchemaError
setup_logging(level="INFO" if not settings.debug else "DEBUG")
# 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()
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Application lifecycle hooks."""
logger.info(f"启动 {settings.app_name} v{settings.app_version}")
logger.info(f"调试模式: {settings.debug}")
logger.info("预加载LLM客户端...")
preload_runtime_dependencies()
# FastMCP-style servers own a session manager that only starts via its own
# lifespan context. app.mount() does NOT propagate nested ASGI lifespans
# automatically (confirmed Starlette/ASGI limitation) — without this,
# every search_regulations call would fail because the MCP session
# manager was never started.
async with AsyncExitStack() as stack:
await stack.enter_async_context(mcp_app.router.lifespan_context(mcp_app))
yield
logger.info(f"启动 {settings.app_name} v{settings.app_version}")
logger.info(f"调试模式: {settings.debug}")
logger.info("预加载LLM客户端...")
preload_runtime_dependencies()
logger.info("应用关闭,执行清理...")
cleanup_runtime_dependencies()
yield
logger.info("应用关闭,执行清理...")
cleanup_runtime_dependencies()
app = FastAPI(
@@ -65,6 +79,7 @@ app.add_middleware(
app.add_middleware(AuditMiddleware)
app.include_router(api_router, prefix="/api/v1")
app.mount("/mcp", mcp_app)
@app.exception_handler(VectorStoreSchemaError)
+8
View File
@@ -0,0 +1,8 @@
"""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.
+88
View File
@@ -0,0 +1,88 @@
"""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
from mcp.server import MCPServer
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_agent_conversation_service, get_jwt_handler
# 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: 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],
}
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.
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="/")
asgi_app.add_middleware(MCPAuthMiddleware)
return asgi_app
+4
View File
@@ -2,6 +2,10 @@
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
python-multipart>=0.0.9
# MCP server module (backend/app/mcp/) — pin >=2.0.0: that release renamed the
# older "FastMCP" class to "MCPServer" (mcp.server.MCPServer), which is the
# class actually used in app/mcp/server.py.
mcp>=2.0.0
# ── Config & utilities ────────────────────────────────────────────────────────
pydantic>=2.0.0
+2
View File
@@ -0,0 +1,2 @@
"""Test package for the MCP module (backend/app/mcp/)."""
# Empty package marker — no shared fixtures needed yet for this small test suite.
@@ -0,0 +1,85 @@
"""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
import sys
from unittest.mock import MagicMock, patch
# app.mcp.server imports app.shared.bootstrap at module scope (needed for
# get_agent_conversation_service/get_jwt_handler), which in turn eagerly
# imports several Postgres store modules that do `import psycopg2` at their
# own module scope. Since this is the only test file in backend/tests/mcp/
# that imports app.mcp.server at module scope (not inside a test function),
# it is the first thing to trigger that chain during pytest collection —
# guard psycopg2 here the same way backend/tests/observability/
# test_model_usage_bootstrap.py and test_model_usage_persistence.py already
# do, so a real (network-connecting) psycopg2 never gets bound first.
mock_psycopg2 = MagicMock()
mock_psycopg2.extras = MagicMock()
sys.modules.setdefault("psycopg2", mock_psycopg2)
sys.modules.setdefault("psycopg2.extras", mock_psycopg2.extras)
sys.modules.setdefault("psycopg2.pool", MagicMock())
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
@@ -0,0 +1,77 @@
"""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):
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