enact-core 0.0.2

//! Runner Callbacks - Before/after hooks for agents, models, and tools
//!
//! This module defines callback traits for observability and instrumentation
//! of the runner execution. Callbacks are invoked before and after agent,
//! model, and tool executions, providing hooks for:
//!
//! - **Telemetry**: Emit spans/metrics to OpenTelemetry
//! - **Logging**: Structured logging of execution events
//! - **Audit**: Record decisions and actions for compliance
//! - **Monitoring**: Track latencies, success rates, token usage
//!
//! ## Architecture
//!
//! ```text
//! Runner
//!   │
//!   ├── BeforeAgentCallback::on_before_agent()
//!   │     └── Agent Execution
//!   ├── AfterAgentCallback::on_after_agent()
//!   │
//!   ├── BeforeModelCallback::on_before_model()
//!   │     └── Model Call
//!   ├── AfterModelCallback::on_after_model()
//!   │
//!   ├── BeforeToolCallback::on_before_tool()
//!   │     └── Tool Execution
//!   └── AfterToolCallback::on_after_tool()
//! ```
//!
//! ## Usage
//!
//! ```rust,ignore
//! use enact_core::runner::callbacks::*;
//!
//! struct MyTelemetryCallback;
//!
//! impl BeforeAgentCallback for MyTelemetryCallback {
//!     fn on_before_agent(&self, ctx: &AgentCallbackContext) {
//!         tracing::info!(
//!             execution_id = %ctx.execution_id,
//!             agent_name = %ctx.agent_name,
//!             "Agent starting"
//!         );
//!     }
//! }
//! ```
//!
//! @see docs/TECHNICAL/01-EXECUTION-TELEMETRY.md
//! @see crate::telemetry::spans - Span attribute types for OpenTelemetry integration

use crate::kernel::{ExecutionId, StepId};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;

// =============================================================================
// Callback Context Types
// =============================================================================

/// Context passed to agent callbacks
///
/// Contains all information needed to identify and trace an agent execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentCallbackContext {
    /// Execution ID for this agent run
    pub execution_id: ExecutionId,
    /// Step ID (if running within a graph)
    pub step_id: Option<StepId>,
    /// Agent name
    pub agent_name: String,
    /// Agent description (if available)
    pub agent_description: Option<String>,
    /// Input provided to the agent (may be truncated for large inputs)
    pub input_preview: Option<String>,
    /// Tenant ID (if multi-tenant)
    pub tenant_id: Option<String>,
    /// User ID (if authenticated)
    pub user_id: Option<String>,
    /// Trace ID for distributed tracing
    pub trace_id: Option<String>,
    /// Parent span ID
    pub parent_span_id: Option<String>,
}

impl AgentCallbackContext {
    /// Create a new agent callback context
    pub fn new(execution_id: ExecutionId, agent_name: impl Into<String>) -> Self {
        Self {
            execution_id,
            step_id: None,
            agent_name: agent_name.into(),
            agent_description: None,
            input_preview: None,
            tenant_id: None,
            user_id: None,
            trace_id: None,
            parent_span_id: None,
        }
    }

    /// Add step context
    pub fn with_step(mut self, step_id: StepId) -> Self {
        self.step_id = Some(step_id);
        self
    }

    /// Add agent description
    pub fn with_description(mut self, description: impl Into<String>) -> Self {
        self.agent_description = Some(description.into());
        self
    }

    /// Add input preview (truncated if needed)
    pub fn with_input_preview(mut self, input: impl Into<String>) -> Self {
        let input = input.into();
        // Truncate long inputs to prevent callback overhead
        self.input_preview = Some(if input.len() > 500 {
            format!("{}...", &input[..497])
        } else {
            input
        });
        self
    }

    /// Add tenant context
    pub fn with_tenant(mut self, tenant_id: impl Into<String>) -> Self {
        self.tenant_id = Some(tenant_id.into());
        self
    }

    /// Add user context
    pub fn with_user(mut self, user_id: impl Into<String>) -> Self {
        self.user_id = Some(user_id.into());
        self
    }

    /// Add trace context
    pub fn with_trace(
        mut self,
        trace_id: impl Into<String>,
        parent_span_id: Option<String>,
    ) -> Self {
        self.trace_id = Some(trace_id.into());
        self.parent_span_id = parent_span_id;
        self
    }
}

/// Result information for agent completion callback
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentCallbackResult {
    /// Whether the agent succeeded
    pub success: bool,
    /// Duration of agent execution
    pub duration: Duration,
    /// Output preview (truncated if needed)
    pub output_preview: Option<String>,
    /// Error message if failed
    pub error: Option<String>,
    /// Number of steps executed
    pub steps_executed: Option<u32>,
    /// Number of tool calls made
    pub tool_calls: Option<u32>,
    /// Number of model calls made
    pub model_calls: Option<u32>,
}

