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, ToolExecutionRecord};
use ai_agents_hitl::{ApprovalRequest, ApprovalResolvedOutcome, ApprovalResult};
use ai_agents_llm::{ChatMessage, LLMResponse};
use ai_agents_memory::{MemoryBudgetEvent, MemoryCompressEvent, MemoryEvictEvent};
use ai_agents_tools::ToolResult;
fn preview_text(text: &str, max_chars: usize) -> String {
if max_chars == 0 {
return String::new();
}
let mut chars = text.chars();
let preview: String = chars.by_ref().take(max_chars).collect();
if chars.next().is_some() {
format!("{}...", preview)
} else {
text.to_string()
}
}
#[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_tool_execution_record(&self, _record: &ToolExecutionRecord) {}
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_approval_resolved(
&self,
_request: &ApprovalRequest,
_raw_result: &ApprovalResult,
_outcome: &ApprovalResolvedOutcome,
) {
}
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) {}
async fn on_delegate_start(&self, _agent_id: &str, _state: &str) {}
async fn on_delegate_complete(&self, _agent_id: &str, _state: &str, _duration_ms: u64) {}
async fn on_concurrent_complete(
&self,
_agent_ids: &[String],
_strategy: &str,
_duration_ms: u64,
) {
}
async fn on_group_chat_round(&self, _round: u32, _speaker: &str, _content: &str) {}
async fn on_pipeline_stage(&self, _stage: usize, _agent_id: &str, _duration_ms: u64) {}
async fn on_pipeline_complete(&self, _stages: usize, _duration_ms: u64) {}
async fn on_handoff_start(&self, _initial_agent: &str) {}
async fn on_handoff(&self, _from: &str, _to: &str, _reason: &str) {}
async fn on_persona_evolve(
&self,
_field: &str,
_old_value: &Value,
_new_value: &Value,
_reason: Option<&str>,
) {
}
async fn on_secret_revealed(&self, _content: &str) {}
async fn on_facts_extracted(&self, _actor_id: &str, _facts: &[ai_agents_core::KeyFact]) {}
async fn on_actor_memory_loaded(&self, _actor_id: &str, _fact_count: usize) {}
async fn on_session_created(&self, _session_id: &str) {}
async fn on_sessions_expired(&self, _count: usize) {}
async fn on_relationship_loaded(
&self,
_actor_id: &str,
_relationship: &ai_agents_relationships::Relationship,
) {
}
async fn on_relationship_change(
&self,
_actor_id: &str,
_changes: &[ai_agents_relationships::DimensionChange],
) {
}
async fn on_notable_event(
&self,
_actor_id: &str,
_event: &ai_agents_relationships::RelationshipEvent,
) {
}
}
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 = preview_text(message, 100);
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 = preview_text(&response.content, 100);
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 = preview_text(content, 80);
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) {
debug!("{}[secret_revealed] {}", self.prefix, content);
}
async fn on_facts_extracted(&self, actor_id: &str, facts: &[ai_agents_core::KeyFact]) {
debug!(
"{}[facts_extracted] actor={} count={}",
self.prefix,
actor_id,
facts.len()
);
}
async fn on_actor_memory_loaded(&self, actor_id: &str, fact_count: usize) {
debug!(
"{}[actor_memory_loaded] actor={} facts={}",
self.prefix, actor_id, fact_count
);
}
async fn on_session_created(&self, session_id: &str) {
debug!("{}[session_created] session={}", self.prefix, session_id);
}
async fn on_sessions_expired(&self, count: usize) {
debug!("{}[sessions_expired] count={}", self.prefix, count);
}
async fn on_relationship_loaded(
&self,
actor_id: &str,
relationship: &ai_agents_relationships::Relationship,
) {
debug!(
"{}[relationship_loaded] actor={} dimensions={}",
self.prefix,
actor_id,
relationship.dimensions.len()
);
}
async fn on_relationship_change(
&self,
actor_id: &str,
changes: &[ai_agents_relationships::DimensionChange],
) {
debug!(
"{}[relationship_change] actor={} changes={}",
self.prefix,
actor_id,
changes.len()
);
}
async fn on_notable_event(
&self,
actor_id: &str,
event: &ai_agents_relationships::RelationshipEvent,
) {
debug!(
"{}[notable_event] actor={} significance={:.2} description={}",
self.prefix, actor_id, event.significance, event.description
);
}
}
pub struct CompositeHooks {
hooks: Vec<Arc<dyn AgentHooks>>,
}
impl CompositeHooks {
pub fn new() -> Self {
Self { hooks: Vec::new() }
}
#[allow(clippy::should_implement_trait)]
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_tool_execution_record(&self, record: &ToolExecutionRecord) {
for hook in &self.hooks {
hook.on_tool_execution_record(record).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_approval_resolved(
&self,
request: &ApprovalRequest,
raw_result: &ApprovalResult,
outcome: &ApprovalResolvedOutcome,
) {
for hook in &self.hooks {
hook.on_approval_resolved(request, raw_result, outcome)
.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;
}
}
async fn on_facts_extracted(&self, actor_id: &str, facts: &[ai_agents_core::KeyFact]) {
for hook in &self.hooks {
hook.on_facts_extracted(actor_id, facts).await;
}
}
async fn on_actor_memory_loaded(&self, actor_id: &str, fact_count: usize) {
for hook in &self.hooks {
hook.on_actor_memory_loaded(actor_id, fact_count).await;
}
}
async fn on_session_created(&self, session_id: &str) {
for hook in &self.hooks {
hook.on_session_created(session_id).await;
}
}
async fn on_sessions_expired(&self, count: usize) {
for hook in &self.hooks {
hook.on_sessions_expired(count).await;
}
}
async fn on_relationship_loaded(
&self,
actor_id: &str,
relationship: &ai_agents_relationships::Relationship,
) {
for hook in &self.hooks {
hook.on_relationship_loaded(actor_id, relationship).await;
}
}
async fn on_relationship_change(
&self,
actor_id: &str,
changes: &[ai_agents_relationships::DimensionChange],
) {
for hook in &self.hooks {
hook.on_relationship_change(actor_id, changes).await;
}
}
async fn on_notable_event(
&self,
actor_id: &str,
event: &ai_agents_relationships::RelationshipEvent,
) {
for hook in &self.hooks {
hook.on_notable_event(actor_id, event).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 = approval_status(result);
self.events
.lock()
.push(format!("approval_result:{}:{}", request_id, status));
}
async fn on_approval_resolved(
&self,
request: &ApprovalRequest,
raw_result: &ApprovalResult,
outcome: &ApprovalResolvedOutcome,
) {
self.events.lock().push(format!(
"approval_resolved:{}:{}:{}",
request.id,
approval_status(raw_result),
resolved_status(outcome)
));
}
}
fn approval_status(result: &ApprovalResult) -> &'static str {
match result {
ApprovalResult::Approved => "approved",
ApprovalResult::Rejected { .. } => "rejected",
ApprovalResult::Modified { .. } => "modified",
ApprovalResult::Timeout => "timeout",
}
}
fn resolved_status(outcome: &ApprovalResolvedOutcome) -> &'static str {
match outcome {
ApprovalResolvedOutcome::Approved => "approved",
ApprovalResolvedOutcome::Rejected { .. } => "rejected",
ApprovalResolvedOutcome::Modified { .. } => "modified",
ApprovalResolvedOutcome::Error { .. } => "error",
}
}
#[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;
}
#[test]
fn test_preview_text_handles_unicode_boundaries() {
let text = "제 이름은 Jay이고 가족관계 관련해서 계약서 내용을 확인하고 싶어서";
let preview = preview_text(text, 34);
assert!(preview.ends_with("..."));
assert!(preview.starts_with("제 이름은 Jay"));
}
#[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;
let request = ApprovalRequest::new(
ai_agents_hitl::ApprovalTrigger::tool("calculator", serde_json::json!({})),
"Approve?",
);
composite
.on_approval_resolved(
&request,
&ApprovalResult::Timeout,
&ApprovalResolvedOutcome::Approved,
)
.await;
assert_eq!(
hooks1.events(),
vec![
"tool_start:calculator".to_string(),
format!("approval_resolved:{}:timeout:approved", request.id)
]
);
assert_eq!(hooks1.events(), hooks2.events());
}
#[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());
}
}