92 lines
3.4 KiB
Python
92 lines
3.4 KiB
Python
|
"""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()
|