Skip to main content

crabtalk_core/agent/
tool.rs

1//! Tool registry (schema store), ToolRequest, and ToolSender.
2//!
3//! [`ToolRegistry`] stores `crabllm_core::Tool` schemas by name — no
4//! handlers, no closures. [`ToolRequest`] and [`ToolSender`] are the
5//! agent-side dispatch primitives: the agent sends a `ToolRequest` per tool
6//! call and awaits a `Result<String, String>` reply — `Ok` carries normal
7//! output, `Err` carries an error message the UI can render distinctly.
8
9use crabllm_core::{FunctionDef, Tool, ToolType};
10use heck::ToSnakeCase;
11use schemars::JsonSchema;
12use std::collections::BTreeMap;
13use tokio::sync::{mpsc, oneshot};
14
15/// Sender half of the agent tool channel.
16///
17/// Captured by `Agent` at construction. When the model returns tool calls,
18/// the agent sends one `ToolRequest` per call and awaits each reply.
19/// `None` means no tools are available (e.g. CLI path without a daemon).
20pub type ToolSender = mpsc::UnboundedSender<ToolRequest>;
21
22/// A single tool call request sent by the agent to the runtime's tool handler.
23pub struct ToolRequest {
24    /// Tool name as returned by the model.
25    pub name: String,
26    /// JSON-encoded arguments string.
27    pub args: String,
28    /// Name of the agent that made this call.
29    pub agent: String,
30    /// Reply channel — the handler sends a `Result<String, String>`:
31    /// `Ok` for normal output, `Err` for an error message the UI can
32    /// render as a failure and the agent can act on for retry decisions.
33    pub reply: oneshot::Sender<Result<String, String>>,
34    /// Task ID of the calling task, if running within a task context.
35    /// Set by the daemon when dispatching task-bound tool calls.
36    pub task_id: Option<u64>,
37    /// Sender identity of the user who triggered this agent run.
38    /// Empty for local/owner conversations.
39    pub sender: String,
40    /// Conversation ID, if running within a conversation.
41    /// Set by the runtime; the agent passes it through as an opaque value.
42    pub conversation_id: Option<u64>,
43}
44
45/// Schema-only registry of named tools.
46///
47/// Stores `crabllm_core::Tool` definitions keyed by function name. Used by
48/// `Runtime` to filter tool schemas per agent at `add_agent` time. No
49/// handlers or closures are stored here.
50#[derive(Default, Clone)]
51pub struct ToolRegistry {
52    tools: BTreeMap<String, Tool>,
53}
54
55impl ToolRegistry {
56    /// Create an empty registry.
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    /// Insert a tool schema, keyed by its function name.
62    pub fn insert(&mut self, tool: Tool) {
63        self.tools.insert(tool.function.name.clone(), tool);
64    }
65
66    /// Insert multiple tool schemas.
67    pub fn insert_all(&mut self, tools: Vec<Tool>) {
68        for tool in tools {
69            self.insert(tool);
70        }
71    }
72
73    /// Remove a tool by name. Returns `true` if it existed.
74    pub fn remove(&mut self, name: &str) -> bool {
75        self.tools.remove(name).is_some()
76    }
77
78    /// Check if a tool is registered.
79    pub fn contains(&self, name: &str) -> bool {
80        self.tools.contains_key(name)
81    }
82
83    /// Number of registered tools.
84    pub fn len(&self) -> usize {
85        self.tools.len()
86    }
87
88    /// Whether the registry is empty.
89    pub fn is_empty(&self) -> bool {
90        self.tools.is_empty()
91    }
92
93    /// Return all tool schemas as a `Vec`.
94    pub fn tools(&self) -> Vec<Tool> {
95        self.tools.values().cloned().collect()
96    }
97
98    /// Build a filtered list of tool schemas matching the given names.
99    ///
100    /// If `names` is empty, all tools are returned. Used by `Runtime::add_agent`
101    /// to build the per-agent schema snapshot stored on `Agent`.
102    pub fn filtered_snapshot(&self, names: &[String]) -> Vec<Tool> {
103        if names.is_empty() {
104            return self.tools();
105        }
106        self.tools
107            .iter()
108            .filter(|(k, _)| names.iter().any(|n| n == *k))
109            .map(|(_, v)| v.clone())
110            .collect()
111    }
112}
113
114/// Trait to provide a description for a tool.
115pub trait ToolDescription {
116    /// The description of the tool.
117    const DESCRIPTION: &'static str;
118}
119
120/// Trait to convert a type into a `crabllm_core::Tool`.
121pub trait AsTool: ToolDescription {
122    /// Convert the type into a `crabllm_core::Tool` (the enveloped
123    /// `{kind, function}` wire shape).
124    fn as_tool() -> Tool;
125}
126
127impl<T> AsTool for T
128where
129    T: JsonSchema + ToolDescription,
130{
131    fn as_tool() -> Tool {
132        // `strict: None` matches the prior wire behavior: the wcore
133        // `Tool.strict: bool` field was set to `true` by every `AsTool` impl
134        // but silently dropped by the converter (old convert::to_ct_tool
135        // hard-coded `strict: None`). Turning on strict-mode validation
136        // here would be a behavior change masquerading as a refactor —
137        // leave any opt-in to a separate commit that validates every tool
138        // schema.
139        Tool {
140            kind: ToolType::Function,
141            function: FunctionDef {
142                name: T::schema_name().to_snake_case(),
143                description: Some(Self::DESCRIPTION.into()),
144                parameters: Some(
145                    serde_json::to_value(schemars::schema_for!(T)).unwrap_or_default(),
146                ),
147            },
148            strict: None,
149        }
150    }
151}