From 574aef1d77b0f0ba74d6a0d077b4cb18de43ec83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 29 Jul 2025 12:47:27 +0200 Subject: [PATCH 1/7] WIP: New "edit" tool with old_text and new_text Is the new default. The replace_in_file tool (using diff blocks with SEARCH/REPLACE) still exists and can be enabled via CLI flag. --- crates/code_assistant/src/agent/tests.rs | 40 +- crates/code_assistant/src/main.rs | 21 +- .../code_assistant/src/tools/core/registry.rs | 28 +- crates/code_assistant/src/tools/impls/edit.rs | 556 ++++++++++++++++++ crates/code_assistant/src/tools/impls/mod.rs | 2 + .../src/tools/integration_tests.rs | 102 ++++ crates/code_assistant/src/tools/mod.rs | 3 + crates/code_assistant/src/tools/parse.rs | 17 +- .../src/ui/gpui/diff_renderer.rs | 114 +++- 9 files changed, 823 insertions(+), 60 deletions(-) create mode 100644 crates/code_assistant/src/tools/impls/edit.rs create mode 100644 crates/code_assistant/src/tools/integration_tests.rs diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index 25eb2e6d..9f01b1f4 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -53,27 +53,16 @@ fn test_replacement_xml_parsing() -> Result<()> { let text = concat!( "I will fix the code formatting.\n", "\n", - "\n", + "\n", "test\n", "src/main.rs\n", - "\n", - "<<<<<<< SEARCH\n", - "function test(){\n", + "function test(){\n", " console.log(\"messy\");\n", - "}\n", - "=======\n", - "function test() {\n", + "}\n", + "function test() {\n", " console.log(\"clean\");\n", - "}\n", - ">>>>>>> REPLACE\n", - "\n", - "<<<<<<< SEARCH\n", - "const x=42\n", - "=======\n", - "const x = 42;\n", - ">>>>>>> REPLACE\n", - "\n", - "\n", + "}\n", + "\n", ) .to_string(); let response = LLMResponse { @@ -89,7 +78,7 @@ fn test_replacement_xml_parsing() -> Result<()> { assert_eq!(tool_requests.len(), 1); let request = &tool_requests[0]; - assert_eq!(request.name, "replace_in_file"); + assert_eq!(request.name, "edit"); assert_eq!( request.input.get("project").unwrap().as_str().unwrap(), "test" @@ -99,14 +88,13 @@ fn test_replacement_xml_parsing() -> Result<()> { "src/main.rs" ); - let diff = request.input.get("diff").unwrap().as_str().unwrap(); - assert!(diff.contains("<<<<<<< SEARCH")); - assert!(diff.contains("function test(){")); - assert!(diff.contains("console.log(\"messy\")")); - assert!(diff.contains("function test() {")); - assert!(diff.contains("console.log(\"clean\")")); - assert!(diff.contains("const x=42")); - assert!(diff.contains("const x = 42;")); + let old_text = request.input.get("old_text").unwrap().as_str().unwrap(); + assert!(old_text.contains("function test(){")); + assert!(old_text.contains("console.log(\"messy\")")); + + let new_text = request.input.get("new_text").unwrap().as_str().unwrap(); + assert!(new_text.contains("function test() {")); + assert!(new_text.contains("console.log(\"clean\")")); Ok(()) } diff --git a/crates/code_assistant/src/main.rs b/crates/code_assistant/src/main.rs index 25022d36..5d1e2bd9 100644 --- a/crates/code_assistant/src/main.rs +++ b/crates/code_assistant/src/main.rs @@ -106,6 +106,10 @@ struct Args { /// Fast playback mode - ignore chunk timing when playing recordings #[arg(long)] fast_playback: bool, + + /// Use the legacy diff format for file editing (enables replace_in_file tool) + #[arg(long)] + use_diff_format: bool, } #[derive(Subcommand, Debug)] @@ -115,6 +119,9 @@ enum Mode { /// Enable verbose logging #[arg(short, long)] verbose: bool, + /// Use the legacy diff format for file editing (enables replace_in_file tool) + #[arg(long)] + use_diff_format: bool, }, } @@ -332,7 +339,12 @@ fn setup_logging(verbose: bool, use_stdout: bool) { subscriber.init(); } -async fn run_mcp_server(verbose: bool) -> Result<()> { +async fn run_mcp_server(verbose: bool, use_diff_format: bool) -> Result<()> { + // Set environment variable for diff format preference before tools are initialized + if use_diff_format { + std::env::set_var("CODE_ASSISTANT_USE_DIFF_FORMAT", "true"); + } + // Setup logging based on verbose flag setup_logging(verbose, false); @@ -636,6 +648,11 @@ fn run_agent_gpui( } async fn run_agent(args: Args) -> Result<()> { + // Set environment variable for diff format preference before tools are initialized + if args.use_diff_format { + std::env::set_var("CODE_ASSISTANT_USE_DIFF_FORMAT", "true"); + } + // Get all the agent options from args let path = args.path.clone().unwrap_or_else(|| PathBuf::from(".")); let task = args.task.clone(); @@ -972,7 +989,7 @@ async fn main() -> Result<()> { match args.mode { // Server mode - Some(Mode::Server { verbose }) => run_mcp_server(verbose).await, + Some(Mode::Server { verbose, use_diff_format }) => run_mcp_server(verbose, use_diff_format).await, // Agent mode (default) None => run_agent(args).await, diff --git a/crates/code_assistant/src/tools/core/registry.rs b/crates/code_assistant/src/tools/core/registry.rs index b1ddc3ad..0414061b 100644 --- a/crates/code_assistant/src/tools/core/registry.rs +++ b/crates/code_assistant/src/tools/core/registry.rs @@ -17,7 +17,11 @@ impl ToolRegistry { static INSTANCE: OnceLock = OnceLock::new(); INSTANCE.get_or_init(|| { let mut registry = ToolRegistry::new(); - registry.register_default_tools(); + // Check environment variable for diff format preference + let use_diff_format = std::env::var("CODE_ASSISTANT_USE_DIFF_FORMAT") + .map(|v| v.to_lowercase() == "true") + .unwrap_or(false); + registry.register_default_tools_impl(use_diff_format); registry }) } @@ -68,15 +72,21 @@ impl ToolRegistry { /// Register all default tools in the system /// This will be expanded as we implement more tools - fn register_default_tools(&mut self) { + #[cfg(test)] + pub fn register_default_tools(&mut self, use_diff_format: bool) { + self.register_default_tools_impl(use_diff_format); + } + + /// Internal implementation of register_default_tools + fn register_default_tools_impl(&mut self, use_diff_format: bool) { // Import all tools use crate::tools::impls::{ - DeleteFilesTool, ExecuteCommandTool, ListFilesTool, ListProjectsTool, NameSessionTool, + DeleteFilesTool, EditTool, ExecuteCommandTool, ListFilesTool, ListProjectsTool, NameSessionTool, PerplexityAskTool, ReadFilesTool, ReplaceInFileTool, SearchFilesTool, WebFetchTool, WebSearchTool, WriteFileTool, }; - // Register tools + // Register core tools self.register(Box::new(DeleteFilesTool)); self.register(Box::new(ExecuteCommandTool)); self.register(Box::new(ListFilesTool)); @@ -84,12 +94,20 @@ impl ToolRegistry { self.register(Box::new(NameSessionTool)); self.register(Box::new(PerplexityAskTool)); self.register(Box::new(ReadFilesTool)); - self.register(Box::new(ReplaceInFileTool)); self.register(Box::new(SearchFilesTool)); self.register(Box::new(WebFetchTool)); self.register(Box::new(WebSearchTool)); self.register(Box::new(WriteFileTool)); + // Register file editing tools based on configuration + if use_diff_format { + // Use legacy diff format (replace_in_file tool) + self.register(Box::new(ReplaceInFileTool)); + } else { + // Use new edit tool (default) + self.register(Box::new(EditTool)); + } + // More tools will be added here as they are implemented } } diff --git a/crates/code_assistant/src/tools/impls/edit.rs b/crates/code_assistant/src/tools/impls/edit.rs new file mode 100644 index 00000000..0d8ed633 --- /dev/null +++ b/crates/code_assistant/src/tools/impls/edit.rs @@ -0,0 +1,556 @@ +use crate::tools::core::{ + Render, ResourcesTracker, Tool, ToolContext, ToolResult, ToolScope, ToolSpec, +}; +use crate::types::{FileReplacement, LoadedResource}; +use crate::utils::FileUpdaterError; +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; +use serde_json::json; +use std::path::PathBuf; + +// Input type for the edit tool +#[derive(Deserialize)] +pub struct EditInput { + pub project: String, + pub path: String, + pub old_text: String, + pub new_text: String, + #[serde(default)] + pub replace_all: Option, +} + +// Output type +#[derive(Serialize, Deserialize)] +pub struct EditOutput { + pub project: String, + pub path: PathBuf, + pub error: Option, +} + +// Render implementation for output formatting +impl Render for EditOutput { + fn status(&self) -> String { + if self.error.is_none() { + format!( + "Successfully edited content in file: {}", + self.path.display() + ) + } else { + format!("Failed to edit content in file: {}", self.path.display()) + } + } + + fn render(&self, _tracker: &mut ResourcesTracker) -> String { + if let Some(error) = &self.error { + match error { + FileUpdaterError::SearchBlockNotFound(idx, _) => { + format!( + "Could not find the text to replace (search block {}). Please check that the old_text matches exactly what's in the file.", + idx + ) + } + FileUpdaterError::MultipleMatches(count, idx, _) => { + format!( + "Found {} occurrences of the text to replace (search block {})\nThe old_text must match exactly one location. Try making the old_text more specific.", + count, idx + ) + } + FileUpdaterError::Other(msg) => { + format!("Failed to edit file '{}': {}", self.path.display(), msg) + } + } + } else { + format!( + "Successfully edited content in file '{}'", + self.path.display() + ) + } + } +} + +// ToolResult implementation +impl ToolResult for EditOutput { + fn is_success(&self) -> bool { + self.error.is_none() + } +} + +// Tool implementation +pub struct EditTool; + +#[async_trait::async_trait] +impl Tool for EditTool { + type Input = EditInput; + type Output = EditOutput; + + fn spec(&self) -> ToolSpec { + let description = concat!( + "Edit a file by replacing specific text content. ", + "This tool finds the exact text specified in old_text and replaces it with new_text. ", + "By default, the old_text must match exactly one location in the file. ", + "Set replace_all to true to replace all occurrences of the pattern.", + ); + ToolSpec { + name: "edit", + description, + parameters_schema: json!({ + "type": "object", + "properties": { + "project": { + "type": "string", + "description": "Name of the project containing the file" + }, + "path": { + "type": "string", + "description": "Path to the file to modify (relative to project root)" + }, + "old_text": { + "type": "string", + "description": "The exact text content to find and replace. This must match exactly what appears in the file, including whitespace and line breaks. The search is case-sensitive and whitespace-sensitive." + }, + "new_text": { + "type": "string", + "description": "The text content to replace the old_text with. Can be empty to delete the old_text. Maintains the same indentation and formatting as needed." + }, + "replace_all": { + "type": "boolean", + "description": "Optional. If true, replace all occurrences of old_text. If false or omitted, old_text must match exactly one location (default: false)." + } + }, + "required": ["project", "path", "old_text", "new_text"] + }), + annotations: Some(json!({ + "readOnlyHint": false, + "destructiveHint": true + })), + supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + hidden: false, + } + } + + async fn execute<'a>( + &self, + context: &mut ToolContext<'a>, + input: Self::Input, + ) -> Result { + // Get explorer for the specified project + let explorer = context + .project_manager + .get_explorer_for_project(&input.project) + .map_err(|e| { + anyhow!( + "Failed to get explorer for project {}: {}", + input.project, + e + ) + })?; + + // Check for absolute path + let path = PathBuf::from(&input.path); + if path.is_absolute() { + return Ok(EditOutput { + project: input.project, + path, + error: Some(FileUpdaterError::Other( + "Absolute paths are not allowed".to_string(), + )), + }); + } + + // Join with root_dir to get full path + let full_path = explorer.root_dir().join(&path); + + // Create a FileReplacement from the input + let replacement = FileReplacement { + search: input.old_text, + replace: input.new_text, + replace_all: input.replace_all.unwrap_or(false), + }; + + // Apply the replacements using the explorer + match explorer.apply_replacements(&full_path, &[replacement]) { + Ok(new_content) => { + // If we have a working memory reference, update it with the modified file + if let Some(working_memory) = &mut context.working_memory { + // Add the file with new content to working memory + working_memory.loaded_resources.insert( + (input.project.clone(), path.clone()), + LoadedResource::File(new_content.clone()), + ); + } + + Ok(EditOutput { + project: input.project, + path, + error: None, + }) + } + Err(e) => { + // Extract FileUpdaterError if present + let error = if let Some(file_err) = e.downcast_ref::() { + file_err.clone() + } else { + FileUpdaterError::Other(e.to_string()) + }; + + Ok(EditOutput { + project: input.project, + path, + error: Some(error), + }) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tests::mocks::{MockCommandExecutor, MockExplorer, MockProjectManager}; + use crate::types::WorkingMemory; + use std::collections::HashMap; + + #[tokio::test] + async fn test_edit_output_rendering() { + // Success case + let output = EditOutput { + project: "test-project".to_string(), + path: PathBuf::from("src/test.rs"), + error: None, + }; + + let mut tracker = ResourcesTracker::new(); + let rendered = output.render(&mut tracker); + assert!(rendered.contains("Successfully edited content")); + assert!(rendered.contains("src/test.rs")); + + // Error case with text not found + let output_error = EditOutput { + project: "test-project".to_string(), + path: PathBuf::from("src/test.rs"), + error: Some(FileUpdaterError::SearchBlockNotFound( + 0, + "missing content".to_string(), + )), + }; + + let rendered_error = output_error.render(&mut tracker); + assert!(rendered_error.contains("Could not find the text to replace")); + assert!(rendered_error.contains("old_text matches exactly")); + + // Error case with multiple matches + let output_multiple = EditOutput { + project: "test-project".to_string(), + path: PathBuf::from("src/test.rs"), + error: Some(FileUpdaterError::MultipleMatches( + 3, + 0, + "common pattern".to_string(), + )), + }; + + let rendered_multiple = output_multiple.render(&mut tracker); + assert!(rendered_multiple.contains("Found 3 occurrences")); + assert!(rendered_multiple.contains("more specific")); + } + + #[tokio::test] + async fn test_edit_basic_replacement() -> Result<()> { + // Create a mock project manager and setup test files + let mut files = HashMap::new(); + files.insert( + PathBuf::from("./root/test.rs"), + "fn original() {\n println!(\"Original\");\n}".to_string(), + ); + + let explorer = MockExplorer::new(files, None); + + let project_manager = Box::new(MockProjectManager::default().with_project( + "test-project", + PathBuf::from("./root"), + Box::new(explorer), + )); + + // Create a command executor + let command_executor = Box::new(MockCommandExecutor::new(vec![])); + + // Create working memory + let mut working_memory = WorkingMemory::default(); + + // Create a tool context with working memory + let mut context = ToolContext { + project_manager: project_manager.as_ref(), + command_executor: command_executor.as_ref(), + working_memory: Some(&mut working_memory), + }; + + // Create input for a valid replacement + let input = EditInput { + project: "test-project".to_string(), + path: "test.rs".to_string(), + old_text: "fn original() {\n println!(\"Original\");\n}".to_string(), + new_text: "fn renamed() {\n println!(\"Updated\");\n}".to_string(), + replace_all: None, + }; + + // Execute the tool + let tool = EditTool; + let result = tool.execute(&mut context, input).await?; + + // Verify the result + assert!(result.error.is_none()); + + // Verify that working memory was updated + assert_eq!(working_memory.loaded_resources.len(), 1); + + // Verify the content in working memory + let key = ("test-project".to_string(), PathBuf::from("test.rs")); + if let Some(LoadedResource::File(content)) = working_memory.loaded_resources.get(&key) { + assert!(content.contains("fn renamed()")); + assert!(content.contains("println!(\"Updated\")")); + } else { + panic!("File not found in working memory or wrong resource type"); + } + + Ok(()) + } + + #[tokio::test] + async fn test_edit_replace_all() -> Result<()> { + // Create a mock project manager with test files + let mut files = HashMap::new(); + files.insert( + PathBuf::from("./root/test.js"), + "console.log('test1');\nconsole.log('test2');\nconsole.log('test3');".to_string(), + ); + + let explorer = MockExplorer::new(files, None); + + let project_manager = Box::new(MockProjectManager::default().with_project( + "test-project", + PathBuf::from("./root"), + Box::new(explorer), + )); + + // Create a command executor + let command_executor = Box::new(MockCommandExecutor::new(vec![])); + + // Create working memory + let mut working_memory = WorkingMemory::default(); + + // Create a tool context with working memory + let mut context = ToolContext { + project_manager: project_manager.as_ref(), + command_executor: command_executor.as_ref(), + working_memory: Some(&mut working_memory), + }; + + // Create input for replace all + let input = EditInput { + project: "test-project".to_string(), + path: "test.js".to_string(), + old_text: "console.log(".to_string(), + new_text: "logger.debug(".to_string(), + replace_all: Some(true), + }; + + // Execute the tool + let tool = EditTool; + let result = tool.execute(&mut context, input).await?; + + // Verify the result + assert!(result.error.is_none()); + + // Verify the content in working memory + let key = ("test-project".to_string(), PathBuf::from("test.js")); + if let Some(LoadedResource::File(content)) = working_memory.loaded_resources.get(&key) { + assert!(content.contains("logger.debug('test1')")); + assert!(content.contains("logger.debug('test2')")); + assert!(content.contains("logger.debug('test3')")); + assert!(!content.contains("console.log")); + } else { + panic!("File not found in working memory or wrong resource type"); + } + + Ok(()) + } + + #[tokio::test] + async fn test_edit_error_handling() -> Result<()> { + // Create a mock project manager with test files + let mut files = HashMap::new(); + files.insert( + PathBuf::from("./root/test.rs"), + "console.log('test');\nconsole.log('test');\nconsole.log('test');".to_string(), + ); + + let explorer = MockExplorer::new(files, None); + + let project_manager = Box::new(MockProjectManager::default().with_project( + "test-project", + PathBuf::from("./root"), + Box::new(explorer), + )); + + // Create a command executor + let command_executor = Box::new(MockCommandExecutor::new(vec![])); + + // Create a tool context + let mut context = ToolContext { + project_manager: project_manager.as_ref(), + command_executor: command_executor.as_ref(), + working_memory: None, + }; + + // Test case with multiple matches but replace_all = false + let input_multiple = EditInput { + project: "test-project".to_string(), + path: "test.rs".to_string(), + old_text: "console.log".to_string(), + new_text: "console.debug".to_string(), + replace_all: None, // defaults to false + }; + + // Execute the tool - should fail with multiple matches + let tool = EditTool; + let result = tool.execute(&mut context, input_multiple).await?; + + // Verify error for multiple matches + assert!(result.error.is_some()); + if let Some(FileUpdaterError::MultipleMatches(count, _, _)) = result.error { + assert_eq!(count, 3); + } else { + panic!("Expected MultipleMatches error"); + } + + // Test case with missing content + let input_missing = EditInput { + project: "test-project".to_string(), + path: "test.rs".to_string(), + old_text: "non_existent_content".to_string(), + new_text: "replacement".to_string(), + replace_all: None, + }; + + // Execute the tool - should fail with content not found + let result = tool.execute(&mut context, input_missing).await?; + + // Verify error for missing content + assert!(result.error.is_some()); + match &result.error { + Some(FileUpdaterError::SearchBlockNotFound(_, _)) => (), + _ => panic!("Expected SearchBlockNotFound error"), + } + + Ok(()) + } + + #[tokio::test] + async fn test_edit_empty_replacement() -> Result<()> { + // Test deleting content by replacing with empty string + let mut files = HashMap::new(); + files.insert( + PathBuf::from("./root/test.rs"), + "fn test() {\n // TODO: Remove this comment\n println!(\"Hello\");\n}" + .to_string(), + ); + + let explorer = MockExplorer::new(files, None); + + let project_manager = Box::new(MockProjectManager::default().with_project( + "test-project", + PathBuf::from("./root"), + Box::new(explorer), + )); + + let command_executor = Box::new(MockCommandExecutor::new(vec![])); + let mut working_memory = WorkingMemory::default(); + + let mut context = ToolContext { + project_manager: project_manager.as_ref(), + command_executor: command_executor.as_ref(), + working_memory: Some(&mut working_memory), + }; + + // Delete the TODO comment + let input = EditInput { + project: "test-project".to_string(), + path: "test.rs".to_string(), + old_text: " // TODO: Remove this comment\n".to_string(), + new_text: "".to_string(), + replace_all: None, + }; + + let tool = EditTool; + let result = tool.execute(&mut context, input).await?; + + // Verify the result + assert!(result.error.is_none()); + + // Verify the content in working memory + let key = ("test-project".to_string(), PathBuf::from("test.rs")); + if let Some(LoadedResource::File(content)) = working_memory.loaded_resources.get(&key) { + assert!(!content.contains("TODO")); + assert!(content.contains("fn test() {")); + assert!(content.contains("println!(\"Hello\");")); + } else { + panic!("File not found in working memory or wrong resource type"); + } + + Ok(()) + } + + #[tokio::test] + async fn test_edit_whitespace_normalization() -> Result<()> { + // Test that whitespace differences are handled correctly + let mut files = HashMap::new(); + files.insert( + PathBuf::from("./root/test.rs"), + "function test() {\r\n console.log('test');\r\n}".to_string(), // CRLF endings + ); + + let explorer = MockExplorer::new(files, None); + + let project_manager = Box::new(MockProjectManager::default().with_project( + "test-project", + PathBuf::from("./root"), + Box::new(explorer), + )); + + let command_executor = Box::new(MockCommandExecutor::new(vec![])); + let mut working_memory = WorkingMemory::default(); + + let mut context = ToolContext { + project_manager: project_manager.as_ref(), + command_executor: command_executor.as_ref(), + working_memory: Some(&mut working_memory), + }; + + // Use LF endings in search text, should still match CRLF in file + let input = EditInput { + project: "test-project".to_string(), + path: "test.rs".to_string(), + old_text: "function test() {\n console.log('test');\n}".to_string(), // LF endings + new_text: "function answer() {\n return 42;\n}".to_string(), + replace_all: None, + }; + + let tool = EditTool; + let result = tool.execute(&mut context, input).await?; + + // Verify the result + assert!(result.error.is_none()); + + // Verify the content in working memory + let key = ("test-project".to_string(), PathBuf::from("test.rs")); + if let Some(LoadedResource::File(content)) = working_memory.loaded_resources.get(&key) { + assert!(content.contains("function answer()")); + assert!(content.contains("return 42;")); + assert!(!content.contains("console.log")); + } else { + panic!("File not found in working memory or wrong resource type"); + } + + Ok(()) + } +} diff --git a/crates/code_assistant/src/tools/impls/mod.rs b/crates/code_assistant/src/tools/impls/mod.rs index 93e39862..3bfada30 100644 --- a/crates/code_assistant/src/tools/impls/mod.rs +++ b/crates/code_assistant/src/tools/impls/mod.rs @@ -1,5 +1,6 @@ // Tool implementations pub mod delete_files; +pub mod edit; pub mod execute_command; pub mod list_files; pub mod list_projects; @@ -14,6 +15,7 @@ pub mod write_file; // Re-export all tools for registration pub use delete_files::DeleteFilesTool; +pub use edit::EditTool; pub use execute_command::ExecuteCommandTool; pub use list_files::ListFilesTool; pub use list_projects::ListProjectsTool; diff --git a/crates/code_assistant/src/tools/integration_tests.rs b/crates/code_assistant/src/tools/integration_tests.rs new file mode 100644 index 00000000..0d76d10e --- /dev/null +++ b/crates/code_assistant/src/tools/integration_tests.rs @@ -0,0 +1,102 @@ +#[cfg(test)] +mod tests { + use crate::tools::core::ToolRegistry; + use std::env; + + #[test] + fn test_default_tool_registration() { + // Clear environment variable to ensure default behavior + env::remove_var("CODE_ASSISTANT_USE_DIFF_FORMAT"); + + // Create a new registry to test default registration + let mut registry = ToolRegistry::new(); + registry.register_default_tools(false); // Default behavior + + // Edit tool should be registered + assert!(registry.get("edit").is_some(), "Edit tool should be registered by default"); + + // Replace in file tool should NOT be registered + assert!(registry.get("replace_in_file").is_none(), "ReplaceInFile tool should not be registered by default"); + + // Other core tools should still be registered + assert!(registry.get("read_files").is_some(), "ReadFiles tool should be registered"); + assert!(registry.get("write_file").is_some(), "WriteFile tool should be registered"); + assert!(registry.get("list_files").is_some(), "ListFiles tool should be registered"); + } + + #[test] + fn test_diff_format_tool_registration() { + // Create a new registry to test diff format registration + let mut registry = ToolRegistry::new(); + registry.register_default_tools(true); // Use diff format + + // Replace in file tool should be registered + assert!(registry.get("replace_in_file").is_some(), "ReplaceInFile tool should be registered with diff format"); + + // Edit tool should NOT be registered + assert!(registry.get("edit").is_none(), "Edit tool should not be registered with diff format"); + + // Other core tools should still be registered + assert!(registry.get("read_files").is_some(), "ReadFiles tool should be registered"); + assert!(registry.get("write_file").is_some(), "WriteFile tool should be registered"); + assert!(registry.get("list_files").is_some(), "ListFiles tool should be registered"); + } + + #[test] + fn test_tool_specs() { + // Test edit tool spec + let mut registry = ToolRegistry::new(); + registry.register_default_tools(false); + + if let Some(edit_tool) = registry.get("edit") { + let spec = edit_tool.spec(); + assert_eq!(spec.name, "edit"); + assert!(spec.description.contains("Edit a file by replacing specific text content")); + + // Check required parameters + let params = &spec.parameters_schema; + let properties = params.get("properties").unwrap().as_object().unwrap(); + assert!(properties.contains_key("project")); + assert!(properties.contains_key("path")); + assert!(properties.contains_key("old_text")); + assert!(properties.contains_key("new_text")); + + // Check that replace_all is optional + let required = params.get("required").unwrap().as_array().unwrap(); + assert!(required.contains(&serde_json::Value::String("project".to_string()))); + assert!(required.contains(&serde_json::Value::String("path".to_string()))); + assert!(required.contains(&serde_json::Value::String("old_text".to_string()))); + assert!(required.contains(&serde_json::Value::String("new_text".to_string()))); + assert!(!required.contains(&serde_json::Value::String("replace_all".to_string()))); + } else { + panic!("Edit tool should be registered"); + } + } + + #[test] + fn test_replace_in_file_tool_specs() { + // Test replace_in_file tool spec + let mut registry = ToolRegistry::new(); + registry.register_default_tools(true); + + if let Some(replace_tool) = registry.get("replace_in_file") { + let spec = replace_tool.spec(); + assert_eq!(spec.name, "replace_in_file"); + assert!(spec.description.contains("Replace sections in a file")); + + // Check required parameters + let params = &spec.parameters_schema; + let properties = params.get("properties").unwrap().as_object().unwrap(); + assert!(properties.contains_key("project")); + assert!(properties.contains_key("path")); + assert!(properties.contains_key("diff")); + + let required = params.get("required").unwrap().as_array().unwrap(); + assert!(required.contains(&serde_json::Value::String("project".to_string()))); + assert!(required.contains(&serde_json::Value::String("path".to_string()))); + assert!(required.contains(&serde_json::Value::String("diff".to_string()))); + } else { + panic!("ReplaceInFile tool should be registered with diff format"); + } + } +} diff --git a/crates/code_assistant/src/tools/mod.rs b/crates/code_assistant/src/tools/mod.rs index c6cc8ced..4af961ac 100644 --- a/crates/code_assistant/src/tools/mod.rs +++ b/crates/code_assistant/src/tools/mod.rs @@ -18,6 +18,9 @@ pub mod impls; #[cfg(test)] mod tests; +#[cfg(test)] +mod integration_tests; + pub use parse::{parse_caret_tool_invocations, parse_xml_tool_invocations}; pub use parser_registry::ParserRegistry; pub use system_message::generate_system_message; diff --git a/crates/code_assistant/src/tools/parse.rs b/crates/code_assistant/src/tools/parse.rs index cba9840f..48793280 100644 --- a/crates/code_assistant/src/tools/parse.rs +++ b/crates/code_assistant/src/tools/parse.rs @@ -1479,24 +1479,25 @@ mod tests { #[tokio::test] async fn test_parse_caret_tool_invocations_multiple_multiline() { let text = concat!( - "^^^replace_in_file\n", + "^^^edit\n", "project: test\n", "path: test.rs\n", - "diff ---\n", + "old_text ---\n", "old code here\n", - "--- diff\n", - "comment ---\n", - "This change fixes a bug\n", - "--- comment\n", + "--- old_text\n", + "new_text ---\n", + "new code here\n", + "--- new_text\n", "^^^" ); let (result, _) = parse_caret_tool_invocations(text, 123, 0, None).unwrap(); assert_eq!(result.len(), 1); - assert_eq!(result[0].name, "replace_in_file"); + assert_eq!(result[0].name, "edit"); assert_eq!(result[0].input["project"], "test"); assert_eq!(result[0].input["path"], "test.rs"); - assert_eq!(result[0].input["diff"], "old code here"); + assert_eq!(result[0].input["old_text"], "old code here"); + assert_eq!(result[0].input["new_text"], "new code here"); } #[tokio::test] diff --git a/crates/code_assistant/src/ui/gpui/diff_renderer.rs b/crates/code_assistant/src/ui/gpui/diff_renderer.rs index 389f6146..79dbbf94 100644 --- a/crates/code_assistant/src/ui/gpui/diff_renderer.rs +++ b/crates/code_assistant/src/ui/gpui/diff_renderer.rs @@ -2,38 +2,74 @@ use crate::ui::gpui::parameter_renderers::ParameterRenderer; use gpui::{div, px, rgb, rgba, Element, FontWeight, IntoElement, ParentElement, Styled}; use similar::{ChangeTag, TextDiff}; -/// Renderer for the "diff" parameter of the replace_in_file tool +/// Renderer for diff-style parameters (replace_in_file "diff" and edit tool "old_text"/"new_text") pub struct DiffParameterRenderer; impl ParameterRenderer for DiffParameterRenderer { fn supported_parameters(&self) -> Vec<(String, String)> { - vec![("replace_in_file".to_string(), "diff".to_string())] + vec![ + ("replace_in_file".to_string(), "diff".to_string()), + ("edit".to_string(), "old_text".to_string()), + ("edit".to_string(), "new_text".to_string()), + ] } fn render( &self, - _tool_name: &str, - _param_name: &str, + tool_name: &str, + param_name: &str, param_value: &str, theme: &gpui_component::theme::Theme, ) -> gpui::AnyElement { - // Container for the diff content - no parameter name shown - div() - .rounded_md() - .bg(if theme.is_dark() { - rgba(0x0A0A0AFF) // Dunklerer Hintergrund im Dark Mode - } else { - rgba(0xEAEAEAFF) // Hellerer Hintergrund im Light Mode - }) - .p_2() - .text_size(px(15.)) - .font_weight(FontWeight(500.0)) - .child(parse_and_render_diff(param_value, theme)) - .into_any() + // Handle different parameter types + match (tool_name, param_name) { + ("replace_in_file", "diff") => { + // Traditional diff format - parse and render as before + div() + .rounded_md() + .bg(if theme.is_dark() { + rgba(0x0A0A0AFF) // Dunklerer Hintergrund im Dark Mode + } else { + rgba(0xEAEAEAFF) // Hellerer Hintergrund im Light Mode + }) + .p_2() + .text_size(px(15.)) + .font_weight(FontWeight(500.0)) + .child(parse_and_render_diff(param_value, theme)) + .into_any() + } + ("edit", "old_text") => { + // Show old_text with deletion styling + render_edit_text_block(param_value, true, theme) + } + ("edit", "new_text") => { + // Show new_text with addition styling + render_edit_text_block(param_value, false, theme) + } + _ => { + // Fallback for unknown parameters + div() + .rounded_md() + .bg(if theme.is_dark() { + rgba(0x0A0A0AFF) + } else { + rgba(0xEAEAEAFF) + }) + .p_2() + .text_size(px(15.)) + .font_weight(FontWeight(500.0)) + .child(param_value.to_string()) + .into_any() + } + } } - fn is_full_width(&self, _tool_name: &str, _param_name: &str) -> bool { - true // Diff parameter is always full-width + fn is_full_width(&self, tool_name: &str, param_name: &str) -> bool { + // All diff-style parameters are full-width + matches!( + (tool_name, param_name), + ("replace_in_file", "diff") | ("edit", "old_text") | ("edit", "new_text") + ) } } @@ -357,3 +393,43 @@ fn render_streaming_diff_section( ]) .into_any() } + +/// Render a text block for the edit tool (old_text or new_text) +fn render_edit_text_block( + text: &str, + is_deletion: bool, + theme: &gpui_component::theme::Theme, +) -> gpui::AnyElement { + // Choose colors based on whether this is old_text (deletion) or new_text (addition) + let (border_color, text_color) = if is_deletion { + // Red for deletions (old_text) + if theme.is_dark() { + (rgb(0xCC5555), rgb(0xFFBBBB)) + } else { + (rgb(0xCC3333), rgb(0xAA0000)) + } + } else { + // Green for additions (new_text) + if theme.is_dark() { + (rgb(0x55CC55), rgb(0xBBFFBB)) + } else { + (rgb(0x33AA33), rgb(0x007700)) + } + }; + + div() + .rounded_md() + .bg(if theme.is_dark() { + rgba(0x0A0A0AFF) // Dunklerer Hintergrund im Dark Mode + } else { + rgba(0xEAEAEAFF) // Hellerer Hintergrund im Light Mode + }) + .p_2() + .border_l_2() + .border_color(border_color) + .text_color(text_color) + .text_size(px(15.)) + .font_weight(FontWeight(500.0)) + .child(text.to_string()) + .into_any() +} From c870ea29df4c69e210d20316c6bc8f3b73d33a24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 29 Jul 2025 14:18:27 +0200 Subject: [PATCH 2/7] Use ToolScope to switch to replace_in_file tool --- .../code_assistant/src/agent/persistence.rs | 12 ++- crates/code_assistant/src/agent/runner.rs | 14 ++- crates/code_assistant/src/agent/tests.rs | 1 + crates/code_assistant/src/main.rs | 34 +++--- crates/code_assistant/src/persistence.rs | 3 + crates/code_assistant/src/session/manager.rs | 13 +++ .../code_assistant/src/tools/core/registry.rs | 33 ++---- crates/code_assistant/src/tools/core/spec.rs | 2 + .../src/tools/impls/delete_files.rs | 6 +- .../src/tools/impls/execute_command.rs | 6 +- .../src/tools/impls/list_files.rs | 6 +- .../src/tools/impls/name_session.rs | 2 +- .../src/tools/impls/perplexity_ask.rs | 6 +- .../src/tools/impls/read_files.rs | 6 +- .../src/tools/impls/replace_in_file.rs | 2 +- .../src/tools/impls/search_files.rs | 6 +- .../src/tools/impls/web_fetch.rs | 6 +- .../src/tools/impls/web_search.rs | 6 +- .../src/tools/impls/write_file.rs | 6 +- .../src/tools/integration_tests.rs | 102 ------------------ crates/code_assistant/src/tools/mod.rs | 3 - 21 files changed, 111 insertions(+), 164 deletions(-) delete mode 100644 crates/code_assistant/src/tools/integration_tests.rs diff --git a/crates/code_assistant/src/agent/persistence.rs b/crates/code_assistant/src/agent/persistence.rs index 97a0041c..5c489e5e 100644 --- a/crates/code_assistant/src/agent/persistence.rs +++ b/crates/code_assistant/src/agent/persistence.rs @@ -17,6 +17,7 @@ pub trait AgentStatePersistence: Send + Sync { /// Save the current agent state fn save_agent_state( &mut self, + name: String, messages: Vec, tool_executions: Vec, working_memory: WorkingMemory, @@ -51,6 +52,7 @@ impl MockStatePersistence { impl AgentStatePersistence for MockStatePersistence { fn save_agent_state( &mut self, + _name: String, messages: Vec, tool_executions: Vec, working_memory: WorkingMemory, @@ -86,6 +88,7 @@ impl SessionStatePersistence { impl AgentStatePersistence for SessionStatePersistence { fn save_agent_state( &mut self, + name: String, messages: Vec, tool_executions: Vec, working_memory: WorkingMemory, @@ -100,6 +103,7 @@ impl AgentStatePersistence for SessionStatePersistence { session_manager.save_session_state( &self.session_id, + name, messages, tool_executions, working_memory, @@ -118,15 +122,17 @@ const STATE_FILE: &str = ".code-assistant.state.json"; pub struct FileStatePersistence { state_file_path: PathBuf, tool_syntax: ToolSyntax, + use_diff_blocks: bool, } impl FileStatePersistence { - pub fn new(working_dir: &Path, tool_syntax: ToolSyntax) -> Self { + pub fn new(working_dir: &Path, tool_syntax: ToolSyntax, use_diff_blocks: bool) -> Self { let state_file_path = working_dir.join(STATE_FILE); info!("Using state file: {}", state_file_path.display()); Self { state_file_path, tool_syntax, + use_diff_blocks, } } @@ -163,6 +169,7 @@ impl FileStatePersistence { impl AgentStatePersistence for FileStatePersistence { fn save_agent_state( &mut self, + name: String, messages: Vec, tool_executions: Vec, working_memory: WorkingMemory, @@ -181,7 +188,7 @@ impl AgentStatePersistence for FileStatePersistence { // Create a ChatSession with the current state let session = ChatSession { id: "terminal-session".to_string(), - name: "Terminal Session".to_string(), + name: name, created_at: SystemTime::now(), updated_at: SystemTime::now(), messages, @@ -190,6 +197,7 @@ impl AgentStatePersistence for FileStatePersistence { init_path, initial_project, tool_syntax: self.tool_syntax, + use_diff_blocks: self.use_diff_blocks, next_request_id, }; diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index 21f6a94b..bda4e243 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -34,6 +34,7 @@ pub struct Agent { working_memory: WorkingMemory, llm_provider: Box, tool_syntax: ToolSyntax, + tool_scope: ToolScope, project_manager: Box, command_executor: Box, ui: Arc>, @@ -91,6 +92,7 @@ impl Agent { working_memory: WorkingMemory::default(), llm_provider, tool_syntax, + tool_scope: ToolScope::Agent, // Default to Agent scope project_manager, ui, command_executor, @@ -108,6 +110,13 @@ impl Agent { } } + /// Enable diff blocks format for file editing (uses replace_in_file tool instead of edit) + pub fn enable_diff_blocks(&mut self) { + self.tool_scope = ToolScope::AgentWithDiffBlocks; + // Clear cached system message so it gets regenerated with the new scope + self.cached_system_message = OnceLock::new(); + } + /// Set the shared pending message reference from SessionInstance pub fn set_pending_message_ref(&mut self, pending_ref: Arc>>) { self.pending_message_ref = Some(pending_ref); @@ -185,6 +194,7 @@ impl Agent { self.message_history.len() ); self.state_persistence.save_agent_state( + self.session_name.clone(), self.message_history.clone(), self.tool_executions.clone(), self.working_memory.clone(), @@ -605,7 +615,7 @@ impl Agent { } // Generate the system message using the tools module - let mut system_message = generate_system_message(self.tool_syntax, ToolScope::Agent); + let mut system_message = generate_system_message(self.tool_syntax, self.tool_scope); // Add project information let mut project_info = String::new(); @@ -784,7 +794,7 @@ impl Agent { tools: match self.tool_syntax { ToolSyntax::Native => { Some(crate::tools::AnnotatedToolDefinition::to_tool_definitions( - ToolRegistry::global().get_tool_definitions_for_scope(ToolScope::Agent), + ToolRegistry::global().get_tool_definitions_for_scope(self.tool_scope), )) } ToolSyntax::Xml => None, diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index 9f01b1f4..a9d9f635 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -785,6 +785,7 @@ fn test_ui_filtering_with_failed_tool_messages() -> Result<()> { init_path: None, initial_project: None, tool_syntax: ToolSyntax::Xml, + use_diff_blocks: false, next_request_id: 1, }; diff --git a/crates/code_assistant/src/main.rs b/crates/code_assistant/src/main.rs index 5d1e2bd9..88c1f957 100644 --- a/crates/code_assistant/src/main.rs +++ b/crates/code_assistant/src/main.rs @@ -107,7 +107,7 @@ struct Args { #[arg(long)] fast_playback: bool, - /// Use the legacy diff format for file editing (enables replace_in_file tool) + /// Use the legacy diff format for file editing (enables replace_in_file tool instead of edit) #[arg(long)] use_diff_format: bool, } @@ -119,9 +119,6 @@ enum Mode { /// Enable verbose logging #[arg(short, long)] verbose: bool, - /// Use the legacy diff format for file editing (enables replace_in_file tool) - #[arg(long)] - use_diff_format: bool, }, } @@ -339,12 +336,7 @@ fn setup_logging(verbose: bool, use_stdout: bool) { subscriber.init(); } -async fn run_mcp_server(verbose: bool, use_diff_format: bool) -> Result<()> { - // Set environment variable for diff format preference before tools are initialized - if use_diff_format { - std::env::set_var("CODE_ASSISTANT_USE_DIFF_FORMAT", "true"); - } - +async fn run_mcp_server(verbose: bool) -> Result<()> { // Setup logging based on verbose flag setup_logging(verbose, false); @@ -363,6 +355,7 @@ async fn run_agent_terminal( aicore_config: Option, num_ctx: usize, tool_syntax: ToolSyntax, + use_diff_format: bool, record: Option, playback: Option, fast_playback: bool, @@ -370,7 +363,7 @@ async fn run_agent_terminal( let root_path = path.canonicalize()?; // Create file persistence for simple state management - let file_persistence = FileStatePersistence::new(&root_path, tool_syntax); + let file_persistence = FileStatePersistence::new(&root_path, tool_syntax, use_diff_format); // Setup dynamic types let project_manager = Box::new(DefaultProjectManager::new()); @@ -403,6 +396,11 @@ async fn run_agent_terminal( Some(root_path.clone()), ); + // Configure diff blocks format if requested + if use_diff_format { + agent.enable_diff_blocks(); + } + // Check if we should continue from previous state or start new if continue_task && file_persistence.has_saved_state() { // Load from saved state @@ -460,6 +458,7 @@ fn run_agent_gpui( aicore_config: Option, num_ctx: usize, tool_syntax: ToolSyntax, + use_diff_format: bool, record: Option, playback: Option, fast_playback: bool, @@ -480,6 +479,7 @@ fn run_agent_gpui( tool_syntax: tool_syntax, init_path: Some(root_path.clone()), initial_project: None, + use_diff_blocks: use_diff_format, }; // Create the new SessionManager @@ -648,11 +648,6 @@ fn run_agent_gpui( } async fn run_agent(args: Args) -> Result<()> { - // Set environment variable for diff format preference before tools are initialized - if args.use_diff_format { - std::env::set_var("CODE_ASSISTANT_USE_DIFF_FORMAT", "true"); - } - // Get all the agent options from args let path = args.path.clone().unwrap_or_else(|| PathBuf::from(".")); let task = args.task.clone(); @@ -664,7 +659,6 @@ async fn run_agent(args: Args) -> Result<()> { let aicore_config = args.aicore_config.clone(); let num_ctx = args.num_ctx.unwrap_or(8192); let tool_syntax = args.tool_syntax.unwrap_or(ToolSyntax::Native); - let use_gui = args.ui; // Setup logging based on verbose flag setup_logging(verbose, true); @@ -675,7 +669,7 @@ async fn run_agent(args: Args) -> Result<()> { } // Run in either GUI or terminal mode - if use_gui { + if args.ui { run_agent_gpui( path.clone(), task, // Can be None - will connect to latest session instead @@ -685,6 +679,7 @@ async fn run_agent(args: Args) -> Result<()> { aicore_config, num_ctx, tool_syntax, + args.use_diff_format, args.record.clone(), args.playback.clone(), args.fast_playback, @@ -701,6 +696,7 @@ async fn run_agent(args: Args) -> Result<()> { aicore_config, num_ctx, tool_syntax, + args.use_diff_format, args.record.clone(), args.playback.clone(), args.fast_playback, @@ -989,7 +985,7 @@ async fn main() -> Result<()> { match args.mode { // Server mode - Some(Mode::Server { verbose, use_diff_format }) => run_mcp_server(verbose, use_diff_format).await, + Some(Mode::Server { verbose }) => run_mcp_server(verbose).await, // Agent mode (default) None => run_agent(args).await, diff --git a/crates/code_assistant/src/persistence.rs b/crates/code_assistant/src/persistence.rs index 5eb2259f..d7e2dd8f 100644 --- a/crates/code_assistant/src/persistence.rs +++ b/crates/code_assistant/src/persistence.rs @@ -32,6 +32,9 @@ pub struct ChatSession { /// Tool syntax used for this session (XML or Native) #[serde(alias = "tool_mode")] pub tool_syntax: ToolSyntax, + /// Whether this session uses diff blocks format (replace_in_file vs edit tool) + #[serde(default)] + pub use_diff_blocks: bool, /// Counter for generating unique request IDs within this session #[serde(default)] pub next_request_id: u64, diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index 07f5208f..4e43b0e4 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -37,6 +37,7 @@ pub struct AgentConfig { pub tool_syntax: ToolSyntax, pub init_path: Option, pub initial_project: Option, + pub use_diff_blocks: bool, } impl SessionManager { @@ -66,6 +67,7 @@ impl SessionManager { init_path: self.agent_config.init_path.clone(), initial_project: self.agent_config.initial_project.clone(), tool_syntax: self.agent_config.tool_syntax, + use_diff_blocks: self.agent_config.use_diff_blocks, next_request_id: 1, }; @@ -154,6 +156,7 @@ impl SessionManager { // Prepare session - need to scope the mutable borrow carefully let ( tool_syntax, + use_diff_blocks, init_path, proxy_ui, session_state, @@ -180,6 +183,7 @@ impl SessionManager { // Clone all needed data to avoid borrowing conflicts let name = session_instance.session.name.clone(); let tool_syntax = session_instance.session.tool_syntax; + let use_diff_blocks = session_instance.session.use_diff_blocks; let init_path = session_instance.session.init_path.clone(); let proxy_ui = session_instance.create_proxy_ui(ui.clone()); let activity_state_ref = session_instance.activity_state.clone(); @@ -207,6 +211,7 @@ impl SessionManager { ( tool_syntax, + use_diff_blocks, init_path, proxy_ui, session_state, @@ -218,6 +223,7 @@ impl SessionManager { // Now save the session state with the user message (outside the borrow scope) self.save_session_state( session_id, + session_state.name.clone(), session_state.messages.clone(), session_state.tool_executions.clone(), session_state.working_memory.clone(), @@ -255,6 +261,11 @@ impl SessionManager { init_path, ); + // Configure diff blocks format based on session setting + if use_diff_blocks { + agent.enable_diff_blocks(); + } + // Set the shared pending message reference agent.set_pending_message_ref(pending_message_ref); @@ -331,6 +342,7 @@ impl SessionManager { pub fn save_session_state( &mut self, session_id: &str, + name: String, messages: Vec, tool_executions: Vec, working_memory: WorkingMemory, @@ -344,6 +356,7 @@ impl SessionManager { .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id))?; // Update session with current state + session.name = name; session.messages = messages; session.tool_executions = tool_executions .into_iter() diff --git a/crates/code_assistant/src/tools/core/registry.rs b/crates/code_assistant/src/tools/core/registry.rs index 0414061b..128935f3 100644 --- a/crates/code_assistant/src/tools/core/registry.rs +++ b/crates/code_assistant/src/tools/core/registry.rs @@ -17,11 +17,7 @@ impl ToolRegistry { static INSTANCE: OnceLock = OnceLock::new(); INSTANCE.get_or_init(|| { let mut registry = ToolRegistry::new(); - // Check environment variable for diff format preference - let use_diff_format = std::env::var("CODE_ASSISTANT_USE_DIFF_FORMAT") - .map(|v| v.to_lowercase() == "true") - .unwrap_or(false); - registry.register_default_tools_impl(use_diff_format); + registry.register_default_tools(); registry }) } @@ -72,42 +68,29 @@ impl ToolRegistry { /// Register all default tools in the system /// This will be expanded as we implement more tools - #[cfg(test)] - pub fn register_default_tools(&mut self, use_diff_format: bool) { - self.register_default_tools_impl(use_diff_format); - } - - /// Internal implementation of register_default_tools - fn register_default_tools_impl(&mut self, use_diff_format: bool) { + fn register_default_tools(&mut self) { // Import all tools use crate::tools::impls::{ - DeleteFilesTool, EditTool, ExecuteCommandTool, ListFilesTool, ListProjectsTool, NameSessionTool, - PerplexityAskTool, ReadFilesTool, ReplaceInFileTool, SearchFilesTool, WebFetchTool, - WebSearchTool, WriteFileTool, + DeleteFilesTool, EditTool, ExecuteCommandTool, ListFilesTool, ListProjectsTool, + NameSessionTool, PerplexityAskTool, ReadFilesTool, ReplaceInFileTool, SearchFilesTool, + WebFetchTool, WebSearchTool, WriteFileTool, }; - // Register core tools + // Register all tools - the ToolScope system will filter which ones are available self.register(Box::new(DeleteFilesTool)); + self.register(Box::new(EditTool)); self.register(Box::new(ExecuteCommandTool)); self.register(Box::new(ListFilesTool)); self.register(Box::new(ListProjectsTool)); self.register(Box::new(NameSessionTool)); self.register(Box::new(PerplexityAskTool)); self.register(Box::new(ReadFilesTool)); + self.register(Box::new(ReplaceInFileTool)); self.register(Box::new(SearchFilesTool)); self.register(Box::new(WebFetchTool)); self.register(Box::new(WebSearchTool)); self.register(Box::new(WriteFileTool)); - // Register file editing tools based on configuration - if use_diff_format { - // Use legacy diff format (replace_in_file tool) - self.register(Box::new(ReplaceInFileTool)); - } else { - // Use new edit tool (default) - self.register(Box::new(EditTool)); - } - // More tools will be added here as they are implemented } } diff --git a/crates/code_assistant/src/tools/core/spec.rs b/crates/code_assistant/src/tools/core/spec.rs index bd378cf5..133450ff 100644 --- a/crates/code_assistant/src/tools/core/spec.rs +++ b/crates/code_assistant/src/tools/core/spec.rs @@ -5,6 +5,8 @@ pub enum ToolScope { McpServer, /// Tool can be used in the message history agent Agent, + /// Tool can be used in the agent when configured for diff blocks format + AgentWithDiffBlocks, } /// Specification for a tool, including metadata diff --git a/crates/code_assistant/src/tools/impls/delete_files.rs b/crates/code_assistant/src/tools/impls/delete_files.rs index 26563586..dfbe7bac 100644 --- a/crates/code_assistant/src/tools/impls/delete_files.rs +++ b/crates/code_assistant/src/tools/impls/delete_files.rs @@ -101,7 +101,11 @@ impl Tool for DeleteFilesTool { "destructiveHint": true, "idempotentHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/execute_command.rs b/crates/code_assistant/src/tools/impls/execute_command.rs index 0272b6e7..daf681d4 100644 --- a/crates/code_assistant/src/tools/impls/execute_command.rs +++ b/crates/code_assistant/src/tools/impls/execute_command.rs @@ -96,7 +96,11 @@ impl Tool for ExecuteCommandTool { "readOnlyHint": false, "idempotentHint": false })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/list_files.rs b/crates/code_assistant/src/tools/impls/list_files.rs index d1e4222b..1363a3bf 100644 --- a/crates/code_assistant/src/tools/impls/list_files.rs +++ b/crates/code_assistant/src/tools/impls/list_files.rs @@ -153,7 +153,11 @@ impl Tool for ListFilesTool { annotations: Some(json!({ "readOnlyHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/name_session.rs b/crates/code_assistant/src/tools/impls/name_session.rs index 37ecb188..e670df80 100644 --- a/crates/code_assistant/src/tools/impls/name_session.rs +++ b/crates/code_assistant/src/tools/impls/name_session.rs @@ -62,7 +62,7 @@ impl Tool for NameSessionTool { "required": ["title"] }), annotations: None, - supported_scopes: &[ToolScope::Agent], + supported_scopes: &[ToolScope::Agent, ToolScope::AgentWithDiffBlocks], hidden: true, // This tool should be hidden from UI } } diff --git a/crates/code_assistant/src/tools/impls/perplexity_ask.rs b/crates/code_assistant/src/tools/impls/perplexity_ask.rs index 6ed40f62..767aea8e 100644 --- a/crates/code_assistant/src/tools/impls/perplexity_ask.rs +++ b/crates/code_assistant/src/tools/impls/perplexity_ask.rs @@ -110,7 +110,11 @@ impl Tool for PerplexityAskTool { "idempotentHint": false, "openWorldHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/read_files.rs b/crates/code_assistant/src/tools/impls/read_files.rs index d481b922..756ba7e3 100644 --- a/crates/code_assistant/src/tools/impls/read_files.rs +++ b/crates/code_assistant/src/tools/impls/read_files.rs @@ -135,7 +135,11 @@ impl Tool for ReadFilesTool { "readOnlyHint": true, "idempotentHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/replace_in_file.rs b/crates/code_assistant/src/tools/impls/replace_in_file.rs index ddb49ea7..4fe745ee 100644 --- a/crates/code_assistant/src/tools/impls/replace_in_file.rs +++ b/crates/code_assistant/src/tools/impls/replace_in_file.rs @@ -116,7 +116,7 @@ impl Tool for ReplaceInFileTool { "readOnlyHint": false, "destructiveHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ToolScope::AgentWithDiffBlocks], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/search_files.rs b/crates/code_assistant/src/tools/impls/search_files.rs index 5ed5cb59..08e5e9a2 100644 --- a/crates/code_assistant/src/tools/impls/search_files.rs +++ b/crates/code_assistant/src/tools/impls/search_files.rs @@ -103,7 +103,11 @@ impl Tool for SearchFilesTool { annotations: Some(json!({ "readOnlyHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/web_fetch.rs b/crates/code_assistant/src/tools/impls/web_fetch.rs index da0345fb..284c0a39 100644 --- a/crates/code_assistant/src/tools/impls/web_fetch.rs +++ b/crates/code_assistant/src/tools/impls/web_fetch.rs @@ -90,7 +90,11 @@ impl Tool for WebFetchTool { "idempotentHint": true, "openWorldHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/web_search.rs b/crates/code_assistant/src/tools/impls/web_search.rs index 0562f7d1..03e5beae 100644 --- a/crates/code_assistant/src/tools/impls/web_search.rs +++ b/crates/code_assistant/src/tools/impls/web_search.rs @@ -102,7 +102,11 @@ impl Tool for WebSearchTool { "idempotentHint": true, "openWorldHint": true })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/impls/write_file.rs b/crates/code_assistant/src/tools/impls/write_file.rs index ae7675e5..293eec77 100644 --- a/crates/code_assistant/src/tools/impls/write_file.rs +++ b/crates/code_assistant/src/tools/impls/write_file.rs @@ -107,7 +107,11 @@ impl Tool for WriteFileTool { "destructiveHint": true, "idempotentHint": false })), - supported_scopes: &[ToolScope::McpServer, ToolScope::Agent], + supported_scopes: &[ + ToolScope::McpServer, + ToolScope::Agent, + ToolScope::AgentWithDiffBlocks, + ], hidden: false, } } diff --git a/crates/code_assistant/src/tools/integration_tests.rs b/crates/code_assistant/src/tools/integration_tests.rs deleted file mode 100644 index 0d76d10e..00000000 --- a/crates/code_assistant/src/tools/integration_tests.rs +++ /dev/null @@ -1,102 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::tools::core::ToolRegistry; - use std::env; - - #[test] - fn test_default_tool_registration() { - // Clear environment variable to ensure default behavior - env::remove_var("CODE_ASSISTANT_USE_DIFF_FORMAT"); - - // Create a new registry to test default registration - let mut registry = ToolRegistry::new(); - registry.register_default_tools(false); // Default behavior - - // Edit tool should be registered - assert!(registry.get("edit").is_some(), "Edit tool should be registered by default"); - - // Replace in file tool should NOT be registered - assert!(registry.get("replace_in_file").is_none(), "ReplaceInFile tool should not be registered by default"); - - // Other core tools should still be registered - assert!(registry.get("read_files").is_some(), "ReadFiles tool should be registered"); - assert!(registry.get("write_file").is_some(), "WriteFile tool should be registered"); - assert!(registry.get("list_files").is_some(), "ListFiles tool should be registered"); - } - - #[test] - fn test_diff_format_tool_registration() { - // Create a new registry to test diff format registration - let mut registry = ToolRegistry::new(); - registry.register_default_tools(true); // Use diff format - - // Replace in file tool should be registered - assert!(registry.get("replace_in_file").is_some(), "ReplaceInFile tool should be registered with diff format"); - - // Edit tool should NOT be registered - assert!(registry.get("edit").is_none(), "Edit tool should not be registered with diff format"); - - // Other core tools should still be registered - assert!(registry.get("read_files").is_some(), "ReadFiles tool should be registered"); - assert!(registry.get("write_file").is_some(), "WriteFile tool should be registered"); - assert!(registry.get("list_files").is_some(), "ListFiles tool should be registered"); - } - - #[test] - fn test_tool_specs() { - // Test edit tool spec - let mut registry = ToolRegistry::new(); - registry.register_default_tools(false); - - if let Some(edit_tool) = registry.get("edit") { - let spec = edit_tool.spec(); - assert_eq!(spec.name, "edit"); - assert!(spec.description.contains("Edit a file by replacing specific text content")); - - // Check required parameters - let params = &spec.parameters_schema; - let properties = params.get("properties").unwrap().as_object().unwrap(); - assert!(properties.contains_key("project")); - assert!(properties.contains_key("path")); - assert!(properties.contains_key("old_text")); - assert!(properties.contains_key("new_text")); - - // Check that replace_all is optional - let required = params.get("required").unwrap().as_array().unwrap(); - assert!(required.contains(&serde_json::Value::String("project".to_string()))); - assert!(required.contains(&serde_json::Value::String("path".to_string()))); - assert!(required.contains(&serde_json::Value::String("old_text".to_string()))); - assert!(required.contains(&serde_json::Value::String("new_text".to_string()))); - assert!(!required.contains(&serde_json::Value::String("replace_all".to_string()))); - } else { - panic!("Edit tool should be registered"); - } - } - - #[test] - fn test_replace_in_file_tool_specs() { - // Test replace_in_file tool spec - let mut registry = ToolRegistry::new(); - registry.register_default_tools(true); - - if let Some(replace_tool) = registry.get("replace_in_file") { - let spec = replace_tool.spec(); - assert_eq!(spec.name, "replace_in_file"); - assert!(spec.description.contains("Replace sections in a file")); - - // Check required parameters - let params = &spec.parameters_schema; - let properties = params.get("properties").unwrap().as_object().unwrap(); - assert!(properties.contains_key("project")); - assert!(properties.contains_key("path")); - assert!(properties.contains_key("diff")); - - let required = params.get("required").unwrap().as_array().unwrap(); - assert!(required.contains(&serde_json::Value::String("project".to_string()))); - assert!(required.contains(&serde_json::Value::String("path".to_string()))); - assert!(required.contains(&serde_json::Value::String("diff".to_string()))); - } else { - panic!("ReplaceInFile tool should be registered with diff format"); - } - } -} diff --git a/crates/code_assistant/src/tools/mod.rs b/crates/code_assistant/src/tools/mod.rs index 4af961ac..c6cc8ced 100644 --- a/crates/code_assistant/src/tools/mod.rs +++ b/crates/code_assistant/src/tools/mod.rs @@ -18,9 +18,6 @@ pub mod impls; #[cfg(test)] mod tests; -#[cfg(test)] -mod integration_tests; - pub use parse::{parse_caret_tool_invocations, parse_xml_tool_invocations}; pub use parser_registry::ParserRegistry; pub use system_message::generate_system_message; From 8f195847af99f2f5c3fc9ed1a16e721272aea0d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 29 Jul 2025 14:30:33 +0200 Subject: [PATCH 3/7] Fix non-macOS builds --- crates/code_assistant/src/ui/gpui/mod.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/code_assistant/src/ui/gpui/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index 366c10ea..405b1a2e 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -376,6 +376,8 @@ impl Gpui { title: Some(gpui::SharedString::from("Code Assistant")), #[cfg(target_os = "macos")] appears_transparent: true, + #[cfg(not(target_os = "macos"))] + appears_transparent: false, traffic_light_position: Some(Point { x: px(16.), y: px(16.), From 5bed85f60140b621175c2ece2c00c1171dd54eb3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 29 Jul 2025 15:53:50 +0200 Subject: [PATCH 4/7] Improve edit tool output --- crates/code_assistant/src/tools/impls/edit.rs | 32 ++++++++----------- 1 file changed, 13 insertions(+), 19 deletions(-) diff --git a/crates/code_assistant/src/tools/impls/edit.rs b/crates/code_assistant/src/tools/impls/edit.rs index 0d8ed633..5ed5071a 100644 --- a/crates/code_assistant/src/tools/impls/edit.rs +++ b/crates/code_assistant/src/tools/impls/edit.rs @@ -31,28 +31,24 @@ pub struct EditOutput { impl Render for EditOutput { fn status(&self) -> String { if self.error.is_none() { - format!( - "Successfully edited content in file: {}", - self.path.display() - ) + format!("Successfully edited file: {}", self.path.display()) } else { - format!("Failed to edit content in file: {}", self.path.display()) + format!("Failed to edit file: {}", self.path.display()) } } fn render(&self, _tracker: &mut ResourcesTracker) -> String { if let Some(error) = &self.error { match error { - FileUpdaterError::SearchBlockNotFound(idx, _) => { + FileUpdaterError::SearchBlockNotFound(_, _) => { format!( - "Could not find the text to replace (search block {}). Please check that the old_text matches exactly what's in the file.", - idx + "Could not find old_text. Make sure it matches exactly what's in the file." ) } - FileUpdaterError::MultipleMatches(count, idx, _) => { + FileUpdaterError::MultipleMatches(count, _, _) => { format!( - "Found {} occurrences of the text to replace (search block {})\nThe old_text must match exactly one location. Try making the old_text more specific.", - count, idx + "Found {} occurrences of old_text\nIt must match exactly one location. Try enlarging old_text to make it unique or use replace_all to replace all occurrences.", + count ) } FileUpdaterError::Other(msg) => { @@ -60,10 +56,7 @@ impl Render for EditOutput { } } } else { - format!( - "Successfully edited content in file '{}'", - self.path.display() - ) + format!("Successfully edited file '{}'", self.path.display()) } } } @@ -221,7 +214,7 @@ mod tests { let mut tracker = ResourcesTracker::new(); let rendered = output.render(&mut tracker); - assert!(rendered.contains("Successfully edited content")); + assert!(rendered.contains("Successfully edited file")); assert!(rendered.contains("src/test.rs")); // Error case with text not found @@ -235,8 +228,8 @@ mod tests { }; let rendered_error = output_error.render(&mut tracker); - assert!(rendered_error.contains("Could not find the text to replace")); - assert!(rendered_error.contains("old_text matches exactly")); + assert!(rendered_error.contains("Could not find old_text")); + assert!(rendered_error.contains("matches exactly")); // Error case with multiple matches let output_multiple = EditOutput { @@ -251,7 +244,8 @@ mod tests { let rendered_multiple = output_multiple.render(&mut tracker); assert!(rendered_multiple.contains("Found 3 occurrences")); - assert!(rendered_multiple.contains("more specific")); + assert!(rendered_multiple.contains("Try enlarging old_text")); + assert!(rendered_multiple.contains("make it unique")); } #[tokio::test] From b10d80f588c874e2f0855280576d344812bcc4d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 29 Jul 2025 16:04:27 +0200 Subject: [PATCH 5/7] Fix system reminders about naming session are added to actual user messages only --- crates/code_assistant/src/agent/runner.rs | 63 ++++++-- crates/code_assistant/src/agent/tests.rs | 188 ++++++++++++++++++++++ 2 files changed, 236 insertions(+), 15 deletions(-) diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index bda4e243..3038232e 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -128,6 +128,12 @@ impl Agent { self.enable_naming_reminders = false; } + /// Set session name (used for tests) + #[cfg(test)] + pub(crate) fn set_session_name(&mut self, name: String) { + self.session_name = name; + } + /// Get and clear the pending message from shared state fn get_and_clear_pending_message(&self) -> Option { if let Some(ref pending_ref) = self.pending_message_ref { @@ -734,28 +740,55 @@ impl Agent { } /// Inject system reminder for session naming if needed - fn inject_naming_reminder_if_needed(&self, mut messages: Vec) -> Vec { + pub(crate) fn inject_naming_reminder_if_needed( + &self, + mut messages: Vec, + ) -> Vec { // Only inject if enabled, session is not named yet, and we have messages if !self.enable_naming_reminders || !self.session_name.is_empty() || messages.is_empty() { return messages; } - // Find the last user message and add system reminder - if let Some(last_msg) = messages.last_mut() { - if matches!(last_msg.role, MessageRole::User) { - let reminder_text = "\n\n\nThis is an automatic reminder from the system. Please use the `name_session` tool first, provided the user has already given you a clear task or question. You can chain additional tools after using the `name_session` tool.\n"; - - trace!("Injecting session naming reminder"); - - match &mut last_msg.content { - MessageContent::Text(text) => { - text.push_str(reminder_text); - } + // Find the last actual user message (not tool results) and add system reminder + // Iterate backwards through messages to find the last user message with actual content + for msg in messages.iter_mut().rev() { + if matches!(msg.role, MessageRole::User) { + let is_actual_user_message = match &msg.content { + MessageContent::Text(_) => true, // Text content is always actual user input MessageContent::Structured(blocks) => { - blocks.push(ContentBlock::Text { - text: reminder_text.to_string(), - }); + // Check if this message contains tool results + // If it contains only ToolResult blocks, it's not an actual user message + blocks + .iter() + .any(|block| !matches!(block, ContentBlock::ToolResult { .. })) + } + }; + + if is_actual_user_message { + let reminder_text = "\nThis is an automatic reminder from the system. Please use the `name_session` tool first, provided the user has already given you a clear task or question. You can chain additional tools after using the `name_session` tool.\n"; + + trace!("Injecting session naming reminder to actual user message"); + + match &mut msg.content { + MessageContent::Text(original_text) => { + // Convert from Text to Structured with two ContentBlocks + let original_block = ContentBlock::Text { + text: original_text.clone(), + }; + let reminder_block = ContentBlock::Text { + text: reminder_text.to_string(), + }; + msg.content = + MessageContent::Structured(vec![original_block, reminder_block]); + } + MessageContent::Structured(blocks) => { + // Add reminder as a new ContentBlock + blocks.push(ContentBlock::Text { + text: reminder_text.to_string(), + }); + } } + break; // Found and updated the last actual user message, we're done } } } diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index a9d9f635..e888da87 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -1111,3 +1111,191 @@ fn test_original_caret_issue_reproduction() -> Result<()> { Ok(()) } + +#[test] +fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { + // This test verifies that: + // 1. Naming reminders are only added to actual user messages, not tool result messages + // 2. Text messages are converted to Structured with separate ContentBlocks for original text and reminder + // 3. Structured messages get the reminder added as an additional ContentBlock + // Create a mock agent for testing + let llm_provider = Box::new(MockLLMProvider::new(vec![])); + let project_manager = Box::new(MockProjectManager::default()); + let command_executor = Box::new(create_command_executor_mock()); + let ui = Arc::new(Box::new(MockUI::default()) as Box); + let state_persistence = Box::new(MockStatePersistence::new()); + + let mut agent = Agent::new( + llm_provider, + ToolSyntax::Xml, + project_manager, + command_executor, + ui, + state_persistence, + None, + ); + + // Test case 1: User message with text content should get reminder + let messages = vec![ + Message { + role: MessageRole::User, + content: MessageContent::Text("Hello, help me with a task".to_string()), + request_id: None, + usage: None, + }, + ]; + + let result_messages = agent.inject_naming_reminder_if_needed(messages.clone()); + assert_eq!(result_messages.len(), 1); + + // The message should now be structured with two ContentBlocks + if let MessageContent::Structured(blocks) = &result_messages[0].content { + assert_eq!(blocks.len(), 2); + + // First block should contain the original user text + if let ContentBlock::Text { text } = &blocks[0] { + assert_eq!(text, "Hello, help me with a task"); + } else { + panic!("Expected first block to be text with original user message"); + } + + // Second block should contain the reminder + if let ContentBlock::Text { text } = &blocks[1] { + assert!(text.contains("")); + assert!(text.contains("name_session")); + } else { + panic!("Expected second block to be text with reminder"); + } + } else { + panic!("Expected structured content after reminder injection"); + } + + // Test case 2: User message with only tool results should be skipped + let messages_with_tool_results = vec![ + Message { + role: MessageRole::User, + content: MessageContent::Text("Hello, help me with a task".to_string()), + request_id: None, + usage: None, + }, + Message { + role: MessageRole::Assistant, + content: MessageContent::Structured(vec![ContentBlock::Text { + text: "I'll help you with that task.".to_string(), + }]), + request_id: Some(1), + usage: Some(Usage::zero()), + }, + Message { + role: MessageRole::User, + content: MessageContent::Structured(vec![ContentBlock::ToolResult { + tool_use_id: "tool-1-1".to_string(), + content: "Tool execution result".to_string(), + is_error: None, + }]), + request_id: None, + usage: None, + }, + ]; + + let result_messages = agent.inject_naming_reminder_if_needed(messages_with_tool_results.clone()); + assert_eq!(result_messages.len(), 3); + + // The reminder should be added to the first user message (with text content), not the tool result message + // The first message should now be structured with two ContentBlocks + if let MessageContent::Structured(blocks) = &result_messages[0].content { + assert_eq!(blocks.len(), 2); + + // First block should contain the original user text + if let ContentBlock::Text { text } = &blocks[0] { + assert_eq!(text, "Hello, help me with a task"); + } else { + panic!("Expected first block to be text with original user message"); + } + + // Second block should contain the reminder + if let ContentBlock::Text { text } = &blocks[1] { + assert!(text.contains("")); + assert!(text.contains("name_session")); + } else { + panic!("Expected second block to be text with reminder"); + } + } else { + panic!("Expected structured content in first message after reminder injection"); + } + + // The tool result message should remain unchanged + if let MessageContent::Structured(blocks) = &result_messages[2].content { + assert_eq!(blocks.len(), 1); + assert!(matches!(blocks[0], ContentBlock::ToolResult { .. })); + // No reminder should be added to this message + for block in blocks { + if let ContentBlock::Text { text } = block { + assert!(!text.contains("")); + } + } + } else { + panic!("Expected structured content in tool result message"); + } + + // Test case 3: User message with mixed content (text + tool results) should get reminder + let mixed_message = vec![ + Message { + role: MessageRole::User, + content: MessageContent::Structured(vec![ + ContentBlock::Text { + text: "Please analyze this file".to_string(), + }, + ContentBlock::ToolResult { + tool_use_id: "tool-1-1".to_string(), + content: "Previous tool result".to_string(), + is_error: None, + }, + ]), + request_id: None, + usage: None, + }, + ]; + + let result_messages = agent.inject_naming_reminder_if_needed(mixed_message.clone()); + assert_eq!(result_messages.len(), 1); + + if let MessageContent::Structured(blocks) = &result_messages[0].content { + assert_eq!(blocks.len(), 3); // Original text + tool result + reminder text + + // First block should be the original text + if let ContentBlock::Text { text } = &blocks[0] { + assert_eq!(text, "Please analyze this file"); + } else { + panic!("Expected first block to be original text"); + } + + // Second block should be the tool result (unchanged) + assert!(matches!(blocks[1], ContentBlock::ToolResult { .. })); + + // Third block should be the reminder + if let ContentBlock::Text { text } = &blocks[2] { + assert!(text.contains("")); + assert!(text.contains("name_session")); + } else { + panic!("Expected third block to be reminder text"); + } + } else { + panic!("Expected structured content"); + } + + // Test case 4: No reminder should be added if session is already named + agent.set_session_name("Test Session".to_string()); + let result_messages = agent.inject_naming_reminder_if_needed(messages); + assert_eq!(result_messages.len(), 1); + + // When session is already named, the message should remain unchanged (Text content) + if let MessageContent::Text(text) = &result_messages[0].content { + assert_eq!(text, "Hello, help me with a task"); + assert!(!text.contains("")); + } else { + panic!("Expected text content to remain unchanged when session is already named"); + } + + Ok(()) +} From 1bb1db1476e2e562d39dd40473fb44e8181b3202 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 29 Jul 2025 16:22:13 +0200 Subject: [PATCH 6/7] New icon for edit tool --- crates/code_assistant/assets/icons/file_icons/file_types.json | 3 +++ crates/code_assistant/assets/icons/square-pen.svg | 1 + crates/code_assistant/src/ui/gpui/file_icons.rs | 3 +++ 3 files changed, 7 insertions(+) create mode 100644 crates/code_assistant/assets/icons/square-pen.svg diff --git a/crates/code_assistant/assets/icons/file_icons/file_types.json b/crates/code_assistant/assets/icons/file_icons/file_types.json index 4f9ecd1a..9f131935 100644 --- a/crates/code_assistant/assets/icons/file_icons/file_types.json +++ b/crates/code_assistant/assets/icons/file_icons/file_types.json @@ -398,6 +398,9 @@ }, "plus": { "icon": "icons/plus.svg" + }, + "edit": { + "icon": "icons/square-pen.svg" } } } diff --git a/crates/code_assistant/assets/icons/square-pen.svg b/crates/code_assistant/assets/icons/square-pen.svg new file mode 100644 index 00000000..06515ddf --- /dev/null +++ b/crates/code_assistant/assets/icons/square-pen.svg @@ -0,0 +1 @@ + diff --git a/crates/code_assistant/src/ui/gpui/file_icons.rs b/crates/code_assistant/src/ui/gpui/file_icons.rs index f24368ae..ddfd1ae9 100644 --- a/crates/code_assistant/src/ui/gpui/file_icons.rs +++ b/crates/code_assistant/src/ui/gpui/file_icons.rs @@ -63,6 +63,7 @@ pub const TOOL_LIST_FILES: &str = "reveal"; // reveal.svg pub const TOOL_EXECUTE_COMMAND: &str = "terminal"; // terminal.svg pub const TOOL_WRITE_FILE: &str = "pencil"; // pencil.svg pub const TOOL_REPLACE_IN_FILE: &str = "replace"; // replace.svg +pub const TOOL_EDIT: &str = "edit"; // square-pen.svg pub const TOOL_SEARCH_FILES: &str = "magnifying_glass"; // magnifying_glass.svg pub const TOOL_WEB_SEARCH: &str = "magnifying_glass"; // magnifying_glass.svg pub const TOOL_WEB_FETCH: &str = "template"; // html.svg (use template/html as fallback) @@ -269,6 +270,7 @@ impl FileIcons { TOOL_EXECUTE_COMMAND => Some(SharedString::from("🖥️")), TOOL_WRITE_FILE => Some(SharedString::from("✏️")), TOOL_REPLACE_IN_FILE => Some(SharedString::from("🔄")), + TOOL_EDIT => Some(SharedString::from("✏️")), // Pen/edit icon TOOL_SEARCH_FILES => Some(SharedString::from("🔍")), // TOOL_WEB_SEARCH uses same icon as SEARCH_FILES - handled above TOOL_WEB_FETCH => Some(SharedString::from("📥")), @@ -290,6 +292,7 @@ impl FileIcons { "execute_command" => TOOL_EXECUTE_COMMAND, "write_file" => TOOL_WRITE_FILE, "replace_in_file" => TOOL_REPLACE_IN_FILE, + "edit" => TOOL_EDIT, "search_files" => TOOL_SEARCH_FILES, "web_search" => TOOL_WEB_SEARCH, "web_fetch" => TOOL_WEB_FETCH, From 6b37bdb23d3d1ac020403bedf6b1aeb6c63849e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Tue, 29 Jul 2025 16:26:25 +0200 Subject: [PATCH 7/7] cargo fmt --- crates/code_assistant/src/agent/tests.rs | 49 +++++++++++------------- 1 file changed, 23 insertions(+), 26 deletions(-) diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index e888da87..b04c15b9 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -1136,14 +1136,12 @@ fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { ); // Test case 1: User message with text content should get reminder - let messages = vec![ - Message { - role: MessageRole::User, - content: MessageContent::Text("Hello, help me with a task".to_string()), - request_id: None, - usage: None, - }, - ]; + let messages = vec![Message { + role: MessageRole::User, + content: MessageContent::Text("Hello, help me with a task".to_string()), + request_id: None, + usage: None, + }]; let result_messages = agent.inject_naming_reminder_if_needed(messages.clone()); assert_eq!(result_messages.len(), 1); @@ -1198,7 +1196,8 @@ fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { }, ]; - let result_messages = agent.inject_naming_reminder_if_needed(messages_with_tool_results.clone()); + let result_messages = + agent.inject_naming_reminder_if_needed(messages_with_tool_results.clone()); assert_eq!(result_messages.len(), 3); // The reminder should be added to the first user message (with text content), not the tool result message @@ -1239,23 +1238,21 @@ fn test_inject_naming_reminder_skips_tool_result_messages() -> Result<()> { } // Test case 3: User message with mixed content (text + tool results) should get reminder - let mixed_message = vec![ - Message { - role: MessageRole::User, - content: MessageContent::Structured(vec![ - ContentBlock::Text { - text: "Please analyze this file".to_string(), - }, - ContentBlock::ToolResult { - tool_use_id: "tool-1-1".to_string(), - content: "Previous tool result".to_string(), - is_error: None, - }, - ]), - request_id: None, - usage: None, - }, - ]; + let mixed_message = vec![Message { + role: MessageRole::User, + content: MessageContent::Structured(vec![ + ContentBlock::Text { + text: "Please analyze this file".to_string(), + }, + ContentBlock::ToolResult { + tool_use_id: "tool-1-1".to_string(), + content: "Previous tool result".to_string(), + is_error: None, + }, + ]), + request_id: None, + usage: None, + }]; let result_messages = agent.inject_naming_reminder_if_needed(mixed_message.clone()); assert_eq!(result_messages.len(), 1);