relay-knowledge 1.1.17

Graph-database-based knowledge graph project.
Documentation
//! Maps MCP method and tool outcomes into bounded agent audit records.

use serde_json::Value;

use crate::{
    api::{AgentProtocolKind, RuntimeIdentity},
    application::AgentDurableAuditInput,
    domain::AuditStatus,
    interfaces::agent::{AgentAuditEvent, AgentAuditQosDecision, AgentAuditStatus},
};

use super::{McpServer, json_rpc::request_id_key};

pub(super) fn record_mcp_qos_rejection(
    server: &McpServer,
    operation: &str,
    id: &Value,
    error_kind: &str,
) {
    let request_id = request_id_key("mcp", id).unwrap_or_else(|| "mcp|invalid-id".to_owned());
    let event = AgentAuditEvent {
        sequence: 0,
        protocol: AgentProtocolKind::Mcp,
        operation: operation.to_owned(),
        request_id: request_id.clone(),
        trace_id: format!("trace-{request_id}"),
        runtime_identity: RuntimeIdentity::mcp(Some(request_id)),
        qos_decision: AgentAuditQosDecision::Rejected,
        status: AgentAuditStatus::Failed,
        source_scope: None,
        freshness: None,
        limit: None,
        result_count: None,
        truncated: false,
        elapsed_ms: 0,
        error_kind: Some(error_kind.to_owned()),
    };
    server.audit.record(event.clone());
    let service = server.service.clone();
    tokio::spawn(async move {
        persist_agent_audit_with_service(service, &event, 0).await;
    });
}

pub(super) async fn record_mcp_tool_audit(
    server: &McpServer,
    operation: &str,
    request_id: &str,
    result: &Value,
    elapsed_ms: u64,
) {
    let structured = &result["structuredContent"];
    let error_kind = structured["error_kind"].as_str().map(str::to_owned);
    let is_error = result["isError"].as_bool().unwrap_or(false);
    let status = match error_kind.as_deref() {
        Some("cancelled") => AgentAuditStatus::Cancelled,
        _ if is_error => AgentAuditStatus::Failed,
        _ => AgentAuditStatus::Completed,
    };

    let event = AgentAuditEvent {
        sequence: 0,
        protocol: AgentProtocolKind::Mcp,
        operation: operation.to_owned(),
        request_id: request_id.to_owned(),
        trace_id: format!("trace-mcp-{request_id}"),
        runtime_identity: RuntimeIdentity::mcp(Some(request_id.to_owned())),
        qos_decision: AgentAuditQosDecision::Admitted,
        status,
        source_scope: audit_source_scope(structured),
        freshness: audit_freshness(structured),
        limit: audit_limit(structured),
        result_count: audit_result_count(structured),
        truncated: audit_truncated(structured),
        elapsed_ms,
        error_kind,
    };
    server.audit.record(event.clone());
    let status_label = match event.status {
        AgentAuditStatus::Completed => "completed",
        AgentAuditStatus::Failed => "failed",
        AgentAuditStatus::Cancelled => "cancelled",
    };
    server
        .metrics
        .record_request("mcp", operation, status_label, elapsed_ms, event.truncated);
    if event.status == AgentAuditStatus::Cancelled {
        server.metrics.record_cancelled("mcp");
    }
    persist_agent_audit(server, &event, audit_graph_version(structured)).await;
}

async fn persist_agent_audit(server: &McpServer, event: &AgentAuditEvent, graph_version: u64) {
    persist_agent_audit_with_service(server.service.clone(), event, graph_version).await;
}

async fn persist_agent_audit_with_service(
    service: crate::application::RelayKnowledgeService,
    event: &AgentAuditEvent,
    graph_version: u64,
) {
    if !service.storage_is_ready() {
        return;
    }
    let detail_json = serde_json::to_string(event).unwrap_or_else(|_| "{}".to_owned());
    let status = match event.status {
        AgentAuditStatus::Completed => AuditStatus::Completed,
        AgentAuditStatus::Failed => AuditStatus::Failed,
        AgentAuditStatus::Cancelled => AuditStatus::Cancelled,
    };
    let _ = service
        .record_agent_audit(AgentDurableAuditInput {
            operation: event.operation.clone(),
            interface: "mcp".to_owned(),
            request_id: event.request_id.clone(),
            trace_id: event.trace_id.clone(),
            status,
            actor: event.runtime_identity.actor_id.clone(),
            source_scope: event.source_scope.clone(),
            graph_version,
            detail_json,
            message: event.error_kind.clone(),
        })
        .await;
}

