loopctl 0.1.0

//! Tool dispatch middleware pipeline.
//!
//! Each middleware wraps the dispatch call and can inspect/modify the request before
//! execution and the result after execution.
//!
//! # Architecture
//!
//! ```text
//! [outer middleware .dispatch()]
//!   → pre-processing (inspect/modify context)
//!   → next.dispatch(ctx)
//!     → [inner middleware .dispatch()]
//!       → ...
//!       → [Innermost: ToolCallMiddleware → Tool::call()]
//!       → ...
//!   → post-processing (inspect/modify result)
//!   → return ToolDispatchResult
//! ```
//!
//! # Provided Middlewares
//!
//! | Middleware                | Description                                                       |
//! |---------------------------|-------------------------------------------------------------------|
//! | [`ToolCallMiddleware`]    | Innermost layer — looks up tool in registry, calls `Tool::call()` |
//! | [`PermissionMiddleware`]  | Checks [`PermissionCheck`] before execution                       |
//! | [`TimeoutMiddleware`]     | Wraps execution in a deadline with retry                          |
//! | [`UnknownToolMiddleware`] | Suggests closest matching tool on "not found"                     |
//!
//! # Example
//!
//! ```rust,ignore
//! use loopctl::middleware::{
//!     ToolPipeline, ToolCallMiddleware, PermissionMiddleware, TimeoutMiddleware,
//! };
//! use loopctl::tool::ToolRegistry;
//! use std::sync::Arc;
//!
//! let registry = Arc::new(ToolRegistry::new());
//!
//! let pipeline = ToolPipeline::builder()
//!     .with(PermissionMiddleware::deny_all())
//!     .with(TimeoutMiddleware::from_secs(120))
//!     .core(registry)
//!     .build()
//!     .expect("pipeline configuration is valid");
//!
//! let result = pipeline.invoke(ctx).await;
//! ```

pub mod output_limit;
pub mod permission;
pub mod timeout;
pub mod tool_call;
pub mod unknown_tool;

use crate::cancel::CancelSignal;
use crate::error::LoopError;
use crate::tool::{PermissionCheck, ToolContext, ToolRegistry};
use serde_json::Value;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

pub use crate::tool::ToolDispatchResult;

pub use output_limit::OutputLimitMiddleware;
pub use permission::{AskResolverFn, PermissionCheckFn, PermissionMiddleware};
pub use timeout::{TimeoutConfig, TimeoutMiddleware};
pub use tool_call::ToolCallMiddleware;
pub use unknown_tool::UnknownToolMiddleware;

// ==================================================
// Dispatch context
// ==================================================

/// Context for a single tool invocation, passed through the middleware chain.
///
/// Built once per tool call by the framework, then threaded through each
/// middleware. Middlewares can read from and write to this struct.
///
/// Session identity is available via [`tool_context.session_id`](ToolContext::session_id)
/// rather than duplicated here.
pub struct ToolDispatchContext {
    /// From the model's tool call. Middlewares can redirect by modifying this.
    pub tool_name: String,
    /// Raw JSON from the model. Middlewares can inspect/transform before invocation.
    pub input: Value,
    /// Generated by the model provider, used to correlate results back to the request.
    pub call_id: String,
    /// Monotonically increasing counter. Useful for per-turn limits.
    pub turn_number: usize,
    /// Checked via `tokio::select!` in long-running middlewares (e.g. [`TimeoutMiddleware`]).
    pub cancel: Arc<CancelSignal>,
    /// Middlewares (e.g. [`PermissionMiddleware`]) can read and modify this.
    pub permission: PermissionCheck,
    /// Augmented by middlewares before [`ToolCallMiddleware`] invokes `Tool::call()`.
    pub tool_context: ToolContext,
}

// ==================================================
// Middleware trait
// ==================================================

/// The trait for a middleware in the tool dispatch pipeline.
///
/// Each middleware wraps the "next" dispatch layer. The outermost middleware
/// is called first; it can pre-process the context, call `next.dispatch()`,
/// then post-process the result. Or it can short-circuit by returning a
/// result without calling `next`.
///
/// # Ordering Conventions
///
/// The typical registration order is:
///
/// 1. [`PermissionMiddleware`] — deny before any work is done
/// 2. [`TimeoutMiddleware`] — wrap execution in a deadline
/// 3. [`UnknownToolMiddleware`] — intercept "not found" errors
/// 4. (innermost) [`ToolCallMiddleware`] — the actual `Tool::call()`
///
/// Post-processing middlewares (e.g. output limiting) are registered
/// *after* the core so they wrap the result on the way out.
///
/// # Example
///
/// ```rust,ignore
/// struct LoggingMiddleware;
///
/// impl ToolMiddleware for LoggingMiddleware {
///     fn name(&self) -> &str { "logging" }
///
///     fn dispatch(
///         &self,
///         ctx: &mut ToolDispatchContext,
///         next: &ToolPipeline,
///     ) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + '_>> {
///         Box::pin(async move {
///             tracing::info!("before tool: {}", ctx.tool_name);
///             let result = next.dispatch(ctx).await;
///             tracing::info!("after tool: {} (error={})", ctx.tool_name, result.is_error);
///             result
///         })
///     }
/// }
/// ```
pub trait ToolMiddleware: Send + Sync {
    /// Human-readable name for this middleware.
    ///
    /// Used in logging, diagnostics, and error messages.
    fn name(&self) -> &str;

