diff --git a/Cargo.toml b/Cargo.toml index 2bfaa17..166e9be 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,7 +22,7 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" serde_repr = "0.1" thiserror = "2" -tokio = { version = "1.52.2", features = ["sync", "macros"] } +tokio = { version = "1.52.3", features = ["sync", "macros", "time"] } uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" diff --git a/src/engine.rs b/src/engine.rs index 72cdf2d..492f590 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -23,5 +23,6 @@ //! ``` mod bare; +pub mod middleware; pub use bare::*; diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 2feee19..51a8321 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -670,45 +670,62 @@ impl BareLoop { self.notify_tool_call(&tc.name, &tc.input.to_string()); let start = Instant::now(); let tool_result = match self.tools.get(&tc.name) { - Some(tool) => match tool.call(tc.input.clone(), &tool_context).await { - Ok(result) => { - let duration = start.elapsed(); - let output_text = result.text_content(); - let success = !result.is_error; - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - &output_text, - duration, - success, - None, - ); - ToolCallResult { - tool_call_id: tc.id.clone(), - output: result.payload, - is_error: result.is_error, - duration, + Some(tool) => { + let cancel = Arc::clone(&self.cancelled); + let call_result = tokio::select! { + r = tool.call(tc.input.clone(), &tool_context) => r, + () = cancel.notified() => { + self.notify_tool_complete( + &tc.name, + &tc.input.to_string(), + "", + start.elapsed(), + false, + Some("cancelled"), + ); + return Err(AgentError::Cancelled); } - } - Err(e) => { - let duration = start.elapsed(); - let error_msg = e.to_string(); - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - &error_msg, - duration, - false, - Some(&error_msg), - ); - ToolCallResult { - tool_call_id: tc.id.clone(), - output: ToolContent::Text(error_msg), - is_error: true, - duration, + }; + match call_result { + Ok(result) => { + let duration = start.elapsed(); + let output_text = result.text_content(); + let success = !result.is_error; + self.notify_tool_complete( + &tc.name, + &tc.input.to_string(), + &output_text, + duration, + success, + None, + ); + ToolCallResult { + tool_call_id: tc.id.clone(), + output: result.payload, + is_error: result.is_error, + duration, + } + } + Err(e) => { + let duration = start.elapsed(); + let error_msg = e.to_string(); + self.notify_tool_complete( + &tc.name, + &tc.input.to_string(), + &error_msg, + duration, + false, + Some(&error_msg), + ); + ToolCallResult { + tool_call_id: tc.id.clone(), + output: ToolContent::Text(error_msg), + is_error: true, + duration, + } } } - }, + } None => { let available: Vec = self.tools.tool_names().clone(); let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); diff --git a/src/engine/middleware.rs b/src/engine/middleware.rs new file mode 100644 index 0000000..086fabd --- /dev/null +++ b/src/engine/middleware.rs @@ -0,0 +1,1699 @@ +//! 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::engine::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; +//! ``` + +use crate::cancel::CancelSignal; +use crate::core::error::AgentError; +use crate::message::ToolContent; +use crate::tool::{PermissionCheck, ToolContext, ToolError, ToolOutput, ToolRegistry}; +use serde_json::Value; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tracing; +use uuid::Uuid; + +// =================================================== +// Dispatch context and result +// =================================================== + +/// Context passed through the middleware chain for a single tool invocation. +/// +/// Built once per tool call by the framework, then threaded through each +/// middleware. Middlewares can read from and write to this struct. +/// +/// # Fields +/// +/// - [`tool_name`](ToolDispatchContext::tool_name) — Which tool is being called. +/// - [`input`](ToolDispatchContext::input) — The JSON input to the tool. +/// - [`call_id`](ToolDispatchContext::call_id) — Unique ID for this tool call. +/// - [`session_id`](ToolDispatchContext::session_id) — Session ID from config. +/// - [`turn_number`](ToolDispatchContext::turn_number) — Turn number within the session. +/// - [`cancel`](ToolDispatchContext::cancel) — Shared cancellation signal. +/// - [`permission`](ToolDispatchContext::permission) — Permission state for this call. +/// - [`tool_context`](ToolDispatchContext::tool_context) — Context passed to the tool. +pub struct ToolDispatchContext { + /// Which tool is being called. + /// + /// Set by the framework from the model's tool call part. + /// Middlewares can redirect to a different tool by modifying + /// this field (e.g. a routing middleware). + pub tool_name: String, + + /// The JSON input to the tool. + /// + /// The raw JSON value from the model's tool call part. + /// Middlewares can inspect or transform this before the tool + /// is invoked (e.g. sanitisation, validation, schema upgrade). + pub input: Value, + + /// Unique ID for this tool call. + /// + /// Generated by the model provider and used to correlate tool + /// results back to the originating tool call request. + pub call_id: String, + + /// Session ID (from `AgentConfig`). + /// + /// Identifies the agent session this call belongs to. Middlewares + /// can use this for logging, rate limiting, or audit trails. + pub session_id: Uuid, + + /// Turn number within the session. + /// + /// Monotonically increasing counter for the current conversation + /// turn. Useful for logging and for middlewares that need to + /// enforce per-turn limits. + pub turn_number: usize, + + /// Shared cancellation signal — middlewares should check this + /// before doing expensive work. + /// + /// Set by the framework when the user or a deadline triggers + /// cancellation. Long-running middlewares (e.g. [`TimeoutMiddleware`]) + /// should cooperative-check this via `tokio::select!`. + pub cancel: Arc, + + /// Permission state for this call. + /// + /// Set by the framework based on tool registration. Middlewares + /// (e.g. [`PermissionMiddleware`]) can read and modify this. + pub permission: PermissionCheck, + + /// Context passed to the underlying tool invocation. + /// + /// Middlewares can augment this (e.g. adding metadata) before + /// the innermost [`ToolCallMiddleware`] invokes `Tool::call()`. + pub tool_context: ToolContext, +} + +/// The outcome of a single tool invocation, produced by the innermost +/// middleware (or by an outer middleware that short-circuits). +/// +/// Carries the output text/error, execution duration, and the resolved +/// tool name (which may differ from the requested name if a router +/// middleware redirected). +pub struct ToolDispatchResult { + /// The tool's output content or error message. + /// + /// Contains the text or multipart payload returned by the tool. + /// When [`is_error`](ToolDispatchResult::is_error) is `true`, this + /// holds a human-readable error description. + pub output: ToolContent, + + /// Whether the output represents an error. + /// + /// `true` when the tool reported failure or a middleware + /// short-circuited with an error. The framework uses this to + /// set the `is_error` flag on the tool result sent back to + /// the model. + pub is_error: bool, + + /// Wall-clock duration of the actual tool execution. + /// + /// Measured by [`ToolCallMiddleware`] from immediately before + /// `Tool::call()` to when it returns. For short-circuited + /// results from outer middlewares this may be [`Duration::ZERO`]. + pub duration: Duration, + + /// The name of the tool that actually ran. + /// + /// May differ from the requested `tool_name` if a routing + /// middleware redirected the call to an alternative tool. + /// When no redirection occurred, this matches the original + /// [`ToolDispatchContext::tool_name`]. + pub resolved_tool_name: String, +} + +impl ToolDispatchResult { + /// Create a successful result with text output. + /// + /// Convenience constructor for the common case where a tool + /// produces a plain-text response. + #[must_use] + pub fn ok(tool_name: &str, output: String, duration: Duration) -> Self { + Self { + output: ToolContent::Text(output), + is_error: false, + duration, + resolved_tool_name: tool_name.to_string(), + } + } + + /// Create an error result with a message. + /// + /// Used when a middleware short-circuits or the tool reports failure. + #[must_use] + pub fn err(tool_name: &str, message: String, duration: Duration) -> Self { + Self { + output: ToolContent::Text(message), + is_error: true, + duration, + resolved_tool_name: tool_name.to_string(), + } + } + + /// Create a result from a [`ToolOutput`]. + /// + /// Converts the tool's output struct into a dispatch result, + /// preserving the error flag and content payload. + #[must_use] + pub fn from_tool_output(tool_name: &str, output: ToolOutput, duration: Duration) -> Self { + Self { + output: output.payload, + is_error: output.is_error, + duration, + resolved_tool_name: tool_name.to_string(), + } + } + + /// Create a result from a [`ToolError`]. + /// + /// Converts the tool's error into a dispatch result with `is_error` set to `true`. + #[must_use] + pub fn from_tool_error(tool_name: &str, error: &ToolError, duration: Duration) -> Self { + Self { + output: ToolContent::Text(error.to_string()), + is_error: true, + duration, + resolved_tool_name: tool_name.to_string(), + } + } +} + +// =================================================== +// 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 + 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 + 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 { + /// Ordered middleware layers (outermost first). + middlewares: Arc<[Arc]>, + + /// The innermost core dispatch that calls `Tool::call()`. + core: Arc, + + /// Cursor position during dispatch. + /// + /// `0` at the entry point. Each middleware receives a pipeline + /// at `index + 1` as its `next`. + 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) -> 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 + 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 [`AgentError::Cancelled`] is returned. + /// + /// Returns results in the same order as the input calls. + /// + /// # Errors + /// + /// Returns [`AgentError::Cancelled`] if the cancellation signal is + /// set before all calls have been dispatched. + pub async fn dispatch_all( + &self, + calls: Vec, + ) -> Result, AgentError> { + let mut results = Vec::with_capacity(calls.len()); + for ctx in calls { + if ctx.cancel.is_cancelled() { + return Err(AgentError::Cancelled); + } + let cancel = Arc::clone(&ctx.cancel); + let result = self.invoke(ctx).await; + if cancel.is_cancelled() { + return Err(AgentError::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>, + core: Option>, +} + +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(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) -> 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. This is the innermost + /// layer of the pipeline. + #[must_use] + pub fn core(mut self, registry: Arc) -> Self { + self.core = Some(registry); + self + } + + /// Build the pipeline. + /// + /// # Errors + /// + /// Returns [`PipelineError::MissingCore`] if no core registry was provided. + pub fn build(self) -> Result { + 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() + } +} + +// =================================================== +// ToolCallMiddleware (innermost core) +// =================================================== + +/// The innermost middleware that performs the actual tool invocation. +/// +/// Looks up the tool by name in the [`ToolRegistry`], calls +/// [`crate::tool::Tool::call()`], and converts the result into a +/// [`ToolDispatchResult`]. If the tool is not found, produces a +/// soft error result (not a hard error) so the model can recover. +/// +/// This middleware is automatically created by the pipeline builder +/// and always occupies the innermost position in the chain. +pub struct ToolCallMiddleware { + registry: Arc, +} + +impl ToolCallMiddleware { + const NAME: &str = "tool_call"; + + /// Create a new core dispatch wrapping the given registry. + #[must_use] + pub fn new(registry: Arc) -> Self { + Self { registry } + } + + /// Execute the tool call — the terminal dispatch. + /// + /// Looks up the tool by name in the registry, calls `Tool::call()`, + /// and converts the result. This is not a middleware — it has no + /// `next` parameter because there is nothing to chain to. + fn dispatch( + &self, + ctx: &mut ToolDispatchContext, + ) -> Pin + Send + '_>> { + let tool_name = ctx.tool_name.clone(); + let input = ctx.input.clone(); + let tool_ctx = ctx.tool_context.clone(); + let registry = Arc::clone(&self.registry); + let cancel = Arc::clone(&ctx.cancel); + + Box::pin(async move { + let start = Instant::now(); + let Some(tool) = registry.get(&tool_name) else { + let available: Vec = registry.tool_names(); + let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); + let error = AgentError::tool_not_found(&tool_name, &available_refs); + return ToolDispatchResult::err(&tool_name, error.to_string(), start.elapsed()); + }; + + let call_result = tokio::select! { + r = tool.call(input, &tool_ctx) => r, + () = cancel.notified() => { + return ToolDispatchResult::err( + &tool_name, + format!("Tool '{tool_name}' cancelled"), + start.elapsed(), + ); + } + }; + + match call_result { + Ok(output) => { + let duration = start.elapsed(); + ToolDispatchResult::from_tool_output(&tool_name, output, duration) + } + Err(e) => { + let duration = start.elapsed(); + ToolDispatchResult::from_tool_error(&tool_name, &e, duration) + } + } + }) + } +} + +// =================================================== +// PermissionMiddleware +// =================================================== + +/// Permission check function type. +pub type PermissionCheckFn = Arc PermissionCheck + Send + Sync>; + +/// Middleware that checks tool permissions before execution. +/// +/// Inspects the [`PermissionCheck`] in the dispatch context. If the +/// permission is `Deny`, short-circuits with an error result. If +/// `Ask`, also denies (interactive prompts are outside the framework's +/// scope — agents should handle this themselves). +/// +/// # Example +/// +/// ```rust,ignore +/// // Deny all by default +/// let mw = PermissionMiddleware::deny_all(); +/// +/// // Custom logic +/// let mw = PermissionMiddleware::with_check(|ctx| { +/// if ctx.tool_name == "safe_read" { +/// PermissionCheck::Allow +/// } else { +/// PermissionCheck::Deny +/// } +/// }); +/// ``` +pub struct PermissionMiddleware { + /// Override function for permission checking. + /// + /// When `Some`, this function is called to determine the permission + /// for each dispatch. When `None`, the middleware reads + /// [`ToolDispatchContext::permission`] directly. + check_fn: Option, +} + +impl PermissionMiddleware { + /// Create a permission middleware that denies all calls. + /// + /// Every tool call will be short-circuited with a permission-denied + /// error. Useful as a safety default in restricted environments. + #[must_use] + pub fn deny_all() -> Self { + Self { + check_fn: Some(Arc::new(|_| PermissionCheck::Deny { + reason: "blocked by policy".into(), + })), + } + } + + /// Create a permission middleware that allows all calls. + /// + /// No permission checks are performed — every tool call passes + /// through to the next layer. This is the runtime-equivalent of + /// having no permission middleware at all, but it can be useful + /// for logging or metrics in permissive environments. + #[must_use] + pub fn allow_all() -> Self { + Self { + check_fn: Some(Arc::new(|_| PermissionCheck::Allow)), + } + } + + /// Create a permission middleware with a custom check function. + /// + /// The function receives a reference to the dispatch context and + /// returns the appropriate [`PermissionCheck`] for that call. + pub fn with_check( + f: impl Fn(&ToolDispatchContext) -> PermissionCheck + Send + Sync + 'static, + ) -> Self { + Self { + check_fn: Some(Arc::new(f)), + } + } + + /// Create a permission middleware that reads from the context. + /// + /// The middleware reads `ctx.permission` directly, without + /// applying any override. This is useful when the permission + /// is set by the framework or a prior middleware. + #[must_use] + pub fn from_context() -> Self { + Self { check_fn: None } + } + + fn resolve_permission(&self, ctx: &ToolDispatchContext) -> PermissionCheck { + match &self.check_fn { + Some(f) => f(ctx), + None => ctx.permission.clone(), + } + } +} + +impl ToolMiddleware for PermissionMiddleware { + fn name(&self) -> &'static str { + "permission" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + let permission = self.resolve_permission(ctx); + match permission { + PermissionCheck::Allow => next.dispatch(ctx), + PermissionCheck::Ask { .. } | PermissionCheck::Deny { .. } => { + let tool_name = ctx.tool_name.clone(); + let reason = match &permission { + PermissionCheck::Deny { reason } => reason.clone(), + PermissionCheck::Ask { prompt } => format!("permission required: {prompt}"), + _ => "blocked".to_string(), + }; + tracing::warn!( + tool = %tool_name, + permission = %reason, + "tool call blocked by permission middleware" + ); + Box::pin(std::future::ready(ToolDispatchResult::err( + &tool_name, + format!("Permission {reason} for tool '{tool_name}'"), + Duration::ZERO, + ))) + } + PermissionCheck::Modify { .. } => { + // Modify is treated as allow — the tool can proceed + // but the agent might apply modifications to the input. + next.dispatch(ctx) + } + } + } +} + +// =================================================== +// TimeoutMiddleware +// =================================================== + +/// Configuration for the [`TimeoutMiddleware`]. +#[derive(Debug, Clone)] +pub struct TimeoutConfig { + /// Timeout for the initial tool execution attempt. + /// + /// If the inner dispatch does not complete within this duration, + /// [`TimeoutMiddleware`] cancels the future and either retries + /// (if [`retry_on_timeout`](TimeoutConfig::retry_on_timeout) is + /// `true`) or returns a timeout error. + pub timeout: Duration, + + /// Whether to retry once on timeout with double the timeout. + /// + /// When `true`, the middleware makes up to + /// [`max_retries`](TimeoutConfig::max_retries) additional attempts, + /// each with double the previous timeout. + pub retry_on_timeout: bool, + + /// Maximum number of retries (0 = no retry, 1 = one retry). + /// + /// Each retry doubles the current timeout. The total number of + /// attempts is `1 + max_retries`. + pub max_retries: u32, +} + +impl Default for TimeoutConfig { + fn default() -> Self { + Self { + timeout: Duration::from_secs(120), + retry_on_timeout: false, + max_retries: 0, + } + } +} + +/// Middleware that wraps tool execution in a timeout. +/// +/// If the tool execution exceeds the configured timeout, returns an +/// error result. Optionally retries once with a longer timeout. +/// Respects the [`CancelSignal`] via `tokio::select!` so that +/// cancellation is not blocked by a slow tool. +/// +/// # Example +/// +/// ```rust,ignore +/// let mw = TimeoutMiddleware::from_secs(120); +/// let mw = TimeoutMiddleware::new(TimeoutConfig { +/// timeout: Duration::from_secs(60), +/// retry_on_timeout: false, +/// max_retries: 0, +/// }); +/// ``` +pub struct TimeoutMiddleware { + config: TimeoutConfig, +} + +impl TimeoutMiddleware { + /// Create a timeout middleware with the given configuration. + /// + /// `config.timeout` controls the per-tool execution deadline. + /// When `config.retry_on_timeout` is `true`, a timed-out call is + /// retried up to `config.max_retries` additional times with an + /// increasing back-off. + /// + /// For simpler construction see [`from_secs`](Self::from_secs) or + /// [`none`](Self::none). + #[must_use] + pub fn new(config: TimeoutConfig) -> Self { + Self { config } + } + + /// Create a timeout middleware with a fixed timeout in seconds. + /// + /// Uses default retry settings (one retry with double timeout). + #[must_use] + pub fn from_secs(secs: u64) -> Self { + Self { + config: TimeoutConfig { + timeout: Duration::from_secs(secs), + ..TimeoutConfig::default() + }, + } + } + + /// Create a timeout middleware with no timeout (pass-through). + /// + /// Useful for testing or when timeouts are handled elsewhere. + #[must_use] + pub fn none() -> Self { + Self { + config: TimeoutConfig { + timeout: Duration::MAX, + retry_on_timeout: false, + max_retries: 0, + }, + } + } +} + +impl ToolMiddleware for TimeoutMiddleware { + fn name(&self) -> &'static str { + "timeout" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + let config = self.config.clone(); + let tool_name = ctx.tool_name.clone(); + let cancel = Arc::clone(&ctx.cancel); + + Box::pin(async move { + let mut attempt = 0u32; + let mut current_timeout = config.timeout; + + loop { + let result_future = next.dispatch(ctx); + let attempt_for_log = attempt; + + tokio::select! { + result = tokio::time::timeout(current_timeout, result_future) => { + if let Ok(dispatch_result) = result { + return dispatch_result; + } + attempt = attempt.saturating_add(1); + if config.retry_on_timeout && attempt <= config.max_retries { + tracing::warn!( + tool = %tool_name, + attempt = attempt_for_log, + timeout_secs = current_timeout.as_secs(), + "tool execution timed out, retrying" + ); + current_timeout = current_timeout.saturating_mul(2); + continue; + } + tracing::error!( + tool = %tool_name, + timeout_secs = current_timeout.as_secs(), + "tool execution timed out" + ); + return ToolDispatchResult::err( + &tool_name, + format!( + "Tool '{}' timed out after {}s", + tool_name, + current_timeout.as_secs() + ), + current_timeout, + ); + } + () = cancel.notified() => { + return ToolDispatchResult::err( + &tool_name, + format!("Tool '{tool_name}' cancelled"), + Duration::ZERO, + ); + } + } + } + }) + } +} + +// =================================================== +// UnknownToolMiddleware +// =================================================== + +/// Middleware that suggests alternatives when a tool is not found. +/// +/// This middleware wraps the core dispatch. When the inner dispatch +/// produces a "tool not found" error, this middleware intercepts it, +/// computes a string-similarity score against all registered tools, +/// and appends a suggestion to the error message. +/// +/// # Similarity Metric +/// +/// Uses a simple normalized longest-common-substring ratio. This is +/// fast and effective for the common case of minor typos (e.g. +/// `"bash"` → `"basj"`). +/// +/// # Example +/// +/// ```rust,ignore +/// let mw = UnknownToolMiddleware::new(); +/// // If tool "basj" is not found, error message will say: +/// // "Tool 'basj' not found. Did you mean 'bash'?" +/// ``` +pub struct UnknownToolMiddleware { + /// The tool registry, used to enumerate available tool names + /// for suggestions. + /// + /// Owned via `Arc` so the middleware can independently list + /// tool names without needing to extract them from the chain. + registry: Arc, + + /// Minimum similarity score (0.0–1.0) to suggest an alternative. + /// + /// Only tools with a score at or above this threshold are suggested. + /// Defaults to `0.4`. + suggestion_threshold: f64, +} + +impl UnknownToolMiddleware { + /// Create a new unknown-tool middleware with default settings. + /// + /// Uses a [`suggestion_threshold`](Self::with_threshold) of `0.4`, + /// which balances catching common typos against false-positive + /// suggestions. + /// + /// For a custom threshold see [`with_threshold`](Self::with_threshold). + #[must_use] + pub fn new(registry: Arc) -> Self { + Self { + registry, + suggestion_threshold: 0.4, + } + } + + /// Create with a custom similarity threshold. + /// + /// Lower values produce more suggestions (more false positives). + /// Higher values require closer matches. + #[must_use] + pub fn with_threshold(registry: Arc, threshold: f64) -> Self { + Self { + registry, + suggestion_threshold: threshold.clamp(0.0, 1.0), + } + } + + /// Compute similarity between two strings using a normalized + /// longest-common-subsequence approach. + /// + /// Returns a value between 0.0 (completely different) and 1.0 + /// (identical). + fn similarity(a: &str, b: &str) -> f64 { + if a.is_empty() && b.is_empty() { + return 1.0; + } + if a.is_empty() || b.is_empty() { + return 0.0; + } + + let a_lower = a.to_lowercase(); + let b_lower = b.to_lowercase(); + + if a_lower == b_lower { + return 1.0; + } + + // Use longest common subsequence length as similarity metric + let a_chars: Vec = a_lower.chars().collect(); + let b_chars: Vec = b_lower.chars().collect(); + let lcs_len = Self::lcs_length(&a_chars, &b_chars); + + // Tool names are short strings, so u32 is sufficient and avoids + // usize→f64 precision loss on 64-bit targets. + let max_len = u32::try_from(a_chars.len().max(b_chars.len())).unwrap_or(u32::MAX); + let lcs_u32 = u32::try_from(lcs_len).unwrap_or(u32::MAX); + + // Check for prefix match bonus + let prefix_len = a_chars + .iter() + .zip(b_chars.iter()) + .take_while(|(a, b)| a == b) + .count(); + let prefix_bonus = if prefix_len > 0 { + let p = u32::try_from(prefix_len).unwrap_or(u32::MAX); + f64::from(p) / f64::from(max_len) * 0.1 + } else { + 0.0 + }; + + f64::from(lcs_u32) / f64::from(max_len) + prefix_bonus + } + + /// Compute the length of the longest common subsequence. + /// Compute the length of the longest common subsequence of two character slices. + /// + /// Uses an iterative dynamic-programming approach with only two rows + /// to keep memory usage O(min(a, b)). + fn lcs_length(a: &[char], b: &[char]) -> usize { + let mut prev = vec![0usize; b.len().saturating_add(1)]; + let mut curr = vec![0usize; b.len().saturating_add(1)]; + + for &a_ch in a { + for (j, &b_ch) in b.iter().enumerate() { + let j_idx = j.saturating_add(1); + *curr.get_mut(j_idx).unwrap_or(&mut 0) = if a_ch == b_ch { + prev.get(j_idx.saturating_sub(1)) + .copied() + .unwrap_or(0) + .saturating_add(1) + } else { + prev.get(j_idx) + .copied() + .unwrap_or(0) + .max(curr.get(j_idx.saturating_sub(1)).copied().unwrap_or(0)) + }; + } + std::mem::swap(&mut prev, &mut curr); + curr.fill(0); + } + + *prev.get(b.len()).unwrap_or(&0) + } + + /// Find the best matching tool name from a list, given a threshold. + /// + /// Returns the name with the highest similarity score that meets or + /// exceeds `threshold`. Returns `None` when no candidate scores high + /// enough. + fn find_best_match_inner<'a>( + requested: &str, + available: &[&'a str], + threshold: f64, + ) -> Option<(&'a str, f64)> { + let mut best: Option<(&'a str, f64)> = None; + for &name in available { + let score = Self::similarity(requested, name); + if score >= threshold { + match best { + Some((_, best_score)) if score <= best_score => {} + _ => best = Some((name, score)), + } + } + } + best + } + + /// Check if a result looks like a "tool not found" error. + /// + /// Only considers single [`Text`](ToolContent::Text) results whose + /// lowercased body contains `"not found"`. + /// [`Multipart`](ToolContent::Multipart) results and non-error results + /// always return `false`. + fn is_tool_not_found(result: &ToolDispatchResult) -> bool { + if !result.is_error { + return false; + } + let msg = match &result.output { + ToolContent::Text(t) => t.to_lowercase(), + ToolContent::Multipart(_) => return false, + }; + msg.contains("not found") + } +} + +impl ToolMiddleware for UnknownToolMiddleware { + fn name(&self) -> &'static str { + "unknown_tool" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + let tool_name = ctx.tool_name.clone(); + let registry_names = self.registry.tool_names(); + let threshold = self.suggestion_threshold; + + Box::pin(async move { + let mut result = next.dispatch(ctx).await; + + if Self::is_tool_not_found(&result) { + let available_refs: Vec<&str> = registry_names.iter().map(String::as_str).collect(); + + if let Some((suggestion, score)) = + Self::find_best_match_inner(&tool_name, &available_refs, threshold) + { + tracing::info!( + requested = %tool_name, + suggestion = %suggestion, + score = %score, + "suggesting alternative tool" + ); + // Append suggestion to the error message + if let ToolContent::Text(ref mut msg) = result.output { + *msg = format!("{msg}. Did you mean '{suggestion}'?"); + } + } + } + + result + }) + } +} + +// =================================================== +// Tests +// =================================================== + +#[cfg(test)] +mod tests { + use super::*; + use crate::tool::{Tool, ToolContext, ToolError, ToolOutput, ToolSchema}; + use serde_json::json; + use std::future::Future; + use std::pin::Pin; + use std::sync::atomic::{AtomicBool, Ordering}; + + // ================================================== + // Test tools + // ================================================== + + /// A simple echo tool for testing. + struct EchoTool; + + impl Tool for EchoTool { + fn name(&self) -> &str { + "echo" + } + fn description(&self) -> &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> + Send + '_>> { + let msg = input["message"].as_str().unwrap_or("").to_string(); + Box::pin(async move { Ok(ToolOutput::text(msg)) }) + } + } + + /// A tool that always errors. + struct ErrorTool; + + impl Tool for ErrorTool { + fn name(&self) -> &str { + "error_tool" + } + fn description(&self) -> &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> + Send + '_>> { + Box::pin(async move { Err(ToolError::Execution("deliberate error".to_string())) }) + } + } + + /// A slow tool for testing timeouts. + struct SlowTool { + delay_ms: u64, + } + + impl Tool for SlowTool { + fn name(&self) -> &str { + "slow" + } + fn description(&self) -> &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> + Send + '_>> { + let delay = self.delay_ms; + Box::pin(async move { + tokio::time::sleep(Duration::from_millis(delay)).await; + Ok(ToolOutput::text("finally done")) + }) + } + } + + // ================================================== + // Helpers + // ================================================== + + fn test_registry() -> Arc { + 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(), + session_id: Uuid::new_v4(), + turn_number: 1, + cancel: Arc::new(CancelSignal::new()), + permission: PermissionCheck::Allow, + tool_context: ToolContext::default(), + } + } + + // ================================================== + // Pipeline builder tests + // ================================================== + + #[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"] + ); + } + + // ================================================== + // Core dispatch tests + // ================================================== + + #[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 => 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 => 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"); + } + + // ================================================== + // PermissionMiddleware tests + // ================================================== + + #[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 => 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::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 => 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); + } + + // ================================================== + // TimeoutMiddleware tests + // ================================================== + + #[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 => panic!("expected Text, got {other:?}"), + } + } + + // ================================================== + // UnknownToolMiddleware tests + // ================================================== + + #[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_eq!(UnknownToolMiddleware::similarity("", ""), 1.0); + assert_eq!(UnknownToolMiddleware::similarity("a", ""), 0.0); + } + + // ================================================== + // dispatch_all tests + // ================================================== + + #[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()); + } + + // ================================================== + // Ordering tests + // ================================================== + + #[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 => panic!("expected Text, got {other:?}"), + } + } + + // ================================================== + // Integration: full pipeline + // ================================================== + + #[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(®istry))) + .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 => 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::with_threshold( + Arc::clone(®istry), + 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 => panic!("expected Text, got {other:?}"), + } + } + + // ================================================== + // Short-circuit test + // ================================================== + + /// A middleware that tracks whether it was reached. + struct ReachTracker { + reached: Arc, + } + + impl ToolMiddleware for ReachTracker { + fn name(&self) -> &str { + "reach_tracker" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + 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" + ); + } +}