impl AgentCallbackResult {
    /// Create a successful result
    pub fn success(duration: Duration) -> Self {
        Self {
            success: true,
            duration,
            output_preview: None,
            error: None,
            steps_executed: None,
            tool_calls: None,
            model_calls: None,
        }
    }

    /// Create a failed result
    pub fn failure(duration: Duration, error: impl Into<String>) -> Self {
        Self {
            success: false,
            duration,
            output_preview: None,
            error: Some(error.into()),
            steps_executed: None,
            tool_calls: None,
            model_calls: None,
        }
    }

    /// Add output preview
    pub fn with_output_preview(mut self, output: impl Into<String>) -> Self {
        let output = output.into();
        self.output_preview = Some(if output.len() > 500 {
            format!("{}...", &output[..497])
        } else {
            output
        });
        self
    }

    /// Add execution stats
    pub fn with_stats(mut self, steps: u32, tool_calls: u32, model_calls: u32) -> Self {
        self.steps_executed = Some(steps);
        self.tool_calls = Some(tool_calls);
        self.model_calls = Some(model_calls);
        self
    }
}

/// Context passed to model callbacks
///
/// Contains information about an LLM model call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelCallbackContext {
    /// Execution ID
    pub execution_id: ExecutionId,
    /// Step ID
    pub step_id: Option<StepId>,
    /// Provider name (e.g., "openai", "anthropic")
    pub provider: String,
    /// Model name (e.g., "gpt-4", "claude-3-opus")
    pub model: String,
    /// Temperature setting
    pub temperature: Option<f32>,
    /// Max tokens setting
    pub max_tokens: Option<u32>,
    /// Number of messages in the request
    pub message_count: usize,
    /// Whether this is a streaming request
    pub streaming: bool,
    /// Whether tools are enabled
    pub tools_enabled: bool,
    /// Trace context
    pub trace_id: Option<String>,
}

impl ModelCallbackContext {
    /// Create a new model callback context
    pub fn new(
        execution_id: ExecutionId,
        provider: impl Into<String>,
        model: impl Into<String>,
    ) -> Self {
        Self {
            execution_id,
            step_id: None,
            provider: provider.into(),
            model: model.into(),
            temperature: None,
            max_tokens: None,
            message_count: 0,
            streaming: false,
            tools_enabled: false,
            trace_id: None,
        }
    }

    /// Add step context
    pub fn with_step(mut self, step_id: StepId) -> Self {
        self.step_id = Some(step_id);
        self
    }

    /// Add model parameters
    pub fn with_params(mut self, temperature: Option<f32>, max_tokens: Option<u32>) -> Self {
        self.temperature = temperature;
        self.max_tokens = max_tokens;
        self
    }

    /// Add request info
    pub fn with_request_info(
        mut self,
        message_count: usize,
        streaming: bool,
        tools_enabled: bool,
    ) -> Self {
        self.message_count = message_count;
        self.streaming = streaming;
        self.tools_enabled = tools_enabled;
        self
    }

    /// Add trace context
    pub fn with_trace(mut self, trace_id: impl Into<String>) -> Self {
        self.trace_id = Some(trace_id.into());
        self
    }
}

/// Result information for model completion callback
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelCallbackResult {
    /// Whether the model call succeeded
    pub success: bool,
    /// Duration of model call
    pub duration: Duration,
    /// Input tokens used
    pub input_tokens: Option<u32>,
    /// Output tokens generated
    pub output_tokens: Option<u32>,
    /// Total tokens
    pub total_tokens: Option<u32>,
    /// Finish reason (e.g., "stop", "tool_use", "length")
    pub finish_reason: Option<String>,
    /// Number of tool calls in response
    pub tool_calls_count: Option<u32>,
    /// Error message if failed
    pub error: Option<String>,
    /// Whether response was cached
    pub cached: bool,
}

impl ModelCallbackResult {
    /// Create a successful result
    pub fn success(duration: Duration) -> Self {
        Self {
            success: true,
            duration,
            input_tokens: None,
            output_tokens: None,
            total_tokens: None,
            finish_reason: None,
            tool_calls_count: None,
            error: None,
            cached: false,
        }
    }

    /// Create a failed result
    pub fn failure(duration: Duration, error: impl Into<String>) -> Self {
        Self {
            success: false,
            duration,
            input_tokens: None,
            output_tokens: None,
            total_tokens: None,
            finish_reason: None,
            tool_calls_count: None,
            error: Some(error.into()),
            cached: false,
        }
    }

    /// Add token usage
    pub fn with_tokens(mut self, input: u32, output: u32) -> Self {
        self.input_tokens = Some(input);
        self.output_tokens = Some(output);
        self.total_tokens = Some(input + output);
        self
    }

    /// Add finish info
    pub fn with_finish_info(mut self, reason: impl Into<String>, tool_calls: u32) -> Self {
        self.finish_reason = Some(reason.into());
        self.tool_calls_count = Some(tool_calls);
        self
    }

