Skip to main content

crabtalk_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    /// Session ID of the calling session, if running within a session.
37    /// Set by the runtime; the agent passes it through as an opaque value.
38    pub session_id: Option<u64>,
39}
40
41/// Schema-only registry of named tools.
42///
43/// Stores `Tool` definitions (name, description, JSON schema) keyed by name.
44/// Used by `Runtime` to filter tool schemas per agent at `add_agent` time.
45/// No handlers or closures are stored here.
46#[derive(Default, Clone)]
47pub struct ToolRegistry {
48    tools: BTreeMap<String, Tool>,
49}
50
51impl ToolRegistry {
52    /// Create an empty registry.
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    /// Insert a tool schema.
58    pub fn insert(&mut self, tool: Tool) {
59        self.tools.insert(tool.name.clone(), tool);
60    }
61
62    /// Insert multiple tool schemas.
63    pub fn insert_all(&mut self, tools: Vec<Tool>) {
64        for tool in tools {
65            self.insert(tool);
66        }
67    }
68
69    /// Remove a tool by name. Returns `true` if it existed.
70    pub fn remove(&mut self, name: &str) -> bool {
71        self.tools.remove(name).is_some()
72    }
73
74    /// Check if a tool is registered.
75    pub fn contains(&self, name: &str) -> bool {
76        self.tools.contains_key(name)
77    }
78
79    /// Number of registered tools.
80    pub fn len(&self) -> usize {
81        self.tools.len()
82    }
83
84    /// Whether the registry is empty.
85    pub fn is_empty(&self) -> bool {
86        self.tools.is_empty()
87    }
88
89    /// Return all tool schemas as a `Vec`.
90    pub fn tools(&self) -> Vec<Tool> {
91        self.tools.values().cloned().collect()
92    }
93
94    /// Build a filtered list of tool schemas matching the given names.
95    ///
96    /// If `names` is empty, all tools are returned. Used by `Runtime::add_agent`
97    /// to build the per-agent schema snapshot stored on `Agent`.
98    pub fn filtered_snapshot(&self, names: &[String]) -> Vec<Tool> {
99        if names.is_empty() {
100            return self.tools();
101        }
102        self.tools
103            .iter()
104            .filter(|(k, _)| names.iter().any(|n| n == *k))
105            .map(|(_, v)| v.clone())
106            .collect()
107    }
108}
109
110/// Trait to provide a description for a tool.
111pub trait ToolDescription {
112    /// The description of the tool.
113    const DESCRIPTION: &'static str;
114}
115
116/// Trait to convert a type into a tool.
117pub trait AsTool: ToolDescription {
118    /// Convert the type into a tool.
119    fn as_tool() -> Tool;
120}
121
122impl<T> AsTool for T
123where
124    T: JsonSchema + ToolDescription,
125{
126    fn as_tool() -> Tool {
127        Tool {
128            name: T::schema_name().to_snake_case(),
129            description: Self::DESCRIPTION.into(),
130            parameters: schemars::schema_for!(T),
131            strict: true,
132        }
133    }
134}