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}