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}