diff --git a/src/crates/core/src/agentic/tools/registry.rs b/src/crates/core/src/agentic/tools/registry.rs index ccf81f3c..14a52b26 100644 --- a/src/crates/core/src/agentic/tools/registry.rs +++ b/src/crates/core/src/agentic/tools/registry.rs @@ -49,7 +49,7 @@ impl ToolRegistry { ); } - self.tools.insert(name.clone(), tool); + self.register_tool(tool); debug!("MCP tool registered: tool_name={}", name); } @@ -146,6 +146,9 @@ impl ToolRegistry { /// Register a single tool pub fn register_tool(&mut self, tool: Arc) { + // Snapshot-aware wrapping happens once at registration time so every + // subsequent lookup returns the same runtime implementation. + let tool = crate::service::snapshot::wrap_tool_for_snapshot_tracking(tool); let name = tool.name().to_string(); self.tools.insert(name, tool); } @@ -173,6 +176,7 @@ impl ToolRegistry { #[cfg(test)] mod tests { use super::create_tool_registry; + use serde_json::json; #[test] fn registry_includes_webfetch_tool() { @@ -185,27 +189,32 @@ mod tests { let registry = create_tool_registry(); assert!(registry.get_tool("Cron").is_some()); } + + #[test] + fn registry_wraps_file_modification_tools_for_snapshot_tracking() { + let registry = create_tool_registry(); + let tool = registry + .get_tool("Write") + .expect("Write tool should be registered"); + + let assistant_text = tool.render_result_for_assistant(&json!({ + "success": true, + "file_path": "E:/Projects/demo.txt" + })); + + assert!( + assistant_text.contains("snapshot system"), + "expected snapshot wrapper text, got: {}", + assistant_text + ); + } } -/// Get all tools -/// If you need **always include** MCP tools, use [get_all_registered_tools] +/// Get all tools from the snapshot-aware global registry. pub async fn get_all_tools() -> Vec> { let registry = get_global_tool_registry(); let registry_lock = registry.read().await; - let all_tools = registry_lock.get_all_tools(); - let wrapped_tools = crate::service::snapshot::get_snapshot_wrapped_tools(); - let file_tool_names: std::collections::HashSet = wrapped_tools - .iter() - .map(|tool| tool.name().to_string()) - .collect(); - - let mut result = wrapped_tools; - for tool in all_tools { - if !file_tool_names.contains(tool.name()) { - result.push(tool); - } - } - result + registry_lock.get_all_tools() } /// Get readonly tools @@ -243,22 +252,9 @@ pub fn get_global_tool_registry() -> Arc> { .clone() } -/// Get all registered tools (**always include** dynamically registered MCP tools) +/// Backward-compatible alias for callers that expect MCP tools to be included. pub async fn get_all_registered_tools() -> Vec> { - let registry = get_global_tool_registry(); - let registry_lock = registry.read().await; - let all_tools = registry_lock.get_all_tools(); - let wrapped_tools = crate::service::snapshot::get_snapshot_wrapped_tools(); - let file_tool_names: std::collections::HashSet = - wrapped_tools.iter().map(|t| t.name().to_string()).collect(); - - let mut result = wrapped_tools; - for tool in all_tools { - if !file_tool_names.contains(tool.name()) { - result.push(tool); - } - } - result + get_all_tools().await } /// Get all registered tool names diff --git a/src/crates/core/src/service/snapshot/manager.rs b/src/crates/core/src/service/snapshot/manager.rs index 208da6e2..b0502e95 100644 --- a/src/crates/core/src/service/snapshot/manager.rs +++ b/src/crates/core/src/service/snapshot/manager.rs @@ -7,7 +7,7 @@ use crate::service::snapshot::types::{ use async_trait::async_trait; use log::{debug, error, info, warn}; use serde_json::Value; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::{Arc, OnceLock, RwLock as StdRwLock}; use tokio::sync::RwLock; @@ -17,9 +17,6 @@ use tokio::sync::RwLock; /// Manages all components of the snapshot system. pub struct SnapshotManager { snapshot_service: Arc>, - original_tools: Vec>, - file_modification_tools: HashSet, - initialized: bool, } impl SnapshotManager { @@ -36,56 +33,7 @@ impl SnapshotManager { let mut snapshot_service = SnapshotService::new(workspace_dir, config); snapshot_service.initialize().await?; let snapshot_service = Arc::new(RwLock::new(snapshot_service)); - - let original_tools = ToolRegistry::new().get_all_tools(); - - let file_modification_tools = [ - "Write", - "Edit", - "Delete", - "write_file", - "edit_file", - "create_file", - "delete_file", - "rename_file", - "move_file", - ] - .iter() - .map(|s| s.to_string()) - .collect(); - - Ok(Self { - snapshot_service, - original_tools, - file_modification_tools, - initialized: true, - }) - } - - /// Returns whether the tool modifies files. - fn is_file_modification_tool(&self, tool_name: &str) -> bool { - self.file_modification_tools.contains(tool_name) - } - - /// Returns wrapped tool list. - pub fn get_wrapped_tools(&self) -> Vec> { - if !self.initialized { - error!("Snapshot manager not initialized"); - return vec![]; - } - - let mut wrapped_tools: Vec> = Vec::new(); - - for tool in &self.original_tools { - if self.is_file_modification_tool(tool.name()) { - let wrapped_tool: Arc = Arc::new(WrappedTool::new(tool.clone())); - wrapped_tools.push(wrapped_tool); - } else { - wrapped_tools.push(tool.clone()); - } - } - - wrapped_tools + Ok(Self { snapshot_service }) } /// Records a file change. @@ -340,18 +288,20 @@ fn snapshot_managers() -> &'static StdRwLock) -> Arc { + if WrappedTool::is_file_modification_tool_name(tool.name()) { + Arc::new(WrappedTool::new(tool)) + } else { + tool + } +} + +/// Compatibility helper that returns a fresh snapshot-aware tool list. pub fn get_snapshot_wrapped_tools() -> Vec> { - ToolRegistry::new() - .get_all_tools() - .into_iter() - .map(|tool| { - if WrappedTool::is_file_modification_tool_name(tool.name()) { - Arc::new(WrappedTool::new(tool)) as Arc - } else { - tool - } - }) - .collect() + ToolRegistry::new().get_all_tools() } /// Wrapped tool diff --git a/src/crates/core/src/service/snapshot/mod.rs b/src/crates/core/src/service/snapshot/mod.rs index 5fa50ca5..64dd4b38 100644 --- a/src/crates/core/src/service/snapshot/mod.rs +++ b/src/crates/core/src/service/snapshot/mod.rs @@ -14,7 +14,7 @@ pub use events::{ pub use manager::{ ensure_snapshot_manager_for_workspace, get_or_create_snapshot_manager, get_snapshot_manager_for_workspace, get_snapshot_wrapped_tools, - initialize_snapshot_manager_for_workspace, SnapshotManager, + initialize_snapshot_manager_for_workspace, wrap_tool_for_snapshot_tracking, SnapshotManager, }; pub use service::{SnapshotService, SystemStats}; pub use snapshot_core::{FileChangeEntry, FileChangeQueue, SessionStats, SnapshotCore}; diff --git a/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx b/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx index b774cfb2..ee35ff2a 100644 --- a/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx +++ b/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx @@ -42,6 +42,8 @@ export const FileOperationToolCard: React.FC = ({ const prevIsParamsStreamingRef = useRef(isParamsStreaming); const userCollapsedRef = useRef(false); + const hasInitializedCompletionEffectRef = useRef(false); + const previousCompletionEndTimeRef = useRef(toolItem.endTime ?? null); useEffect(() => { const prevIsParamsStreaming = prevIsParamsStreamingRef.current; @@ -108,13 +110,33 @@ export const FileOperationToolCard: React.FC = ({ const currentFile = files.find(f => f.filePath === currentFilePath); useEffect(() => { - if (status === 'completed' && toolResult?.success && sessionId && currentFilePath) { - eventBus.emit(SNAPSHOT_EVENTS.FILE_OPERATION_COMPLETED, { - toolName: toolItem.toolName, - toolResult - }, sessionId, currentFilePath); + const completionEndTime = toolItem.endTime ?? null; + const isCompletedSuccess = status === 'completed' && Boolean(toolResult?.success); + + if (!hasInitializedCompletionEffectRef.current) { + hasInitializedCompletionEffectRef.current = true; + previousCompletionEndTimeRef.current = completionEndTime; + return; + } + + const shouldEmitCompletionEvent = + isCompletedSuccess && + completionEndTime !== null && + previousCompletionEndTimeRef.current !== completionEndTime && + Boolean(sessionId) && + Boolean(currentFilePath); + + previousCompletionEndTimeRef.current = completionEndTime; + + if (!shouldEmitCompletionEvent || !sessionId || !currentFilePath) { + return; } - }, [status, toolResult, sessionId, currentFilePath, toolItem.toolName, eventBus]); + + eventBus.emit(SNAPSHOT_EVENTS.FILE_OPERATION_COMPLETED, { + toolName: toolItem.toolName, + toolResult + }, sessionId, currentFilePath); + }, [status, toolResult, sessionId, currentFilePath, toolItem.toolName, toolItem.endTime, eventBus]); const getToolDisplayInfo = () => { const toolMap: Record = {