From 379d8a5ffbb4060c543464b2c02e6cab87a46a36 Mon Sep 17 00:00:00 2001 From: wsp Date: Wed, 18 Mar 2026 18:00:00 +0800 Subject: [PATCH] fix: unify snapshot-aware file tools and avoid replaying file completion events Make the global tool registry the single source of truth for snapshot-aware file modification tools so tool definitions, permission checks, and runtime execution all use the same wrapped implementations. This restores snapshot tracking for new file changes and prevents restored file operation cards from re-emitting completion events that spam diff refresh errors. --- src/crates/core/src/agentic/tools/registry.rs | 60 +++++++------- .../core/src/service/snapshot/manager.rs | 80 ++++--------------- src/crates/core/src/service/snapshot/mod.rs | 2 +- .../tool-cards/FileOperationToolCard.tsx | 34 ++++++-- 4 files changed, 72 insertions(+), 104 deletions(-) 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 = {