wasm_agent/
tools.rs

1// SPDX-License-Identifier: MIT
2//! Tool registry and dispatch for WASM-safe function signatures.
3
4use std::collections::{BTreeMap, HashMap};
5use crate::error::AgentError;
6use crate::types::ToolResult;
7
8/// A tool handler: takes a JSON string input, returns a [`ToolResult`].
9///
10/// Uses `Box<dyn Fn>` for WASM-compatible dispatch: no async, no `Send`
11/// requirement in single-threaded WASM environments.
12pub type ToolFn = Box<dyn Fn(&str) -> ToolResult>;
13
14/// Metadata describing a tool for system-prompt generation and introspection.
15#[derive(Debug, Clone)]
16pub struct ToolSpec {
17    /// Unique name used to invoke the tool from the ReAct loop.
18    pub name: String,
19    /// Human-readable description included in the agent's system prompt.
20    pub description: String,
21    /// JSON Schema string describing the tool's input format.
22    pub input_schema: String,
23}
24
25impl ToolSpec {
26    /// Creates a new [`ToolSpec`].
27    ///
28    /// # Arguments
29    /// * `name`: Unique tool name (must be non-empty).
30    /// * `description`: Human-readable description.
31    /// * `schema`: JSON Schema for the tool's input.
32    pub fn new(name: impl Into<String>, description: impl Into<String>, schema: impl Into<String>) -> Self {
33        Self { name: name.into(), description: description.into(), input_schema: schema.into() }
34    }
35}
36
37/// Registry of available tools, keyed by name.
38pub struct ToolRegistry {
39    // Sorted, so tools_prompt() is identical on every run (a prompt that
40    // changes order between runs defeats provider prompt caching).
41    specs: BTreeMap<String, ToolSpec>,
42    handlers: HashMap<String, ToolFn>,
43}
44
45impl ToolRegistry {
46    /// Creates an empty registry.
47    pub fn new() -> Self {
48        Self { specs: BTreeMap::new(), handlers: HashMap::new() }
49    }
50
51    /// Registers a tool with its specification and handler.
52    ///
53    /// # Errors
54    /// Returns [`AgentError::InvalidToolSignature`] if the tool name is empty.
55    pub fn register(&mut self, spec: ToolSpec, handler: ToolFn) -> Result<(), AgentError> {
56        if spec.name.is_empty() {
57            return Err(AgentError::InvalidToolSignature("tool name cannot be empty".into()));
58        }
59        self.specs.insert(spec.name.clone(), spec.clone());
60        self.handlers.insert(spec.name, handler);
61        Ok(())
62    }
63
64    /// Dispatches a tool call by name with the given JSON input string.
65    ///
66    /// # Errors
67    /// Returns [`AgentError::ToolNotFound`] if no tool with that name is registered.
68    pub fn dispatch(&self, tool_name: &str, input: &str) -> Result<ToolResult, AgentError> {
69        let handler = self.handlers.get(tool_name)
70            .ok_or_else(|| AgentError::ToolNotFound { name: tool_name.to_string() })?;
71        Ok(handler(input))
72    }
73
74    /// Dispatches a native tool call (see [`crate::native`]). Its arguments
75    /// arrive as JSON. For a tool with an object schema they are passed on
76    /// unchanged; for a tool with any other schema the API saw
77    /// `{"input": ...}`, so that value is unwrapped (a JSON string is passed as
78    /// plain text). Arguments that are not valid JSON produce a failed
79    /// [`ToolResult`] the model can see and correct, not an error.
80    ///
81    /// # Errors
82    /// [`AgentError::ToolNotFound`] if no tool with that name is registered.
83    pub fn dispatch_call(&self, call: &crate::native::ToolCall) -> Result<ToolResult, AgentError> {
84        let spec = self.specs.get(&call.name).ok_or_else(|| AgentError::ToolNotFound { name: call.name.clone() })?;
85        let wrapped = crate::native::api_schema(&spec.input_schema).map(|(_, w)| w).unwrap_or(true);
86        if !wrapped {
87            return self.dispatch(&call.name, &call.arguments);
88        }
89        let input = match serde_json::from_str::<serde_json::Value>(&call.arguments) {
90            Ok(v) => match v.get("input") {
91                Some(serde_json::Value::String(s)) => s.clone(),
92                Some(other) => other.to_string(),
93                None => {
94                    return Ok(ToolResult {
95                        tool_name: call.name.clone(),
96                        output: "missing \"input\" argument".into(),
97                        success: false,
98                    })
99                }
100            },
101            Err(e) => {
102                return Ok(ToolResult { tool_name: call.name.clone(), output: format!("arguments are not valid JSON: {e}"), success: false })
103            }
104        };
105        self.dispatch(&call.name, &input)
106    }
107
108    /// The `tools` array for OpenAI Chat Completions (and OpenAI-compatible
109    /// APIs: Mistral, Groq, DeepSeek, OpenRouter, Ollama's `/v1`), sorted by name.
110    ///
111    /// # Errors
112    /// [`AgentError::InvalidToolSignature`] if a tool's schema is not JSON.
113    pub fn openai_tools(&self) -> Result<Vec<serde_json::Value>, AgentError> {
114        self.specs
115            .values()
116            .map(|s| {
117                let (params, _) = crate::native::api_schema(&s.input_schema)?;
118                Ok(serde_json::json!({"type": "function", "function": {"name": s.name, "description": s.description, "parameters": params}}))
119            })
120            .collect()
121    }
122
123    /// The `tools` array for Anthropic Messages, sorted by name.
124    ///
125    /// # Errors
126    /// [`AgentError::InvalidToolSignature`] if a tool's schema is not JSON.
127    pub fn anthropic_tools(&self) -> Result<Vec<serde_json::Value>, AgentError> {
128        self.specs
129            .values()
130            .map(|s| {
131                let (schema, _) = crate::native::api_schema(&s.input_schema)?;
132                Ok(serde_json::json!({"name": s.name, "description": s.description, "input_schema": schema}))
133            })
134            .collect()
135    }
136
137    /// Register a tool whose input is a Rust type: the JSON Schema comes from
138    /// [schemars](https://crates.io/crates/schemars) and the input is parsed
139    /// with serde before `f` runs. Input that does not match the type gives
140    /// the model a failed result explaining why, so it can retry. Needs the
141    /// `schemars` feature.
142    ///
143    /// ```
144    /// use schemars::JsonSchema;
145    /// use serde::Deserialize;
146    ///
147    /// #[derive(Deserialize, JsonSchema)]
148    /// struct Add { a: i64, b: i64 }
149    ///
150    /// let mut tools = wasm_agent::ToolRegistry::new();
151    /// tools.register_typed("add", "Adds two integers", |x: Add| Ok((x.a + x.b).to_string()))?;
152    /// assert_eq!(tools.dispatch("add", r#"{"a": 2, "b": 3}"#)?.output, "5");
153    /// assert!(!tools.dispatch("add", r#"{"a": "two"}"#)?.success);
154    /// # Ok::<(), wasm_agent::AgentError>(())
155    /// ```
156    ///
157    /// # Errors
158    /// [`AgentError::InvalidToolSignature`] if the name is empty.
159    #[cfg(feature = "schemars")]
160    #[cfg_attr(docsrs, doc(cfg(feature = "schemars")))]
161    pub fn register_typed<A, F>(&mut self, name: &str, description: &str, f: F) -> Result<(), AgentError>
162    where
163        A: serde::de::DeserializeOwned + schemars::JsonSchema + 'static,
164        F: Fn(A) -> Result<String, String> + 'static,
165    {
166        let schema = serde_json::to_string(&schemars::schema_for!(A)).map_err(|e| AgentError::Serialization(e.to_string()))?;
167        let tool = name.to_string();
168        self.register(
169            ToolSpec::new(name, description, schema),
170            Box::new(move |input: &str| match serde_json::from_str::<A>(input) {
171                Ok(args) => match f(args) {
172                    Ok(output) => ToolResult { tool_name: tool.clone(), output, success: true },
173                    Err(output) => ToolResult { tool_name: tool.clone(), output, success: false },
174                },
175                Err(e) => ToolResult { tool_name: tool.clone(), output: format!("invalid arguments: {e}"), success: false },
176            }),
177        )
178    }
179
180    /// Returns the spec for a registered tool, or `None` if not found.
181    pub fn spec(&self, name: &str) -> Option<&ToolSpec> { self.specs.get(name) }
182
183    /// Returns the number of registered tools.
184    pub fn tool_count(&self) -> usize { self.specs.len() }
185
186    /// Returns the names of all registered tools, sorted.
187    pub fn tool_names(&self) -> Vec<&str> { self.specs.keys().map(|s| s.as_str()).collect() }
188
189    /// Generates a system-prompt snippet listing all registered tools, sorted
190    /// by name, with each tool's input schema when it has one.
191    pub fn tools_prompt(&self) -> String {
192        let mut lines = vec!["Available tools:".to_string()];
193        for spec in self.specs.values() {
194            let schema = spec.input_schema.trim();
195            if schema.is_empty() || schema == "{}" {
196                lines.push(format!("- {}: {}", spec.name, spec.description));
197            } else {
198                lines.push(format!("- {}: {} (input: {})", spec.name, spec.description, schema));
199            }
200        }
201        lines.join("\n")
202    }
203}
204
205impl Default for ToolRegistry {
206    fn default() -> Self { Self::new() }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    fn echo_tool() -> (ToolSpec, ToolFn) {
214        let spec = ToolSpec::new("echo", "Echoes input back", r#"{"type":"string"}"#);
215        let handler: ToolFn = Box::new(|input: &str| ToolResult {
216            tool_name: "echo".into(),
217            output: format!("echo: {input}"),
218            success: true,
219        });
220        (spec, handler)
221    }
222
223    #[test]
224    fn test_registry_register_and_dispatch_ok() {
225        let mut reg = ToolRegistry::new();
226        let (spec, handler) = echo_tool();
227        reg.register(spec, handler).unwrap();
228        let result = reg.dispatch("echo", "hello").unwrap();
229        assert!(result.success);
230        assert_eq!(result.output, "echo: hello");
231    }
232
233    #[test]
234    fn test_registry_dispatch_unknown_tool_returns_error() {
235        let reg = ToolRegistry::new();
236        let err = reg.dispatch("nonexistent", "").unwrap_err();
237        assert!(matches!(err, AgentError::ToolNotFound { .. }));
238    }
239
240    #[test]
241    fn test_registry_register_empty_name_returns_error() {
242        let mut reg = ToolRegistry::new();
243        let spec = ToolSpec::new("", "bad", "{}");
244        let err = reg.register(spec, Box::new(|_| ToolResult {
245            tool_name: "".into(), output: "".into(), success: false,
246        })).unwrap_err();
247        assert!(matches!(err, AgentError::InvalidToolSignature(_)));
248    }
249
250    #[test]
251    fn test_registry_tool_count_increments() {
252        let mut reg = ToolRegistry::new();
253        assert_eq!(reg.tool_count(), 0);
254        let (spec, handler) = echo_tool();
255        reg.register(spec, handler).unwrap();
256        assert_eq!(reg.tool_count(), 1);
257    }
258
259    #[test]
260    fn test_registry_tools_prompt_contains_tool_name() {
261        let mut reg = ToolRegistry::new();
262        let (spec, handler) = echo_tool();
263        reg.register(spec, handler).unwrap();
264        assert!(reg.tools_prompt().contains("echo"));
265    }
266
267    #[test]
268    fn test_registry_spec_retrieval_present_and_absent() {
269        let mut reg = ToolRegistry::new();
270        let (spec, handler) = echo_tool();
271        reg.register(spec, handler).unwrap();
272        assert!(reg.spec("echo").is_some());
273        assert!(reg.spec("missing").is_none());
274    }
275
276    #[test]
277    fn test_registry_tool_names_lists_all() {
278        let mut reg = ToolRegistry::new();
279        let (spec, handler) = echo_tool();
280        reg.register(spec, handler).unwrap();
281        let names = reg.tool_names();
282        assert!(names.contains(&"echo"));
283    }
284
285    #[test]
286    fn test_registry_default_is_empty() {
287        let reg = ToolRegistry::default();
288        assert_eq!(reg.tool_count(), 0);
289    }
290
291    #[test]
292    fn test_tools_prompt_is_sorted_and_stable() {
293        // 0.1.x iterated a HashMap, so the tool order changed between runs.
294        let mut reg = ToolRegistry::new();
295        for name in ["zeta", "alpha", "mid"] {
296            reg.register(
297                ToolSpec::new(name, "d", r#"{"type":"string"}"#),
298                Box::new(|_: &str| ToolResult { tool_name: String::new(), output: String::new(), success: true }),
299            )
300            .unwrap();
301        }
302        let p = reg.tools_prompt();
303        let a = p.find("- alpha").unwrap();
304        let m = p.find("- mid").unwrap();
305        let z = p.find("- zeta").unwrap();
306        assert!(a < m && m < z);
307        assert!(p.contains(r#"(input: {"type":"string"})"#));
308        assert_eq!(reg.tool_names(), vec!["alpha", "mid", "zeta"]);
309    }
310
311    #[test]
312    fn test_dispatch_call_unwraps_non_object_schemas() {
313        let reg = registry_with_echo_schema(r#"{"type":"string"}"#);
314        let call = |args: &str| crate::native::ToolCall { id: None, name: "echo".into(), arguments: args.into() };
315        assert_eq!(reg.dispatch_call(&call(r#"{"input":"hi"}"#)).unwrap().output, "hi");
316        assert_eq!(reg.dispatch_call(&call(r#"{"input":[1,2]}"#)).unwrap().output, "[1,2]");
317        assert!(!reg.dispatch_call(&call("{bad")).unwrap().success);
318        assert!(!reg.dispatch_call(&call("{}")).unwrap().success);
319        let obj = registry_with_echo_schema(r#"{"type":"object","properties":{"q":{"type":"string"}}}"#);
320        assert_eq!(obj.dispatch_call(&call(r#"{"q":"x"}"#)).unwrap().output, r#"{"q":"x"}"#);
321        assert!(matches!(reg.dispatch_call(&crate::native::ToolCall { id: None, name: "nope".into(), arguments: "{}".into() }), Err(AgentError::ToolNotFound { .. })));
322    }
323
324    #[test]
325    fn test_openai_and_anthropic_tool_definitions() {
326        let reg = registry_with_echo_schema(r#"{"type":"string"}"#);
327        let o = reg.openai_tools().unwrap();
328        assert_eq!(o[0]["type"], "function");
329        assert_eq!(o[0]["function"]["parameters"]["type"], "object");
330        assert_eq!(o[0]["function"]["parameters"]["required"][0], "input");
331        let a = reg.anthropic_tools().unwrap();
332        assert_eq!(a[0]["input_schema"]["properties"]["input"]["type"], "string");
333        let bad = registry_with_echo_schema("{not json");
334        assert!(bad.openai_tools().is_err());
335    }
336
337    fn registry_with_echo_schema(schema: &str) -> ToolRegistry {
338        let mut reg = ToolRegistry::new();
339        reg.register(
340            ToolSpec::new("echo", "Echoes", schema),
341            Box::new(|i: &str| ToolResult { tool_name: "echo".into(), output: i.into(), success: true }),
342        )
343        .unwrap();
344        reg
345    }
346}
wasm_agent/tools.rs

wasm_agent/
tools.rs
