WIP: Start on a test of AssistantDiff

Co-authored-by: Nathan Sobo <nathan@zed.dev>
Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>

Cargo.lock

@@ -693,17 +693,22 @@ name = "assistant_tool"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "async-watch",
  "buffer_diff",
  "clock",
  "collections",
+ "ctor",
  "derive_more",
+ "env_logger 0.11.7",
+ "futures 0.3.31",
  "gpui",
  "icons",
  "language",
  "language_model",
+ "log",
  "parking_lot",
+ "pretty_assertions",
  "project",
+ "rand 0.8.5",
  "serde",
  "serde_json",
  "settings",
@@ -718,7 +723,6 @@ dependencies = [
  "anyhow",
  "assistant_tool",
  "chrono",
- "clock",
  "collections",
  "feature_flags",
  "futures 0.3.31",

assets/keymaps/default-linux.json

@@ -149,7 +149,7 @@
   {
     "context": "AssistantDiff",
     "bindings": {
-      "ctrl-y": "agent::ToggleKeep",
+      "ctrl-y": "agent::Keep",
       "ctrl-k ctrl-r": "agent::Reject"
     }
   },

assets/keymaps/default-macos.json

@@ -241,7 +241,7 @@
     "context": "AssistantDiff",
     "use_key_equivalents": true,
     "bindings": {
-      "cmd-y": "agent::ToggleKeep",
+      "cmd-y": "agent::Keep",
       "cmd-alt-z": "agent::Reject"
     }
   },

crates/assistant2/src/assistant.rs

