cerebro 1.1.8

//! # Execution Trace
//!
//! Full audit trail of every action taken during a swarm run.
//! Captures agent decisions, handoffs, tool calls, memory operations,
//! and performance metrics for observability and debugging.

use serde::{Deserialize, Serialize};
use std::sync::{Arc, RwLock};

/// A complete trace of a swarm execution run.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionTrace {
    /// Unique identifier for this swarm run.
    pub run_id: String,
    /// Ordered list of every step taken during execution.
    pub steps: Vec<TraceStep>,
    /// When the run started.
    pub started_at: String,
    /// When the run completed (None if still running).
    pub completed_at: Option<String>,
    /// Overall status.
    pub status: RunStatus,
}

/// Status of a swarm run.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RunStatus {
    Running,
    Completed,
    Failed,
    CircuitBroken,
}

/// A single step in the execution trace.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraceStep {
    /// Which agent performed this step.
    pub agent_id: String,
    /// Human-readable agent name.
    pub agent_name: String,
    /// Step number within this agent's execution.
    pub step_number: usize,
    /// What kind of action was taken.
    pub action: TraceAction,
    /// The input that triggered this step.
    pub input_summary: String,
    /// The output produced by this step.
    pub output_summary: String,
    /// How long this step took in milliseconds.
    pub duration_ms: u64,
    /// Approximate tokens consumed (input + output).
    pub tokens_used: usize,
    /// ISO 8601 timestamp.
    pub timestamp: String,
}

/// The type of action recorded in a trace step.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum TraceAction {
    /// Agent called its LLM.
    LlmCall,
    /// Agent invoked a tool.
    ToolCall { tool_name: String },
    /// Control was transferred from one agent to another.
    Handoff {
        from_agent: String,
        to_agent: String,
    },
    /// Agent queried semantic memory for context.
    MemoryQuery { query: String, results_count: usize },
    /// Agent committed its output to semantic memory.
    MemoryCommit { doc_id: String },
    /// Agent started execution.
    AgentStart,
    /// Agent finished execution.
    AgentComplete,
}

/// Thread-safe tracer that collects steps during a live swarm execution.
pub struct ExecutionTracer {
    run_id: String,
    steps: Arc<RwLock<Vec<TraceStep>>>,
    started_at: String,
}

impl ExecutionTracer {
    /// Create a new tracer for a swarm run.
    pub fn new(run_id: impl Into<String>) -> Self {
        Self {
            run_id: run_id.into(),
            steps: Arc::new(RwLock::new(Vec::new())),
            started_at: chrono::Utc::now().to_rfc3339(),
        }
    }

    /// Record a step in the trace.
    pub fn record(&self, step: TraceStep) {
        if let Ok(mut steps) = self.steps.write() {
            steps.push(step);
        }
    }

    /// Helper to record an LLM call step.
    #[allow(clippy::too_many_arguments)]
    pub fn record_llm_call(
        &self,
        agent_id: &str,
        agent_name: &str,
        step_number: usize,
        input: &str,
        output: &str,
        duration_ms: u64,
        tokens_used: usize,
    ) {
        self.record(TraceStep {
            agent_id: agent_id.to_string(),
            agent_name: agent_name.to_string(),
            step_number,
            action: TraceAction::LlmCall,
            input_summary: truncate_str(input, 200),
            output_summary: truncate_str(output, 500),
            duration_ms,
            tokens_used,
            timestamp: chrono::Utc::now().to_rfc3339(),
        });
    }

    /// Helper to record a handoff between agents.
    pub fn record_handoff(&self, from_id: &str, from_name: &str, to_id: &str) {
        self.record(TraceStep {
            agent_id: from_id.to_string(),
            agent_name: from_name.to_string(),
            step_number: 0,
            action: TraceAction::Handoff {
                from_agent: from_id.to_string(),
                to_agent: to_id.to_string(),
            },
            input_summary: String::new(),
            output_summary: format!("Handing off to {}", to_id),
            duration_ms: 0,
            tokens_used: 0,
            timestamp: chrono::Utc::now().to_rfc3339(),
        });
    }

    /// Helper to record a memory query.
    pub fn record_memory_query(
        &self,
        agent_id: &str,
        agent_name: &str,
        query: &str,
        results_count: usize,
    ) {
        self.record(TraceStep {
            agent_id: agent_id.to_string(),
            agent_name: agent_name.to_string(),
            step_number: 0,
            action: TraceAction::MemoryQuery {
                query: truncate_str(query, 200),
                results_count,
            },
            input_summary: truncate_str(query, 200),
            output_summary: format!("{} results found", results_count),
            duration_ms: 0,
            tokens_used: 0,
            timestamp: chrono::Utc::now().to_rfc3339(),
        });
    }