    /// Process a tool dispatch through this middleware layer.
    ///
    /// Call `next.dispatch(ctx)` to invoke the next layer in the pipeline.
    /// Return a [`ToolDispatchResult`] directly to short-circuit.
    fn dispatch<'a>(
        &'a self,
        ctx: &'a mut ToolDispatchContext,
        next: &'a ToolPipeline,
    ) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>>;
}

/// Errors that can occur during pipeline construction.
#[derive(Debug, thiserror::Error)]
pub enum PipelineError {
    /// No core dispatch (tool registry) was provided.
    #[error("pipeline requires a core dispatch (call .core() with a ToolRegistry)")]
    MissingCore,
    /// The pipeline is empty — no middlewares and no core.
    #[error("pipeline has no middlewares and no core dispatch")]
    Empty,
}

// ===================================================
// Pipeline
// ===================================================

/// An assembled middleware pipeline for tool dispatch.
///
/// Holds an ordered list of [`ToolMiddleware`] layers around an innermost
/// [`ToolCallMiddleware`] that performs the actual tool invocation via a
/// [`ToolRegistry`]. An internal `index` cursor tracks the current position
/// during dispatch — the entry point starts at index 0, and each middleware
/// receives a pipeline advanced by one position as its `next`.
///
/// # Construction
///
/// Use [`ToolPipeline::builder()`] to create a [`ToolPipelineBuilder`],
/// add middlewares, set the core registry, and call [`build()`](ToolPipelineBuilder::build).
///
/// # Dispatch
///
/// - [`dispatch()`](ToolPipeline::dispatch) — execute a single tool call
///   through the full chain.
/// - [`dispatch_all()`](ToolPipeline::dispatch_all) — execute multiple tool
///   calls sequentially (parallel dispatch is planned).
///
/// # Example
///
/// ```rust,ignore
/// let pipeline = ToolPipeline::builder()
///     .with(PermissionMiddleware::deny_all())
///     .with(TimeoutMiddleware::from_secs(120))
///     .core(registry)
///     .build()
///     .expect("pipeline configuration is valid");
///
/// let result = pipeline.invoke(ctx).await;
/// ```
pub struct ToolPipeline {
    /// Outermost first. Index 0 runs first.
    middlewares: Arc<[Arc<dyn ToolMiddleware>]>,
    /// Created from the [`ToolRegistry`] when the pipeline is built.
    core: Arc<ToolCallMiddleware>,
    /// `0` at entry. Each middleware receives a pipeline at `index + 1`.
    index: usize,
}

impl std::fmt::Debug for ToolPipeline {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ToolPipeline")
            .field("middleware_count", &self.middlewares.len())
            .field("middleware_names", &self.middleware_names())
            .finish_non_exhaustive()
    }
}

impl ToolPipeline {
    /// Create a new pipeline builder.
    ///
    /// Returns a [`ToolPipelineBuilder`] for fluent configuration.
    #[must_use]
    pub fn builder() -> ToolPipelineBuilder {
        ToolPipelineBuilder::new()
    }

    /// Create a minimal pipeline with only the core dispatch.
    ///
    /// Equivalent to `ToolPipeline::builder().core(registry).build().unwrap()`.
    /// Useful for cases where no middleware is needed (e.g. testing).
    #[must_use]
    pub fn new(registry: Arc<ToolRegistry>) -> Self {
        Self {
            middlewares: Arc::new([]),
            core: Arc::new(ToolCallMiddleware::new(registry)),
            index: 0,
        }
    }

    /// Dispatch to the middleware layer at the current cursor index.
    ///
    /// If `index` points past the middleware list, falls through to the
    /// core [`ToolCallMiddleware`]. Middlewares receive a pipeline advanced
    /// by one position as their `next` and call this method to continue
    /// the chain.
    pub fn dispatch<'a>(
        &'a self,
        ctx: &'a mut ToolDispatchContext,
    ) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>> {
        let Some(middleware) = self.middlewares.get(self.index).cloned() else {
            let core = Arc::clone(&self.core);
            return Box::pin(async move { core.dispatch(ctx).await });
        };
        let next = ToolPipeline {
            middlewares: Arc::clone(&self.middlewares),
            core: Arc::clone(&self.core),
            index: self.index.saturating_add(1),
        };
        Box::pin(async move { middleware.dispatch(ctx, &next).await })
    }

    /// Invoke the full pipeline from the beginning.
    ///
    /// Always starts at index 0 regardless of the stored cursor — this
    /// ensures correct behavior even if called on a non-root pipeline.
    /// The cost is two `Arc::clone`s, which is negligible.
    pub async fn invoke(&self, mut ctx: ToolDispatchContext) -> ToolDispatchResult {
        let root = ToolPipeline {
            middlewares: Arc::clone(&self.middlewares),
            core: Arc::clone(&self.core),
            index: 0,
        };
        root.dispatch(&mut ctx).await
    }

    /// Dispatch multiple tool calls sequentially.
    ///
    /// Each call goes through the full middleware chain. Cancellation is
    /// checked between calls — if the signal is set, remaining calls
    /// are skipped and an [`LoopError::Cancelled`] is returned.
    ///
    /// Returns results in the same order as the input calls.
    ///
    /// # Errors
    ///
    /// Returns [`LoopError::Cancelled`] if the cancellation signal is
    /// set before all calls have been dispatched.
    pub async fn dispatch_all(
        &self,
        calls: Vec<ToolDispatchContext>,
    ) -> Result<Vec<ToolDispatchResult>, LoopError> {
        let mut results = Vec::with_capacity(calls.len());
        for ctx in calls {
            if ctx.cancel.is_cancelled() {
                return Err(LoopError::Cancelled);
            }
            let cancel = Arc::clone(&ctx.cancel);
            let result = self.invoke(ctx).await;
            if cancel.is_cancelled() {
                return Err(LoopError::Cancelled);
            }
            results.push(result);
        }
        Ok(results)
    }

    /// Get the names of all middleware layers in order.
    ///
    /// Useful for diagnostics and logging.
    #[must_use]
    pub fn middleware_names(&self) -> Vec<&str> {
        let mut names: Vec<&str> = self.middlewares.iter().map(|m| m.name()).collect();
        names.push(ToolCallMiddleware::NAME);
        names
    }
}

