ai_agents_hooks/

lib.rs

use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
use std::time::Instant;
use tracing::{debug, error, info, warn};

use ai_agents_core::{AgentError, AgentResponse};
use ai_agents_hitl::{ApprovalRequest, ApprovalResult};
use ai_agents_llm::{ChatMessage, LLMResponse};
use ai_agents_memory::{MemoryBudgetEvent, MemoryCompressEvent, MemoryEvictEvent};
use ai_agents_tools::ToolResult;

#[async_trait]
pub trait AgentHooks: Send + Sync {
    async fn on_message_received(&self, _message: &str) {}

    async fn on_llm_start(&self, _messages: &[ChatMessage]) {}

    async fn on_llm_complete(&self, _response: &LLMResponse, _duration_ms: u64) {}

    async fn on_tool_start(&self, _tool: &str, _args: &Value) {}

    async fn on_tool_complete(&self, _tool: &str, _result: &ToolResult, _duration_ms: u64) {}

    async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {}

    async fn on_error(&self, _error: &AgentError) {}

    async fn on_response(&self, _response: &AgentResponse) {}

    async fn on_approval_requested(&self, _request: &ApprovalRequest) {}

    async fn on_approval_result(&self, _request_id: &str, _result: &ApprovalResult) {}

    async fn on_memory_compress(&self, _event: &MemoryCompressEvent) {}

    async fn on_memory_evict(&self, _event: &MemoryEvictEvent) {}

    async fn on_memory_budget_warning(&self, _event: &MemoryBudgetEvent) {}

    /// Fired when a delegated state starts forwarding to a registry agent.
    async fn on_delegate_start(&self, _agent_id: &str, _state: &str) {}

    /// Fired when a delegated state completes.
    async fn on_delegate_complete(&self, _agent_id: &str, _state: &str, _duration_ms: u64) {}

    /// Fired when a concurrent state completes aggregation.
    async fn on_concurrent_complete(
        &self,
        _agent_ids: &[String],
        _strategy: &str,
        _duration_ms: u64,
    ) {
    }

    /// Fired when a group chat round completes.
    async fn on_group_chat_round(&self, _round: u32, _speaker: &str, _content: &str) {}

    /// Fired after each pipeline stage completes.
    async fn on_pipeline_stage(&self, _stage: usize, _agent_id: &str, _duration_ms: u64) {}

    /// Fired when a pipeline completes all stages.
    async fn on_pipeline_complete(&self, _stages: usize, _duration_ms: u64) {}

    /// Fired when a handoff chain starts.
    async fn on_handoff_start(&self, _initial_agent: &str) {}

    /// Fired on each agent-to-agent control transfer.
    async fn on_handoff(&self, _from: &str, _to: &str, _reason: &str) {}

    /// Fired when a persona field is mutated via evolve().
    async fn on_persona_evolve(
        &self,
        _field: &str,
        _old_value: &Value,
        _new_value: &Value,
        _reason: Option<&str>,
    ) {
    }

    /// Fired when a secret's reveal conditions are satisfied for the first time.
    async fn on_secret_revealed(&self, _content: &str) {}
}

pub struct NoopHooks;

#[async_trait]
impl AgentHooks for NoopHooks {}

pub struct LoggingHooks {
    prefix: String,
}

impl LoggingHooks {
    pub fn new() -> Self {
        Self {
            prefix: "[Agent]".to_string(),
        }
    }

    pub fn with_prefix(prefix: impl Into<String>) -> Self {
        Self {
            prefix: prefix.into(),
        }
    }
}

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

#[async_trait]
impl AgentHooks for LoggingHooks {
    async fn on_message_received(&self, message: &str) {
        let preview = if message.len() > 100 {
            format!("{}...", &message[..100])
        } else {
            message.to_string()
        };
        info!("{} Message received: {}", self.prefix, preview);
    }

    async fn on_llm_start(&self, messages: &[ChatMessage]) {
        debug!(
            "{} LLM starting with {} messages",
            self.prefix,
            messages.len()
        );
    }

    async fn on_llm_complete(&self, response: &LLMResponse, duration_ms: u64) {
        info!(
            "{} LLM complete in {}ms, tokens: {:?}",
            self.prefix, duration_ms, response.usage
        );
    }

    async fn on_tool_start(&self, tool: &str, args: &Value) {
        debug!("{} Tool {} starting with args: {}", self.prefix, tool, args);
    }

