flashmind_types/

tool.rs

//! Tool trait, registry, context, and result types for agent function calling.
//!
//! # Key types
//!
//! | Type | Role |
//! |------|------|
//! | [`Tool`] | Trait for tools callable by the agent (name, description, params, execute) |
//! | [`ToolRegistry`] | Registry of available tools with aliases |
//! | [`ToolContext`] | Per-invocation context passed to `execute()` |
//! | [`ToolResult`] | Result of a tool execution returned to the LLM |
//! | [`FileDiff`] | A file diff produced by a tool |
//! | [`ForbiddenCmd`] | Command restriction pattern (regex-based) |

use async_trait::async_trait;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::any::Any;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio_util::sync::CancellationToken;

use crate::event::Source;
use crate::llm::ToolDefinition;
use crate::{AgentLlmConfig, LlmProvider};

// ---------------------------------------------------------------------------
// ForbiddenCmd
// ---------------------------------------------------------------------------

/// A forbidden command configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForbiddenCmd {
    /// Regex pattern to match against the full command line (command + args).
    pub command: String,
    /// Reason shown to user when blocked.
    pub reason: String,
    /// If true, inject a warning message and let the LLM reconsider instead of rejecting.
    /// Like tool-aware RAG, this will retry the turn with additional context.
    #[serde(default)]
    pub reconsider: bool,
}

impl ForbiddenCmd {
    /// Validate the forbidden command config.
    pub fn validate(&self) -> std::result::Result<(), Vec<String>> {
        let mut errors = Vec::new();
        if self.command.is_empty() {
            errors.push("forbidden command name cannot be empty".to_string());
        }
        if regex::Regex::new(&self.command).is_err() {
            errors.push(format!("invalid regex for name: '{}'", self.command));
        }
        if errors.is_empty() {
            Ok(())
        } else {
            Err(errors)
        }
    }
}

// ---------------------------------------------------------------------------
// Command allow-list
// ---------------------------------------------------------------------------

/// Trait for checking whether a command is allowed to execute.
///
/// Implementations live in the app layer; tools receive this as
/// `Arc<dyn CommandAllowList>`. When no allow-list is set (e.g. in tests
/// or desktop), all commands execute freely.
pub trait CommandAllowList: Send + Sync {
    /// Check whether a command is pre-approved (permanent config or session).
    fn is_allowed(&self, command: &str) -> bool;

    /// Record a session-level approval pattern (glob-style).
    ///
    /// The `pattern` argument expects a glob-style string (e.g. `"git commit -m '*'"`,
    /// `"cargo build"`). Patterns are matched against the full command line at execution time.
    fn add_session_pattern(&self, pattern: &str);
}

// ---------------------------------------------------------------------------
// InterruptPayload
// ---------------------------------------------------------------------------

/// Typed payload for [`ToolResult::Interrupt`].
///
/// Tools define concrete structs implementing this trait. The CLI/UI layer
/// downcasts via [`as_any`](Self::as_any) to handle each type specifically.
pub trait InterruptPayload: Send + Sync + std::fmt::Debug {
    /// Downcast support — return `self` as `&dyn Any`.
    fn as_any(&self) -> &dyn Any;

    /// Human-readable description of the interrupt, used as the tool result
    /// string in the conversation so the LLM knows what happened.
    fn display_output(&self) -> String;
}

// ---------------------------------------------------------------------------
// FileDiff
// ---------------------------------------------------------------------------

/// A single line in a unified diff output.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DiffLine {
    /// A line that was added to the file.
    Added { line: u64, content: String },
    /// A line that was removed from the file.
    Removed { line: u64, content: String },
    /// An unchanged context line surrounding the actual changes.
    /// Provides locality so a reviewer can see *where* the edit landed.
    Context { line: u64, content: String },
}

/// A file diff produced by a tool.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileDiff {
    /// Absolute or relative path to the modified file.
    pub path: String,
    /// List of added/removed lines that make up this diff.
    pub diff: Vec<DiffLine>,
}

// ---------------------------------------------------------------------------
// ToolResult (tool execution result — distinct from message::ToolResult)
// ---------------------------------------------------------------------------

