diff --git a/crates/agent/src/edit_agent.rs b/crates/agent/src/edit_agent.rs index 288a3178f3c4501ae9de65d19624b66cbda2548d..ef95eee07378438686aff688fdaf2d7fa98e036b 100644 --- a/crates/agent/src/edit_agent.rs +++ b/crates/agent/src/edit_agent.rs @@ -2,6 +2,7 @@ mod create_file_parser; mod edit_parser; #[cfg(test)] mod evals; +pub mod reindent; pub mod streaming_fuzzy_matcher; use crate::{Template, Templates}; @@ -24,9 +25,10 @@ use language_model::{ LanguageModelToolChoice, MessageContent, Role, }; use project::{AgentLocation, Project}; +use reindent::{IndentDelta, Reindenter}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use std::{cmp, iter, mem, ops::Range, pin::Pin, sync::Arc, task::Poll}; +use std::{mem, ops::Range, pin::Pin, sync::Arc, task::Poll}; use streaming_diff::{CharOperation, StreamingDiff}; use streaming_fuzzy_matcher::StreamingFuzzyMatcher; @@ -553,15 +555,8 @@ impl EditAgent { let compute_edits = cx.background_spawn(async move { let buffer_start_indent = snapshot .line_indent_for_row(snapshot.offset_to_point(resolved_old_text.range.start).row); - let indent_delta = if buffer_start_indent.tabs > 0 { - IndentDelta::Tabs( - buffer_start_indent.tabs as isize - resolved_old_text.indent.tabs as isize, - ) - } else { - IndentDelta::Spaces( - buffer_start_indent.spaces as isize - resolved_old_text.indent.spaces as isize, - ) - }; + let indent_delta = + reindent::compute_indent_delta(buffer_start_indent, resolved_old_text.indent); let old_text = snapshot .text_for_range(resolved_old_text.range.clone()) @@ -608,8 +603,7 @@ impl EditAgent { delta: IndentDelta, mut stream: impl Unpin + Stream>, ) -> impl Stream> { - let mut buffer = String::new(); - let mut in_leading_whitespace = true; + let mut reindenter = Reindenter::new(delta); let mut done = false; futures::stream::poll_fn(move |cx| { while !done { @@ -622,55 +616,10 @@ impl EditAgent { _ => return Poll::Ready(None), }; - buffer.push_str(&chunk); - - let mut indented_new_text = String::new(); - let mut start_ix = 0; - let mut newlines = buffer.match_indices('\n').peekable(); - loop { - let (line_end, is_pending_line) = match newlines.next() { - Some((ix, _)) => (ix, false), - None => (buffer.len(), true), - }; - let line = &buffer[start_ix..line_end]; - - if in_leading_whitespace { - if let Some(non_whitespace_ix) = line.find(|c| delta.character() != c) { - // We found a non-whitespace character, adjust - // indentation based on the delta. - let new_indent_len = - cmp::max(0, non_whitespace_ix as isize + delta.len()) as usize; - indented_new_text - .extend(iter::repeat(delta.character()).take(new_indent_len)); - indented_new_text.push_str(&line[non_whitespace_ix..]); - in_leading_whitespace = false; - } else if is_pending_line { - // We're still in leading whitespace and this line is incomplete. - // Stop processing until we receive more input. - break; - } else { - // This line is entirely whitespace. Push it without indentation. - indented_new_text.push_str(line); - } - } else { - indented_new_text.push_str(line); - } - - if is_pending_line { - start_ix = line_end; - break; - } else { - in_leading_whitespace = true; - indented_new_text.push('\n'); - start_ix = line_end + 1; - } - } - buffer.replace_range(..start_ix, ""); - + let mut indented_new_text = reindenter.push(&chunk); // This was the last chunk, push all the buffered content as-is. if is_last_chunk { - indented_new_text.push_str(&buffer); - buffer.clear(); + indented_new_text.push_str(&reindenter.finish()); done = true; } @@ -761,28 +710,6 @@ struct ResolvedOldText { indent: LineIndent, } -#[derive(Copy, Clone, Debug)] -enum IndentDelta { - Spaces(isize), - Tabs(isize), -} - -impl IndentDelta { - fn character(&self) -> char { - match self { - IndentDelta::Spaces(_) => ' ', - IndentDelta::Tabs(_) => '\t', - } - } - - fn len(&self) -> isize { - match self { - IndentDelta::Spaces(n) => *n, - IndentDelta::Tabs(n) => *n, - } - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/agent/src/edit_agent/reindent.rs b/crates/agent/src/edit_agent/reindent.rs new file mode 100644 index 0000000000000000000000000000000000000000..7f08749e475f6acfcf63013abd9139574112e4b5 --- /dev/null +++ b/crates/agent/src/edit_agent/reindent.rs @@ -0,0 +1,214 @@ +use language::LineIndent; +use std::{cmp, iter}; + +#[derive(Copy, Clone, Debug)] +pub enum IndentDelta { + Spaces(isize), + Tabs(isize), +} + +impl IndentDelta { + pub fn character(&self) -> char { + match self { + IndentDelta::Spaces(_) => ' ', + IndentDelta::Tabs(_) => '\t', + } + } + + pub fn len(&self) -> isize { + match self { + IndentDelta::Spaces(n) => *n, + IndentDelta::Tabs(n) => *n, + } + } +} + +pub fn compute_indent_delta(buffer_indent: LineIndent, query_indent: LineIndent) -> IndentDelta { + if buffer_indent.tabs > 0 { + IndentDelta::Tabs(buffer_indent.tabs as isize - query_indent.tabs as isize) + } else { + IndentDelta::Spaces(buffer_indent.spaces as isize - query_indent.spaces as isize) + } +} + +/// Synchronous re-indentation adapter. Buffers incomplete lines and applies +/// an `IndentDelta` to each line's leading whitespace before emitting it. +pub struct Reindenter { + delta: IndentDelta, + buffer: String, + in_leading_whitespace: bool, +} + +impl Reindenter { + pub fn new(delta: IndentDelta) -> Self { + Self { + delta, + buffer: String::new(), + in_leading_whitespace: true, + } + } + + /// Feed a chunk of text and return the re-indented portion that is + /// ready to emit. Incomplete trailing lines are buffered internally. + pub fn push(&mut self, chunk: &str) -> String { + self.buffer.push_str(chunk); + self.drain(false) + } + + /// Flush any remaining buffered content (call when the stream is done). + pub fn finish(&mut self) -> String { + self.drain(true) + } + + fn drain(&mut self, is_final: bool) -> String { + let mut indented = String::new(); + let mut start_ix = 0; + let mut newlines = self.buffer.match_indices('\n'); + loop { + let (line_end, is_pending_line) = match newlines.next() { + Some((ix, _)) => (ix, false), + None => (self.buffer.len(), true), + }; + let line = &self.buffer[start_ix..line_end]; + + if self.in_leading_whitespace { + if let Some(non_whitespace_ix) = line.find(|c| self.delta.character() != c) { + // We found a non-whitespace character, adjust indentation + // based on the delta. + let new_indent_len = + cmp::max(0, non_whitespace_ix as isize + self.delta.len()) as usize; + indented.extend(iter::repeat(self.delta.character()).take(new_indent_len)); + indented.push_str(&line[non_whitespace_ix..]); + self.in_leading_whitespace = false; + } else if is_pending_line && !is_final { + // We're still in leading whitespace and this line is incomplete. + // Stop processing until we receive more input. + break; + } else { + // This line is entirely whitespace. Push it without indentation. + indented.push_str(line); + } + } else { + indented.push_str(line); + } + + if is_pending_line { + start_ix = line_end; + break; + } else { + self.in_leading_whitespace = true; + indented.push('\n'); + start_ix = line_end + 1; + } + } + self.buffer.replace_range(..start_ix, ""); + if is_final { + indented.push_str(&self.buffer); + self.buffer.clear(); + } + indented + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_indent_single_chunk() { + let mut r = Reindenter::new(IndentDelta::Spaces(2)); + let out = r.push(" abc\n def\n ghi"); + // All three lines are emitted: "ghi" starts with spaces but + // contains non-whitespace, so it's processed immediately. + assert_eq!(out, " abc\n def\n ghi"); + let out = r.finish(); + assert_eq!(out, ""); + } + + #[test] + fn test_outdent_tabs() { + let mut r = Reindenter::new(IndentDelta::Tabs(-2)); + let out = r.push("\t\t\t\tabc\n\t\tdef\n\t\t\t\t\t\tghi"); + assert_eq!(out, "\t\tabc\ndef\n\t\t\t\tghi"); + let out = r.finish(); + assert_eq!(out, ""); + } + + #[test] + fn test_incremental_chunks() { + let mut r = Reindenter::new(IndentDelta::Spaces(2)); + // Feed " ab" — the `a` is non-whitespace, so the line is + // processed immediately even without a trailing newline. + let out = r.push(" ab"); + assert_eq!(out, " ab"); + // Feed "c\n" — appended to the already-processed line (no longer + // in leading whitespace). + let out = r.push("c\n"); + assert_eq!(out, "c\n"); + let out = r.finish(); + assert_eq!(out, ""); + } + + #[test] + fn test_zero_delta() { + let mut r = Reindenter::new(IndentDelta::Spaces(0)); + let out = r.push(" hello\n world\n"); + assert_eq!(out, " hello\n world\n"); + let out = r.finish(); + assert_eq!(out, ""); + } + + #[test] + fn test_clamp_negative_indent() { + let mut r = Reindenter::new(IndentDelta::Spaces(-10)); + let out = r.push(" abc\n"); + // max(0, 2 - 10) = 0, so no leading spaces. + assert_eq!(out, "abc\n"); + let out = r.finish(); + assert_eq!(out, ""); + } + + #[test] + fn test_whitespace_only_lines() { + let mut r = Reindenter::new(IndentDelta::Spaces(2)); + let out = r.push(" \n code\n"); + // First line is all whitespace — emitted verbatim. Second line is indented. + assert_eq!(out, " \n code\n"); + let out = r.finish(); + assert_eq!(out, ""); + } + + #[test] + fn test_compute_indent_delta_spaces() { + let buffer = LineIndent { + tabs: 0, + spaces: 8, + line_blank: false, + }; + let query = LineIndent { + tabs: 0, + spaces: 4, + line_blank: false, + }; + let delta = compute_indent_delta(buffer, query); + assert_eq!(delta.len(), 4); + assert_eq!(delta.character(), ' '); + } + + #[test] + fn test_compute_indent_delta_tabs() { + let buffer = LineIndent { + tabs: 2, + spaces: 0, + line_blank: false, + }; + let query = LineIndent { + tabs: 3, + spaces: 0, + line_blank: false, + }; + let delta = compute_indent_delta(buffer, query); + assert_eq!(delta.len(), -1); + assert_eq!(delta.character(), '\t'); + } +} diff --git a/crates/agent/src/tools.rs b/crates/agent/src/tools.rs index b2724801befc7459ad37494d298819f4b7ca6b27..446472e0c459aa15fa57bb8b49178b08e6781d11 100644 --- a/crates/agent/src/tools.rs +++ b/crates/agent/src/tools.rs @@ -17,6 +17,7 @@ mod save_file_tool; mod spawn_agent_tool; mod streaming_edit_file_tool; mod terminal_tool; +mod tool_edit_parser; mod tool_permissions; mod web_search_tool; diff --git a/crates/agent/src/tools/streaming_edit_file_tool.rs b/crates/agent/src/tools/streaming_edit_file_tool.rs index 20dfe0ab18aa05e6b90125f1c50a1b8a66ab25f9..2658e372d77044b60648d8fab39e458f02dba23d 100644 --- a/crates/agent/src/tools/streaming_edit_file_tool.rs +++ b/crates/agent/src/tools/streaming_edit_file_tool.rs @@ -1,13 +1,17 @@ use super::edit_file_tool::EditFileTool; use super::restore_file_from_disk_tool::RestoreFileFromDiskTool; use super::save_file_tool::SaveFileTool; +use super::tool_edit_parser::{ToolEditEvent, ToolEditParser}; use crate::{ AgentTool, Thread, ToolCallEventStream, ToolInput, - edit_agent::streaming_fuzzy_matcher::StreamingFuzzyMatcher, + edit_agent::{ + reindent::{Reindenter, compute_indent_delta}, + streaming_fuzzy_matcher::StreamingFuzzyMatcher, + }, }; use acp_thread::Diff; use agent_client_protocol::{self as acp, ToolCallLocation, ToolCallUpdateFields}; -use anyhow::{Context as _, Result, anyhow}; +use anyhow::{Context as _, Result}; use collections::HashSet; use futures::FutureExt as _; use gpui::{App, AppContext, AsyncApp, Entity, Task, WeakEntity}; @@ -15,16 +19,15 @@ use language::language_settings::{self, FormatOnSave}; use language::{Buffer, LanguageRegistry}; use language_model::LanguageModelToolResultContent; use project::lsp_store::{FormatTrigger, LspFormatTarget}; -use project::{Project, ProjectPath}; +use project::{AgentLocation, Project, ProjectPath}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use std::ops::Range; use std::path::PathBuf; use std::sync::Arc; -use text::{BufferSnapshot, ToOffset as _}; +use streaming_diff::{CharOperation, StreamingDiff}; use ui::SharedString; use util::rel_path::RelPath; -use util::{Deferred, ResultExt, debug_panic}; +use util::{Deferred, ResultExt}; const DEFAULT_UI_TEXT: &str = "Editing file"; @@ -70,14 +73,13 @@ pub struct StreamingEditFileToolInput { pub path: String, /// The mode of operation on the file. Possible values: - /// - 'create': Create a new file if it doesn't exist. Requires 'content' field. - /// - 'overwrite': Replace the entire contents of an existing file. Requires 'content' field. + /// - 'write': Replace the entire contents of the file. If the file doesn't exist, it will be created. Requires 'content' field. /// - 'edit': Make granular edits to an existing file. Requires 'edits' field. /// /// When a file already exists or you just created it, prefer editing it as opposed to recreating it from scratch. pub mode: StreamingEditFileMode, - /// The complete content for the new file (required for 'create' and 'overwrite' modes). + /// The complete content for the new file (required for 'write' mode). /// This field should contain the entire file content. #[serde(default, skip_serializing_if = "Option::is_none")] pub content: Option, @@ -85,23 +87,22 @@ pub struct StreamingEditFileToolInput { /// List of edit operations to apply sequentially (required for 'edit' mode). /// Each edit finds `old_text` in the file and replaces it with `new_text`. #[serde(default, skip_serializing_if = "Option::is_none")] - pub edits: Option>, + pub edits: Option>, } #[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub enum StreamingEditFileMode { - /// Create a new file if it doesn't exist - Create, - /// Replace the entire contents of an existing file - Overwrite, + /// Overwrite the file with new content (replacing any existing content). + /// If the file does not exist, it will be created. + Write, /// Make granular edits to an existing file Edit, } /// A single edit operation that replaces old text with new text #[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)] -pub struct EditOperation { +pub struct Edit { /// The exact text to find in the file. This will be matched using fuzzy matching /// to handle minor differences in whitespace or formatting. pub old_text: String, @@ -118,271 +119,328 @@ struct StreamingEditFileToolPartialInput { #[serde(default)] mode: Option, #[serde(default)] - #[allow(dead_code)] content: Option, #[serde(default)] - edits: Option>, + edits: Option>, } #[derive(Default, Debug, Deserialize)] -struct PartialEditOperation { +pub struct PartialEdit { #[serde(default)] - old_text: Option, + pub old_text: Option, #[serde(default)] - new_text: Option, + pub new_text: Option, } -enum StreamingEditState { - Idle, - BufferResolved { - abs_path: PathBuf, - buffer: Entity, +#[derive(Debug, Serialize, Deserialize)] +#[serde(untagged)] +pub enum StreamingEditFileToolOutput { + Success { + #[serde(alias = "original_path")] + input_path: PathBuf, + new_text: String, old_text: Arc, - diff: Entity, - mode: StreamingEditFileMode, - last_content_len: usize, - edit_state: IncrementalEditState, - _finalize_diff_guard: Deferred>, + #[serde(default)] + diff: String, + }, + Error { + error: String, }, } -#[derive(Default)] -struct IncrementalEditState { - in_progress_matcher: Option, - last_old_text_len: usize, - applied_ranges: Vec>, +impl StreamingEditFileToolOutput { + pub fn error(error: impl Into) -> Self { + Self::Error { + error: error.into(), + } + } +} + +impl std::fmt::Display for StreamingEditFileToolOutput { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + StreamingEditFileToolOutput::Success { + diff, input_path, .. + } => { + if diff.is_empty() { + write!(f, "No edits were made.") + } else { + write!( + f, + "Edited {}:\n\n```diff\n{diff}\n```", + input_path.display() + ) + } + } + StreamingEditFileToolOutput::Error { error } => write!(f, "{error}"), + } + } } -impl IncrementalEditState { - fn applied_count(&self) -> usize { - self.applied_ranges.len() +impl From for LanguageModelToolResultContent { + fn from(output: StreamingEditFileToolOutput) -> Self { + output.to_string().into() } } -impl StreamingEditState { - async fn finalize( - &mut self, - input: StreamingEditFileToolInput, - tool: &StreamingEditFileTool, +pub struct StreamingEditFileTool { + thread: WeakEntity, + language_registry: Arc, + project: Entity, +} + +impl StreamingEditFileTool { + pub fn new( + project: Entity, + thread: WeakEntity, + language_registry: Arc, + ) -> Self { + Self { + project, + thread, + language_registry, + } + } + + fn authorize( + &self, + path: &PathBuf, + description: &str, event_stream: &ToolCallEventStream, - cx: &mut AsyncApp, - ) -> Result { - let remaining_edits_start_ix = match self { - StreamingEditState::Idle => { - *self = Self::transition_to_buffer_resolved( - &input.path, - &input.display_description, - input.mode.clone(), - tool, - event_stream, - cx, - ) - .await?; - 0 - } - StreamingEditState::BufferResolved { edit_state, .. } => edit_state.applied_count(), - }; + cx: &mut App, + ) -> Task> { + super::tool_permissions::authorize_file_edit( + EditFileTool::NAME, + path, + description, + &self.thread, + event_stream, + cx, + ) + } - let StreamingEditState::BufferResolved { - buffer, - old_text, - diff, - abs_path, - .. - } = self - else { - debug_panic!("Invalid state"); - return Ok(StreamingEditFileToolOutput::Error { - error: "Internal error. Try to apply the edits again".to_string(), - }); - }; + fn set_agent_location(&self, buffer: WeakEntity, position: text::Anchor, cx: &mut App) { + self.project.update(cx, |project, cx| { + project.set_agent_location(Some(AgentLocation { buffer, position }), cx); + }); + } +} - let result: anyhow::Result = async { - let action_log = tool - .thread - .read_with(cx, |thread, _cx| thread.action_log().clone())?; +impl AgentTool for StreamingEditFileTool { + type Input = StreamingEditFileToolInput; + type Output = StreamingEditFileToolOutput; - match input.mode { - StreamingEditFileMode::Create | StreamingEditFileMode::Overwrite => { - action_log.update(cx, |log, cx| { - log.buffer_created(buffer.clone(), cx); - }); - let content = input.content.ok_or_else(|| { - anyhow!("'content' field is required for create and overwrite modes") - })?; - cx.update(|cx| { - buffer.update(cx, |buffer, cx| { - buffer.edit([(0..buffer.len(), content.as_str())], None, cx); - }); - action_log.update(cx, |log, cx| { - log.buffer_edited(buffer.clone(), cx); - }); - }); - } - StreamingEditFileMode::Edit => { - let edits = input - .edits - .ok_or_else(|| anyhow!("'edits' field is required for edit mode"))?; - - let remaining_edits = &edits[remaining_edits_start_ix..]; - apply_edits( - &buffer, - &action_log, - remaining_edits, - &diff, - event_stream, - &abs_path, - cx, - )?; - } - } + const NAME: &'static str = "streaming_edit_file"; - let format_on_save_enabled = buffer.read_with(cx, |buffer, cx| { - let settings = language_settings::language_settings( - buffer.language().map(|l| l.name()), - buffer.file(), - cx, - ); - settings.format_on_save != FormatOnSave::Off - }); + fn supports_input_streaming() -> bool { + true + } - if format_on_save_enabled { - action_log.update(cx, |log, cx| { - log.buffer_edited(buffer.clone(), cx); - }); + fn kind() -> acp::ToolKind { + acp::ToolKind::Edit + } - let format_task = tool.project.update(cx, |project, cx| { - project.format( - HashSet::from_iter([buffer.clone()]), - LspFormatTarget::Buffers, - false, - FormatTrigger::Save, - cx, - ) - }); - futures::select! { - result = format_task.fuse() => { result.log_err(); }, - _ = event_stream.cancelled_by_user().fuse() => { - anyhow::bail!("Edit cancelled by user"); + fn initial_title( + &self, + input: Result, + cx: &mut App, + ) -> SharedString { + match input { + Ok(input) => self + .project + .read(cx) + .find_project_path(&input.path, cx) + .and_then(|project_path| { + self.project + .read(cx) + .short_full_path_for_project_path(&project_path, cx) + }) + .unwrap_or(input.path) + .into(), + Err(raw_input) => { + if let Some(input) = + serde_json::from_value::(raw_input).ok() + { + let path = input.path.unwrap_or_default(); + let path = path.trim(); + if !path.is_empty() { + return self + .project + .read(cx) + .find_project_path(&path, cx) + .and_then(|project_path| { + self.project + .read(cx) + .short_full_path_for_project_path(&project_path, cx) + }) + .unwrap_or_else(|| path.to_string()) + .into(); } - }; - } - let save_task = tool - .project - .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx)); - futures::select! { - result = save_task.fuse() => { result?; }, - _ = event_stream.cancelled_by_user().fuse() => { - anyhow::bail!("Edit cancelled by user"); + let description = input.display_description.unwrap_or_default(); + let description = description.trim(); + if !description.is_empty() { + return description.to_string().into(); + } } - }; - - action_log.update(cx, |log, cx| { - log.buffer_edited(buffer.clone(), cx); - }); - if let Some(new_mtime) = buffer.read_with(cx, |buffer, _| { - buffer.file().and_then(|file| file.disk_state().mtime()) - }) { - tool.thread.update(cx, |thread, _| { - thread - .file_read_times - .insert(abs_path.to_path_buf(), new_mtime); - })?; + DEFAULT_UI_TEXT.into() } + } + } - let new_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); - let (new_text, unified_diff) = cx - .background_spawn({ - let new_snapshot = new_snapshot.clone(); - let old_text = old_text.clone(); - async move { - let new_text = new_snapshot.text(); - let diff = language::unified_diff(&old_text, &new_text); - (new_text, diff) + fn run( + self: Arc, + mut input: ToolInput, + event_stream: ToolCallEventStream, + cx: &mut App, + ) -> Task> { + cx.spawn(async move |cx: &mut AsyncApp| { + let mut state: Option = None; + loop { + futures::select! { + partial = input.recv_partial().fuse() => { + let Some(partial_value) = partial else { break }; + if let Ok(parsed) = serde_json::from_value::(partial_value) { + if state.is_none() && let Some(path_str) = &parsed.path + && let Some(display_description) = &parsed.display_description + && let Some(mode) = parsed.mode.clone() { + state = Some( + EditSession::new( + path_str, + display_description, + mode, + &self, + &event_stream, + cx, + ) + .await?, + ); + } + + if let Some(state) = &mut state { + state.process(parsed, &self, &event_stream, cx)?; + } + } } - }) - .await; + _ = event_stream.cancelled_by_user().fuse() => { + return Err(StreamingEditFileToolOutput::error("Edit cancelled by user")); + } + } + } + let full_input = + input + .recv() + .await + .map_err(|e| StreamingEditFileToolOutput::error(format!("Failed to receive tool input: {e}")))?; - let output = StreamingEditFileToolOutput::Success { - input_path: PathBuf::from(input.path), - new_text, - old_text: old_text.clone(), - diff: unified_diff, + let mut state = if let Some(state) = state { + state + } else { + EditSession::new( + &full_input.path, + &full_input.display_description, + full_input.mode.clone(), + &self, + &event_stream, + cx, + ) + .await? }; - Ok(output) - } - .await; - result.map_err(|e| StreamingEditFileToolOutput::Error { - error: e.to_string(), + state.finalize(full_input, &self, &event_stream, cx).await }) } - async fn process( - &mut self, - partial: StreamingEditFileToolPartialInput, - tool: &StreamingEditFileTool, - event_stream: &ToolCallEventStream, - cx: &mut AsyncApp, - ) -> Result<(), StreamingEditFileToolOutput> { - match self { - Self::Idle => { - if let Some(path_str) = partial.path - && let Some(display_description) = partial.display_description - && let Some(mode) = partial.mode - { - *self = Self::transition_to_buffer_resolved( - &path_str, - &display_description, - mode, - tool, - event_stream, + fn replay( + &self, + _input: Self::Input, + output: Self::Output, + event_stream: ToolCallEventStream, + cx: &mut App, + ) -> Result<()> { + match output { + StreamingEditFileToolOutput::Success { + input_path, + old_text, + new_text, + .. + } => { + event_stream.update_diff(cx.new(|cx| { + Diff::finalized( + input_path.to_string_lossy().into_owned(), + Some(old_text.to_string()), + new_text, + self.language_registry.clone(), cx, ) - .await?; - } + })); + Ok(()) } - Self::BufferResolved { - abs_path, - buffer, - edit_state, - diff, - mode, - last_content_len, - .. - } => match mode { - StreamingEditFileMode::Create | StreamingEditFileMode::Overwrite => { - if let Some(content) = &partial.content { - Self::process_streaming_content( - buffer, - diff, - last_content_len, - content, - cx, - )?; - } - } - StreamingEditFileMode::Edit => { - if let Some(edits) = partial.edits { - Self::process_streaming_edits( - buffer, - diff, - edit_state, - &edits, - abs_path, - tool, - event_stream, - cx, - )?; - } - } - }, + StreamingEditFileToolOutput::Error { .. } => Ok(()), + } + } +} + +pub struct EditSession { + abs_path: PathBuf, + buffer: Entity, + old_text: Arc, + diff: Entity, + mode: StreamingEditFileMode, + parser: ToolEditParser, + pipeline: EditPipeline, + _finalize_diff_guard: Deferred>, +} + +struct EditPipeline { + edits: Vec, + content_written: bool, +} + +enum EditPipelineEntry { + ResolvingOldText { + matcher: StreamingFuzzyMatcher, + }, + StreamingNewText { + streaming_diff: StreamingDiff, + edit_cursor: usize, + reindenter: Reindenter, + original_snapshot: text::BufferSnapshot, + }, + Done, +} + +impl EditPipeline { + fn new() -> Self { + Self { + edits: Vec::new(), + content_written: false, } - Ok(()) } - async fn transition_to_buffer_resolved( + fn ensure_resolving_old_text( + &mut self, + edit_index: usize, + buffer: &Entity, + cx: &mut AsyncApp, + ) { + while self.edits.len() <= edit_index { + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.text_snapshot()); + self.edits.push(EditPipelineEntry::ResolvingOldText { + matcher: StreamingFuzzyMatcher::new(snapshot), + }); + } + } +} + +/// Compute the `LineIndent` of the first line in a set of query lines. +fn query_first_line_indent(query_lines: &[String]) -> text::LineIndent { + let first_line = query_lines.first().map(|s| s.as_str()).unwrap_or(""); + text::LineIndent::from_iter(first_line.chars()) +} + +impl EditSession { + async fn new( path_str: &str, display_description: &str, mode: StreamingEditFileMode, @@ -393,15 +451,13 @@ impl StreamingEditState { let path = PathBuf::from(path_str); let project_path = cx .update(|cx| resolve_path(mode.clone(), &path, &tool.project, cx)) - .map_err(|e| StreamingEditFileToolOutput::Error { - error: e.to_string(), - })?; + .map_err(|e| StreamingEditFileToolOutput::error(e.to_string()))?; let Some(abs_path) = cx.update(|cx| tool.project.read(cx).absolute_path(&project_path, cx)) else { - return Err(StreamingEditFileToolOutput::Error { - error: format!("File '{path_str}' does not exist"), - }); + return Err(StreamingEditFileToolOutput::error(format!( + "Worktree at '{path_str}' does not exist" + ))); }; event_stream.update_fields( @@ -410,17 +466,13 @@ impl StreamingEditState { cx.update(|cx| tool.authorize(&path, &display_description, event_stream, cx)) .await - .map_err(|e| StreamingEditFileToolOutput::Error { - error: e.to_string(), - })?; + .map_err(|e| StreamingEditFileToolOutput::error(e.to_string()))?; let buffer = tool .project .update(cx, |project, cx| project.open_buffer(project_path, cx)) .await - .map_err(|e| StreamingEditFileToolOutput::Error { - error: e.to_string(), - })?; + .map_err(|e| StreamingEditFileToolOutput::error(e.to_string()))?; ensure_buffer_saved(&buffer, &abs_path, tool, cx)?; @@ -434,6 +486,14 @@ impl StreamingEditState { } }) as Box); + tool.thread + .update(cx, |thread, cx| { + thread + .action_log() + .update(cx, |log, cx| log.buffer_read(buffer.clone(), cx)) + }) + .ok(); + let old_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); let old_text = cx .background_spawn({ @@ -442,205 +502,497 @@ impl StreamingEditState { }) .await; - Ok(Self::BufferResolved { + Ok(Self { abs_path, buffer, old_text, diff, mode, - last_content_len: 0, - edit_state: IncrementalEditState::default(), + parser: ToolEditParser::default(), + pipeline: EditPipeline::new(), _finalize_diff_guard: finalize_diff_guard, }) } - fn process_streaming_content( - buffer: &Entity, - diff: &Entity, - last_content_len: &mut usize, - content: &str, + async fn finalize( + &mut self, + input: StreamingEditFileToolInput, + tool: &StreamingEditFileTool, + event_stream: &ToolCallEventStream, cx: &mut AsyncApp, - ) -> Result<(), StreamingEditFileToolOutput> { - let new_len = content.len(); - if new_len > *last_content_len { - let new_chunk = &content[*last_content_len..]; - cx.update(|cx| { - buffer.update(cx, |buffer, cx| { - // On the first update, replace the entire buffer (handles Overwrite - // clearing existing content). For Create the buffer is already empty - // so 0..0 is a no-op range prefix. - let insert_at = if *last_content_len == 0 { - 0..buffer.len() - } else { - let len = buffer.len(); - len..len - }; - buffer.edit([(insert_at, new_chunk)], None, cx); + ) -> Result { + let Self { + buffer, + old_text, + diff, + abs_path, + parser, + pipeline, + .. + } = self; + + let action_log = tool + .thread + .read_with(cx, |thread, _cx| thread.action_log().clone()) + .map_err(|e| StreamingEditFileToolOutput::error(e.to_string()))?; + + match input.mode { + StreamingEditFileMode::Write => { + action_log.update(cx, |log, cx| { + log.buffer_created(buffer.clone(), cx); }); + let content = input.content.ok_or_else(|| { + StreamingEditFileToolOutput::error("'content' field is required for write mode") + })?; + + let events = parser.finalize_content(&content); + Self::process_events( + &events, + buffer, + diff, + pipeline, + abs_path, + tool, + event_stream, + cx, + )?; + } + StreamingEditFileMode::Edit => { + let edits = input.edits.ok_or_else(|| { + StreamingEditFileToolOutput::error("'edits' field is required for edit mode") + })?; + + let final_edits = edits + .into_iter() + .map(|e| Edit { + old_text: e.old_text, + new_text: e.new_text, + }) + .collect::>(); + let events = parser.finalize_edits(&final_edits); + Self::process_events( + &events, + buffer, + diff, + pipeline, + abs_path, + tool, + event_stream, + cx, + )?; + } + } + + let format_on_save_enabled = buffer.read_with(cx, |buffer, cx| { + let settings = language_settings::language_settings( + buffer.language().map(|l| l.name()), + buffer.file(), + cx, + ); + settings.format_on_save != FormatOnSave::Off + }); + + if format_on_save_enabled { + action_log.update(cx, |log, cx| { + log.buffer_edited(buffer.clone(), cx); }); - *last_content_len = new_len; - let anchor_range = buffer.read_with(cx, |buffer, _cx| { - buffer.anchor_range_between(0..buffer.len()) + let format_task = tool.project.update(cx, |project, cx| { + project.format( + HashSet::from_iter([buffer.clone()]), + LspFormatTarget::Buffers, + false, + FormatTrigger::Save, + cx, + ) }); - diff.update(cx, |diff, cx| diff.reveal_range(anchor_range, cx)); + futures::select! { + result = format_task.fuse() => { result.log_err(); }, + _ = event_stream.cancelled_by_user().fuse() => { + return Err(StreamingEditFileToolOutput::error("Edit cancelled by user")); + } + }; + } + + let save_task = tool + .project + .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx)); + futures::select! { + result = save_task.fuse() => { result.map_err(|e| StreamingEditFileToolOutput::error(e.to_string()))?; }, + _ = event_stream.cancelled_by_user().fuse() => { + return Err(StreamingEditFileToolOutput::error("Edit cancelled by user")); + } + }; + + action_log.update(cx, |log, cx| { + log.buffer_edited(buffer.clone(), cx); + }); + + if let Some(new_mtime) = buffer.read_with(cx, |buffer, _| { + buffer.file().and_then(|file| file.disk_state().mtime()) + }) { + tool.thread + .update(cx, |thread, _| { + thread + .file_read_times + .insert(abs_path.to_path_buf(), new_mtime); + }) + .map_err(|e| StreamingEditFileToolOutput::error(e.to_string()))?; + } + + let new_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + let (new_text, unified_diff) = cx + .background_spawn({ + let new_snapshot = new_snapshot.clone(); + let old_text = old_text.clone(); + async move { + let new_text = new_snapshot.text(); + let diff = language::unified_diff(&old_text, &new_text); + (new_text, diff) + } + }) + .await; + + let output = StreamingEditFileToolOutput::Success { + input_path: PathBuf::from(input.path), + new_text, + old_text: old_text.clone(), + diff: unified_diff, + }; + Ok(output) + } + + fn process( + &mut self, + partial: StreamingEditFileToolPartialInput, + tool: &StreamingEditFileTool, + event_stream: &ToolCallEventStream, + cx: &mut AsyncApp, + ) -> Result<(), StreamingEditFileToolOutput> { + match &self.mode { + StreamingEditFileMode::Write => { + if let Some(content) = &partial.content { + let events = self.parser.push_content(content); + Self::process_events( + &events, + &self.buffer, + &self.diff, + &mut self.pipeline, + &self.abs_path, + tool, + event_stream, + cx, + )?; + } + } + StreamingEditFileMode::Edit => { + if let Some(edits) = partial.edits { + let events = self.parser.push_edits(&edits); + Self::process_events( + &events, + &self.buffer, + &self.diff, + &mut self.pipeline, + &self.abs_path, + tool, + event_stream, + cx, + )?; + } + } } Ok(()) } - fn process_streaming_edits( + fn process_events( + events: &[ToolEditEvent], buffer: &Entity, diff: &Entity, - edit_state: &mut IncrementalEditState, - edits: &[PartialEditOperation], + pipeline: &mut EditPipeline, abs_path: &PathBuf, tool: &StreamingEditFileTool, event_stream: &ToolCallEventStream, cx: &mut AsyncApp, ) -> Result<(), StreamingEditFileToolOutput> { - if edits.is_empty() { - return Ok(()); - } + for event in events { + match event { + ToolEditEvent::ContentChunk { chunk } => { + cx.update(|cx| { + buffer.update(cx, |buffer, cx| { + let insert_at = if !pipeline.content_written && buffer.len() > 0 { + 0..buffer.len() + } else { + let len = buffer.len(); + len..len + }; + buffer.edit([(insert_at, chunk.as_str())], None, cx); + }); + let buffer_id = buffer.read(cx).remote_id(); + tool.set_agent_location( + buffer.downgrade(), + text::Anchor::max_for_buffer(buffer_id), + cx, + ); + }); + pipeline.content_written = true; + } - // Edits at indices applied_count..edits.len()-1 are newly complete - // (a subsequent edit exists, proving the LLM moved on). - // The last edit (edits.len()-1) is potentially still in progress. - let completed_count = edits.len().saturating_sub(1); + ToolEditEvent::OldTextChunk { + edit_index, + chunk, + done: false, + } => { + pipeline.ensure_resolving_old_text(*edit_index, buffer, cx); + + if let EditPipelineEntry::ResolvingOldText { matcher } = + &mut pipeline.edits[*edit_index] + { + if !chunk.is_empty() { + if let Some(match_range) = matcher.push(chunk, None) { + let anchor_range = buffer.read_with(cx, |buffer, _cx| { + buffer.anchor_range_between(match_range.clone()) + }); + diff.update(cx, |diff, cx| diff.reveal_range(anchor_range, cx)); + + cx.update(|cx| { + let position = buffer.read(cx).anchor_before(match_range.end); + tool.set_agent_location(buffer.downgrade(), position, cx); + }); + } + } + } + } - // Apply newly-complete edits - while edit_state.applied_count() < completed_count { - let edit_index = edit_state.applied_count(); - let partial_edit = &edits[edit_index]; + ToolEditEvent::OldTextChunk { + edit_index, + chunk, + done: true, + } => { + pipeline.ensure_resolving_old_text(*edit_index, buffer, cx); + + let EditPipelineEntry::ResolvingOldText { matcher } = + &mut pipeline.edits[*edit_index] + else { + continue; + }; - let old_text = partial_edit.old_text.clone().ok_or_else(|| { - StreamingEditFileToolOutput::Error { - error: format!("Edit at index {} is missing old_text.", edit_index), - } - })?; - let new_text = partial_edit.new_text.clone().unwrap_or_default(); + if !chunk.is_empty() { + matcher.push(chunk, None); + } + let matches = matcher.finish(); + + if matches.is_empty() { + return Err(StreamingEditFileToolOutput::error(format!( + "Could not find matching text for edit at index {}. \ + The old_text did not match any content in the file. \ + Please read the file again to get the current content.", + edit_index, + ))); + } + if matches.len() > 1 { + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + let lines = matches + .iter() + .map(|r| (snapshot.offset_to_point(r.start).row + 1).to_string()) + .collect::>() + .join(", "); + return Err(StreamingEditFileToolOutput::error(format!( + "Edit {} matched multiple locations in the file at lines: {}. \ + Please provide more context in old_text to uniquely \ + identify the location.", + edit_index, lines + ))); + } - edit_state.in_progress_matcher = None; - edit_state.last_old_text_len = 0; + let range = matches.into_iter().next().expect("checked len above"); - let edit_op = EditOperation { - old_text: old_text.clone(), - new_text: new_text.clone(), - }; + let anchor_range = buffer + .read_with(cx, |buffer, _cx| buffer.anchor_range_between(range.clone())); + diff.update(cx, |diff, cx| diff.reveal_range(anchor_range, cx)); + + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); - let action_log = tool - .thread - .read_with(cx, |thread, _cx| thread.action_log().clone()) - .ok(); + let line = snapshot.offset_to_point(range.start).row; + event_stream.update_fields( + ToolCallUpdateFields::new() + .locations(vec![ToolCallLocation::new(abs_path).line(Some(line))]), + ); + + let EditPipelineEntry::ResolvingOldText { matcher } = + &pipeline.edits[*edit_index] + else { + continue; + }; + let buffer_indent = + snapshot.line_indent_for_row(snapshot.offset_to_point(range.start).row); + let query_indent = query_first_line_indent(matcher.query_lines()); + let indent_delta = compute_indent_delta(buffer_indent, query_indent); + + let old_text_in_buffer = + snapshot.text_for_range(range.clone()).collect::(); + + let text_snapshot = buffer.read_with(cx, |buffer, _cx| buffer.text_snapshot()); + pipeline.edits[*edit_index] = EditPipelineEntry::StreamingNewText { + streaming_diff: StreamingDiff::new(old_text_in_buffer), + edit_cursor: range.start, + reindenter: Reindenter::new(indent_delta), + original_snapshot: text_snapshot, + }; - // On the first edit, mark the buffer as read - if edit_state.applied_count() == 0 { - if let Some(action_log) = &action_log { - action_log.update(cx, |log, cx| { - log.buffer_read(buffer.clone(), cx); + cx.update(|cx| { + let position = buffer.read(cx).anchor_before(range.end); + tool.set_agent_location(buffer.downgrade(), position, cx); }); } - } - let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); - - let (range, new_text) = - match resolve_and_reveal_edit(buffer, diff, &snapshot, &edit_op, cx) { - Ok(resolved) => resolved, - Err(EditResolveError::NotFound) => { - return Err(StreamingEditFileToolOutput::Error { - error: format!( - "Could not find matching text for edit at index {}. \ - The old_text did not match any content in the file. \ - Please read the file again to get the current content.", - edit_index - ), - }); + ToolEditEvent::NewTextChunk { + edit_index, + chunk, + done: false, + } => { + if *edit_index >= pipeline.edits.len() { + continue; } - Err(EditResolveError::Ambiguous(ranges)) => { - let lines = ranges - .iter() - .map(|r| (snapshot.offset_to_point(r.start).row + 1).to_string()) - .collect::>() - .join(", "); - return Err(StreamingEditFileToolOutput::Error { - error: format!( - "Edit {} matched multiple locations in the file at lines: {}. \ - Please provide more context in old_text to uniquely \ - identify the location.", - edit_index, lines - ), - }); + let EditPipelineEntry::StreamingNewText { + streaming_diff, + edit_cursor, + reindenter, + original_snapshot, + .. + } = &mut pipeline.edits[*edit_index] + else { + continue; + }; + + let reindented = reindenter.push(chunk); + if reindented.is_empty() { + continue; } - }; - for previous_range in &edit_state.applied_ranges { - let previous_start = previous_range.start.to_offset(&snapshot); - let previous_end = previous_range.end.to_offset(&snapshot); - if range.start < previous_end && previous_start < range.end { - let earlier_start_line = snapshot.offset_to_point(previous_start).row + 1; - let earlier_end_line = snapshot.offset_to_point(previous_end).row + 1; - let later_start_line = snapshot.offset_to_point(range.start).row + 1; - let later_end_line = snapshot.offset_to_point(range.end).row + 1; - return Err(StreamingEditFileToolOutput::Error { - error: format!( - "Conflicting edit ranges detected: lines {}-{} \ - conflicts with lines {}-{}. Conflicting edit \ - ranges are not allowed, as they would overwrite \ - each other.", - earlier_start_line, earlier_end_line, later_start_line, later_end_line, - ), + let char_ops = streaming_diff.push_new(&reindented); + Self::apply_char_operations( + &char_ops, + buffer, + original_snapshot, + edit_cursor, + cx, + ); + + let position = original_snapshot.anchor_before(*edit_cursor); + cx.update(|cx| { + tool.set_agent_location(buffer.downgrade(), position, cx); }); + + let action_log = tool + .thread + .read_with(cx, |thread, _cx| thread.action_log().clone()) + .ok(); + if let Some(action_log) = action_log { + action_log.update(cx, |log, cx| { + log.buffer_edited(buffer.clone(), cx); + }); + } } - } - let anchor_range = - buffer.read_with(cx, |buffer, _cx| buffer.anchor_range_between(range.clone())); - edit_state.applied_ranges.push(anchor_range); + ToolEditEvent::NewTextChunk { + edit_index, + chunk, + done: true, + } => { + if *edit_index >= pipeline.edits.len() { + continue; + } - let line = snapshot.offset_to_point(range.start).row; - event_stream.update_fields( - ToolCallUpdateFields::new() - .locations(vec![ToolCallLocation::new(abs_path).line(Some(line))]), - ); + let EditPipelineEntry::StreamingNewText { + mut streaming_diff, + mut edit_cursor, + mut reindenter, + original_snapshot, + } = std::mem::replace( + &mut pipeline.edits[*edit_index], + EditPipelineEntry::Done, + ) + else { + continue; + }; - if let Some(action_log) = action_log { - cx.update(|cx| { - buffer.update(cx, |buffer, cx| { - buffer.edit([(range, new_text.as_str())], None, cx); - }); - action_log.update(cx, |log, cx| { - log.buffer_edited(buffer.clone(), cx); - }); - }); - } - } + // Flush any remaining reindent buffer + final chunk. + let mut final_text = reindenter.push(chunk); + final_text.push_str(&reindenter.finish()); - // Feed the in-progress last edit's old_text to the matcher for live preview - if let Some(partial_edit) = edits.last() { - if let Some(old_text) = &partial_edit.old_text { - let old_text_len = old_text.len(); - if old_text_len > edit_state.last_old_text_len { - let new_chunk = &old_text[edit_state.last_old_text_len..]; + if !final_text.is_empty() { + let char_ops = streaming_diff.push_new(&final_text); + Self::apply_char_operations( + &char_ops, + buffer, + &original_snapshot, + &mut edit_cursor, + cx, + ); + } + + let remaining_ops = streaming_diff.finish(); + Self::apply_char_operations( + &remaining_ops, + buffer, + &original_snapshot, + &mut edit_cursor, + cx, + ); - let matcher = edit_state.in_progress_matcher.get_or_insert_with(|| { - let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.text_snapshot()); - StreamingFuzzyMatcher::new(snapshot) + let position = original_snapshot.anchor_before(edit_cursor); + cx.update(|cx| { + tool.set_agent_location(buffer.downgrade(), position, cx); }); - if let Some(match_range) = matcher.push(new_chunk, None) { - let anchor_range = buffer.read_with(cx, |buffer, _cx| { - buffer.anchor_range_between(match_range.clone()) + let action_log = tool + .thread + .read_with(cx, |thread, _cx| thread.action_log().clone()) + .ok(); + if let Some(action_log) = action_log { + action_log.update(cx, |log, cx| { + log.buffer_edited(buffer.clone(), cx); }); - diff.update(cx, |card, cx| card.reveal_range(anchor_range, cx)); } - - edit_state.last_old_text_len = old_text_len; } } } - Ok(()) } + + fn apply_char_operations( + ops: &[CharOperation], + buffer: &Entity, + snapshot: &text::BufferSnapshot, + edit_cursor: &mut usize, + cx: &mut AsyncApp, + ) { + for op in ops { + match op { + CharOperation::Insert { text } => { + let anchor = snapshot.anchor_after(*edit_cursor); + cx.update(|cx| { + buffer.update(cx, |buffer, cx| { + buffer.edit([(anchor..anchor, text.as_str())], None, cx); + }); + }); + } + CharOperation::Delete { bytes } => { + let delete_end = *edit_cursor + bytes; + let anchor_range = snapshot.anchor_range_around(*edit_cursor..delete_end); + cx.update(|cx| { + buffer.update(cx, |buffer, cx| { + buffer.edit([(anchor_range, "")], None, cx); + }); + }); + *edit_cursor = delete_end; + } + CharOperation::Keep { bytes } => { + *edit_cursor += bytes; + } + } + } + } } fn ensure_buffer_saved( @@ -670,396 +1022,40 @@ fn ensure_buffer_saved( if is_dirty { let message = match (has_save_tool, has_restore_tool) { (true, true) => { - "This file has unsaved changes. Ask the user whether they want to keep or discard those changes. \ - If they want to keep them, ask for confirmation then use the save_file tool to save the file, then retry this edit. \ - If they want to discard them, ask for confirmation then use the restore_file_from_disk tool to restore the on-disk contents, then retry this edit." - } - (true, false) => { - "This file has unsaved changes. Ask the user whether they want to keep or discard those changes. \ - If they want to keep them, ask for confirmation then use the save_file tool to save the file, then retry this edit. \ - If they want to discard them, ask the user to manually revert the file, then inform you when it's ok to proceed." - } - (false, true) => { - "This file has unsaved changes. Ask the user whether they want to keep or discard those changes. \ - If they want to keep them, ask the user to manually save the file, then inform you when it's ok to proceed. \ - If they want to discard them, ask for confirmation then use the restore_file_from_disk tool to restore the on-disk contents, then retry this edit." - } - (false, false) => { - "This file has unsaved changes. Ask the user whether they want to keep or discard those changes, \ - then ask them to save or revert the file manually and inform you when it's ok to proceed." - } - }; - return Err(StreamingEditFileToolOutput::Error { - error: message.to_string(), - }); - } - - if let (Some(last_read), Some(current)) = (last_read_mtime, current_mtime) { - if current != last_read { - return Err(StreamingEditFileToolOutput::Error { - error: "The file has been modified since you last read it. \ - Please read the file again to get the current state before editing it." - .to_string(), - }); - } - } - - Ok(()) -} - -#[derive(Debug, Serialize, Deserialize)] -#[serde(untagged)] -pub enum StreamingEditFileToolOutput { - Success { - #[serde(alias = "original_path")] - input_path: PathBuf, - new_text: String, - old_text: Arc, - #[serde(default)] - diff: String, - }, - Error { - error: String, - }, -} - -impl std::fmt::Display for StreamingEditFileToolOutput { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - StreamingEditFileToolOutput::Success { - diff, input_path, .. - } => { - if diff.is_empty() { - write!(f, "No edits were made.") - } else { - write!( - f, - "Edited {}:\n\n```diff\n{diff}\n```", - input_path.display() - ) - } - } - StreamingEditFileToolOutput::Error { error } => write!(f, "{error}"), - } - } -} - -impl From for LanguageModelToolResultContent { - fn from(output: StreamingEditFileToolOutput) -> Self { - output.to_string().into() - } -} - -pub struct StreamingEditFileTool { - thread: WeakEntity, - language_registry: Arc, - project: Entity, -} - -impl StreamingEditFileTool { - pub fn new( - project: Entity, - thread: WeakEntity, - language_registry: Arc, - ) -> Self { - Self { - project, - thread, - language_registry, - } - } - - fn authorize( - &self, - path: &PathBuf, - description: &str, - event_stream: &ToolCallEventStream, - cx: &mut App, - ) -> Task> { - super::tool_permissions::authorize_file_edit( - EditFileTool::NAME, - path, - description, - &self.thread, - event_stream, - cx, - ) - } -} - -impl AgentTool for StreamingEditFileTool { - type Input = StreamingEditFileToolInput; - type Output = StreamingEditFileToolOutput; - - const NAME: &'static str = "streaming_edit_file"; - - fn supports_input_streaming() -> bool { - true - } - - fn kind() -> acp::ToolKind { - acp::ToolKind::Edit - } - - fn initial_title( - &self, - input: Result, - cx: &mut App, - ) -> SharedString { - match input { - Ok(input) => self - .project - .read(cx) - .find_project_path(&input.path, cx) - .and_then(|project_path| { - self.project - .read(cx) - .short_full_path_for_project_path(&project_path, cx) - }) - .unwrap_or(input.path) - .into(), - Err(raw_input) => { - if let Some(input) = - serde_json::from_value::(raw_input).ok() - { - let path = input.path.unwrap_or_default(); - let path = path.trim(); - if !path.is_empty() { - return self - .project - .read(cx) - .find_project_path(&path, cx) - .and_then(|project_path| { - self.project - .read(cx) - .short_full_path_for_project_path(&project_path, cx) - }) - .unwrap_or_else(|| path.to_string()) - .into(); - } - - let description = input.display_description.unwrap_or_default(); - let description = description.trim(); - if !description.is_empty() { - return description.to_string().into(); - } - } - - DEFAULT_UI_TEXT.into() - } - } - } - - fn run( - self: Arc, - mut input: ToolInput, - event_stream: ToolCallEventStream, - cx: &mut App, - ) -> Task> { - cx.spawn(async move |cx: &mut AsyncApp| { - let mut state = StreamingEditState::Idle; - loop { - futures::select! { - partial = input.recv_partial().fuse() => { - let Some(partial_value) = partial else { break }; - if let Ok(parsed) = serde_json::from_value::(partial_value) { - state.process(parsed, &self, &event_stream, cx).await?; - } - } - _ = event_stream.cancelled_by_user().fuse() => { - return Err(StreamingEditFileToolOutput::Error { - error: "Edit cancelled by user".to_string(), - }); - } - } - } - let full_input = - input - .recv() - .await - .map_err(|e| StreamingEditFileToolOutput::Error { - error: format!("Failed to receive tool input: {e}"), - })?; - - state.finalize(full_input, &self, &event_stream, cx).await - }) - } - - fn replay( - &self, - _input: Self::Input, - output: Self::Output, - event_stream: ToolCallEventStream, - cx: &mut App, - ) -> Result<()> { - match output { - StreamingEditFileToolOutput::Success { - input_path, - old_text, - new_text, - .. - } => { - event_stream.update_diff(cx.new(|cx| { - Diff::finalized( - input_path.to_string_lossy().into_owned(), - Some(old_text.to_string()), - new_text, - self.language_registry.clone(), - cx, - ) - })); - Ok(()) - } - StreamingEditFileToolOutput::Error { .. } => Ok(()), - } - } -} - -fn apply_edits( - buffer: &Entity, - action_log: &Entity, - edits: &[EditOperation], - diff: &Entity, - event_stream: &ToolCallEventStream, - abs_path: &PathBuf, - cx: &mut AsyncApp, -) -> Result<()> { - let mut failed_edits = Vec::new(); - let mut ambiguous_edits = Vec::new(); - let mut resolved_edits: Vec<(Range, String)> = Vec::new(); - - let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); - for (index, edit) in edits.iter().enumerate() { - match resolve_and_reveal_edit(buffer, diff, &snapshot, edit, cx) { - Ok((range, new_text)) => { - resolved_edits.push((range, new_text)); + "This file has unsaved changes. Ask the user whether they want to keep or discard those changes. \ + If they want to keep them, ask for confirmation then use the save_file tool to save the file, then retry this edit. \ + If they want to discard them, ask for confirmation then use the restore_file_from_disk tool to restore the on-disk contents, then retry this edit." } - Err(EditResolveError::NotFound) => { - failed_edits.push(index); + (true, false) => { + "This file has unsaved changes. Ask the user whether they want to keep or discard those changes. \ + If they want to keep them, ask for confirmation then use the save_file tool to save the file, then retry this edit. \ + If they want to discard them, ask the user to manually revert the file, then inform you when it's ok to proceed." } - Err(EditResolveError::Ambiguous(ranges)) => { - ambiguous_edits.push((index, ranges)); + (false, true) => { + "This file has unsaved changes. Ask the user whether they want to keep or discard those changes. \ + If they want to keep them, ask the user to manually save the file, then inform you when it's ok to proceed. \ + If they want to discard them, ask for confirmation then use the restore_file_from_disk tool to restore the on-disk contents, then retry this edit." } - } - } - - if !failed_edits.is_empty() { - let indices = failed_edits - .iter() - .map(|i| i.to_string()) - .collect::>() - .join(", "); - anyhow::bail!( - "Could not find matching text for edit(s) at index(es): {}. \ - The old_text did not match any content in the file. \ - Please read the file again to get the current content.", - indices - ); - } - - if !ambiguous_edits.is_empty() { - let details: Vec = ambiguous_edits - .iter() - .map(|(index, ranges)| { - let lines = ranges - .iter() - .map(|r| (snapshot.offset_to_point(r.start).row + 1).to_string()) - .collect::>() - .join(", "); - format!("edit {}: matches at lines {}", index, lines) - }) - .collect(); - anyhow::bail!( - "Some edits matched multiple locations in the file:\n{}. \ - Please provide more context in old_text to uniquely identify the location.", - details.join("\n") - ); - } - - let mut edits_sorted = resolved_edits; - edits_sorted.sort_by(|a, b| a.0.start.cmp(&b.0.start)); - - if let Some((first_range, _)) = edits_sorted.first() { - let line = snapshot.offset_to_point(first_range.start).row; - event_stream.update_fields( - ToolCallUpdateFields::new() - .locations(vec![ToolCallLocation::new(abs_path).line(Some(line))]), - ); + (false, false) => { + "This file has unsaved changes. Ask the user whether they want to keep or discard those changes, \ + then ask them to save or revert the file manually and inform you when it's ok to proceed." + } + }; + return Err(StreamingEditFileToolOutput::error(message)); } - for window in edits_sorted.windows(2) { - if let [(earlier_range, _), (later_range, _)] = window - && (earlier_range.end > later_range.start || earlier_range.start == later_range.start) - { - let earlier_start_line = snapshot.offset_to_point(earlier_range.start).row + 1; - let earlier_end_line = snapshot.offset_to_point(earlier_range.end).row + 1; - let later_start_line = snapshot.offset_to_point(later_range.start).row + 1; - let later_end_line = snapshot.offset_to_point(later_range.end).row + 1; - anyhow::bail!( - "Conflicting edit ranges detected: lines {}-{} conflicts with lines {}-{}. \ - Conflicting edit ranges are not allowed, as they would overwrite each other.", - earlier_start_line, - earlier_end_line, - later_start_line, - later_end_line, - ); + if let (Some(last_read), Some(current)) = (last_read_mtime, current_mtime) { + if current != last_read { + return Err(StreamingEditFileToolOutput::error( + "The file has been modified since you last read it. \ + Please read the file again to get the current state before editing it.", + )); } } - if !edits_sorted.is_empty() { - cx.update(|cx| { - buffer.update(cx, |buffer, cx| { - buffer.edit( - edits_sorted - .iter() - .map(|(range, new_text)| (range.clone(), new_text.as_str())), - None, - cx, - ); - }); - action_log.update(cx, |log, cx| { - log.buffer_edited(buffer.clone(), cx); - }); - }); - } - Ok(()) } -enum EditResolveError { - NotFound, - Ambiguous(Vec>), -} - -/// Resolves an edit operation by finding matching text in the buffer, -/// reveals the matched range in the diff view, and returns the resolved -/// range and replacement text. -fn resolve_and_reveal_edit( - buffer: &Entity, - diff: &Entity, - snapshot: &BufferSnapshot, - edit: &EditOperation, - cx: &mut AsyncApp, -) -> std::result::Result<(Range, String), EditResolveError> { - let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone()); - matcher.push(&edit.old_text, None); - let matches = matcher.finish(); - if matches.is_empty() { - return Err(EditResolveError::NotFound); - } - if matches.len() > 1 { - return Err(EditResolveError::Ambiguous(matches)); - } - - let range = matches.into_iter().next().expect("checked len above"); - - let anchor_range = - buffer.read_with(cx, |buffer, _cx| buffer.anchor_range_between(range.clone())); - diff.update(cx, |card, cx| card.reveal_range(anchor_range, cx)); - - Ok((range, edit.new_text.clone())) -} - fn resolve_path( mode: StreamingEditFileMode, path: &PathBuf, @@ -1069,7 +1065,7 @@ fn resolve_path( let project = project.read(cx); match mode { - StreamingEditFileMode::Edit | StreamingEditFileMode::Overwrite => { + StreamingEditFileMode::Edit => { let path = project .find_project_path(&path, cx) .context("Can't edit file: path not found")?; @@ -1081,13 +1077,12 @@ fn resolve_path( anyhow::ensure!(entry.is_file(), "Can't edit file: path is a directory"); Ok(path) } - - StreamingEditFileMode::Create => { - if let Some(path) = project.find_project_path(&path, cx) { - anyhow::ensure!( - project.entry_for_path(&path, cx).is_none(), - "Can't create file: file already exists" - ); + StreamingEditFileMode::Write => { + if let Some(path) = project.find_project_path(&path, cx) + && let Some(entry) = project.entry_for_path(&path, cx) + { + anyhow::ensure!(entry.is_file(), "Can't write to file: path is a directory"); + return Ok(path); } let parent_path = path.parent().context("Can't create file: incorrect path")?; @@ -1162,7 +1157,7 @@ mod tests { let input = StreamingEditFileToolInput { display_description: "Create new file".into(), path: "root/dir/new_file.txt".into(), - mode: StreamingEditFileMode::Create, + mode: StreamingEditFileMode::Write, content: Some("Hello, World!".into()), edits: None, }; @@ -1214,7 +1209,7 @@ mod tests { let input = StreamingEditFileToolInput { display_description: "Overwrite file".into(), path: "root/file.txt".into(), - mode: StreamingEditFileMode::Overwrite, + mode: StreamingEditFileMode::Write, content: Some("new content".into()), edits: None, }; @@ -1276,7 +1271,7 @@ mod tests { path: "root/file.txt".into(), mode: StreamingEditFileMode::Edit, content: None, - edits: Some(vec![EditOperation { + edits: Some(vec![Edit { old_text: "line 2".into(), new_text: "modified line 2".into(), }]), @@ -1301,7 +1296,7 @@ mod tests { } #[gpui::test] - async fn test_streaming_edit_multiple_nonoverlapping_edits(cx: &mut TestAppContext) { + async fn test_streaming_edit_multiple_edits(cx: &mut TestAppContext) { init_test(cx); let fs = project::FakeFs::new(cx.executor()); @@ -1336,11 +1331,11 @@ mod tests { mode: StreamingEditFileMode::Edit, content: None, edits: Some(vec![ - EditOperation { + Edit { old_text: "line 5".into(), new_text: "modified line 5".into(), }, - EditOperation { + Edit { old_text: "line 1".into(), new_text: "modified line 1".into(), }, @@ -1404,11 +1399,11 @@ mod tests { mode: StreamingEditFileMode::Edit, content: None, edits: Some(vec![ - EditOperation { + Edit { old_text: "line 2".into(), new_text: "modified line 2".into(), }, - EditOperation { + Edit { old_text: "line 3".into(), new_text: "modified line 3".into(), }, @@ -1472,11 +1467,11 @@ mod tests { mode: StreamingEditFileMode::Edit, content: None, edits: Some(vec![ - EditOperation { + Edit { old_text: "line 1".into(), new_text: "modified line 1".into(), }, - EditOperation { + Edit { old_text: "line 5".into(), new_text: "modified line 5".into(), }, @@ -1533,7 +1528,7 @@ mod tests { path: "root/nonexistent_file.txt".into(), mode: StreamingEditFileMode::Edit, content: None, - edits: Some(vec![EditOperation { + edits: Some(vec![Edit { old_text: "foo".into(), new_text: "bar".into(), }]), @@ -1587,7 +1582,7 @@ mod tests { path: "root/file.txt".into(), mode: StreamingEditFileMode::Edit, content: None, - edits: Some(vec![EditOperation { + edits: Some(vec![Edit { old_text: "nonexistent text that is not in the file".into(), new_text: "replacement".into(), }]), @@ -1614,79 +1609,6 @@ mod tests { ); } - #[gpui::test] - async fn test_streaming_edit_overlapping_edits_out_of_order(cx: &mut TestAppContext) { - init_test(cx); - - let fs = project::FakeFs::new(cx.executor()); - // Multi-line file so the line-based fuzzy matcher can resolve each edit. - fs.insert_tree( - "/root", - json!({ - "file.txt": "line 1\nline 2\nline 3\nline 4\nline 5\n" - }), - ) - .await; - let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await; - let language_registry = project.read_with(cx, |project, _cx| project.languages().clone()); - let context_server_registry = - cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx)); - let model = Arc::new(FakeLanguageModel::default()); - let thread = cx.new(|cx| { - crate::Thread::new( - project.clone(), - cx.new(|_cx| ProjectContext::default()), - context_server_registry, - Templates::new(), - Some(model), - cx, - ) - }); - - // Edit A spans lines 3-4, edit B spans lines 2-3. They overlap on - // "line 3" and are given in descending file order so the ascending - // sort must reorder them before the pairwise overlap check can - // detect them correctly. - let result = cx - .update(|cx| { - let input = StreamingEditFileToolInput { - display_description: "Overlapping edits".into(), - path: "root/file.txt".into(), - mode: StreamingEditFileMode::Edit, - content: None, - edits: Some(vec![ - EditOperation { - old_text: "line 3\nline 4".into(), - new_text: "SECOND".into(), - }, - EditOperation { - old_text: "line 2\nline 3".into(), - new_text: "FIRST".into(), - }, - ]), - }; - Arc::new(StreamingEditFileTool::new( - project, - thread.downgrade(), - language_registry, - )) - .run( - ToolInput::resolved(input), - ToolCallEventStream::test().0, - cx, - ) - }) - .await; - - let StreamingEditFileToolOutput::Error { error } = result.unwrap_err() else { - panic!("expected error"); - }; - assert!( - error.contains("Conflicting edit ranges detected"), - "Expected 'Conflicting edit ranges detected' but got: {error}" - ); - } - #[gpui::test] async fn test_streaming_early_buffer_open(cx: &mut TestAppContext) { init_test(cx); @@ -1809,7 +1731,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite" + "mode": "write" })); cx.run_until_parked(); @@ -1817,7 +1739,7 @@ mod tests { sender.send_final(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite", + "mode": "write", "content": "new content" })); @@ -2026,14 +1948,14 @@ mod tests { sender.send_partial(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create" + "mode": "write" })); cx.run_until_parked(); sender.send_partial(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create", + "mode": "write", "content": "Hello, " })); cx.run_until_parked(); @@ -2042,7 +1964,7 @@ mod tests { sender.send_final(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create", + "mode": "write", "content": "Hello, World!" })); @@ -2304,14 +2226,16 @@ mod tests { })); cx.run_until_parked(); - // Verify edit 1 applied + // Verify edit 1 fully applied. Edit 2's new_text is being + // streamed: "CCC" is inserted but the old "ccc" isn't deleted + // yet (StreamingDiff::finish runs when edit 3 marks edit 2 done). let buffer_text = project.update(cx, |project, cx| { let pp = project .find_project_path(&PathBuf::from("root/file.txt"), cx) .unwrap(); project.get_open_buffer(&pp, cx).map(|b| b.read(cx).text()) }); - assert_eq!(buffer_text.as_deref(), Some("AAA\nbbb\nccc\nddd\neee\n")); + assert_eq!(buffer_text.as_deref(), Some("AAA\nbbb\nCCCccc\nddd\neee\n")); // Edit 3 appears — edit 2 is now complete and should be applied sender.send_partial(json!({ @@ -2326,14 +2250,15 @@ mod tests { })); cx.run_until_parked(); - // Verify edits 1 and 2 both applied + // Verify edits 1 and 2 fully applied. Edit 3's new_text is being + // streamed: "EEE" is inserted but old "eee" isn't deleted yet. let buffer_text = project.update(cx, |project, cx| { let pp = project .find_project_path(&PathBuf::from("root/file.txt"), cx) .unwrap(); project.get_open_buffer(&pp, cx).map(|b| b.read(cx).text()) }); - assert_eq!(buffer_text.as_deref(), Some("AAA\nbbb\nCCC\nddd\neee\n")); + assert_eq!(buffer_text.as_deref(), Some("AAA\nbbb\nCCC\nddd\nEEEeee\n")); // Send final sender.send_final(json!({ @@ -2466,82 +2391,6 @@ mod tests { ); } - #[gpui::test] - async fn test_streaming_overlapping_edits_detected_naturally(cx: &mut TestAppContext) { - init_test(cx); - - let fs = project::FakeFs::new(cx.executor()); - fs.insert_tree( - "/root", - json!({ - "file.txt": "line 1\nline 2\nline 3\n" - }), - ) - .await; - let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await; - let language_registry = project.read_with(cx, |project, _cx| project.languages().clone()); - let context_server_registry = - cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx)); - let model = Arc::new(FakeLanguageModel::default()); - let thread = cx.new(|cx| { - crate::Thread::new( - project.clone(), - cx.new(|_cx| ProjectContext::default()), - context_server_registry, - Templates::new(), - Some(model), - cx, - ) - }); - - let (sender, input) = ToolInput::::test(); - let (event_stream, _receiver) = ToolCallEventStream::test(); - - let tool = Arc::new(StreamingEditFileTool::new( - project.clone(), - thread.downgrade(), - language_registry, - )); - - let task = cx.update(|cx| tool.run(input, event_stream, cx)); - - // Setup - sender.send_partial(json!({ - "display_description": "Overlapping edits", - "path": "root/file.txt", - "mode": "edit" - })); - cx.run_until_parked(); - - // Edit 1 targets "line 1\nline 2" and replaces it. - // Edit 2 targets "line 2\nline 3" — but after edit 1 is applied, - // "line 2" has been removed so this should fail to match. - // Edit 3 exists to make edit 2 "complete" during streaming. - sender.send_partial(json!({ - "display_description": "Overlapping edits", - "path": "root/file.txt", - "mode": "edit", - "edits": [ - {"old_text": "line 1\nline 2", "new_text": "REPLACED"}, - {"old_text": "line 2\nline 3", "new_text": "ALSO REPLACED"}, - {"old_text": "line 3", "new_text": "DUMMY"} - ] - })); - cx.run_until_parked(); - - // Edit 1 was applied, edit 2 should fail since "line 2" no longer exists - drop(sender); - - let result = task.await; - let StreamingEditFileToolOutput::Error { error } = result.unwrap_err() else { - panic!("expected error"); - }; - assert!( - error.contains("Could not find matching text for edit at index 1"), - "Expected overlapping edit to fail naturally, got: {error}" - ); - } - #[gpui::test] async fn test_streaming_single_edit_no_incremental(cx: &mut TestAppContext) { init_test(cx); @@ -2590,7 +2439,10 @@ mod tests { })); cx.run_until_parked(); - // Buffer should NOT be modified — the single edit is still in-progress + // The edit's old_text and new_text both arrived in one partial, so + // the old_text is resolved and new_text is being streamed via + // StreamingDiff. The buffer reflects the in-progress diff (new text + // inserted, old text not yet fully removed until finalization). let buffer_text = project.update(cx, |project, cx| { let pp = project .find_project_path(&PathBuf::from("root/file.txt"), cx) @@ -2599,8 +2451,8 @@ mod tests { }); assert_eq!( buffer_text.as_deref(), - Some("hello world\n"), - "Single in-progress edit should not be applied during streaming" + Some("goodbye worldhello world\n"), + "In-progress streaming diff: new text inserted, old text not yet removed" ); // Send final — the edit is applied during finalization @@ -2795,12 +2647,12 @@ mod tests { sender.send_partial(json!({ "display_description": "Create", "path": "root/dir/new.txt", - "mode": "create" + "mode": "write" })); sender.send_final(json!({ "display_description": "Create", "path": "root/dir/new.txt", - "mode": "create", + "mode": "write", "content": "streamed content" })); @@ -2813,7 +2665,7 @@ mod tests { #[gpui::test] async fn test_streaming_resolve_path_for_creating_file(cx: &mut TestAppContext) { - let mode = StreamingEditFileMode::Create; + let mode = StreamingEditFileMode::Write; let result = test_resolve_path(&mode, "root/new.txt", cx); assert_resolved_path_eq(result.await, rel_path("new.txt")); @@ -2825,9 +2677,12 @@ mod tests { assert_resolved_path_eq(result.await, rel_path("dir/new.txt")); let result = test_resolve_path(&mode, "root/dir/subdir/existing.txt", cx); + assert_resolved_path_eq(result.await, rel_path("dir/subdir/existing.txt")); + + let result = test_resolve_path(&mode, "root/dir/subdir", cx); assert_eq!( result.await.unwrap_err().to_string(), - "Can't create file: file already exists" + "Can't write to file: path is a directory" ); let result = test_resolve_path(&mode, "root/dir/nonexistent_dir/new.txt", cx); @@ -3003,14 +2858,14 @@ mod tests { sender.send_partial(json!({ "display_description": "Create main function", "path": "root/src/main.rs", - "mode": "overwrite" + "mode": "write" })); cx.run_until_parked(); sender.send_final(json!({ "display_description": "Create main function", "path": "root/src/main.rs", - "mode": "overwrite", + "mode": "write", "content": UNFORMATTED_CONTENT })); @@ -3060,14 +2915,14 @@ mod tests { sender.send_partial(json!({ "display_description": "Update main function", "path": "root/src/main.rs", - "mode": "overwrite" + "mode": "write" })); cx.run_until_parked(); sender.send_final(json!({ "display_description": "Update main function", "path": "root/src/main.rs", - "mode": "overwrite", + "mode": "write", "content": UNFORMATTED_CONTENT })); @@ -3136,7 +2991,7 @@ mod tests { let input = StreamingEditFileToolInput { display_description: "Create main function".into(), path: "root/src/main.rs".into(), - mode: StreamingEditFileMode::Overwrite, + mode: StreamingEditFileMode::Write, content: Some(CONTENT_WITH_TRAILING_WHITESPACE.into()), edits: None, }; @@ -3183,7 +3038,7 @@ mod tests { let input = StreamingEditFileToolInput { display_description: "Update main function".into(), path: "root/src/main.rs".into(), - mode: StreamingEditFileMode::Overwrite, + mode: StreamingEditFileMode::Write, content: Some(CONTENT_WITH_TRAILING_WHITESPACE.into()), edits: None, }; @@ -3904,11 +3759,7 @@ mod tests { language_registry, )); - let modes = vec![ - StreamingEditFileMode::Edit, - StreamingEditFileMode::Create, - StreamingEditFileMode::Overwrite, - ]; + let modes = vec![StreamingEditFileMode::Edit, StreamingEditFileMode::Write]; for _mode in modes { // Test .zed path with different modes @@ -4061,7 +3912,7 @@ mod tests { ToolInput::resolved(StreamingEditFileToolInput { display_description: "Edit file".into(), path: path!("/main.rs").into(), - mode: StreamingEditFileMode::Overwrite, + mode: StreamingEditFileMode::Write, content: Some("new content".into()), edits: None, }), @@ -4090,7 +3941,7 @@ mod tests { ToolInput::resolved(StreamingEditFileToolInput { display_description: "Edit file".into(), path: path!("/main.rs").into(), - mode: StreamingEditFileMode::Overwrite, + mode: StreamingEditFileMode::Write, content: Some("dropped content".into()), edits: None, }), @@ -4171,7 +4022,7 @@ mod tests { path: "root/test.txt".into(), mode: StreamingEditFileMode::Edit, content: None, - edits: Some(vec![EditOperation { + edits: Some(vec![Edit { old_text: "original content".into(), new_text: "modified content".into(), }]), @@ -4196,7 +4047,7 @@ mod tests { path: "root/test.txt".into(), mode: StreamingEditFileMode::Edit, content: None, - edits: Some(vec![EditOperation { + edits: Some(vec![Edit { old_text: "modified content".into(), new_text: "further modified content".into(), }]), @@ -4305,7 +4156,7 @@ mod tests { path: "root/test.txt".into(), mode: StreamingEditFileMode::Edit, content: None, - edits: Some(vec![EditOperation { + edits: Some(vec![Edit { old_text: "externally modified content".into(), new_text: "new content".into(), }]), @@ -4409,7 +4260,7 @@ mod tests { path: "root/test.txt".into(), mode: StreamingEditFileMode::Edit, content: None, - edits: Some(vec![EditOperation { + edits: Some(vec![Edit { old_text: "original content".into(), new_text: "new content".into(), }]), @@ -4441,15 +4292,14 @@ mod tests { } #[gpui::test] - async fn test_streaming_overlapping_edits_detected_early(cx: &mut TestAppContext) { + async fn test_streaming_overlapping_edits_resolved_sequentially(cx: &mut TestAppContext) { init_test(cx); let fs = project::FakeFs::new(cx.executor()); - // The file content is crafted so that edit 1's replacement still - // contains the old_text of edit 2 as a contiguous substring. - // Without early overlap detection, edit 2 would silently match - // inside the already-modified region and corrupt the file instead - // of producing a clear "Conflicting edit ranges" error. + // Edit 1's replacement introduces text that contains edit 2's + // old_text as a substring. Because edits resolve sequentially + // against the current buffer, edit 2 finds a unique match in + // the modified buffer and succeeds. fs.insert_tree( "/root", json!({ @@ -4492,17 +4342,10 @@ mod tests { })); cx.run_until_parked(); - // Edit 1 targets "bbb\nccc" (lines 2-3) and replaces it with - // text that preserves "ccc\nddd" as a contiguous substring in the - // buffer — so edit 2's old_text will still match after edit 1 is - // applied. - // - // Edit 2 targets "ccc\nddd" (lines 3-4), overlapping with edit 1 on - // line 3 ("ccc"). After edit 1 runs, the buffer becomes: - // "aaa\nXXX\nccc\nddd\nddd\neee\n" - // and "ccc\nddd" is still present, so edit 2 would silently - // succeed without early overlap detection. - // + // Edit 1 replaces "bbb\nccc" with "XXX\nccc\nddd", so the + // buffer becomes "aaa\nXXX\nccc\nddd\nddd\neee\n". + // Edit 2's old_text "ccc\nddd" matches the first occurrence + // in the modified buffer and replaces it with "ZZZ". // Edit 3 exists only to mark edit 2 as "complete" during streaming. sender.send_partial(json!({ "display_description": "Overlapping edits", @@ -4529,23 +4372,10 @@ mod tests { })); let result = task.await; - // We expect a "Conflicting edit ranges" error. Currently the overlap - // goes undetected during streaming and the file gets silently - // corrupted, so this assertion will fail until we add early overlap - // detection. - match result { - Err(StreamingEditFileToolOutput::Error { error }) - if error.contains("Conflicting edit ranges") => {} - Err(StreamingEditFileToolOutput::Error { error }) => { - panic!("Expected 'Conflicting edit ranges' error, got different error: {error}"); - } - Ok(output) => { - panic!("Expected 'Conflicting edit ranges' error, but got success: {output}"); - } - Err(other) => { - panic!("Expected 'Conflicting edit ranges' error, got unexpected output: {other}"); - } - } + let StreamingEditFileToolOutput::Success { new_text, .. } = result.unwrap() else { + panic!("expected success"); + }; + assert_eq!(new_text, "aaa\nXXX\nZZZ\nddd\nDUMMY\n"); } #[gpui::test] @@ -4585,7 +4415,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create" + "mode": "write" })); cx.run_until_parked(); @@ -4593,7 +4423,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create", + "mode": "write", "content": "line 1\n" })); cx.run_until_parked(); @@ -4611,7 +4441,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create", + "mode": "write", "content": "line 1\nline 2\n" })); cx.run_until_parked(); @@ -4621,7 +4451,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create", + "mode": "write", "content": "line 1\nline 2\nline 3\n" })); cx.run_until_parked(); @@ -4634,7 +4464,7 @@ mod tests { sender.send_final(json!({ "display_description": "Create new file", "path": "root/dir/new_file.txt", - "mode": "create", + "mode": "write", "content": "line 1\nline 2\nline 3\n" })); @@ -4688,7 +4518,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite" + "mode": "write" })); cx.run_until_parked(); @@ -4706,7 +4536,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite", + "mode": "write", "content": "new line 1\n" })); cx.run_until_parked(); @@ -4720,7 +4550,7 @@ mod tests { sender.send_final(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite", + "mode": "write", "content": "new line 1\nnew line 2\n" })); @@ -4781,7 +4611,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite" + "mode": "write" })); cx.run_until_parked(); @@ -4799,7 +4629,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite", + "mode": "write", "content": "new line 1\n" })); cx.run_until_parked(); @@ -4809,7 +4639,7 @@ mod tests { sender.send_partial(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite", + "mode": "write", "content": "new line 1\nnew line 2\n" })); cx.run_until_parked(); @@ -4822,7 +4652,7 @@ mod tests { sender.send_final(json!({ "display_description": "Overwrite file", "path": "root/file.txt", - "mode": "overwrite", + "mode": "write", "content": "new line 1\nnew line 2\nnew line 3\n" })); @@ -4837,6 +4667,89 @@ mod tests { assert_eq!(*old_text, "old line 1\nold line 2\nold line 3\n"); } + #[gpui::test] + async fn test_streaming_edit_json_fixer_escape_corruption(cx: &mut TestAppContext) { + init_test(cx); + + let fs = project::FakeFs::new(cx.executor()); + fs.insert_tree( + "/root", + json!({ + "file.txt": "hello\nworld\nfoo\n" + }), + ) + .await; + let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await; + let language_registry = project.read_with(cx, |project, _cx| project.languages().clone()); + let context_server_registry = + cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx)); + let model = Arc::new(FakeLanguageModel::default()); + let thread = cx.new(|cx| { + crate::Thread::new( + project.clone(), + cx.new(|_cx| ProjectContext::default()), + context_server_registry, + Templates::new(), + Some(model), + cx, + ) + }); + + let (sender, input) = ToolInput::::test(); + let (event_stream, _receiver) = ToolCallEventStream::test(); + + let tool = Arc::new(StreamingEditFileTool::new( + project.clone(), + thread.downgrade(), + language_registry, + )); + + let task = cx.update(|cx| tool.run(input, event_stream, cx)); + + sender.send_partial(json!({ + "display_description": "Edit", + "path": "root/file.txt", + "mode": "edit" + })); + cx.run_until_parked(); + + // Simulate JSON fixer producing a literal backslash when the LLM + // stream cuts in the middle of a \n escape sequence. + // The old_text "hello\nworld" would be streamed as: + // partial 1: old_text = "hello\\" (fixer closes incomplete \n as \\) + // partial 2: old_text = "hello\nworld" (fixer corrected the escape) + sender.send_partial(json!({ + "display_description": "Edit", + "path": "root/file.txt", + "mode": "edit", + "edits": [{"old_text": "hello\\"}] + })); + cx.run_until_parked(); + + // Now the fixer corrects it to the real newline. + sender.send_partial(json!({ + "display_description": "Edit", + "path": "root/file.txt", + "mode": "edit", + "edits": [{"old_text": "hello\nworld"}] + })); + cx.run_until_parked(); + + // Send final. + sender.send_final(json!({ + "display_description": "Edit", + "path": "root/file.txt", + "mode": "edit", + "edits": [{"old_text": "hello\nworld", "new_text": "HELLO\nWORLD"}] + })); + + let result = task.await; + let StreamingEditFileToolOutput::Success { new_text, .. } = result.unwrap() else { + panic!("expected success"); + }; + assert_eq!(new_text, "HELLO\nWORLD\nfoo\n"); + } + fn init_test(cx: &mut TestAppContext) { cx.update(|cx| { let settings_store = SettingsStore::test(cx); diff --git a/crates/agent/src/tools/tool_edit_parser.rs b/crates/agent/src/tools/tool_edit_parser.rs new file mode 100644 index 0000000000000000000000000000000000000000..86259db916f49c07bbecc63625a93a9ebb955539 --- /dev/null +++ b/crates/agent/src/tools/tool_edit_parser.rs @@ -0,0 +1,941 @@ +use smallvec::SmallVec; + +use crate::{Edit, PartialEdit}; + +/// Events emitted by `ToolEditParser` as tool call input streams in. +#[derive(Debug, PartialEq, Eq)] +pub enum ToolEditEvent { + /// A chunk of `old_text` for an edit operation. + OldTextChunk { + edit_index: usize, + chunk: String, + done: bool, + }, + /// A chunk of `new_text` for an edit operation. + NewTextChunk { + edit_index: usize, + chunk: String, + done: bool, + }, + /// A chunk of content for write/overwrite mode. + ContentChunk { chunk: String }, +} + +/// Tracks the streaming state of a single edit to detect deltas. +#[derive(Default, Debug)] +struct EditStreamState { + old_text_emitted_len: usize, + old_text_done: bool, + new_text_emitted_len: usize, + new_text_done: bool, +} + +/// Converts incrementally-growing tool call JSON into a stream of chunk events. +/// +/// The tool call streaming infrastructure delivers partial JSON objects where +/// string fields grow over time. This parser compares consecutive partials, +/// computes the deltas, and emits `ToolEditEvent`s that downstream pipeline +/// stages (`StreamingFuzzyMatcher` for old_text, `StreamingDiff` for new_text) +/// can consume incrementally. +/// +/// Because partial JSON comes through a fixer (`partial-json-fixer`) that +/// closes incomplete escape sequences, a string can temporarily contain wrong +/// trailing characters (e.g. a literal `\` instead of `\n`). We handle this +/// by holding back trailing backslash characters in non-finalized chunks: if +/// a partial string ends with `\` (0x5C), that byte is not emitted until the +/// next partial confirms or corrects it. This avoids feeding corrupted bytes +/// to downstream consumers. +#[derive(Default, Debug)] +pub struct ToolEditParser { + edit_states: Vec, + content_emitted_len: usize, +} + +impl ToolEditParser { + /// Push a new set of partial edits (from edit mode) and return any events. + /// + /// Each call should pass the *entire current* edits array as seen in the + /// latest partial input. The parser will diff it against its internal state + /// to produce only the new events. + pub fn push_edits(&mut self, edits: &[PartialEdit]) -> SmallVec<[ToolEditEvent; 4]> { + let mut events = SmallVec::new(); + + for (index, partial) in edits.iter().enumerate() { + if index >= self.edit_states.len() { + // A new edit appeared — finalize the previous one if there was one. + if let Some(previous) = self.finalize_previous_edit(index) { + events.extend(previous); + } + self.edit_states.push(EditStreamState::default()); + } + + let state = &mut self.edit_states[index]; + + // Process old_text changes. + if let Some(old_text) = &partial.old_text + && !state.old_text_done + { + if partial.new_text.is_some() { + // new_text appeared, so old_text is done — emit everything. + let start = state.old_text_emitted_len.min(old_text.len()); + let chunk = old_text[start..].to_string(); + state.old_text_done = true; + state.old_text_emitted_len = old_text.len(); + events.push(ToolEditEvent::OldTextChunk { + edit_index: index, + chunk, + done: true, + }); + } else { + let safe_end = safe_emit_end(old_text); + if safe_end > state.old_text_emitted_len { + let chunk = old_text[state.old_text_emitted_len..safe_end].to_string(); + state.old_text_emitted_len = safe_end; + events.push(ToolEditEvent::OldTextChunk { + edit_index: index, + chunk, + done: false, + }); + } + } + } + + // Process new_text changes. + if let Some(new_text) = &partial.new_text + && !state.new_text_done + { + let safe_end = safe_emit_end(new_text); + if safe_end > state.new_text_emitted_len { + let chunk = new_text[state.new_text_emitted_len..safe_end].to_string(); + state.new_text_emitted_len = safe_end; + events.push(ToolEditEvent::NewTextChunk { + edit_index: index, + chunk, + done: false, + }); + } + } + } + + events + } + + /// Push new content and return any events. + /// + /// Each call should pass the *entire current* content string. The parser + /// will diff it against its internal state to emit only the new chunk. + pub fn push_content(&mut self, content: &str) -> SmallVec<[ToolEditEvent; 1]> { + let mut events = SmallVec::new(); + + let safe_end = safe_emit_end(content); + if safe_end > self.content_emitted_len { + let chunk = content[self.content_emitted_len..safe_end].to_string(); + self.content_emitted_len = safe_end; + events.push(ToolEditEvent::ContentChunk { chunk }); + } + + events + } + + /// Finalize all edits with the complete input. This emits `done: true` + /// events for any in-progress old_text or new_text that hasn't been + /// finalized yet. + /// + /// `final_edits` should be the fully deserialized final edits array. The + /// parser compares against its tracked state and emits any remaining deltas + /// with `done: true`. + pub fn finalize_edits(&mut self, edits: &[Edit]) -> SmallVec<[ToolEditEvent; 4]> { + let mut events = SmallVec::new(); + + for (index, edit) in edits.iter().enumerate() { + if index >= self.edit_states.len() { + // This edit was never seen in partials — emit it fully. + if let Some(previous) = self.finalize_previous_edit(index) { + events.extend(previous); + } + self.edit_states.push(EditStreamState::default()); + } + + let state = &mut self.edit_states[index]; + + if !state.old_text_done { + let start = state.old_text_emitted_len.min(edit.old_text.len()); + let chunk = edit.old_text[start..].to_string(); + state.old_text_done = true; + state.old_text_emitted_len = edit.old_text.len(); + events.push(ToolEditEvent::OldTextChunk { + edit_index: index, + chunk, + done: true, + }); + } + + if !state.new_text_done { + let start = state.new_text_emitted_len.min(edit.new_text.len()); + let chunk = edit.new_text[start..].to_string(); + state.new_text_done = true; + state.new_text_emitted_len = edit.new_text.len(); + events.push(ToolEditEvent::NewTextChunk { + edit_index: index, + chunk, + done: true, + }); + } + } + + events + } + + /// Finalize content with the complete input. + pub fn finalize_content(&mut self, content: &str) -> SmallVec<[ToolEditEvent; 1]> { + let mut events = SmallVec::new(); + + let start = self.content_emitted_len.min(content.len()); + if content.len() > start { + let chunk = content[start..].to_string(); + self.content_emitted_len = content.len(); + events.push(ToolEditEvent::ContentChunk { chunk }); + } + + events + } + + /// When a new edit appears at `index`, finalize the edit at `index - 1` + /// by emitting a `NewTextChunk { done: true }` if it hasn't been finalized. + fn finalize_previous_edit(&mut self, new_index: usize) -> Option> { + if new_index == 0 || self.edit_states.is_empty() { + return None; + } + + let previous_index = new_index - 1; + if previous_index >= self.edit_states.len() { + return None; + } + + let state = &mut self.edit_states[previous_index]; + let mut events = SmallVec::new(); + + // If old_text was never finalized, finalize it now with an empty done chunk. + if !state.old_text_done { + state.old_text_done = true; + events.push(ToolEditEvent::OldTextChunk { + edit_index: previous_index, + chunk: String::new(), + done: true, + }); + } + + // Emit a done event for new_text if not already finalized. + if !state.new_text_done { + state.new_text_done = true; + events.push(ToolEditEvent::NewTextChunk { + edit_index: previous_index, + chunk: String::new(), + done: true, + }); + } + + Some(events) + } +} + +/// Returns the byte position up to which it is safe to emit from a partial +/// string. If the string ends with a backslash (`\`, 0x5C), that byte is +/// held back because it may be an artifact of the partial JSON fixer closing +/// an incomplete escape sequence (e.g. turning a half-received `\n` into `\\`). +/// The next partial will reveal the correct character. +fn safe_emit_end(text: &str) -> usize { + if text.as_bytes().last() == Some(&b'\\') { + text.len() - 1 + } else { + text.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_single_edit_streamed_incrementally() { + let mut parser = ToolEditParser::default(); + + // old_text arrives in chunks: "hell" → "hello w" → "hello world" + let events = parser.push_edits(&[PartialEdit { + old_text: Some("hell".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "hell".into(), + done: false, + }] + ); + + let events = parser.push_edits(&[PartialEdit { + old_text: Some("hello w".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "o w".into(), + done: false, + }] + ); + + // new_text appears → old_text finalizes + let events = parser.push_edits(&[PartialEdit { + old_text: Some("hello world".into()), + new_text: Some("good".into()), + }]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "orld".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "good".into(), + done: false, + }, + ] + ); + + // new_text grows + let events = parser.push_edits(&[PartialEdit { + old_text: Some("hello world".into()), + new_text: Some("goodbye world".into()), + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "bye world".into(), + done: false, + }] + ); + + // Finalize + let events = parser.finalize_edits(&[Edit { + old_text: "hello world".into(), + new_text: "goodbye world".into(), + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "".into(), + done: true, + }] + ); + } + + #[test] + fn test_multiple_edits_sequential() { + let mut parser = ToolEditParser::default(); + + // First edit streams in + let events = parser.push_edits(&[PartialEdit { + old_text: Some("first old".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "first old".into(), + done: false, + }] + ); + + let events = parser.push_edits(&[PartialEdit { + old_text: Some("first old".into()), + new_text: Some("first new".into()), + }]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "first new".into(), + done: false, + }, + ] + ); + + // Second edit appears → first edit's new_text is finalized + let events = parser.push_edits(&[ + PartialEdit { + old_text: Some("first old".into()), + new_text: Some("first new".into()), + }, + PartialEdit { + old_text: Some("second".into()), + new_text: None, + }, + ]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "".into(), + done: true, + }, + ToolEditEvent::OldTextChunk { + edit_index: 1, + chunk: "second".into(), + done: false, + }, + ] + ); + + // Finalize everything + let events = parser.finalize_edits(&[ + Edit { + old_text: "first old".into(), + new_text: "first new".into(), + }, + Edit { + old_text: "second old".into(), + new_text: "second new".into(), + }, + ]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 1, + chunk: " old".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 1, + chunk: "second new".into(), + done: true, + }, + ] + ); + } + + #[test] + fn test_content_streamed_incrementally() { + let mut parser = ToolEditParser::default(); + + let events = parser.push_content("hello"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { + chunk: "hello".into(), + }] + ); + + let events = parser.push_content("hello world"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { + chunk: " world".into(), + }] + ); + + // No change + let events = parser.push_content("hello world"); + assert!(events.is_empty()); + + let events = parser.push_content("hello world!"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { chunk: "!".into() }] + ); + + // Finalize with no additional content + let events = parser.finalize_content("hello world!"); + assert!(events.is_empty()); + } + + #[test] + fn test_finalize_content_with_remaining() { + let mut parser = ToolEditParser::default(); + + parser.push_content("partial"); + let events = parser.finalize_content("partial content here"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { + chunk: " content here".into(), + }] + ); + } + + #[test] + fn test_content_trailing_backslash_held_back() { + let mut parser = ToolEditParser::default(); + + // Partial JSON fixer turns incomplete \n into \\ (literal backslash). + // The trailing backslash is held back. + let events = parser.push_content("hello,\\"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { + chunk: "hello,".into(), + }] + ); + + // Next partial corrects the escape to an actual newline. + // The held-back byte was wrong; the correct newline is emitted. + let events = parser.push_content("hello,\n"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { chunk: "\n".into() }] + ); + + // Normal growth. + let events = parser.push_content("hello,\nworld"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { + chunk: "world".into(), + }] + ); + } + + #[test] + fn test_content_finalize_with_trailing_backslash() { + let mut parser = ToolEditParser::default(); + + // Stream a partial with a fixer-corrupted trailing backslash. + // The backslash is held back. + parser.push_content("abc\\"); + + // Finalize reveals the correct character. + let events = parser.finalize_content("abc\n"); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::ContentChunk { chunk: "\n".into() }] + ); + } + + #[test] + fn test_no_partials_direct_finalize() { + let mut parser = ToolEditParser::default(); + + let events = parser.finalize_edits(&[Edit { + old_text: "old".into(), + new_text: "new".into(), + }]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "old".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "new".into(), + done: true, + }, + ] + ); + } + + #[test] + fn test_no_partials_direct_finalize_multiple() { + let mut parser = ToolEditParser::default(); + + let events = parser.finalize_edits(&[ + Edit { + old_text: "first old".into(), + new_text: "first new".into(), + }, + Edit { + old_text: "second old".into(), + new_text: "second new".into(), + }, + ]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "first old".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "first new".into(), + done: true, + }, + ToolEditEvent::OldTextChunk { + edit_index: 1, + chunk: "second old".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 1, + chunk: "second new".into(), + done: true, + }, + ] + ); + } + + #[test] + fn test_old_text_no_growth() { + let mut parser = ToolEditParser::default(); + + let events = parser.push_edits(&[PartialEdit { + old_text: Some("same".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "same".into(), + done: false, + }] + ); + + // Same old_text, no new_text → no events + let events = parser.push_edits(&[PartialEdit { + old_text: Some("same".into()), + new_text: None, + }]); + assert!(events.is_empty()); + } + + #[test] + fn test_old_text_none_then_appears() { + let mut parser = ToolEditParser::default(); + + // Edit exists but old_text is None (field hasn't arrived yet) + let events = parser.push_edits(&[PartialEdit { + old_text: None, + new_text: None, + }]); + assert!(events.is_empty()); + + // old_text appears + let events = parser.push_edits(&[PartialEdit { + old_text: Some("text".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "text".into(), + done: false, + }] + ); + } + + #[test] + fn test_empty_old_text_with_new_text() { + let mut parser = ToolEditParser::default(); + + // old_text is empty, new_text appears immediately + let events = parser.push_edits(&[PartialEdit { + old_text: Some("".into()), + new_text: Some("inserted".into()), + }]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "inserted".into(), + done: false, + }, + ] + ); + } + + #[test] + fn test_three_edits_streamed() { + let mut parser = ToolEditParser::default(); + + // Stream first edit + parser.push_edits(&[PartialEdit { + old_text: Some("a".into()), + new_text: Some("A".into()), + }]); + + // Second edit appears + parser.push_edits(&[ + PartialEdit { + old_text: Some("a".into()), + new_text: Some("A".into()), + }, + PartialEdit { + old_text: Some("b".into()), + new_text: Some("B".into()), + }, + ]); + + // Third edit appears + let events = parser.push_edits(&[ + PartialEdit { + old_text: Some("a".into()), + new_text: Some("A".into()), + }, + PartialEdit { + old_text: Some("b".into()), + new_text: Some("B".into()), + }, + PartialEdit { + old_text: Some("c".into()), + new_text: None, + }, + ]); + + // Should finalize edit 1 (index=1) and start edit 2 (index=2) + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::NewTextChunk { + edit_index: 1, + chunk: "".into(), + done: true, + }, + ToolEditEvent::OldTextChunk { + edit_index: 2, + chunk: "c".into(), + done: false, + }, + ] + ); + + // Finalize + let events = parser.finalize_edits(&[ + Edit { + old_text: "a".into(), + new_text: "A".into(), + }, + Edit { + old_text: "b".into(), + new_text: "B".into(), + }, + Edit { + old_text: "c".into(), + new_text: "C".into(), + }, + ]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 2, + chunk: "".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 2, + chunk: "C".into(), + done: true, + }, + ] + ); + } + + #[test] + fn test_finalize_with_unseen_old_text() { + let mut parser = ToolEditParser::default(); + + // Only saw partial old_text, never saw new_text in partials + parser.push_edits(&[PartialEdit { + old_text: Some("partial".into()), + new_text: None, + }]); + + let events = parser.finalize_edits(&[Edit { + old_text: "partial old text".into(), + new_text: "replacement".into(), + }]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: " old text".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "replacement".into(), + done: true, + }, + ] + ); + } + + #[test] + fn test_finalize_with_partially_seen_new_text() { + let mut parser = ToolEditParser::default(); + + parser.push_edits(&[PartialEdit { + old_text: Some("old".into()), + new_text: Some("partial".into()), + }]); + + let events = parser.finalize_edits(&[Edit { + old_text: "old".into(), + new_text: "partial new text".into(), + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: " new text".into(), + done: true, + }] + ); + } + + #[test] + fn test_repeated_pushes_with_no_change() { + let mut parser = ToolEditParser::default(); + + let events = parser.push_edits(&[PartialEdit { + old_text: Some("stable".into()), + new_text: Some("also stable".into()), + }]); + assert_eq!(events.len(), 2); // old done + new chunk + + // Push the exact same data again + let events = parser.push_edits(&[PartialEdit { + old_text: Some("stable".into()), + new_text: Some("also stable".into()), + }]); + assert!(events.is_empty()); + + // And again + let events = parser.push_edits(&[PartialEdit { + old_text: Some("stable".into()), + new_text: Some("also stable".into()), + }]); + assert!(events.is_empty()); + } + + #[test] + fn test_old_text_trailing_backslash_held_back() { + let mut parser = ToolEditParser::default(); + + // Partial-json-fixer produces a literal backslash when the JSON stream + // cuts in the middle of an escape sequence like \n. The parser holds + // back the trailing backslash instead of emitting it. + let events = parser.push_edits(&[PartialEdit { + old_text: Some("hello,\\".into()), // fixer closed incomplete \n as \\ + new_text: None, + }]); + // The trailing `\` is held back — only "hello," is emitted. + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "hello,".into(), + done: false, + }] + ); + + // Next partial: the fixer corrects the escape to \n. + // The held-back byte was wrong, but we never emitted it. Now the + // correct newline at that position is emitted normally. + let events = parser.push_edits(&[PartialEdit { + old_text: Some("hello,\n".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "\n".into(), + done: false, + }] + ); + + // Continue normally. + let events = parser.push_edits(&[PartialEdit { + old_text: Some("hello,\nworld".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "world".into(), + done: false, + }] + ); + } + + #[test] + fn test_multiline_old_and_new_text() { + let mut parser = ToolEditParser::default(); + + let events = parser.push_edits(&[PartialEdit { + old_text: Some("line1\nline2".into()), + new_text: None, + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "line1\nline2".into(), + done: false, + }] + ); + + let events = parser.push_edits(&[PartialEdit { + old_text: Some("line1\nline2\nline3".into()), + new_text: Some("LINE1\n".into()), + }]); + assert_eq!( + events.as_slice(), + &[ + ToolEditEvent::OldTextChunk { + edit_index: 0, + chunk: "\nline3".into(), + done: true, + }, + ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "LINE1\n".into(), + done: false, + }, + ] + ); + + let events = parser.push_edits(&[PartialEdit { + old_text: Some("line1\nline2\nline3".into()), + new_text: Some("LINE1\nLINE2\nLINE3".into()), + }]); + assert_eq!( + events.as_slice(), + &[ToolEditEvent::NewTextChunk { + edit_index: 0, + chunk: "LINE2\nLINE3".into(), + done: false, + }] + ); + } +}