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

walrus_core/agent/
tool.rs