/// Result of a tool execution, returned by [`Tool::execute`] and fed back to
/// the LLM as a `"role": "tool"` message.
///
/// Three variants cover the full lifecycle of a tool call:
///
/// - **`Success`** — The tool ran without error. The `output` field contains
///   the text payload sent back to the model. Optional `sources` (URLs the
///   agent can cite) and `diffs` (unified diffs for file-mutating tools) may
///   be attached.
/// - **`Failure`** — The tool encountered an error (invalid arguments, I/O
///   failure, etc.). The `output` field carries a human- and LLM-readable
///   error description. Unlike `Success`, no sources or diffs are included.
/// - **`Interrupt`** — The tool needs user interaction before continuing
///   (e.g. plan approval, choice picker). The agent loop should halt and
///   surface the `output` (structured as JSON) to the caller for handling.
///
/// # Construction helpers
///
/// Prefer the associated constructors on the `impl` block rather than building
/// variants directly:
///
/// ```ignore
/// use flashmind_types::tool::ToolResult;
///
/// // Basic success
/// ToolResult::success(call_id, "file written")
///
/// // Failure with error message
/// ToolResult::failure(call_id, "permission denied")
///
/// // Success with file diffs attached
/// ToolResult::success_with_diffs(call_id, "replaced", vec![diff])
///
/// // Interrupt for interactive prompt
/// ToolResult::interrupt(call_id, Arc::new(MyPayload { ... }))
///
/// // Chain .with_sources() onto any Success
/// ToolResult::success(call_id, results).with_sources(srcs)
/// ```
///
/// # Relation to [`message::ToolResult`](crate::message::ToolResult)
///
/// This type lives in the *tool* module and represents the raw outcome of
/// executing a tool. It is converted into a [`Message`](crate::message::Message)
/// with `role: "tool"` (via [`AgentEvent::ToolResult`](crate::event::AgentEvent::ToolResult))
/// before being appended to the conversation history sent to the LLM.
#[derive(Clone)]
pub enum ToolResult {
    /// Tool executed successfully.
    Success {
        tool_call_id: String,
        output: String,
        sources: Vec<Source>,
        diffs: Vec<FileDiff>,
        /// New agent working directory after a successful workspace switch.
        working_dir: Option<PathBuf>,
    },
    /// Tool execution failed.
    Failure {
        tool_call_id: String,
        output: String,
    },
    /// The agent loop should stop and return this result to the caller for
    /// interactive handling. The typed payload is downcast by the caller to
    /// determine the interrupt kind.
    Interrupt {
        tool_call_id: String,
        payload: Arc<dyn InterruptPayload>,
    },
}

impl std::fmt::Debug for ToolResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Success {
                tool_call_id,
                output,
                ..
            } => f
                .debug_struct("Success")
                .field("tool_call_id", tool_call_id)
                .field("output", output)
                .finish(),
            Self::Failure {
                tool_call_id,
                output,
            } => f
                .debug_struct("Failure")
                .field("tool_call_id", tool_call_id)
                .field("output", output)
                .finish(),
            Self::Interrupt {
                tool_call_id,
                payload,
            } => f
                .debug_struct("Interrupt")
                .field("tool_call_id", tool_call_id)
                .field("output", &payload.display_output())
                .finish(),
        }
    }
}

impl ToolResult {
    /// Create a basic success result with no sources or diffs.
    pub fn success(tool_call_id: &str, output: impl Into<String>) -> Self {
        Self::Success {
            tool_call_id: tool_call_id.into(),
            output: output.into(),
            sources: Vec::new(),
            diffs: Vec::new(),
            working_dir: None,
        }
    }

    /// Create a failure result with an error message.
    pub fn failure(tool_call_id: &str, output: impl Into<String>) -> Self {
        Self::Failure {
            tool_call_id: tool_call_id.into(),
            output: output.into(),
        }
    }

    /// Create a success result that includes file diffs for tools that
    /// mutate content (file write, text replace, etc.).
    pub fn success_with_diffs(
        tool_call_id: &str,
        output: impl Into<String>,
        diffs: Vec<FileDiff>,
    ) -> Self {
        Self::Success {
            tool_call_id: tool_call_id.into(),
            output: output.into(),
            sources: Vec::new(),
            diffs,
            working_dir: None,
        }
    }

    /// Create an interrupt result that pauses the agent loop for interactive
    /// user handling (command approval, OAuth, etc.).
    ///
    /// NOTE: the agent turn ends without a tool result in the conversation,
    /// leaving a dangling tool call. Callers must add a `ConversationEntry::tool`
    /// with the same `tool_call_id` after handling the interrupt.
    pub fn interrupt(tool_call_id: &str, payload: Arc<dyn InterruptPayload>) -> Self {
        Self::Interrupt {
            tool_call_id: tool_call_id.into(),
            payload,
        }
    }

