Skip to main content

aither_core/llm/
tool.rs

1//! # LLM Tool Calling Framework
2//!
3//!
4//! Type-safe tool calling system for Large Language Models. Enables LLMs to execute external
5//! functions, access APIs, and interact with systems through well-defined interfaces.
6//!
7//! ## How a model calls a tool
8//!
9//! A [`Tool`] declares a name, a description, and an `Arguments` type whose
10//! JSON schema is sent to the provider. When the model decides to use a tool it
11//! emits a [`ToolCall`](crate::llm::ToolCall) event carrying the tool name and
12//! arguments as JSON; this crate does not execute it. Executing the call and
13//! feeding the result back is the job of a higher layer such as `aither-agent`,
14//! which keeps tool execution under the caller's control.
15//!
16//! ## Core Components
17//!
18//! - [`Tool`] - Trait for defining executable tools
19//! - [`Tools`] - Registry for managing multiple tools  
20//! - [`tool::ToolDefinition`] - Metadata and schema for LLM consumption
21//!
22//! ## Quick Start
23//!
24//! ```rust
25//! use aither_core::llm::{Tool, ToolResult};
26//! use schemars::JsonSchema;
27//! use serde::Deserialize;
28//! use std::borrow::Cow;
29//!
30//! /// Performs basic math operations.
31//! #[derive(JsonSchema, Deserialize)]
32//! struct MathArgs {
33//!     /// Operation: "add", "subtract", "multiply", "divide"
34//!     operation: String,
35//!     /// First number
36//!     a: f64,
37//!     /// Second number
38//!     b: f64,
39//! }
40//!
41//! struct Calculator;
42//!
43//! impl Tool for Calculator {
44//!     fn name(&self) -> Cow<'static, str> {
45//!         Cow::Borrowed("calculator")
46//!     }
47//!
48//!     type Arguments = MathArgs;
49//!     type Res = ToolResult;
50//!
51//!     async fn call(&self, args: Self::Arguments) -> aither_core::Result<Self::Res> {
52//!         let result = match args.operation.as_str() {
53//!             "add" => args.a + args.b,
54//!             "subtract" => args.a - args.b,
55//!             "multiply" => args.a * args.b,
56//!             "divide" if args.b != 0.0 => args.a / args.b,
57//!             "divide" => return Err(anyhow::Error::msg("Division by zero")),
58//!             _ => return Err(anyhow::Error::msg("Unknown operation")),
59//!         };
60//!         Ok(ToolResult::text(result.to_string()))
61//!     }
62//! }
63//! ```
64//!
65//! ## Schema Design Best Practices
66//!
67//! ### 1. Use Clear Documentation Comments
68//! Doc comments automatically become schema descriptions:
69//!
70//! ```rust,ignore
71//! use schemars::JsonSchema;
72//! use serde::Deserialize;
73//!
74//! #[derive(JsonSchema, Deserialize)]
75//! struct WeatherArgs {
76//!     /// City name (e.g., "London", "Tokyo", "New York")
77//!     city: String,
78//!     /// Temperature unit: "celsius" or "fahrenheit"
79//!     #[serde(default = "default_celsius")]
80//!     unit: String,
81//! }
82//!
83//! fn default_celsius() -> String { "celsius".to_string() }
84//! ```
85//!
86//! ### 2. Prefer Enums Over Strings
87//! Enums provide clear constraints for LLMs:
88//!
89//! ```rust,ignore
90//! use schemars::JsonSchema;
91//! use serde::Deserialize;
92//!
93//! #[derive(JsonSchema, Deserialize)]
94//! enum Priority { Low, Medium, High, Critical }
95//!
96//! #[derive(JsonSchema, Deserialize)]  
97//! struct TaskArgs {
98//!     /// Task description
99//!     description: String,
100//!     /// Task priority level
101//!     priority: Priority,
102//! }
103//! ```
104//!
105//! ### 3. Add Validation Constraints
106//! Use schemars attributes for validation:
107//!
108//! ```rust,ignore
109//! use schemars::JsonSchema;
110//! use serde::Deserialize;
111//!
112//! #[derive(JsonSchema, Deserialize)]
113//! struct UserArgs {
114//!     /// Valid email address
115//!     #[schemars(regex(pattern = "^[^@]+@[^@]+\\.[^@]+$"))]
116//!     email: String,
117//!     /// Age between 13 and 120
118//!     #[schemars(range(min = 13, max = 120))]
119//!     age: u8,
120//!     /// Bio text, max 500 characters
121//!     #[schemars(length(max = 500))]
122//!     bio: Option<String>,
123//! }
124//! ```
125//!
126//! ### 4. Structure Complex Data
127//! Break down complex parameters into nested types:
128//!
129//! ```rust,ignore
130//! use schemars::JsonSchema;
131//! use serde::Deserialize;
132//!
133//! #[derive(JsonSchema, Deserialize)]
134//! struct Address {
135//!     street: String,
136//!     city: String,
137//!     /// Two-letter country code (e.g., "US", "GB", "JP")
138//!     country: String,
139//! }
140//!
141//! #[derive(JsonSchema, Deserialize)]
142//! struct CreateUserArgs {
143//!     name: String,
144//!     address: Address,
145//!     /// List of user interests
146//!     #[schemars(length(max = 10))]
147//!     interests: Vec<String>,
148//! }
149//! ```
150//!
151
152// Re-export procedural macros
153#[cfg(feature = "derive")]
154pub use aither_derive::tool;
155use alloc::borrow::Cow;
156use serde_json::Value;
157
158use crate::Result;
159use alloc::format;
160use alloc::string::{String, ToString};
161use alloc::vec::Vec;
162use alloc::{boxed::Box, collections::BTreeMap};
163use core::any::Any;
164use core::fmt::{Debug, Display};
165use core::{future::Future, pin::Pin};
166pub use mime::Mime;
167use schemars::{JsonSchema, Schema, schema_for};
168use serde::{Serialize, de::DeserializeOwned};
169
170/// Final structured result from a tool execution.
171///
172/// This keeps successful content and tool-level failures in distinct variants so
173/// UIs and runtimes do not need to infer errors from free-form strings.
174///
175/// # Example
176///
177/// ```rust,ignore
178/// use aither::llm::tool::ToolResult;
179///
180/// fn search_tool() -> ToolResult {
181///     ToolResult::text("Found 3 results...")
182/// }
183///
184/// fn delete_tool() -> ToolResult {
185///     ToolResult::Done
186/// }
187/// ```
188#[derive(Debug, Clone, PartialEq, Eq)]
189#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
190#[cfg_attr(feature = "serde", serde(tag = "kind", rename_all = "snake_case"))]
191pub enum ToolResult {
192    /// Tool completed with no output to return.
193    Done,
194
195    /// UTF-8 plain text output.
196    Text {
197        /// Plain text content.
198        text: String,
199    },
200
201    /// Tab-separated tabular output.
202    Tsv {
203        /// TSV content.
204        text: String,
205    },
206
207    /// Structured JSON output.
208    Json {
209        /// JSON content.
210        value: Value,
211    },
212
213    /// Binary or media payload.
214    Binary {
215        /// MIME type of the payload (for example `image/png`).
216        mime: String,
217        /// Raw content bytes.
218        content: Vec<u8>,
219    },
220
221    /// Tool-level error. This is distinct from transport/runtime failures.
222    Error {
223        /// Tool-level error message.
224        message: String,
225    },
226
227    /// Several content items produced by one tool call.
228    ///
229    /// A single call can legitimately yield more than one payload — for
230    /// example a relayed MCP tool result that pairs explanatory text with a
231    /// screenshot. Parts never nest and never carry errors: tool-level
232    /// failures stay on [`ToolResult::Error`].
233    Parts {
234        /// Ordered content items.
235        parts: Vec<ToolResultPart>,
236    },
237}
238
239/// One content item inside [`ToolResult::Parts`].
240///
241/// Mirrors the payload-carrying [`ToolResult`] variants. `Done` and `Error`
242/// have no part equivalent: they describe the call as a whole.
243#[derive(Debug, Clone, PartialEq, Eq)]
244#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
245#[cfg_attr(feature = "serde", serde(tag = "kind", rename_all = "snake_case"))]
246pub enum ToolResultPart {
247    /// UTF-8 plain text output.
248    Text {
249        /// Plain text content.
250        text: String,
251    },
252
253    /// Tab-separated tabular output.
254    Tsv {
255        /// TSV content.
256        text: String,
257    },
258
259    /// Structured JSON output.
260    Json {
261        /// JSON content.
262        value: Value,
263    },
264
265    /// Binary or media payload.
266    Binary {
267        /// MIME type of the payload (for example `image/png`).
268        mime: String,
269        /// Raw content bytes.
270        content: Vec<u8>,
271    },
272}
273
274impl ToolResult {
275    /// Creates a plain text result.
276    #[must_use]
277    pub fn text(s: impl Into<String>) -> Self {
278        Self::Text { text: s.into() }
279    }
280
281    /// Creates a TSV result.
282    #[must_use]
283    pub fn tsv(s: impl Into<String>) -> Self {
284        Self::Tsv { text: s.into() }
285    }
286
287    /// Creates a JSON result from a serializable value.
288    ///
289    /// # Errors
290    ///
291    /// Returns an error if serialization fails.
292    pub fn json<T: Serialize>(value: &T) -> Result<Self> {
293        Ok(Self::Json {
294            value: serde_json::to_value(value)?,
295        })
296    }
297
298    /// Creates a JSON result from an already-materialized JSON value.
299    #[must_use]
300    pub const fn json_value(value: Value) -> Self {
301        Self::Json { value }
302    }
303
304    /// Creates an image result.
305    #[must_use]
306    pub fn image(data: Vec<u8>, media_type: &str) -> Self {
307        Self::Binary {
308            mime: parse_media_type_or_octet_stream(media_type),
309            content: data,
310        }
311    }
312
313    /// Creates a binary result.
314    #[must_use]
315    pub fn binary(data: Vec<u8>) -> Self {
316        Self::Binary {
317            mime: mime::APPLICATION_OCTET_STREAM.essence_str().to_string(),
318            content: data,
319        }
320    }
321
322    /// Creates a typed tool error result.
323    #[must_use]
324    pub fn error(message: impl Into<String>) -> Self {
325        Self::Error {
326            message: message.into(),
327        }
328    }
329
330    /// Creates a multi-part result.
331    ///
332    /// An empty part list collapses to [`ToolResult::Done`] and a single part
333    /// collapses to the matching single-payload variant, so callers that
334    /// conditionally accumulate parts never produce a degenerate `Parts`.
335    #[must_use]
336    pub fn parts(parts: Vec<ToolResultPart>) -> Self {
337        match <[ToolResultPart; 1]>::try_from(parts) {
338            Ok([part]) => part.into_result(),
339            Err(parts) if parts.is_empty() => Self::Done,
340            Err(parts) => Self::Parts { parts },
341        }
342    }
343
344    /// Returns `true` if this is a `Done` variant.
345    #[must_use]
346    pub const fn is_done(&self) -> bool {
347        matches!(self, Self::Done)
348    }
349
350    /// Returns `true` if this is a typed tool error.
351    #[must_use]
352    pub const fn is_error(&self) -> bool {
353        matches!(self, Self::Error { .. })
354    }
355
356    /// Returns plain textual content when every payload this result carries is text-like.
357    #[must_use]
358    pub fn as_text(&self) -> Option<&str> {
359        match self {
360            Self::Text { text } | Self::Tsv { text } => Some(text),
361            Self::Error { message } => Some(message),
362            Self::Done | Self::Json { .. } | Self::Binary { .. } | Self::Parts { .. } => None,
363        }
364    }
365
366    /// Returns the error message if this is a typed tool error.
367    #[must_use]
368    pub fn error_message(&self) -> Option<&str> {
369        match self {
370            Self::Error { message } => Some(message),
371            Self::Done
372            | Self::Text { .. }
373            | Self::Tsv { .. }
374            | Self::Json { .. }
375            | Self::Binary { .. }
376            | Self::Parts { .. } => None,
377        }
378    }
379
380    /// Projects the result into a textual representation safe to re-inject into model context.
381    ///
382    /// # Errors
383    ///
384    /// Returns an error if JSON serialization fails.
385    pub fn render_for_model(&self) -> Result<String> {
386        match self {
387            Self::Done => Ok(String::new()),
388            Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
389            Self::Json { value } => Ok(serde_json::to_string(value)?),
390            Self::Binary { mime, content } => Ok(render_binary_placeholder(mime, content)),
391            Self::Error { message } => Ok(message.clone()),
392            Self::Parts { parts } => join_part_renders(parts, false),
393        }
394    }
395
396    /// Renders the result for CLI display.
397    ///
398    /// # Errors
399    ///
400    /// Returns an error if JSON serialization fails.
401    pub fn render_for_cli(&self) -> Result<String> {
402        match self {
403            Self::Done => Ok(String::new()),
404            Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
405            Self::Json { value } => Ok(serde_json::to_string_pretty(value)?),
406            Self::Binary { mime, content } => Ok(render_binary_placeholder(mime, content)),
407            Self::Error { message } => Ok(message.clone()),
408            Self::Parts { parts } => join_part_renders(parts, true),
409        }
410    }
411
412    /// Parses and returns the MIME type when this result carries binary content.
413    ///
414    /// A multi-part result yields a MIME type only when it holds exactly one
415    /// binary part; mixed content has no single type to report.
416    #[must_use]
417    pub fn mime(&self) -> Option<Mime> {
418        match self {
419            Self::Binary { mime, .. } => mime.parse().ok(),
420            Self::Parts { parts } => match parts.as_slice() {
421                [ToolResultPart::Binary { mime, .. }] => mime.parse().ok(),
422                _ => None,
423            },
424            Self::Done
425            | Self::Text { .. }
426            | Self::Tsv { .. }
427            | Self::Json { .. }
428            | Self::Error { .. } => None,
429        }
430    }
431
432    /// Returns raw bytes when this result carries binary content.
433    ///
434    /// Follows the same single-binary rule as [`Self::mime`].
435    #[must_use]
436    pub fn content(&self) -> Option<&[u8]> {
437        match self {
438            Self::Binary { content, .. } => Some(content),
439            Self::Parts { parts } => match parts.as_slice() {
440                [ToolResultPart::Binary { content, .. }] => Some(content),
441                _ => None,
442            },
443            Self::Done
444            | Self::Text { .. }
445            | Self::Tsv { .. }
446            | Self::Json { .. }
447            | Self::Error { .. } => None,
448        }
449    }
450}
451
452impl ToolResultPart {
453    /// Creates a plain text part.
454    #[must_use]
455    pub fn text(s: impl Into<String>) -> Self {
456        Self::Text { text: s.into() }
457    }
458
459    /// Creates a TSV part.
460    #[must_use]
461    pub fn tsv(s: impl Into<String>) -> Self {
462        Self::Tsv { text: s.into() }
463    }
464
465    /// Creates a JSON part from a serializable value.
466    ///
467    /// # Errors
468    ///
469    /// Returns an error if serialization fails.
470    pub fn json<T: Serialize>(value: &T) -> Result<Self> {
471        Ok(Self::Json {
472            value: serde_json::to_value(value)?,
473        })
474    }
475
476    /// Creates a JSON part from an already-materialized JSON value.
477    #[must_use]
478    pub const fn json_value(value: Value) -> Self {
479        Self::Json { value }
480    }
481
482    /// Creates an image part.
483    #[must_use]
484    pub fn image(data: Vec<u8>, media_type: &str) -> Self {
485        Self::Binary {
486            mime: parse_media_type_or_octet_stream(media_type),
487            content: data,
488        }
489    }
490
491    /// Creates a binary part.
492    #[must_use]
493    pub fn binary(data: Vec<u8>) -> Self {
494        Self::Binary {
495            mime: mime::APPLICATION_OCTET_STREAM.essence_str().to_string(),
496            content: data,
497        }
498    }
499
500    /// Returns plain textual content for text-like parts.
501    #[must_use]
502    pub fn as_text(&self) -> Option<&str> {
503        match self {
504            Self::Text { text } | Self::Tsv { text } => Some(text),
505            Self::Json { .. } | Self::Binary { .. } => None,
506        }
507    }
508
509    /// Projects the part into a textual representation safe to re-inject into model context.
510    ///
511    /// # Errors
512    ///
513    /// Returns an error if JSON serialization fails.
514    pub fn render_for_model(&self) -> Result<String> {
515        self.render(false)
516    }
517
518    /// Renders the part for CLI display.
519    ///
520    /// # Errors
521    ///
522    /// Returns an error if JSON serialization fails.
523    pub fn render_for_cli(&self) -> Result<String> {
524        self.render(true)
525    }
526
527    fn render(&self, pretty: bool) -> Result<String> {
528        match self {
529            Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
530            Self::Json { value } => Ok(if pretty {
531                serde_json::to_string_pretty(value)?
532            } else {
533                serde_json::to_string(value)?
534            }),
535            Self::Binary { mime, content } => Ok(render_binary_placeholder(mime, content)),
536        }
537    }
538
539    /// Parses and returns the MIME type when this part carries binary content.
540    #[must_use]
541    pub fn mime(&self) -> Option<Mime> {
542        match self {
543            Self::Binary { mime, .. } => mime.parse().ok(),
544            Self::Text { .. } | Self::Tsv { .. } | Self::Json { .. } => None,
545        }
546    }
547
548    /// Returns raw bytes for binary parts.
549    #[must_use]
550    pub fn content(&self) -> Option<&[u8]> {
551        match self {
552            Self::Binary { content, .. } => Some(content),
553            Self::Text { .. } | Self::Tsv { .. } | Self::Json { .. } => None,
554        }
555    }
556
557    /// Lifts the part into its single-payload [`ToolResult`] equivalent.
558    #[must_use]
559    pub fn into_result(self) -> ToolResult {
560        match self {
561            Self::Text { text } | Self::Tsv { text } => ToolResult::Text { text },
562            Self::Json { value } => ToolResult::Json { value },
563            Self::Binary { mime, content } => ToolResult::Binary { mime, content },
564        }
565    }
566}
567
568fn render_binary_placeholder(mime: &str, content: &[u8]) -> String {
569    let mut rendered = String::new();
570    rendered.push_str("[binary tool result: ");
571    rendered.push_str(mime);
572    rendered.push_str(", ");
573    rendered.push_str(content.len().to_string().as_str());
574    rendered.push_str(" bytes]");
575    rendered
576}
577
578fn join_part_renders(parts: &[ToolResultPart], pretty: bool) -> Result<String> {
579    let mut rendered = String::new();
580    for part in parts {
581        let text = if pretty {
582            part.render_for_cli()?
583        } else {
584            part.render_for_model()?
585        };
586        if text.is_empty() {
587            continue;
588        }
589        if !rendered.is_empty() {
590            rendered.push('\n');
591        }
592        rendered.push_str(&text);
593    }
594    Ok(rendered)
595}
596
597/// Conversion trait for values returned by [`Tool::call`].
598///
599/// The conversion itself is fallible so tool authors can return types like
600/// `Result<T: Serialize, E: Error>` and still surface serialization failures as
601/// framework errors while preserving typed tool errors inside [`ToolResult`].
602pub trait IntoToolResult {
603    /// Converts the value into a final [`ToolResult`].
604    ///
605    /// # Errors
606    ///
607    /// Returns an error if the conversion cannot be completed.
608    fn into_tool_result(self) -> Result<ToolResult>;
609}
610
611impl IntoToolResult for ToolResult {
612    fn into_tool_result(self) -> Result<ToolResult> {
613        Ok(self)
614    }
615}
616
617impl IntoToolResult for () {
618    fn into_tool_result(self) -> Result<ToolResult> {
619        Ok(ToolResult::Done)
620    }
621}
622
623impl IntoToolResult for String {
624    fn into_tool_result(self) -> Result<ToolResult> {
625        Ok(ToolResult::text(self))
626    }
627}
628
629impl IntoToolResult for &str {
630    fn into_tool_result(self) -> Result<ToolResult> {
631        Ok(ToolResult::text(self))
632    }
633}
634
635impl IntoToolResult for Cow<'_, str> {
636    fn into_tool_result(self) -> Result<ToolResult> {
637        Ok(ToolResult::text(self.into_owned()))
638    }
639}
640
641impl IntoToolResult for Value {
642    fn into_tool_result(self) -> Result<ToolResult> {
643        Ok(ToolResult::json_value(self))
644    }
645}
646
647impl<T> IntoToolResult for Option<T>
648where
649    T: IntoToolResult,
650{
651    fn into_tool_result(self) -> Result<ToolResult> {
652        self.map_or_else(|| Ok(ToolResult::Done), IntoToolResult::into_tool_result)
653    }
654}
655
656impl IntoToolResult for Vec<ToolResultPart> {
657    fn into_tool_result(self) -> Result<ToolResult> {
658        Ok(ToolResult::parts(self))
659    }
660}
661
662impl<T, E> IntoToolResult for core::result::Result<T, E>
663where
664    T: Serialize,
665    E: Display,
666{
667    fn into_tool_result(self) -> Result<ToolResult> {
668        match self {
669            Ok(value) => serialize_success_value(&value),
670            Err(error) => Ok(ToolResult::error(error.to_string())),
671        }
672    }
673}
674
675fn parse_media_type_or_octet_stream(media_type: &str) -> String {
676    media_type
677        .parse::<Mime>()
678        .unwrap_or(mime::APPLICATION_OCTET_STREAM)
679        .essence_str()
680        .to_string()
681}
682
683fn serialize_success_value<T: Serialize>(value: &T) -> Result<ToolResult> {
684    let value = serde_json::to_value(value)?;
685    if let Some(tsv) = json_value_to_tsv(&value) {
686        return Ok(ToolResult::tsv(tsv));
687    }
688
689    match value {
690        Value::String(text) => Ok(ToolResult::text(text)),
691        other => Ok(ToolResult::json_value(other)),
692    }
693}
694
695/// Converts a JSON value into TSV when it represents an object or non-empty array.
696#[must_use]
697pub fn json_value_to_tsv(value: &Value) -> Option<String> {
698    let rows = match value {
699        Value::Array(arr) if !arr.is_empty() => arr
700            .iter()
701            .map(|value| flatten_json_value(value, ""))
702            .collect::<Vec<_>>(),
703        Value::Object(_) => alloc::vec![flatten_json_value(value, "")],
704        Value::Array(_) | Value::String(_) | Value::Number(_) | Value::Bool(_) | Value::Null => {
705            return None;
706        }
707    };
708
709    if rows.is_empty() {
710        return None;
711    }
712
713    let mut columns: Vec<String> = Vec::new();
714    let mut seen: alloc::collections::BTreeSet<String> = alloc::collections::BTreeSet::new();
715    for row in &rows {
716        for (key, _) in row {
717            if seen.insert(key.clone()) {
718                columns.push(key.clone());
719            }
720        }
721    }
722
723    if columns.is_empty() {
724        return None;
725    }
726
727    let mut tsv = String::new();
728    for (index, column) in columns.iter().enumerate() {
729        if index > 0 {
730            tsv.push('\t');
731        }
732        tsv.push_str(&escape_tsv_field(column));
733    }
734    tsv.push('\n');
735
736    for row in &rows {
737        let row_map: alloc::collections::BTreeMap<&str, &str> = row
738            .iter()
739            .map(|(key, value)| (key.as_str(), value.as_str()))
740            .collect::<alloc::collections::BTreeMap<&str, &str>>();
741        for (index, column) in columns.iter().enumerate() {
742            if index > 0 {
743                tsv.push('\t');
744            }
745            if let Some(value) = row_map.get(column.as_str()) {
746                tsv.push_str(&escape_tsv_field(value));
747            }
748        }
749        tsv.push('\n');
750    }
751
752    Some(tsv)
753}
754
755fn flatten_json_value(value: &Value, prefix: &str) -> Vec<(String, String)> {
756    let mut flattened = Vec::new();
757    match value {
758        Value::Object(map) => {
759            for (key, child) in map {
760                let full_key = if prefix.is_empty() {
761                    key.clone()
762                } else {
763                    format!("{prefix}.{key}")
764                };
765                flattened.extend(flatten_json_value(child, &full_key));
766            }
767        }
768        Value::Array(_) => {
769            let serialized = serde_json::to_string(value).unwrap_or_default();
770            flattened.push((prefix.to_string(), serialized));
771        }
772        Value::String(text) => {
773            flattened.push((prefix.to_string(), text.clone()));
774        }
775        Value::Number(number) => {
776            flattened.push((prefix.to_string(), number.to_string()));
777        }
778        Value::Bool(boolean) => {
779            flattened.push((prefix.to_string(), boolean.to_string()));
780        }
781        Value::Null => {
782            flattened.push((prefix.to_string(), String::new()));
783        }
784    }
785    flattened
786}
787
788fn escape_tsv_field(value: &str) -> String {
789    value.replace(['\t', '\n', '\r'], " ")
790}
791
792/// Tools that can be called by language models.
793///
794/// # Example
795///
796/// ```rust,ignore
797/// use aither::llm::{Tool, ToolResult};
798/// use schemars::JsonSchema;
799/// use serde::Deserialize;
800///
801/// #[derive(JsonSchema, Deserialize)]
802/// struct CalculatorArgs {
803///     operation: String,
804///     a: f64,
805///     b: f64,
806/// }
807///
808/// struct Calculator;
809///
810/// impl Tool for Calculator {
811///     type Arguments = CalculatorArgs;
812///     type Res = ToolResult;
813///
814///     async fn call(&mut self, args: Self::Arguments) -> aither::Result<Self::Res> {
815///         match args.operation.as_str() {
816///             "add" => Ok(ToolResult::text((args.a + args.b).to_string())),
817///             "subtract" => Ok(ToolResult::text((args.a - args.b).to_string())),
818///             "multiply" => Ok(ToolResult::text((args.a * args.b).to_string())),
819///             "divide" => {
820///                 if args.b != 0.0 {
821///                     Ok(ToolResult::text((args.a / args.b).to_string()))
822///                 } else {
823///                     Err(anyhow::Error::msg("Division by zero"))
824///                 }
825///             }
826///             _ => Err(anyhow::Error::msg("Unknown operation")),
827///         }
828///     }
829/// }
830/// ```
831pub trait Tool: Send + Sync {
832    /// Tool name. Must be unique.
833    fn name(&self) -> Cow<'static, str>;
834
835    /// What the tool does, as shown to the model.
836    ///
837    /// This is the single most important thing a model uses to decide whether
838    /// to call a tool, so it must not be empty — [`Tools::register`] rejects a
839    /// tool whose description is blank.
840    ///
841    /// The default implementation reads the rustdoc comment on
842    /// [`Self::Arguments`], which `schemars` records in the generated schema.
843    /// Override it to supply the description directly.
844    fn description(&self) -> Cow<'static, str> {
845        description_from_schema::<Self::Arguments>().unwrap_or_default()
846    }
847
848    /// Tool arguments type. Must implement [`schemars::JsonSchema`] and [`serde::de::DeserializeOwned`].
849    /// Its rustdoc becomes the default tool description.
850    type Arguments: Send + JsonSchema + DeserializeOwned;
851
852    /// Raw return type from the tool implementation.
853    type Res: IntoToolResult + Send;
854
855    /// Executes the tool with the provided arguments.
856    ///
857    /// Returns a value that can be converted into a final [`ToolResult`].
858    ///
859    /// Tools that need mutable state should use interior mutability (e.g., `Mutex`).
860    fn call(&self, arguments: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send;
861}
862
863/// Utility to convert a serializable value to a pretty-printed JSON string.
864///
865/// # Example
866/// ```rust,ignore
867/// use aither::llm::tool::json;
868/// use serde::Serialize;
869/// #[derive(Serialize)]
870/// struct Data {
871///     name: String,
872///     value: u32,
873/// }
874/// let data = Data {
875///     name: "example".to_string(),
876///     value: 42,
877/// };
878/// let json_str = json(&data);
879/// println!("{}", json_str);
880/// ```
881///
882/// # Errors
883///
884/// Returns an error if the value cannot be serialized to JSON, which happens
885/// for types such as maps with non-string keys.
886pub fn json<T: Serialize>(value: &T) -> Result<String> {
887    let value = serde_json::to_value(value)?;
888
889    Ok(value
890        .as_str()
891        .map_or_else(|| format!("{value:#}"), ToString::to_string))
892}
893
894trait ToolImpl: Send + Sync + Any {
895    fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>>;
896
897    /// The cached definition. Borrowed, so registering a tool does not have to
898    /// clone its argument schema.
899    fn definition(&self) -> &ToolDefinition;
900
901    /// Upcast to [`Any`] so [`Tools::get`] can recover the concrete tool.
902    ///
903    /// Casting the `Box<dyn ToolImpl>` itself would downcast the box rather
904    /// than the tool inside it, and so never match.
905    fn as_any(&self) -> &dyn Any;
906
907    /// Mutable counterpart of [`Self::as_any`].
908    fn as_any_mut(&mut self) -> &mut dyn Any;
909}
910
911/// Dynamic tool implementation for type-erased tools.
912struct DynToolImpl<F>
913where
914    F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>> + Send + Sync,
915{
916    definition: ToolDefinition,
917    handler: F,
918}
919
920impl<F> ToolImpl for DynToolImpl<F>
921where
922    F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>> + Send + Sync + 'static,
923{
924    fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>> {
925        (self.handler)(args)
926    }
927
928    fn definition(&self) -> &ToolDefinition {
929        &self.definition
930    }
931
932    fn as_any(&self) -> &dyn Any {
933        self
934    }
935
936    fn as_any_mut(&mut self) -> &mut dyn Any {
937        self
938    }
939}
940
941/// Whether a schema describes a JSON object, and so can be sent to a provider
942/// as tool arguments without being wrapped in [`ToolArgument`].
943fn schema_is_object(value: &Value) -> bool {
944    matches!(value.get("type").and_then(Value::as_str), Some("object"))
945        || value.get("properties").is_some()
946        || value.get("oneOf").is_some()
947        || value.get("anyOf").is_some()
948        || value.get("$defs").is_some()
949}
950
951fn is_object<T: JsonSchema>() -> bool {
952    schema_is_object(&schema_for!(T).to_value())
953}
954
955/// Builds the argument schema for a tool, wrapping scalars so the root is
956/// always an object as providers require.
957fn arguments_schema<T: JsonSchema>() -> Schema {
958    if is_object::<T>() {
959        schema_for!(T)
960    } else {
961        schema_for!(ToolArgument<T>)
962    }
963}
964
965/// Reads the `description` a `JsonSchema` derive records from a type's rustdoc.
966fn description_from_schema<T: JsonSchema>() -> Option<Cow<'static, str>> {
967    schema_for!(T)
968        .to_value()
969        .get("description")
970        .and_then(Value::as_str)
971        .filter(|text| !text.trim().is_empty())
972        .map(|text| Cow::Owned(text.to_string()))
973}
974
975/// A registered tool together with everything derived from its type.
976///
977/// The argument schema and the "are these arguments an object?" decision are
978/// properties of `T::Arguments` alone, so they are computed once here rather
979/// than rebuilt on every invocation.
980struct RegisteredTool<T: Tool> {
981    tool: T,
982    definition: ToolDefinition,
983    args_are_object: bool,
984}
985
986impl<T: Tool> RegisteredTool<T> {
987    fn new(tool: T) -> Self {
988        let definition = ToolDefinition::new(&tool);
989        let args_are_object = is_object::<T::Arguments>();
990        Self {
991            tool,
992            definition,
993            args_are_object,
994        }
995    }
996}
997
998impl<T: Tool + 'static> ToolImpl for RegisteredTool<T> {
999    fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>> {
1000        let result = if self.args_are_object {
1001            serde_json::from_str::<T::Arguments>(args)
1002        } else {
1003            serde_json::from_str::<ToolArgument<T::Arguments>>(args).map(|wrapper| wrapper.value)
1004        };
1005
1006        let Ok(arguments) = result else {
1007            // Cold path: spelling the schema out for the model is worth the
1008            // allocation only when it has actually got the call wrong.
1009            let name = self.definition.name().to_string();
1010            let schema_str =
1011                serde_json::to_string_pretty(&self.definition.arguments_openai_schema())
1012                    .unwrap_or_else(|_| "{}".to_string());
1013            return Box::pin(async move {
1014                Err(anyhow::Error::msg(format!(
1015                    "Invalid arguments for tool '{name}'. Expected schema:\n{schema_str}"
1016                )))
1017            });
1018        };
1019
1020        Box::pin(async move { Tool::call(&self.tool, arguments).await?.into_tool_result() })
1021    }
1022
1023    fn definition(&self) -> &ToolDefinition {
1024        &self.definition
1025    }
1026
1027    fn as_any(&self) -> &dyn Any {
1028        &self.tool
1029    }
1030
1031    fn as_any_mut(&mut self) -> &mut dyn Any {
1032        &mut self.tool
1033    }
1034}
1035
1036/// A tool definition carried something that is not a JSON schema.
1037///
1038/// JSON Schema allows an object or a bare boolean; anything else — a string, an
1039/// array, a number — describes nothing a model could fill in.
1040#[derive(Debug, Clone, PartialEq, Eq)]
1041pub struct InvalidSchema {
1042    /// The tool whose schema was rejected.
1043    name: Cow<'static, str>,
1044}
1045
1046impl InvalidSchema {
1047    /// The name of the tool whose schema was rejected.
1048    #[must_use]
1049    pub fn name(&self) -> &str {
1050        &self.name
1051    }
1052}
1053
1054impl Display for InvalidSchema {
1055    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1056        write!(
1057            f,
1058            "tool '{}' has an argument schema that is neither an object nor a boolean",
1059            self.name
1060        )
1061    }
1062}
1063
1064impl core::error::Error for InvalidSchema {}
1065
1066/// Why a tool could not be added to a [`Tools`] registry.
1067#[derive(Debug, Clone, PartialEq, Eq)]
1068pub enum RegisterError {
1069    /// A tool with this name is already registered.
1070    ///
1071    /// Names address tools in a model's tool-call, so they must be unique.
1072    DuplicateName(Cow<'static, str>),
1073
1074    /// The tool's description is empty.
1075    ///
1076    /// A description is what a model uses to decide whether to call a tool, so
1077    /// an empty one makes the tool unusable rather than merely undocumented.
1078    EmptyDescription(Cow<'static, str>),
1079}
1080
1081impl Display for RegisterError {
1082    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1083        match self {
1084            Self::DuplicateName(name) => {
1085                write!(f, "a tool named '{name}' is already registered")
1086            }
1087            Self::EmptyDescription(name) => write!(
1088                f,
1089                "tool '{name}' has an empty description; add a rustdoc comment to its \
1090                 Arguments type or implement Tool::description"
1091            ),
1092        }
1093    }
1094}
1095
1096impl core::error::Error for RegisterError {}
1097
1098/// Tool registry for managing and calling tools by name.
1099///
1100///
1101/// # Example
1102///
1103/// ```rust,ignore
1104/// use aither::llm::tool::Tools;
1105///
1106/// let mut tools = Tools::new();
1107/// // tools.register(Calculator);
1108/// let definitions = tools.definitions();
1109/// // let result = tools.call("calculator", r#"{"operation": "add", "a": 5, "b": 3}"#).await;
1110/// ```
1111pub struct Tools {
1112    tools: BTreeMap<Cow<'static, str>, Box<dyn ToolImpl>>,
1113}
1114
1115impl Debug for Tools {
1116    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1117        f.debug_struct("Tools")
1118            .field("tools", &self.tools.keys().collect::<Vec<_>>())
1119            .finish()
1120    }
1121}
1122
1123/// Tool definition including schema for language models.
1124///
1125/// Used to provide language models with information about available [`Tool`]s.
1126#[derive(Debug, Clone, PartialEq)]
1127#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
1128pub struct ToolDefinition {
1129    /// Tool name.
1130    name: Cow<'static, str>,
1131    /// Tool description.
1132    description: Cow<'static, str>,
1133    /// JSON schema for tool arguments.
1134    arguments: Schema,
1135}
1136
1137impl ToolDefinition {
1138    /// Creates a tool definition for a given tool type.
1139    ///
1140    /// The description comes from [`Tool::description`], which by default reads
1141    /// the rustdoc on the tool's `Arguments` type.
1142    #[must_use]
1143    pub fn new<T: Tool>(tool: &T) -> Self {
1144        Self {
1145            name: tool.name(),
1146            description: tool.description(),
1147            arguments: arguments_schema::<T::Arguments>(),
1148        }
1149    }
1150
1151    /// Creates a tool definition from raw parts.
1152    ///
1153    /// This is useful for creating definitions from external sources like MCP servers.
1154    ///
1155    /// # Errors
1156    ///
1157    /// Returns [`InvalidSchema`] if the value is not a JSON schema — that is,
1158    /// anything other than an object or a boolean. The schema often comes from
1159    /// a remote server, so this is a rejection to report, not a bug to panic
1160    /// on.
1161    pub fn from_parts(
1162        name: Cow<'static, str>,
1163        description: Cow<'static, str>,
1164        schema: Value,
1165    ) -> core::result::Result<Self, InvalidSchema> {
1166        let arguments: Schema = schema
1167            .try_into()
1168            .map_err(|_| InvalidSchema { name: name.clone() })?;
1169
1170        Ok(Self {
1171            name,
1172            description,
1173            arguments,
1174        })
1175    }
1176
1177    /// Returns the tool's name.
1178    #[must_use]
1179    pub fn name(&self) -> &str {
1180        &self.name
1181    }
1182
1183    /// Returns the tool's description.
1184    #[must_use]
1185    pub fn description(&self) -> &str {
1186        &self.description
1187    }
1188
1189    /// Return an OpenAI-compatible JSON schema for the tool's arguments.
1190    ///
1191    /// This schema would have an object type at the root, as required by `OpenAI`.
1192    #[must_use]
1193    pub fn arguments_openai_schema(&self) -> serde_json::Value {
1194        let mut inner = self.arguments.clone().to_value();
1195        clean_schema(&mut inner);
1196
1197        inner
1198    }
1199}
1200
1201#[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
1202struct ToolArgument<T> {
1203    value: T,
1204}
1205
1206fn clean_schema(value: &mut Value) {
1207    // First pass: extract $defs for reference resolution
1208    let defs = extract_defs(value);
1209
1210    // Second pass: resolve refs and clean
1211    resolve_and_clean(value, &defs);
1212
1213    // Clean up root-level schema
1214    if let Value::Object(map) = value {
1215        // Remove root-level description - it's already used as tool description
1216        // Keeping it duplicates the description in the API request
1217        map.remove("description");
1218
1219        // Ensure root has type: object (required by OpenAI function calling)
1220        if map.contains_key("properties") && !map.contains_key("type") {
1221            map.insert("type".to_string(), Value::String("object".to_string()));
1222        }
1223    }
1224}
1225
1226/// Extracts `$defs` or `definitions` from the root schema.
1227fn extract_defs(value: &Value) -> serde_json::Map<String, Value> {
1228    if let Value::Object(map) = value
1229        && let Some(Value::Object(defs)) = map.get("$defs").or_else(|| map.get("definitions"))
1230    {
1231        return defs.clone();
1232    }
1233    serde_json::Map::new()
1234}
1235
1236/// Resolves `$ref` and cleans the schema recursively.
1237#[allow(clippy::too_many_lines)]
1238fn resolve_and_clean(value: &mut Value, defs: &serde_json::Map<String, Value>) {
1239    resolve_and_clean_inner(value, defs, false);
1240}
1241
1242/// Inner recursive function with flag to track if we're inside a properties object.
1243#[allow(clippy::too_many_lines)]
1244fn resolve_and_clean_inner(
1245    value: &mut Value,
1246    defs: &serde_json::Map<String, Value>,
1247    inside_properties: bool,
1248) {
1249    match value {
1250        Value::Object(map) => {
1251            // Handle $ref - inline the referenced definition, preserving sibling properties
1252            if let Some(Value::String(ref_path)) = map.remove("$ref")
1253                && let Some(Value::Object(resolved_map)) = resolve_ref(&ref_path, defs)
1254            {
1255                // Merge resolved definition with any existing properties (like description)
1256                // Resolved definition takes precedence for conflicts except description
1257                let existing_description = map.remove("description");
1258                for (k, v) in resolved_map {
1259                    map.entry(k).or_insert(v);
1260                }
1261                // Preserve the field-level description if it exists
1262                if let Some(desc) = existing_description {
1263                    map.insert("description".to_string(), desc);
1264                }
1265            }
1266
1267            // Convert "const" to "enum" with single value (before filtering)
1268            if let Some(const_val) = map.remove("const") {
1269                map.insert("enum".to_string(), Value::Array(alloc::vec![const_val]));
1270            }
1271
1272            // Flatten oneOf/anyOf variants (before filtering, since oneOf is not in allowed list)
1273            if let Some(Value::Array(variants)) =
1274                map.remove("oneOf").or_else(|| map.remove("anyOf"))
1275            {
1276                // Check if this is a simple string enum (variants have const/type but no properties)
1277                let is_simple_enum = variants.iter().all(|v| {
1278                    if let Value::Object(vm) = v {
1279                        (vm.contains_key("const") || vm.contains_key("enum"))
1280                            && !vm.contains_key("properties")
1281                    } else {
1282                        false
1283                    }
1284                });
1285
1286                if is_simple_enum {
1287                    // Collect all const/enum values into a single enum array
1288                    let mut enum_values: alloc::vec::Vec<Value> = alloc::vec::Vec::new();
1289                    let mut variant_type: Option<String> = None;
1290
1291                    for variant in &variants {
1292                        if let Value::Object(vm) = variant {
1293                            if let Some(const_val) = vm.get("const")
1294                                && !enum_values.contains(const_val)
1295                            {
1296                                enum_values.push(const_val.clone());
1297                            }
1298                            if let Some(Value::Array(arr)) = vm.get("enum") {
1299                                for val in arr {
1300                                    if !enum_values.contains(val) {
1301                                        enum_values.push(val.clone());
1302                                    }
1303                                }
1304                            }
1305                            if variant_type.is_none()
1306                                && let Some(Value::String(t)) = vm.get("type")
1307                            {
1308                                variant_type = Some(t.clone());
1309                            }
1310                        }
1311                    }
1312
1313                    if !enum_values.is_empty() {
1314                        map.insert("enum".to_string(), Value::Array(enum_values));
1315                        if let Some(t) = variant_type {
1316                            map.insert("type".to_string(), Value::String(t));
1317                        }
1318                    }
1319                } else {
1320                    // Complex variants with properties - merge them
1321                    let mut all_properties = serde_json::Map::new();
1322
1323                    for variant in variants {
1324                        if let Value::Object(variant_map) = variant
1325                            && let Some(Value::Object(props)) = variant_map.get("properties")
1326                        {
1327                            for (key, val) in props {
1328                                // Extract enum value - handle both "enum" and "const"
1329                                let new_values: Option<alloc::vec::Vec<Value>> =
1330                                    if let Value::Object(val_obj) = val {
1331                                        if let Some(Value::Array(arr)) = val_obj.get("enum") {
1332                                            Some(arr.clone())
1333                                        } else {
1334                                            val_obj
1335                                                .get("const")
1336                                                .map(|const_val| alloc::vec![const_val.clone()])
1337                                        }
1338                                    } else {
1339                                        None
1340                                    };
1341
1342                                if all_properties.contains_key(key) {
1343                                    // Merge enum/const values into existing
1344                                    if let Some(values) = new_values
1345                                        && let Some(Value::Object(existing_obj)) =
1346                                            all_properties.get_mut(key)
1347                                        && let Some(Value::Array(existing_enum)) =
1348                                            existing_obj.get_mut("enum")
1349                                    {
1350                                        for e in values {
1351                                            if !existing_enum.contains(&e) {
1352                                                existing_enum.push(e);
1353                                            }
1354                                        }
1355                                    }
1356                                } else {
1357                                    // First time seeing this property - convert const to enum
1358                                    let mut val_clone = val.clone();
1359                                    if let Value::Object(obj) = &mut val_clone
1360                                        && let Some(const_val) = obj.remove("const")
1361                                    {
1362                                        obj.insert(
1363                                            "enum".to_string(),
1364                                            Value::Array(alloc::vec![const_val]),
1365                                        );
1366                                    }
1367                                    all_properties.insert(key.clone(), val_clone);
1368                                }
1369                            }
1370                        }
1371                    }
1372
1373                    // Set type as object if we have properties
1374                    if !all_properties.is_empty() {
1375                        map.insert("type".to_string(), Value::String("object".to_string()));
1376                        map.insert("properties".to_string(), Value::Object(all_properties));
1377                    }
1378                }
1379            }
1380
1381            // Only filter schema keywords, not property names inside "properties"
1382            // OpenAPI schema subset supported by most LLM providers
1383            if !inside_properties {
1384                let allowed = [
1385                    "type",
1386                    "description",
1387                    "properties",
1388                    "required",
1389                    "items",
1390                    "enum",
1391                    "nullable",
1392                ];
1393                map.retain(|k, _| allowed.contains(&k.as_str()));
1394            }
1395
1396            // Simplify "type" arrays like ["string", "null"] to single type
1397            if let Some(Value::Array(types)) = map.get("type") {
1398                // Filter out "null" and take the first non-null type
1399                let non_null: Vec<&Value> = types
1400                    .iter()
1401                    .filter(|t| !matches!(t, Value::String(s) if s == "null"))
1402                    .collect();
1403                if non_null.len() == 1 {
1404                    map.insert("type".to_string(), non_null[0].clone());
1405                }
1406            }
1407
1408            // Recursively clean all values
1409            for (key, v) in map.iter_mut() {
1410                // When entering "properties", its children are property definitions
1411                let child_inside_props = key == "properties";
1412                resolve_and_clean_inner(v, defs, child_inside_props);
1413            }
1414        }
1415        Value::Array(arr) => {
1416            for v in arr {
1417                resolve_and_clean_inner(v, defs, false);
1418            }
1419        }
1420        _ => {}
1421    }
1422}
1423
1424/// Resolves a `$ref` path like `#/$defs/FsOperation` to its definition.
1425fn resolve_ref(ref_path: &str, defs: &serde_json::Map<String, Value>) -> Option<Value> {
1426    // Handle common patterns: #/$defs/Name or #/definitions/Name
1427    let name = ref_path
1428        .strip_prefix("#/$defs/")
1429        .or_else(|| ref_path.strip_prefix("#/definitions/"))?;
1430
1431    defs.get(name).cloned()
1432}
1433
1434impl Default for Tools {
1435    fn default() -> Self {
1436        Self::new()
1437    }
1438}
1439
1440impl Tools {
1441    /// Creates a new empty tools registry.
1442    #[must_use]
1443    pub const fn new() -> Self {
1444        Self {
1445            tools: BTreeMap::new(),
1446        }
1447    }
1448
1449    /// Retrieves a tool by type.
1450    ///
1451    /// Returns `None` if the tool is not found.
1452    #[must_use]
1453    pub fn get<T>(&self) -> Option<&T>
1454    where
1455        T: Tool + 'static,
1456    {
1457        self.tools
1458            .values()
1459            .find_map(|tool| tool.as_any().downcast_ref::<T>())
1460    }
1461
1462    /// Retrieves a mutable reference to a tool by type.
1463    ///
1464    /// Returns `None` if the tool is not found.
1465    #[must_use]
1466    pub fn get_mut<T>(&mut self) -> Option<&mut T>
1467    where
1468        T: Tool + 'static,
1469    {
1470        self.tools
1471            .values_mut()
1472            .find_map(|tool| tool.as_any_mut().downcast_mut::<T>())
1473    }
1474
1475    /// Returns definitions of all registered tools.
1476    #[must_use]
1477    pub fn definitions(&self) -> Vec<ToolDefinition> {
1478        self.tools
1479            .values()
1480            .map(|tool| tool.definition().clone())
1481            .collect()
1482    }
1483
1484    /// Registers a new tool.
1485    ///
1486    /// The tool must implement [`Tool`] and be `'static`.
1487    ///
1488    /// # Errors
1489    ///
1490    /// Returns [`RegisterError::DuplicateName`] if a tool of that name is
1491    /// already registered, or [`RegisterError::EmptyDescription`] if the tool
1492    /// has no description — a model cannot use a tool it cannot read about, so
1493    /// this is rejected rather than silently passed on.
1494    pub fn register<T: Tool + 'static>(
1495        &mut self,
1496        tool: T,
1497    ) -> core::result::Result<(), RegisterError> {
1498        self.insert(Box::new(RegisteredTool::new(tool)))
1499    }
1500
1501    /// Registers a dynamic tool with a pre-made definition and handler.
1502    ///
1503    /// This is useful for type-erased tools (e.g., child terminal tools for subagents)
1504    /// where the concrete type isn't known at compile time.
1505    ///
1506    /// # Errors
1507    ///
1508    /// Same conditions as [`Self::register`].
1509    pub fn register_dyn<F>(
1510        &mut self,
1511        definition: ToolDefinition,
1512        handler: F,
1513    ) -> core::result::Result<(), RegisterError>
1514    where
1515        F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>>
1516            + Send
1517            + Sync
1518            + 'static,
1519    {
1520        self.insert(Box::new(DynToolImpl {
1521            definition,
1522            handler,
1523        }))
1524    }
1525
1526    fn insert(&mut self, tool: Box<dyn ToolImpl>) -> core::result::Result<(), RegisterError> {
1527        let name = tool.definition().name.clone();
1528        if self.tools.contains_key(&name) {
1529            return Err(RegisterError::DuplicateName(name));
1530        }
1531        if tool.definition().description().trim().is_empty() {
1532            return Err(RegisterError::EmptyDescription(name));
1533        }
1534        self.tools.insert(name, tool);
1535        Ok(())
1536    }
1537
1538    /// Removes a tool from the registry.
1539    pub fn unregister(&mut self, name: &str) {
1540        self.tools.remove(name);
1541    }
1542
1543    /// Calls a tool by name with JSON arguments.
1544    ///
1545    /// # Errors
1546    ///
1547    /// Returns an error if the tool is not found, arguments cannot be parsed,
1548    /// or tool execution fails.
1549    pub async fn call(&self, name: &str, args: &str) -> Result<ToolResult> {
1550        if let Some(tool) = self.tools.get(name) {
1551            tool.call(args).await
1552        } else {
1553            Err(anyhow::Error::msg(format!("Tool '{name}' not found")))
1554        }
1555    }
1556}
1557
1558#[cfg(test)]
1559mod tests {
1560    use super::*;
1561    use alloc::{format, string::ToString, vec};
1562    use schemars::JsonSchema;
1563    use serde::{Deserialize, Serialize};
1564
1565    /// Performs basic mathematical operations.
1566    #[derive(JsonSchema, Deserialize, Debug, PartialEq)]
1567    struct CalculatorArgs {
1568        operation: String,
1569        a: f64,
1570        b: f64,
1571    }
1572
1573    struct Calculator;
1574
1575    impl Tool for Calculator {
1576        fn name(&self) -> Cow<'static, str> {
1577            "calculator".into()
1578        }
1579        type Arguments = CalculatorArgs;
1580        type Res = ToolResult;
1581
1582        fn call(&self, args: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send {
1583            core::future::ready(match args.operation.as_str() {
1584                "add" => Ok(ToolResult::text((args.a + args.b).to_string())),
1585                "subtract" => Ok(ToolResult::text((args.a - args.b).to_string())),
1586                "multiply" => Ok(ToolResult::text((args.a * args.b).to_string())),
1587                "divide" => {
1588                    if args.b == 0.0 {
1589                        Err(anyhow::Error::msg("Division by zero"))
1590                    } else {
1591                        Ok(ToolResult::text((args.a / args.b).to_string()))
1592                    }
1593                }
1594                _ => Err(anyhow::Error::msg(format!(
1595                    "Unknown operation: {}",
1596                    args.operation
1597                ))),
1598            })
1599        }
1600    }
1601
1602    /// Greets a person by name.
1603    #[derive(JsonSchema, Deserialize)]
1604    struct GreetArgs {
1605        name: String,
1606    }
1607
1608    struct Greeter;
1609
1610    impl Tool for Greeter {
1611        fn name(&self) -> Cow<'static, str> {
1612            "greeter".into()
1613        }
1614        type Arguments = GreetArgs;
1615        type Res = ToolResult;
1616
1617        fn call(&self, args: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send {
1618            core::future::ready(Ok(ToolResult::text(format!("Hello, {}!", args.name))))
1619        }
1620    }
1621
1622    #[test]
1623    fn from_parts_accepts_object_and_boolean_schemas() {
1624        // JSON Schema allows a bare boolean as well as an object.
1625        for schema in [
1626            serde_json::json!({"type": "object"}),
1627            serde_json::json!(true),
1628        ] {
1629            assert!(
1630                ToolDefinition::from_parts("t".into(), "does a thing".into(), schema.clone())
1631                    .is_ok(),
1632                "{schema} should be accepted"
1633            );
1634        }
1635    }
1636
1637    #[test]
1638    fn from_parts_rejects_non_schema_values() {
1639        // A server that sends any of these describes nothing a model could
1640        // fill in. Rejecting must not panic: the value came off the wire.
1641        for schema in [
1642            serde_json::json!("a string"),
1643            serde_json::json!([1, 2, 3]),
1644            serde_json::json!(7),
1645            serde_json::json!(null),
1646        ] {
1647            let result = ToolDefinition::from_parts("weird".into(), "d".into(), schema.clone());
1648            let Err(err) = result else {
1649                panic!("{schema} should be rejected");
1650            };
1651            assert_eq!(err.name(), "weird");
1652        }
1653    }
1654
1655    #[test]
1656    fn json_utility() {
1657        let value = serde_json::json!({
1658            "name": "test",
1659            "value": 42
1660        });
1661
1662        let json_str = json(&value).expect("a JSON value always serializes");
1663        assert!(json_str.contains("\"name\": \"test\""));
1664        assert!(json_str.contains("\"value\": 42"));
1665    }
1666
1667    #[test]
1668    fn tool_definition_creation() {
1669        let calculator = Calculator;
1670        let definition = ToolDefinition::new(&calculator);
1671
1672        assert_eq!(definition.name, "calculator");
1673        assert_eq!(
1674            definition.description,
1675            "Performs basic mathematical operations."
1676        );
1677        // Schema should be present - just check it exists
1678        // The exact structure of schemars::Schema is implementation detail
1679    }
1680
1681    #[test]
1682    fn tools_creation() {
1683        let tools = Tools::new();
1684        assert_eq!(tools.definitions().len(), 0);
1685    }
1686
1687    #[test]
1688    fn tools_default() {
1689        let tools = Tools::default();
1690        assert_eq!(tools.definitions().len(), 0);
1691    }
1692
1693    #[tokio::test]
1694    async fn tools_register_and_call() {
1695        let mut tools = Tools::new();
1696        tools.register(Calculator).expect("calculator registers");
1697
1698        let definitions = tools.definitions();
1699        assert_eq!(definitions.len(), 1);
1700        assert_eq!(definitions[0].name, "calculator");
1701
1702        let result = tools
1703            .call("calculator", r#"{"operation": "add", "a": 5, "b": 3}"#)
1704            .await;
1705        assert!(result.is_ok());
1706        assert_eq!(result.unwrap().as_text(), Some("8"));
1707    }
1708
1709    #[tokio::test]
1710    async fn calculator_operations() {
1711        let mut tools = Tools::new();
1712        tools.register(Calculator).expect("calculator registers");
1713
1714        // Test addition
1715        let result = tools
1716            .call("calculator", r#"{"operation": "add", "a": 10, "b": 5}"#)
1717            .await;
1718        assert_eq!(result.unwrap().as_text(), Some("15"));
1719
1720        // Test subtraction
1721        let result = tools
1722            .call(
1723                "calculator",
1724                r#"{"operation": "subtract", "a": 10, "b": 3}"#,
1725            )
1726            .await;
1727        assert_eq!(result.unwrap().as_text(), Some("7"));
1728
1729        // Test multiplication
1730        let result = tools
1731            .call("calculator", r#"{"operation": "multiply", "a": 4, "b": 3}"#)
1732            .await;
1733        assert_eq!(result.unwrap().as_text(), Some("12"));
1734
1735        // Test division
1736        let result = tools
1737            .call("calculator", r#"{"operation": "divide", "a": 15, "b": 3}"#)
1738            .await;
1739        assert_eq!(result.unwrap().as_text(), Some("5"));
1740    }
1741
1742    #[tokio::test]
1743    async fn calculator_division_by_zero() {
1744        let mut tools = Tools::new();
1745        tools.register(Calculator).expect("calculator registers");
1746
1747        let result = tools
1748            .call("calculator", r#"{"operation": "divide", "a": 10, "b": 0}"#)
1749            .await;
1750        assert!(result.is_err());
1751        assert!(result.unwrap_err().to_string().contains("Division by zero"));
1752    }
1753
1754    #[tokio::test]
1755    async fn calculator_unknown_operation() {
1756        let mut tools = Tools::new();
1757        tools.register(Calculator).expect("calculator registers");
1758
1759        let result = tools
1760            .call("calculator", r#"{"operation": "modulo", "a": 10, "b": 3}"#)
1761            .await;
1762        assert!(result.is_err());
1763        assert!(
1764            result
1765                .unwrap_err()
1766                .to_string()
1767                .contains("Unknown operation")
1768        );
1769    }
1770
1771    #[tokio::test]
1772    async fn multiple_tools() {
1773        let mut tools = Tools::new();
1774        tools.register(Calculator).expect("calculator registers");
1775        tools.register(Greeter).expect("greeter registers");
1776
1777        let definitions = tools.definitions();
1778        assert_eq!(definitions.len(), 2);
1779
1780        // Find calculator and greeter in definitions
1781        let calc_def = definitions.iter().find(|d| d.name == "calculator").unwrap();
1782        let greet_def = definitions.iter().find(|d| d.name == "greeter").unwrap();
1783
1784        assert_eq!(
1785            calc_def.description,
1786            "Performs basic mathematical operations."
1787        );
1788        assert_eq!(greet_def.description, "Greets a person by name.");
1789
1790        // Test both tools
1791        let calc_result = tools
1792            .call("calculator", r#"{"operation": "add", "a": 2, "b": 3}"#)
1793            .await;
1794        assert_eq!(calc_result.unwrap().as_text(), Some("5"));
1795
1796        let greet_result = tools.call("greeter", r#"{"name": "Alice"}"#).await;
1797        assert_eq!(greet_result.unwrap().as_text(), Some("Hello, Alice!"));
1798    }
1799
1800    #[derive(Debug, Serialize)]
1801    struct TableRow {
1802        name: &'static str,
1803        count: u32,
1804    }
1805
1806    #[derive(Debug, Serialize)]
1807    struct NestedTableRow {
1808        user: TableRow,
1809        ok: bool,
1810    }
1811
1812    #[derive(Debug)]
1813    struct ToolFailure(&'static str);
1814
1815    impl core::fmt::Display for ToolFailure {
1816        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1817            f.write_str(self.0)
1818        }
1819    }
1820
1821    impl core::error::Error for ToolFailure {}
1822
1823    #[test]
1824    fn into_tool_result_string_is_plain_text() {
1825        assert_eq!(
1826            String::from("hello").into_tool_result().unwrap(),
1827            ToolResult::text("hello")
1828        );
1829    }
1830
1831    #[test]
1832    fn into_tool_result_str_is_plain_text() {
1833        assert_eq!(
1834            "hello".into_tool_result().unwrap(),
1835            ToolResult::text("hello")
1836        );
1837    }
1838
1839    #[test]
1840    fn into_tool_result_option_none_is_done() {
1841        let result = Option::<String>::None.into_tool_result().unwrap();
1842        assert_eq!(result, ToolResult::Done);
1843    }
1844
1845    #[test]
1846    fn into_tool_result_option_some_delegates() {
1847        let result = Some("hello").into_tool_result().unwrap();
1848        assert_eq!(result, ToolResult::text("hello"));
1849    }
1850
1851    #[test]
1852    fn into_tool_result_result_ok_string_is_text() {
1853        let result = core::result::Result::<String, ToolFailure>::Ok(String::from("hello"))
1854            .into_tool_result()
1855            .unwrap();
1856        assert_eq!(result, ToolResult::text("hello"));
1857    }
1858
1859    #[test]
1860    fn into_tool_result_result_ok_object_is_tsv() {
1861        let result = core::result::Result::<TableRow, ToolFailure>::Ok(TableRow {
1862            name: "alpha",
1863            count: 3,
1864        })
1865        .into_tool_result()
1866        .unwrap();
1867
1868        assert_eq!(result, ToolResult::tsv("count\tname\n3\talpha\n"));
1869    }
1870
1871    #[test]
1872    fn into_tool_result_result_ok_array_of_objects_is_tsv() {
1873        let result = core::result::Result::<Vec<TableRow>, ToolFailure>::Ok(vec![
1874            TableRow {
1875                name: "alpha",
1876                count: 3,
1877            },
1878            TableRow {
1879                name: "beta",
1880                count: 5,
1881            },
1882        ])
1883        .into_tool_result()
1884        .unwrap();
1885
1886        assert_eq!(result, ToolResult::tsv("count\tname\n3\talpha\n5\tbeta\n"));
1887    }
1888
1889    #[test]
1890    fn into_tool_result_result_ok_scalar_is_json() {
1891        let result = core::result::Result::<bool, ToolFailure>::Ok(true)
1892            .into_tool_result()
1893            .unwrap();
1894        assert_eq!(result, ToolResult::json_value(Value::Bool(true)));
1895    }
1896
1897    #[test]
1898    fn into_tool_result_result_err_is_typed_error() {
1899        let result = core::result::Result::<TableRow, ToolFailure>::Err(ToolFailure("boom"))
1900            .into_tool_result()
1901            .unwrap();
1902        assert_eq!(result, ToolResult::error("boom"));
1903        assert!(result.is_error());
1904        assert_eq!(result.error_message(), Some("boom"));
1905    }
1906
1907    #[test]
1908    fn json_value_to_tsv_flattens_nested_objects() {
1909        let value = serde_json::to_value(NestedTableRow {
1910            user: TableRow {
1911                name: "alpha",
1912                count: 3,
1913            },
1914            ok: true,
1915        })
1916        .unwrap();
1917
1918        assert_eq!(
1919            json_value_to_tsv(&value),
1920            Some("ok\tuser.count\tuser.name\ntrue\t3\talpha\n".to_string())
1921        );
1922    }
1923
1924    #[tokio::test]
1925    async fn tool_not_found() {
1926        let tools = Tools::new();
1927
1928        let result = tools.call("nonexistent", "{}").await;
1929        assert!(result.is_err());
1930        assert!(
1931            result
1932                .unwrap_err()
1933                .to_string()
1934                .contains("Tool 'nonexistent' not found")
1935        );
1936    }
1937
1938    #[tokio::test]
1939    async fn invalid_json() {
1940        let mut tools = Tools::new();
1941        tools.register(Calculator).expect("calculator registers");
1942
1943        let result = tools.call("calculator", "invalid json").await;
1944        assert!(result.is_err());
1945    }
1946
1947    #[test]
1948    fn tools_unregister() {
1949        let mut tools = Tools::new();
1950        tools.register(Calculator).expect("calculator registers");
1951        tools.register(Greeter).expect("greeter registers");
1952
1953        assert_eq!(tools.definitions().len(), 2);
1954
1955        tools.unregister("calculator");
1956        assert_eq!(tools.definitions().len(), 1);
1957
1958        let remaining = &tools.definitions()[0];
1959        assert_eq!(remaining.name, "greeter");
1960
1961        tools.unregister("greeter");
1962        assert_eq!(tools.definitions().len(), 0);
1963    }
1964
1965    #[test]
1966    fn tools_debug() {
1967        let mut tools = Tools::new();
1968        tools.register(Calculator).expect("calculator registers");
1969        tools.register(Greeter).expect("greeter registers");
1970
1971        let debug_str = format!("{tools:?}");
1972        assert!(debug_str.contains("Tools"));
1973        assert!(debug_str.contains("calculator"));
1974        assert!(debug_str.contains("greeter"));
1975    }
1976
1977    #[test]
1978    fn tool_definition_debug() {
1979        let calculator = Calculator;
1980        let definition = ToolDefinition::new(&calculator);
1981        let debug_str = format!("{definition:?}");
1982
1983        assert!(debug_str.contains("ToolDefinition"));
1984        assert!(debug_str.contains("calculator"));
1985        assert!(debug_str.contains("Performs basic mathematical operations"));
1986    }
1987
1988    #[test]
1989    fn tool_definition_clone() {
1990        let calculator = Calculator;
1991        let original = ToolDefinition::new(&calculator);
1992        let cloned = original.clone();
1993
1994        assert_eq!(original.name, cloned.name);
1995        assert_eq!(original.description, cloned.description);
1996    }
1997
1998    #[test]
1999    fn schema_preserves_enum() {
2000        #[derive(JsonSchema, Deserialize)]
2001        #[serde(rename_all = "snake_case")]
2002        enum Status {
2003            Pending,
2004            InProgress,
2005            Completed,
2006        }
2007
2008        #[allow(dead_code)]
2009        #[derive(JsonSchema, Deserialize)]
2010        struct Item {
2011            status: Status,
2012        }
2013
2014        #[allow(dead_code)]
2015        #[derive(JsonSchema, Deserialize)]
2016        struct Args {
2017            items: Vec<Item>,
2018        }
2019
2020        struct TestTool;
2021
2022        impl Tool for TestTool {
2023            fn name(&self) -> Cow<'static, str> {
2024                "test".into()
2025            }
2026            type Arguments = Args;
2027            type Res = ToolResult;
2028
2029            fn call(
2030                &self,
2031                _args: Self::Arguments,
2032            ) -> impl Future<Output = Result<Self::Res>> + Send {
2033                core::future::ready(Ok(ToolResult::text("ok")))
2034            }
2035        }
2036
2037        let tool = TestTool;
2038        let def = ToolDefinition::new(&tool);
2039        let schema = def.arguments_openai_schema();
2040
2041        // Check that status has enum values
2042        let schema_obj = schema.as_object().expect("schema should be object");
2043        let properties = schema_obj
2044            .get("properties")
2045            .expect("should have properties")
2046            .as_object()
2047            .unwrap();
2048        let items = properties
2049            .get("items")
2050            .expect("should have items")
2051            .as_object()
2052            .unwrap();
2053        let item_props = items
2054            .get("items")
2055            .expect("items should have items schema")
2056            .as_object()
2057            .unwrap();
2058        let item_properties = item_props
2059            .get("properties")
2060            .expect("item should have properties")
2061            .as_object()
2062            .unwrap();
2063        let status = item_properties
2064            .get("status")
2065            .expect("should have status")
2066            .as_object()
2067            .unwrap();
2068
2069        // Status should have enum
2070        assert!(
2071            status.contains_key("enum"),
2072            "Status should have enum field. Full schema: {}",
2073            serde_json::to_string_pretty(&schema).unwrap()
2074        );
2075    }
2076
2077    #[test]
2078    fn schema_ref_resolution() {
2079        // Test that $ref schemas get resolved and enum values preserved
2080        let raw_schema = serde_json::json!({
2081            "type": "object",
2082            "properties": {
2083                "status": {
2084                    "$ref": "#/$defs/Status"
2085                }
2086            },
2087            "$defs": {
2088                "Status": {
2089                    "type": "string",
2090                    "enum": ["pending", "in_progress", "completed"]
2091                }
2092            }
2093        });
2094
2095        let mut schema = raw_schema;
2096        clean_schema(&mut schema);
2097
2098        let props = schema.get("properties").unwrap().as_object().unwrap();
2099        let status = props.get("status").unwrap().as_object().unwrap();
2100
2101        assert!(
2102            status.contains_key("enum"),
2103            "Status should have enum after ref resolution. Got: {}",
2104            serde_json::to_string_pretty(&schema).unwrap()
2105        );
2106    }
2107
2108    #[test]
2109    fn schema_nested_ref_in_array() {
2110        // Test nested $ref inside array items (like TodoWriteArgs)
2111        let raw_schema = serde_json::json!({
2112            "type": "object",
2113            "properties": {
2114                "todos": {
2115                    "type": "array",
2116                    "items": {
2117                        "$ref": "#/$defs/TodoItem"
2118                    }
2119                }
2120            },
2121            "$defs": {
2122                "TodoItem": {
2123                    "type": "object",
2124                    "properties": {
2125                        "content": { "type": "string" },
2126                        "status": { "$ref": "#/$defs/TodoStatus" }
2127                    },
2128                    "required": ["content", "status"]
2129                },
2130                "TodoStatus": {
2131                    "type": "string",
2132                    "enum": ["pending", "in_progress", "completed"]
2133                }
2134            }
2135        });
2136
2137        let mut schema = raw_schema;
2138        clean_schema(&mut schema);
2139
2140        // Navigate to status
2141        let props = schema.get("properties").unwrap().as_object().unwrap();
2142        let todos = props.get("todos").unwrap().as_object().unwrap();
2143        let items = todos.get("items").unwrap().as_object().unwrap();
2144        let item_props = items.get("properties").unwrap().as_object().unwrap();
2145        let status = item_props.get("status").unwrap().as_object().unwrap();
2146
2147        assert!(
2148            status.contains_key("enum"),
2149            "Nested status should have enum. Full schema: {}",
2150            serde_json::to_string_pretty(&schema).unwrap()
2151        );
2152    }
2153
2154    #[test]
2155    fn parts_collapses_degenerate_inputs() {
2156        assert_eq!(ToolResult::parts(vec![]), ToolResult::Done);
2157        assert_eq!(
2158            ToolResult::parts(vec![ToolResultPart::text("only")]),
2159            ToolResult::text("only")
2160        );
2161        let multi = ToolResult::parts(vec![
2162            ToolResultPart::text("a"),
2163            ToolResultPart::image(vec![1, 2], "image/png"),
2164        ]);
2165        let ToolResult::Parts { parts } = &multi else {
2166            panic!("two parts must stay a Parts result");
2167        };
2168        assert_eq!(parts.len(), 2);
2169    }
2170
2171    #[test]
2172    fn parts_render_joins_text_and_marks_binary() {
2173        let result = ToolResult::parts(vec![
2174            ToolResultPart::text("header"),
2175            ToolResultPart::json_value(serde_json::json!({"n": 1})),
2176            ToolResultPart::image(vec![0x89], "image/png"),
2177        ]);
2178        assert_eq!(
2179            result.render_for_model().unwrap(),
2180            "header\n{\"n\":1}\n[binary tool result: image/png, 1 bytes]"
2181        );
2182        assert!(result.mime().is_none() && result.content().is_none());
2183    }
2184
2185    #[test]
2186    fn parts_with_one_binary_expose_mime_and_content() {
2187        let result = ToolResult::Parts {
2188            parts: vec![ToolResultPart::image(vec![9, 9], "image/png")],
2189        };
2190        assert_eq!(result.mime().unwrap().essence_str(), "image/png");
2191        assert_eq!(result.content().unwrap(), &[9, 9]);
2192    }
2193
2194    #[cfg(feature = "serde")]
2195    #[test]
2196    fn parts_serde_roundtrip_uses_kind_tag() {
2197        let result = ToolResult::parts(vec![
2198            ToolResultPart::text("note"),
2199            ToolResultPart::json_value(serde_json::json!({"k": true})),
2200        ]);
2201        let value = serde_json::to_value(&result).unwrap();
2202        assert_eq!(value["kind"], "parts");
2203        assert_eq!(value["parts"][0]["kind"], "text");
2204        let back: ToolResult = serde_json::from_value(value).unwrap();
2205        assert_eq!(back, result);
2206    }
2207}