PicoBot/src/observability/mod.rs

//! Observability module for tracking agent and tool events.
//!
//! This module provides an Observer pattern for emitting and collecting
//! telemetry events during agent execution.

use std::time::Duration;

/// Events emitted during agent and tool execution.
#[derive(Debug, Clone)]
pub enum ObserverEvent {
    /// Emitted before a tool starts executing.
    ToolCallStart {
        tool: String,
        arguments: Option<String>,
    },
    /// Emitted after a tool completes execution.
    ToolCall {
        tool: String,
        duration: Duration,
        success: bool,
    },
    /// Emitted when the agent starts processing.
    AgentStart {
        provider: String,
        model: String,
    },
    /// Emitted when the agent finishes processing.
    AgentEnd {
        provider: String,
        model: String,
        duration: Duration,
        tokens_used: Option<u64>,
    },
}

/// Observer trait for receiving events.
///
/// Implement this trait to receive events during agent execution.
/// Observers are shared across async tasks, so implementations must be
/// Send + Sync.
pub trait Observer: Send + Sync + 'static {
    /// Record a single event.
    fn record_event(&self, event: &ObserverEvent);

    /// Get the observer's name for identification.
    fn name(&self) -> &str;

    /// Flush any buffered events (default no-op).
    fn flush(&self) {}
}

/// Outcome of a single tool execution.
#[derive(Debug, Clone)]
pub struct ToolExecutionOutcome {
    /// The output from the tool execution.
    pub output: String,
    /// Whether the tool executed successfully.
    pub success: bool,
    /// The error reason if the tool failed.
    pub error_reason: Option<String>,
    /// How long the tool took to execute.
    pub duration: Duration,
}

impl ToolExecutionOutcome {
    /// Create a successful outcome with zero duration.
    pub fn success(output: String) -> Self {
        Self {
            output,
            success: true,
            error_reason: None,
            duration: Duration::ZERO,
        }
    }

    /// Create a successful outcome with duration.
    pub fn success_with_duration(output: String, duration: Duration) -> Self {
        Self {
            output,
            success: true,
            error_reason: None,
            duration,
        }
    }

    /// Create a failed outcome with zero duration.
    pub fn failure(output: String, error_reason: Option<String>) -> Self {
        Self {
            output,
            success: false,
            error_reason,
            duration: Duration::ZERO,
        }
    }

    /// Create a failed outcome with duration.
    pub fn failure_with_duration(output: String, error_reason: Option<String>, duration: Duration) -> Self {
        Self {
            output,
            success: false,
            error_reason,
            duration,
        }
    }
}

/// MultiObserver broadcasts events to multiple observers.
pub struct MultiObserver {
    observers: Vec<Box<dyn Observer>>,
}

impl MultiObserver {
    /// Create a new MultiObserver.
    pub fn new() -> Self {
        Self {
            observers: Vec::new(),
        }
    }

    /// Add an observer.
    pub fn add_observer(&mut self, observer: Box<dyn Observer>) {
        self.observers.push(observer);
    }

    /// Get the number of registered observers.
    pub fn len(&self) -> usize {
        self.observers.len()
    }

    /// Check if there are no observers.
    pub fn is_empty(&self) -> bool {
        self.observers.is_empty()
    }
}

impl Default for MultiObserver {
    fn default() -> Self {
        Self::new()
    }
}

impl Observer for MultiObserver {
    fn record_event(&self, event: &ObserverEvent) {
        for observer in &self.observers {
            observer.record_event(event);
        }
    }

    fn flush(&self) {
        for observer in &self.observers {
            observer.flush();
        }
    }

    fn name(&self) -> &str {
        "multi_observer"
    }
}

/// Truncate arguments for logging to avoid oversized events.
pub fn truncate_args(args: &serde_json::Value, max_len: usize) -> String {
    let args_str = args.to_string();
    if args_str.len() <= max_len {
        return args_str;
    }
    format!("{}...truncated", &args_str[..max_len])
}

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

    struct TestObserver {
        name: String,
        events: std::sync::Mutex<Vec<ObserverEvent>>,
    }

    impl TestObserver {
        fn new(name: &str) -> Self {
            Self {
                name: name.to_string(),
                events: std::sync::Mutex::new(Vec::new()),
            }
        }
    }

    impl Observer for TestObserver {
        fn record_event(&self, event: &ObserverEvent) {
            self.events.lock().unwrap().push(event.clone());
        }

        fn name(&self) -> &str {
            &self.name
        }
    }

    #[test]
    fn test_tool_execution_outcome_success() {
        let outcome = ToolExecutionOutcome::success("output content".to_string());
        assert!(outcome.success);
        assert_eq!(outcome.output, "output content");
        assert!(outcome.error_reason.is_none());
        assert_eq!(outcome.duration, Duration::ZERO);
    }

    #[test]
    fn test_tool_execution_outcome_success_with_duration() {
        let outcome = ToolExecutionOutcome::success_with_duration(
            "output content".to_string(),
            Duration::from_millis(100),
        );
        assert!(outcome.success);
        assert_eq!(outcome.duration, Duration::from_millis(100));
    }

    #[test]
    fn test_tool_execution_outcome_failure() {
        let outcome = ToolExecutionOutcome::failure(
            "error output".to_string(),
            Some("error reason".to_string()),
        );
        assert!(!outcome.success);
        assert_eq!(outcome.output, "error output");
        assert_eq!(outcome.error_reason, Some("error reason".to_string()));
        assert_eq!(outcome.duration, Duration::ZERO);
    }

    #[test]
    fn test_multi_observer_broadcasts() {
        let mut multi = MultiObserver::new();
        let obs1 = Box::new(TestObserver::new("obs1"));
        let obs2 = Box::new(TestObserver::new("obs2"));
        multi.add_observer(obs1);
        multi.add_observer(obs2);

        let event = ObserverEvent::ToolCallStart {
            tool: "test_tool".to_string(),
            arguments: Some("{}".to_string()),
        };

        multi.record_event(&event);

        // Both observers should have received the event
        assert_eq!(multi.len(), 2);
    }

    #[test]
    fn test_truncate_args() {
        let args = serde_json::json!({"key": "value"});
        assert_eq!(truncate_args(&args, 100), args.to_string());

        let long_args = serde_json::json!({"key": "a".repeat(200)});
        let truncated = truncate_args(&long_args, 50);
        assert!(truncated.ends_with("...truncated"));
        assert!(truncated.len() < long_args.to_string().len());
    }
}