    /// Attach source URLs to a `Success` result. Returns `self` unchanged if
    /// the variant is not `Success`.
    ///
    /// This is designed for method chaining:
    /// ```ignore
    /// ToolResult::success(id, "results").with_sources(srcs)
    /// ```
    pub fn with_sources(mut self, sources: Vec<Source>) -> Self {
        if let Self::Success {
            sources: ref mut s, ..
        } = self
        {
            *s = sources;
        }
        self
    }

    /// Change this agent's working directory after the call succeeds.
    /// Tools returning this effect must set `concurrent()` to false.
    pub fn with_working_dir(mut self, path: PathBuf) -> Self {
        if let Self::Success { working_dir, .. } = &mut self {
            *working_dir = Some(path);
        }
        self
    }

    /// Working directory selected by a successful tool call.
    pub fn working_dir(&self) -> Option<&Path> {
        match self {
            Self::Success { working_dir, .. } => working_dir.as_deref(),
            _ => None,
        }
    }

    /// Return the tool call ID echoed from the original invocation.
    pub fn tool_call_id(&self) -> &str {
        match self {
            Self::Success { tool_call_id, .. }
            | Self::Failure { tool_call_id, .. }
            | Self::Interrupt { tool_call_id, .. } => tool_call_id,
        }
    }

    /// Return the text output payload (success result, error message, or
    /// interrupt display string).
    pub fn output(&self) -> String {
        match self {
            Self::Success { output, .. } | Self::Failure { output, .. } => output.clone(),
            Self::Interrupt { payload, .. } => payload.display_output(),
        }
    }

    /// Return the typed interrupt payload, if this is an Interrupt variant.
    pub fn payload(&self) -> Option<&Arc<dyn InterruptPayload>> {
        match self {
            Self::Interrupt { payload, .. } => Some(payload),
            _ => None,
        }
    }

    /// Check whether this result represents a successful execution.
    pub fn is_success(&self) -> bool {
        matches!(self, Self::Success { .. })
    }

    /// Check whether this result is an interrupt requiring user interaction.
    pub fn is_interrupt(&self) -> bool {
        matches!(self, Self::Interrupt { .. })
    }

    /// Return the attached source references, or an empty slice for
    /// non-success variants.
    pub fn sources(&self) -> &[Source] {
        match self {
            Self::Success { sources, .. } => sources,
            _ => &[],
        }
    }

    /// Return the attached file diffs, or an empty slice for non-success
    /// variants.
    pub fn diffs(&self) -> &[FileDiff] {
        match self {
            Self::Success { diffs, .. } => diffs,
            _ => &[],
        }
    }

    /// Truncate output to fit within byte and line limits.
    /// Appends `\n[truncated]` when the output was cut.
    pub fn truncate_output(&mut self, max_bytes: usize, max_lines: usize) {
        let output = match self {
            Self::Success { output, .. } | Self::Failure { output, .. } => output,
            Self::Interrupt { .. } => return,
        };

        if output.len() <= max_bytes && max_lines == usize::MAX {
            return;
        }

        let byte_cut = output.floor_char_boundary(max_bytes);

        let line_cut = output
            .match_indices('\n')
            .nth(max_lines.saturating_sub(1))
            .map(|(i, _)| i)
            .unwrap_or(output.len());

        let cut = byte_cut.min(line_cut);
        if cut < output.len() {
            output.truncate(cut);
            output.push_str("\n[truncated]");
        }
    }
}

// ---------------------------------------------------------------------------
// parse_args
// ---------------------------------------------------------------------------

/// Deserialize tool arguments from a serde_json::Value into a typed struct.
/// Returns a typed error mentioning which tool failed to parse.
pub fn parse_args<T: DeserializeOwned>(tool_name: &str, args: Value) -> anyhow::Result<T> {
    serde_json::from_value(args)
        .map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
}

// ---------------------------------------------------------------------------
// Tool trait
// ---------------------------------------------------------------------------

