Skip to main content

flashmind_types/
tool.rs

1//! Tool trait, registry, context, and result types for agent function calling.
2//!
3//! # Key types
4//!
5//! | Type | Role |
6//! |------|------|
7//! | [`Tool`] | Trait for tools callable by the agent (name, description, params, execute) |
8//! | [`ToolRegistry`] | Registry of available tools with aliases |
9//! | [`ToolContext`] | Per-invocation context passed to `execute()` |
10//! | [`ToolResult`] | Result of a tool execution returned to the LLM |
11//! | [`FileDiff`] | A file diff produced by a tool |
12//! | [`ForbiddenCmd`] | Command restriction pattern (regex-based) |
13//! | [`ToolGate`] | Host check every call passes before it runs |
14
15use async_trait::async_trait;
16use futures_util::future::BoxFuture;
17use serde::de::DeserializeOwned;
18use serde::{Deserialize, Serialize};
19use serde_json::Value;
20use std::any::Any;
21use std::collections::HashMap;
22use std::path::{Path, PathBuf};
23use std::sync::Arc;
24use std::time::Duration;
25use tempfile::TempDir;
26use tokio::fs::OpenOptions;
27use tokio::io::AsyncWriteExt;
28use tokio::sync::OnceCell;
29use tokio_util::sync::CancellationToken;
30use uuid::Uuid;
31
32use crate::event::Source;
33use crate::llm::ToolDefinition;
34use crate::{AgentLlmConfig, ContentPart, LlmProvider};
35
36// ---------------------------------------------------------------------------
37// ForbiddenCmd
38// ---------------------------------------------------------------------------
39
40/// A forbidden command configuration.
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct ForbiddenCmd {
43    /// Regex pattern to match against the full command line (command + args).
44    pub command: String,
45    /// Reason shown to user when blocked.
46    pub reason: String,
47    /// If true, inject a warning message and let the LLM reconsider instead of rejecting.
48    /// Like tool-aware RAG, this will retry the turn with additional context.
49    #[serde(default)]
50    pub reconsider: bool,
51}
52
53impl ForbiddenCmd {
54    /// Validate the forbidden command config.
55    pub fn validate(&self) -> std::result::Result<(), Vec<String>> {
56        let mut errors = Vec::new();
57        if self.command.is_empty() {
58            errors.push("forbidden command name cannot be empty".to_string());
59        }
60        if regex::Regex::new(&self.command).is_err() {
61            errors.push(format!("invalid regex for name: '{}'", self.command));
62        }
63        if errors.is_empty() {
64            Ok(())
65        } else {
66            Err(errors)
67        }
68    }
69}
70
71// ---------------------------------------------------------------------------
72// Command allow-list
73// ---------------------------------------------------------------------------
74
75/// Trait for checking whether a command is allowed to execute.
76///
77/// Implementations live in the app layer; tools receive this as
78/// `Arc<dyn CommandAllowList>`. When no allow-list is set (e.g. in tests
79/// or desktop), all commands execute freely.
80pub trait CommandAllowList: Send + Sync {
81    /// Check whether a command is pre-approved (permanent config or session).
82    fn is_allowed(&self, command: &str) -> bool;
83
84    /// Record a session-level approval pattern (glob-style).
85    ///
86    /// The `pattern` argument expects a glob-style string (e.g. `"git commit -m '*'"`,
87    /// `"cargo build"`). Patterns are matched against the full command line at execution time.
88    fn add_session_pattern(&self, pattern: &str);
89}
90
91// ---------------------------------------------------------------------------
92// ToolGate
93// ---------------------------------------------------------------------------
94
95/// A host's check on every tool call, made where the registry executes it.
96///
97/// A registry with a gate asks it before each call. Sub-agents get the
98/// gate of the registry they are given, or their manager's, so the same
99/// rules hold for every agent the host runs.
100#[async_trait]
101pub trait ToolGate: Send + Sync {
102    /// Decide about `call` before it runs.
103    async fn check(&self, call: &crate::ToolCall) -> GateDecision;
104}
105
106/// What a [`ToolGate`] decides about one call.
107pub enum GateDecision {
108    /// Run it.
109    Allow,
110    /// Do not run it; the reason is the call's result.
111    Deny(String),
112    /// Ask someone and wait for the answer. Only this call waits.
113    Ask(GateAsk),
114}
115
116/// A question a gate asked, and how long the call waits for its answer.
117pub struct GateAsk {
118    /// Resolves to the answer, such as a oneshot the host completes.
119    pub answer: BoxFuture<'static, GateAnswer>,
120    /// A call left unanswered this long is refused.
121    pub timeout: Duration,
122}
123
124/// The answer to a [`GateAsk`].
125#[derive(Debug, Clone, PartialEq, Eq)]
126pub enum GateAnswer {
127    Allow,
128    /// Refused; the reason is the call's result.
129    Deny(String),
130}
131
132impl std::fmt::Debug for GateDecision {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        match self {
135            Self::Allow => f.write_str("Allow"),
136            Self::Deny(reason) => f.debug_tuple("Deny").field(reason).finish(),
137            Self::Ask(ask) => f
138                .debug_struct("Ask")
139                .field("timeout", &ask.timeout)
140                .finish(),
141        }
142    }
143}
144
145// ---------------------------------------------------------------------------
146// InterruptPayload
147// ---------------------------------------------------------------------------
148
149/// Typed payload for [`ToolResult::Interrupt`].
150///
151/// Tools define concrete structs implementing this trait. The CLI/UI layer
152/// downcasts via [`as_any`](Self::as_any) to handle each type specifically.
153pub trait InterruptPayload: Send + Sync + std::fmt::Debug {
154    /// Downcast support — return `self` as `&dyn Any`.
155    fn as_any(&self) -> &dyn Any;
156
157    /// Human-readable description of the interrupt, used as the tool result
158    /// string in the conversation so the LLM knows what happened.
159    fn display_output(&self) -> String;
160}
161
162// ---------------------------------------------------------------------------
163// FileDiff
164// ---------------------------------------------------------------------------
165
166/// A single line in a unified diff output.
167#[derive(Debug, Clone, Serialize, Deserialize)]
168pub enum DiffLine {
169    /// A line that was added to the file.
170    Added { line: u64, content: String },
171    /// A line that was removed from the file.
172    Removed { line: u64, content: String },
173    /// An unchanged context line surrounding the actual changes.
174    /// Provides locality so a reviewer can see *where* the edit landed.
175    Context { line: u64, content: String },
176}
177
178/// A file diff produced by a tool.
179#[derive(Debug, Clone, Serialize, Deserialize)]
180pub struct FileDiff {
181    /// Absolute or relative path to the modified file.
182    pub path: String,
183    /// List of added/removed lines that make up this diff.
184    pub diff: Vec<DiffLine>,
185}
186
187// ---------------------------------------------------------------------------
188// ToolResult (tool execution result — distinct from message::ToolResult)
189// ---------------------------------------------------------------------------
190
191/// Result of a tool execution, returned by [`Tool::execute`] and fed back to
192/// the LLM as a `"role": "tool"` message.
193///
194/// Three variants cover the full lifecycle of a tool call:
195///
196/// - **`Success`** — The tool ran without error. The `output` field contains
197///   the text payload sent back to the model. Optional `sources` (URLs the
198///   agent can cite) and `diffs` (unified diffs for file-mutating tools) may
199///   be attached.
200/// - **`Failure`** — The tool encountered an error (invalid arguments, I/O
201///   failure, etc.). The `output` field carries a human- and LLM-readable
202///   error description. Unlike `Success`, no sources or diffs are included.
203/// - **`Interrupt`** — The tool needs user interaction before continuing
204///   (e.g. plan approval, choice picker). The agent loop should halt and
205///   surface the `output` (structured as JSON) to the caller for handling.
206///
207/// # Construction helpers
208///
209/// Prefer the associated constructors on the `impl` block rather than building
210/// variants directly:
211///
212/// ```ignore
213/// use flashmind_types::tool::ToolResult;
214///
215/// // Basic success
216/// ToolResult::success(call_id, "file written")
217///
218/// // Failure with error message
219/// ToolResult::failure(call_id, "permission denied")
220///
221/// // Success with file diffs attached
222/// ToolResult::success_with_diffs(call_id, "replaced", vec![diff])
223///
224/// // Interrupt for interactive prompt
225/// ToolResult::interrupt(call_id, Arc::new(MyPayload { ... }))
226///
227/// // Chain .with_sources() onto any Success
228/// ToolResult::success(call_id, results).with_sources(srcs)
229/// ```
230///
231/// # Relation to [`message::ToolResult`](crate::message::ToolResult)
232///
233/// This type lives in the *tool* module and represents the raw outcome of
234/// executing a tool. It is converted into a [`Message`](crate::message::Message)
235/// with `role: "tool"` (via [`AgentEvent::ToolResult`](crate::event::AgentEvent::ToolResult))
236/// before being appended to the conversation history sent to the LLM.
237#[derive(Clone)]
238pub enum ToolResult {
239    /// Tool executed successfully.
240    Success {
241        tool_call_id: String,
242        content: Vec<ContentPart>,
243        sources: Vec<Source>,
244        diffs: Vec<FileDiff>,
245        /// New agent working directory after a successful workspace switch.
246        working_dir: Option<PathBuf>,
247    },
248    /// Tool execution failed.
249    Failure {
250        tool_call_id: String,
251        content: Vec<ContentPart>,
252    },
253    /// The agent loop should stop and return this result to the caller for
254    /// interactive handling. The typed payload is downcast by the caller to
255    /// determine the interrupt kind.
256    Interrupt {
257        tool_call_id: String,
258        payload: Arc<dyn InterruptPayload>,
259    },
260}
261
262impl std::fmt::Debug for ToolResult {
263    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
264        match self {
265            Self::Success { tool_call_id, .. } => f
266                .debug_struct("Success")
267                .field("tool_call_id", tool_call_id)
268                .field("output", &self.output())
269                .finish(),
270            Self::Failure { tool_call_id, .. } => f
271                .debug_struct("Failure")
272                .field("tool_call_id", tool_call_id)
273                .field("output", &self.output())
274                .finish(),
275            Self::Interrupt {
276                tool_call_id,
277                payload,
278            } => f
279                .debug_struct("Interrupt")
280                .field("tool_call_id", tool_call_id)
281                .field("output", &payload.display_output())
282                .finish(),
283        }
284    }
285}
286
287impl ToolResult {
288    /// Create a basic success result with no sources or diffs.
289    pub fn success(tool_call_id: &str, output: impl Into<String>) -> Self {
290        Self::success_with_content(
291            tool_call_id,
292            vec![ContentPart::Text {
293                text: output.into(),
294            }],
295        )
296    }
297
298    /// Create a successful result containing ordered text and media parts.
299    pub fn success_with_content(tool_call_id: &str, content: Vec<ContentPart>) -> Self {
300        Self::Success {
301            tool_call_id: tool_call_id.into(),
302            content,
303            sources: Vec::new(),
304            diffs: Vec::new(),
305            working_dir: None,
306        }
307    }
308
309    /// Append content to a successful result, preserving its order.
310    pub fn with_parts(mut self, parts: Vec<ContentPart>) -> Self {
311        if let Self::Success { content, .. } = &mut self {
312            content.extend(parts);
313        }
314        self
315    }
316
317    /// Return the ordered content of a completed tool call.
318    pub fn content(&self) -> &[ContentPart] {
319        match self {
320            Self::Success { content, .. } | Self::Failure { content, .. } => content,
321            Self::Interrupt { .. } => &[],
322        }
323    }
324
325    /// Append text to a completed result.
326    pub fn append_text(&mut self, text: &str) {
327        match self {
328            Self::Success { content, .. } | Self::Failure { content, .. } => {
329                match content.last_mut() {
330                    Some(ContentPart::Text { text: last }) => last.push_str(text),
331                    _ => content.push(ContentPart::Text {
332                        text: text.to_owned(),
333                    }),
334                }
335            }
336            Self::Interrupt { .. } => {}
337        }
338    }
339
340    /// Create a failure result with an error message.
341    pub fn failure(tool_call_id: &str, output: impl Into<String>) -> Self {
342        Self::failure_with_content(
343            tool_call_id,
344            vec![ContentPart::Text {
345                text: output.into(),
346            }],
347        )
348    }
349
350    /// Create a failed result containing ordered text and media parts.
351    pub fn failure_with_content(tool_call_id: &str, content: Vec<ContentPart>) -> Self {
352        Self::Failure {
353            tool_call_id: tool_call_id.into(),
354            content,
355        }
356    }
357
358    /// Create a success result that includes file diffs for tools that
359    /// mutate content (file write, text replace, etc.).
360    pub fn success_with_diffs(
361        tool_call_id: &str,
362        output: impl Into<String>,
363        diffs: Vec<FileDiff>,
364    ) -> Self {
365        Self::Success {
366            tool_call_id: tool_call_id.into(),
367            content: vec![ContentPart::Text {
368                text: output.into(),
369            }],
370            sources: Vec::new(),
371            diffs,
372            working_dir: None,
373        }
374    }
375
376    /// Create an interrupt result that pauses the agent loop for interactive
377    /// user handling (command approval, OAuth, etc.).
378    ///
379    /// NOTE: the agent turn ends without a tool result in the conversation,
380    /// leaving a dangling tool call. Callers must add a `ConversationEntry::tool`
381    /// with the same `tool_call_id` after handling the interrupt.
382    pub fn interrupt(tool_call_id: &str, payload: Arc<dyn InterruptPayload>) -> Self {
383        Self::Interrupt {
384            tool_call_id: tool_call_id.into(),
385            payload,
386        }
387    }
388
389    /// Attach source URLs to a `Success` result. Returns `self` unchanged if
390    /// the variant is not `Success`.
391    ///
392    /// This is designed for method chaining:
393    /// ```ignore
394    /// ToolResult::success(id, "results").with_sources(srcs)
395    /// ```
396    pub fn with_sources(mut self, sources: Vec<Source>) -> Self {
397        if let Self::Success {
398            sources: ref mut s, ..
399        } = self
400        {
401            *s = sources;
402        }
403        self
404    }
405
406    /// Change this agent's working directory after the call succeeds.
407    /// Tools returning this effect must set `concurrent()` to false.
408    pub fn with_working_dir(mut self, path: PathBuf) -> Self {
409        if let Self::Success { working_dir, .. } = &mut self {
410            *working_dir = Some(path);
411        }
412        self
413    }
414
415    /// Working directory selected by a successful tool call.
416    pub fn working_dir(&self) -> Option<&Path> {
417        match self {
418            Self::Success { working_dir, .. } => working_dir.as_deref(),
419            _ => None,
420        }
421    }
422
423    /// Return the tool call ID echoed from the original invocation.
424    pub fn tool_call_id(&self) -> &str {
425        match self {
426            Self::Success { tool_call_id, .. }
427            | Self::Failure { tool_call_id, .. }
428            | Self::Interrupt { tool_call_id, .. } => tool_call_id,
429        }
430    }
431
432    /// Return the text output payload (success result, error message, or
433    /// interrupt display string).
434    pub fn output(&self) -> String {
435        match self {
436            Self::Success { content, .. } | Self::Failure { content, .. } => content
437                .iter()
438                .filter_map(|part| match part {
439                    ContentPart::Text { text } => Some(text.as_str()),
440                    _ => None,
441                })
442                .collect(),
443            Self::Interrupt { payload, .. } => payload.display_output(),
444        }
445    }
446
447    /// Return the typed interrupt payload, if this is an Interrupt variant.
448    pub fn payload(&self) -> Option<&Arc<dyn InterruptPayload>> {
449        match self {
450            Self::Interrupt { payload, .. } => Some(payload),
451            _ => None,
452        }
453    }
454
455    /// Check whether this result represents a successful execution.
456    pub fn is_success(&self) -> bool {
457        matches!(self, Self::Success { .. })
458    }
459
460    /// Check whether this result is an interrupt requiring user interaction.
461    pub fn is_interrupt(&self) -> bool {
462        matches!(self, Self::Interrupt { .. })
463    }
464
465    /// Return the attached source references, or an empty slice for
466    /// non-success variants.
467    pub fn sources(&self) -> &[Source] {
468        match self {
469            Self::Success { sources, .. } => sources,
470            _ => &[],
471        }
472    }
473
474    /// Return the attached file diffs, or an empty slice for non-success
475    /// variants.
476    pub fn diffs(&self) -> &[FileDiff] {
477        match self {
478            Self::Success { diffs, .. } => diffs,
479            _ => &[],
480        }
481    }
482
483    /// Truncate output to fit within byte and line limits.
484    /// Appends `\n[truncated]` when the output was cut.
485    pub fn truncate_output(&mut self, max_bytes: usize, max_lines: usize) {
486        let output = self.output();
487        let content = match self {
488            Self::Success { content, .. } | Self::Failure { content, .. } => content,
489            Self::Interrupt { .. } => return,
490        };
491        if output.len() <= max_bytes && max_lines == usize::MAX {
492            return;
493        }
494        let byte_cut = output.floor_char_boundary(max_bytes);
495        let line_cut = if max_lines == 0 {
496            0
497        } else {
498            output
499                .match_indices('\n')
500                .nth(max_lines - 1)
501                .map(|(index, _)| index)
502                .unwrap_or(output.len())
503        };
504        let cut = byte_cut.min(line_cut);
505        if cut >= output.len() {
506            return;
507        }
508        let mut remaining = cut;
509        let mut marked = false;
510        for part in content {
511            if let ContentPart::Text { text } = part {
512                if marked {
513                    text.clear();
514                } else if text.len() > remaining {
515                    text.truncate(remaining);
516                    text.push_str("\n[truncated]");
517                    marked = true;
518                } else {
519                    remaining -= text.len();
520                }
521            }
522        }
523    }
524}
525
526// ---------------------------------------------------------------------------
527// parse_args
528// ---------------------------------------------------------------------------
529
530/// Deserialize tool arguments from a serde_json::Value into a typed struct.
531/// Returns a typed error mentioning which tool failed to parse.
532pub fn parse_args<T: DeserializeOwned>(tool_name: &str, args: Value) -> anyhow::Result<T> {
533    serde_json::from_value(args)
534        .map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
535}
536
537// ---------------------------------------------------------------------------
538// Tool trait
539// ---------------------------------------------------------------------------
540
541/// Trait for tools that can be invoked by the agent.
542///
543/// Implementations define a name, description, JSON schema parameters, and async execution
544/// logic. Tools are registered in a [`ToolRegistry`] and become available to the LLM during
545/// conversation turns via function calling.
546///
547/// # Required methods
548///
549/// - [`name`](Self::name) — unique identifier used by the LLM to invoke the tool
550/// - [`description`](Self::description) — shown to the LLM to decide when to call the tool
551/// - [`parameters`](Self::parameters) — JSON Schema object describing expected arguments
552/// - [`execute`](Self::execute) — async logic that runs when the tool is called
553///
554/// # Optional overrides
555///
556/// - [`humanize`](Self::humanize) — display-friendly summary (default: comma-separated arg keys)
557/// - [`timeout_secs`](Self::timeout_secs) — per-tool timeout (default: uses global default)
558/// - [`max_output_bytes`](Self::max_output_bytes) — output size limit (default: 256 KiB)
559/// - [`max_output_lines`](Self::max_output_lines) — output line count limit (default: 10,000)
560#[async_trait]
561pub trait Tool: Send + Sync {
562    /// Unique tool name (used for invocation). Must match the name the LLM sees in its function definitions.
563    fn name(&self) -> &str;
564    /// Human-readable description for the LLM. Used to help the model decide when to invoke this tool.
565    fn description(&self) -> &str;
566    /// JSON Schema object describing the tool's parameters.
567    fn parameters(&self) -> Value;
568    /// Execute the tool with given arguments. Returns a [`ToolResult`] on success.
569    async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult>;
570    /// Optional per-tool timeout in seconds. Returns `None` to use the global default.
571    fn timeout_secs(&self) -> Option<u64> {
572        None
573    }
574
575    /// Maximum allowed output size in bytes. Tools that produce large payloads
576    /// (e.g. base64-encoded images, arbitrary MCP responses) should return
577    /// `usize::MAX` to bypass the limit.
578    fn max_output_bytes(&self) -> usize {
579        256 * 1024
580    }
581
582    /// Maximum allowed output line count. Tools that produce large payloads
583    /// should return `usize::MAX` to bypass the limit.
584    fn max_output_lines(&self) -> usize {
585        10_000
586    }
587
588    /// Whether this tool may run at the same time as other tools from the
589    /// same model reply, when the agent allows parallel tool calls. Return
590    /// `false` for tools that share state or act on one device, so they run
591    /// alone between concurrent batches.
592    fn concurrent(&self) -> bool {
593        true
594    }
595
596    /// Whether a running call can continue as a managed background task.
597    fn supports_background(&self) -> bool {
598        false
599    }
600
601    /// Generate a human-readable description of what this tool call does,
602    /// suitable for display in a TUI or chat card (e.g. "Reading file Cargo.toml").
603    fn humanize(&self, args: &Value) -> String {
604        let Some(obj) = args.as_object() else {
605            return self.name().to_string();
606        };
607        let keys: Vec<&str> = obj.keys().map(|k| k.as_str()).collect();
608        if keys.is_empty() {
609            return self.name().to_string();
610        }
611        let summary = keys.join(", ");
612        if summary.len() > 80 {
613            let truncated: String = summary.chars().take(77).collect();
614            format!("{truncated}…")
615        } else {
616            summary
617        }
618    }
619}
620
621// ---------------------------------------------------------------------------
622// ToolContext
623// ---------------------------------------------------------------------------
624
625/// A request to move a running tool call into the background.
626#[derive(Clone, Debug, Default)]
627pub struct BackgroundRequest {
628    requested: CancellationToken,
629}
630
631impl BackgroundRequest {
632    pub fn request(&self) {
633        self.requested.cancel();
634    }
635
636    pub async fn requested(&self) {
637        self.requested.cancelled().await;
638    }
639}
640
641/// Per-invocation context passed to [`Tool::execute`] methods.
642///
643/// Carries everything a tool needs: the call ID for building results, raw JSON arguments,
644/// cancellation support, and the working directory for path resolution.
645///
646/// # Key fields
647///
648/// | Field | Purpose |
649/// |-------|---------|
650/// | [`tool_call_id`](Self::tool_call_id) | Echo into [`ToolResult::tool_call_id`] so results match calls |
651/// | [`args`](Self::args) | Raw JSON arguments sent by the LLM — parse with [`parse_args`](Self::parse_args) |
652/// | [`working_dir`](Self::working_dir) | Chat workspace root; relative paths resolve against this |
653pub struct ToolContext<'a> {
654    /// The unique ID for this tool call. Echo into [`ToolResult::tool_call_id`] when returning.
655    pub tool_call_id: &'a str,
656    /// Raw JSON arguments from the LLM. Use [`parse_args`](Self::parse_args) to deserialize.
657    pub args: Value,
658    /// Current parent configuration, including its cache routing identity.
659    pub llm: Option<&'a AgentLlmConfig>,
660    /// Current provider for tools that delegate model requests.
661    pub provider: Option<Arc<dyn LlmProvider>>,
662    /// Working directory for the current agent (typically chat workspace). Relative paths resolve here.
663    pub working_dir: Option<&'a PathBuf>,
664    /// Cancellation token for cooperative cancellation. Access via [`cancel_token`](Self::cancel_token) or [`child_token`](Self::child_token).
665    cancel_token: &'a CancellationToken,
666    /// Optional channel for streaming incremental output during execution.
667    progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
668    background_request: Option<&'a BackgroundRequest>,
669}
670
671impl<'a> ToolContext<'a> {
672    /// Create a new tool execution context.
673    pub fn new(
674        tool_call_id: &'a str,
675        args: Value,
676        working_dir: Option<&'a PathBuf>,
677        cancel_token: &'a CancellationToken,
678    ) -> Self {
679        Self {
680            tool_call_id,
681            args,
682            llm: None,
683            provider: None,
684            working_dir,
685            cancel_token,
686            progress_tx: None,
687            background_request: None,
688        }
689    }
690
691    /// Pass the current agent configuration to the tool.
692    pub fn with_llm(mut self, llm: &'a AgentLlmConfig) -> Self {
693        self.llm = Some(llm);
694        self
695    }
696
697    /// Pass the current provider to tools that delegate model requests.
698    pub fn with_provider(mut self, provider: Arc<dyn LlmProvider>) -> Self {
699        self.provider = Some(provider);
700        self
701    }
702
703    /// Attach a progress channel for streaming incremental output.
704    pub fn with_progress(mut self, tx: tokio::sync::mpsc::UnboundedSender<String>) -> Self {
705        self.progress_tx = Some(tx);
706        self
707    }
708
709    /// Attach the user's control for backgrounding this running call.
710    pub fn with_background_request(mut self, request: &'a BackgroundRequest) -> Self {
711        self.background_request = Some(request);
712        self
713    }
714
715    pub fn background_request(&self) -> Option<&BackgroundRequest> {
716        self.background_request
717    }
718
719    /// Clone the live output destination when execution moves to an owned task.
720    pub fn progress_sender(&self) -> Option<tokio::sync::mpsc::UnboundedSender<String>> {
721        self.progress_tx.clone()
722    }
723
724    /// Send a progress line through the channel (no-op if no channel is set).
725    pub fn progress(&self, line: &str) {
726        if let Some(ref tx) = self.progress_tx {
727            let _ = tx.send(line.to_string());
728        }
729    }
730
731    /// Parse the raw JSON arguments into a typed value without consuming them.
732    ///
733    /// The tool name is included in errors, along with the deserialization error.
734    /// This does not validate the tool's JSON schema, so keep the schema in sync
735    /// with the argument type.
736    ///
737    /// ```
738    /// use std::path::PathBuf;
739    ///
740    /// use flashmind_types::tool::ToolContext;
741    /// use serde::Deserialize;
742    /// use serde_json::json;
743    /// use tokio_util::sync::CancellationToken;
744    ///
745    /// #[derive(Deserialize)]
746    /// struct ReadArgs {
747    ///     path: PathBuf,
748    /// }
749    ///
750    /// let cancel = CancellationToken::new();
751    /// let context = ToolContext::new("call", json!({"path": "src/main.rs"}), None, &cancel);
752    /// let args: ReadArgs = context.parse_args("read")?;
753    /// assert_eq!(args.path, PathBuf::from("src/main.rs"));
754    /// # Ok::<(), anyhow::Error>(())
755    /// ```
756    pub fn parse_args<T: DeserializeOwned>(&self, tool_name: &str) -> anyhow::Result<T> {
757        T::deserialize(&self.args)
758            .map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
759    }
760
761    /// Access the cancellation token for cooperative cancellation checks.
762    ///
763    /// For long-running operations, prefer [`child_token`](Self::child_token) to get
764    /// a derived token that can be passed to async sub-operations.
765    pub fn cancel_token(&self) -> &CancellationToken {
766        self.cancel_token
767    }
768
769    /// Create a child cancellation token. Cancels when the parent turn is cancelled.
770    /// Use this for spawned tasks so they're automatically cancelled when the user aborts.
771    pub fn child_token(&self) -> CancellationToken {
772        self.cancel_token.child_token()
773    }
774}
775
776// ---------------------------------------------------------------------------
777// ToolRegistry
778// ---------------------------------------------------------------------------
779
780/// Registry of tools available to an agent.
781///
782/// Stores tools in a HashMap for O(1) lookup by name and supports aliases
783/// (e.g., `"bash_exec"` → `"exec"`) that are resolved transparently but not
784/// exposed in `definitions()` or `list()`.
785#[derive(Clone)]
786pub struct ToolRegistry {
787    tools: HashMap<String, Arc<dyn Tool>>,
788    aliases: HashMap<String, String>,
789    truncate_on_overflow: bool,
790    output_dir: Arc<OutputDir>,
791    gate: Option<Arc<dyn ToolGate>>,
792}
793
794/// Temp directory for oversized tool output, created on first use and
795/// deleted with its contents when the last registry clone sharing it drops.
796#[derive(Default)]
797struct OutputDir(OnceCell<TempDir>);
798
799impl OutputDir {
800    async fn path(&self) -> std::io::Result<&Path> {
801        let dir = self
802            .0
803            .get_or_try_init(|| async {
804                tempfile::Builder::new()
805                    .prefix("flashmind-output-")
806                    .tempdir()
807            })
808            .await?;
809        Ok(dir.path())
810    }
811}
812
813impl ToolRegistry {
814    /// Create an empty tool registry.
815    pub fn new() -> Self {
816        Self {
817            tools: HashMap::new(),
818            aliases: HashMap::new(),
819            truncate_on_overflow: false,
820            output_dir: Arc::default(),
821            gate: None,
822        }
823    }
824
825    /// Check every call with `gate` before it runs. Clones of the registry
826    /// keep it.
827    pub fn set_gate(&mut self, gate: Arc<dyn ToolGate>) {
828        self.gate = Some(gate);
829    }
830
831    /// The gate calls pass, if any.
832    pub fn gate(&self) -> Option<&Arc<dyn ToolGate>> {
833        self.gate.as_ref()
834    }
835
836    /// When `true`, oversized output is truncated and dropped. Default:
837    /// `false` (the full output is saved to a temp file and the agent gets a
838    /// truncated preview with the file's path). Those files are deleted when
839    /// the registry and all its clones are dropped.
840    pub fn set_truncate_on_overflow(&mut self, enabled: bool) {
841        self.truncate_on_overflow = enabled;
842    }
843
844    /// Register a tool (replaces existing tool with same name).
845    pub fn register(&mut self, tool: Arc<dyn Tool>) {
846        tracing::debug!(tool = tool.name(), "registered tool");
847        self.tools.insert(tool.name().to_string(), tool);
848    }
849
850    /// Remove a tool by name. Returns true if the tool was present.
851    pub fn remove(&mut self, name: &str) -> bool {
852        let removed = self.tools.remove(name).is_some();
853        if removed {
854            let tools = &self.tools;
855            self.aliases.retain(|_, target| tools.contains_key(target));
856        }
857        removed
858    }
859
860    /// Remove every tool whose name starts with one of `prefixes`. Aliases
861    /// pointing at removed tools are also dropped.
862    pub fn strip_prefixes(&mut self, prefixes: &[&str]) {
863        let to_remove: Vec<String> = self
864            .tools
865            .keys()
866            .filter(|name| prefixes.iter().any(|p| name.starts_with(p)))
867            .cloned()
868            .collect();
869        for name in to_remove {
870            self.tools.remove(&name);
871        }
872        let tools = &self.tools;
873        self.aliases.retain(|_, target| tools.contains_key(target));
874    }
875
876    /// Register an alias so that lookups for `alias` resolve to `target`.
877    /// Aliases are not exposed in `definitions()` or `list()`.
878    pub fn alias(&mut self, alias: &str, target: &str) {
879        tracing::debug!(alias, target, "registered tool alias");
880        self.aliases.insert(alias.to_string(), target.to_string());
881    }
882
883    /// Look up a tool by name (checks aliases if direct lookup fails).
884    /// Always resolves any registered tool, even if not in the active set.
885    pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
886        let result = self.tools.get(name).or_else(|| {
887            if let Some(target) = self.aliases.get(name) {
888                tracing::debug!(
889                    alias = name,
890                    target = target.as_str(),
891                    "resolved tool alias"
892                );
893                self.tools.get(target)
894            } else {
895                None
896            }
897        });
898
899        if result.is_none() {
900            tracing::warn!(tool = name, "tool not found in registry");
901        }
902
903        result
904    }
905
906    /// Returns true if a tool with the given name is registered (without logging).
907    pub fn contains(&self, name: &str) -> bool {
908        self.resolve(name).is_some()
909    }
910
911    /// Look up a tool directly or through an alias, without logging.
912    fn resolve(&self, name: &str) -> Option<&Arc<dyn Tool>> {
913        self.tools.get(name).or_else(|| {
914            self.aliases
915                .get(name)
916                .and_then(|target| self.tools.get(target))
917        })
918    }
919
920    /// Keep only tools whose name satisfies the predicate.
921    pub fn retain(&mut self, mut f: impl FnMut(&str) -> bool) {
922        self.tools.retain(|name, _| f(name));
923        let tools = &self.tools;
924        self.aliases.retain(|_, target| tools.contains_key(target));
925    }
926
927    /// List all registered tool names (sorted).
928    pub fn list(&self) -> Vec<&str> {
929        let mut names: Vec<&str> = self.tools.keys().map(|s| s.as_str()).collect();
930        names.sort();
931        names
932    }
933
934    /// List all registered tools with their names and descriptions (sorted by name).
935    pub fn list_with_descriptions(&self) -> Vec<(&str, &str)> {
936        let mut items: Vec<(&str, &str)> = self
937            .tools
938            .values()
939            .map(|t| (t.name(), t.description()))
940            .collect();
941        items.sort_by_key(|(name, _)| *name);
942        items
943    }
944
945    /// Execute a tool call from the LLM.
946    ///
947    /// Looks up the tool by name, builds a [`ToolContext`], and runs it. Returns
948    /// a [`ToolResult::Failure`] if the tool is not found or if execution errors.
949    /// If the tool defines [`timeout_secs`](Tool::timeout_secs), the call is
950    /// wrapped in a timeout. Tools without a timeout are cancelled immediately
951    /// when `cancel_token` fires.
952    pub async fn execute(
953        &self,
954        call: &crate::ToolCall,
955        working_dir: Option<&PathBuf>,
956        cancel_token: &CancellationToken,
957    ) -> ToolResult {
958        self.execute_inner(call, working_dir, cancel_token, None)
959            .await
960    }
961
962    /// Like [`Self::execute`] but attaches a progress channel so tools can stream
963    /// incremental output (e.g. stdout lines from bash).
964    pub async fn execute_with_progress(
965        &self,
966        call: &crate::ToolCall,
967        working_dir: Option<&PathBuf>,
968        cancel_token: &CancellationToken,
969        progress_tx: tokio::sync::mpsc::UnboundedSender<String>,
970    ) -> ToolResult {
971        self.execute_inner(call, working_dir, cancel_token, Some(progress_tx))
972            .await
973    }
974
975    async fn execute_inner(
976        &self,
977        call: &crate::ToolCall,
978        working_dir: Option<&PathBuf>,
979        cancel_token: &CancellationToken,
980        progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
981    ) -> ToolResult {
982        let mut ctx = ToolContext::new(&call.id, call.arguments.clone(), working_dir, cancel_token);
983        if let Some(tx) = progress_tx {
984            ctx = ctx.with_progress(tx);
985        }
986        self.execute_in_context(call, ctx).await
987    }
988
989    /// Execute with parent configuration and optional progress reporting.
990    pub async fn execute_in_context(
991        &self,
992        call: &crate::ToolCall,
993        ctx: ToolContext<'_>,
994    ) -> ToolResult {
995        let Some(tool) = self.get(&call.name) else {
996            return ToolResult::failure(&call.id, format!("Unknown tool: {}", call.name));
997        };
998        if let Some(gate) = &self.gate
999            && let Some(refused) = pass_gate(gate.as_ref(), call, ctx.cancel_token).await
1000        {
1001            return refused;
1002        }
1003        let cancel_token = ctx.cancel_token;
1004        let child_token = cancel_token.child_token();
1005        let ctx = ToolContext {
1006            cancel_token: &child_token,
1007            ..ctx
1008        };
1009        let fut = tool.execute(ctx);
1010
1011        // Always race against cancel_token so user-initiated cancellation
1012        // works regardless of whether the tool defined a hard timeout.
1013        let timeout_secs = tool.timeout_secs();
1014        let run = async {
1015            match timeout_secs {
1016                Some(secs) => tokio::time::timeout(std::time::Duration::from_secs(secs), fut)
1017                    .await
1018                    .map_err(|_| secs),
1019                None => Ok(fut.await),
1020            }
1021        };
1022        let outcome = tokio::select! {
1023            biased;
1024            _ = cancel_token.cancelled() => {
1025                child_token.cancel();
1026                return ToolResult::failure(
1027                    &call.id,
1028                    format!("Tool '{}' cancelled", call.name),
1029                );
1030            }
1031            r = run => match r {
1032                Ok(result) => result,
1033                Err(secs) => {
1034                    child_token.cancel();
1035                    return ToolResult::failure(
1036                        &call.id,
1037                        format!("Tool '{}' timed out after {secs}s", call.name),
1038                    );
1039                }
1040            },
1041        };
1042
1043        let result = match outcome {
1044            Ok(result) => result,
1045            Err(e) => ToolResult::failure(&call.id, format!("Tool error: {e:#}")),
1046        };
1047
1048        self.enforce_output_limits(
1049            &call.id,
1050            result,
1051            tool.max_output_bytes(),
1052            tool.max_output_lines(),
1053        )
1054        .await
1055    }
1056
1057    /// Whether `call` may run alongside other calls. Unknown tools fail at
1058    /// once, so they count as concurrent.
1059    pub fn is_concurrent(&self, call: &crate::ToolCall) -> bool {
1060        self.resolve(&call.name)
1061            .is_none_or(|tool| tool.concurrent())
1062    }
1063
1064    pub fn supports_background(&self, call: &crate::ToolCall) -> bool {
1065        self.resolve(&call.name)
1066            .is_some_and(|tool| tool.supports_background())
1067    }
1068
1069    /// Human-readable summary of a tool call's arguments.
1070    pub fn humanize(&self, call: &crate::ToolCall) -> String {
1071        self.get(&call.name)
1072            .map(|t| t.humanize(&call.arguments))
1073            .unwrap_or_else(|| call.name.clone())
1074    }
1075
1076    /// Generate LLM tool definitions for all registered tools.
1077    pub fn definitions(&self) -> Vec<ToolDefinition> {
1078        let mut defs: Vec<ToolDefinition> = self
1079            .tools
1080            .values()
1081            .map(|t| ToolDefinition {
1082                name: t.name().to_string(),
1083                description: t.description().to_string(),
1084                parameters: t.parameters(),
1085            })
1086            .collect();
1087
1088        defs.sort_by(|a, b| a.name.cmp(&b.name));
1089
1090        tracing::debug!(count = defs.len(), "generated tool definitions");
1091
1092        defs
1093    }
1094}
1095
1096impl Default for ToolRegistry {
1097    fn default() -> Self {
1098        Self::new()
1099    }
1100}
1101
1102/// Asks `gate` about `call`, waiting for an answer when it asks someone.
1103/// `None` lets the call run; otherwise the refusal is its result.
1104async fn pass_gate(
1105    gate: &dyn ToolGate,
1106    call: &crate::ToolCall,
1107    cancel: &CancellationToken,
1108) -> Option<ToolResult> {
1109    let ask = match gate.check(call).await {
1110        GateDecision::Allow => return None,
1111        GateDecision::Deny(reason) => return Some(ToolResult::failure(&call.id, reason)),
1112        GateDecision::Ask(ask) => ask,
1113    };
1114    let answer = tokio::select! {
1115        _ = cancel.cancelled() => {
1116            return Some(ToolResult::failure(&call.id, format!("Tool '{}' cancelled", call.name)));
1117        }
1118        answer = tokio::time::timeout(ask.timeout, ask.answer) => answer,
1119    };
1120    match answer {
1121        Ok(GateAnswer::Allow) => None,
1122        Ok(GateAnswer::Deny(reason)) => Some(ToolResult::failure(&call.id, reason)),
1123        Err(_) => Some(ToolResult::failure(
1124            &call.id,
1125            format!(
1126                "Not done: nobody answered in time to allow '{}'.",
1127                call.name
1128            ),
1129        )),
1130    }
1131}
1132
1133impl ToolRegistry {
1134    /// Enforce output size limits on a tool result.
1135    ///
1136    /// Oversized output is cut to the limits. Unless the registry opts into
1137    /// [`set_truncate_on_overflow`](ToolRegistry::set_truncate_on_overflow), the
1138    /// full output is first saved to a new file in the registry's temp directory,
1139    /// whose path ends the preview,
1140    /// so the agent can read or search it. If that file cannot be written, a
1141    /// successful result becomes a failure asking for narrower arguments. A
1142    /// failure keeps its start (such as an exit code) either way.
1143    async fn enforce_output_limits(
1144        &self,
1145        tool_call_id: &str,
1146        mut result: ToolResult,
1147        max_bytes: usize,
1148        max_lines: usize,
1149    ) -> ToolResult {
1150        if result.is_interrupt() {
1151            return result;
1152        }
1153        let output = result.output();
1154        let bytes = output.len();
1155        let lines = output.lines().count();
1156        let over_bytes = bytes > max_bytes;
1157        let over_lines = lines > max_lines;
1158
1159        if !over_bytes && !over_lines {
1160            return result;
1161        }
1162
1163        if self.truncate_on_overflow {
1164            result.truncate_output(max_bytes, max_lines);
1165            return result;
1166        }
1167
1168        match self.save_full_output(&output).await {
1169            Ok(path) => {
1170                result.truncate_output(max_bytes, max_lines);
1171                result.append_text(&format!(
1172                    "\n[full output ({bytes} bytes, {lines} lines) saved to {}]",
1173                    path.display()
1174                ));
1175                result
1176            }
1177            Err(e) => {
1178                tracing::warn!(error = %e, "could not save oversized tool output");
1179                if !result.is_success() {
1180                    result.truncate_output(max_bytes, max_lines);
1181                    return result;
1182                }
1183                if over_bytes {
1184                    ToolResult::failure(
1185                        tool_call_id,
1186                        format!(
1187                            "Output too large ({bytes} bytes, limit {max_bytes}). \
1188                             Use more specific arguments to reduce output size."
1189                        ),
1190                    )
1191                } else {
1192                    ToolResult::failure(
1193                        tool_call_id,
1194                        format!(
1195                            "Output too many lines ({lines}, limit {max_lines}). \
1196                             Use more specific arguments to reduce output size."
1197                        ),
1198                    )
1199                }
1200            }
1201        }
1202    }
1203
1204    /// Write `output` to a new, randomly named file in the registry's temp
1205    /// directory, readable only by the current user.
1206    async fn save_full_output(&self, output: &str) -> std::io::Result<PathBuf> {
1207        let path = self
1208            .output_dir
1209            .path()
1210            .await?
1211            .join(format!("{}.txt", Uuid::new_v4()));
1212        let mut options = OpenOptions::new();
1213        options.write(true).create_new(true);
1214        #[cfg(unix)]
1215        options.mode(0o600);
1216        let mut file = options.open(&path).await?;
1217        file.write_all(output.as_bytes()).await?;
1218        file.flush().await?;
1219        Ok(path)
1220    }
1221}
1222
1223#[cfg(test)]
1224mod tests {
1225    use super::*;
1226
1227    /// Return the path named in the `saved to` note at the end of `output`.
1228    fn saved_path(output: &str) -> PathBuf {
1229        let start = output
1230            .rfind("saved to ")
1231            .expect("output names a saved file")
1232            + 9;
1233        PathBuf::from(output[start..].trim_end_matches(']'))
1234    }
1235
1236    #[tokio::test]
1237    async fn failed_output_is_truncated_not_passed_through() {
1238        let output = "x\n".repeat(100);
1239        let mut registry = ToolRegistry::new();
1240        registry.set_truncate_on_overflow(true);
1241        let result = registry
1242            .enforce_output_limits(
1243                "call",
1244                ToolResult::failure("call", output.clone()),
1245                50,
1246                usize::MAX,
1247            )
1248            .await;
1249        assert!(!result.is_success());
1250        let output = result.output();
1251        assert!(output.len() < 70, "{output}");
1252        assert!(output.ends_with("[truncated]"), "{output}");
1253    }
1254
1255    #[tokio::test]
1256    async fn oversized_output_is_saved_to_temp_file() {
1257        let registry = ToolRegistry::new();
1258        let full = "line\n".repeat(100);
1259        let mut paths = Vec::new();
1260        for result in [
1261            ToolResult::success("call", full.clone()),
1262            ToolResult::failure("call", full.clone()),
1263        ] {
1264            let was_success = result.is_success();
1265            let result = registry
1266                .enforce_output_limits("call", result, usize::MAX, 10)
1267                .await;
1268            assert_eq!(result.is_success(), was_success);
1269
1270            let output = result.output();
1271            assert!(output.starts_with(&"line\n".repeat(9)), "{output}");
1272            assert!(output.contains("[truncated]"), "{output}");
1273            assert!(output.contains("(500 bytes, 100 lines)"), "{output}");
1274
1275            let path = saved_path(&output);
1276            let saved = std::fs::read_to_string(&path).expect("saved file is readable");
1277            assert_eq!(saved, full);
1278            paths.push(path);
1279        }
1280
1281        // A clone shares the directory, so only the last drop removes it.
1282        let clone = registry.clone();
1283        drop(registry);
1284        assert!(paths.iter().all(|p| p.exists()));
1285        drop(clone);
1286        assert!(paths.iter().all(|p| !p.exists()));
1287    }
1288
1289    #[test]
1290    fn mixed_result_projects_text_and_keeps_ordered_content() {
1291        let mut result = ToolResult::success_with_content(
1292            "call",
1293            vec![
1294                ContentPart::Text {
1295                    text: "before ".into(),
1296                },
1297                ContentPart::Image {
1298                    media_type: "image/png".into(),
1299                    data: "encoded".into(),
1300                },
1301                ContentPart::Text {
1302                    text: "after".into(),
1303                },
1304            ],
1305        );
1306        assert_eq!(result.output(), "before after");
1307        assert!(
1308            matches!(&result.content()[1], ContentPart::Image { data, .. } if data == "encoded")
1309        );
1310        result.append_text("!");
1311        assert_eq!(result.output(), "before after!");
1312        assert!(!format!("{result:?}").contains("encoded"));
1313    }
1314
1315    #[tokio::test]
1316    async fn limits_apply_to_combined_text_and_leave_images_untouched() {
1317        let registry = ToolRegistry::new();
1318        let result = ToolResult::success_with_content(
1319            "call",
1320            vec![
1321                ContentPart::Text { text: "ab".into() },
1322                ContentPart::Image {
1323                    media_type: "image/png".into(),
1324                    data: "encoded".into(),
1325                },
1326                ContentPart::Text {
1327                    text: "éxyz".into(),
1328                },
1329            ],
1330        );
1331        let result = registry
1332            .enforce_output_limits("call", result, 3, usize::MAX)
1333            .await;
1334        assert!(result.is_success());
1335        assert!(result.output().starts_with("ab\n[truncated]"));
1336        assert!(
1337            matches!(&result.content()[1], ContentPart::Image { data, .. } if data == "encoded")
1338        );
1339        assert_eq!(
1340            std::fs::read_to_string(saved_path(&result.output())).unwrap(),
1341            "abéxyz"
1342        );
1343    }
1344
1345    #[test]
1346    fn failure_content_preserves_media_and_text_projection() {
1347        let mut result = ToolResult::failure_with_content(
1348            "call",
1349            vec![
1350                ContentPart::Text {
1351                    text: "failed".into(),
1352                },
1353                ContentPart::Image {
1354                    media_type: "image/png".into(),
1355                    data: "encoded".into(),
1356                },
1357            ],
1358        );
1359        result.append_text(" details");
1360        assert!(!result.is_success());
1361        assert_eq!(result.output(), "failed details");
1362        assert_eq!(result.content().len(), 3);
1363    }
1364
1365    #[tokio::test]
1366    async fn output_within_limits_is_unchanged() {
1367        let result = ToolRegistry::new()
1368            .enforce_output_limits("call", ToolResult::success("call", "ok"), 10, 10)
1369            .await;
1370        assert_eq!(result.output(), "ok");
1371    }
1372
1373    struct DummyTool(&'static str);
1374
1375    #[async_trait]
1376    impl Tool for DummyTool {
1377        fn name(&self) -> &str {
1378            self.0
1379        }
1380        fn description(&self) -> &str {
1381            "dummy"
1382        }
1383        fn parameters(&self) -> Value {
1384            serde_json::json!({})
1385        }
1386        async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
1387            Ok(ToolResult::success(ctx.tool_call_id, "ok"))
1388        }
1389    }
1390
1391    #[test]
1392    fn test_registry_empty() {
1393        let registry = ToolRegistry::new();
1394        assert_eq!(registry.list().len(), 0);
1395    }
1396
1397    #[test]
1398    fn test_registry_definitions_empty() {
1399        let registry = ToolRegistry::new();
1400        assert!(registry.definitions().is_empty());
1401    }
1402
1403    #[test]
1404    fn test_registry_lookup() {
1405        let mut registry = ToolRegistry::new();
1406        registry.register(Arc::new(DummyTool("exec")));
1407        registry.register(Arc::new(DummyTool("slack_react")));
1408        registry.register(Arc::new(DummyTool("browser")));
1409
1410        assert_eq!(registry.definitions().len(), 3);
1411
1412        // get() resolves tools by name
1413        assert!(registry.get("slack_react").is_some());
1414    }
1415
1416    #[test]
1417    fn test_registry_alias_contains_and_remove_cleanup() {
1418        let mut registry = ToolRegistry::new();
1419        registry.register(Arc::new(DummyTool("exec")));
1420        registry.alias("bash_exec", "exec");
1421
1422        assert!(registry.contains("exec"));
1423        assert!(registry.contains("bash_exec"));
1424        assert!(registry.get("bash_exec").is_some());
1425
1426        assert!(registry.remove("exec"));
1427        assert!(!registry.contains("exec"));
1428        assert!(!registry.contains("bash_exec"));
1429        assert!(registry.get("bash_exec").is_none());
1430    }
1431
1432    /// Tool that sleeps forever (until cancelled). Declares a long
1433    /// timeout_secs so it takes the "timeout" branch of ToolRegistry::execute.
1434    struct SlowTimeoutTool;
1435
1436    #[async_trait]
1437    impl Tool for SlowTimeoutTool {
1438        fn name(&self) -> &str {
1439            "slow"
1440        }
1441        fn description(&self) -> &str {
1442            "sleeps forever"
1443        }
1444        fn parameters(&self) -> Value {
1445            serde_json::json!({})
1446        }
1447        fn timeout_secs(&self) -> Option<u64> {
1448            Some(600)
1449        }
1450        async fn execute(&self, _ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
1451            // Sleep beyond the test budget; cancellation must short-circuit.
1452            tokio::time::sleep(std::time::Duration::from_secs(60)).await;
1453            Ok(ToolResult::success("unused", "should not return"))
1454        }
1455    }
1456
1457    #[tokio::test]
1458    async fn registry_cancel_short_circuits_tool_with_timeout_secs() {
1459        // Regression: a tool that sets timeout_secs() used to bypass the
1460        // cancel-token race entirely. Cancelling now must still abort it.
1461        let mut registry = ToolRegistry::new();
1462        registry.register(Arc::new(SlowTimeoutTool));
1463
1464        let cancel = CancellationToken::new();
1465        let cancel_clone = cancel.clone();
1466
1467        let call = crate::ToolCall {
1468            id: "c1".into(),
1469            name: "slow".into(),
1470            arguments: serde_json::json!({}),
1471        };
1472
1473        let registry = Arc::new(registry);
1474        let registry2 = registry.clone();
1475        let exec = tokio::spawn(async move { registry2.execute(&call, None, &cancel).await });
1476
1477        // Let the tool start, then cancel.
1478        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1479        cancel_clone.cancel();
1480
1481        // Must return within 2s of cancellation, NOT after 600s timeout.
1482        let result = tokio::time::timeout(std::time::Duration::from_secs(2), exec)
1483            .await
1484            .expect("registry did not return within 2s of cancellation")
1485            .unwrap();
1486
1487        assert!(!result.is_success(), "expected failure on cancel");
1488        assert!(
1489            result.output().contains("cancelled"),
1490            "expected 'cancelled' in output, got: {}",
1491            result.output()
1492        );
1493    }
1494
1495    /// Counts its runs.
1496    struct CountedTool(Arc<std::sync::atomic::AtomicUsize>);
1497
1498    #[async_trait]
1499    impl Tool for CountedTool {
1500        fn name(&self) -> &str {
1501            "act"
1502        }
1503        fn description(&self) -> &str {
1504            "acts"
1505        }
1506        fn parameters(&self) -> Value {
1507            serde_json::json!({})
1508        }
1509        async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
1510            self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1511            Ok(ToolResult::success(ctx.tool_call_id, "acted"))
1512        }
1513    }
1514
1515    /// Hands out one prepared decision per call.
1516    struct ScriptedGate(std::sync::Mutex<Vec<GateDecision>>);
1517
1518    #[async_trait]
1519    impl ToolGate for ScriptedGate {
1520        async fn check(&self, _call: &crate::ToolCall) -> GateDecision {
1521            self.0
1522                .lock()
1523                .expect("the test does not panic while holding it")
1524                .pop()
1525                .unwrap_or(GateDecision::Allow)
1526        }
1527    }
1528
1529    fn gated(decision: GateDecision) -> (ToolRegistry, Arc<std::sync::atomic::AtomicUsize>) {
1530        let runs = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1531        let mut registry = ToolRegistry::new();
1532        registry.register(Arc::new(CountedTool(runs.clone())));
1533        registry.set_gate(Arc::new(ScriptedGate(std::sync::Mutex::new(vec![
1534            decision,
1535        ]))));
1536        (registry, runs)
1537    }
1538
1539    fn act_call() -> crate::ToolCall {
1540        crate::ToolCall {
1541            id: "c1".into(),
1542            name: "act".into(),
1543            arguments: serde_json::json!({}),
1544        }
1545    }
1546
1547    fn ask(
1548        timeout: std::time::Duration,
1549    ) -> (GateDecision, tokio::sync::oneshot::Sender<GateAnswer>) {
1550        let (answer, asked) = tokio::sync::oneshot::channel();
1551        let decision = GateDecision::Ask(GateAsk {
1552            answer: Box::pin(async move {
1553                asked
1554                    .await
1555                    .unwrap_or(GateAnswer::Deny("nobody answered".into()))
1556            }),
1557            timeout,
1558        });
1559        (decision, answer)
1560    }
1561
1562    #[tokio::test]
1563    async fn a_denied_call_does_not_run_and_returns_the_reason() {
1564        let (registry, runs) = gated(GateDecision::Deny("not allowed here".into()));
1565        let result = registry
1566            .execute(&act_call(), None, &CancellationToken::new())
1567            .await;
1568        assert!(!result.is_success());
1569        assert_eq!(result.output(), "not allowed here");
1570        assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1571    }
1572
1573    #[tokio::test]
1574    async fn an_asked_call_waits_for_the_answer_then_runs() {
1575        let (decision, answer) = ask(std::time::Duration::from_secs(30));
1576        let (registry, runs) = gated(decision);
1577        let registry = Arc::new(registry);
1578        let running = registry.clone();
1579        let call = tokio::spawn(async move {
1580            running
1581                .execute(&act_call(), None, &CancellationToken::new())
1582                .await
1583        });
1584        tokio::task::yield_now().await;
1585        assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1586        answer
1587            .send(GateAnswer::Allow)
1588            .expect("the call is still waiting");
1589        let result = call.await.expect("the call does not panic");
1590        assert!(result.is_success());
1591        assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 1);
1592    }
1593
1594    #[tokio::test]
1595    async fn an_asked_call_refused_by_the_answer_does_not_run() {
1596        let (decision, answer) = ask(std::time::Duration::from_secs(30));
1597        let (registry, runs) = gated(decision);
1598        answer
1599            .send(GateAnswer::Deny("the user said no".into()))
1600            .expect("the receiver is held by the gate's question");
1601        let result = registry
1602            .execute(&act_call(), None, &CancellationToken::new())
1603            .await;
1604        assert_eq!(result.output(), "the user said no");
1605        assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1606    }
1607
1608    #[tokio::test]
1609    async fn an_unanswered_call_is_refused_when_its_time_runs_out() {
1610        let (decision, _answer) = ask(std::time::Duration::from_millis(20));
1611        let (registry, runs) = gated(decision);
1612        let result = registry
1613            .execute(&act_call(), None, &CancellationToken::new())
1614            .await;
1615        assert!(!result.is_success());
1616        assert!(result.output().contains("nobody answered in time"));
1617        assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1618    }
1619
1620    #[tokio::test]
1621    async fn clones_keep_the_gate() {
1622        let (registry, runs) = gated(GateDecision::Deny("no".into()));
1623        let copy = registry.clone();
1624        assert!(copy.gate().is_some());
1625        let result = copy
1626            .execute(&act_call(), None, &CancellationToken::new())
1627            .await;
1628        assert!(!result.is_success());
1629        assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1630    }
1631}