    /// Mark as cached
    pub fn cached(mut self) -> Self {
        self.cached = true;
        self
    }
}

/// Context passed to tool callbacks
///
/// Contains information about a tool execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallbackContext {
    /// Execution ID
    pub execution_id: ExecutionId,
    /// Step ID
    pub step_id: Option<StepId>,
    /// Tool name
    pub tool_name: String,
    /// Tool description
    pub tool_description: Option<String>,
    /// Whether tool requires network access
    pub requires_network: bool,
    /// Arguments preview (truncated if needed)
    pub args_preview: Option<String>,
    /// Trace context
    pub trace_id: Option<String>,
}

impl ToolCallbackContext {
    /// Create a new tool callback context
    pub fn new(execution_id: ExecutionId, tool_name: impl Into<String>) -> Self {
        Self {
            execution_id,
            step_id: None,
            tool_name: tool_name.into(),
            tool_description: None,
            requires_network: true,
            args_preview: None,
            trace_id: None,
        }
    }

    /// Add step context
    pub fn with_step(mut self, step_id: StepId) -> Self {
        self.step_id = Some(step_id);
        self
    }

    /// Add tool info
    pub fn with_tool_info(mut self, description: Option<String>, requires_network: bool) -> Self {
        self.tool_description = description;
        self.requires_network = requires_network;
        self
    }

    /// Add arguments preview
    pub fn with_args_preview(mut self, args: impl Into<String>) -> Self {
        let args = args.into();
        self.args_preview = Some(if args.len() > 500 {
            format!("{}...", &args[..497])
        } else {
            args
        });
        self
    }

    /// Add trace context
    pub fn with_trace(mut self, trace_id: impl Into<String>) -> Self {
        self.trace_id = Some(trace_id.into());
        self
    }
}

/// Result information for tool completion callback
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallbackResult {
    /// Whether the tool succeeded
    pub success: bool,
    /// Duration of tool execution
    pub duration: Duration,
    /// Output preview (truncated if needed)
    pub output_preview: Option<String>,
    /// Error message if failed
    pub error: Option<String>,
    /// Whether the tool was blocked by policy
    pub blocked: bool,
    /// Blocking reason if blocked
    pub blocked_reason: Option<String>,
}

impl ToolCallbackResult {
    /// Create a successful result
    pub fn success(duration: Duration) -> Self {
        Self {
            success: true,
            duration,
            output_preview: None,
            error: None,
            blocked: false,
            blocked_reason: None,
        }
    }

    /// Create a failed result
    pub fn failure(duration: Duration, error: impl Into<String>) -> Self {
        Self {
            success: false,
            duration,
            output_preview: None,
            error: Some(error.into()),
            blocked: false,
            blocked_reason: None,
        }
    }

    /// Create a blocked result
    pub fn blocked(duration: Duration, reason: impl Into<String>) -> Self {
        Self {
            success: false,
            duration,
            output_preview: None,
            error: None,
            blocked: true,
            blocked_reason: Some(reason.into()),
        }
    }

    /// Add output preview
    pub fn with_output_preview(mut self, output: impl Into<String>) -> Self {
        let output = output.into();
        self.output_preview = Some(if output.len() > 500 {
            format!("{}...", &output[..497])
        } else {
            output
        });
        self
    }
}

// =============================================================================
// Callback Traits
// =============================================================================

/// Callback invoked before agent execution starts
///
/// Use this for:
/// - Starting spans/traces
/// - Logging agent invocation
/// - Recording audit events
pub trait BeforeAgentCallback: Send + Sync {
    /// Called before an agent starts executing
    fn on_before_agent(&self, ctx: &AgentCallbackContext);
}

/// Callback invoked after agent execution completes
///
/// Use this for:
/// - Ending spans/traces with status
/// - Logging completion/failure
/// - Recording metrics (latency, token usage)
pub trait AfterAgentCallback: Send + Sync {
    /// Called after an agent finishes (success or failure)
    fn on_after_agent(&self, ctx: &AgentCallbackContext, result: &AgentCallbackResult);
}

/// Callback invoked before a model call
///
/// Use this for:
/// - Starting LLM spans
/// - Logging model requests
/// - Request auditing
pub trait BeforeModelCallback: Send + Sync {
    /// Called before a model call starts
    fn on_before_model(&self, ctx: &ModelCallbackContext);
}

/// Callback invoked after a model call completes
///
/// Use this for:
/// - Ending LLM spans with token counts
/// - Recording token usage metrics
/// - Cost tracking
pub trait AfterModelCallback: Send + Sync {
    /// Called after a model call finishes
    fn on_after_model(&self, ctx: &ModelCallbackContext, result: &ModelCallbackResult);
}

/// Callback invoked before tool execution
///
/// Use this for:
/// - Starting tool spans
/// - Logging tool invocations
/// - Security auditing
pub trait BeforeToolCallback: Send + Sync {
    /// Called before a tool executes
    fn on_before_tool(&self, ctx: &ToolCallbackContext);
}