    async fn on_tool_complete(&self, tool: &str, result: &ToolResult, duration_ms: u64) {
        if result.success {
            info!(
                "{} Tool {} completed in {}ms",
                self.prefix, tool, duration_ms
            );
        } else {
            warn!(
                "{} Tool {} failed in {}ms: {}",
                self.prefix, tool, duration_ms, result.output
            );
        }
    }

    async fn on_state_transition(&self, from: Option<&str>, to: &str, reason: &str) {
        info!(
            "{} State transition: {:?} -> {} ({})",
            self.prefix, from, to, reason
        );
    }

    async fn on_error(&self, err: &AgentError) {
        error!("{} Error: {}", self.prefix, err);
    }

    async fn on_response(&self, response: &AgentResponse) {
        let preview = if response.content.len() > 100 {
            format!("{}...", &response.content[..100])
        } else {
            response.content.clone()
        };
        debug!("{} Response: {}", self.prefix, preview);
    }

    async fn on_approval_requested(&self, request: &ApprovalRequest) {
        info!(
            "{} Approval requested [{}]: {}",
            self.prefix, request.id, request.message
        );
    }

    async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
        match result {
            ApprovalResult::Approved => {
                info!("{} Approval [{}]: approved", self.prefix, request_id);
            }
            ApprovalResult::Rejected { reason } => {
                warn!(
                    "{} Approval [{}]: rejected ({:?})",
                    self.prefix, request_id, reason
                );
            }
            ApprovalResult::Modified { .. } => {
                info!(
                    "{} Approval [{}]: approved with modifications",
                    self.prefix, request_id
                );
            }
            ApprovalResult::Timeout => {
                warn!("{} Approval [{}]: timeout", self.prefix, request_id);
            }
        }
    }

    async fn on_memory_compress(&self, event: &MemoryCompressEvent) {
        info!(
            "{} Memory compressed: {} messages, ratio: {:.2}",
            self.prefix, event.messages_compressed, event.compression_ratio
        );
    }

    async fn on_memory_evict(&self, event: &MemoryEvictEvent) {
        warn!(
            "{} Memory evicted: {} messages, reason: {:?}",
            self.prefix, event.messages_evicted, event.reason
        );
    }

    async fn on_memory_budget_warning(&self, event: &MemoryBudgetEvent) {
        warn!(
            "{} Memory budget warning: {} at {:.1}% ({}/{} tokens)",
            self.prefix,
            event.component,
            event.usage_percent,
            event.used_tokens,
            event.budget_tokens
        );
    }

    async fn on_delegate_start(&self, agent_id: &str, state: &str) {
        info!(
            "{} Delegation started: agent={}, state={}",
            self.prefix, agent_id, state
        );
    }

    async fn on_delegate_complete(&self, agent_id: &str, state: &str, duration_ms: u64) {
        info!(
            "{} Delegation complete: agent={}, state={}, duration={}ms",
            self.prefix, agent_id, state, duration_ms
        );
    }

    async fn on_concurrent_complete(&self, agent_ids: &[String], strategy: &str, duration_ms: u64) {
        info!(
            "{} Concurrent complete: agents={:?}, strategy={}, duration={}ms",
            self.prefix, agent_ids, strategy, duration_ms
        );
    }

    async fn on_group_chat_round(&self, round: u32, speaker: &str, content: &str) {
        let preview = if content.len() > 80 {
            format!("{}...", &content[..80])
        } else {
            content.to_string()
        };
        debug!(
            "{} Group chat round {}: {} said: {}",
            self.prefix, round, speaker, preview
        );
    }

    async fn on_pipeline_stage(&self, stage: usize, agent_id: &str, duration_ms: u64) {
        info!(
            "{} Pipeline stage {}: agent={}, duration={}ms",
            self.prefix, stage, agent_id, duration_ms
        );
    }

    async fn on_pipeline_complete(&self, stages: usize, duration_ms: u64) {
        info!(
            "{} Pipeline complete: {} stages, duration={}ms",
            self.prefix, stages, duration_ms
        );
    }

    async fn on_handoff_start(&self, initial_agent: &str) {
        info!(
            "{} Handoff chain started: initial_agent={}",
            self.prefix, initial_agent
        );
    }

    async fn on_handoff(&self, from: &str, to: &str, reason: &str) {
        info!("{} Handoff: {} -> {} ({})", self.prefix, from, to, reason);
    }

    async fn on_persona_evolve(
        &self,
        field: &str,
        _old_value: &Value,
        new_value: &Value,
        reason: Option<&str>,
    ) {
        info!(
            "{} Persona evolved: field={}, new_value={}, reason={}",
            self.prefix,
            field,
            new_value,
            reason.unwrap_or("(none)")
        );
    }

    async fn on_secret_revealed(&self, content: &str) {
        info!("{} Secret revealed: {}", self.prefix, content);
    }
}

