gitdataai/lib/ai/agent/events.rs

use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Fine-grained agent lifecycle events, inspired by pi's event system.
///
/// Covers the full agent execution lifecycle with enough granularity
/// for UI rendering, telemetry, and extension hooks.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentEvent {
    // === Agent lifecycle ===
    AgentStart,
    AgentEnd {
        messages: Vec<AgentEventMessage>,
        total_input_tokens: u64,
        total_output_tokens: u64,
    },

    // === Turn lifecycle ===
    TurnStart {
        turn_index: usize,
    },
    TurnEnd {
        turn_index: usize,
        assistant_text: Option<String>,
        tool_call_count: usize,
    },

    // === Message lifecycle ===
    MessageStart {
        role: MessageRole,
    },
    MessageTextDelta {
        index: usize,
        delta: String,
    },
    MessageThinkingDelta {
        index: usize,
        delta: String,
    },
    MessageEnd {
        role: MessageRole,
    },

    // === Tool execution lifecycle ===
    ToolExecutionStart {
        tool_call_id: String,
        tool_name: String,
        arguments: Value,
    },
    ToolExecutionUpdate {
        tool_call_id: String,
        tool_name: String,
        partial_output: String,
    },
    ToolExecutionEnd {
        tool_call_id: String,
        tool_name: String,
        output: Option<Value>,
        error: Option<String>,
        elapsed_ms: i64,
    },

    // === Steering / follow-up ===
    SteeringMessagesInjected {
        count: usize,
    },
    FollowUpMessagesInjected {
        count: usize,
    },

    // === Context management ===
    ContextCompacted {
        messages_compacted: usize,
        tokens_saved: i64,
    },
    BranchSummaryCreated {
        entry_count: usize,
        summary_length: usize,
    },

    // === Model switching ===
    ModelSwitched {
        from_model: String,
        to_model: String,
        reason: String,
    },

    // === Error and retry ===
    ErrorClassified {
        category: String,
        message: String,
        will_retry: bool,
        retry_delay_ms: Option<u64>,
    },
    RetryAttempt {
        attempt: usize,
        max_attempts: usize,
        delay_ms: u64,
    },
}

/// Simplified message role for event display.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum MessageRole {
    User,
    Assistant,
    ToolResult,
    System,
}

/// A simplified message representation for `AgentEnd` events.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentEventMessage {
    pub role: MessageRole,
    pub content: String,
    pub tool_calls: Vec<EventToolCall>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EventToolCall {
    pub id: String,
    pub name: String,
    pub arguments: Value,
    pub output: Option<Value>,
    pub error: Option<String>,
}

/// An async-friendly event sink that collects or broadcasts events.
pub struct EventSink {
    senders: Vec<tokio::sync::mpsc::UnboundedSender<AgentEvent>>,
}

impl EventSink {
    pub fn new() -> Self {
        Self {
            senders: Vec::new(),
        }
    }

    /// Subscribe to events, returns a receiver.
    pub fn subscribe(&mut self) -> tokio::sync::mpsc::UnboundedReceiver<AgentEvent> {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        self.senders.push(tx);
        rx
    }

    /// Emit an event to all subscribers. Non-blocking; drops if receiver disconnected.
    pub fn emit(&self, event: AgentEvent) {
        for sender in &self.senders {
            let _ = sender.send(event.clone());
        }
    }

    /// Check if there are any active subscribers.
    pub fn has_subscribers(&self) -> bool {
        !self.senders.is_empty()
    }

    /// Remove disconnected senders.
    pub fn cleanup(&mut self) {
        self.senders.retain(|s| !s.is_closed());
    }
}

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

impl Clone for EventSink {
    fn clone(&self) -> Self {
        Self {
            senders: self.senders.clone(),
        }
    }
}