/// Callback invoked after tool execution completes
///
/// Use this for:
/// - Ending tool spans
/// - Recording tool metrics
/// - Audit logging
pub trait AfterToolCallback: Send + Sync {
    /// Called after a tool finishes
    fn on_after_tool(&self, ctx: &ToolCallbackContext, result: &ToolCallbackResult);
}

// =============================================================================
// Combined Callback Trait
// =============================================================================

/// Combined callback trait for all execution events
///
/// Implement this trait to handle all callback types with a single implementation.
/// This is useful for telemetry backends that need to track all execution events.
pub trait ExecutionCallbacks:
    BeforeAgentCallback
    + AfterAgentCallback
    + BeforeModelCallback
    + AfterModelCallback
    + BeforeToolCallback
    + AfterToolCallback
{
}

// Blanket implementation for any type that implements all callback traits
impl<T> ExecutionCallbacks for T where
    T: BeforeAgentCallback
        + AfterAgentCallback
        + BeforeModelCallback
        + AfterModelCallback
        + BeforeToolCallback
        + AfterToolCallback
{
}

// =============================================================================
// Callback Registry
// =============================================================================

/// Registry for callback handlers
///
/// Allows registering multiple callbacks that will all be invoked.
#[derive(Default)]
pub struct CallbackRegistry {
    before_agent: Vec<Arc<dyn BeforeAgentCallback>>,
    after_agent: Vec<Arc<dyn AfterAgentCallback>>,
    before_model: Vec<Arc<dyn BeforeModelCallback>>,
    after_model: Vec<Arc<dyn AfterModelCallback>>,
    before_tool: Vec<Arc<dyn BeforeToolCallback>>,
    after_tool: Vec<Arc<dyn AfterToolCallback>>,
}

impl CallbackRegistry {
    /// Create a new empty registry
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a before-agent callback
    pub fn on_before_agent<C: BeforeAgentCallback + 'static>(&mut self, callback: C) -> &mut Self {
        self.before_agent.push(Arc::new(callback));
        self
    }

    /// Register an after-agent callback
    pub fn on_after_agent<C: AfterAgentCallback + 'static>(&mut self, callback: C) -> &mut Self {
        self.after_agent.push(Arc::new(callback));
        self
    }

    /// Register a before-model callback
    pub fn on_before_model<C: BeforeModelCallback + 'static>(&mut self, callback: C) -> &mut Self {
        self.before_model.push(Arc::new(callback));
        self
    }

    /// Register an after-model callback
    pub fn on_after_model<C: AfterModelCallback + 'static>(&mut self, callback: C) -> &mut Self {
        self.after_model.push(Arc::new(callback));
        self
    }

    /// Register a before-tool callback
    pub fn on_before_tool<C: BeforeToolCallback + 'static>(&mut self, callback: C) -> &mut Self {
        self.before_tool.push(Arc::new(callback));
        self
    }

    /// Register an after-tool callback
    pub fn on_after_tool<C: AfterToolCallback + 'static>(&mut self, callback: C) -> &mut Self {
        self.after_tool.push(Arc::new(callback));
        self
    }

    /// Register a combined callback handler for all events
    pub fn register_all<C>(&mut self, callback: Arc<C>) -> &mut Self
    where
        C: ExecutionCallbacks + 'static,
    {
        self.before_agent.push(callback.clone());
        self.after_agent.push(callback.clone());
        self.before_model.push(callback.clone());
        self.after_model.push(callback.clone());
        self.before_tool.push(callback.clone());
        self.after_tool.push(callback);
        self
    }

    /// Invoke all before-agent callbacks
    pub fn invoke_before_agent(&self, ctx: &AgentCallbackContext) {
        for callback in &self.before_agent {
            callback.on_before_agent(ctx);
        }
    }

    /// Invoke all after-agent callbacks
    pub fn invoke_after_agent(&self, ctx: &AgentCallbackContext, result: &AgentCallbackResult) {
        for callback in &self.after_agent {
            callback.on_after_agent(ctx, result);
        }
    }

    /// Invoke all before-model callbacks
    pub fn invoke_before_model(&self, ctx: &ModelCallbackContext) {
        for callback in &self.before_model {
            callback.on_before_model(ctx);
        }
    }

    /// Invoke all after-model callbacks
    pub fn invoke_after_model(&self, ctx: &ModelCallbackContext, result: &ModelCallbackResult) {
        for callback in &self.after_model {
            callback.on_after_model(ctx, result);
        }
    }

    /// Invoke all before-tool callbacks
    pub fn invoke_before_tool(&self, ctx: &ToolCallbackContext) {
        for callback in &self.before_tool {
            callback.on_before_tool(ctx);
        }
    }

    /// Invoke all after-tool callbacks
    pub fn invoke_after_tool(&self, ctx: &ToolCallbackContext, result: &ToolCallbackResult) {
        for callback in &self.after_tool {
            callback.on_after_tool(ctx, result);
        }
    }
}

