Skip to main content

agent_sdk_tools/
tools.rs

1//! Tool definition and registry.
2//!
3//! Tools allow the LLM to perform actions in the real world. This module provides:
4//!
5//! - [`Tool`] trait - Define custom tools the LLM can call
6//! - [`ToolName`] trait - Marker trait for strongly-typed tool names
7//! - [`PrimitiveToolName`] - Tool names for SDK's built-in tools
8//! - [`DynamicToolName`] - Tool names created at runtime (MCP bridges)
9//! - [`ToolRegistry`] - Collection of available tools
10//! - [`ToolContext`] - Context passed to tool execution
11//! - [`ListenExecuteTool`] - Tools that listen for updates, then execute later
12//!
13//! # Implementing a Tool
14//!
15//! ```
16//! use agent_sdk_tools::tools::{Tool, ToolContext, DynamicToolName};
17//! use agent_sdk_foundation::types::{ToolResult, ToolTier};
18//! use serde_json::{json, Value};
19//! use std::future::Future;
20//!
21//! struct MyTool;
22//!
23//! // No #[async_trait] needed - Rust 1.75+ supports native async traits
24//! impl Tool<()> for MyTool {
25//!     type Name = DynamicToolName;
26//!
27//!     fn name(&self) -> DynamicToolName { DynamicToolName::new("my_tool") }
28//!     // `display_name` defaults to "" — override it for nicer UI.
29//!     fn description(&self) -> &'static str { "Does something useful" }
30//!     fn input_schema(&self) -> Value { json!({ "type": "object" }) }
31//!     fn tier(&self) -> ToolTier { ToolTier::Observe }
32//!
33//!     fn execute(
34//!         &self,
35//!         _ctx: &ToolContext<()>,
36//!         _input: Value,
37//!     ) -> impl Future<Output = anyhow::Result<ToolResult>> + Send {
38//!         async move { Ok(ToolResult::success("Done!")) }
39//!     }
40//! }
41//! ```
42
43use crate::authority::{EventAuthority, LocalEventAuthority};
44use crate::seed::{HostDependencies, ToolContextSeed};
45use crate::stores::EventStore;
46use agent_sdk_foundation::events::AgentEvent;
47use agent_sdk_foundation::llm;
48use agent_sdk_foundation::types::{ToolOutcome, ToolResult, ToolTier};
49use anyhow::Result;
50use async_trait::async_trait;
51use futures::Stream;
52use serde::{Deserialize, Serialize, de::DeserializeOwned};
53use serde_json::Value;
54use std::collections::HashMap;
55use std::future::Future;
56use std::marker::PhantomData;
57use std::pin::Pin;
58use std::sync::Arc;
59use time::OffsetDateTime;
60use tokio_util::sync::CancellationToken;
61
62// ============================================================================
63// Tool Name Types
64// ============================================================================
65
66/// Marker trait for tool names.
67///
68/// Tool names must be serializable (for storage/logging) and deserializable
69/// (for parsing from LLM responses). The string representation is derived
70/// from serde serialization.
71///
72/// # Example
73///
74/// ```ignore
75/// #[derive(Serialize, Deserialize)]
76/// #[serde(rename_all = "snake_case")]
77/// pub enum MyToolName {
78///     Read,
79///     Write,
80/// }
81///
82/// impl ToolName for MyToolName {}
83/// ```
84pub trait ToolName: Send + Sync + Serialize + DeserializeOwned + 'static {}
85
86/// Helper to get string representation of a tool name via serde.
87///
88/// Returns `"<unknown_tool>"` if serialization fails (should never happen
89/// with properly implemented `ToolName` types that use `#[derive(Serialize)]`).
90#[must_use]
91pub fn tool_name_to_string<N: ToolName>(name: &N) -> String {
92    serde_json::to_string(name)
93        .unwrap_or_else(|_| "\"<unknown_tool>\"".to_string())
94        .trim_matches('"')
95        .to_string()
96}
97
98/// Parse a tool name from string via serde.
99///
100/// The input is encoded as a JSON string with `serde_json::to_string` (not
101/// interpolated with `format!`) so names containing quotes or backslashes —
102/// possible for [`DynamicToolName`]s bridged from remote MCP servers — are
103/// escaped correctly and round-trip with [`tool_name_to_string`].
104///
105/// # Errors
106/// Returns error if the string doesn't match a valid tool name.
107pub fn tool_name_from_str<N: ToolName>(s: &str) -> Result<N, serde_json::Error> {
108    let json = serde_json::to_string(s)?;
109    serde_json::from_str(&json)
110}
111
112/// Tool names for SDK's built-in primitive tools.
113#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
114#[serde(rename_all = "snake_case")]
115pub enum PrimitiveToolName {
116    Read,
117    Write,
118    Edit,
119    MultiEdit,
120    Bash,
121    Glob,
122    Grep,
123    NotebookRead,
124    NotebookEdit,
125    TodoRead,
126    TodoWrite,
127    AskUser,
128    LinkFetch,
129    WebSearch,
130}
131
132impl ToolName for PrimitiveToolName {}
133
134/// Dynamic tool name for runtime-created tools (MCP bridges, subagents).
135#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
136#[serde(transparent)]
137pub struct DynamicToolName(String);
138
139impl DynamicToolName {
140    #[must_use]
141    pub fn new(name: impl Into<String>) -> Self {
142        Self(name.into())
143    }
144
145    #[must_use]
146    pub fn as_str(&self) -> &str {
147        &self.0
148    }
149}
150
151impl ToolName for DynamicToolName {}
152
153// ============================================================================
154// Progress Stage Types (for AsyncTool)
155// ============================================================================
156
157/// Marker trait for tool progress stages (type-safe, like [`ToolName`]).
158///
159/// Progress stages are used by async tools to indicate the current phase
160/// of a long-running operation. They must be serializable for event streaming.
161///
162/// # Example
163///
164/// ```ignore
165/// #[derive(Clone, Debug, Serialize, Deserialize)]
166/// #[serde(rename_all = "snake_case")]
167/// pub enum PixTransferStage {
168///     Initiated,
169///     Processing,
170///     SentToBank,
171/// }
172///
173/// impl ProgressStage for PixTransferStage {}
174/// ```
175pub trait ProgressStage: Clone + Send + Sync + Serialize + DeserializeOwned + 'static {}
176
177/// Helper to get string representation of a progress stage via serde.
178///
179/// Returns `"<unknown_stage>"` if serialization fails (should never happen with
180/// properly implemented `ProgressStage` types). This mirrors
181/// [`tool_name_to_string`]'s non-panicking fallback so a failing `Serialize`
182/// impl cannot panic the turn loop on the async-tool progress hot path.
183#[must_use]
184pub fn stage_to_string<S: ProgressStage>(stage: &S) -> String {
185    serde_json::to_string(stage)
186        .unwrap_or_else(|_| "\"<unknown_stage>\"".to_string())
187        .trim_matches('"')
188        .to_string()
189}
190
191/// Status update from an async tool operation.
192#[derive(Clone, Debug, Serialize)]
193pub enum ToolStatus<S: ProgressStage> {
194    /// Operation is making progress
195    Progress {
196        stage: S,
197        message: String,
198        data: Option<serde_json::Value>,
199    },
200
201    /// Operation completed successfully
202    Completed(ToolResult),
203
204    /// Operation failed
205    Failed(ToolResult),
206}
207
208/// Type-erased status for the agent loop.
209#[derive(Clone, Debug, Serialize, Deserialize)]
210pub enum ErasedToolStatus {
211    /// Operation is making progress
212    Progress {
213        stage: String,
214        message: String,
215        data: Option<serde_json::Value>,
216    },
217    /// Operation completed successfully
218    Completed(ToolResult),
219    /// Operation failed
220    Failed(ToolResult),
221}
222
223/// Update emitted from a `listen()` stream.
224///
225/// This models workflows where a runtime prepares an operation over time, and
226/// execution happens later using an operation identifier and revision.
227#[derive(Clone, Debug, Serialize, Deserialize)]
228pub enum ListenToolUpdate {
229    /// Preparation is still running and should keep listening.
230    Listening {
231        /// Opaque operation identifier used for later execute/cancel calls.
232        operation_id: String,
233        /// Monotonic revision number for optimistic concurrency.
234        revision: u64,
235        /// Human-readable status message.
236        message: String,
237        /// Optional current snapshot for UI rendering.
238        snapshot: Option<serde_json::Value>,
239        /// Optional expiration timestamp (RFC3339).
240        #[serde(with = "time::serde::rfc3339::option")]
241        expires_at: Option<OffsetDateTime>,
242    },
243
244    /// Preparation is complete and execution can be confirmed.
245    Ready {
246        /// Opaque operation identifier used for later execute/cancel calls.
247        operation_id: String,
248        /// Monotonic revision number for optimistic concurrency.
249        revision: u64,
250        /// Human-readable status message.
251        message: String,
252        /// Snapshot shown in confirmation UI.
253        snapshot: serde_json::Value,
254        /// Optional expiration timestamp (RFC3339).
255        #[serde(with = "time::serde::rfc3339::option")]
256        expires_at: Option<OffsetDateTime>,
257    },
258
259    /// Operation is no longer valid.
260    Invalidated {
261        /// Opaque operation identifier.
262        operation_id: String,
263        /// Human-readable reason.
264        message: String,
265        /// Whether caller may recover by starting a new listen operation.
266        recoverable: bool,
267    },
268}
269
270/// Reason for stopping a listen session.
271#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
272pub enum ListenStopReason {
273    /// User explicitly rejected confirmation.
274    UserRejected,
275    /// Agent policy/hook blocked execution before confirmation.
276    Blocked,
277    /// Consumer disconnected while listen stream was active.
278    StreamDisconnected,
279    /// Listen stream ended unexpectedly before terminal state.
280    StreamEnded,
281}
282
283impl<S: ProgressStage> From<ToolStatus<S>> for ErasedToolStatus {
284    fn from(status: ToolStatus<S>) -> Self {
285        match status {
286            ToolStatus::Progress {
287                stage,
288                message,
289                data,
290            } => Self::Progress {
291                stage: stage_to_string(&stage),
292                message,
293                data,
294            },
295            ToolStatus::Completed(r) => Self::Completed(r),
296            ToolStatus::Failed(r) => Self::Failed(r),
297        }
298    }
299}
300
301/// Context passed to tool execution
302#[derive(Clone)]
303pub struct ToolContext<Ctx> {
304    /// Application-specific context (e.g., `user_id`, db connection)
305    pub app: Ctx,
306    /// Tool-specific metadata
307    pub metadata: HashMap<String, Value>,
308    /// Optional event store for tools to emit turn-scoped events.
309    event_store: Option<Arc<dyn EventStore>>,
310    /// Thread associated with the bound event store.
311    event_thread_id: Option<agent_sdk_foundation::types::ThreadId>,
312    /// Turn associated with the bound event store.
313    event_turn: Option<usize>,
314    /// Optional event authority for wrapping events in envelopes
315    event_authority: Option<Arc<dyn EventAuthority>>,
316    /// Optional cancellation token for propagating cancellation to subtasks
317    cancel_token: Option<CancellationToken>,
318    /// Optional semaphore for limiting concurrent subagent threads.
319    subagent_semaphore: Option<Arc<tokio::sync::Semaphore>>,
320    /// Optional per-tool execution timeout enforced at the SDK boundary.
321    ///
322    /// When set, the agent loop races each tool's `execute()` future
323    /// against this duration. A tool that does not finish within the
324    /// budget is stopped at the boundary and reported with a synthetic
325    /// timeout [`ToolResult`] so the `tool_use` / `tool_result` pair stays
326    /// balanced. Tools that hold OS resources (subprocesses, sockets) must
327    /// observe the [cooperative-cancel contract](Tool#cooperative-cancellation)
328    /// so the timeout actually reclaims them.
329    tool_timeout: Option<std::time::Duration>,
330    /// Optional per-thread spill store enforcing the shared inline output
331    /// budget on tool results (see [`crate::artifacts`]).
332    ///
333    /// When set, the agent loop spills over-budget tool output to this
334    /// store and splices the `[raw output: artifact://<id>]` recovery
335    /// footer into the inline result; the `read` tool resolves
336    /// `artifact://<id>` URIs against the same store.
337    artifact_store: Option<Arc<crate::artifacts::ArtifactStore>>,
338}
339
340impl<Ctx> ToolContext<Ctx> {
341    #[must_use]
342    pub fn new(app: Ctx) -> Self {
343        Self {
344            app,
345            metadata: HashMap::new(),
346            event_store: None,
347            event_thread_id: None,
348            event_turn: None,
349            event_authority: None,
350            cancel_token: None,
351            subagent_semaphore: None,
352            tool_timeout: None,
353            artifact_store: None,
354        }
355    }
356
357    /// Reconstruct a `ToolContext` from a durable seed and host-provided
358    /// runtime dependencies.
359    ///
360    /// This is the authoritative reconstruction path.  Workers should use
361    /// this (or a host's [`crate::seed::ExecutionContextFactory`]) instead
362    /// of chaining builder methods, so that the context shape is
363    /// deterministic and auditable.
364    ///
365    /// The event authority is constructed internally from
366    /// [`ToolContextSeed::sequence_offset`] to guarantee monotonic
367    /// sequencing — callers cannot accidentally supply a misaligned
368    /// authority.
369    #[must_use]
370    pub fn from_seed(seed: &ToolContextSeed, app: Ctx, deps: HostDependencies) -> Self {
371        let authority: Arc<dyn EventAuthority> =
372            Arc::new(LocalEventAuthority::with_offset(seed.sequence_offset));
373        Self {
374            app,
375            metadata: seed.metadata.clone(),
376            event_store: Some(deps.event_store),
377            event_thread_id: Some(seed.thread_id.clone()),
378            event_turn: Some(seed.turn),
379            event_authority: Some(authority),
380            cancel_token: Some(deps.cancel_token),
381            subagent_semaphore: deps.subagent_semaphore,
382            tool_timeout: None,
383            artifact_store: deps.artifact_store,
384        }
385    }
386
387    #[must_use]
388    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
389        self.metadata.insert(key.into(), value);
390        self
391    }
392
393    /// Bind the tool context to the event store for a specific thread/turn.
394    #[must_use]
395    pub fn with_event_store(
396        mut self,
397        store: Arc<dyn EventStore>,
398        thread_id: agent_sdk_foundation::types::ThreadId,
399        turn: usize,
400        authority: Arc<dyn EventAuthority>,
401    ) -> Self {
402        self.event_store = Some(store);
403        self.event_thread_id = Some(thread_id);
404        self.event_turn = Some(turn);
405        self.event_authority = Some(authority);
406        self
407    }
408
409    /// Emit an event through the configured event store (if set).
410    ///
411    /// The event is wrapped in an [`agent_sdk_foundation::AgentEventEnvelope`] with a unique ID,
412    /// sequence number, and timestamp before publishing.
413    ///
414    /// # Errors
415    /// Returns an error if the configured event store cannot persist the event.
416    pub async fn emit_event(&self, event: AgentEvent) -> Result<()>
417    where
418        Ctx: Sync,
419    {
420        let Some((store, authority, thread_id, turn)) = self
421            .event_store
422            .as_ref()
423            .zip(self.event_authority.as_ref())
424            .zip(self.event_thread_id.as_ref())
425            .zip(self.event_turn)
426            .map(|(((store, authority), thread_id), turn)| (store, authority, thread_id, turn))
427        else {
428            // Surface the misconfiguration instead of silently dropping the
429            // event: a tool written for the durable host but run under a
430            // hand-built `ToolContext::new()` would otherwise lose every
431            // emitted event with no trace, undermining the audit trail.
432            let kind = serde_json::to_value(&event)
433                .ok()
434                .and_then(|v| {
435                    v.get("type")
436                        .and_then(|t| t.as_str().map(ToOwned::to_owned))
437                })
438                .unwrap_or_else(|| "unknown".to_string());
439            log::warn!(
440                "ToolContext::emit_event called on an unbound context; dropping {kind} event \
441                 (no event store/authority/thread/turn bound)"
442            );
443            return Ok(());
444        };
445        let envelope = authority.wrap(event);
446        store.append(thread_id, turn, envelope).await
447    }
448
449    /// Get a clone of the event authority (if set).
450    ///
451    /// This is useful for tools that spawn subprocesses (like subagents)
452    /// and need to wrap events with the same sequencing authority as the
453    /// parent's turn log.
454    #[must_use]
455    pub fn event_authority(&self) -> Option<Arc<dyn EventAuthority>> {
456        self.event_authority.clone()
457    }
458
459    /// Set the cancellation token for propagating cancellation to subtasks.
460    #[must_use]
461    pub fn with_cancel_token(mut self, token: CancellationToken) -> Self {
462        self.cancel_token = Some(token);
463        self
464    }
465
466    /// Get the cancellation token (if set).
467    ///
468    /// Used by tools that spawn long-running subtasks (like subagents)
469    /// to propagate cancellation from the parent.
470    #[must_use]
471    pub fn cancel_token(&self) -> Option<CancellationToken> {
472        self.cancel_token.clone()
473    }
474
475    /// Set the per-tool execution timeout enforced at the SDK boundary.
476    ///
477    /// The agent loop populates this from `AgentConfig::tool_timeout_ms`;
478    /// callers can also set it directly when constructing a context.
479    #[must_use]
480    pub const fn with_tool_timeout(mut self, timeout: std::time::Duration) -> Self {
481        self.tool_timeout = Some(timeout);
482        self
483    }
484
485    /// Get the per-tool execution timeout (if set).
486    ///
487    /// Read by the agent loop's SDK-boundary execution race; tools do not
488    /// normally need to consult this themselves.
489    #[must_use]
490    pub const fn tool_timeout(&self) -> Option<std::time::Duration> {
491        self.tool_timeout
492    }
493
494    /// Set a shared semaphore for limiting concurrent subagent threads.
495    #[must_use]
496    pub fn with_subagent_semaphore(mut self, semaphore: Arc<tokio::sync::Semaphore>) -> Self {
497        self.subagent_semaphore = Some(semaphore);
498        self
499    }
500
501    /// Get the subagent thread-limiting semaphore (if set).
502    #[must_use]
503    pub fn subagent_semaphore(&self) -> Option<Arc<tokio::sync::Semaphore>> {
504        self.subagent_semaphore.clone()
505    }
506
507    /// Attach the per-thread artifact spill store.
508    #[must_use]
509    pub fn with_artifact_store(mut self, store: Arc<crate::artifacts::ArtifactStore>) -> Self {
510        self.artifact_store = Some(store);
511        self
512    }
513
514    /// Get the per-thread artifact spill store (if set).
515    ///
516    /// The agent loop consults this to enforce the shared inline output
517    /// budget; the `read` tool consults it to resolve `artifact://` URIs.
518    #[must_use]
519    pub const fn artifact_store(&self) -> Option<&Arc<crate::artifacts::ArtifactStore>> {
520        self.artifact_store.as_ref()
521    }
522}
523
524// ============================================================================
525// Tool Trait
526// ============================================================================
527
528/// Definition of a tool that can be called by the agent.
529///
530/// Tools have a strongly-typed `Name` associated type that determines
531/// how the tool name is serialized for LLM communication.
532///
533/// # Native Async Support
534///
535/// This trait uses Rust's native async functions in traits (stabilized in Rust 1.75).
536/// You do NOT need the `async_trait` crate to implement this trait.
537///
538/// # Cooperative cancellation
539///
540/// The agent loop races every tool's `execute()` future against the run's
541/// [`ToolContext::cancel_token`] and, when configured, against
542/// [`ToolContext::tool_timeout`]. If either fires the SDK drops the
543/// in-flight `execute()` future and synthesises a balanced `tool_result`
544/// (`"Cancelled by user"` or a timeout message). Dropping a future runs
545/// its destructors but cannot, on its own, reclaim OS resources a tool
546/// has handed to the kernel.
547///
548/// **Subprocess contract:** a tool that spawns a child process MUST make
549/// the process die when its `execute()` future is dropped. The two
550/// supported ways to satisfy this are:
551///
552/// * Build the command with `tokio::process::Command::kill_on_drop(true)`,
553///   so the child is killed when the `Child` handle is dropped together
554///   with the cancelled future (this is what the SDK's MCP stdio transport
555///   does), or
556/// * Observe [`ToolContext::cancel_token`] directly and `kill()` the child
557///   when it fires.
558///
559/// A tool that holds a subprocess open without either of these will leak
560/// the process when cancelled or timed out — the synthesised `tool_result`
561/// keeps the conversation balanced, but the orphaned OS process is the
562/// tool author's bug, not the SDK's.
563pub trait Tool<Ctx>: Send + Sync {
564    /// The type of name for this tool.
565    type Name: ToolName;
566
567    /// Returns the tool's strongly-typed name.
568    fn name(&self) -> Self::Name;
569
570    /// Human-readable display name for UI (e.g., "Read File" vs "read").
571    ///
572    /// Defaults to the empty string. Override for better UX.
573    fn display_name(&self) -> &'static str {
574        ""
575    }
576
577    /// Human-readable description of what the tool does.
578    fn description(&self) -> &'static str;
579
580    /// JSON schema for the tool's input parameters.
581    fn input_schema(&self) -> Value;
582
583    /// Permission tier for this tool.
584    ///
585    /// Defaults to [`ToolTier::Confirm`] (fail-closed): a tool author who
586    /// forgets to declare a tier gets confirmation gating, not silent
587    /// auto-execution. Read-only tools should explicitly opt in to
588    /// [`ToolTier::Observe`].
589    fn tier(&self) -> ToolTier {
590        ToolTier::Confirm
591    }
592
593    /// Execute the tool with the given input.
594    ///
595    /// # Errors
596    /// Returns an error if tool execution fails.
597    fn execute(
598        &self,
599        ctx: &ToolContext<Ctx>,
600        input: Value,
601    ) -> impl Future<Output = Result<ToolResult>> + Send;
602}
603
604// ============================================================================
605// TypedTool Trait (typed input + runtime validation / self-correction)
606// ============================================================================
607
608/// A tool whose model-emitted arguments are validated against a typed,
609/// deserializable [`Input`](TypedTool::Input) **before** [`execute`](TypedTool::execute)
610/// runs.
611///
612/// Today a raw [`serde_json::Value`] is handed straight to [`Tool::execute`],
613/// so a malformed tool call reaches tool code unvalidated. `TypedTool` closes
614/// that gap: you declare a `Serialize` / `Deserialize` argument struct as
615/// [`Input`](TypedTool::Input), and the runtime deserializes the model's args
616/// into it at the dispatch boundary. On a deserialization/validation failure
617/// the runtime synthesises a structured error [`ToolResult`] (carrying the
618/// serde error message) so the model can self-correct on its next turn —
619/// `execute` is **never** called with invalid arguments.
620///
621/// # Relationship to [`Tool`]
622///
623/// `TypedTool` is the typed, opt-in *sugar* layer; [`Tool`] remains the
624/// untyped baseline. A [`TypedTool`] becomes a full [`Tool`] through
625/// [`TypedToolAdapter`] (mirroring how [`SimpleTool`] becomes a [`Tool`] via
626/// [`SimpleToolAdapter`]). Register one with
627/// [`ToolRegistry::register_typed`], which wraps it in the adapter for you;
628/// the adapter performs the deserialize-then-dispatch (or
629/// deserialize-then-synthesise-error) described above.
630///
631/// # Back-compat / migration
632///
633/// Existing [`Tool`] impls (and [`SimpleTool`] / [`DynamicToolName`] tools)
634/// keep compiling and running unchanged — they stay on the `Value`-in
635/// baseline, which is the identity passthrough (a `Value` always
636/// "deserializes" into a `Value`). Migrate a tool to typed args by moving its
637/// `impl Tool<Ctx>` to `impl TypedTool<Ctx>`, setting `type Input = MyArgs`,
638/// and changing `execute`'s signature from `input: Value` to `input: MyArgs`.
639/// The hand-written [`input_schema`](TypedTool::input_schema) JSON stays
640/// user-declared; this trait does **not** auto-derive a schema from `Input`.
641///
642/// # Example
643///
644/// ```
645/// use agent_sdk_tools::tools::{TypedTool, ToolContext};
646/// use agent_sdk_foundation::types::ToolResult;
647/// use serde::{Deserialize, Serialize};
648/// use serde_json::{json, Value};
649/// use std::future::Future;
650///
651/// #[derive(Debug, Serialize, Deserialize)]
652/// struct WeatherArgs {
653///     city: String,
654/// }
655///
656/// struct WeatherTool;
657///
658/// impl TypedTool<()> for WeatherTool {
659///     type Input = WeatherArgs;
660///
661///     fn name(&self) -> &'static str { "get_weather" }
662///     fn description(&self) -> &'static str { "Get current weather for a city" }
663///     fn input_schema(&self) -> Value {
664///         json!({
665///             "type": "object",
666///             "properties": { "city": { "type": "string" } },
667///             "required": ["city"]
668///         })
669///     }
670///
671///     fn execute(
672///         &self,
673///         _ctx: &ToolContext<()>,
674///         input: WeatherArgs,
675///     ) -> impl Future<Output = anyhow::Result<ToolResult>> + Send {
676///         async move { Ok(ToolResult::success(format!("Weather in {}: Sunny", input.city))) }
677///     }
678/// }
679/// ```
680///
681/// Like [`SimpleTool`], a `TypedTool` has a single fixed `&'static str`
682/// [`name`](TypedTool::name) (mapping to [`DynamicToolName`] via
683/// [`TypedToolAdapter`]). Reach for a hand-written [`Tool`] with a
684/// strongly-typed [`ToolName`] when the name must be computed at runtime or
685/// constrained to an enum.
686pub trait TypedTool<Ctx>: Send + Sync {
687    /// The typed input the model's arguments are deserialized into before
688    /// [`execute`](TypedTool::execute) runs.
689    ///
690    /// Must be [`DeserializeOwned`] (to parse model args), [`Serialize`] (so
691    /// the typed value round-trips for logging/storage), and `Send + 'static`
692    /// (to cross the async dispatch boundary).
693    type Input: DeserializeOwned + Serialize + Send + 'static;
694
695    /// The tool's name as sent to (and parsed from) the model.
696    fn name(&self) -> &'static str;
697
698    /// Human-readable display name for UI. Defaults to an empty string.
699    fn display_name(&self) -> &'static str {
700        ""
701    }
702
703    /// Human-readable description of what the tool does.
704    fn description(&self) -> &'static str;
705
706    /// User-declared JSON schema for the tool's input parameters.
707    ///
708    /// This stays hand-written JSON — it is **not** auto-derived from
709    /// [`Input`](TypedTool::Input). Keeping the schema explicit lets the
710    /// declared provider-facing contract diverge from the Rust type when that
711    /// is useful (descriptions, examples, provider-specific keywords).
712    fn input_schema(&self) -> Value;
713
714    /// Permission tier for this tool. Defaults to [`ToolTier::Confirm`]
715    /// (fail-closed); read-only tools should opt in to [`ToolTier::Observe`].
716    fn tier(&self) -> ToolTier {
717        ToolTier::Confirm
718    }
719
720    /// Execute the tool with the already-validated, typed input.
721    ///
722    /// The runtime guarantees `input` deserialized cleanly from the model's
723    /// arguments; a malformed call is turned into a structured error
724    /// [`ToolResult`] before this method is reached, so implementations never
725    /// see invalid arguments.
726    ///
727    /// # Errors
728    /// Returns an error if tool execution fails.
729    fn execute(
730        &self,
731        ctx: &ToolContext<Ctx>,
732        input: Self::Input,
733    ) -> impl Future<Output = Result<ToolResult>> + Send;
734}
735
736/// Synthesise the structured validation-error [`ToolResult`] returned to the
737/// model when its arguments fail to deserialize into a [`TypedTool::Input`].
738///
739/// Factored out (and `pub`) so the exact self-correction wording is
740/// consistent with [`TypedToolAdapter`] and is directly unit-testable. The
741/// error is an *error* [`ToolResult`] (not a thrown `anyhow::Error`): it flows
742/// through the normal balanced `tool_use` / `tool_result` path so history
743/// stays balanced and the model gets a concrete, machine-actionable hint on
744/// its next turn.
745#[must_use]
746pub fn invalid_tool_input_result(tool_name: &str, error: &serde_json::Error) -> ToolResult {
747    ToolResult::error(format!(
748        "Invalid arguments for tool `{tool_name}`: {error}. \
749         The arguments did not match the tool's input schema — \
750         re-read the schema and call the tool again with corrected arguments."
751    ))
752}
753
754/// Deserialize raw model args into a typed `Input`, or synthesise the
755/// structured validation-error result.
756///
757/// Returns `Ok(typed)` for the happy path and `Err(result)` carrying the
758/// balanced error [`ToolResult`] for the self-correction path.
759/// [`TypedToolAdapter`] uses this to ensure [`TypedTool::execute`] is never
760/// reached with invalid arguments.
761///
762/// # Errors
763/// Returns the synthesised error [`ToolResult`] when `raw` does not
764/// deserialize into `Input`.
765pub fn validate_tool_input<Input>(tool_name: &str, raw: Value) -> Result<Input, ToolResult>
766where
767    Input: DeserializeOwned,
768{
769    serde_json::from_value(raw).map_err(|error| invalid_tool_input_result(tool_name, &error))
770}
771
772/// Adapter that turns any [`TypedTool`] into a full [`Tool`].
773///
774/// It gives the wrapped tool `Name = DynamicToolName`, deserializes the
775/// model's `Value` arguments into [`TypedTool::Input`] before dispatching, and
776/// synthesises a structured validation-error [`ToolResult`] when that fails.
777///
778/// You rarely name this type directly — register a [`TypedTool`] with
779/// [`ToolRegistry::register_typed`], which wraps it for you. The adapter
780/// pattern (rather than a blanket `impl Tool for T: TypedTool`) is required
781/// for coherence: a blanket impl would conflict with the existing
782/// [`SimpleToolAdapter`] impl, because the compiler cannot rule out a
783/// downstream `TypedTool` impl for `SimpleToolAdapter`.
784///
785/// This adapter is also where the typed `Input` is threaded through the
786/// erased-tool machinery without leaking the generic into trait objects: the
787/// registry's [`ErasedTool`] wrapper still only ever sees `Value`, while the
788/// concrete `Input` type (and the deserialize) live here, inside the adapter's
789/// concrete `T`.
790pub struct TypedToolAdapter<T> {
791    inner: T,
792}
793
794impl<T> TypedToolAdapter<T> {
795    /// Wrap a [`TypedTool`] so it can be used anywhere a [`Tool`] is expected.
796    pub const fn new(tool: T) -> Self {
797        Self { inner: tool }
798    }
799
800    /// Unwrap the inner [`TypedTool`].
801    pub fn into_inner(self) -> T {
802        self.inner
803    }
804}
805
806impl<Ctx, T> Tool<Ctx> for TypedToolAdapter<T>
807where
808    T: TypedTool<Ctx>,
809    Ctx: Send + Sync,
810{
811    type Name = DynamicToolName;
812
813    fn name(&self) -> DynamicToolName {
814        DynamicToolName::new(TypedTool::name(&self.inner))
815    }
816
817    fn display_name(&self) -> &'static str {
818        TypedTool::display_name(&self.inner)
819    }
820
821    fn description(&self) -> &'static str {
822        TypedTool::description(&self.inner)
823    }
824
825    fn input_schema(&self) -> Value {
826        TypedTool::input_schema(&self.inner)
827    }
828
829    fn tier(&self) -> ToolTier {
830        TypedTool::tier(&self.inner)
831    }
832
833    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolResult> {
834        match validate_tool_input::<<T as TypedTool<Ctx>>::Input>(
835            TypedTool::name(&self.inner),
836            input,
837        ) {
838            Ok(typed) => TypedTool::execute(&self.inner, ctx, typed).await,
839            // A validation failure is returned as an error `ToolResult`,
840            // never `?`-bailed: it must reach the model as a balanced
841            // `tool_result` for self-correction. `execute` is not called.
842            Err(result) => Ok(result),
843        }
844    }
845}
846
847// ============================================================================
848// ToolLogic Trait (execute-only companion for the derive macros)
849// ============================================================================
850
851/// The `execute`-only half of a tool, used as the target of the
852/// `#[derive(Tool)]` / `#[derive(TypedTool)]` ergonomics macros.
853///
854/// The derives generate everything *except* the behaviour — `name`,
855/// `description`, `input_schema`, `tier` come from `#[tool(...)]` attributes —
856/// and delegate execution to this trait. You implement `ToolLogic` to supply
857/// the one thing a macro cannot: the `execute` body.
858///
859/// It is deliberately a **trait** (not an inherent method): a trait-method
860/// `async fn` that performs no `await` is fine, whereas an inherent one trips
861/// `clippy::unused_async`. Writing the body here keeps trivial, fully
862/// synchronous tools lint-clean without an `#[allow]`.
863///
864/// You rarely name this trait in prose — the derive docs show it in context —
865/// but the shape is:
866///
867/// ```
868/// use agent_sdk_tools::tools::{ToolLogic, ToolContext};
869/// use agent_sdk_foundation::types::ToolResult;
870/// use serde_json::Value;
871///
872/// struct MyTool;
873///
874/// impl ToolLogic<()> for MyTool {
875///     type Input = Value; // typed tools set this to their `Input` struct
876///
877///     async fn execute(&self, _ctx: &ToolContext<()>, input: Value) -> anyhow::Result<ToolResult> {
878///         Ok(ToolResult::success(format!("got {input}")))
879///     }
880/// }
881/// ```
882pub trait ToolLogic<Ctx>: Send + Sync {
883    /// The input the tool's `execute` receives. For `#[derive(Tool)]` this is
884    /// [`serde_json::Value`]; for `#[derive(TypedTool)]` it is the typed
885    /// `Input` (validated before `execute` runs).
886    type Input;
887
888    /// The tool's behaviour. Receives the (already-validated, for typed tools)
889    /// input.
890    ///
891    /// # Errors
892    /// Returns an error if tool execution fails.
893    fn execute(
894        &self,
895        ctx: &ToolContext<Ctx>,
896        input: Self::Input,
897    ) -> impl Future<Output = Result<ToolResult>> + Send;
898}
899
900// ============================================================================
901// SimpleTool Trait
902// ============================================================================
903
904/// An ergonomic [`Tool`] whose name is a plain string.
905///
906/// Most custom tools don't need a strongly-typed [`ToolName`] enum — they have
907/// a single, fixed name. `SimpleTool` lets you write a tool by returning a
908/// `&str` from [`name`](SimpleTool::name) instead of defining a `ToolName`
909/// type and an associated [`Tool::Name`].
910///
911/// Any `SimpleTool` is automatically a [`Tool`] (via a blanket impl) with
912/// `Name = DynamicToolName`, so it can be registered and used exactly like a
913/// hand-written `Tool`.
914///
915/// # Example
916///
917/// ```
918/// use agent_sdk_tools::tools::{SimpleTool, ToolContext};
919/// use agent_sdk_foundation::types::ToolResult;
920/// use serde_json::{json, Value};
921/// use std::future::Future;
922///
923/// struct WeatherTool;
924///
925/// impl SimpleTool<()> for WeatherTool {
926///     fn name(&self) -> &'static str { "get_weather" }
927///     fn description(&self) -> &'static str { "Get current weather for a city" }
928///     fn input_schema(&self) -> Value {
929///         json!({ "type": "object", "properties": { "city": { "type": "string" } } })
930///     }
931///
932///     fn execute(
933///         &self,
934///         _ctx: &ToolContext<()>,
935///         input: Value,
936///     ) -> impl Future<Output = anyhow::Result<ToolResult>> + Send {
937///         async move {
938///             let city = input["city"].as_str().unwrap_or("Unknown");
939///             Ok(ToolResult::success(format!("Weather in {city}: Sunny")))
940///         }
941///     }
942/// }
943/// ```
944pub trait SimpleTool<Ctx>: Send + Sync {
945    /// The tool's name as sent to (and parsed from) the LLM.
946    ///
947    /// Returns `&'static str` because a simple tool has one fixed name; reach
948    /// for the full [`Tool`] trait with a [`DynamicToolName`] when the name is
949    /// computed at runtime.
950    fn name(&self) -> &'static str;
951
952    /// Human-readable display name for UI.
953    ///
954    /// Defaults to an empty string; override for a friendlier label.
955    fn display_name(&self) -> &'static str {
956        ""
957    }
958
959    /// Human-readable description of what the tool does.
960    fn description(&self) -> &'static str;
961
962    /// JSON schema for the tool's input parameters.
963    fn input_schema(&self) -> Value;
964
965    /// Permission tier for this tool. Defaults to [`ToolTier::Confirm`]
966    /// (fail-closed); read-only tools should opt in to [`ToolTier::Observe`].
967    fn tier(&self) -> ToolTier {
968        ToolTier::Confirm
969    }
970
971    /// Execute the tool with the given input.
972    ///
973    /// # Errors
974    /// Returns an error if tool execution fails.
975    fn execute(
976        &self,
977        ctx: &ToolContext<Ctx>,
978        input: Value,
979    ) -> impl Future<Output = Result<ToolResult>> + Send;
980}
981
982/// Adapter that turns any [`SimpleTool`] into a full [`Tool`] with
983/// `Name = DynamicToolName`.
984///
985/// You rarely name this type directly — register a [`SimpleTool`] with
986/// [`ToolRegistry::register_simple`], which wraps it for you. Use this adapter
987/// explicitly only when you need a `Tool` value (e.g. to pass to code that is
988/// generic over [`Tool`]).
989pub struct SimpleToolAdapter<T> {
990    inner: T,
991}
992
993impl<T> SimpleToolAdapter<T> {
994    /// Wrap a [`SimpleTool`] so it can be used anywhere a [`Tool`] is expected.
995    pub const fn new(tool: T) -> Self {
996        Self { inner: tool }
997    }
998
999    /// Unwrap the inner [`SimpleTool`].
1000    pub fn into_inner(self) -> T {
1001        self.inner
1002    }
1003}
1004
1005impl<Ctx, T> Tool<Ctx> for SimpleToolAdapter<T>
1006where
1007    T: SimpleTool<Ctx>,
1008{
1009    type Name = DynamicToolName;
1010
1011    fn name(&self) -> DynamicToolName {
1012        DynamicToolName::new(SimpleTool::name(&self.inner))
1013    }
1014
1015    fn display_name(&self) -> &'static str {
1016        SimpleTool::display_name(&self.inner)
1017    }
1018
1019    fn description(&self) -> &'static str {
1020        SimpleTool::description(&self.inner)
1021    }
1022
1023    fn input_schema(&self) -> Value {
1024        SimpleTool::input_schema(&self.inner)
1025    }
1026
1027    fn tier(&self) -> ToolTier {
1028        SimpleTool::tier(&self.inner)
1029    }
1030
1031    fn execute(
1032        &self,
1033        ctx: &ToolContext<Ctx>,
1034        input: Value,
1035    ) -> impl Future<Output = Result<ToolResult>> + Send {
1036        SimpleTool::execute(&self.inner, ctx, input)
1037    }
1038}
1039
1040// ============================================================================
1041// AsyncTool Trait
1042// ============================================================================
1043
1044/// A tool that performs long-running async operations.
1045///
1046/// `AsyncTool`s have two phases:
1047/// 1. `execute()` - Start the operation (lightweight, returns quickly)
1048/// 2. `check_status()` - Stream progress until completion
1049///
1050/// The actual work should happen externally (background task, external service)
1051/// and persist results to a durable store. The tool is just an orchestrator.
1052///
1053/// # Example
1054///
1055/// ```ignore
1056/// impl AsyncTool<MyCtx> for ExecutePixTransferTool {
1057///     type Name = PixToolName;
1058///     type Stage = PixTransferStage;
1059///
1060///     async fn execute(&self, ctx: &ToolContext<MyCtx>, input: Value) -> Result<ToolOutcome> {
1061///         let params = parse_input(&input)?;
1062///         let operation_id = ctx.app.pix_service.start_transfer(params).await?;
1063///         Ok(ToolOutcome::in_progress(
1064///             operation_id,
1065///             format!("PIX transfer of {} initiated", params.amount),
1066///         ))
1067///     }
1068///
1069///     fn check_status(&self, ctx: &ToolContext<MyCtx>, operation_id: &str)
1070///         -> impl Stream<Item = ToolStatus<PixTransferStage>> + Send
1071///     {
1072///         async_stream::stream! {
1073///             loop {
1074///                 let status = ctx.app.pix_service.get_status(operation_id).await;
1075///                 match status {
1076///                     PixStatus::Success { id } => {
1077///                         yield ToolStatus::Completed(ToolResult::success(id));
1078///                         break;
1079///                     }
1080///                     _ => yield ToolStatus::Progress { ... };
1081///                 }
1082///                 tokio::time::sleep(Duration::from_millis(500)).await;
1083///             }
1084///         }
1085///     }
1086/// }
1087/// ```
1088pub trait AsyncTool<Ctx>: Send + Sync {
1089    /// The type of name for this tool.
1090    type Name: ToolName;
1091    /// The type of progress stages for this tool.
1092    type Stage: ProgressStage;
1093
1094    /// Returns the tool's strongly-typed name.
1095    fn name(&self) -> Self::Name;
1096
1097    /// Human-readable display name for UI. Defaults to the empty string.
1098    fn display_name(&self) -> &'static str {
1099        ""
1100    }
1101
1102    /// Human-readable description of what the tool does.
1103    fn description(&self) -> &'static str;
1104
1105    /// JSON schema for the tool's input parameters.
1106    fn input_schema(&self) -> Value;
1107
1108    /// Permission tier for this tool. Defaults to [`ToolTier::Confirm`]
1109    /// (fail-closed); read-only tools should opt in to [`ToolTier::Observe`].
1110    fn tier(&self) -> ToolTier {
1111        ToolTier::Confirm
1112    }
1113
1114    /// Execute the tool. Returns immediately with one of:
1115    /// - Success/Failed: Operation completed synchronously
1116    /// - `InProgress`: Operation started, use `check_status()` to stream updates
1117    ///
1118    /// # Errors
1119    /// Returns an error if tool execution fails.
1120    fn execute(
1121        &self,
1122        ctx: &ToolContext<Ctx>,
1123        input: Value,
1124    ) -> impl Future<Output = Result<ToolOutcome>> + Send;
1125
1126    /// Stream status updates for an in-progress operation.
1127    /// Must yield until Completed or Failed.
1128    fn check_status(
1129        &self,
1130        ctx: &ToolContext<Ctx>,
1131        operation_id: &str,
1132    ) -> impl Stream<Item = ToolStatus<Self::Stage>> + Send;
1133}
1134
1135// ============================================================================
1136// ListenExecuteTool Trait
1137// ============================================================================
1138
1139/// A tool whose runtime has two phases:
1140/// 1. `listen()` - starts preparation and streams updates
1141/// 2. `execute()` - performs final execution after confirmation
1142///
1143/// This abstraction is useful when runtime state can expire or evolve before
1144/// execution (quotes, challenge windows, leases, approvals).
1145///
1146/// Ordering note: the agent loop consumes `listen()` updates before
1147/// `AgentHooks::pre_tool_use()` runs. Hooks can therefore block `execute()`, but
1148/// any side effects done during `listen()` have already happened.
1149pub trait ListenExecuteTool<Ctx>: Send + Sync {
1150    /// The type of name for this tool.
1151    type Name: ToolName;
1152
1153    /// Returns the tool's strongly-typed name.
1154    fn name(&self) -> Self::Name;
1155
1156    /// Human-readable display name for UI. Defaults to the empty string.
1157    fn display_name(&self) -> &'static str {
1158        ""
1159    }
1160
1161    /// Human-readable description of what the tool does.
1162    fn description(&self) -> &'static str;
1163
1164    /// JSON schema for the tool's input parameters.
1165    fn input_schema(&self) -> Value;
1166
1167    /// Permission tier for this tool.
1168    fn tier(&self) -> ToolTier {
1169        ToolTier::Confirm
1170    }
1171
1172    /// Start and stream runtime preparation updates.
1173    fn listen(
1174        &self,
1175        ctx: &ToolContext<Ctx>,
1176        input: Value,
1177    ) -> impl Stream<Item = ListenToolUpdate> + Send;
1178
1179    /// Execute using operation ID and optimistic concurrency revision.
1180    ///
1181    /// # Errors
1182    /// Returns an error if execution fails or revision is stale.
1183    fn execute(
1184        &self,
1185        ctx: &ToolContext<Ctx>,
1186        operation_id: &str,
1187        expected_revision: u64,
1188    ) -> impl Future<Output = Result<ToolResult>> + Send;
1189
1190    /// Stop a listen operation (best effort).
1191    ///
1192    /// # Errors
1193    /// Returns an error if cancellation fails.
1194    fn cancel(
1195        &self,
1196        _ctx: &ToolContext<Ctx>,
1197        _operation_id: &str,
1198        _reason: ListenStopReason,
1199    ) -> impl Future<Output = Result<()>> + Send {
1200        async { Ok(()) }
1201    }
1202}
1203
1204// ============================================================================
1205// Type-Erased Tool (for Registry)
1206// ============================================================================
1207
1208/// Type-erased tool trait for registry storage.
1209///
1210/// This allows tools with different `Name` associated types to be stored
1211/// in the same registry by erasing the type information.
1212///
1213/// # Example
1214///
1215/// ```ignore
1216/// for tool in registry.all() {
1217///     println!("Tool: {} - {}", tool.name_str(), tool.description());
1218/// }
1219/// ```
1220#[async_trait]
1221pub trait ErasedTool<Ctx>: Send + Sync {
1222    /// Get the tool name as a string.
1223    fn name_str(&self) -> &str;
1224    /// Get a human-friendly display name for the tool.
1225    fn display_name(&self) -> &'static str;
1226    /// Get the tool description.
1227    fn description(&self) -> &'static str;
1228    /// Get the JSON schema for tool inputs.
1229    fn input_schema(&self) -> Value;
1230    /// Get the tool's permission tier.
1231    fn tier(&self) -> ToolTier;
1232    /// Execute the tool with the given input.
1233    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolResult>;
1234}
1235
1236/// Wrapper that erases the Name associated type from a Tool.
1237struct ToolWrapper<T, Ctx>
1238where
1239    T: Tool<Ctx>,
1240{
1241    inner: T,
1242    name_cache: String,
1243    _marker: PhantomData<Ctx>,
1244}
1245
1246impl<T, Ctx> ToolWrapper<T, Ctx>
1247where
1248    T: Tool<Ctx>,
1249{
1250    fn new(tool: T) -> Self {
1251        let name_cache = tool_name_to_string(&tool.name());
1252        Self {
1253            inner: tool,
1254            name_cache,
1255            _marker: PhantomData,
1256        }
1257    }
1258}
1259
1260#[async_trait]
1261impl<T, Ctx> ErasedTool<Ctx> for ToolWrapper<T, Ctx>
1262where
1263    T: Tool<Ctx> + 'static,
1264    Ctx: Send + Sync + 'static,
1265{
1266    fn name_str(&self) -> &str {
1267        &self.name_cache
1268    }
1269
1270    fn display_name(&self) -> &'static str {
1271        self.inner.display_name()
1272    }
1273
1274    fn description(&self) -> &'static str {
1275        self.inner.description()
1276    }
1277
1278    fn input_schema(&self) -> Value {
1279        self.inner.input_schema()
1280    }
1281
1282    fn tier(&self) -> ToolTier {
1283        self.inner.tier()
1284    }
1285
1286    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolResult> {
1287        self.inner.execute(ctx, input).await
1288    }
1289}
1290
1291// ============================================================================
1292// Type-Erased AsyncTool (for Registry)
1293// ============================================================================
1294
1295/// Type-erased async tool trait for registry storage.
1296///
1297/// This allows async tools with different `Name` and `Stage` associated types
1298/// to be stored in the same registry by erasing the type information.
1299#[async_trait]
1300pub trait ErasedAsyncTool<Ctx>: Send + Sync {
1301    /// Get the tool name as a string.
1302    fn name_str(&self) -> &str;
1303    /// Get a human-friendly display name for the tool.
1304    fn display_name(&self) -> &'static str;
1305    /// Get the tool description.
1306    fn description(&self) -> &'static str;
1307    /// Get the JSON schema for tool inputs.
1308    fn input_schema(&self) -> Value;
1309    /// Get the tool's permission tier.
1310    fn tier(&self) -> ToolTier;
1311    /// Execute the tool with the given input.
1312    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolOutcome>;
1313    /// Stream status updates for an in-progress operation (type-erased).
1314    fn check_status_stream<'a>(
1315        &'a self,
1316        ctx: &'a ToolContext<Ctx>,
1317        operation_id: &'a str,
1318    ) -> Pin<Box<dyn Stream<Item = ErasedToolStatus> + Send + 'a>>;
1319}
1320
1321/// Wrapper that erases the Name and Stage associated types from an [`AsyncTool`].
1322struct AsyncToolWrapper<T, Ctx>
1323where
1324    T: AsyncTool<Ctx>,
1325{
1326    inner: T,
1327    name_cache: String,
1328    _marker: PhantomData<Ctx>,
1329}
1330
1331impl<T, Ctx> AsyncToolWrapper<T, Ctx>
1332where
1333    T: AsyncTool<Ctx>,
1334{
1335    fn new(tool: T) -> Self {
1336        let name_cache = tool_name_to_string(&tool.name());
1337        Self {
1338            inner: tool,
1339            name_cache,
1340            _marker: PhantomData,
1341        }
1342    }
1343}
1344
1345#[async_trait]
1346impl<T, Ctx> ErasedAsyncTool<Ctx> for AsyncToolWrapper<T, Ctx>
1347where
1348    T: AsyncTool<Ctx> + 'static,
1349    Ctx: Send + Sync + 'static,
1350{
1351    fn name_str(&self) -> &str {
1352        &self.name_cache
1353    }
1354
1355    fn display_name(&self) -> &'static str {
1356        self.inner.display_name()
1357    }
1358
1359    fn description(&self) -> &'static str {
1360        self.inner.description()
1361    }
1362
1363    fn input_schema(&self) -> Value {
1364        self.inner.input_schema()
1365    }
1366
1367    fn tier(&self) -> ToolTier {
1368        self.inner.tier()
1369    }
1370
1371    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolOutcome> {
1372        self.inner.execute(ctx, input).await
1373    }
1374
1375    fn check_status_stream<'a>(
1376        &'a self,
1377        ctx: &'a ToolContext<Ctx>,
1378        operation_id: &'a str,
1379    ) -> Pin<Box<dyn Stream<Item = ErasedToolStatus> + Send + 'a>> {
1380        use futures::StreamExt;
1381        let stream = self.inner.check_status(ctx, operation_id);
1382        Box::pin(stream.map(ErasedToolStatus::from))
1383    }
1384}
1385
1386// ============================================================================
1387// Type-Erased ListenExecuteTool (for Registry)
1388// ============================================================================
1389
1390/// Type-erased listen/execute tool trait for registry storage.
1391#[async_trait]
1392pub trait ErasedListenTool<Ctx>: Send + Sync {
1393    /// Get the tool name as a string.
1394    fn name_str(&self) -> &str;
1395    /// Get a human-friendly display name for the tool.
1396    fn display_name(&self) -> &'static str;
1397    /// Get the tool description.
1398    fn description(&self) -> &'static str;
1399    /// Get the JSON schema for tool inputs.
1400    fn input_schema(&self) -> Value;
1401    /// Get the tool's permission tier.
1402    fn tier(&self) -> ToolTier;
1403    /// Start listen stream.
1404    fn listen_stream<'a>(
1405        &'a self,
1406        ctx: &'a ToolContext<Ctx>,
1407        input: Value,
1408    ) -> Pin<Box<dyn Stream<Item = ListenToolUpdate> + Send + 'a>>;
1409    /// Execute using a prepared operation.
1410    async fn execute(
1411        &self,
1412        ctx: &ToolContext<Ctx>,
1413        operation_id: &str,
1414        expected_revision: u64,
1415    ) -> Result<ToolResult>;
1416    /// Cancel operation.
1417    async fn cancel(
1418        &self,
1419        ctx: &ToolContext<Ctx>,
1420        operation_id: &str,
1421        reason: ListenStopReason,
1422    ) -> Result<()>;
1423}
1424
1425/// Wrapper that erases the Name associated type from a [`ListenExecuteTool`].
1426struct ListenToolWrapper<T, Ctx>
1427where
1428    T: ListenExecuteTool<Ctx>,
1429{
1430    inner: T,
1431    name_cache: String,
1432    _marker: PhantomData<Ctx>,
1433}
1434
1435impl<T, Ctx> ListenToolWrapper<T, Ctx>
1436where
1437    T: ListenExecuteTool<Ctx>,
1438{
1439    fn new(tool: T) -> Self {
1440        let name_cache = tool_name_to_string(&tool.name());
1441        Self {
1442            inner: tool,
1443            name_cache,
1444            _marker: PhantomData,
1445        }
1446    }
1447}
1448
1449#[async_trait]
1450impl<T, Ctx> ErasedListenTool<Ctx> for ListenToolWrapper<T, Ctx>
1451where
1452    T: ListenExecuteTool<Ctx> + 'static,
1453    Ctx: Send + Sync + 'static,
1454{
1455    fn name_str(&self) -> &str {
1456        &self.name_cache
1457    }
1458
1459    fn display_name(&self) -> &'static str {
1460        self.inner.display_name()
1461    }
1462
1463    fn description(&self) -> &'static str {
1464        self.inner.description()
1465    }
1466
1467    fn input_schema(&self) -> Value {
1468        self.inner.input_schema()
1469    }
1470
1471    fn tier(&self) -> ToolTier {
1472        self.inner.tier()
1473    }
1474
1475    fn listen_stream<'a>(
1476        &'a self,
1477        ctx: &'a ToolContext<Ctx>,
1478        input: Value,
1479    ) -> Pin<Box<dyn Stream<Item = ListenToolUpdate> + Send + 'a>> {
1480        let stream = self.inner.listen(ctx, input);
1481        Box::pin(stream)
1482    }
1483
1484    async fn execute(
1485        &self,
1486        ctx: &ToolContext<Ctx>,
1487        operation_id: &str,
1488        expected_revision: u64,
1489    ) -> Result<ToolResult> {
1490        self.inner
1491            .execute(ctx, operation_id, expected_revision)
1492            .await
1493    }
1494
1495    async fn cancel(
1496        &self,
1497        ctx: &ToolContext<Ctx>,
1498        operation_id: &str,
1499        reason: ListenStopReason,
1500    ) -> Result<()> {
1501        self.inner.cancel(ctx, operation_id, reason).await
1502    }
1503}
1504
1505/// Registry of available tools.
1506///
1507/// Tools are stored with their names erased to allow different `Name` types
1508/// in the same registry. The registry uses string-based lookup for LLM
1509/// compatibility.
1510///
1511/// Supports both synchronous [`Tool`]s and asynchronous [`AsyncTool`]s.
1512pub struct ToolRegistry<Ctx> {
1513    tools: HashMap<String, Arc<dyn ErasedTool<Ctx>>>,
1514    async_tools: HashMap<String, Arc<dyn ErasedAsyncTool<Ctx>>>,
1515    listen_tools: HashMap<String, Arc<dyn ErasedListenTool<Ctx>>>,
1516}
1517
1518impl<Ctx> Clone for ToolRegistry<Ctx> {
1519    fn clone(&self) -> Self {
1520        Self {
1521            tools: self.tools.clone(),
1522            async_tools: self.async_tools.clone(),
1523            listen_tools: self.listen_tools.clone(),
1524        }
1525    }
1526}
1527
1528impl<Ctx: Send + Sync + 'static> Default for ToolRegistry<Ctx> {
1529    fn default() -> Self {
1530        Self::new()
1531    }
1532}
1533
1534impl<Ctx: Send + Sync + 'static> ToolRegistry<Ctx> {
1535    #[must_use]
1536    pub fn new() -> Self {
1537        Self {
1538            tools: HashMap::new(),
1539            async_tools: HashMap::new(),
1540            listen_tools: HashMap::new(),
1541        }
1542    }
1543
1544    /// Evict any existing registration for `name` across **all three** maps so
1545    /// a name lives in exactly one map, then warn about the replacement.
1546    ///
1547    /// Without this, re-registering a name silently replaced the tool, and the
1548    /// same name registered as both (say) a sync and a listen tool coexisted —
1549    /// [`len`](ToolRegistry::len) double-counted it and
1550    /// [`to_llm_tools`](ToolRegistry::to_llm_tools) emitted two definitions with
1551    /// identical names (which providers reject). A remote MCP server could also
1552    /// silently shadow a vetted built-in (`read`, `bash`). We keep the
1553    /// non-breaking last-registration-wins behavior but make it loud; callers
1554    /// that need fail-closed semantics should use the `try_register*` variants.
1555    fn evict_existing(&mut self, name: &str, new_kind: &str) {
1556        // Evict from all three maps (each is side-effecting and must run); a
1557        // name lives in at most one map, so the listen > async > sync ordering
1558        // only disambiguates the pathological double-registration case.
1559        let previous_kind = [
1560            (self.listen_tools.remove(name).is_some(), "listen"),
1561            (self.async_tools.remove(name).is_some(), "async"),
1562            (self.tools.remove(name).is_some(), "sync"),
1563        ]
1564        .into_iter()
1565        .find_map(|(removed, kind)| removed.then_some(kind));
1566        if let Some(previous_kind) = previous_kind {
1567            log::warn!(
1568                "tool registry: name {name:?} already registered as a {previous_kind} tool; \
1569                 replacing it with a {new_kind} tool (last registration wins)"
1570            );
1571        }
1572    }
1573
1574    /// Error if `name` is already registered in any of the three maps.
1575    fn ensure_unique(&self, name: &str) -> Result<()> {
1576        anyhow::ensure!(
1577            !self.tools.contains_key(name)
1578                && !self.async_tools.contains_key(name)
1579                && !self.listen_tools.contains_key(name),
1580            "tool {name:?} is already registered",
1581        );
1582        Ok(())
1583    }
1584
1585    /// Register a synchronous tool in the registry.
1586    ///
1587    /// The tool's name is converted to a string via serde serialization
1588    /// and used as the lookup key. If the name is already registered (in any
1589    /// map), the previous tool is evicted and a warning is logged; use
1590    /// [`try_register`](ToolRegistry::try_register) for fail-closed semantics.
1591    pub fn register<T>(&mut self, tool: T) -> &mut Self
1592    where
1593        T: Tool<Ctx> + 'static,
1594    {
1595        let wrapper = ToolWrapper::new(tool);
1596        let name = wrapper.name_str().to_string();
1597        self.evict_existing(&name, "sync");
1598        self.tools.insert(name, Arc::new(wrapper));
1599        self
1600    }
1601
1602    /// Register a synchronous tool, returning an error on name collision.
1603    ///
1604    /// Unlike [`register`](ToolRegistry::register), this never silently
1605    /// replaces an existing tool — it checks all three maps and fails if the
1606    /// name is taken. Useful for registering untrusted (e.g. MCP-supplied)
1607    /// tools without letting them squat over vetted built-ins.
1608    ///
1609    /// # Errors
1610    /// Returns an error if a tool with the same name is already registered.
1611    pub fn try_register<T>(&mut self, tool: T) -> Result<&mut Self>
1612    where
1613        T: Tool<Ctx> + 'static,
1614    {
1615        let wrapper = ToolWrapper::new(tool);
1616        let name = wrapper.name_str().to_string();
1617        self.ensure_unique(&name)?;
1618        self.tools.insert(name, Arc::new(wrapper));
1619        Ok(self)
1620    }
1621
1622    /// Register a [`SimpleTool`] — a tool whose name is a plain `&str` and
1623    /// which needs no [`ToolName`] type.
1624    ///
1625    /// The tool is wrapped in a [`SimpleToolAdapter`] (giving it
1626    /// `Name = DynamicToolName`) and registered like any other [`Tool`].
1627    /// This is the lowest-ceremony way to add a first custom tool.
1628    pub fn register_simple<T>(&mut self, tool: T) -> &mut Self
1629    where
1630        T: SimpleTool<Ctx> + 'static,
1631    {
1632        self.register(SimpleToolAdapter::new(tool))
1633    }
1634
1635    /// Register a [`TypedTool`] — a tool whose model-emitted arguments are
1636    /// deserialized into a typed [`TypedTool::Input`] and validated **before**
1637    /// `execute` runs.
1638    ///
1639    /// The tool is wrapped in a [`TypedToolAdapter`] (giving it
1640    /// `Name = DynamicToolName`) and registered like any other [`Tool`]. A
1641    /// malformed tool call is turned into a structured validation-error
1642    /// [`ToolResult`] at the dispatch boundary so the model can self-correct;
1643    /// `execute` is never reached with invalid arguments.
1644    pub fn register_typed<T>(&mut self, tool: T) -> &mut Self
1645    where
1646        T: TypedTool<Ctx> + 'static,
1647    {
1648        self.register(TypedToolAdapter::new(tool))
1649    }
1650
1651    /// Register an async tool in the registry.
1652    ///
1653    /// Async tools have two phases: execute (lightweight, starts operation)
1654    /// and `check_status` (streams progress until completion).
1655    pub fn register_async<T>(&mut self, tool: T) -> &mut Self
1656    where
1657        T: AsyncTool<Ctx> + 'static,
1658    {
1659        let wrapper = AsyncToolWrapper::new(tool);
1660        let name = wrapper.name_str().to_string();
1661        self.evict_existing(&name, "async");
1662        self.async_tools.insert(name, Arc::new(wrapper));
1663        self
1664    }
1665
1666    /// Register an async tool, returning an error on name collision.
1667    ///
1668    /// The fail-closed counterpart to [`register_async`](ToolRegistry::register_async).
1669    ///
1670    /// # Errors
1671    /// Returns an error if a tool with the same name is already registered.
1672    pub fn try_register_async<T>(&mut self, tool: T) -> Result<&mut Self>
1673    where
1674        T: AsyncTool<Ctx> + 'static,
1675    {
1676        let wrapper = AsyncToolWrapper::new(tool);
1677        let name = wrapper.name_str().to_string();
1678        self.ensure_unique(&name)?;
1679        self.async_tools.insert(name, Arc::new(wrapper));
1680        Ok(self)
1681    }
1682
1683    /// Register a listen/execute tool in the registry.
1684    ///
1685    /// Listen/execute tools start by streaming updates via `listen()`, then run
1686    /// final execution with `execute()` once confirmed.
1687    pub fn register_listen<T>(&mut self, tool: T) -> &mut Self
1688    where
1689        T: ListenExecuteTool<Ctx> + 'static,
1690    {
1691        let wrapper = ListenToolWrapper::new(tool);
1692        let name = wrapper.name_str().to_string();
1693        self.evict_existing(&name, "listen");
1694        self.listen_tools.insert(name, Arc::new(wrapper));
1695        self
1696    }
1697
1698    /// Register a listen/execute tool, returning an error on name collision.
1699    ///
1700    /// The fail-closed counterpart to [`register_listen`](ToolRegistry::register_listen).
1701    ///
1702    /// # Errors
1703    /// Returns an error if a tool with the same name is already registered.
1704    pub fn try_register_listen<T>(&mut self, tool: T) -> Result<&mut Self>
1705    where
1706        T: ListenExecuteTool<Ctx> + 'static,
1707    {
1708        let wrapper = ListenToolWrapper::new(tool);
1709        let name = wrapper.name_str().to_string();
1710        self.ensure_unique(&name)?;
1711        self.listen_tools.insert(name, Arc::new(wrapper));
1712        Ok(self)
1713    }
1714
1715    /// Get a synchronous tool by name.
1716    #[must_use]
1717    pub fn get(&self, name: &str) -> Option<&Arc<dyn ErasedTool<Ctx>>> {
1718        self.tools.get(name)
1719    }
1720
1721    /// Get an async tool by name.
1722    #[must_use]
1723    pub fn get_async(&self, name: &str) -> Option<&Arc<dyn ErasedAsyncTool<Ctx>>> {
1724        self.async_tools.get(name)
1725    }
1726
1727    /// Get a listen/execute tool by name.
1728    #[must_use]
1729    pub fn get_listen(&self, name: &str) -> Option<&Arc<dyn ErasedListenTool<Ctx>>> {
1730        self.listen_tools.get(name)
1731    }
1732
1733    /// Check if a tool name refers to an async tool.
1734    #[must_use]
1735    pub fn is_async(&self, name: &str) -> bool {
1736        self.async_tools.contains_key(name)
1737    }
1738
1739    /// Check if a tool name refers to a listen/execute tool.
1740    #[must_use]
1741    pub fn is_listen(&self, name: &str) -> bool {
1742        self.listen_tools.contains_key(name)
1743    }
1744
1745    /// Get all registered synchronous tools.
1746    pub fn all(&self) -> impl Iterator<Item = &Arc<dyn ErasedTool<Ctx>>> {
1747        self.tools.values()
1748    }
1749
1750    /// Get all registered async tools.
1751    pub fn all_async(&self) -> impl Iterator<Item = &Arc<dyn ErasedAsyncTool<Ctx>>> {
1752        self.async_tools.values()
1753    }
1754
1755    /// Get all registered listen/execute tools.
1756    pub fn all_listen(&self) -> impl Iterator<Item = &Arc<dyn ErasedListenTool<Ctx>>> {
1757        self.listen_tools.values()
1758    }
1759
1760    /// Get the number of registered tools (sync + async).
1761    #[must_use]
1762    pub fn len(&self) -> usize {
1763        self.tools.len() + self.async_tools.len() + self.listen_tools.len()
1764    }
1765
1766    /// Check if the registry is empty.
1767    #[must_use]
1768    pub fn is_empty(&self) -> bool {
1769        self.tools.is_empty() && self.async_tools.is_empty() && self.listen_tools.is_empty()
1770    }
1771
1772    /// Filter tools by a predicate.
1773    ///
1774    /// Removes tools for which the predicate returns false.
1775    /// The predicate receives the tool name.
1776    /// Applies to both sync and async tools.
1777    ///
1778    /// # Example
1779    ///
1780    /// ```ignore
1781    /// registry.filter(|name| name != "bash");
1782    /// ```
1783    pub fn filter<F>(&mut self, predicate: F)
1784    where
1785        F: Fn(&str) -> bool,
1786    {
1787        self.tools.retain(|name, _| predicate(name));
1788        self.async_tools.retain(|name, _| predicate(name));
1789        self.listen_tools.retain(|name, _| predicate(name));
1790    }
1791
1792    /// Convert all tools (sync + async + listen) to LLM tool
1793    /// definitions. The output is sorted by tool name so the order
1794    /// is deterministic across builds and across calls.
1795    ///
1796    /// Determinism matters for **prompt caching**. Anthropic's
1797    /// `cache_control: ephemeral` keys on the byte content of the
1798    /// system + tool list. Anything that perturbs the order of the
1799    /// tool list invalidates the cache. The three backing maps are
1800    /// `HashMap`s, whose `values()` order is randomized (DoS-safe
1801    /// `RandomState` by default), so two consecutive turns with the
1802    /// same registered tool set were producing different orderings
1803    /// and silently zeroing the cache hit rate.
1804    ///
1805    /// Sorting by name is the cheapest fix that holds across
1806    /// insertion order, internal map type changes, and concurrent
1807    /// builds. The tool count is small (tens, not thousands) so the
1808    /// sort cost is negligible compared to a single LLM call.
1809    #[must_use]
1810    pub fn to_llm_tools(&self) -> Vec<llm::Tool> {
1811        /// Build the LLM tool descriptor from the accessors every erased tool
1812        /// trait shares. Extracted so the five-field `llm::Tool` literal —
1813        /// whose byte content is prompt-cache load-bearing — exists in exactly
1814        /// one place across the sync / async / listen iterators.
1815        fn descriptor(
1816            name: &str,
1817            display_name: &str,
1818            description: &str,
1819            input_schema: Value,
1820            tier: ToolTier,
1821        ) -> llm::Tool {
1822            llm::Tool {
1823                name: name.to_string(),
1824                description: description.to_string(),
1825                input_schema,
1826                display_name: display_name.to_string(),
1827                tier,
1828            }
1829        }
1830
1831        let mut tools: Vec<_> = self
1832            .tools
1833            .values()
1834            .map(|tool| {
1835                descriptor(
1836                    tool.name_str(),
1837                    tool.display_name(),
1838                    tool.description(),
1839                    tool.input_schema(),
1840                    tool.tier(),
1841                )
1842            })
1843            .collect();
1844
1845        tools.extend(self.async_tools.values().map(|tool| {
1846            descriptor(
1847                tool.name_str(),
1848                tool.display_name(),
1849                tool.description(),
1850                tool.input_schema(),
1851                tool.tier(),
1852            )
1853        }));
1854
1855        tools.extend(self.listen_tools.values().map(|tool| {
1856            descriptor(
1857                tool.name_str(),
1858                tool.display_name(),
1859                tool.description(),
1860                tool.input_schema(),
1861                tool.tier(),
1862            )
1863        }));
1864
1865        tools.sort_by(|a, b| a.name.cmp(&b.name));
1866        tools
1867    }
1868}
1869
1870#[cfg(test)]
1871mod tests {
1872    use super::*;
1873    use anyhow::Context;
1874
1875    // Test tool name enum for tests
1876    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1877    #[serde(rename_all = "snake_case")]
1878    enum TestToolName {
1879        MockTool,
1880        AnotherTool,
1881    }
1882
1883    impl ToolName for TestToolName {}
1884
1885    struct MockTool;
1886
1887    impl Tool<()> for MockTool {
1888        type Name = TestToolName;
1889
1890        fn name(&self) -> TestToolName {
1891            TestToolName::MockTool
1892        }
1893
1894        fn display_name(&self) -> &'static str {
1895            "Mock Tool"
1896        }
1897
1898        fn description(&self) -> &'static str {
1899            "A mock tool for testing"
1900        }
1901
1902        fn input_schema(&self) -> Value {
1903            serde_json::json!({
1904                "type": "object",
1905                "properties": {
1906                    "message": { "type": "string" }
1907                }
1908            })
1909        }
1910
1911        async fn execute(&self, _ctx: &ToolContext<()>, input: Value) -> Result<ToolResult> {
1912            let message = input
1913                .get("message")
1914                .and_then(|v| v.as_str())
1915                .unwrap_or("no message");
1916            Ok(ToolResult::success(format!("Received: {message}")))
1917        }
1918    }
1919
1920    #[test]
1921    fn test_tool_name_serialization() {
1922        let name = TestToolName::MockTool;
1923        assert_eq!(tool_name_to_string(&name), "mock_tool");
1924
1925        let parsed: TestToolName = tool_name_from_str("mock_tool").unwrap();
1926        assert_eq!(parsed, TestToolName::MockTool);
1927    }
1928
1929    #[test]
1930    fn test_dynamic_tool_name() {
1931        let name = DynamicToolName::new("my_mcp_tool");
1932        assert_eq!(tool_name_to_string(&name), "my_mcp_tool");
1933        assert_eq!(name.as_str(), "my_mcp_tool");
1934    }
1935
1936    #[test]
1937    fn test_tool_registry() {
1938        let mut registry = ToolRegistry::new();
1939        registry.register(MockTool);
1940
1941        assert_eq!(registry.len(), 1);
1942        assert!(registry.get("mock_tool").is_some());
1943        assert!(registry.get("nonexistent").is_none());
1944    }
1945
1946    #[test]
1947    fn test_to_llm_tools() {
1948        let mut registry = ToolRegistry::new();
1949        registry.register(MockTool);
1950
1951        let llm_tools = registry.to_llm_tools();
1952        assert_eq!(llm_tools.len(), 1);
1953        assert_eq!(llm_tools[0].name, "mock_tool");
1954    }
1955
1956    #[test]
1957    fn to_llm_tools_returns_alphabetical_order() {
1958        let mut registry = ToolRegistry::new();
1959        // Register in non-alphabetical order so the assertion would
1960        // fail if we ever returned insertion order again.
1961        registry.register(MockTool); // "mock_tool"
1962        registry.register(AnotherTool); // "another_tool"
1963
1964        let names: Vec<String> = registry
1965            .to_llm_tools()
1966            .into_iter()
1967            .map(|t| t.name)
1968            .collect();
1969        assert_eq!(names, vec!["another_tool", "mock_tool"]);
1970    }
1971
1972    #[test]
1973    fn to_llm_tools_is_deterministic_across_calls() {
1974        // Regression: prompt caching depends on byte-stable tool list
1975        // ordering. The `HashMap` behind the registry randomizes its
1976        // `values()` order, so without an explicit sort two consecutive
1977        // builds with the same registered set could ship different
1978        // tool orderings to the LLM and silently invalidate the cache.
1979        let mut registry = ToolRegistry::new();
1980        registry.register(MockTool);
1981        registry.register(AnotherTool);
1982
1983        let first: Vec<String> = registry
1984            .to_llm_tools()
1985            .into_iter()
1986            .map(|t| t.name)
1987            .collect();
1988
1989        for _ in 0..32 {
1990            let next: Vec<String> = registry
1991                .to_llm_tools()
1992                .into_iter()
1993                .map(|t| t.name)
1994                .collect();
1995            assert_eq!(next, first, "tool ordering must be stable across calls");
1996        }
1997    }
1998
1999    struct AnotherTool;
2000
2001    impl Tool<()> for AnotherTool {
2002        type Name = TestToolName;
2003
2004        fn name(&self) -> TestToolName {
2005            TestToolName::AnotherTool
2006        }
2007
2008        fn display_name(&self) -> &'static str {
2009            "Another Tool"
2010        }
2011
2012        fn description(&self) -> &'static str {
2013            "Another tool for testing"
2014        }
2015
2016        fn input_schema(&self) -> Value {
2017            serde_json::json!({ "type": "object" })
2018        }
2019
2020        async fn execute(&self, _ctx: &ToolContext<()>, _input: Value) -> Result<ToolResult> {
2021            Ok(ToolResult::success("Done"))
2022        }
2023    }
2024
2025    #[test]
2026    fn test_filter_tools() {
2027        let mut registry = ToolRegistry::new();
2028        registry.register(MockTool);
2029        registry.register(AnotherTool);
2030
2031        assert_eq!(registry.len(), 2);
2032
2033        // Filter out mock_tool
2034        registry.filter(|name| name != "mock_tool");
2035
2036        assert_eq!(registry.len(), 1);
2037        assert!(registry.get("mock_tool").is_none());
2038        assert!(registry.get("another_tool").is_some());
2039    }
2040
2041    #[test]
2042    fn test_filter_tools_keep_all() {
2043        let mut registry = ToolRegistry::new();
2044        registry.register(MockTool);
2045        registry.register(AnotherTool);
2046
2047        registry.filter(|_| true);
2048
2049        assert_eq!(registry.len(), 2);
2050    }
2051
2052    #[test]
2053    fn test_filter_tools_remove_all() {
2054        let mut registry = ToolRegistry::new();
2055        registry.register(MockTool);
2056        registry.register(AnotherTool);
2057
2058        registry.filter(|_| false);
2059
2060        assert!(registry.is_empty());
2061    }
2062
2063    #[test]
2064    fn test_display_name() {
2065        let mut registry = ToolRegistry::new();
2066        registry.register(MockTool);
2067
2068        let tool = registry.get("mock_tool").unwrap();
2069        assert_eq!(tool.display_name(), "Mock Tool");
2070    }
2071
2072    struct ListenMockTool;
2073
2074    impl ListenExecuteTool<()> for ListenMockTool {
2075        type Name = TestToolName;
2076
2077        fn name(&self) -> TestToolName {
2078            TestToolName::MockTool
2079        }
2080
2081        fn display_name(&self) -> &'static str {
2082            "Listen Mock Tool"
2083        }
2084
2085        fn description(&self) -> &'static str {
2086            "A listen/execute mock tool for testing"
2087        }
2088
2089        fn input_schema(&self) -> Value {
2090            serde_json::json!({ "type": "object" })
2091        }
2092
2093        fn listen(
2094            &self,
2095            _ctx: &ToolContext<()>,
2096            _input: Value,
2097        ) -> impl futures::Stream<Item = ListenToolUpdate> + Send {
2098            futures::stream::iter(vec![ListenToolUpdate::Ready {
2099                operation_id: "op_1".to_string(),
2100                revision: 1,
2101                message: "ready".to_string(),
2102                snapshot: serde_json::json!({"ok": true}),
2103                expires_at: None,
2104            }])
2105        }
2106
2107        async fn execute(
2108            &self,
2109            _ctx: &ToolContext<()>,
2110            _operation_id: &str,
2111            _expected_revision: u64,
2112        ) -> Result<ToolResult> {
2113            Ok(ToolResult::success("Executed"))
2114        }
2115    }
2116
2117    #[test]
2118    fn test_listen_tool_registry() {
2119        let mut registry = ToolRegistry::new();
2120        registry.register_listen(ListenMockTool);
2121
2122        assert_eq!(registry.len(), 1);
2123        assert!(registry.get_listen("mock_tool").is_some());
2124        assert!(registry.is_listen("mock_tool"));
2125    }
2126
2127    // ── TypedTool: typed input + validation / self-correction ───────────
2128
2129    use std::sync::atomic::{AtomicBool, Ordering};
2130
2131    #[derive(Debug, Serialize, Deserialize)]
2132    struct GreetArgs {
2133        name: String,
2134        // Required so a missing/typo'd field is a hard validation error.
2135        greeting: String,
2136    }
2137
2138    /// A typed tool that records whether `execute` was reached, so tests can
2139    /// assert the validation boundary never calls `execute` with bad args.
2140    struct GreetTool {
2141        executed: Arc<AtomicBool>,
2142    }
2143
2144    impl TypedTool<()> for GreetTool {
2145        type Input = GreetArgs;
2146
2147        fn name(&self) -> &'static str {
2148            "greet"
2149        }
2150
2151        fn description(&self) -> &'static str {
2152            "Greet someone by name"
2153        }
2154
2155        fn input_schema(&self) -> Value {
2156            serde_json::json!({
2157                "type": "object",
2158                "properties": {
2159                    "name": { "type": "string" },
2160                    "greeting": { "type": "string" }
2161                },
2162                "required": ["name", "greeting"]
2163            })
2164        }
2165
2166        async fn execute(&self, _ctx: &ToolContext<()>, input: GreetArgs) -> Result<ToolResult> {
2167            self.executed.store(true, Ordering::SeqCst);
2168            Ok(ToolResult::success(format!(
2169                "{}, {}!",
2170                input.greeting, input.name
2171            )))
2172        }
2173    }
2174
2175    #[tokio::test]
2176    async fn typed_tool_happy_path_receives_typed_input() -> Result<()> {
2177        let executed = Arc::new(AtomicBool::new(false));
2178        let adapter = TypedToolAdapter::new(GreetTool {
2179            executed: executed.clone(),
2180        });
2181        let ctx = ToolContext::new(());
2182
2183        let result = Tool::execute(
2184            &adapter,
2185            &ctx,
2186            serde_json::json!({ "name": "Ada", "greeting": "Hello" }),
2187        )
2188        .await?;
2189
2190        assert!(executed.load(Ordering::SeqCst), "execute must be called");
2191        assert!(result.success);
2192        assert_eq!(result.output, "Hello, Ada!");
2193        Ok(())
2194    }
2195
2196    #[tokio::test]
2197    async fn typed_tool_invalid_args_self_correct_without_executing() -> Result<()> {
2198        let executed = Arc::new(AtomicBool::new(false));
2199        let adapter = TypedToolAdapter::new(GreetTool {
2200            executed: executed.clone(),
2201        });
2202        let ctx = ToolContext::new(());
2203
2204        // `greeting` is missing — must not deserialize into `GreetArgs`.
2205        let result = Tool::execute(&adapter, &ctx, serde_json::json!({ "name": "Ada" })).await?;
2206
2207        assert!(
2208            !executed.load(Ordering::SeqCst),
2209            "execute must NOT be called with invalid arguments"
2210        );
2211        assert!(!result.success, "validation failure is an error result");
2212        assert!(
2213            result.output.contains("Invalid arguments for tool `greet`"),
2214            "error must identify the tool: {}",
2215            result.output
2216        );
2217        assert!(
2218            result.output.contains("greeting"),
2219            "error must surface the serde message naming the bad field: {}",
2220            result.output
2221        );
2222        Ok(())
2223    }
2224
2225    #[tokio::test]
2226    async fn typed_tool_wrong_type_self_corrects() -> Result<()> {
2227        let executed = Arc::new(AtomicBool::new(false));
2228        let adapter = TypedToolAdapter::new(GreetTool {
2229            executed: executed.clone(),
2230        });
2231        let ctx = ToolContext::new(());
2232
2233        // `name` is a number, not a string.
2234        let result = Tool::execute(
2235            &adapter,
2236            &ctx,
2237            serde_json::json!({ "name": 42, "greeting": "Hi" }),
2238        )
2239        .await?;
2240
2241        assert!(!executed.load(Ordering::SeqCst));
2242        assert!(!result.success);
2243        Ok(())
2244    }
2245
2246    /// Back-compat: a `TypedTool` whose `Input = Value` is the identity
2247    /// passthrough — any JSON deserializes, mirroring today's untyped tools.
2248    struct ValueTypedTool;
2249
2250    impl TypedTool<()> for ValueTypedTool {
2251        type Input = Value;
2252
2253        fn name(&self) -> &'static str {
2254            "value_typed"
2255        }
2256
2257        fn description(&self) -> &'static str {
2258            "Accepts any JSON, like an untyped tool"
2259        }
2260
2261        fn input_schema(&self) -> Value {
2262            serde_json::json!({ "type": "object" })
2263        }
2264
2265        async fn execute(&self, _ctx: &ToolContext<()>, input: Value) -> Result<ToolResult> {
2266            Ok(ToolResult::success(input.to_string()))
2267        }
2268    }
2269
2270    #[tokio::test]
2271    async fn typed_tool_value_input_is_identity_passthrough() -> Result<()> {
2272        let adapter = TypedToolAdapter::new(ValueTypedTool);
2273        let ctx = ToolContext::new(());
2274
2275        // Arbitrary shape — Value always "deserializes".
2276        let result = Tool::execute(
2277            &adapter,
2278            &ctx,
2279            serde_json::json!({ "anything": [1, 2, 3], "nested": { "ok": true } }),
2280        )
2281        .await?;
2282
2283        assert!(result.success);
2284        Ok(())
2285    }
2286
2287    #[test]
2288    fn register_typed_exposes_tool_via_registry() -> Result<()> {
2289        let mut registry = ToolRegistry::new();
2290        registry.register_typed(GreetTool {
2291            executed: Arc::new(AtomicBool::new(false)),
2292        });
2293
2294        assert_eq!(registry.len(), 1);
2295        let tool = registry.get("greet").context("typed tool registered")?;
2296        // The user-declared schema flows through unchanged.
2297        assert_eq!(tool.input_schema()["required"][0], "name");
2298        Ok(())
2299    }
2300
2301    #[test]
2302    fn invalid_tool_input_result_is_balanced_error() -> Result<()> {
2303        let Err(err) = serde_json::from_str::<GreetArgs>("{}") else {
2304            anyhow::bail!("empty object must fail to deserialize GreetArgs");
2305        };
2306        let result = invalid_tool_input_result("greet", &err);
2307
2308        assert!(!result.success);
2309        assert!(result.output.contains("greet"));
2310        assert!(result.output.contains("call the tool again"));
2311        Ok(())
2312    }
2313
2314    // ── Fail-closed tier + display_name defaults (findings 8 & 19) ───────
2315
2316    /// A tool that overrides neither `tier()` nor `display_name()`, exercising
2317    /// the trait defaults.
2318    struct DefaultsTool;
2319
2320    impl Tool<()> for DefaultsTool {
2321        type Name = DynamicToolName;
2322
2323        fn name(&self) -> DynamicToolName {
2324            DynamicToolName::new("defaults")
2325        }
2326
2327        fn description(&self) -> &'static str {
2328            "uses trait defaults"
2329        }
2330
2331        fn input_schema(&self) -> Value {
2332            serde_json::json!({ "type": "object" })
2333        }
2334
2335        async fn execute(&self, _ctx: &ToolContext<()>, _input: Value) -> Result<ToolResult> {
2336            Ok(ToolResult::success("ok"))
2337        }
2338    }
2339
2340    #[test]
2341    fn tool_trait_defaults_are_fail_closed() {
2342        let tool = DefaultsTool;
2343        // display_name defaults to "" (finding 19: the doc-claimed default now
2344        // actually exists).
2345        assert_eq!(Tool::display_name(&tool), "");
2346        // tier defaults to Confirm so a side-effecting tool whose author forgot
2347        // to declare a tier is gated, not auto-run (finding 8).
2348        assert_eq!(Tool::tier(&tool), ToolTier::Confirm);
2349    }
2350
2351    // ── Registry name-collision handling (findings 9 & 10) ───────────────
2352
2353    #[test]
2354    fn re_registering_same_name_replaces_without_duplicates() {
2355        let mut registry = ToolRegistry::new();
2356        registry.register(MockTool);
2357        registry.register(MockTool); // same name "mock_tool"
2358
2359        assert_eq!(registry.len(), 1, "re-register must replace, not add");
2360        let names: Vec<String> = registry
2361            .to_llm_tools()
2362            .into_iter()
2363            .map(|t| t.name)
2364            .collect();
2365        assert_eq!(names, vec!["mock_tool"]);
2366    }
2367
2368    #[test]
2369    fn cross_kind_name_collision_keeps_single_entry() {
2370        let mut registry = ToolRegistry::new();
2371        registry.register(MockTool); // sync "mock_tool"
2372        registry.register_listen(ListenMockTool); // listen "mock_tool"
2373
2374        // The listen registration evicts the sync one — a name lives in exactly
2375        // one map, so `len()` and `to_llm_tools()` never double-count it.
2376        assert_eq!(registry.len(), 1);
2377        assert!(registry.is_listen("mock_tool"));
2378        assert!(
2379            registry.get("mock_tool").is_none(),
2380            "the shadowed sync tool must be evicted"
2381        );
2382        let names: Vec<String> = registry
2383            .to_llm_tools()
2384            .into_iter()
2385            .map(|t| t.name)
2386            .collect();
2387        assert_eq!(names, vec!["mock_tool"], "no duplicate LLM definitions");
2388    }
2389
2390    #[test]
2391    fn try_register_rejects_name_collision() {
2392        let mut registry = ToolRegistry::new();
2393        registry.register(MockTool); // "mock_tool"
2394
2395        assert!(
2396            registry.try_register(MockTool).is_err(),
2397            "duplicate sync name must be rejected"
2398        );
2399        assert!(
2400            registry.try_register_listen(ListenMockTool).is_err(),
2401            "cross-map duplicate (squatting) must be rejected"
2402        );
2403        assert_eq!(
2404            registry.len(),
2405            1,
2406            "rejected registrations must not be stored"
2407        );
2408    }
2409
2410    // ── Non-panicking serde helpers (findings 16 & 17) ───────────────────
2411
2412    #[derive(Clone)]
2413    struct FailingStage;
2414
2415    impl Serialize for FailingStage {
2416        fn serialize<S>(&self, _serializer: S) -> core::result::Result<S::Ok, S::Error>
2417        where
2418            S: serde::Serializer,
2419        {
2420            Err(serde::ser::Error::custom("intentionally unserializable"))
2421        }
2422    }
2423
2424    impl<'de> Deserialize<'de> for FailingStage {
2425        fn deserialize<D>(_deserializer: D) -> core::result::Result<Self, D::Error>
2426        where
2427            D: serde::Deserializer<'de>,
2428        {
2429            Ok(Self)
2430        }
2431    }
2432
2433    impl ProgressStage for FailingStage {}
2434
2435    #[test]
2436    fn stage_to_string_falls_back_instead_of_panicking() {
2437        // A ProgressStage whose Serialize impl fails must not panic the turn
2438        // loop on the async-tool progress hot path.
2439        assert_eq!(stage_to_string(&FailingStage), "<unknown_stage>");
2440    }
2441
2442    #[test]
2443    fn tool_name_from_str_round_trips_special_characters() -> Result<()> {
2444        // Names with quotes/backslashes (possible from remote MCP servers) must
2445        // be JSON-escaped, not interpolated raw, so parsing succeeds.
2446        let name: DynamicToolName =
2447            tool_name_from_str("weird\"name\\with-escapes").context("must parse escaped name")?;
2448        assert_eq!(name.as_str(), "weird\"name\\with-escapes");
2449        Ok(())
2450    }
2451
2452    // ── emit_event surfaces unbound misuse instead of silently dropping ──
2453
2454    #[tokio::test]
2455    async fn emit_event_persists_when_bound_and_is_noop_when_unbound() -> Result<()> {
2456        use crate::stores::InMemoryEventStore;
2457        use agent_sdk_foundation::types::ThreadId;
2458
2459        let store = Arc::new(InMemoryEventStore::new());
2460        let thread_id = ThreadId::new();
2461        let authority: Arc<dyn EventAuthority> = Arc::new(LocalEventAuthority::new());
2462
2463        let bound =
2464            ToolContext::new(()).with_event_store(store.clone(), thread_id.clone(), 1, authority);
2465        bound.emit_event(AgentEvent::text("m1", "hi")).await?;
2466        assert_eq!(
2467            store.event_count(&thread_id).await?,
2468            1,
2469            "a bound context persists the event"
2470        );
2471
2472        // An unbound context is a no-op (the fix also logs a warning) — it must
2473        // not silently append elsewhere or error.
2474        let unbound = ToolContext::new(());
2475        unbound.emit_event(AgentEvent::text("m2", "lost")).await?;
2476        assert_eq!(
2477            store.event_count(&thread_id).await?,
2478            1,
2479            "an unbound context changes nothing"
2480        );
2481        Ok(())
2482    }
2483}