walrus_core/agent/
tool.rs

1//! Dispatcher trait, Handler type, and ToolRegistry.
2//!
3//! [`Dispatcher`] is a generic async trait for tool dispatch, passed to
4//! `Agent::step()`. [`ToolRegistry`] is the canonical implementation — it
5//! holds `(Tool, Handler)` pairs keyed by name and implements `Dispatcher`
6//! directly, removing the need for `ClosureDispatcher` or `DispatchFn`.
7
8use crate::model::Tool;
9use anyhow::Result;
10use compact_str::CompactString;
11use std::{collections::BTreeMap, future::Future, pin::Pin, sync::Arc};
12
13/// Type-erased async tool handler.
14///
15/// Takes JSON-encoded arguments, returns a result string. Captured state
16/// (e.g. `Arc<M>`) must be `Send + Sync + 'static`.
17pub type Handler =
18    Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = String> + Send>> + Send + Sync>;
19
20/// Generic tool dispatcher.
21///
22/// Passed as a method param to `Agent::step()`. Uses RPITIT for async
23/// without boxing — callers monomorphize over concrete dispatcher types.
24pub trait Dispatcher: Send + Sync {
25    /// Dispatch a batch of tool calls. Each entry is `(method, params)`.
26    ///
27    /// Returns one result per call in the same order.
28    fn dispatch(&self, calls: &[(&str, &str)]) -> impl Future<Output = Vec<Result<String>>> + Send;
29
30    /// Return the tool schemas this dispatcher can handle.
31    ///
32    /// Agent uses this to populate `Request.tools` before calling the model.
33    fn tools(&self) -> Vec<Tool>;
34}
35
36/// Registry of named tools with their async handlers.
37///
38/// Implements [`Dispatcher`] directly — `tools()` returns schemas and
39/// `dispatch()` looks up handlers by name. Used as both the runtime's
40/// shared store and the per-agent dispatcher snapshot.
41#[derive(Default)]
42pub struct ToolRegistry {
43    tools: BTreeMap<CompactString, (Tool, Handler)>,
44}
45
46impl ToolRegistry {
47    /// Create an empty registry.
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    /// Insert a tool and its handler.
53    pub fn insert(&mut self, tool: Tool, handler: Handler) {
54        self.tools.insert(tool.name.clone(), (tool, handler));
55    }
56
57    /// Remove a tool by name. Returns `true` if it existed.
58    pub fn remove(&mut self, name: &str) -> bool {
59        self.tools.remove(name).is_some()
60    }
61
62    /// Check if a tool is registered.
63    pub fn contains(&self, name: &str) -> bool {
64        self.tools.contains_key(name)
65    }
66
67    /// Iterate over all `(Tool, Handler)` pairs.
68    pub fn iter(&self) -> impl Iterator<Item = (&Tool, &Handler)> {
69        self.tools.values().map(|(t, h)| (t, h))
70    }
71
72    /// Number of registered tools.
73    pub fn len(&self) -> usize {
74        self.tools.len()
75    }
76
77    /// Whether the registry is empty.
78    pub fn is_empty(&self) -> bool {
79        self.tools.is_empty()
80    }
81
82    /// Build a filtered snapshot containing only the named tools.
83    ///
84    /// If `names` is empty, all tools are included. Used by runtime to
85    /// build a per-agent dispatcher.
86    pub fn filtered_snapshot(&self, names: &[CompactString]) -> Self {
87        let tools = if names.is_empty() {
88            self.tools.clone()
89        } else {
90            self.tools
91                .iter()
92                .filter(|(k, _)| names.iter().any(|n| n == *k))
93                .map(|(k, v)| (k.clone(), v.clone()))
94                .collect()
95        };
96        Self { tools }
97    }
98}
99
100impl Clone for ToolRegistry {
101    fn clone(&self) -> Self {
102        Self {
103            tools: self.tools.clone(),
104        }
105    }
106}
107
108impl Dispatcher for ToolRegistry {
109    fn dispatch(&self, calls: &[(&str, &str)]) -> impl Future<Output = Vec<Result<String>>> + Send {
110        let owned: Vec<(String, String, Option<Handler>)> = calls
111            .iter()
112            .map(|(m, p)| {
113                let handler = self.tools.get(*m).map(|(_, h)| Arc::clone(h));
114                (m.to_string(), p.to_string(), handler)
115            })
116            .collect();
117
118        async move {
119            let mut results = Vec::with_capacity(owned.len());
120            for (method, params, handler) in owned {
121                let output = if let Some(h) = handler {
122                    Ok(h(params).await)
123                } else {
124                    Ok(format!("function {method} not available"))
125                };
126                results.push(output);
127            }
128            results
129        }
130    }
131
132    fn tools(&self) -> Vec<Tool> {
133        self.tools.values().map(|(t, _)| t.clone()).collect()
134    }
135}
walrus_core/agent/tool.rs

walrus_core/agent/
tool.rs