pub struct CompositeHooks {
    hooks: Vec<Arc<dyn AgentHooks>>,
}

impl CompositeHooks {
    pub fn new() -> Self {
        Self { hooks: Vec::new() }
    }

    pub fn add(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
        self.hooks.push(hooks);
        self
    }

    pub fn with_hooks(hooks: Vec<Arc<dyn AgentHooks>>) -> Self {
        Self { hooks }
    }
}

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

#[async_trait]
impl AgentHooks for CompositeHooks {
    async fn on_message_received(&self, message: &str) {
        for hook in &self.hooks {
            hook.on_message_received(message).await;
        }
    }

    async fn on_llm_start(&self, messages: &[ChatMessage]) {
        for hook in &self.hooks {
            hook.on_llm_start(messages).await;
        }
    }

    async fn on_llm_complete(&self, response: &LLMResponse, duration_ms: u64) {
        for hook in &self.hooks {
            hook.on_llm_complete(response, duration_ms).await;
        }
    }

    async fn on_tool_start(&self, tool: &str, args: &Value) {
        for hook in &self.hooks {
            hook.on_tool_start(tool, args).await;
        }
    }

    async fn on_tool_complete(&self, tool: &str, result: &ToolResult, duration_ms: u64) {
        for hook in &self.hooks {
            hook.on_tool_complete(tool, result, duration_ms).await;
        }
    }

    async fn on_state_transition(&self, from: Option<&str>, to: &str, reason: &str) {
        for hook in &self.hooks {
            hook.on_state_transition(from, to, reason).await;
        }
    }

    async fn on_error(&self, error: &AgentError) {
        for hook in &self.hooks {
            hook.on_error(error).await;
        }
    }

    async fn on_response(&self, response: &AgentResponse) {
        for hook in &self.hooks {
            hook.on_response(response).await;
        }
    }

    async fn on_approval_requested(&self, request: &ApprovalRequest) {
        for hook in &self.hooks {
            hook.on_approval_requested(request).await;
        }
    }

    async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
        for hook in &self.hooks {
            hook.on_approval_result(request_id, result).await;
        }
    }

    async fn on_memory_compress(&self, event: &MemoryCompressEvent) {
        for hook in &self.hooks {
            hook.on_memory_compress(event).await;
        }
    }

    async fn on_memory_evict(&self, event: &MemoryEvictEvent) {
        for hook in &self.hooks {
            hook.on_memory_evict(event).await;
        }
    }

    async fn on_memory_budget_warning(&self, event: &MemoryBudgetEvent) {
        for hook in &self.hooks {
            hook.on_memory_budget_warning(event).await;
        }
    }

    async fn on_delegate_start(&self, agent_id: &str, state: &str) {
        for hook in &self.hooks {
            hook.on_delegate_start(agent_id, state).await;
        }
    }

    async fn on_delegate_complete(&self, agent_id: &str, state: &str, duration_ms: u64) {
        for hook in &self.hooks {
            hook.on_delegate_complete(agent_id, state, duration_ms)
                .await;
        }
    }

    async fn on_concurrent_complete(&self, agent_ids: &[String], strategy: &str, duration_ms: u64) {
        for hook in &self.hooks {
            hook.on_concurrent_complete(agent_ids, strategy, duration_ms)
                .await;
        }
    }

    async fn on_group_chat_round(&self, round: u32, speaker: &str, content: &str) {
        for hook in &self.hooks {
            hook.on_group_chat_round(round, speaker, content).await;
        }
    }

    async fn on_pipeline_stage(&self, stage: usize, agent_id: &str, duration_ms: u64) {
        for hook in &self.hooks {
            hook.on_pipeline_stage(stage, agent_id, duration_ms).await;
        }
    }

    async fn on_pipeline_complete(&self, stages: usize, duration_ms: u64) {
        for hook in &self.hooks {
            hook.on_pipeline_complete(stages, duration_ms).await;
        }
    }

    async fn on_handoff_start(&self, initial_agent: &str) {
        for hook in &self.hooks {
            hook.on_handoff_start(initial_agent).await;
        }
    }

    async fn on_handoff(&self, _from: &str, _to: &str, _reason: &str) {
        for hook in &self.hooks {
            hook.on_handoff(_from, _to, _reason).await;
        }
    }

    async fn on_persona_evolve(
        &self,
        _field: &str,
        _old_value: &Value,
        _new_value: &Value,
        _reason: Option<&str>,
    ) {
        for hook in &self.hooks {
            hook.on_persona_evolve(_field, _old_value, _new_value, _reason)
                .await;
        }
    }

    async fn on_secret_revealed(&self, _content: &str) {
        for hook in &self.hooks {
            hook.on_secret_revealed(_content).await;
        }
    }
}