// ===================================================
// Pipeline builder
// ===================================================

/// Builds an ordered middleware pipeline for tool dispatch.
///
/// # Example
///
/// ```rust,ignore
/// let pipeline = ToolPipeline::builder()
///     .with(PermissionMiddleware::deny_all())
///     .with(TimeoutMiddleware::from_secs(120))
///     .core(Arc::new(tool_registry))
///     .build()
///     .expect("pipeline configuration is valid");
/// ```
pub struct ToolPipelineBuilder {
    middlewares: Vec<Arc<dyn ToolMiddleware>>,
    core: Option<Arc<ToolRegistry>>,
}

impl ToolPipelineBuilder {
    /// Create a new, empty builder.
    #[must_use]
    pub fn new() -> Self {
        Self {
            middlewares: Vec::new(),
            core: None,
        }
    }

    /// Add a middleware layer to the pipeline.
    ///
    /// Middlewares are executed in the order they are added (outermost first).
    /// The core dispatch is always the innermost layer.
    #[must_use]
    pub fn with<M: ToolMiddleware + 'static>(mut self, middleware: M) -> Self {
        self.middlewares.push(Arc::new(middleware));
        self
    }

    /// Add an already-`Arc`-wrapped middleware.
    ///
    /// Useful when the same middleware instance needs to be shared
    /// across multiple pipelines.
    #[must_use]
    pub fn with_arc(mut self, middleware: Arc<dyn ToolMiddleware>) -> Self {
        self.middlewares.push(middleware);
        self
    }

    /// Set the core tool registry for the pipeline.
    ///
    /// The registry is wrapped in a [`ToolCallMiddleware`] that performs
    /// the actual tool lookup and invocation — the innermost
    /// layer of the pipeline.
    #[must_use]
    pub fn core(mut self, registry: Arc<ToolRegistry>) -> Self {
        self.core = Some(registry);
        self
    }

    /// Build the pipeline.
    ///
    /// # Errors
    ///
    /// Returns [`PipelineError::MissingCore`] if no core registry was provided.
    pub fn build(self) -> Result<ToolPipeline, PipelineError> {
        let registry = self.core.ok_or(PipelineError::MissingCore)?;
        Ok(ToolPipeline {
            middlewares: self.middlewares.into(),
            core: Arc::new(ToolCallMiddleware::new(registry)),
            index: 0,
        })
    }
}

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

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cancel::CancelSignal;
    use crate::message::ToolContent;
    use crate::message::ToolContentPart;
    use crate::tool::{PermissionCheck, Tool, ToolContext, ToolError, ToolOutput, ToolSchema};
    use serde_json::json;
    use std::future::Future;
    use std::pin::Pin;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::time::{Duration, Instant};

    struct EchoTool;

    impl Tool for EchoTool {
        fn name(&self) -> &'static str {
            "echo"
        }
        fn description(&self) -> &'static str {
            "Echoes back the input"
        }
        fn schema(&self) -> ToolSchema {
            ToolSchema {
                tool: "echo".into(),
                description: "Echoes back the input".into(),
                input_schema: json!({
                    "type": "object",
                    "properties": { "message": { "type": "string" } },
                    "required": ["message"]
                }),
            }
        }
        fn call(
            &self,
            input: Value,
            _ctx: &ToolContext,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            let msg = input["message"].as_str().unwrap_or("").to_string();
            Box::pin(async move { Ok(ToolOutput::text(msg)) })
        }
    }

    struct ErrorTool;

    impl Tool for ErrorTool {
        fn name(&self) -> &'static str {
            "error_tool"
        }
        fn description(&self) -> &'static str {
            "Always returns an error"
        }
        fn schema(&self) -> ToolSchema {
            ToolSchema {
                tool: "error_tool".into(),
                description: "Always returns an error".into(),
                input_schema: json!({"type": "object", "properties": {}}),
            }
        }
        fn call(
            &self,
            _input: Value,
            _ctx: &ToolContext,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            Box::pin(async move { Err(ToolError::Execution("deliberate error".to_string())) })
        }
    }

    struct SlowTool {
        delay_ms: u64,
    }

    impl Tool for SlowTool {
        fn name(&self) -> &'static str {
            "slow"
        }
        fn description(&self) -> &'static str {
            "Takes a long time"
        }
        fn schema(&self) -> ToolSchema {
            ToolSchema {
                tool: "slow".into(),
                description: "Takes a long time".into(),
                input_schema: json!({"type": "object", "properties": {}}),
            }
        }
        fn call(
            &self,
            _input: Value,
            _ctx: &ToolContext,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            let delay = self.delay_ms;
            Box::pin(async move {
                tokio::time::sleep(Duration::from_millis(delay)).await;
                Ok(ToolOutput::text("finally done"))
            })
        }
    }

    fn test_registry() -> Arc<ToolRegistry> {
        let mut reg = ToolRegistry::new();
        reg.register(EchoTool);
        reg.register(ErrorTool);
        reg.register(SlowTool { delay_ms: 5000 });
        Arc::new(reg)
    }

    fn test_ctx(name: &str) -> ToolDispatchContext {
        ToolDispatchContext {
            tool_name: name.to_string(),
            input: json!({ "message": "hello" }),
            call_id: "call_123".to_string(),
            turn_number: 1,
            cancel: Arc::new(CancelSignal::new()),
            permission: PermissionCheck::Allow,
            tool_context: ToolContext::default(),
        }
    }

    #[test]
    fn test_builder_requires_core() {
        let result = ToolPipeline::builder().build();
        assert!(result.is_err(), "should fail without core");
        match result {
            Err(PipelineError::MissingCore) => {}
            other => panic!("expected MissingCore, got {other:?}"),
        }
    }

    #[test]
    fn test_builder_succeeds_with_core() {
        let result = ToolPipeline::builder().core(test_registry()).build();
        assert!(result.is_ok());
    }

    #[test]
    fn test_pipeline_names_no_middleware() {
        let pipeline = ToolPipeline::builder()
            .core(test_registry())
            .build()
            .expect("valid");
        assert_eq!(pipeline.middleware_names(), vec!["tool_call"]);
    }

    #[test]
    fn test_pipeline_names_with_middleware() {
        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::allow_all())
            .with(TimeoutMiddleware::none())
            .core(test_registry())
            .build()
            .expect("valid");
        assert_eq!(
            pipeline.middleware_names(),
            vec!["permission", "timeout", "tool_call"]
        );
    }

    #[tokio::test]
    async fn test_core_dispatch_echo() {
        let pipeline = ToolPipeline::new(test_registry());
        let result = pipeline.invoke(test_ctx("echo")).await;
        assert!(!result.is_error);
        assert_eq!(result.resolved_tool_name, "echo");
        match result.output {
            ToolContent::Text(ref t) => assert_eq!(t, "hello"),
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_core_dispatch_not_found() {
        let pipeline = ToolPipeline::new(test_registry());
        let result = pipeline.invoke(test_ctx("nonexistent")).await;
        assert!(result.is_error);
        assert_eq!(result.resolved_tool_name, "nonexistent");
        match result.output {
            ToolContent::Text(ref t) => assert!(t.contains("not found")),
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_core_dispatch_error_tool() {
        let pipeline = ToolPipeline::new(test_registry());
        let result = pipeline.invoke(test_ctx("error_tool")).await;
        assert!(result.is_error);
        assert_eq!(result.resolved_tool_name, "error_tool");
    }

    #[tokio::test]
    async fn test_permission_deny_all() {
        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::deny_all())
            .core(test_registry())
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("echo")).await;
        assert!(result.is_error);
        match result.output {
            ToolContent::Text(ref t) => assert!(
                t.contains("Permission") && t.contains("blocked"),
                "got: {t}"
            ),
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_permission_allow_all() {
        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::allow_all())
            .core(test_registry())
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("echo")).await;
        assert!(!result.is_error);
    }

    #[tokio::test]
    async fn test_permission_custom_check() {
        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::from_context().with_check(|ctx| {
                if ctx.tool_name == "echo" {
                    PermissionCheck::Allow
                } else {
                    PermissionCheck::Deny {
                        reason: "blocked".to_string(),
                    }
                }
            }))
            .core(test_registry())
            .build()
            .expect("valid");

        let echo_result = pipeline.invoke(test_ctx("echo")).await;
        assert!(!echo_result.is_error);

        let error_result = pipeline.invoke(test_ctx("error_tool")).await;
        assert!(error_result.is_error);
        match error_result.output {
            ToolContent::Text(ref t) => assert!(
                t.contains("Permission") && t.contains("blocked"),
                "got: {t}"
            ),
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_permission_from_context() {
        let mut ctx = test_ctx("echo");
        ctx.permission = PermissionCheck::Deny {
            reason: "test deny".to_string(),
        };

        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::from_context())
            .core(test_registry())
            .build()
            .expect("valid");

        let result = pipeline.invoke(ctx).await;
        assert!(result.is_error);
    }

    #[tokio::test]
    async fn test_permission_ask_without_resolver_denies() {
        let mut ctx = test_ctx("echo");
        ctx.permission = PermissionCheck::Ask {
            prompt: "Allow echo?".to_string(),
        };

        // from_context() has no ask_resolver configured.
        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::from_context())
            .core(test_registry())
            .build()
            .expect("valid");

        let result = pipeline.invoke(ctx).await;
        assert!(
            result.is_error,
            "Ask without resolver should deny the tool call"
        );
        let msg = result.output.to_string();
        assert!(
            msg.contains("permission"),
            "error should mention permission: {msg}"
        );
    }

    #[tokio::test]
    async fn test_timeout_fast_tool_succeeds() {
        let registry = {
            let mut reg = ToolRegistry::new();
            reg.register(EchoTool);
            Arc::new(reg)
        };

        let pipeline = ToolPipeline::builder()
            .with(TimeoutMiddleware::from_secs(5))
            .core(registry)
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("echo")).await;
        assert!(!result.is_error);
    }

    #[tokio::test]
    async fn test_timeout_slow_tool_times_out() {
        let registry = {
            let mut reg = ToolRegistry::new();
            reg.register(SlowTool { delay_ms: 5000 });
            Arc::new(reg)
        };

        let pipeline = ToolPipeline::builder()
            .with(TimeoutMiddleware::new(TimeoutConfig {
                timeout: Duration::from_millis(50),
                retry_on_timeout: false,
                max_retries: 0,
            }))
            .core(registry)
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("slow")).await;
        assert!(result.is_error);
        match result.output {
            ToolContent::Text(ref t) => assert!(t.contains("timed out")),
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    #[test]
    fn test_similarity_identical() {
        let score = UnknownToolMiddleware::similarity("bash", "bash");
        assert!((score - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_similarity_different() {
        let score = UnknownToolMiddleware::similarity("bash", "read_file");
        assert!(score < 0.5);
    }

    #[test]
    fn test_similarity_close() {
        let score = UnknownToolMiddleware::similarity("basj", "bash");
        assert!(score > 0.6, "expected close match, got {score}");
    }

    #[test]
    fn test_similarity_empty() {
        assert!((UnknownToolMiddleware::similarity("", "") - 1.0).abs() < f64::EPSILON);
        assert!((UnknownToolMiddleware::similarity("a", "") - 0.0).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_dispatch_all_sequential() {
        let pipeline = ToolPipeline::new(test_registry());

        let calls = vec![test_ctx("echo"), test_ctx("echo")];
        let results = pipeline.dispatch_all(calls).await.expect("should succeed");

        assert_eq!(results.len(), 2);
        for result in &results {
            assert!(!result.is_error);
        }
    }

    #[tokio::test]
    async fn test_dispatch_all_cancellation() {
        let cancel = Arc::new(CancelSignal::new());
        cancel.cancel();

        let mut ctx = test_ctx("echo");
        ctx.cancel = Arc::clone(&cancel);

        let pipeline = ToolPipeline::new(test_registry());
        let result = pipeline.dispatch_all(vec![ctx]).await;

        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_middleware_ordering_permission_before_timeout() {
        // Permission denies → timeout middleware is never reached
        let registry = {
            let mut reg = ToolRegistry::new();
            reg.register(SlowTool { delay_ms: 5000 });
            Arc::new(reg)
        };

        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::deny_all())
            .with(TimeoutMiddleware::new(TimeoutConfig {
                timeout: Duration::from_millis(50),
                retry_on_timeout: false,
                max_retries: 0,
            }))
            .core(registry)
            .build()
            .expect("valid");

        let start = Instant::now();
        let result = pipeline.invoke(test_ctx("slow")).await;
        let elapsed = start.elapsed();

        assert!(result.is_error);
        // Should return immediately with permission denied, not wait for timeout
        assert!(
            elapsed < Duration::from_millis(200),
            "permission should short-circuit before timeout, took {elapsed:?}",
        );
        match result.output {
            ToolContent::Text(ref t) => assert!(
                t.contains("Permission") && t.contains("blocked"),
                "got: {t}"
            ),
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_full_pipeline_echo() {
        let registry = {
            let mut reg = ToolRegistry::new();
            reg.register(EchoTool);
            Arc::new(reg)
        };

        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::allow_all())
            .with(TimeoutMiddleware::from_secs(30))
            .with(UnknownToolMiddleware::new(Arc::clone(&registry)))
            .core(registry)
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("echo")).await;
        assert!(!result.is_error);
        match result.output {
            ToolContent::Text(ref t) => assert_eq!(t, "hello"),
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_full_pipeline_not_found_with_suggestion() {
        let registry = {
            let mut reg = ToolRegistry::new();
            reg.register(EchoTool);
            reg.register(ErrorTool);
            Arc::new(reg)
        };

        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::allow_all())
            .with(UnknownToolMiddleware::new(Arc::clone(&registry)).with_threshold(0.3))
            .core(registry)
            .build()
            .expect("valid");

        // "errr_tool" is close to "error_tool"
        let result = pipeline.invoke(test_ctx("errr_tool")).await;
        assert!(result.is_error);
        // Should now get a suggestion because the middleware has access to the registry
        match result.output {
            ToolContent::Text(ref msg) => {
                assert!(
                    msg.contains("Did you mean"),
                    "expected suggestion, got: {msg}"
                );
            }
            other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"),
        }
    }

    struct ReachTracker {
        reached: Arc<AtomicBool>,
    }

    impl ToolMiddleware for ReachTracker {
        fn name(&self) -> &'static str {
            "reach_tracker"
        }

        fn dispatch<'a>(
            &'a self,
            ctx: &'a mut ToolDispatchContext,
            next: &'a ToolPipeline,
        ) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>> {
            self.reached.store(true, Ordering::SeqCst);
            next.dispatch(ctx)
        }
    }

    #[tokio::test]
    async fn test_permission_short_circuits_prevents_later_middleware() {
        let reached = Arc::new(AtomicBool::new(false));

        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::deny_all())
            .with(ReachTracker {
                reached: Arc::clone(&reached),
            })
            .core(test_registry())
            .build()
            .expect("valid");

        let _ = pipeline.invoke(test_ctx("echo")).await;
        assert!(
            !reached.load(Ordering::SeqCst),
            "middleware after deny should not be reached"
        );
    }

    struct LongOutputTool;

    impl Tool for LongOutputTool {
        fn name(&self) -> &'static str {
            "long_output"
        }

        fn description(&self) -> &'static str {
            "Returns a long string for testing"
        }

        fn schema(&self) -> ToolSchema {
            ToolSchema {
                tool: self.name().to_string(),
                description: self.description().to_string(),
                input_schema: serde_json::json!({"type": "object", "properties": {}}),
            }
        }

        fn call(
            &self,
            input: Value,
            _ctx: &ToolContext,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            Box::pin(async move {
                let count = usize::try_from(
                    input
                        .get("count")
                        .and_then(serde_json::Value::as_u64)
                        .unwrap_or(100),
                )
                .unwrap();
                let ch = input.get("char").and_then(|v| v.as_str()).unwrap_or("a");
                Ok(ToolOutput::text(ch.repeat(count)))
            })
        }
    }

    struct MultipartTool;

    impl Tool for MultipartTool {
        fn name(&self) -> &'static str {
            "multipart"
        }

        fn description(&self) -> &'static str {
            "Returns multipart content"
        }

        fn schema(&self) -> ToolSchema {
            ToolSchema {
                tool: self.name().to_string(),
                description: self.description().to_string(),
                input_schema: serde_json::json!({"type": "object", "properties": {}}),
            }
        }

        fn call(
            &self,
            _input: Value,
            _ctx: &ToolContext,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            Box::pin(async move {
                Ok(ToolOutput {
                    payload: ToolContent::Multipart(vec![
                        ToolContentPart::Text {
                            text: "part1".to_string(),
                        },
                        ToolContentPart::Text {
                            text: "part2".to_string(),
                        },
                    ]),
                    is_error: false,
                })
            })
        }
    }

    fn long_output_registry() -> Arc<ToolRegistry> {
        let mut registry = ToolRegistry::new();
        registry.register(LongOutputTool);
        registry.register(MultipartTool);
        Arc::new(registry)
    }

    fn long_output_ctx(tool_name: &str, count: usize) -> ToolDispatchContext {
        let mut ctx = test_ctx(tool_name);
        ctx.input = serde_json::json!({"count": count});
        ctx
    }

    #[tokio::test]
    async fn test_output_limit_short_text_passes_through() {
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(1000))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(long_output_ctx("long_output", 50)).await;
        assert!(!result.is_error);
        assert_eq!(result.output, ToolContent::Text("a".repeat(50)));
    }

    #[tokio::test]
    async fn test_output_limit_exact_limit_passes_through() {
        let registry = long_output_registry();
        let limit = 100;
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(limit))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(long_output_ctx("long_output", limit)).await;
        assert!(!result.is_error);
        assert_eq!(result.output, ToolContent::Text("a".repeat(limit)));
    }

    #[tokio::test]
    async fn test_output_limit_truncates_long_text() {
        let registry = long_output_registry();
        let limit = 50;
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(limit))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(long_output_ctx("long_output", 200)).await;
        assert!(!result.is_error);
        let expected = format!("{}\n[truncated]", "a".repeat(limit));
        assert_eq!(result.output, ToolContent::Text(expected));
    }

    #[tokio::test]
    async fn test_output_limit_truncation_respects_char_boundary() {
        // Verify truncation operates on char boundaries, not bytes.
        let registry = long_output_registry();
        let limit = 5;
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(limit))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");
        let mut ctx = test_ctx("long_output");
        ctx.input = serde_json::json!({"count": 10, "char": "é"});

        let result = pipeline.invoke(ctx).await;
        assert!(!result.is_error);
        let expected = format!("{}\n[truncated]", "é".repeat(limit));
        assert_eq!(result.output, ToolContent::Text(expected));
    }

    #[tokio::test]
    async fn test_output_limit_multipart_text_parts_are_truncated() {
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(3))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("multipart")).await;
        assert!(!result.is_error);
        // The shared character budget (3) is consumed by the first text
        // part. "part1" (5 chars > 3) → "par\n[truncated]", remaining=0.
        // "part2" gets zero remaining budget → "[truncated]".
        match result.output {
            ToolContent::Multipart(parts) => {
                assert_eq!(parts.len(), 2, "should still have 2 parts");
                // Part 0 consumes the entire budget.
                match &parts[0] {
                    ToolContentPart::Text { text } => {
                        assert!(
                            text.starts_with("par"),
                            "part 0 should start with first 3 chars: got {text:?}"
                        );
                        assert!(
                            text.contains("[truncated]"),
                            "part 0 should be truncated: got {text:?}"
                        );
                    }
                    ToolContentPart::Image { .. } => {
                        panic!("unexpected image part in MultipartTool output");
                    }
                }
                // Part 1 gets zero remaining budget.
                match &parts[1] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(
                            text, "[truncated]",
                            "part 1 should be fully truncated with zero budget: got {text:?}"
                        );
                    }
                    ToolContentPart::Image { .. } => {
                        panic!("unexpected image part in MultipartTool output");
                    }
                }
            }
            other @ ToolContent::Text(_) => panic!("expected Multipart, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_output_limit_multipart_short_parts_pass_through() {
        // When all text parts are within the limit, multipart should
        // pass through unchanged.
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(100))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("multipart")).await;
        assert!(!result.is_error);
        match result.output {
            ToolContent::Multipart(parts) => {
                assert_eq!(parts.len(), 2);
                match &parts[0] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(text, "part1", "short part should not be truncated");
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part[0], got Image"),
                }
                match &parts[1] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(text, "part2", "short part should not be truncated");
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part[1], got Image"),
                }
            }
            other @ ToolContent::Text(_) => panic!("expected Multipart, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_output_limit_error_output_is_truncated() {
        // When a tool produces an error with a long message, the
        // OutputLimitMiddleware should still truncate it.
        let registry = test_registry(); // has EchoTool
        let pipeline = ToolPipeline::builder()
            .with(PermissionMiddleware::deny_all())
            .with(OutputLimitMiddleware::new(10))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("echo")).await;
        assert!(result.is_error);
        // Permission denied messages are short, so they may or may not
        // be truncated. The key assertion is that the pipeline ran both
        // middlewares without panicking.
    }

    #[tokio::test]
    async fn test_output_limit_preserves_duration() {
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(10))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(long_output_ctx("long_output", 5000)).await;
        // Duration should reflect actual tool execution time, not be
        // zeroed out by the truncation middleware.
        assert!(
            result.duration > Duration::ZERO,
            "duration should be preserved after truncation"
        );
    }

    #[tokio::test]
    async fn test_output_limit_preserves_resolved_tool_name() {
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(10))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(long_output_ctx("long_output", 500)).await;
        assert_eq!(result.resolved_tool_name, "long_output");
    }

    #[tokio::test]
    async fn test_output_limit_stacked_with_timeout() {
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(TimeoutMiddleware::from_secs(120))
            .with(OutputLimitMiddleware::new(20))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(long_output_ctx("long_output", 1000)).await;
        assert!(!result.is_error);
        let expected = format!("{}\n[truncated]", "a".repeat(20));
        assert_eq!(result.output, ToolContent::Text(expected));
    }

    #[tokio::test]
    async fn test_output_limit_zero_chars_truncates_everything() {
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(0))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(long_output_ctx("long_output", 100)).await;
        assert!(!result.is_error);
        assert_eq!(
            result.output,
            ToolContent::Text("\n[truncated]".to_string())
        );
    }

    #[tokio::test]
    async fn test_output_limit_multibyte_under_char_limit_not_truncated() {
        // "日" is 3 bytes. 5 repetitions = 5 chars, 15 bytes.
        // Set max_chars = 8: 5 chars < 8, but 15 bytes > 8.
        let registry = long_output_registry();
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(8))
            .core(Arc::clone(&registry))
            .build()
            .expect("valid");
        let ctx = ToolDispatchContext {
            tool_name: "long_output".to_string(),
            input: json!({"count": 5, "char": "日"}),
            call_id: "mb-1".to_string(),
            turn_number: 1,
            cancel: Arc::new(CancelSignal::new()),
            permission: PermissionCheck::Allow,
            tool_context: ToolContext::default(),
        };

        let result = pipeline.invoke(ctx).await;
        assert!(!result.is_error);
        // Must NOT be truncated: 5 chars <= 8 max_chars, even though 15 bytes > 8.
        assert_eq!(result.output, ToolContent::Text("日日日日日".to_string()));
    }

    #[test]
    fn from_result_maps_ok_output() {
        let output = ToolOutput::text("hello");
        let result = ToolDispatchResult::from_result("echo", Ok(output), Duration::from_millis(50));
        assert!(!result.is_error);
        assert_eq!(result.output, ToolContent::Text("hello".to_string()));
        assert_eq!(result.resolved_tool_name, "echo");
        assert_eq!(result.duration, Duration::from_millis(50));
    }

    #[test]
    fn from_result_maps_err_to_error_result() {
        let error = ToolError::not_found("missing_tool", &[]);
        let result =
            ToolDispatchResult::from_result("missing_tool", Err(error), Duration::from_millis(10));
        assert!(result.is_error);
        assert_eq!(result.resolved_tool_name, "missing_tool");
    }

    #[test]
    fn from_tool_output_defaults() {
        let output = ToolOutput::text("ok");
        let result = ToolDispatchResult::from(output);
        assert!(!result.is_error);
        assert_eq!(result.output, ToolContent::Text("ok".to_string()));
        assert_eq!(result.duration, Duration::ZERO);
        assert_eq!(result.resolved_tool_name, "");
    }

    #[test]
    fn from_tool_output_error_preserves_is_error() {
        let output = ToolOutput::error_text("boom");
        let result = ToolDispatchResult::from(output);
        assert!(result.is_error);
        assert_eq!(result.output, ToolContent::Text("boom".to_string()));
    }

    #[test]
    fn builder_methods_chain() {
        let output = ToolOutput::text("done");
        let result = ToolDispatchResult::from(output)
            .with_tool_name("bash")
            .with_duration(Duration::from_millis(99));
        assert_eq!(result.resolved_tool_name, "bash");
        assert_eq!(result.duration, Duration::from_millis(99));
        assert!(!result.is_error);
    }

    /// A tool that returns 3 text parts, each individually under the limit
    /// but collectively exceeding it.
    struct ThreePartTextTool;

    impl Tool for ThreePartTextTool {
        fn name(&self) -> &'static str {
            "three_part_text"
        }

        fn description(&self) -> &'static str {
            "Returns 3 short text parts"
        }

        fn schema(&self) -> ToolSchema {
            ToolSchema {
                tool: self.name().to_string(),
                description: self.description().to_string(),
                input_schema: json!({"type": "object", "properties": {}}),
            }
        }

        fn call(
            &self,
            _input: Value,
            _ctx: &ToolContext,
        ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + '_>> {
            Box::pin(async move {
                Ok(ToolOutput {
                    payload: ToolContent::Multipart(vec![
                        ToolContentPart::Text {
                            text: "aaaa".to_string(),
                        },
                        ToolContentPart::Text {
                            text: "bbbb".to_string(),
                        },
                        ToolContentPart::Text {
                            text: "cccc".to_string(),
                        },
                    ]),
                    is_error: false,
                })
            })
        }
    }

    #[tokio::test]
    async fn output_limit_multipart_shared_budget_truncates_later_parts() {
        // max_chars = 6: part 0 (4) fits, leaving 2.
        // Part 1 (4 > 2) → "bb\n[truncated]". Part 2 → 0 remaining.
        let mut registry = ToolRegistry::new();
        registry.register(ThreePartTextTool);
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(6))
            .core(Arc::new(registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("three_part_text")).await;
        assert!(!result.is_error);

        match result.output {
            ToolContent::Multipart(parts) => {
                assert_eq!(parts.len(), 3);
                match &parts[0] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(text, "aaaa", "part 0 should be unmodified");
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part 0"),
                }
                match &parts[1] {
                    ToolContentPart::Text { text } => {
                        assert!(
                            text.starts_with("bb"),
                            "part 1 should start with 2 chars: got {text:?}"
                        );
                        assert!(
                            text.contains("[truncated]"),
                            "part 1 should be truncated: got {text:?}"
                        );
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part 1"),
                }
                match &parts[2] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(
                            text, "[truncated]",
                            "part 2 should be fully truncated: got {text:?}"
                        );
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part 2"),
                }
            }
            ToolContent::Text(t) => panic!("expected Multipart, got Text({t})"),
        }
    }

    #[tokio::test]
    async fn output_limit_multipart_shared_budget_all_fit() {
        // max_chars = 20: all three 4-char parts (12 total) fit.
        let mut registry = ToolRegistry::new();
        registry.register(ThreePartTextTool);
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(20))
            .core(Arc::new(registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("three_part_text")).await;
        assert!(!result.is_error);

        match result.output {
            ToolContent::Multipart(parts) => {
                assert_eq!(parts.len(), 3);
                for (i, part) in parts.iter().enumerate() {
                    match part {
                        ToolContentPart::Text { text } => {
                            assert!(
                                !text.contains("[truncated]"),
                                "part {i} should not be truncated: got {text:?}"
                            );
                        }
                        ToolContentPart::Image { .. } => panic!("expected Text part {i}"),
                    }
                }
            }
            ToolContent::Text(t) => panic!("expected Multipart, got Text({t})"),
        }
    }

    #[tokio::test]
    async fn output_limit_multipart_shared_budget_exactly_consumed() {
        // max_chars = 8: parts 0+1 (4+4) exactly consume the budget.
        // Part 2 → 0 remaining → "[truncated]".
        let mut registry = ToolRegistry::new();
        registry.register(ThreePartTextTool);
        let pipeline = ToolPipeline::builder()
            .with(OutputLimitMiddleware::new(8))
            .core(Arc::new(registry))
            .build()
            .expect("valid");

        let result = pipeline.invoke(test_ctx("three_part_text")).await;
        assert!(!result.is_error);

        match result.output {
            ToolContent::Multipart(parts) => {
                assert_eq!(parts.len(), 3);
                match &parts[0] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(text, "aaaa", "part 0 unmodified");
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part 0"),
                }
                match &parts[1] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(text, "bbbb", "part 1 unmodified");
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part 1"),
                }
                match &parts[2] {
                    ToolContentPart::Text { text } => {
                        assert_eq!(
                            text, "[truncated]",
                            "part 2 should be fully truncated: got {text:?}"
                        );
                    }
                    ToolContentPart::Image { .. } => panic!("expected Text part 2"),
                }
            }
            ToolContent::Text(t) => panic!("expected Multipart, got Text({t})"),
        }
    }
}