// =============================================================================
// No-op Implementation
// =============================================================================

/// A no-op callback implementation that does nothing
///
/// Useful as a default or for testing.
#[derive(Debug, Clone, Copy, Default)]
pub struct NoOpCallbacks;

impl BeforeAgentCallback for NoOpCallbacks {
    fn on_before_agent(&self, _ctx: &AgentCallbackContext) {}
}

impl AfterAgentCallback for NoOpCallbacks {
    fn on_after_agent(&self, _ctx: &AgentCallbackContext, _result: &AgentCallbackResult) {}
}

impl BeforeModelCallback for NoOpCallbacks {
    fn on_before_model(&self, _ctx: &ModelCallbackContext) {}
}

impl AfterModelCallback for NoOpCallbacks {
    fn on_after_model(&self, _ctx: &ModelCallbackContext, _result: &ModelCallbackResult) {}
}

impl BeforeToolCallback for NoOpCallbacks {
    fn on_before_tool(&self, _ctx: &ToolCallbackContext) {}
}

impl AfterToolCallback for NoOpCallbacks {
    fn on_after_tool(&self, _ctx: &ToolCallbackContext, _result: &ToolCallbackResult) {}
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::time::Duration;

    // =========================================================================
    // AgentCallbackContext Tests
    // =========================================================================

    #[test]
    fn test_agent_callback_context_new() {
        let ctx = AgentCallbackContext::new(ExecutionId::new(), "test_agent");
        assert_eq!(ctx.agent_name, "test_agent");
        assert!(ctx.step_id.is_none());
        assert!(ctx.agent_description.is_none());
        assert!(ctx.input_preview.is_none());
    }

    #[test]
    fn test_agent_callback_context_with_step() {
        let step_id = StepId::new();
        let ctx =
            AgentCallbackContext::new(ExecutionId::new(), "test_agent").with_step(step_id.clone());
        assert!(ctx.step_id.is_some());
        assert_eq!(ctx.step_id.unwrap().as_str(), step_id.as_str());
    }

    #[test]
    fn test_agent_callback_context_with_description() {
        let ctx = AgentCallbackContext::new(ExecutionId::new(), "test_agent")
            .with_description("A test agent");
        assert_eq!(ctx.agent_description, Some("A test agent".to_string()));
    }

    #[test]
    fn test_agent_callback_context_input_preview_truncation() {
        let long_input = "x".repeat(1000);
        let ctx = AgentCallbackContext::new(ExecutionId::new(), "test_agent")
            .with_input_preview(&long_input);
        let preview = ctx.input_preview.unwrap();
        assert!(preview.len() <= 500);
        assert!(preview.ends_with("..."));
    }

    #[test]
    fn test_agent_callback_context_short_input_not_truncated() {
        let short_input = "hello world";
        let ctx = AgentCallbackContext::new(ExecutionId::new(), "test_agent")
            .with_input_preview(short_input);
        assert_eq!(ctx.input_preview, Some("hello world".to_string()));
    }

    #[test]
    fn test_agent_callback_context_with_tenant() {
        let ctx =
            AgentCallbackContext::new(ExecutionId::new(), "test_agent").with_tenant("tenant_123");
        assert_eq!(ctx.tenant_id, Some("tenant_123".to_string()));
    }

    #[test]
    fn test_agent_callback_context_with_user() {
        let ctx = AgentCallbackContext::new(ExecutionId::new(), "test_agent").with_user("user_456");
        assert_eq!(ctx.user_id, Some("user_456".to_string()));
    }

    #[test]
    fn test_agent_callback_context_with_trace() {
        let ctx = AgentCallbackContext::new(ExecutionId::new(), "test_agent")
            .with_trace("trace_abc", Some("span_xyz".to_string()));
        assert_eq!(ctx.trace_id, Some("trace_abc".to_string()));
        assert_eq!(ctx.parent_span_id, Some("span_xyz".to_string()));
    }

    #[test]
    fn test_agent_callback_context_builder_chain() {
        let ctx = AgentCallbackContext::new(ExecutionId::new(), "test_agent")
            .with_step(StepId::new())
            .with_description("Description")
            .with_input_preview("Input")
            .with_tenant("tenant")
            .with_user("user")
            .with_trace("trace", None);

        assert!(ctx.step_id.is_some());
        assert!(ctx.agent_description.is_some());
        assert!(ctx.input_preview.is_some());
        assert!(ctx.tenant_id.is_some());
        assert!(ctx.user_id.is_some());
        assert!(ctx.trace_id.is_some());
    }

    #[test]
    fn test_agent_callback_context_serde() {
        let ctx = AgentCallbackContext::new(ExecutionId::from_string("exec_test"), "test_agent")
            .with_description("Test description");
        let json = serde_json::to_string(&ctx).unwrap();
        let parsed: AgentCallbackContext = serde_json::from_str(&json).unwrap();
        assert_eq!(ctx.agent_name, parsed.agent_name);
        assert_eq!(ctx.agent_description, parsed.agent_description);
    }