    /// Helper to record a memory commit.
    pub fn record_memory_commit(&self, agent_id: &str, agent_name: &str, doc_id: &str) {
        self.record(TraceStep {
            agent_id: agent_id.to_string(),
            agent_name: agent_name.to_string(),
            step_number: 0,
            action: TraceAction::MemoryCommit {
                doc_id: doc_id.to_string(),
            },
            input_summary: String::new(),
            output_summary: format!("Committed to memory: {}", doc_id),
            duration_ms: 0,
            tokens_used: 0,
            timestamp: chrono::Utc::now().to_rfc3339(),
        });
    }

    /// Finalize the trace into an immutable ExecutionTrace.
    pub fn finalize(&self, status: RunStatus) -> ExecutionTrace {
        let steps = self.steps.read().map(|s| s.clone()).unwrap_or_default();
        ExecutionTrace {
            run_id: self.run_id.clone(),
            steps,
            started_at: self.started_at.clone(),
            completed_at: Some(chrono::Utc::now().to_rfc3339()),
            status,
        }
    }

    /// Get the current step count.
    pub fn step_count(&self) -> usize {
        self.steps.read().map(|s| s.len()).unwrap_or(0)
    }

    /// Get the run ID.
    pub fn run_id(&self) -> &str {
        &self.run_id
    }
}

/// Truncate a string for trace summaries.
fn truncate_str(s: &str, max_len: usize) -> String {
    if s.len() <= max_len {
        s.to_string()
    } else {
        format!("{}...", &s[..max_len])
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_tracer_record_and_finalize() {
        let tracer = ExecutionTracer::new("run-001");
        assert_eq!(tracer.step_count(), 0);

        tracer.record_llm_call(
            "agent-1",
            "Security Agent",
            1,
            "input text",
            "output text",
            150,
            100,
        );
        assert_eq!(tracer.step_count(), 1);

        tracer.record_handoff("agent-1", "Security Agent", "agent-2");
        assert_eq!(tracer.step_count(), 2);

        let trace = tracer.finalize(RunStatus::Completed);
        assert_eq!(trace.run_id, "run-001");
        assert_eq!(trace.steps.len(), 2);
        assert_eq!(trace.status, RunStatus::Completed);
        assert!(trace.completed_at.is_some());
    }

    #[test]
    fn test_tracer_memory_operations() {
        let tracer = ExecutionTracer::new("run-002");

        tracer.record_memory_query("agent-1", "Research Agent", "Rust safety features", 3);
        tracer.record_memory_commit("agent-1", "Research Agent", "doc-xyz");

        let trace = tracer.finalize(RunStatus::Completed);
        assert_eq!(trace.steps.len(), 2);

        match &trace.steps[0].action {
            TraceAction::MemoryQuery {
                query,
                results_count,
            } => {
                assert!(query.contains("Rust safety"));
                assert_eq!(*results_count, 3);
            }
            _ => panic!("Expected MemoryQuery"),
        }

        match &trace.steps[1].action {
            TraceAction::MemoryCommit { doc_id } => {
                assert_eq!(doc_id, "doc-xyz");
            }
            _ => panic!("Expected MemoryCommit"),
        }
    }

    #[test]
    fn test_trace_serialization() {
        let tracer = ExecutionTracer::new("run-003");
        tracer.record_llm_call("a1", "Agent", 0, "in", "out", 50, 10);
        let trace = tracer.finalize(RunStatus::Completed);
        let json = serde_json::to_string_pretty(&trace).unwrap();
        let deser: ExecutionTrace = serde_json::from_str(&json).unwrap();
        assert_eq!(deser.run_id, "run-003");
        assert_eq!(deser.steps.len(), 1);
    }

    #[test]
    fn test_truncate_str() {
        assert_eq!(truncate_str("hello", 10), "hello");
        assert_eq!(truncate_str("hello world", 5), "hello...");
    }

    #[test]
    fn test_run_status_variants() {
        let statuses = vec![
            RunStatus::Running,
            RunStatus::Completed,
            RunStatus::Failed,
            RunStatus::CircuitBroken,
        ];
        for status in statuses {
            let json = serde_json::to_string(&status).unwrap();
            let deser: RunStatus = serde_json::from_str(&json).unwrap();
            assert_eq!(deser, status);
        }
    }
}