"""MCP audit and daily usage statistics. Every MCP provider call should leave a durable trace that can be joined back to incidents, flywheel nodes, and provider health. This service intentionally uses raw SQL because the audit schema is additive and may be created by migration before every runtime model has been refreshed. """ from __future__ import annotations import json import time import uuid from typing import Any import structlog from sqlalchemy import text from src.db.base import get_db_context logger = structlog.get_logger(__name__) _REDACT_KEYS = {"token", "password", "secret", "api_key", "authorization", "key"} def infer_flywheel_node(mcp_server: str, tool_name: str) -> str: """Infer the flywheel node for a provider/tool pair.""" name = f"{mcp_server}:{tool_name}".lower() if any(k in name for k in ("prometheus", "alert", "metric", "query_range")): return "detect" if any(k in name for k in ("logs", "describe", "events", "node", "hpa", "ssh_get")): return "sense" if any(k in name for k in ("database", "rag", "knowledge", "history")): return "reason" if any(k in name for k in ("blast", "risk", "approval")): return "decide" if any(k in name for k in ("restart", "delete", "scale", "rollout", "ssh_restart")): return "execute" if any(k in name for k in ("watch", "status", "health", "grafana")): return "verify" if any(k in name for k in ("playbook", "km", "embedding")): return "learn" return "govern" def _redact(value: Any) -> Any: if isinstance(value, dict): redacted = {} for key, item in value.items(): if any(marker in str(key).lower() for marker in _REDACT_KEYS): redacted[key] = "" else: redacted[key] = _redact(item) return redacted if isinstance(value, list): return [_redact(item) for item in value] return value def _json_dumps(value: Any) -> str: return json.dumps(_redact(value), ensure_ascii=False, default=str) def _extract_incident_id(parameters: dict[str, Any]) -> str | None: audit_context = parameters.get("_mcp_audit") if isinstance(audit_context, dict) and audit_context.get("incident_id"): return str(audit_context["incident_id"]) for key in ("incident_id", "incidentId"): value = parameters.get(key) if value: return str(value) return None def _extract_audit_context(parameters: dict[str, Any]) -> dict[str, Any]: audit_context = parameters.get("_mcp_audit") return audit_context if isinstance(audit_context, dict) else {} async def record_mcp_call( *, mcp_server: str, tool_name: str, input_params: dict[str, Any], output_result: Any | None, duration_ms: int, success: bool, error_message: str | None, session_id: str | None = None, flywheel_node: str | None = None, incident_id: str | None = None, agent_role: str | None = None, ) -> None: """Persist one MCP tool call and update daily aggregate stats.""" audit_context = _extract_audit_context(input_params) session_id = session_id or audit_context.get("session_id") or str(uuid.uuid4()) flywheel_node = flywheel_node or infer_flywheel_node(mcp_server, tool_name) incident_id = incident_id or _extract_incident_id(input_params) agent_role = agent_role or audit_context.get("agent_role") try: async with get_db_context() as db: await db.execute( text( """ INSERT INTO mcp_audit_log ( session_id, flywheel_node, mcp_server, tool_name, input_params, output_result, duration_ms, success, error_message, incident_id, agent_role ) VALUES ( :session_id, :flywheel_node, :mcp_server, :tool_name, CAST(:input_params AS jsonb), CAST(:output_result AS jsonb), :duration_ms, :success, :error_message, :incident_id, :agent_role ) """ ), { "session_id": session_id, "flywheel_node": flywheel_node, "mcp_server": mcp_server, "tool_name": tool_name, "input_params": _json_dumps(input_params), "output_result": _json_dumps(output_result), "duration_ms": duration_ms, "success": success, "error_message": error_message, "incident_id": incident_id, "agent_role": agent_role, }, ) await db.execute( text( """ INSERT INTO mcp_daily_stats ( date, mcp_server, tool_name, call_count, success_count, avg_duration_ms ) VALUES ( CURRENT_DATE, :mcp_server, :tool_name, 1, CASE WHEN :success THEN 1 ELSE 0 END, :duration_ms ) ON CONFLICT (date, mcp_server, tool_name) DO UPDATE SET call_count = mcp_daily_stats.call_count + 1, success_count = mcp_daily_stats.success_count + CASE WHEN EXCLUDED.success_count > 0 THEN 1 ELSE 0 END, avg_duration_ms = ( (mcp_daily_stats.avg_duration_ms * mcp_daily_stats.call_count) + EXCLUDED.avg_duration_ms ) / (mcp_daily_stats.call_count + 1) """ ), { "mcp_server": mcp_server, "tool_name": tool_name, "success": success, "duration_ms": duration_ms, }, ) except Exception as exc: logger.warning( "mcp_audit_write_failed", mcp_server=mcp_server, tool_name=tool_name, error=str(exc), ) def monotonic_ms() -> int: return int(time.monotonic() * 1000)