feat: surface MCP server status in System Status page
Add per-tool in-memory call counters to the MCP module and a GET /api/v1/status/mcp endpoint that joins them with the live tool registry and endpoint config, then render it as a new card on the System Status page with a one-click client-config copy button. - app/mcp/stats.py: lock-guarded MCPStatsTracker (the mcp SDK runs sync tool bodies via anyio.to_thread.run_sync, so this is genuinely multi-threaded, unlike the async REST routes) - app/mcp/server.py: instrument search_regulations, add get_mcp_status() - app/config/settings.py: optional MCP_PUBLIC_URL override, required because the Vite proxy and reverse proxies rewrite the Host header - StatusPage.tsx: MCP Server card, joins the existing parallel fetch Counters are process-local by design; token usage is already persisted by ModelUsageTracker since MCP calls route through ask(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
@@ -9,6 +9,7 @@ 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
|
||||
@@ -18,6 +19,7 @@ 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__)
|
||||
@@ -50,7 +52,24 @@ def search_regulations(
|
||||
# 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)
|
||||
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],
|
||||
@@ -99,6 +118,13 @@ class MCPAuthMiddleware:
|
||||
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.
|
||||
|
||||
@@ -106,7 +132,7 @@ def _build_transport_security() -> TransportSecuritySettings:
|
||||
defaults to 127.0.0.1 and therefore rejects every remote client with HTTP
|
||||
421 — fatal for a remotely deployed backend.
|
||||
"""
|
||||
allowed = [h.strip() for h in settings.mcp_allowed_hosts.split(",") if h.strip()]
|
||||
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.
|
||||
@@ -140,3 +166,36 @@ def build_mcp_asgi_app() -> ASGIApp:
|
||||
)
|
||||
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)
|
||||
],
|
||||
}
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
"""In-memory per-tool call counters for the MCP server.
|
||||
|
||||
Lives in `app/mcp/` rather than `app/shared/` because these counters are
|
||||
meaningful only for the MCP transport: they answer "is anything actually
|
||||
calling our MCP endpoint, and does it work?" for the System Status page.
|
||||
Token consumption is deliberately not tracked here — MCP tool calls route
|
||||
through AgentConversationService.ask() like every other caller, so the
|
||||
existing ModelUsageTracker already accounts for it.
|
||||
|
||||
Counters are process-local and reset on restart. That is an accepted
|
||||
tradeoff, recorded in the design spec: nothing billable depends on them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from functools import lru_cache
|
||||
|
||||
from loguru import logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPToolStats:
|
||||
"""Accumulated call outcomes for a single MCP tool."""
|
||||
|
||||
calls: int = 0
|
||||
errors: int = 0
|
||||
total_duration_ms: float = 0.0
|
||||
last_called_at: datetime | None = None
|
||||
|
||||
@property
|
||||
def avg_duration_ms(self) -> float | None:
|
||||
"""Mean call duration, or None when the tool has never been called.
|
||||
|
||||
Returning None rather than 0.0 keeps "never called" distinguishable
|
||||
from "called, but instantaneous" in the status UI.
|
||||
"""
|
||||
if self.calls == 0:
|
||||
return None
|
||||
return self.total_duration_ms / self.calls
|
||||
|
||||
|
||||
class MCPStatsTracker:
|
||||
"""Thread-safe registry of per-tool MCP call statistics.
|
||||
|
||||
The lock is load-bearing, not defensive habit: the mcp SDK dispatches
|
||||
synchronous tool functions through anyio.to_thread.run_sync, so tool
|
||||
bodies genuinely run on multiple worker threads at once — unlike the
|
||||
async REST routes, which are serialized by the event loop.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize an empty registry guarded by a single lock."""
|
||||
self._tools: dict[str, MCPToolStats] = {}
|
||||
# One coarse lock is enough: record() runs once per MCP tool call and
|
||||
# snapshot() is only read by the low-traffic status endpoint.
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def record(self, *, tool: str, duration_ms: float, success: bool) -> None:
|
||||
"""Record the outcome of one MCP tool invocation.
|
||||
|
||||
Never raises: a defect in observability code must not turn a working
|
||||
tool call into a protocol error for the client.
|
||||
"""
|
||||
try:
|
||||
# Coerce outside the lock so a bad argument cannot abort mid-update
|
||||
# and leave calls incremented but duration unaccounted for.
|
||||
duration = float(duration_ms)
|
||||
now = datetime.now(timezone.utc)
|
||||
with self._lock:
|
||||
stats = self._tools.setdefault(tool, MCPToolStats())
|
||||
stats.calls += 1
|
||||
if not success:
|
||||
stats.errors += 1
|
||||
stats.total_duration_ms += duration
|
||||
stats.last_called_at = now
|
||||
except Exception as exc: # noqa: BLE001 - tracking must never break a real call
|
||||
logger.warning("MCPStatsTracker.record failed for tool {} - {}", tool, exc)
|
||||
|
||||
def snapshot(self) -> dict[str, MCPToolStats]:
|
||||
"""Return a shallow copy of all tracked tools, safe to read outside the lock."""
|
||||
with self._lock:
|
||||
return dict(self._tools)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def get_mcp_stats_tracker() -> MCPStatsTracker:
|
||||
"""Return the process-wide singleton tracker (mirrors get_model_usage_tracker())."""
|
||||
return MCPStatsTracker()
|
||||
Reference in New Issue
Block a user