awoooi/apps/api/src/services/mcp_audit_service.py

"""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] = "<redacted>"
            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)