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 heck::ToSnakeCase;
10use schemars::JsonSchema;
11use std::collections::BTreeMap;
12use tokio::sync::{mpsc, oneshot};
13
14/// Sender half of the agent tool channel.
15///
16/// Captured by `Agent` at construction. When the model returns tool calls,
17/// the agent sends one `ToolRequest` per call and awaits each reply.
18/// `None` means no tools are available (e.g. CLI path without a daemon).
19pub type ToolSender = mpsc::UnboundedSender<ToolRequest>;
20
21/// A single tool call request sent by the agent to the runtime's tool handler.
22pub struct ToolRequest {
23    /// Tool name as returned by the model.
24    pub name: String,
25    /// JSON-encoded arguments string.
26    pub args: String,
27    /// Name of the agent that made this call.
28    pub agent: String,
29    /// Reply channel — the handler sends the result string here.
30    pub reply: oneshot::Sender<String>,
31    /// Task ID of the calling task, if running within a task context.
32    /// Set by the daemon when dispatching task-bound tool calls.
33    pub task_id: Option<u64>,
34    /// Sender identity of the user who triggered this agent run.
35    /// Empty for local/owner sessions.
36    pub sender: CompactString,
37}
38
39/// Schema-only registry of named tools.
40///
41/// Stores `Tool` definitions (name, description, JSON schema) keyed by name.
42/// Used by `Runtime` to filter tool schemas per agent at `add_agent` time.
43/// No handlers or closures are stored here.
44#[derive(Default, Clone)]
45pub struct ToolRegistry {
46    tools: BTreeMap<CompactString, Tool>,
47}
48
49impl ToolRegistry {
50    /// Create an empty registry.
51    pub fn new() -> Self {
52        Self::default()
53    }
54
55    /// Insert a tool schema.
56    pub fn insert(&mut self, tool: Tool) {
57        self.tools.insert(tool.name.clone(), tool);
58    }
59
60    /// Insert multiple tool schemas.
61    pub fn insert_all(&mut self, tools: Vec<Tool>) {
62        for tool in tools {
63            self.insert(tool);
64        }
65    }
66
67    /// Remove a tool by name. Returns `true` if it existed.
68    pub fn remove(&mut self, name: &str) -> bool {
69        self.tools.remove(name).is_some()
70    }
71
72    /// Check if a tool is registered.
73    pub fn contains(&self, name: &str) -> bool {
74        self.tools.contains_key(name)
75    }
76
77    /// Number of registered tools.
78    pub fn len(&self) -> usize {
79        self.tools.len()
80    }
81
82    /// Whether the registry is empty.
83    pub fn is_empty(&self) -> bool {
84        self.tools.is_empty()
85    }
86
87    /// Return all tool schemas as a `Vec`.
88    pub fn tools(&self) -> Vec<Tool> {
89        self.tools.values().cloned().collect()
90    }
91
92    /// Build a filtered list of tool schemas matching the given names.
93    ///
94    /// If `names` is empty, all tools are returned. Used by `Runtime::add_agent`
95    /// to build the per-agent schema snapshot stored on `Agent`.
96    pub fn filtered_snapshot(&self, names: &[CompactString]) -> Vec<Tool> {
97        if names.is_empty() {
98            return self.tools();
99        }
100        self.tools
101            .iter()
102            .filter(|(k, _)| names.iter().any(|n| n == *k))
103            .map(|(_, v)| v.clone())
104            .collect()
105    }
106}
107
108/// Trait to provide a description for a tool.
109pub trait ToolDescription {
110    /// The description of the tool.
111    const DESCRIPTION: &'static str;
112}
113
114/// Trait to convert a type into a tool.
115pub trait AsTool: ToolDescription {
116    /// Convert the type into a tool.
117    fn as_tool() -> Tool;
118}
119
120impl<T> AsTool for T
121where
122    T: JsonSchema + ToolDescription,
123{
124    fn as_tool() -> Tool {
125        Tool {
126            name: T::schema_name().to_snake_case().into(),
127            description: Self::DESCRIPTION.into(),
128            parameters: schemars::schema_for!(T),
129            strict: true,
130        }
131    }
132}