@@ -66,7 +66,7 @@ actions!(
         AcceptSuggestedContext,
         OpenActiveThreadAsMarkdown,
         OpenAssistantDiff,
-        ToggleKeep,
+        Keep,
         Reject,
         RejectAll,
         KeepAll

crates/assistant2/src/assistant_diff.rs

@@ -1,10 +1,11 @@
-use crate::{Thread, ThreadEvent, ToggleKeep};
+use crate::{Thread, ThreadEvent};
 use anyhow::Result;
 use buffer_diff::DiffHunkStatus;
 use collections::HashSet;
 use editor::{
     Direction, Editor, EditorEvent, MultiBuffer, ToPoint,
     actions::{GoToHunk, GoToPreviousHunk},
+    scroll::Autoscroll,
 };
 use gpui::{
     Action, AnyElement, AnyView, App, Entity, EventEmitter, FocusHandle, Focusable, SharedString,
@@ -78,7 +79,7 @@ impl AssistantDiff {
                   hunk_range,
                   is_created_file,
                   line_height,
-                  _editor: &Entity<Editor>,
+                  editor: &Entity<Editor>,
                   window: &mut Window,
                   cx: &mut App| {
                 render_diff_hunk_controls(
@@ -88,6 +89,7 @@ impl AssistantDiff {
                     is_created_file,
                     line_height,
                     &assistant_diff,
+                    editor,
                     window,
                     cx,
                 )
@@ -130,7 +132,7 @@ impl AssistantDiff {
         let changed_buffers = thread.action_log().read(cx).changed_buffers(cx);
         let mut paths_to_delete = self.multibuffer.read(cx).paths().collect::<HashSet<_>>();
 
-        for (buffer, changed) in changed_buffers {
+        for (buffer, diff_handle) in changed_buffers {
             let Some(file) = buffer.read(cx).file().cloned() else {
                 continue;
             };
@@ -139,7 +141,8 @@ impl AssistantDiff {
             paths_to_delete.remove(&path_key);
 
             let snapshot = buffer.read(cx).snapshot();
-            let diff = changed.diff.read(cx);
+            let diff = diff_handle.read(cx);
+
             let diff_hunk_ranges = diff
                 .hunks_intersecting_range(
                     language::Anchor::MIN..language::Anchor::MAX,
@@ -159,15 +162,31 @@ impl AssistantDiff {
                         editor::DEFAULT_MULTIBUFFER_CONTEXT,
                         cx,
                     );
-                    multibuffer.add_diff(changed.diff.clone(), cx);
+                    multibuffer.add_diff(diff_handle, cx);
                     (was_empty, is_excerpt_newly_added)
                 });
 
             self.editor.update(cx, |editor, cx| {
                 if was_empty {
-                    editor.change_selections(None, window, cx, |selections| {
-                        selections.select_ranges([0..0])
-                    });
+                    let first_hunk = editor
+                        .diff_hunks_in_ranges(
+                            &[editor::Anchor::min()..editor::Anchor::max()],
+                            &self.multibuffer.read(cx).read(cx),
+                        )
+                        .next();
+
+                    if let Some(first_hunk) = first_hunk {
+                        let first_hunk_start = first_hunk.multi_buffer_range().start;
+                        editor.change_selections(
+                            Some(Autoscroll::fit()),
+                            window,
+                            cx,
+                            |selections| {
+                                selections
+                                    .select_anchor_ranges([first_hunk_start..first_hunk_start]);
+                            },
+                        )
+                    }
                 }
 
                 if is_excerpt_newly_added
@@ -221,7 +240,7 @@ impl AssistantDiff {
         }
     }
 
-    fn toggle_keep(&mut self, _: &crate::ToggleKeep, _window: &mut Window, cx: &mut Context<Self>) {
+    fn keep(&mut self, _: &crate::Keep, window: &mut Window, cx: &mut Context<Self>) {
         let ranges = self
             .editor
             .read(cx)
@@ -236,15 +255,40 @@ impl AssistantDiff {
             .diff_hunks_in_ranges(&ranges, &snapshot)
             .collect::<Vec<_>>();
 
-        for hunk in diff_hunks_in_ranges {
+        for hunk in &diff_hunks_in_ranges {
             let buffer = self.multibuffer.read(cx).buffer(hunk.buffer_id);
             if let Some(buffer) = buffer {
                 self.thread.update(cx, |thread, cx| {
-                    let accept = hunk.status().has_secondary_hunk();
-                    thread.review_edits_in_range(buffer, hunk.buffer_range, accept, cx)
+                    thread.keep_edits_in_range(buffer, hunk.buffer_range.clone(), cx)
                 });
             }
         }
+
+        if let Some(last_kept_hunk) = diff_hunks_in_ranges.last() {
+            let last_kept_hunk_end = editor::Anchor::in_buffer(
+                last_kept_hunk.excerpt_id,
+                last_kept_hunk.buffer_id,
+                last_kept_hunk.buffer_range.start,
+            );
+            let next_hunk = self
+                .editor
+                .read(cx)
+                .diff_hunks_in_ranges(&[last_kept_hunk_end..editor::Anchor::max()], &snapshot)
+                .skip(1)
+                .next();
+            if let Some(next_hunk) = next_hunk {
+                self.editor.update(cx, |editor, cx| {
+                    editor.change_selections(Some(Autoscroll::fit()), window, cx, |selections| {
+                        let next_hunk_start = editor::Anchor::in_buffer(
+                            next_hunk.excerpt_id,
+                            next_hunk.buffer_id,
+                            next_hunk.buffer_range.start,
+                        );
+                        selections.select_anchor_ranges([next_hunk_start..next_hunk_start]);
+                    })
+                })
+            }
+        }
     }
 
     fn reject(&mut self, _: &crate::Reject, window: &mut Window, cx: &mut Context<Self>) {
@@ -268,10 +312,9 @@ impl AssistantDiff {
             .update(cx, |thread, cx| thread.keep_all_edits(cx));
     }
 
-    fn review_diff_hunks(
+    fn keep_edits_in_ranges(
         &mut self,
         hunk_ranges: Vec<Range<editor::Anchor>>,
-        accept: bool,
         cx: &mut Context<Self>,
     ) {
         let snapshot = self.multibuffer.read(cx).snapshot(cx);
@@ -285,7 +328,7 @@ impl AssistantDiff {
             let buffer = self.multibuffer.read(cx).buffer(hunk.buffer_id);
             if let Some(buffer) = buffer {
                 self.thread.update(cx, |thread, cx| {
-                    thread.review_edits_in_range(buffer, hunk.buffer_range, accept, cx)
+                    thread.keep_edits_in_range(buffer, hunk.buffer_range, cx)
                 });
             }
         }
@@ -479,7 +522,7 @@ impl Render for AssistantDiff {
             } else {
                 "AssistantDiff"
             })
-            .on_action(cx.listener(Self::toggle_keep))
+            .on_action(cx.listener(Self::keep))
             .on_action(cx.listener(Self::reject))
             .on_action(cx.listener(Self::reject_all))
             .on_action(cx.listener(Self::keep_all))
@@ -495,16 +538,16 @@ impl Render for AssistantDiff {
 
 fn render_diff_hunk_controls(
     row: u32,
-    status: &DiffHunkStatus,
+    _status: &DiffHunkStatus,
     hunk_range: Range<editor::Anchor>,
     is_created_file: bool,
     line_height: Pixels,
     assistant_diff: &Entity<AssistantDiff>,
+    editor: &Entity<Editor>,
     window: &mut Window,
     cx: &mut App,
 ) -> AnyElement {
-    let editor = assistant_diff.read(cx).editor.clone();
-
+    let editor = editor.clone();
     h_flex()
         .h(line_height)
         .mr_0p5()
@@ -519,75 +562,47 @@ fn render_diff_hunk_controls(
         .gap_1()
         .occlude()
         .shadow_md()
-        .children(if status.has_secondary_hunk() {
-            vec![
-                Button::new("reject", "Reject")
-                    .disabled(is_created_file)
-                    .key_binding(
-                        KeyBinding::for_action_in(
-                            &crate::Reject,
-                            &editor.read(cx).focus_handle(cx),
-                            window,
-                            cx,
-                        )
-                        .map(|kb| kb.size(rems_from_px(12.))),
-                    )
-                    .on_click({
-                        let editor = editor.clone();
-                        move |_event, window, cx| {
-                            editor.update(cx, |editor, cx| {
-                                let snapshot = editor.snapshot(window, cx);
-                                let point = hunk_range.start.to_point(&snapshot.buffer_snapshot);
-                                editor.restore_hunks_in_ranges(vec![point..point], window, cx);
-                            });
-                        }
-                    }),
-                Button::new(("keep", row as u64), "Keep")
-                    .key_binding(
-                        KeyBinding::for_action_in(
-                            &crate::ToggleKeep,
-                            &editor.read(cx).focus_handle(cx),
-                            window,
-                            cx,
-                        )
-                        .map(|kb| kb.size(rems_from_px(12.))),
-                    )
-                    .on_click({
-                        let assistant_diff = assistant_diff.clone();
-                        move |_event, _window, cx| {
-                            assistant_diff.update(cx, |diff, cx| {
-                                diff.review_diff_hunks(
-                                    vec![hunk_range.start..hunk_range.start],
-                                    true,
-                                    cx,
-                                );
-                            });
-                        }
-                    }),
-            ]
-        } else {
-            vec![
-                Button::new(("review", row as u64), "Review")
-                    .key_binding(KeyBinding::for_action_in(
-                        &ToggleKeep,
+        .children(vec![
+            Button::new("reject", "Reject")
+                .disabled(is_created_file)
+                .key_binding(
+                    KeyBinding::for_action_in(
+                        &crate::Reject,
                         &editor.read(cx).focus_handle(cx),
                         window,
                         cx,
-                    ))
-                    .on_click({
-                        let assistant_diff = assistant_diff.clone();
-                        move |_event, _window, cx| {
-                            assistant_diff.update(cx, |diff, cx| {
-                                diff.review_diff_hunks(
-                                    vec![hunk_range.start..hunk_range.start],
-                                    false,
-                                    cx,
-                                );
-                            });
-                        }
-                    }),
-            ]
-        })
+                    )
+                    .map(|kb| kb.size(rems_from_px(12.))),
+                )
+                .on_click({
+                    let editor = editor.clone();
+                    move |_event, window, cx| {
+                        editor.update(cx, |editor, cx| {
+                            let snapshot = editor.snapshot(window, cx);
+                            let point = hunk_range.start.to_point(&snapshot.buffer_snapshot);
+                            editor.restore_hunks_in_ranges(vec![point..point], window, cx);
+                        });
+                    }
+                }),
+            Button::new(("keep", row as u64), "Keep")
+                .key_binding(
+                    KeyBinding::for_action_in(
+                        &crate::Keep,
+                        &editor.read(cx).focus_handle(cx),
+                        window,
+                        cx,
+                    )
+                    .map(|kb| kb.size(rems_from_px(12.))),
+                )
+                .on_click({
+                    let assistant_diff = assistant_diff.clone();
+                    move |_event, _window, cx| {
+                        assistant_diff.update(cx, |diff, cx| {
+                            diff.keep_edits_in_ranges(vec![hunk_range.start..hunk_range.start], cx);
+                        });
+                    }
+                }),
+        ])
         .when(
             !editor.read(cx).buffer().read(cx).all_diff_hunks_expanded(),
             |el| {
@@ -770,3 +785,93 @@ impl Render for AssistantDiffToolbar {
             )
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{ThreadStore, thread_store};
+    use assistant_settings::AssistantSettings;
+    use context_server::ContextServerSettings;
+    use editor::EditorSettings;
+    use gpui::TestAppContext;
+    use project::{FakeFs, Project};
+    use prompt_store::PromptBuilder;
+    use serde_json::json;
+    use settings::{Settings, SettingsStore};
+    use std::{path::Path, sync::Arc};
+    use theme::ThemeSettings;
+    use util::path;
+
+    #[gpui::test]
+    async fn test_assistant_diff(cx: &mut TestAppContext) {
+        cx.update(|cx| {
+            let settings_store = SettingsStore::test(cx);
+            cx.set_global(settings_store);
+            language::init(cx);
+            Project::init_settings(cx);
+            AssistantSettings::register(cx);
+            thread_store::init(cx);
+            workspace::init_settings(cx);
+            ThemeSettings::register(cx);
+            ContextServerSettings::register(cx);
+            EditorSettings::register(cx);
+        });
+
+        let fs = FakeFs::new(cx.executor());
+        fs.insert_tree(path!("/test"), json!({"file1": "abc\ndef\nghi\njkl\nmno"}))
+            .await;
+        let project = Project::test(fs, [Path::new("/test")], cx).await;
+        let buffer_path = project
+            .read_with(cx, |project, cx| {
+                project.find_project_path("test/file1", cx)
+            })
+            .unwrap();
+
+        let thread_store = cx.update(|cx| {
+            ThreadStore::new(
+                project.clone(),
+                Arc::default(),
+                Arc::new(PromptBuilder::new(None).unwrap()),
+                cx,
+            )
+            .unwrap()
+        });
+        let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
+        let action_log = thread.read_with(cx, |thread, _| thread.action_log().clone());
+
+        let (workspace, cx) =
+            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
+        let assistant_diff = cx.new_window_entity(|window, cx| {
+            AssistantDiff::new(thread.clone(), workspace.downgrade(), window, cx)
+        });
+        let editor = assistant_diff.read_with(cx, |diff, _cx| diff.editor.clone());
+
+        let buffer = project
+            .update(cx, |project, cx| project.open_buffer(buffer_path, cx))
+            .await
+            .unwrap();
+        cx.update(|_, cx| {
+            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+            buffer.update(cx, |buffer, cx| {
+                buffer
+                    .edit([(Point::new(1, 1)..Point::new(1, 2), "E")], None, cx)
+                    .unwrap()
+            });
+            buffer.update(cx, |buffer, cx| {
+                buffer
+                    .edit([(Point::new(4, 2)..Point::new(4, 3), "O")], None, cx)
+                    .unwrap()
+            });
+            action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
+        });
+        cx.run_until_parked();
+
+        // When opening the assistant diff, the cursor is positioned on the first hunk.
+        assert_eq!(
+            editor
+                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
+                .range(),
+            Point::new(1, 0)..Point::new(1, 0)
+        );
+    }
+}

crates/assistant2/src/message_editor.rs

@@ -245,9 +245,6 @@ impl MessageEditor {
             thread
                 .update(cx, |thread, cx| {
                     let context = context_store.read(cx).context().clone();
-                    thread.action_log().update(cx, |action_log, cx| {
-                        action_log.clear_reviewed_changes(cx);
-                    });
                     thread.insert_user_message(user_message, context, checkpoint, cx);
                 })
                 .ok();
@@ -546,7 +543,7 @@ impl Render for MessageEditor {
                             parent.child(
                                 v_flex().bg(cx.theme().colors().editor_background).children(
                                     changed_buffers.into_iter().enumerate().flat_map(
-                                        |(index, (buffer, changed))| {
+                                        |(index, (buffer, _diff))| {
                                             let file = buffer.read(cx).file()?;
                                             let path = file.path();
 
@@ -619,25 +616,13 @@ impl Render for MessageEditor {
                                                                         .color(Color::Deleted),
                                                                 ),
                                                         )
-                                                        .when(!changed.needs_review, |parent| {
-                                                            parent.child(
-                                                                Icon::new(IconName::Check)
-                                                                    .color(Color::Success),
-                                                            )
-                                                        })
                                                         .child(
                                                             div()
                                                                 .h_full()
                                                                 .absolute()
                                                                 .w_8()
                                                                 .bottom_0()
-                                                                .map(|this| {
-                                                                    if !changed.needs_review {
-                                                                        this.right_4()
-                                                                    } else {
-                                                                        this.right_0()
-                                                                    }
-                                                                })
+                                                                .right_0()
                                                                 .bg(linear_gradient(
                                                                     90.,
                                                                     linear_color_stop(

crates/assistant2/src/thread.rs

@@ -1679,23 +1679,20 @@ impl Thread {
         Ok(String::from_utf8_lossy(&markdown).to_string())
     }
 
-    pub fn review_edits_in_range(
+    pub fn keep_edits_in_range(
         &mut self,
         buffer: Entity<language::Buffer>,
         buffer_range: Range<language::Anchor>,
-        accept: bool,
         cx: &mut Context<Self>,
     ) {
         self.action_log.update(cx, |action_log, cx| {
-            action_log.review_edits_in_range(buffer, buffer_range, accept, cx)
+            action_log.keep_edits_in_range(buffer, buffer_range, cx)
         });
     }
 
-    /// Keeps all edits across all buffers at once.
-    /// This provides a more performant alternative to calling review_edits_in_range for each buffer.
     pub fn keep_all_edits(&mut self, cx: &mut Context<Self>) {
         self.action_log
-            .update(cx, |action_log, _cx| action_log.keep_all_edits());
+            .update(cx, |action_log, cx| action_log.keep_all_edits(cx));
     }
 
     pub fn action_log(&self) -> &Entity<ActionLog> {

crates/assistant_tool/Cargo.toml

@@ -13,11 +13,11 @@ path = "src/assistant_tool.rs"
 
 [dependencies]
 anyhow.workspace = true
-async-watch.workspace = true
 buffer_diff.workspace = true
 clock.workspace = true
 collections.workspace = true
 derive_more.workspace = true
+futures.workspace = true
 gpui.workspace = true
 icons.workspace = true
 language.workspace = true
@@ -27,15 +27,21 @@ project.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 text.workspace = true
+util.workspace = true
 
 [dev-dependencies]
 buffer_diff = { workspace = true, features = ["test-support"] }
 collections = { workspace = true, features = ["test-support"] }
 clock = { workspace = true, features = ["test-support"] }
+ctor.workspace = true
+env_logger.workspace = true
 gpui = { workspace = true, features = ["test-support"] }
 language = { workspace = true, features = ["test-support"] }
 language_model = { workspace = true, features = ["test-support"] }
+log.workspace = true
+pretty_assertions.workspace = true
 project = { workspace = true, features = ["test-support"] }
+rand.workspace = true
 settings = { workspace = true, features = ["test-support"] }
 text = { workspace = true, features = ["test-support"] }
 util = { workspace = true, features = ["test-support"] }

crates/assistant_tools/Cargo.toml

@@ -14,7 +14,6 @@ path = "src/assistant_tools.rs"
 [dependencies]
 anyhow.workspace = true
 assistant_tool.workspace = true
-clock.workspace = true
 chrono.workspace = true
 collections.workspace = true
 feature_flags.workspace = true
@@ -42,7 +41,6 @@ worktree.workspace = true
 open = { workspace = true }
 
 [dev-dependencies]
-clock = { workspace = true, features = ["test-support"] }
 collections = { workspace = true, features = ["test-support"] }
 gpui = { workspace = true, features = ["test-support"] }
 language = { workspace = true, features = ["test-support"] }

crates/assistant_tools/src/create_file_tool.rs

@@ -92,10 +92,11 @@ impl Tool for CreateFileTool {
                 })?
                 .await
                 .map_err(|err| anyhow!("Unable to open buffer for {destination_path}: {err}"))?;
-            let edit_id = buffer.update(cx, |buffer, cx| buffer.set_text(contents, cx))?;
-
-            action_log.update(cx, |action_log, cx| {
-                action_log.will_create_buffer(buffer.clone(), edit_id, cx)
+            cx.update(|cx| {
+                buffer.update(cx, |buffer, cx| buffer.set_text(contents, cx));
+                action_log.update(cx, |action_log, cx| {
+                    action_log.will_create_buffer(buffer.clone(), cx)
+                });
             })?;
 
             project

crates/assistant_tools/src/edit_files_tool.rs

@@ -174,7 +174,6 @@ enum EditorResponse {
 struct AppliedAction {
     source: String,
     buffer: Entity<language::Buffer>,
-    edit_ids: Vec<clock::Lamport>,
 }
 
 #[derive(Debug)]
@@ -339,18 +338,17 @@ impl EditToolRequest {
                 self.push_search_error(error);
             }
             DiffResult::Diff(diff) => {
-                let edit_ids = buffer.update(cx, |buffer, cx| {
-                    buffer.finalize_last_transaction();
-                    buffer.apply_diff(diff, false, cx);
-                    let transaction = buffer.finalize_last_transaction();
-                    transaction.map_or(Vec::new(), |transaction| transaction.edit_ids.clone())
+                cx.update(|cx| {
+                    buffer.update(cx, |buffer, cx| {
+                        buffer.finalize_last_transaction();
+                        buffer.apply_diff(diff, cx);
+                        buffer.finalize_last_transaction();
+                    });
+                    self.action_log
+                        .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
                 })?;
 
-                self.push_applied_action(AppliedAction {
-                    source,
-                    buffer,
-                    edit_ids,
-                });
+                self.push_applied_action(AppliedAction { source, buffer });
             }
         }
 
@@ -473,9 +471,6 @@ impl EditToolRequest {
 
                 for action in applied {
                     changed_buffers.insert(action.buffer.clone());
-                    self.action_log.update(cx, |log, cx| {
-                        log.buffer_edited(action.buffer, action.edit_ids, cx)
-                    })?;
                     write!(&mut output, "\n\n{}", action.source)?;
                 }
 

crates/assistant_tools/src/find_replace_file_tool.rs

@@ -1,7 +1,7 @@
 use crate::{replace::replace_with_flexible_indent, schema::json_schema_for};
 use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool};
-use gpui::{App, AppContext, Entity, Task};
+use gpui::{App, AppContext, AsyncApp, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -165,7 +165,7 @@ impl Tool for FindReplaceFileTool {
             Err(err) => return Task::ready(Err(anyhow!(err))),
         };
 
-        cx.spawn(async move |cx| {
+        cx.spawn(async move |cx: &mut AsyncApp| {
             let project_path = project.read_with(cx, |project, cx| {
                 project
                     .find_project_path(&input.path, cx)
@@ -225,20 +225,18 @@ impl Tool for FindReplaceFileTool {
                 return Err(err)
             };
 
-            let (edit_ids, snapshot) = buffer.update(cx, |buffer, cx| {
+            let snapshot = buffer.update(cx, |buffer, cx| {
                 buffer.finalize_last_transaction();
-                buffer.apply_diff(diff, false, cx);
-                let transaction = buffer.finalize_last_transaction();
-                let edit_ids = transaction.map_or(Vec::new(), |transaction| transaction.edit_ids.clone());
-
-                (edit_ids, buffer.snapshot())
+                buffer.apply_diff(diff, cx);
+                buffer.finalize_last_transaction();
+                buffer.snapshot()
             })?;
 
             action_log.update(cx, |log, cx| {
-                log.buffer_edited(buffer.clone(), edit_ids, cx)
+                log.buffer_edited(buffer.clone(), cx)
             })?;
 
-            project.update(cx, |project, cx| {
+            project.update( cx, |project, cx| {
                 project.save_buffer(buffer, cx)
             })?.await?;
 
@@ -249,6 +247,7 @@ impl Tool for FindReplaceFileTool {
 
 
             Ok(format!("Edited {}:\n\n```diff\n{}\n```", input.path.display(), diff_str))
+
         })
     }
 }

crates/assistant_tools/src/replace.rs

@@ -518,7 +518,7 @@ mod tests {
         // Call replace_flexible and transform the result
         replace_with_flexible_indent(old, new, &buffer_snapshot).map(|diff| {
             buffer.update(cx, |buffer, cx| {
-                let _ = buffer.apply_diff(diff, false, cx);
+                let _ = buffer.apply_diff(diff, cx);
                 buffer.text()
             })
         })

crates/git/src/repository.rs

@@ -377,7 +377,6 @@ impl RealGitRepository {
 #[derive(Clone, Debug)]
 pub struct GitRepositoryCheckpoint {
     ref_name: String,
-    head_sha: Option<Oid>,
     commit_sha: Oid,
 }
 
@@ -1213,11 +1212,6 @@ impl GitRepository for RealGitRepository {
 
                     Ok(GitRepositoryCheckpoint {
                         ref_name,
-                        head_sha: if let Some(head_sha) = head_sha {
-                            Some(head_sha.parse()?)
-                        } else {
-                            None
-                        },
                         commit_sha: checkpoint_sha.parse()?,
                     })
                 })
@@ -1252,13 +1246,6 @@ impl GitRepository for RealGitRepository {
                 })
                 .await?;
 
-                if let Some(head_sha) = checkpoint.head_sha {
-                    git.run(&["reset", "--mixed", &head_sha.to_string()])
-                        .await?;
-                } else {
-                    git.run(&["update-ref", "-d", "HEAD"]).await?;
-                }
-
                 Ok(())
             })
             .boxed()
@@ -1269,10 +1256,6 @@ impl GitRepository for RealGitRepository {
         left: GitRepositoryCheckpoint,
         right: GitRepositoryCheckpoint,
     ) -> BoxFuture<Result<bool>> {
-        if left.head_sha != right.head_sha {
-            return future::ready(Ok(false)).boxed();
-        }
-
         let working_directory = self.working_directory();
         let git_binary_path = self.git_binary_path.clone();
 
@@ -1768,7 +1751,6 @@ fn checkpoint_author_envs() -> HashMap<String, String> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::status::FileStatus;
     use gpui::TestAppContext;
 
     #[gpui::test]
@@ -1803,7 +1785,6 @@ mod tests {
         smol::fs::write(repo_dir.path().join("new_file_before_checkpoint"), "1")
             .await
             .unwrap();
-        let sha_before_checkpoint = repo.head_sha().unwrap();
         let checkpoint = repo.checkpoint().await.unwrap();
 
         // Ensure the user can't see any branches after creating a checkpoint.
@@ -1837,7 +1818,6 @@ mod tests {
         repo.gc().await.unwrap();
         repo.restore_checkpoint(checkpoint.clone()).await.unwrap();
 
-        assert_eq!(repo.head_sha().unwrap(), sha_before_checkpoint);
         assert_eq!(
             smol::fs::read_to_string(&file_path).await.unwrap(),
             "modified before checkpoint"
@@ -1901,83 +1881,6 @@ mod tests {
         );
     }
 
-    #[gpui::test]
-    async fn test_undoing_commit_via_checkpoint(cx: &mut TestAppContext) {
-        cx.executor().allow_parking();
-
-        let repo_dir = tempfile::tempdir().unwrap();
-
-        git2::Repository::init(repo_dir.path()).unwrap();
-        let file_path = repo_dir.path().join("file");
-        smol::fs::write(&file_path, "initial").await.unwrap();
-
-        let repo =
-            RealGitRepository::new(&repo_dir.path().join(".git"), None, cx.executor()).unwrap();
-        repo.stage_paths(
-            vec![RepoPath::from_str("file")],
-            Arc::new(HashMap::default()),
-        )
-        .await
-        .unwrap();
-        repo.commit(
-            "Initial commit".into(),
-            None,
-            Arc::new(checkpoint_author_envs()),
-        )
-        .await
-        .unwrap();
-
-        let initial_commit_sha = repo.head_sha().unwrap();
-
-        smol::fs::write(repo_dir.path().join("new_file1"), "content1")
-            .await
-            .unwrap();
-        smol::fs::write(repo_dir.path().join("new_file2"), "content2")
-            .await
-            .unwrap();
-
-        let checkpoint = repo.checkpoint().await.unwrap();
-
-        repo.stage_paths(
-            vec![
-                RepoPath::from_str("new_file1"),
-                RepoPath::from_str("new_file2"),
-            ],
-            Arc::new(HashMap::default()),
-        )
-        .await
-        .unwrap();
-        repo.commit(
-            "Commit new files".into(),
-            None,
-            Arc::new(checkpoint_author_envs()),
-        )
-        .await
-        .unwrap();
-
-        repo.restore_checkpoint(checkpoint).await.unwrap();
-        assert_eq!(repo.head_sha().unwrap(), initial_commit_sha);
-        assert_eq!(
-            smol::fs::read_to_string(repo_dir.path().join("new_file1"))
-                .await
-                .unwrap(),
-            "content1"
-        );
-        assert_eq!(
-            smol::fs::read_to_string(repo_dir.path().join("new_file2"))
-                .await
-                .unwrap(),
-            "content2"
-        );
-        assert_eq!(
-            repo.status(&[]).await.unwrap().entries.as_ref(),
-            &[
-                (RepoPath::from_str("new_file1"), FileStatus::Untracked),
-                (RepoPath::from_str("new_file2"), FileStatus::Untracked)
-            ]
-        );
-    }
-
     #[gpui::test]
     async fn test_compare_checkpoints(cx: &mut TestAppContext) {
         cx.executor().allow_parking();

crates/language/src/buffer.rs

@@ -1321,7 +1321,7 @@ impl Buffer {
             this.update(cx, |this, cx| {
                 if this.version() == diff.base_version {
                     this.finalize_last_transaction();
-                    this.apply_diff(diff, true, cx);
+                    this.apply_diff(diff, cx);
                     tx.send(this.finalize_last_transaction().cloned()).ok();
                     this.has_conflict = false;
                     this.did_reload(this.version(), this.line_ending(), new_mtime, cx);
@@ -1882,14 +1882,7 @@ impl Buffer {
     /// Applies a diff to the buffer. If the buffer has changed since the given diff was
     /// calculated, then adjust the diff to account for those changes, and discard any
     /// parts of the diff that conflict with those changes.
-    ///
-    /// If `atomic` is true, the diff will be applied as a single edit.
-    pub fn apply_diff(
-        &mut self,
-        diff: Diff,
-        atomic: bool,
-        cx: &mut Context<Self>,
-    ) -> Option<TransactionId> {
+    pub fn apply_diff(&mut self, diff: Diff, cx: &mut Context<Self>) -> Option<TransactionId> {
         let snapshot = self.snapshot();
         let mut edits_since = snapshot.edits_since::<usize>(&diff.base_version).peekable();
         let mut delta = 0;
@@ -1919,17 +1912,7 @@ impl Buffer {
 
         self.start_transaction();
         self.text.set_line_ending(diff.line_ending);
-        if atomic {
-            self.edit(adjusted_edits, None, cx);
-        } else {
-            let mut delta = 0isize;
-            for (range, new_text) in adjusted_edits {
-                let adjusted_range =
-                    (range.start as isize + delta) as usize..(range.end as isize + delta) as usize;
-                delta += new_text.len() as isize - range.len() as isize;
-                self.edit([(adjusted_range, new_text)], None, cx);
-            }
-        }
+        self.edit(adjusted_edits, None, cx);
         self.end_transaction(cx)
     }
 

crates/language/src/buffer_tests.rs

@@ -376,7 +376,7 @@ async fn test_apply_diff(cx: &mut TestAppContext) {
 
     let diff = buffer.update(cx, |b, cx| b.diff(text.clone(), cx)).await;
     buffer.update(cx, |buffer, cx| {
-        buffer.apply_diff(diff, true, cx).unwrap();
+        buffer.apply_diff(diff, cx).unwrap();
         assert_eq!(buffer.text(), text);
         let actual_offsets = anchors
             .iter()
@@ -390,7 +390,7 @@ async fn test_apply_diff(cx: &mut TestAppContext) {
 
     let diff = buffer.update(cx, |b, cx| b.diff(text.clone(), cx)).await;
     buffer.update(cx, |buffer, cx| {
-        buffer.apply_diff(diff, true, cx).unwrap();
+        buffer.apply_diff(diff, cx).unwrap();
         assert_eq!(buffer.text(), text);
         let actual_offsets = anchors
             .iter()
@@ -435,7 +435,7 @@ async fn test_normalize_whitespace(cx: &mut gpui::TestAppContext) {
     let format_diff = format.await;
     buffer.update(cx, |buffer, cx| {
         let version_before_format = format_diff.base_version.clone();
-        buffer.apply_diff(format_diff, true, cx);
+        buffer.apply_diff(format_diff, cx);
 
         // The outcome depends on the order of concurrent tasks.
         //

crates/multi_buffer/src/multi_buffer.rs

@@ -163,6 +163,12 @@ impl MultiBufferDiffHunk {
         self.diff_base_byte_range == (0..0)
             && self.buffer_range == (text::Anchor::MIN..text::Anchor::MAX)
     }
+
+    pub fn multi_buffer_range(&self) -> Range<Anchor> {
+        let start = Anchor::in_buffer(self.excerpt_id, self.buffer_id, self.buffer_range.start);
+        let end = Anchor::in_buffer(self.excerpt_id, self.buffer_id, self.buffer_range.end);
+        start..end
+    }
 }
 
 #[derive(PartialEq, Eq, Ord, PartialOrd, Clone, Hash, Debug)]

crates/project/src/lsp_store.rs

@@ -1230,7 +1230,7 @@ impl LocalLspStore {
                         .await;
                     buffer.handle.update(cx, |buffer, cx| {
                         buffer.start_transaction();
-                        buffer.apply_diff(diff, true, cx);
+                        buffer.apply_diff(diff, cx);
                         transaction_id_format =
                             transaction_id_format.or(buffer.end_transaction(cx));
                         if let Some(transaction_id) = transaction_id_format {
@@ -1364,7 +1364,7 @@ impl LocalLspStore {
                         zlog::trace!(logger => "Applying changes");
                         buffer.handle.update(cx, |buffer, cx| {
                             buffer.start_transaction();
-                            buffer.apply_diff(diff, true, cx);
+                            buffer.apply_diff(diff, cx);
                             transaction_id_format =
                                 transaction_id_format.or(buffer.end_transaction(cx));
                             if let Some(transaction_id) = transaction_id_format {
@@ -1407,7 +1407,7 @@ impl LocalLspStore {
                         zlog::trace!(logger => "Applying changes");
                         buffer.handle.update(cx, |buffer, cx| {
                             buffer.start_transaction();
-                            buffer.apply_diff(diff, true, cx);
+                            buffer.apply_diff(diff, cx);
                             transaction_id_format =
                                 transaction_id_format.or(buffer.end_transaction(cx));
                             if let Some(transaction_id) = transaction_id_format {

crates/text/src/patch.rs

@@ -224,6 +224,15 @@ where
     }
 }
 
+impl<T> Patch<T> {
+    pub fn retain_mut<F>(&mut self, f: F)
+    where
+        F: FnMut(&mut Edit<T>) -> bool,
+    {
+        self.0.retain_mut(f);
+    }
+}
+
 impl<T: Clone> IntoIterator for Patch<T> {
     type Item = Edit<T>;
     type IntoIter = std::vec::IntoIter<Edit<T>>;