/// Trait for tools that can be invoked by the agent.
///
/// Implementations define a name, description, JSON schema parameters, and async execution
/// logic. Tools are registered in a [`ToolRegistry`] and become available to the LLM during
/// conversation turns via function calling.
///
/// # Required methods
///
/// - [`name`](Self::name) — unique identifier used by the LLM to invoke the tool
/// - [`description`](Self::description) — shown to the LLM to decide when to call the tool
/// - [`parameters`](Self::parameters) — JSON Schema object describing expected arguments
/// - [`execute`](Self::execute) — async logic that runs when the tool is called
///
/// # Optional overrides
///
/// - [`humanize`](Self::humanize) — display-friendly summary (default: comma-separated arg keys)
/// - [`timeout_secs`](Self::timeout_secs) — per-tool timeout (default: uses global default)
/// - [`max_output_bytes`](Self::max_output_bytes) — output size limit (default: 256 KiB)
/// - [`max_output_lines`](Self::max_output_lines) — output line count limit (default: 10,000)
#[async_trait]
pub trait Tool: Send + Sync {
    /// Unique tool name (used for invocation). Must match the name the LLM sees in its function definitions.
    fn name(&self) -> &str;
    /// Human-readable description for the LLM. Used to help the model decide when to invoke this tool.
    fn description(&self) -> &str;
    /// JSON Schema object describing the tool's parameters.
    fn parameters(&self) -> Value;
    /// Execute the tool with given arguments. Returns a [`ToolResult`] on success.
    async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult>;
    /// Optional per-tool timeout in seconds. Returns `None` to use the global default.
    fn timeout_secs(&self) -> Option<u64> {
        None
    }

    /// Maximum allowed output size in bytes. Tools that produce large payloads
    /// (e.g. base64-encoded images, arbitrary MCP responses) should return
    /// `usize::MAX` to bypass the limit.
    fn max_output_bytes(&self) -> usize {
        256 * 1024
    }

    /// Maximum allowed output line count. Tools that produce large payloads
    /// should return `usize::MAX` to bypass the limit.
    fn max_output_lines(&self) -> usize {
        10_000
    }

    /// Whether this tool may run at the same time as other tools from the
    /// same model reply, when the agent allows parallel tool calls. Return
    /// `false` for tools that share state or act on one device, so they run
    /// alone between concurrent batches.
    fn concurrent(&self) -> bool {
        true
    }

    /// Whether a running call can continue as a managed background task.
    fn supports_background(&self) -> bool {
        false
    }

    /// Generate a human-readable description of what this tool call does,
    /// suitable for display in a TUI or chat card (e.g. "Reading file Cargo.toml").
    fn humanize(&self, args: &Value) -> String {
        let Some(obj) = args.as_object() else {
            return self.name().to_string();
        };
        let keys: Vec<&str> = obj.keys().map(|k| k.as_str()).collect();
        if keys.is_empty() {
            return self.name().to_string();
        }
        let summary = keys.join(", ");
        if summary.len() > 80 {
            let truncated: String = summary.chars().take(77).collect();
            format!("{truncated}…")
        } else {
            summary
        }
    }
}

// ---------------------------------------------------------------------------
// ToolContext
// ---------------------------------------------------------------------------

/// A request to move a running tool call into the background.
#[derive(Clone, Debug, Default)]
pub struct BackgroundRequest {
    requested: CancellationToken,
}

impl BackgroundRequest {
    pub fn request(&self) {
        self.requested.cancel();
    }

    pub async fn requested(&self) {
        self.requested.cancelled().await;
    }
}

/// Per-invocation context passed to [`Tool::execute`] methods.
///
/// Carries everything a tool needs: the call ID for building results, raw JSON arguments,
/// cancellation support, and the working directory for path resolution.
///
/// # Key fields
///
/// | Field | Purpose |
/// |-------|---------|
/// | [`tool_call_id`](Self::tool_call_id) | Echo into [`ToolResult::tool_call_id`] so results match calls |
/// | [`args`](Self::args) | Raw JSON arguments sent by the LLM — parse with [`parse_args`](Self::parse_args) |
/// | [`working_dir`](Self::working_dir) | Chat workspace root; relative paths resolve against this |
pub struct ToolContext<'a> {
    /// The unique ID for this tool call. Echo into [`ToolResult::tool_call_id`] when returning.
    pub tool_call_id: &'a str,
    /// Raw JSON arguments from the LLM. Use [`parse_args`](Self::parse_args) to deserialize.
    pub args: Value,
    /// Current parent configuration, including its cache routing identity.
    pub llm: Option<&'a AgentLlmConfig>,
    /// Current provider for tools that delegate model requests.
    pub provider: Option<Arc<dyn LlmProvider>>,
    /// Working directory for the current agent (typically chat workspace). Relative paths resolve here.
    pub working_dir: Option<&'a PathBuf>,
    /// Cancellation token for cooperative cancellation. Access via [`cancel_token`](Self::cancel_token) or [`child_token`](Self::child_token).
    cancel_token: &'a CancellationToken,
    /// Optional channel for streaming incremental output during execution.
    progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
    background_request: Option<&'a BackgroundRequest>,
}