    // =========================================================================
    // AgentCallbackResult Tests
    // =========================================================================

    #[test]
    fn test_agent_callback_result_success() {
        let result = AgentCallbackResult::success(Duration::from_millis(100));
        assert!(result.success);
        assert_eq!(result.duration, Duration::from_millis(100));
        assert!(result.error.is_none());
    }

    #[test]
    fn test_agent_callback_result_failure() {
        let result =
            AgentCallbackResult::failure(Duration::from_millis(50), "Something went wrong");
        assert!(!result.success);
        assert_eq!(result.error, Some("Something went wrong".to_string()));
    }

    #[test]
    fn test_agent_callback_result_with_output_preview() {
        let result = AgentCallbackResult::success(Duration::from_millis(100))
            .with_output_preview("Output here");
        assert_eq!(result.output_preview, Some("Output here".to_string()));
    }

    #[test]
    fn test_agent_callback_result_output_truncation() {
        let long_output = "y".repeat(1000);
        let result = AgentCallbackResult::success(Duration::from_millis(100))
            .with_output_preview(&long_output);
        let preview = result.output_preview.unwrap();
        assert!(preview.len() <= 500);
        assert!(preview.ends_with("..."));
    }

    #[test]
    fn test_agent_callback_result_with_stats() {
        let result = AgentCallbackResult::success(Duration::from_millis(100)).with_stats(5, 3, 2);
        assert_eq!(result.steps_executed, Some(5));
        assert_eq!(result.tool_calls, Some(3));
        assert_eq!(result.model_calls, Some(2));
    }

    // =========================================================================
    // ModelCallbackContext Tests
    // =========================================================================

    #[test]
    fn test_model_callback_context_new() {
        let ctx = ModelCallbackContext::new(ExecutionId::new(), "openai", "gpt-4");
        assert_eq!(ctx.provider, "openai");
        assert_eq!(ctx.model, "gpt-4");
        assert!(ctx.step_id.is_none());
    }

    #[test]
    fn test_model_callback_context_with_params() {
        let ctx = ModelCallbackContext::new(ExecutionId::new(), "anthropic", "claude-3-opus")
            .with_params(Some(0.7), Some(4096));
        assert_eq!(ctx.temperature, Some(0.7));
        assert_eq!(ctx.max_tokens, Some(4096));
    }

    #[test]
    fn test_model_callback_context_with_request_info() {
        let ctx = ModelCallbackContext::new(ExecutionId::new(), "openai", "gpt-4")
            .with_request_info(5, true, true);
        assert_eq!(ctx.message_count, 5);
        assert!(ctx.streaming);
        assert!(ctx.tools_enabled);
    }

    #[test]
    fn test_model_callback_context_with_step() {
        let step_id = StepId::new();
        let ctx = ModelCallbackContext::new(ExecutionId::new(), "openai", "gpt-4")
            .with_step(step_id.clone());
        assert!(ctx.step_id.is_some());
    }

    #[test]
    fn test_model_callback_context_serde() {
        let ctx =
            ModelCallbackContext::new(ExecutionId::from_string("exec_test"), "openai", "gpt-4")
                .with_params(Some(0.5), Some(1000));
        let json = serde_json::to_string(&ctx).unwrap();
        let parsed: ModelCallbackContext = serde_json::from_str(&json).unwrap();
        assert_eq!(ctx.provider, parsed.provider);
        assert_eq!(ctx.model, parsed.model);
        assert_eq!(ctx.temperature, parsed.temperature);
    }

    // =========================================================================
    // ModelCallbackResult Tests
    // =========================================================================

    #[test]
    fn test_model_callback_result_success() {
        let result = ModelCallbackResult::success(Duration::from_millis(500));
        assert!(result.success);
        assert!(result.error.is_none());
    }

    #[test]
    fn test_model_callback_result_failure() {
        let result =
            ModelCallbackResult::failure(Duration::from_millis(100), "Rate limit exceeded");
        assert!(!result.success);
        assert_eq!(result.error, Some("Rate limit exceeded".to_string()));
    }

    #[test]
    fn test_model_callback_result_with_tokens() {
        let result =
            ModelCallbackResult::success(Duration::from_millis(500)).with_tokens(1000, 500);
        assert_eq!(result.input_tokens, Some(1000));
        assert_eq!(result.output_tokens, Some(500));
        assert_eq!(result.total_tokens, Some(1500));
    }

    #[test]
    fn test_model_callback_result_with_finish_info() {
        let result = ModelCallbackResult::success(Duration::from_millis(500))
            .with_finish_info("tool_use", 2);
        assert_eq!(result.finish_reason, Some("tool_use".to_string()));
        assert_eq!(result.tool_calls_count, Some(2));
    }

    #[test]
    fn test_model_callback_result_cached() {
        let result = ModelCallbackResult::success(Duration::from_millis(10)).cached();
        assert!(result.cached);
    }

