Skip to main content

walrus_core/agent/
tool.rs

1//! Tool registry (schema store), ToolRequest, and ToolSender.
2//!
3//! [`ToolRegistry`] stores tool schemas by name — no handlers, no closures.
4//! [`ToolRequest`] and [`ToolSender`] are the agent-side dispatch primitives:
5//! the agent sends a `ToolRequest` per tool call and awaits a `String` reply.
6
7use crate::model::Tool;
8use compact_str::CompactString;
9use std::collections::BTreeMap;
10use tokio::sync::{mpsc, oneshot};
11
12/// A single tool call request sent by the agent to the runtime's tool handler.
13pub struct ToolRequest {
14    /// Tool name as returned by the model.
15    pub name: String,
16    /// JSON-encoded arguments string.
17    pub args: String,
18    /// Reply channel — the handler sends the result string here.
19    pub reply: oneshot::Sender<String>,
20}
21
22/// Sender half of the agent tool channel.
23///
24/// Captured by `Agent` at construction. When the model returns tool calls,
25/// the agent sends one `ToolRequest` per call and awaits each reply.
26/// `None` means no tools are available (e.g. CLI path without a daemon).
27pub type ToolSender = mpsc::UnboundedSender<ToolRequest>;
28
29/// Schema-only registry of named tools.
30///
31/// Stores `Tool` definitions (name, description, JSON schema) keyed by name.
32/// Used by `Runtime` to filter tool schemas per agent at `add_agent` time.
33/// No handlers or closures are stored here.
34#[derive(Default, Clone)]
35pub struct ToolRegistry {
36    tools: BTreeMap<CompactString, Tool>,
37}
38
39impl ToolRegistry {
40    /// Create an empty registry.
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    /// Insert a tool schema.
46    pub fn insert(&mut self, tool: Tool) {
47        self.tools.insert(tool.name.clone(), tool);
48    }
49
50    /// Remove a tool by name. Returns `true` if it existed.
51    pub fn remove(&mut self, name: &str) -> bool {
52        self.tools.remove(name).is_some()
53    }
54
55    /// Check if a tool is registered.
56    pub fn contains(&self, name: &str) -> bool {
57        self.tools.contains_key(name)
58    }
59
60    /// Number of registered tools.
61    pub fn len(&self) -> usize {
62        self.tools.len()
63    }
64
65    /// Whether the registry is empty.
66    pub fn is_empty(&self) -> bool {
67        self.tools.is_empty()
68    }
69
70    /// Return all tool schemas as a `Vec`.
71    pub fn tools(&self) -> Vec<Tool> {
72        self.tools.values().cloned().collect()
73    }
74
75    /// Build a filtered list of tool schemas matching the given names.
76    ///
77    /// If `names` is empty, all tools are returned. Used by `Runtime::add_agent`
78    /// to build the per-agent schema snapshot stored on `Agent`.
79    pub fn filtered_snapshot(&self, names: &[CompactString]) -> Vec<Tool> {
80        if names.is_empty() {
81            return self.tools();
82        }
83        self.tools
84            .iter()
85            .filter(|(k, _)| names.iter().any(|n| n == *k))
86            .map(|(_, v)| v.clone())
87            .collect()
88    }
89}