impl<'a> ToolContext<'a> {
    /// Create a new tool execution context.
    pub fn new(
        tool_call_id: &'a str,
        args: Value,
        working_dir: Option<&'a PathBuf>,
        cancel_token: &'a CancellationToken,
    ) -> Self {
        Self {
            tool_call_id,
            args,
            llm: None,
            provider: None,
            working_dir,
            cancel_token,
            progress_tx: None,
            background_request: None,
        }
    }

    /// Pass the current agent configuration to the tool.
    pub fn with_llm(mut self, llm: &'a AgentLlmConfig) -> Self {
        self.llm = Some(llm);
        self
    }

    /// Pass the current provider to tools that delegate model requests.
    pub fn with_provider(mut self, provider: Arc<dyn LlmProvider>) -> Self {
        self.provider = Some(provider);
        self
    }

    /// Attach a progress channel for streaming incremental output.
    pub fn with_progress(mut self, tx: tokio::sync::mpsc::UnboundedSender<String>) -> Self {
        self.progress_tx = Some(tx);
        self
    }

    /// Attach the user's control for backgrounding this running call.
    pub fn with_background_request(mut self, request: &'a BackgroundRequest) -> Self {
        self.background_request = Some(request);
        self
    }

    pub fn background_request(&self) -> Option<&BackgroundRequest> {
        self.background_request
    }

    /// Clone the live output destination when execution moves to an owned task.
    pub fn progress_sender(&self) -> Option<tokio::sync::mpsc::UnboundedSender<String>> {
        self.progress_tx.clone()
    }

    /// Send a progress line through the channel (no-op if no channel is set).
    pub fn progress(&self, line: &str) {
        if let Some(ref tx) = self.progress_tx {
            let _ = tx.send(line.to_string());
        }
    }

    /// Parse the raw JSON args into a typed struct. Fails with a descriptive error if invalid.
    pub fn parse_args<T: DeserializeOwned>(&self, tool_name: &str) -> anyhow::Result<T> {
        T::deserialize(&self.args)
            .map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
    }

    /// Access the cancellation token for cooperative cancellation checks.
    ///
    /// For long-running operations, prefer [`child_token`](Self::child_token) to get
    /// a derived token that can be passed to async sub-operations.
    pub fn cancel_token(&self) -> &CancellationToken {
        self.cancel_token
    }

    /// Create a child cancellation token. Cancels when the parent turn is cancelled.
    /// Use this for spawned tasks so they're automatically cancelled when the user aborts.
    pub fn child_token(&self) -> CancellationToken {
        self.cancel_token.child_token()
    }
}

// ---------------------------------------------------------------------------
// ToolRegistry
// ---------------------------------------------------------------------------

/// Registry of tools available to an agent.
///
/// Stores tools in a HashMap for O(1) lookup by name and supports aliases
/// (e.g., `"bash_exec"` → `"exec"`) that are resolved transparently but not
/// exposed in `definitions()` or `list()`.
#[derive(Clone)]
pub struct ToolRegistry {
    tools: HashMap<String, Arc<dyn Tool>>,
    aliases: HashMap<String, String>,
    truncate_on_overflow: bool,
}