fn audit_graph_version(structured: &Value) -> u64 {
    structured["metadata"]["graph_version"]
        .as_u64()
        .or_else(|| structured["graph_version"].as_u64())
        .or_else(|| structured["graph"]["graph_version"].as_u64())
        .unwrap_or(0)
}

pub(super) struct McpMethodAudit<'a> {
    pub(super) operation: &'a str,
    pub(super) request_id: &'a str,
    pub(super) status: AgentAuditStatus,
    pub(super) source_scope: Option<String>,
    pub(super) result_count: Option<usize>,
    pub(super) elapsed_ms: u64,
    pub(super) error_kind: Option<&'a str>,
}

pub(super) async fn record_mcp_method_audit(server: &McpServer, input: McpMethodAudit<'_>) {
    let event = AgentAuditEvent {
        sequence: 0,
        protocol: AgentProtocolKind::Mcp,
        operation: input.operation.to_owned(),
        request_id: input.request_id.to_owned(),
        trace_id: format!("trace-mcp-{}", input.request_id),
        runtime_identity: RuntimeIdentity::mcp(Some(input.request_id.to_owned())),
        qos_decision: AgentAuditQosDecision::Admitted,
        status: input.status,
        source_scope: input.source_scope,
        freshness: None,
        limit: None,
        result_count: input.result_count,
        truncated: false,
        elapsed_ms: input.elapsed_ms,
        error_kind: input.error_kind.map(str::to_owned),
    };
    server.audit.record(event.clone());
    let status_label = match event.status {
        AgentAuditStatus::Completed => "completed",
        AgentAuditStatus::Failed => "failed",
        AgentAuditStatus::Cancelled => "cancelled",
    };
    server.metrics.record_request(
        "mcp",
        input.operation,
        status_label,
        input.elapsed_ms,
        false,
    );
    persist_agent_audit(server, &event, 0).await;
}

fn audit_source_scope(structured: &Value) -> Option<String> {
    structured["source_scope"]
        .as_str()
        .or_else(|| structured["scope"]["alias"].as_str())
        .or_else(|| structured["status"]["repository_set"]["alias"].as_str())
        .or_else(|| structured["request"]["repository"]["repository"].as_str())
        .or_else(|| structured["request"]["set_alias"].as_str())
        .map(str::to_owned)
}

fn audit_freshness(structured: &Value) -> Option<String> {
    structured["freshness"]
        .as_str()
        .or_else(|| structured["request"]["freshness_policy"].as_str())
        .map(normalize_freshness_label)
}

fn normalize_freshness_label(value: &str) -> String {
    match value {
        "allow_stale" => "allow-stale",
        "wait_until_fresh" => "wait-until-fresh",
        "graph_only" => "graph-only",
        other => other,
    }
    .to_owned()
}

fn audit_limit(structured: &Value) -> Option<usize> {
    structured["budget_used"]["limit"]
        .as_u64()
        .or_else(|| structured["request"]["limit"].as_u64())
        .and_then(|value| usize::try_from(value).ok())
}

fn audit_truncated(structured: &Value) -> bool {
    structured["truncated"].as_bool().unwrap_or(false)
        || structured["budget"]["snapshot_truncated"]
            .as_bool()
            .unwrap_or(false)
        || structured["budget"]["nodes_truncated"]
            .as_bool()
            .unwrap_or(false)
        || structured["budget"]["edges_truncated"]
            .as_bool()
            .unwrap_or(false)
        || structured["budget"]["sections_truncated"]
            .as_bool()
            .unwrap_or(false)
        || structured["budget"]["evidence_truncated"]
            .as_bool()
            .unwrap_or(false)
}

fn audit_result_count(structured: &Value) -> Option<usize> {
    if let Some(returned) = structured["budget_used"]["returned_count"].as_u64() {
        return usize::try_from(returned).ok();
    }

    structured["results"]
        .as_array()
        .map(Vec::len)
        .or_else(|| structured["flags"].as_array().map(Vec::len))
        .or_else(|| structured["terms"].as_array().map(Vec::len))
        .or_else(|| structured["sections"].as_array().map(Vec::len))
        .or_else(|| software_projection_result_count(structured))
}

fn software_projection_result_count(structured: &Value) -> Option<usize> {
    let fields = [
        "components",
        "dependency_usages",
        "sdk_usages",
        "files",
        "topics",
        "relationships",
        "build_targets",
        "iac_resources",
        "design_elements",
        "entities",
        "statements",
        "diagnostics",
    ];
    let mut count = 0;
    let mut found = false;
    for field in fields {
        if let Some(values) = structured[field].as_array() {
            count += values.len();
            found = true;
        }
    }

    found.then_some(count)
}

#[cfg(test)]
mod mod_tests;