robit_agent/tool/mod.rs
1//! Tool system: trait, registry, result types, and context.
2
3pub mod bash;
4pub mod read;
5pub mod write;
6pub mod edit;
7pub mod generate_image;
8pub mod load_skill;
9pub mod ls;
10pub mod find;
11pub mod grep;
12pub mod memory;
13pub mod search_history;
14pub mod async_runner;
15pub mod task_registry;
16pub mod query_task;
17
18use async_trait::async_trait;
19use robit_ai::ChatCompletionTools;
20use serde_json::Value;
21use std::collections::HashMap;
22use std::any::Any;
23use std::path::{Path, PathBuf};
24use std::sync::Arc;
25use tokio_util::sync::CancellationToken;
26
27use crate::error::Result;
28use crate::event::SessionId;
29use crate::frontend::Frontend;
30use async_runner::AsyncTaskRunner;
31use task_registry::TaskRegistry;
32
33// ============================================================================
34// Tool trait
35// ============================================================================
36
37/// A tool that can be called by the LLM and executed by the Agent.
38#[async_trait]
39pub trait Tool: Send + Sync {
40 /// Tool name — LLM calls the tool by this name.
41 fn name(&self) -> &str;
42
43 /// Tool description — injected into system prompt for LLM understanding.
44 fn description(&self) -> &str;
45
46 /// JSON Schema for tool parameters — LLM generates arguments based on this.
47 fn parameters_schema(&self) -> Value;
48
49 /// Whether this tool requires user confirmation before execution.
50 fn requires_confirmation(&self) -> bool;
51
52 /// Execute the tool with parsed arguments. Returns ToolResult for LLM consumption.
53 async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<ToolResult>;
54
55 /// Whether this tool is capable of running asynchronously (returning a
56 /// pending placeholder and finishing its work in the background).
57 ///
58 /// This is **advisory only** - used by frontends/Agent for UI hints (e.g.
59 /// showing a "task in progress" affordance). Whether a *given invocation*
60 /// actually runs async is decided at runtime inside `execute` (e.g. based
61 /// on the provider protocol or input size), by calling
62 /// `ctx.async_runner.submit(..)` and returning `ToolResult::pending(..)`.
63 /// Tools that never run async should leave the default `false`.
64 fn supports_async(&self) -> bool {
65 false
66 }
67}
68
69// ============================================================================
70// ToolResult
71// ============================================================================
72
73/// A single image attached to a tool result.
74///
75/// When a tool (e.g. `read` on an image file) produces images and the model
76/// supports image inputs, the agent injects them as a multimodal user message
77/// after all tool messages of the batch (OpenAI protocol restricts tool
78/// message content to text, so images cannot travel in the tool result
79/// itself, and nothing may interleave with the tool responses).
80#[derive(Debug, Clone)]
81pub struct ToolImage {
82 /// Base64 data URL, e.g. "data:image/png;base64,...".
83 pub data_url: String,
84 /// Human-readable label for log / fallback text.
85 pub label: String,
86}
87
88/// Result returned to the LLM after tool execution.
89#[derive(Debug, Clone)]
90pub struct ToolResult {
91 /// Text content - LLM will read this.
92 pub content: String,
93 /// Whether this is an error (LLM can see errors and adjust strategy).
94 pub is_error: bool,
95 /// Images attached to this result (e.g. from `read` tool reading an image
96 /// file). Most tools leave this empty.
97 pub images: Vec<ToolImage>,
98 /// `true` when this is a *placeholder* for an async task: `content` tells
99 /// the LLM the work is in progress, and the real result is reinjected
100 /// later (by the Agent) when the background task finishes. The Agent uses
101 /// this flag to emit `AsyncToolStarted` instead of treating the call as
102 /// finished. The placeholder content is still added to history as the tool
103 /// message so the LLM can continue other work while waiting.
104 pub is_pending: bool,
105 /// Task id of the background task, set iff `is_pending`. Used by the Agent
106 /// to track/cancel the task and by the frontend to correlate progress.
107 pub pending_task_id: Option<String>,
108}
109
110impl ToolResult {
111 pub fn success(content: impl Into<String>) -> Self {
112 Self {
113 content: content.into(),
114 is_error: false,
115 images: Vec::new(),
116 is_pending: false,
117 pending_task_id: None,
118 }
119 }
120
121 pub fn error(content: impl Into<String>) -> Self {
122 Self {
123 content: content.into(),
124 is_error: true,
125 images: Vec::new(),
126 is_pending: false,
127 pending_task_id: None,
128 }
129 }
130
131 /// Build a pending placeholder for an async task. `content` should tell the
132 /// LLM what is happening and the `task_id` it can reference later.
133 pub fn pending(content: impl Into<String>, task_id: String) -> Self {
134 Self {
135 content: content.into(),
136 is_error: false,
137 images: Vec::new(),
138 is_pending: true,
139 pending_task_id: Some(task_id),
140 }
141 }
142}
143
144// ============================================================================
145// Shared helpers
146// ============================================================================
147
148/// Resolve a file path relative to the working directory.
149pub fn resolve_path(file_path: &str, working_dir: &Path) -> PathBuf {
150 let p = PathBuf::from(file_path);
151 if p.is_absolute() {
152 p
153 } else {
154 working_dir.join(p)
155 }
156}
157
158// ============================================================================
159// ToolContext
160// ============================================================================
161
162/// Runtime context passed to tools during execution.
163pub struct ToolContext {
164 /// Current working directory.
165 pub working_dir: PathBuf,
166 /// Current session ID.
167 pub session_id: SessionId,
168 /// The tool call id this execution was triggered by. Needed by async tools
169 /// to correlate their background task with the originating call.
170 pub tool_call_id: String,
171 /// Frontend for user interaction (e.g., asking for input during tool execution).
172 pub frontend: Arc<dyn Frontend>,
173 /// Platform-specific extensions, keyed by extension ID.
174 /// Chatbot platforms populate this; GUI/TUI leave it empty.
175 /// Keys like "chatbot.platform_ext" map to Arc<dyn PlatformExt>.
176 pub extensions: HashMap<String, Arc<dyn Any + Send + Sync>>,
177 /// Whether the configured LLM supports image inputs.
178 /// Tools (e.g. `read`) use this to decide whether to encode images.
179 pub supports_images: bool,
180 /// Handle for submitting async background work. A tool that decides (at
181 /// runtime) a call should run async calls `async_runner.submit(..)` and
182 /// returns `ToolResult::pending(..)`. Cheap to clone.
183 pub async_runner: AsyncTaskRunner,
184 /// Cancellation token for this tool call. Async tools pass a clone into
185 /// `async_runner.submit` so the Agent can cancel the background work.
186 /// Sync tools ignore it.
187 pub cancel_token: CancellationToken,
188 /// Shared registry tracking all async tasks for the current Agent. The
189 /// `query_task` tool reads this; async tools register themselves here via
190 /// the Agent when they submit. Cheap to clone (shared `Arc`).
191 pub task_registry: TaskRegistry,
192}
193
194// ============================================================================
195// ToolCallInfo (for confirmation requests)
196// ============================================================================
197
198/// Information about a tool call, used for confirmation requests.
199#[derive(Debug, Clone)]
200pub struct ToolCallInfo {
201 pub id: String,
202 pub name: String,
203 pub arguments: String,
204}
205
206// ============================================================================
207// ToolRegistry
208// ============================================================================
209
210/// Registry that manages all available tools.
211pub struct ToolRegistry {
212 tools: HashMap<String, Box<dyn Tool>>,
213}
214
215impl ToolRegistry {
216 pub fn new() -> Self {
217 Self {
218 tools: HashMap::new(),
219 }
220 }
221
222 /// Register a tool. Overwrites any existing tool with the same name.
223 pub fn register(&mut self, tool: impl Tool + 'static) {
224 self.tools.insert(tool.name().to_string(), Box::new(tool));
225 }
226
227 /// Get a list of all registered tool names.
228 pub fn tool_names(&self) -> Vec<&str> {
229 self.tools.keys().map(|s| s.as_str()).collect()
230 }
231
232 /// Check if a tool exists.
233 pub fn contains(&self, name: &str) -> bool {
234 self.tools.contains_key(name)
235 }
236
237 /// Generate OpenAI function calling schemas for all registered tools.
238 pub fn tool_schemas(&self) -> Vec<ChatCompletionTools> {
239 self.tools
240 .values()
241 .map(|tool| {
242 let function = serde_json::json!({
243 "name": tool.name(),
244 "description": tool.description(),
245 "parameters": tool.parameters_schema(),
246 });
247
248 // Construct ChatCompletionTool via JSON deserialization
249 let tool_json = serde_json::json!({
250 "type": "function",
251 "function": function,
252 });
253
254 serde_json::from_value(tool_json)
255 .expect("tool schema should be valid ChatCompletionTools")
256 })
257 .collect()
258 }
259
260 /// Execute a tool by name. Returns an error ToolResult if the tool doesn't exist.
261 pub async fn execute(
262 &self,
263 name: &str,
264 args: Value,
265 ctx: &ToolContext,
266 ) -> ToolResult {
267 // Truncate args for logging: tool arguments (e.g. `write` file
268 // content) can be huge and would flood the log.
269 let args_str = args.to_string();
270 let args_preview: String = {
271 let chars: String = args_str.chars().take(120).collect();
272 if args_str.chars().count() > 120 {
273 format!("{}...", chars)
274 } else {
275 chars
276 }
277 };
278 tracing::debug!("ToolRegistry.execute: name='{}', args={}", name, args_preview);
279
280 match self.tools.get(name) {
281 Some(tool) => {
282 let started = std::time::Instant::now();
283 let outcome = tool.execute(args, ctx).await;
284 let elapsed = started.elapsed();
285 match &outcome {
286 Ok(result) => tracing::trace!(
287 "[tool:{}] execution finished in {:?}: is_error={}, content_len={}",
288 name,
289 elapsed,
290 result.is_error,
291 result.content.len()
292 ),
293 Err(e) => tracing::warn!(
294 "[tool:{}] execution returned error after {:?}: {}",
295 name,
296 elapsed,
297 e
298 ),
299 }
300 match outcome {
301 Ok(result) => result,
302 Err(e) => ToolResult::error(format!("Tool execution error: {}", e)),
303 }
304 },
305 None => {
306 let available: Vec<&str> = self.tools.keys().map(|s| s.as_str()).collect();
307 tracing::error!("Tool '{}' not found! Available tools: {:?}", name, available);
308 ToolResult::error(format!(
309 "Tool '{}' not found. Available tools: {:?}",
310 name, available
311 ))
312 }
313 }
314 }
315
316 /// Check if a tool requires confirmation.
317 pub fn requires_confirmation(&self, name: &str) -> bool {
318 self.tools
319 .get(name)
320 .map(|t| t.requires_confirmation())
321 .unwrap_or(false)
322 }
323
324 /// Get references to all tools (for prompt building).
325 pub fn tools(&self) -> Vec<&dyn Tool> {
326 self.tools.values().map(|t| t.as_ref()).collect()
327 }
328
329 /// Names of tools that declare async capability (`supports_async() == true`).
330 /// Advisory: frontends use this for UI hints (e.g. a progress affordance).
331 /// Whether an invocation actually runs async is still decided at runtime
332 /// inside `execute`.
333 pub fn async_capable_tools(&self) -> Vec<&str> {
334 self.tools
335 .values()
336 .filter(|t| t.supports_async())
337 .map(|t| t.name())
338 .collect()
339 }
340}
341
342impl Default for ToolRegistry {
343 fn default() -> Self {
344 Self::new()
345 }
346}