impl ToolRegistry {
    /// Create an empty tool registry.
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
            aliases: HashMap::new(),
            truncate_on_overflow: false,
        }
    }

    /// When `true`, oversized output is truncated instead of failing the tool
    /// call. Default: `false` (fail with an error so the agent can retry).
    pub fn set_truncate_on_overflow(&mut self, enabled: bool) {
        self.truncate_on_overflow = enabled;
    }

    /// Register a tool (replaces existing tool with same name).
    pub fn register(&mut self, tool: Arc<dyn Tool>) {
        tracing::debug!(tool = tool.name(), "registered tool");
        self.tools.insert(tool.name().to_string(), tool);
    }

    /// Remove a tool by name. Returns true if the tool was present.
    pub fn remove(&mut self, name: &str) -> bool {
        let removed = self.tools.remove(name).is_some();
        if removed {
            let tools = &self.tools;
            self.aliases.retain(|_, target| tools.contains_key(target));
        }
        removed
    }

    /// Remove every tool whose name starts with one of `prefixes`. Aliases
    /// pointing at removed tools are also dropped.
    pub fn strip_prefixes(&mut self, prefixes: &[&str]) {
        let to_remove: Vec<String> = self
            .tools
            .keys()
            .filter(|name| prefixes.iter().any(|p| name.starts_with(p)))
            .cloned()
            .collect();
        for name in to_remove {
            self.tools.remove(&name);
        }
        let tools = &self.tools;
        self.aliases.retain(|_, target| tools.contains_key(target));
    }

    /// Register an alias so that lookups for `alias` resolve to `target`.
    /// Aliases are not exposed in `definitions()` or `list()`.
    pub fn alias(&mut self, alias: &str, target: &str) {
        tracing::debug!(alias, target, "registered tool alias");
        self.aliases.insert(alias.to_string(), target.to_string());
    }

    /// Look up a tool by name (checks aliases if direct lookup fails).
    /// Always resolves any registered tool, even if not in the active set.
    pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
        let result = self.tools.get(name).or_else(|| {
            if let Some(target) = self.aliases.get(name) {
                tracing::debug!(
                    alias = name,
                    target = target.as_str(),
                    "resolved tool alias"
                );
                self.tools.get(target)
            } else {
                None
            }
        });

        if result.is_none() {
            tracing::warn!(tool = name, "tool not found in registry");
        }

        result
    }

    /// Returns true if a tool with the given name is registered (without logging).
    pub fn contains(&self, name: &str) -> bool {
        self.resolve(name).is_some()
    }

    /// Look up a tool directly or through an alias, without logging.
    fn resolve(&self, name: &str) -> Option<&Arc<dyn Tool>> {
        self.tools.get(name).or_else(|| {
            self.aliases
                .get(name)
                .and_then(|target| self.tools.get(target))
        })
    }

    /// Keep only tools whose name satisfies the predicate.
    pub fn retain(&mut self, mut f: impl FnMut(&str) -> bool) {
        self.tools.retain(|name, _| f(name));
        let tools = &self.tools;
        self.aliases.retain(|_, target| tools.contains_key(target));
    }

    /// List all registered tool names (sorted).
    pub fn list(&self) -> Vec<&str> {
        let mut names: Vec<&str> = self.tools.keys().map(|s| s.as_str()).collect();
        names.sort();
        names
    }

    /// List all registered tools with their names and descriptions (sorted by name).
    pub fn list_with_descriptions(&self) -> Vec<(&str, &str)> {
        let mut items: Vec<(&str, &str)> = self
            .tools
            .values()
            .map(|t| (t.name(), t.description()))
            .collect();
        items.sort_by_key(|(name, _)| *name);
        items
    }

    /// Execute a tool call from the LLM.
    ///
    /// Looks up the tool by name, builds a [`ToolContext`], and runs it. Returns
    /// a [`ToolResult::Failure`] if the tool is not found or if execution errors.
    /// If the tool defines [`timeout_secs`](Tool::timeout_secs), the call is
    /// wrapped in a timeout. Tools without a timeout are cancelled immediately
    /// when `cancel_token` fires.
    pub async fn execute(
        &self,
        call: &crate::ToolCall,
        working_dir: Option<&PathBuf>,
        cancel_token: &CancellationToken,
    ) -> ToolResult {
        self.execute_inner(call, working_dir, cancel_token, None)
            .await
    }

    /// Like [`Self::execute`] but attaches a progress channel so tools can stream
    /// incremental output (e.g. stdout lines from bash).
    pub async fn execute_with_progress(
        &self,
        call: &crate::ToolCall,
        working_dir: Option<&PathBuf>,
        cancel_token: &CancellationToken,
        progress_tx: tokio::sync::mpsc::UnboundedSender<String>,
    ) -> ToolResult {
        self.execute_inner(call, working_dir, cancel_token, Some(progress_tx))
            .await
    }

    async fn execute_inner(
        &self,
        call: &crate::ToolCall,
        working_dir: Option<&PathBuf>,
        cancel_token: &CancellationToken,
        progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
    ) -> ToolResult {
        let mut ctx = ToolContext::new(&call.id, call.arguments.clone(), working_dir, cancel_token);
        if let Some(tx) = progress_tx {
            ctx = ctx.with_progress(tx);
        }
        self.execute_in_context(call, ctx).await
    }

    /// Execute with parent configuration and optional progress reporting.
    pub async fn execute_in_context(
        &self,
        call: &crate::ToolCall,
        ctx: ToolContext<'_>,
    ) -> ToolResult {
        let Some(tool) = self.get(&call.name) else {
            return ToolResult::failure(&call.id, format!("Unknown tool: {}", call.name));
        };
        let cancel_token = ctx.cancel_token;
        let child_token = cancel_token.child_token();
        let ctx = ToolContext {
            cancel_token: &child_token,
            ..ctx
        };
        let fut = tool.execute(ctx);

        // Always race against cancel_token so user-initiated cancellation
        // works regardless of whether the tool defined a hard timeout.
        let timeout_secs = tool.timeout_secs();
        let run = async {
            match timeout_secs {
                Some(secs) => tokio::time::timeout(std::time::Duration::from_secs(secs), fut)
                    .await
                    .map_err(|_| secs),
                None => Ok(fut.await),
            }
        };
        let outcome = tokio::select! {
            biased;
            _ = cancel_token.cancelled() => {
                child_token.cancel();
                return ToolResult::failure(
                    &call.id,
                    format!("Tool '{}' cancelled", call.name),
                );
            }
            r = run => match r {
                Ok(result) => result,
                Err(secs) => {
                    child_token.cancel();
                    return ToolResult::failure(
                        &call.id,
                        format!("Tool '{}' timed out after {secs}s", call.name),
                    );
                }
            },
        };

        let result = match outcome {
            Ok(result) => result,
            Err(e) => ToolResult::failure(&call.id, format!("Tool error: {e:#}")),
        };

        enforce_output_limits(
            &call.id,
            result,
            tool.max_output_bytes(),
            tool.max_output_lines(),
            self.truncate_on_overflow,
        )
    }

    /// Whether `call` may run alongside other calls. Unknown tools fail at
    /// once, so they count as concurrent.
    pub fn is_concurrent(&self, call: &crate::ToolCall) -> bool {
        self.resolve(&call.name)
            .is_none_or(|tool| tool.concurrent())
    }

    pub fn supports_background(&self, call: &crate::ToolCall) -> bool {
        self.resolve(&call.name)
            .is_some_and(|tool| tool.supports_background())
    }

    /// Human-readable summary of a tool call's arguments.
    pub fn humanize(&self, call: &crate::ToolCall) -> String {
        self.get(&call.name)
            .map(|t| t.humanize(&call.arguments))
            .unwrap_or_else(|| call.name.clone())
    }

    /// Generate LLM tool definitions for all registered tools.
    pub fn definitions(&self) -> Vec<ToolDefinition> {
        let mut defs: Vec<ToolDefinition> = self
            .tools
            .values()
            .map(|t| ToolDefinition {
                name: t.name().to_string(),
                description: t.description().to_string(),
                parameters: t.parameters(),
            })
            .collect();

        defs.sort_by(|a, b| a.name.cmp(&b.name));

        tracing::debug!(count = defs.len(), "generated tool definitions");

        defs
    }
}

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

