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