pub struct HookTimer {
    start: Instant,
}

impl HookTimer {
    pub fn start() -> Self {
        Self {
            start: Instant::now(),
        }
    }

    pub fn elapsed_ms(&self) -> u64 {
        self.start.elapsed().as_millis() as u64
    }
}

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

    struct RecordingHooks {
        events: Arc<Mutex<Vec<String>>>,
    }

    impl RecordingHooks {
        fn new() -> Self {
            Self {
                events: Arc::new(Mutex::new(Vec::new())),
            }
        }

        fn events(&self) -> Vec<String> {
            self.events.lock().clone()
        }
    }

    #[async_trait]
    impl AgentHooks for RecordingHooks {
        async fn on_message_received(&self, message: &str) {
            self.events
                .lock()
                .push(format!("message_received:{}", message));
        }

        async fn on_llm_start(&self, messages: &[ChatMessage]) {
            self.events
                .lock()
                .push(format!("llm_start:{}", messages.len()));
        }

        async fn on_llm_complete(&self, _response: &LLMResponse, duration_ms: u64) {
            self.events
                .lock()
                .push(format!("llm_complete:{}", duration_ms));
        }

        async fn on_tool_start(&self, tool: &str, _args: &Value) {
            self.events.lock().push(format!("tool_start:{}", tool));
        }

        async fn on_tool_complete(&self, tool: &str, result: &ToolResult, duration_ms: u64) {
            self.events.lock().push(format!(
                "tool_complete:{}:{}:{}",
                tool, result.success, duration_ms
            ));
        }

        async fn on_state_transition(&self, from: Option<&str>, to: &str, reason: &str) {
            self.events
                .lock()
                .push(format!("state_transition:{:?}:{}:{}", from, to, reason));
        }

        async fn on_error(&self, error: &AgentError) {
            self.events.lock().push(format!("error:{}", error));
        }

        async fn on_response(&self, response: &AgentResponse) {
            self.events
                .lock()
                .push(format!("response:{}", response.content.len()));
        }

        async fn on_approval_requested(&self, request: &ApprovalRequest) {
            self.events
                .lock()
                .push(format!("approval_requested:{}", request.id));
        }

        async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
            let status = match result {
                ApprovalResult::Approved => "approved",
                ApprovalResult::Rejected { .. } => "rejected",
                ApprovalResult::Modified { .. } => "modified",
                ApprovalResult::Timeout => "timeout",
            };
            self.events
                .lock()
                .push(format!("approval_result:{}:{}", request_id, status));
        }
    }

    #[tokio::test]
    async fn test_noop_hooks() {
        let hooks = NoopHooks;
        hooks.on_message_received("test").await;
        hooks.on_llm_start(&[]).await;
    }

    #[tokio::test]
    async fn test_logging_hooks() {
        let hooks = LoggingHooks::new();
        hooks.on_message_received("test message").await;
        hooks.on_llm_start(&[ChatMessage::user("hello")]).await;
    }

    #[tokio::test]
    async fn test_recording_hooks() {
        let hooks = RecordingHooks::new();

        hooks.on_message_received("hello").await;
        hooks.on_llm_start(&[ChatMessage::user("test")]).await;

        let events = hooks.events();
        assert_eq!(events.len(), 2);
        assert!(events[0].contains("message_received"));
        assert!(events[1].contains("llm_start"));
    }

    #[tokio::test]
    async fn test_composite_hooks_with_vec() {
        let hooks1 = Arc::new(RecordingHooks::new());
        let hooks2 = Arc::new(RecordingHooks::new());

        let composite = CompositeHooks::with_hooks(vec![
            hooks1.clone() as Arc<dyn AgentHooks>,
            hooks2.clone() as Arc<dyn AgentHooks>,
        ]);

        composite
            .on_tool_start("calculator", &serde_json::json!({}))
            .await;

        assert_eq!(hooks1.events().len(), 1);
        assert_eq!(hooks2.events().len(), 1);
    }

    #[tokio::test]
    async fn test_hook_timer() {
        let timer = HookTimer::start();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        let elapsed = timer.elapsed_ms();
        assert!(elapsed >= 10);
    }

    #[test]
    fn test_composite_hooks_default() {
        let hooks = CompositeHooks::default();
        assert!(hooks.hooks.is_empty());
    }
}