/// Enforce output size limits on a tool result.
///
/// When the tool opts into [`truncate_on_overflow`](Tool::truncate_on_overflow),
/// oversized output is truncated with a `\n[truncated]` suffix. Otherwise the
/// result is replaced with a failure so the agent can retry with narrower args.
/// A failure is always truncated, keeping its start (such as an exit code).
fn enforce_output_limits(
    tool_call_id: &str,
    mut result: ToolResult,
    max_bytes: usize,
    max_lines: usize,
    truncate: bool,
) -> ToolResult {
    let output = match &result {
        ToolResult::Success { output, .. } => output,
        ToolResult::Failure { .. } => {
            result.truncate_output(max_bytes, max_lines);
            return result;
        }
        ToolResult::Interrupt { .. } => return result,
    };
    let bytes = output.len();
    let lines = output.lines().count();
    let over_bytes = bytes > max_bytes;
    let over_lines = lines > max_lines;

    if !over_bytes && !over_lines {
        return result;
    }

    if truncate {
        result.truncate_output(max_bytes, max_lines);
        return result;
    }

    if over_bytes {
        ToolResult::failure(
            tool_call_id,
            format!(
                "Output too large ({bytes} bytes, limit {max_bytes}). \
                 Use more specific arguments to reduce output size."
            ),
        )
    } else {
        ToolResult::failure(
            tool_call_id,
            format!(
                "Output too many lines ({lines}, limit {max_lines}). \
                 Use more specific arguments to reduce output size."
            ),
        )
    }
}

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

    #[test]
    fn failed_output_is_truncated_not_passed_through() {
        let output = "x\n".repeat(100);
        for truncate in [false, true] {
            let result = enforce_output_limits(
                "call",
                ToolResult::failure("call", output.clone()),
                50,
                usize::MAX,
                truncate,
            );
            let ToolResult::Failure { output, .. } = result else {
                panic!("a failure stays a failure");
            };
            assert!(output.len() < 70, "{output}");
            assert!(output.ends_with("[truncated]"), "{output}");
        }
    }

    struct DummyTool(&'static str);

    #[async_trait]
    impl Tool for DummyTool {
        fn name(&self) -> &str {
            self.0
        }
        fn description(&self) -> &str {
            "dummy"
        }
        fn parameters(&self) -> Value {
            serde_json::json!({})
        }
        async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
            Ok(ToolResult::success(ctx.tool_call_id, "ok"))
        }
    }

    #[test]
    fn test_registry_empty() {
        let registry = ToolRegistry::new();
        assert_eq!(registry.list().len(), 0);
    }

    #[test]
    fn test_registry_definitions_empty() {
        let registry = ToolRegistry::new();
        assert!(registry.definitions().is_empty());
    }

    #[test]
    fn test_registry_lookup() {
        let mut registry = ToolRegistry::new();
        registry.register(Arc::new(DummyTool("exec")));
        registry.register(Arc::new(DummyTool("slack_react")));
        registry.register(Arc::new(DummyTool("browser")));

        assert_eq!(registry.definitions().len(), 3);

        // get() resolves tools by name
        assert!(registry.get("slack_react").is_some());
    }

    #[test]
    fn test_registry_alias_contains_and_remove_cleanup() {
        let mut registry = ToolRegistry::new();
        registry.register(Arc::new(DummyTool("exec")));
        registry.alias("bash_exec", "exec");

        assert!(registry.contains("exec"));
        assert!(registry.contains("bash_exec"));
        assert!(registry.get("bash_exec").is_some());

        assert!(registry.remove("exec"));
        assert!(!registry.contains("exec"));
        assert!(!registry.contains("bash_exec"));
        assert!(registry.get("bash_exec").is_none());
    }

    /// Tool that sleeps forever (until cancelled). Declares a long
    /// timeout_secs so it takes the "timeout" branch of ToolRegistry::execute.
    struct SlowTimeoutTool;

    #[async_trait]
    impl Tool for SlowTimeoutTool {
        fn name(&self) -> &str {
            "slow"
        }
        fn description(&self) -> &str {
            "sleeps forever"
        }
        fn parameters(&self) -> Value {
            serde_json::json!({})
        }
        fn timeout_secs(&self) -> Option<u64> {
            Some(600)
        }
        async fn execute(&self, _ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
            // Sleep beyond the test budget; cancellation must short-circuit.
            tokio::time::sleep(std::time::Duration::from_secs(60)).await;
            Ok(ToolResult::success("unused", "should not return"))
        }
    }

    #[tokio::test]
    async fn registry_cancel_short_circuits_tool_with_timeout_secs() {
        // Regression: a tool that sets timeout_secs() used to bypass the
        // cancel-token race entirely. Cancelling now must still abort it.
        let mut registry = ToolRegistry::new();
        registry.register(Arc::new(SlowTimeoutTool));

        let cancel = CancellationToken::new();
        let cancel_clone = cancel.clone();

        let call = crate::ToolCall {
            id: "c1".into(),
            name: "slow".into(),
            arguments: serde_json::json!({}),
        };

        let registry = Arc::new(registry);
        let registry2 = registry.clone();
        let exec = tokio::spawn(async move { registry2.execute(&call, None, &cancel).await });

        // Let the tool start, then cancel.
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        cancel_clone.cancel();

        // Must return within 2s of cancellation, NOT after 600s timeout.
        let result = tokio::time::timeout(std::time::Duration::from_secs(2), exec)
            .await
            .expect("registry did not return within 2s of cancellation")
            .unwrap();

        assert!(!result.is_success(), "expected failure on cancel");
        assert!(
            result.output().contains("cancelled"),
            "expected 'cancelled' in output, got: {}",
            result.output()
        );
    }
}