Skip to main content

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 the tool message (OpenAI protocol restricts tool message content to
78/// text, so images cannot travel in the tool result itself).
79#[derive(Debug, Clone)]
80pub struct ToolImage {
81    /// Base64 data URL, e.g. "data:image/png;base64,...".
82    pub data_url: String,
83    /// Human-readable label for log / fallback text.
84    pub label: String,
85}
86
87/// Result returned to the LLM after tool execution.
88#[derive(Debug, Clone)]
89pub struct ToolResult {
90    /// Text content - LLM will read this.
91    pub content: String,
92    /// Whether this is an error (LLM can see errors and adjust strategy).
93    pub is_error: bool,
94    /// Images attached to this result (e.g. from `read` tool reading an image
95    /// file). Most tools leave this empty.
96    pub images: Vec<ToolImage>,
97    /// `true` when this is a *placeholder* for an async task: `content` tells
98    /// the LLM the work is in progress, and the real result is reinjected
99    /// later (by the Agent) when the background task finishes. The Agent uses
100    /// this flag to emit `AsyncToolStarted` instead of treating the call as
101    /// finished. The placeholder content is still added to history as the tool
102    /// message so the LLM can continue other work while waiting.
103    pub is_pending: bool,
104    /// Task id of the background task, set iff `is_pending`. Used by the Agent
105    /// to track/cancel the task and by the frontend to correlate progress.
106    pub pending_task_id: Option<String>,
107}
108
109impl ToolResult {
110    pub fn success(content: impl Into<String>) -> Self {
111        Self {
112            content: content.into(),
113            is_error: false,
114            images: Vec::new(),
115            is_pending: false,
116            pending_task_id: None,
117        }
118    }
119
120    pub fn error(content: impl Into<String>) -> Self {
121        Self {
122            content: content.into(),
123            is_error: true,
124            images: Vec::new(),
125            is_pending: false,
126            pending_task_id: None,
127        }
128    }
129
130    /// Build a pending placeholder for an async task. `content` should tell the
131    /// LLM what is happening and the `task_id` it can reference later.
132    pub fn pending(content: impl Into<String>, task_id: String) -> Self {
133        Self {
134            content: content.into(),
135            is_error: false,
136            images: Vec::new(),
137            is_pending: true,
138            pending_task_id: Some(task_id),
139        }
140    }
141}
142
143// ============================================================================
144// Shared helpers
145// ============================================================================
146
147/// Resolve a file path relative to the working directory.
148pub fn resolve_path(file_path: &str, working_dir: &Path) -> PathBuf {
149    let p = PathBuf::from(file_path);
150    if p.is_absolute() {
151        p
152    } else {
153        working_dir.join(p)
154    }
155}
156
157// ============================================================================
158// ToolContext
159// ============================================================================
160
161/// Runtime context passed to tools during execution.
162pub struct ToolContext {
163    /// Current working directory.
164    pub working_dir: PathBuf,
165    /// Current session ID.
166    pub session_id: SessionId,
167    /// The tool call id this execution was triggered by. Needed by async tools
168    /// to correlate their background task with the originating call.
169    pub tool_call_id: String,
170    /// Frontend for user interaction (e.g., asking for input during tool execution).
171    pub frontend: Arc<dyn Frontend>,
172    /// Platform-specific extensions, keyed by extension ID.
173    /// Chatbot platforms populate this; GUI/TUI leave it empty.
174    /// Keys like "chatbot.platform_ext" map to Arc<dyn PlatformExt>.
175    pub extensions: HashMap<String, Arc<dyn Any + Send + Sync>>,
176    /// Whether the configured LLM supports image inputs.
177    /// Tools (e.g. `read`) use this to decide whether to encode images.
178    pub supports_images: bool,
179    /// Handle for submitting async background work. A tool that decides (at
180    /// runtime) a call should run async calls `async_runner.submit(..)` and
181    /// returns `ToolResult::pending(..)`. Cheap to clone.
182    pub async_runner: AsyncTaskRunner,
183    /// Cancellation token for this tool call. Async tools pass a clone into
184    /// `async_runner.submit` so the Agent can cancel the background work.
185    /// Sync tools ignore it.
186    pub cancel_token: CancellationToken,
187    /// Shared registry tracking all async tasks for the current Agent. The
188    /// `query_task` tool reads this; async tools register themselves here via
189    /// the Agent when they submit. Cheap to clone (shared `Arc`).
190    pub task_registry: TaskRegistry,
191}
192
193// ============================================================================
194// ToolCallInfo (for confirmation requests)
195// ============================================================================
196
197/// Information about a tool call, used for confirmation requests.
198#[derive(Debug, Clone)]
199pub struct ToolCallInfo {
200    pub id: String,
201    pub name: String,
202    pub arguments: String,
203}
204
205// ============================================================================
206// ToolRegistry
207// ============================================================================
208
209/// Registry that manages all available tools.
210pub struct ToolRegistry {
211    tools: HashMap<String, Box<dyn Tool>>,
212}
213
214impl ToolRegistry {
215    pub fn new() -> Self {
216        Self {
217            tools: HashMap::new(),
218        }
219    }
220
221    /// Register a tool. Overwrites any existing tool with the same name.
222    pub fn register(&mut self, tool: impl Tool + 'static) {
223        self.tools.insert(tool.name().to_string(), Box::new(tool));
224    }
225
226    /// Get a list of all registered tool names.
227    pub fn tool_names(&self) -> Vec<&str> {
228        self.tools.keys().map(|s| s.as_str()).collect()
229    }
230
231    /// Check if a tool exists.
232    pub fn contains(&self, name: &str) -> bool {
233        self.tools.contains_key(name)
234    }
235
236    /// Generate OpenAI function calling schemas for all registered tools.
237    pub fn tool_schemas(&self) -> Vec<ChatCompletionTools> {
238        self.tools
239            .values()
240            .map(|tool| {
241                let function = serde_json::json!({
242                    "name": tool.name(),
243                    "description": tool.description(),
244                    "parameters": tool.parameters_schema(),
245                });
246
247                // Construct ChatCompletionTool via JSON deserialization
248                let tool_json = serde_json::json!({
249                    "type": "function",
250                    "function": function,
251                });
252
253                serde_json::from_value(tool_json)
254                    .expect("tool schema should be valid ChatCompletionTools")
255            })
256            .collect()
257    }
258
259    /// Execute a tool by name. Returns an error ToolResult if the tool doesn't exist.
260    pub async fn execute(
261        &self,
262        name: &str,
263        args: Value,
264        ctx: &ToolContext,
265    ) -> ToolResult {
266        tracing::info!("ToolRegistry.execute called: name='{}', args={:?}", name, args);
267        tracing::debug!("Available tools: {:?}", self.tool_names());
268
269        match self.tools.get(name) {
270            Some(tool) => {
271                tracing::debug!("Found tool '{}', executing...", name);
272                let started = std::time::Instant::now();
273                let outcome = tool.execute(args, ctx).await;
274                let elapsed = started.elapsed();
275                match &outcome {
276                    Ok(result) => tracing::trace!(
277                        "[tool:{}] execution finished in {:?}: is_error={}, content_len={}",
278                        name,
279                        elapsed,
280                        result.is_error,
281                        result.content.len()
282                    ),
283                    Err(e) => tracing::warn!(
284                        "[tool:{}] execution returned error after {:?}: {}",
285                        name,
286                        elapsed,
287                        e
288                    ),
289                }
290                match outcome {
291                    Ok(result) => result,
292                    Err(e) => ToolResult::error(format!("Tool execution error: {}", e)),
293                }
294            },
295            None => {
296                let available: Vec<&str> = self.tools.keys().map(|s| s.as_str()).collect();
297                tracing::error!("Tool '{}' not found! Available tools: {:?}", name, available);
298                ToolResult::error(format!(
299                    "Tool '{}' not found. Available tools: {:?}",
300                    name, available
301                ))
302            }
303        }
304    }
305
306    /// Check if a tool requires confirmation.
307    pub fn requires_confirmation(&self, name: &str) -> bool {
308        self.tools
309            .get(name)
310            .map(|t| t.requires_confirmation())
311            .unwrap_or(false)
312    }
313
314    /// Get references to all tools (for prompt building).
315    pub fn tools(&self) -> Vec<&dyn Tool> {
316        self.tools.values().map(|t| t.as_ref()).collect()
317    }
318
319    /// Names of tools that declare async capability (`supports_async() == true`).
320    /// Advisory: frontends use this for UI hints (e.g. a progress affordance).
321    /// Whether an invocation actually runs async is still decided at runtime
322    /// inside `execute`.
323    pub fn async_capable_tools(&self) -> Vec<&str> {
324        self.tools
325            .values()
326            .filter(|t| t.supports_async())
327            .map(|t| t.name())
328            .collect()
329    }
330}
331
332impl Default for ToolRegistry {
333    fn default() -> Self {
334        Self::new()
335    }
336}