    // =========================================================================
    // ToolCallbackContext Tests
    // =========================================================================

    #[test]
    fn test_tool_callback_context_new() {
        let ctx = ToolCallbackContext::new(ExecutionId::new(), "read_file");
        assert_eq!(ctx.tool_name, "read_file");
        assert!(ctx.requires_network); // Default is true
    }

    #[test]
    fn test_tool_callback_context_with_tool_info() {
        let ctx = ToolCallbackContext::new(ExecutionId::new(), "calculator")
            .with_tool_info(Some("Performs calculations".to_string()), false);
        assert_eq!(
            ctx.tool_description,
            Some("Performs calculations".to_string())
        );
        assert!(!ctx.requires_network);
    }

    #[test]
    fn test_tool_callback_context_with_args_preview() {
        let ctx = ToolCallbackContext::new(ExecutionId::new(), "search")
            .with_args_preview(r#"{"query": "rust programming"}"#);
        assert!(ctx.args_preview.is_some());
    }

    #[test]
    fn test_tool_callback_context_args_truncation() {
        let long_args = "z".repeat(1000);
        let ctx =
            ToolCallbackContext::new(ExecutionId::new(), "tool").with_args_preview(&long_args);
        let preview = ctx.args_preview.unwrap();
        assert!(preview.len() <= 500);
        assert!(preview.ends_with("..."));
    }

    #[test]
    fn test_tool_callback_context_serde() {
        let ctx = ToolCallbackContext::new(ExecutionId::from_string("exec_test"), "my_tool")
            .with_tool_info(Some("A tool".to_string()), false);
        let json = serde_json::to_string(&ctx).unwrap();
        let parsed: ToolCallbackContext = serde_json::from_str(&json).unwrap();
        assert_eq!(ctx.tool_name, parsed.tool_name);
        assert_eq!(ctx.tool_description, parsed.tool_description);
    }

    // =========================================================================
    // ToolCallbackResult Tests
    // =========================================================================

    #[test]
    fn test_tool_callback_result_success() {
        let result = ToolCallbackResult::success(Duration::from_millis(50));
        assert!(result.success);
        assert!(!result.blocked);
    }

    #[test]
    fn test_tool_callback_result_failure() {
        let result = ToolCallbackResult::failure(Duration::from_millis(20), "File not found");
        assert!(!result.success);
        assert_eq!(result.error, Some("File not found".to_string()));
    }

    #[test]
    fn test_tool_callback_result_blocked() {
        let result =
            ToolCallbackResult::blocked(Duration::from_millis(5), "Tool disabled by policy");
        assert!(!result.success);
        assert!(result.blocked);
        assert_eq!(
            result.blocked_reason,
            Some("Tool disabled by policy".to_string())
        );
    }

    #[test]
    fn test_tool_callback_result_with_output_preview() {
        let result = ToolCallbackResult::success(Duration::from_millis(50))
            .with_output_preview("Result: 42");
        assert_eq!(result.output_preview, Some("Result: 42".to_string()));
    }

    // =========================================================================
    // NoOpCallbacks Tests
    // =========================================================================

    #[test]
    fn test_noop_callbacks_compiles() {
        let callbacks = NoOpCallbacks;
        let agent_ctx = AgentCallbackContext::new(ExecutionId::new(), "test");
        let agent_result = AgentCallbackResult::success(Duration::from_millis(100));
        let model_ctx = ModelCallbackContext::new(ExecutionId::new(), "openai", "gpt-4");
        let model_result = ModelCallbackResult::success(Duration::from_millis(500));
        let tool_ctx = ToolCallbackContext::new(ExecutionId::new(), "tool");
        let tool_result = ToolCallbackResult::success(Duration::from_millis(50));

        // These should compile and not panic
        callbacks.on_before_agent(&agent_ctx);
        callbacks.on_after_agent(&agent_ctx, &agent_result);
        callbacks.on_before_model(&model_ctx);
        callbacks.on_after_model(&model_ctx, &model_result);
        callbacks.on_before_tool(&tool_ctx);
        callbacks.on_after_tool(&tool_ctx, &tool_result);
    }

    // =========================================================================
    // CallbackRegistry Tests
    // =========================================================================

    struct CountingCallback {
        before_agent_count: AtomicU32,
        after_agent_count: AtomicU32,
        before_model_count: AtomicU32,
        after_model_count: AtomicU32,
        before_tool_count: AtomicU32,
        after_tool_count: AtomicU32,
    }

    impl CountingCallback {
        fn new() -> Self {
            Self {
                before_agent_count: AtomicU32::new(0),
                after_agent_count: AtomicU32::new(0),
                before_model_count: AtomicU32::new(0),
                after_model_count: AtomicU32::new(0),
                before_tool_count: AtomicU32::new(0),
                after_tool_count: AtomicU32::new(0),
            }
        }
    }

    impl BeforeAgentCallback for CountingCallback {
        fn on_before_agent(&self, _ctx: &AgentCallbackContext) {
            self.before_agent_count.fetch_add(1, Ordering::SeqCst);
        }
    }

    impl AfterAgentCallback for CountingCallback {
        fn on_after_agent(&self, _ctx: &AgentCallbackContext, _result: &AgentCallbackResult) {
            self.after_agent_count.fetch_add(1, Ordering::SeqCst);
        }
    }

    impl BeforeModelCallback for CountingCallback {
        fn on_before_model(&self, _ctx: &ModelCallbackContext) {
            self.before_model_count.fetch_add(1, Ordering::SeqCst);
        }
    }

    impl AfterModelCallback for CountingCallback {
        fn on_after_model(&self, _ctx: &ModelCallbackContext, _result: &ModelCallbackResult) {
            self.after_model_count.fetch_add(1, Ordering::SeqCst);
        }
    }

    impl BeforeToolCallback for CountingCallback {
        fn on_before_tool(&self, _ctx: &ToolCallbackContext) {
            self.before_tool_count.fetch_add(1, Ordering::SeqCst);
        }
    }

    impl AfterToolCallback for CountingCallback {
        fn on_after_tool(&self, _ctx: &ToolCallbackContext, _result: &ToolCallbackResult) {
            self.after_tool_count.fetch_add(1, Ordering::SeqCst);
        }
    }

    #[test]
    fn test_callback_registry_new() {
        let registry = CallbackRegistry::new();
        // Should not panic
        registry.invoke_before_agent(&AgentCallbackContext::new(ExecutionId::new(), "test"));
    }

    #[test]
    fn test_callback_registry_register_all() {
        let callback = Arc::new(CountingCallback::new());
        let mut registry = CallbackRegistry::new();
        registry.register_all(callback.clone());

        let agent_ctx = AgentCallbackContext::new(ExecutionId::new(), "test");
        let agent_result = AgentCallbackResult::success(Duration::from_millis(100));
        let model_ctx = ModelCallbackContext::new(ExecutionId::new(), "openai", "gpt-4");
        let model_result = ModelCallbackResult::success(Duration::from_millis(500));
        let tool_ctx = ToolCallbackContext::new(ExecutionId::new(), "tool");
        let tool_result = ToolCallbackResult::success(Duration::from_millis(50));

        registry.invoke_before_agent(&agent_ctx);
        registry.invoke_after_agent(&agent_ctx, &agent_result);
        registry.invoke_before_model(&model_ctx);
        registry.invoke_after_model(&model_ctx, &model_result);
        registry.invoke_before_tool(&tool_ctx);
        registry.invoke_after_tool(&tool_ctx, &tool_result);

        assert_eq!(callback.before_agent_count.load(Ordering::SeqCst), 1);
        assert_eq!(callback.after_agent_count.load(Ordering::SeqCst), 1);
        assert_eq!(callback.before_model_count.load(Ordering::SeqCst), 1);
        assert_eq!(callback.after_model_count.load(Ordering::SeqCst), 1);
        assert_eq!(callback.before_tool_count.load(Ordering::SeqCst), 1);
        assert_eq!(callback.after_tool_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_callback_registry_multiple_callbacks() {
        let callback1 = Arc::new(CountingCallback::new());
        let callback2 = Arc::new(CountingCallback::new());
        let mut registry = CallbackRegistry::new();
        registry.register_all(callback1.clone());
        registry.register_all(callback2.clone());

        let agent_ctx = AgentCallbackContext::new(ExecutionId::new(), "test");
        registry.invoke_before_agent(&agent_ctx);

        // Both callbacks should have been invoked
        assert_eq!(callback1.before_agent_count.load(Ordering::SeqCst), 1);
        assert_eq!(callback2.before_agent_count.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_callback_registry_individual_registration() {
        struct SimpleBeforeAgent;
        impl BeforeAgentCallback for SimpleBeforeAgent {
            fn on_before_agent(&self, _ctx: &AgentCallbackContext) {}
        }

        let mut registry = CallbackRegistry::new();
        registry.on_before_agent(SimpleBeforeAgent);

        // Should not panic
        registry.invoke_before_agent(&AgentCallbackContext::new(ExecutionId::new(), "test"));
    }

    // =========================================================================
    // Combined Trait Tests
    // =========================================================================

    #[test]
    fn test_execution_callbacks_trait() {
        // NoOpCallbacks should implement ExecutionCallbacks
        fn accept_execution_callbacks<T: ExecutionCallbacks>(_: &T) {}
        let noop = NoOpCallbacks;
        accept_execution_callbacks(&noop);
    }

    #[test]
    fn test_counting_callback_implements_execution_callbacks() {
        // CountingCallback should implement ExecutionCallbacks
        fn accept_execution_callbacks<T: ExecutionCallbacks>(_: &T) {}
        let counting = CountingCallback::new();
        accept_execution_callbacks(&counting);
    }
}