aster/tools/
mod.rs

1// =============================================================================
2// Tool System Module
3// =============================================================================
4//
5// This module provides a unified tool system for aster-rust, aligned with
6// - Tool trait and base types
7// - Tool registry for managing native and MCP tools
8// - Core tool implementations (Bash, File, Search, etc.)
9// - Permission integration
10// - Audit logging
11
12// Core modules
13pub mod base;
14pub mod context;
15pub mod error;
16pub mod hooks;
17pub mod registry;
18pub mod task;
19
20// Tool implementations
21pub mod analyze_image;
22pub mod ask;
23pub mod bash;
24pub mod file;
25pub mod kill_shell_tool;
26pub mod lsp;
27pub mod notebook_edit_tool;
28pub mod plan_mode_tool;
29pub mod search;
30pub mod task_output_tool;
31pub mod task_tool;
32pub mod three_files_tool;
33pub mod todo_write_tool;
34pub mod web;
35pub mod workflow_integration;
36
37// Skills integration
38
39// =============================================================================
40// Core Type Exports
41// =============================================================================
42
43// Error types
44pub use error::ToolError;
45
46// Context and configuration types
47pub use context::{ToolContext, ToolDefinition, ToolOptions, ToolResult};
48
49// Base trait and permission types
50pub use base::{PermissionBehavior, PermissionCheckResult, Tool};
51
52// Registry types
53pub use registry::{McpToolWrapper, PermissionRequestCallback, ToolRegistry};
54
55// Hook system types
56pub use hooks::{
57    ErrorTrackingHook, FileOperationHook, HookContext, HookTrigger, LoggingHook, ToolHook,
58    ToolHookManager,
59};
60
61// Task management types
62pub use task::{
63    TaskManager, TaskState, TaskStatus, DEFAULT_MAX_CONCURRENT, DEFAULT_MAX_RUNTIME_SECS,
64};
65
66// Tool implementations
67pub use bash::{BashTool, SafetyCheckResult, SandboxConfig, MAX_OUTPUT_LENGTH};
68
69// File tools
70pub use file::{
71    compute_content_hash, create_shared_history, EditTool, FileReadHistory, FileReadRecord,
72    ReadTool, SharedFileReadHistory, WriteTool,
73};
74
75// Search tools
76pub use search::{
77    GlobTool, GrepOutputMode, GrepTool, SearchResult, DEFAULT_MAX_CONTEXT_LINES,
78    DEFAULT_MAX_RESULTS, MAX_OUTPUT_SIZE,
79};
80
81// Ask tool
82pub use ask::{AskCallback, AskOption, AskResult, AskTool, DEFAULT_ASK_TIMEOUT_SECS};
83
84// LSP tool
85pub use lsp::{
86    CompletionItem, CompletionItemKind, Diagnostic, DiagnosticSeverity, HoverInfo, Location,
87    LspCallback, LspOperation, LspResult, LspTool, Position, Range,
88};
89
90// Skill tool
91pub use crate::skills::SkillTool;
92
93// Task tools
94pub use kill_shell_tool::KillShellTool;
95pub use notebook_edit_tool::{NotebookCell, NotebookContent, NotebookEditInput, NotebookEditTool};
96pub use plan_mode_tool::{EnterPlanModeTool, ExitPlanModeTool, PlanModeState, SavedPlan};
97pub use task_output_tool::TaskOutputTool;
98pub use task_tool::TaskTool;
99pub use three_files_tool::{
100    DecisionInfo, ErrorInfo, PhaseUpdate, ThreeStageWorkflowTool, WorkflowParams,
101};
102pub use todo_write_tool::{TodoItem, TodoStatus, TodoStorage, TodoWriteTool};
103
104// Web tools
105pub use web::{clear_web_caches, get_web_cache_stats, WebCache, WebFetchTool, WebSearchTool};
106
107// Image analysis tools
108// Image analysis tools
109pub use analyze_image::AnalyzeImageTool;
110pub use analyze_image::{AnalyzeImageInput, AnalyzeImageResult, ImageDimensions};
111
112// Workflow integration
113pub use workflow_integration::{WorkflowIntegratedTool, WorkflowIntegratedToolBuilder};
114
115// =============================================================================
116// Tool Registration (Requirements: 11.3)
117// =============================================================================
118
119/// Configuration for tool registration
120#[derive(Default)]
121pub struct ToolRegistrationConfig {
122    /// Callback for AskTool user interaction
123    pub ask_callback: Option<AskCallback>,
124    /// Callback for LSPTool operations
125    pub lsp_callback: Option<LspCallback>,
126    /// Whether to enable PDF reading in ReadTool
127    pub pdf_enabled: bool,
128    /// Whether to enable hook system
129    pub hooks_enabled: bool,
130}
131
132impl std::fmt::Debug for ToolRegistrationConfig {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        f.debug_struct("ToolRegistrationConfig")
135            .field(
136                "ask_callback",
137                &self.ask_callback.as_ref().map(|_| "<callback>"),
138            )
139            .field(
140                "lsp_callback",
141                &self.lsp_callback.as_ref().map(|_| "<callback>"),
142            )
143            .field("pdf_enabled", &self.pdf_enabled)
144            .field("hooks_enabled", &self.hooks_enabled)
145            .finish()
146    }
147}
148
149impl Clone for ToolRegistrationConfig {
150    fn clone(&self) -> Self {
151        Self {
152            ask_callback: self.ask_callback.clone(),
153            lsp_callback: self.lsp_callback.clone(),
154            pdf_enabled: self.pdf_enabled,
155            hooks_enabled: self.hooks_enabled,
156        }
157    }
158}
159
160impl ToolRegistrationConfig {
161    /// Create a new configuration with default settings
162    pub fn new() -> Self {
163        Self::default()
164    }
165
166    /// Set the AskTool callback
167    pub fn with_ask_callback(mut self, callback: AskCallback) -> Self {
168        self.ask_callback = Some(callback);
169        self
170    }
171
172    /// Set the LSPTool callback
173    pub fn with_lsp_callback(mut self, callback: LspCallback) -> Self {
174        self.lsp_callback = Some(callback);
175        self
176    }
177
178    /// Enable PDF reading
179    pub fn with_pdf_enabled(mut self, enabled: bool) -> Self {
180        self.pdf_enabled = enabled;
181        self
182    }
183
184    /// Enable hook system
185    pub fn with_hooks_enabled(mut self, enabled: bool) -> Self {
186        self.hooks_enabled = enabled;
187        self
188    }
189}
190
191/// Register all native tools with the registry
192///
193/// This function registers all built-in tools:
194/// - BashTool: Shell command execution
195/// - ReadTool: File reading (text, images, PDF, notebooks)
196/// - WriteTool: File writing with validation
197/// - EditTool: Smart file editing
198/// - GlobTool: File search with glob patterns
199/// - GrepTool: Content search with regex
200/// - AskTool: User interaction (if callback provided)
201/// - LSPTool: Code intelligence (if callback provided)
202/// - SkillTool: Skill execution and management
203///
204/// # Arguments
205/// * `registry` - The ToolRegistry to register tools with
206/// * `config` - Configuration for tool registration
207///
208/// # Returns
209/// A tuple containing (shared file read history, hook manager)
210///
211/// Requirements: 11.3
212pub fn register_all_tools(
213    registry: &mut ToolRegistry,
214    config: ToolRegistrationConfig,
215) -> (SharedFileReadHistory, Option<ToolHookManager>) {
216    // Create shared file read history for file tools
217    let shared_history = create_shared_history();
218
219    // Initialize hook manager if enabled
220    let hook_manager = if config.hooks_enabled {
221        let manager = ToolHookManager::new(true);
222        // Register default hooks in a blocking context
223        tokio::task::block_in_place(|| {
224            tokio::runtime::Handle::current().block_on(async {
225                manager.register_default_hooks().await;
226            })
227        });
228        Some(manager)
229    } else {
230        None
231    };
232
233    // Register BashTool
234    registry.register(Box::new(BashTool::new()));
235
236    // Register file tools with shared history
237    let read_tool = ReadTool::new(shared_history.clone()).with_pdf_enabled(config.pdf_enabled);
238    registry.register(Box::new(read_tool));
239
240    let write_tool = WriteTool::new(shared_history.clone());
241    registry.register(Box::new(write_tool));
242
243    let edit_tool = EditTool::new(shared_history.clone());
244    registry.register(Box::new(edit_tool));
245
246    // Register search tools
247    registry.register(Box::new(GlobTool::new()));
248    registry.register(Box::new(GrepTool::new()));
249
250    // Register AskTool if callback is provided
251    if let Some(callback) = config.ask_callback {
252        let ask_tool = AskTool::new().with_callback(callback);
253        registry.register(Box::new(ask_tool));
254    }
255
256    // Register LSPTool if callback is provided
257    if let Some(callback) = config.lsp_callback {
258        let lsp_tool = LspTool::new().with_callback(callback);
259        registry.register(Box::new(lsp_tool));
260    }
261
262    // Register SkillTool
263    registry.register(Box::new(SkillTool::new()));
264
265    // Register TaskTool and TaskOutputTool
266    registry.register(Box::new(TaskTool::new()));
267    registry.register(Box::new(TaskOutputTool::new()));
268    registry.register(Box::new(KillShellTool::new()));
269    registry.register(Box::new(TodoWriteTool::new()));
270    registry.register(Box::new(NotebookEditTool::new()));
271
272    // Register Plan Mode tools
273    registry.register(Box::new(EnterPlanModeTool::new()));
274    registry.register(Box::new(ExitPlanModeTool::new()));
275
276    // Register Web tools
277    registry.register(Box::new(WebFetchTool::new()));
278    registry.register(Box::new(WebSearchTool::new()));
279
280    // Register Image Analysis tools
281    registry.register(Box::new(AnalyzeImageTool::new()));
282
283    // Register Three-Stage Workflow tool
284    registry.register(Box::new(ThreeStageWorkflowTool::default()));
285
286    (shared_history, hook_manager)
287}
288
289/// Register all native tools with default configuration
290///
291/// This is a convenience function that registers all tools with default settings.
292/// AskTool and LSPTool are not registered since they require callbacks.
293///
294/// # Arguments
295/// * `registry` - The ToolRegistry to register tools with
296///
297/// # Returns
298/// A tuple containing (shared file read history, hook manager)
299///
300/// Requirements: 11.3
301pub fn register_default_tools(
302    registry: &mut ToolRegistry,
303) -> (SharedFileReadHistory, Option<ToolHookManager>) {
304    register_all_tools(registry, ToolRegistrationConfig::default())
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use std::path::PathBuf;
311
312    #[test]
313    fn test_register_default_tools() {
314        let mut registry = ToolRegistry::new();
315        let (_history, _hook_manager) = register_default_tools(&mut registry);
316
317        // Verify core tools are registered
318        assert!(registry.contains("bash"));
319        assert!(registry.contains("read"));
320        assert!(registry.contains("write"));
321        assert!(registry.contains("edit"));
322        assert!(registry.contains("glob"));
323        assert!(registry.contains("grep"));
324        assert!(registry.contains("Skill"));
325        assert!(registry.contains("Task"));
326        assert!(registry.contains("TaskOutput"));
327        assert!(registry.contains("KillShell"));
328        assert!(registry.contains("TodoWrite"));
329        assert!(registry.contains("NotebookEdit"));
330        assert!(registry.contains("EnterPlanMode"));
331        assert!(registry.contains("ExitPlanMode"));
332        assert!(registry.contains("WebFetch"));
333        assert!(registry.contains("WebSearch"));
334        assert!(registry.contains("analyze_image"));
335        assert!(registry.contains("three_stage_workflow"));
336
337        // AskTool and LSPTool should not be registered without callbacks
338        assert!(!registry.contains("ask"));
339        assert!(!registry.contains("lsp"));
340    }
341
342    #[test]
343    fn test_register_all_tools_with_config() {
344        use std::future::Future;
345        use std::pin::Pin;
346        use std::sync::Arc;
347
348        let mut registry = ToolRegistry::new();
349
350        // Create mock callbacks
351        let ask_callback: AskCallback = Arc::new(|_question, _options| {
352            Box::pin(async { Some("test response".to_string()) })
353                as Pin<Box<dyn Future<Output = Option<String>> + Send>>
354        });
355
356        let lsp_callback: LspCallback = Arc::new(|_operation, _path: PathBuf, _position| {
357            Box::pin(async { Ok(LspResult::Definition { locations: vec![] }) })
358                as Pin<Box<dyn Future<Output = Result<LspResult, String>> + Send>>
359        });
360
361        let config = ToolRegistrationConfig::new()
362            .with_ask_callback(ask_callback)
363            .with_lsp_callback(lsp_callback)
364            .with_pdf_enabled(true);
365
366        let (_history, _hook_manager) = register_all_tools(&mut registry, config);
367
368        // Verify all tools are registered
369        assert!(registry.contains("bash"));
370        assert!(registry.contains("read"));
371        assert!(registry.contains("write"));
372        assert!(registry.contains("edit"));
373        assert!(registry.contains("glob"));
374        assert!(registry.contains("grep"));
375        assert!(registry.contains("ask"));
376        assert!(registry.contains("lsp"));
377        assert!(registry.contains("Skill"));
378        assert!(registry.contains("Task"));
379        assert!(registry.contains("TaskOutput"));
380        assert!(registry.contains("KillShell"));
381        assert!(registry.contains("TodoWrite"));
382        assert!(registry.contains("NotebookEdit"));
383        assert!(registry.contains("EnterPlanMode"));
384        assert!(registry.contains("ExitPlanMode"));
385        assert!(registry.contains("WebFetch"));
386        assert!(registry.contains("WebSearch"));
387        assert!(registry.contains("analyze_image"));
388        assert!(registry.contains("three_stage_workflow"));
389    }
390
391    #[test]
392    fn test_shared_history_is_shared() {
393        let mut registry = ToolRegistry::new();
394        let (history, _hook_manager) = register_default_tools(&mut registry);
395
396        // The history should be empty initially
397        assert!(history.read().unwrap().is_empty());
398
399        // We can write to it
400        {
401            let mut write_guard = history.write().unwrap();
402            write_guard.record_read(FileReadRecord::new(
403                std::path::PathBuf::from("/tmp/test.txt"),
404                "hash123".to_string(),
405                100,
406            ));
407        }
408
409        // And read from it
410        assert!(history
411            .read()
412            .unwrap()
413            .has_read(&std::path::PathBuf::from("/tmp/test.txt")));
414    }
415
416    #[test]
417    fn test_tool_registration_config_builder() {
418        let config = ToolRegistrationConfig::new().with_pdf_enabled(true);
419
420        assert!(config.pdf_enabled);
421        assert!(config.ask_callback.is_none());
422        assert!(config.lsp_callback.is_none());
423    }
424}
aster/tools/mod.rs

aster/tools/
mod.rs
