Start tracking edits performed by the agent (#27064)

Antonio Scandurra , Danilo Leal , and Agus Zubiaga created

Release Notes:

- N/A

---------

Co-authored-by: Danilo Leal <daniloleal09@gmail.com>
Co-authored-by: Agus Zubiaga <hi@aguz.me>

Change summary

Cargo.lock                                    |   2 
crates/assistant2/Cargo.toml                  |   2 
crates/assistant2/src/assistant.rs            |   2 
crates/assistant2/src/assistant_diff.rs       | 625 +++++++++++++++++++++
crates/assistant2/src/message_editor.rs       | 203 ++++--
crates/assistant2/src/thread.rs               |  17 
crates/assistant_tool/Cargo.toml              |  10 
crates/assistant_tool/src/action_log.rs       | 398 +++++++++++++
crates/assistant_tool/src/assistant_tool.rs   |  62 -
crates/assistant_tools/src/edit_files_tool.rs |  16 
crates/editor/src/editor.rs                   | 211 +++++++
crates/editor/src/element.rs                  | 205 ------
crates/text/src/text.rs                       |  16 
13 files changed, 1,424 insertions(+), 345 deletions(-)

Detailed changes

Cargo.lock 🔗

@@ -453,6 +453,7 @@ dependencies = [
  "assistant_slash_command",
  "assistant_tool",
  "async-watch",
+ "buffer_diff",
  "chrono",
  "client",
  "clock",
@@ -691,6 +692,7 @@ name = "assistant_tool"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "buffer_diff",
  "clock",
  "collections",
  "derive_more",

crates/assistant2/Cargo.toml 🔗

@@ -25,6 +25,7 @@ assistant_settings.workspace = true
 assistant_slash_command.workspace = true
 assistant_tool.workspace = true
 async-watch.workspace = true
+buffer_diff.workspace = true
 chrono.workspace = true
 client.workspace = true
 clock.workspace = true
@@ -82,6 +83,7 @@ workspace.workspace = true
 zed_actions.workspace = true
 
 [dev-dependencies]
+buffer_diff = { workspace = true, features = ["test-support"] }
 editor = { workspace = true, features = ["test-support"] }
 gpui = { workspace = true, "features" = ["test-support"] }
 indoc.workspace = true

crates/assistant2/src/assistant.rs 🔗

@@ -1,5 +1,6 @@
 mod active_thread;
 mod assistant_configuration;
+mod assistant_diff;
 mod assistant_model_selector;
 mod assistant_panel;
 mod buffer_codegen;
@@ -36,6 +37,7 @@ pub use crate::assistant_panel::{AssistantPanel, ConcreteAssistantPanelDelegate}
 pub use crate::inline_assistant::InlineAssistant;
 pub use crate::thread::{Message, RequestKind, Thread, ThreadEvent};
 pub use crate::thread_store::ThreadStore;
+pub use assistant_diff::AssistantDiff;
 
 actions!(
     assistant2,

crates/assistant2/src/assistant_diff.rs 🔗

@@ -0,0 +1,625 @@
+use crate::{Thread, ThreadEvent};
+use anyhow::Result;
+use buffer_diff::DiffHunkStatus;
+use collections::HashSet;
+use editor::{Editor, EditorEvent, MultiBuffer};
+use futures::future;
+use gpui::{
+    prelude::*, AnyElement, AnyView, App, Entity, EventEmitter, FocusHandle, Focusable,
+    SharedString, Subscription, Task, WeakEntity, Window,
+};
+use language::{Capability, OffsetRangeExt};
+use multi_buffer::PathKey;
+use project::{Project, ProjectPath};
+use std::{
+    any::{Any, TypeId},
+    ops::Range,
+    sync::Arc,
+};
+use ui::{prelude::*, IconButtonShape};
+use util::TryFutureExt;
+use workspace::{
+    item::{BreadcrumbText, ItemEvent, TabContentParams},
+    searchable::SearchableItemHandle,
+    Item, ItemHandle, ItemNavHistory, ToolbarItemLocation, Workspace,
+};
+
+pub struct AssistantDiff {
+    multibuffer: Entity<MultiBuffer>,
+    editor: Entity<Editor>,
+    thread: Entity<Thread>,
+    focus_handle: FocusHandle,
+    workspace: WeakEntity<Workspace>,
+    title: SharedString,
+    _subscriptions: Vec<Subscription>,
+}
+
+impl AssistantDiff {
+    pub fn deploy(
+        thread: Entity<Thread>,
+        workspace: WeakEntity<Workspace>,
+        window: &mut Window,
+        cx: &mut App,
+    ) -> Result<()> {
+        let existing_diff = workspace.update(cx, |workspace, cx| {
+            workspace
+                .items_of_type::<AssistantDiff>(cx)
+                .find(|diff| diff.read(cx).thread == thread)
+        })?;
+        if let Some(existing_diff) = existing_diff {
+            workspace.update(cx, |workspace, cx| {
+                workspace.activate_item(&existing_diff, true, true, window, cx);
+            })
+        } else {
+            let assistant_diff =
+                cx.new(|cx| AssistantDiff::new(thread.clone(), workspace.clone(), window, cx));
+            workspace.update(cx, |workspace, cx| {
+                workspace.add_item_to_center(Box::new(assistant_diff), window, cx);
+            })
+        }
+    }
+
+    pub fn new(
+        thread: Entity<Thread>,
+        workspace: WeakEntity<Workspace>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Self {
+        let focus_handle = cx.focus_handle();
+        let multibuffer = cx.new(|_| MultiBuffer::new(Capability::ReadWrite));
+
+        let project = thread.read(cx).project().clone();
+        let render_diff_hunk_controls = Arc::new({
+            let assistant_diff = cx.entity();
+            move |row,
+                  status: &DiffHunkStatus,
+                  hunk_range,
+                  is_created_file,
+                  line_height,
+                  _editor: &Entity<Editor>,
+                  cx: &mut App| {
+                render_diff_hunk_controls(
+                    row,
+                    status,
+                    hunk_range,
+                    is_created_file,
+                    line_height,
+                    &assistant_diff,
+                    cx,
+                )
+            }
+        });
+        let editor = cx.new(|cx| {
+            let mut editor =
+                Editor::for_multibuffer(multibuffer.clone(), Some(project.clone()), window, cx);
+            editor.disable_inline_diagnostics();
+            editor.set_expand_all_diff_hunks(cx);
+            editor.set_render_diff_hunk_controls(render_diff_hunk_controls, cx);
+            editor
+        });
+
+        let action_log = thread.read(cx).action_log().clone();
+        let mut this = Self {
+            _subscriptions: vec![
+                cx.observe_in(&action_log, window, |this, _action_log, window, cx| {
+                    this.update_excerpts(window, cx)
+                }),
+                cx.subscribe(&thread, |this, _thread, event, cx| {
+                    this.handle_thread_event(event, cx)
+                }),
+            ],
+            title: SharedString::default(),
+            multibuffer,
+            editor,
+            thread,
+            focus_handle,
+            workspace,
+        };
+        this.update_excerpts(window, cx);
+        this.update_title(cx);
+        this
+    }
+
+    fn update_excerpts(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        let thread = self.thread.read(cx);
+        let unreviewed_buffers = thread.action_log().read(cx).unreviewed_buffers();
+        let mut paths_to_delete = self.multibuffer.read(cx).paths().collect::<HashSet<_>>();
+
+        for (buffer, tracked) in unreviewed_buffers {
+            let Some(file) = buffer.read(cx).file().cloned() else {
+                continue;
+            };
+
+            let path_key = PathKey::namespaced("", file.full_path(cx).into());
+            paths_to_delete.remove(&path_key);
+
+            let snapshot = buffer.read(cx).snapshot();
+            let diff = tracked.diff().read(cx);
+            let diff_hunk_ranges = diff
+                .hunks_intersecting_range(
+                    language::Anchor::MIN..language::Anchor::MAX,
+                    &snapshot,
+                    cx,
+                )
+                .map(|diff_hunk| diff_hunk.buffer_range.to_point(&snapshot))
+                .collect::<Vec<_>>();
+
+            let was_empty = self.multibuffer.update(cx, |multibuffer, cx| {
+                let was_empty = multibuffer.is_empty();
+                multibuffer.set_excerpts_for_path(
+                    path_key.clone(),
+                    buffer,
+                    diff_hunk_ranges,
+                    editor::DEFAULT_MULTIBUFFER_CONTEXT,
+                    cx,
+                );
+                multibuffer.add_diff(tracked.diff().clone(), cx);
+                was_empty
+            });
+
+            self.editor.update(cx, |editor, cx| {
+                if was_empty {
+                    editor.change_selections(None, window, cx, |selections| {
+                        // TODO select the very beginning (possibly inside a deletion)
+                        selections.select_ranges([0..0])
+                    });
+                }
+            });
+        }
+
+        self.multibuffer.update(cx, |multibuffer, cx| {
+            for path in paths_to_delete {
+                multibuffer.remove_excerpts_for_path(path, cx);
+            }
+        });
+
+        if self.multibuffer.read(cx).is_empty()
+            && self
+                .editor
+                .read(cx)
+                .focus_handle(cx)
+                .contains_focused(window, cx)
+        {
+            self.focus_handle.focus(window);
+        } else if self.focus_handle.is_focused(window) && !self.multibuffer.read(cx).is_empty() {
+            self.editor.update(cx, |editor, cx| {
+                editor.focus_handle(cx).focus(window);
+            });
+        }
+    }
+
+    fn update_title(&mut self, cx: &mut Context<Self>) {
+        let new_title = self
+            .thread
+            .read(cx)
+            .summary()
+            .unwrap_or("Assistant Changes".into());
+        if new_title != self.title {
+            self.title = new_title;
+            cx.emit(EditorEvent::TitleChanged);
+        }
+    }
+
+    fn handle_thread_event(&mut self, event: &ThreadEvent, cx: &mut Context<Self>) {
+        match event {
+            ThreadEvent::SummaryChanged => self.update_title(cx),
+            _ => {}
+        }
+    }
+
+    fn review_diff_hunks(
+        &mut self,
+        hunk_ranges: Vec<Range<editor::Anchor>>,
+        accept: bool,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let snapshot = self.multibuffer.read(cx).snapshot(cx);
+        let diff_hunks_in_ranges = self
+            .editor
+            .read(cx)
+            .diff_hunks_in_ranges(&hunk_ranges, &snapshot)
+            .collect::<Vec<_>>();
+
+        let mut tasks = Vec::new();
+        for hunk in diff_hunks_in_ranges {
+            let buffer = self.multibuffer.read(cx).buffer(hunk.buffer_id);
+            if let Some(buffer) = buffer {
+                let task = self.thread.update(cx, |thread, cx| {
+                    thread.review_edits_in_range(buffer, hunk.buffer_range, accept, cx)
+                });
+                tasks.push(task.log_err());
+            }
+        }
+
+        cx.spawn_in(window, async move |this, cx| {
+            future::join_all(tasks).await;
+            this.update_in(cx, |this, window, cx| this.update_excerpts(window, cx))
+        })
+        .detach_and_log_err(cx);
+    }
+}
+
+impl EventEmitter<EditorEvent> for AssistantDiff {}
+
+impl Focusable for AssistantDiff {
+    fn focus_handle(&self, cx: &App) -> FocusHandle {
+        if self.multibuffer.read(cx).is_empty() {
+            self.focus_handle.clone()
+        } else {
+            self.editor.focus_handle(cx)
+        }
+    }
+}
+
+impl Item for AssistantDiff {
+    type Event = EditorEvent;
+
+    fn tab_icon(&self, _window: &Window, _cx: &App) -> Option<Icon> {
+        Some(Icon::new(IconName::ZedAssistant).color(Color::Muted))
+    }
+
+    fn to_item_events(event: &EditorEvent, f: impl FnMut(ItemEvent)) {
+        Editor::to_item_events(event, f)
+    }
+
+    fn deactivated(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        self.editor
+            .update(cx, |editor, cx| editor.deactivated(window, cx));
+    }
+
+    fn navigate(
+        &mut self,
+        data: Box<dyn Any>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> bool {
+        self.editor
+            .update(cx, |editor, cx| editor.navigate(data, window, cx))
+    }
+
+    fn tab_tooltip_text(&self, _: &App) -> Option<SharedString> {
+        Some("Project Diff".into())
+    }
+
+    fn tab_content(&self, params: TabContentParams, _window: &Window, cx: &App) -> AnyElement {
+        let summary = self
+            .thread
+            .read(cx)
+            .summary()
+            .unwrap_or("Assistant Changes".into());
+        Label::new(format!("Review: {}", summary))
+            .color(if params.selected {
+                Color::Default
+            } else {
+                Color::Muted
+            })
+            .into_any_element()
+    }
+
+    fn telemetry_event_text(&self) -> Option<&'static str> {
+        Some("Project Diff Opened")
+    }
+
+    fn as_searchable(&self, _: &Entity<Self>) -> Option<Box<dyn SearchableItemHandle>> {
+        Some(Box::new(self.editor.clone()))
+    }
+
+    fn for_each_project_item(
+        &self,
+        cx: &App,
+        f: &mut dyn FnMut(gpui::EntityId, &dyn project::ProjectItem),
+    ) {
+        self.editor.for_each_project_item(cx, f)
+    }
+
+    fn is_singleton(&self, _: &App) -> bool {
+        false
+    }
+
+    fn set_nav_history(
+        &mut self,
+        nav_history: ItemNavHistory,
+        _: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        self.editor.update(cx, |editor, _| {
+            editor.set_nav_history(Some(nav_history));
+        });
+    }
+
+    fn clone_on_split(
+        &self,
+        _workspace_id: Option<workspace::WorkspaceId>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Option<Entity<Self>>
+    where
+        Self: Sized,
+    {
+        Some(cx.new(|cx| Self::new(self.thread.clone(), self.workspace.clone(), window, cx)))
+    }
+
+    fn is_dirty(&self, cx: &App) -> bool {
+        self.multibuffer.read(cx).is_dirty(cx)
+    }
+
+    fn has_conflict(&self, cx: &App) -> bool {
+        self.multibuffer.read(cx).has_conflict(cx)
+    }
+
+    fn can_save(&self, _: &App) -> bool {
+        true
+    }
+
+    fn save(
+        &mut self,
+        format: bool,
+        project: Entity<Project>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<()>> {
+        self.editor.save(format, project, window, cx)
+    }
+
+    fn save_as(
+        &mut self,
+        _: Entity<Project>,
+        _: ProjectPath,
+        _window: &mut Window,
+        _: &mut Context<Self>,
+    ) -> Task<Result<()>> {
+        unreachable!()
+    }
+
+    fn reload(
+        &mut self,
+        project: Entity<Project>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<()>> {
+        self.editor.reload(project, window, cx)
+    }
+
+    fn act_as_type<'a>(
+        &'a self,
+        type_id: TypeId,
+        self_handle: &'a Entity<Self>,
+        _: &'a App,
+    ) -> Option<AnyView> {
+        if type_id == TypeId::of::<Self>() {
+            Some(self_handle.to_any())
+        } else if type_id == TypeId::of::<Editor>() {
+            Some(self.editor.to_any())
+        } else {
+            None
+        }
+    }
+
+    fn breadcrumb_location(&self, _: &App) -> ToolbarItemLocation {
+        ToolbarItemLocation::PrimaryLeft
+    }
+
+    fn breadcrumbs(&self, theme: &theme::Theme, cx: &App) -> Option<Vec<BreadcrumbText>> {
+        self.editor.breadcrumbs(theme, cx)
+    }
+
+    fn added_to_workspace(
+        &mut self,
+        workspace: &mut Workspace,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        self.editor.update(cx, |editor, cx| {
+            editor.added_to_workspace(workspace, window, cx)
+        });
+    }
+}
+
+impl Render for AssistantDiff {
+    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        let is_empty = self.multibuffer.read(cx).is_empty();
+        div()
+            .track_focus(&self.focus_handle)
+            .key_context(if is_empty {
+                "EmptyPane"
+            } else {
+                "AssistantDiff"
+            })
+            .bg(cx.theme().colors().editor_background)
+            .flex()
+            .items_center()
+            .justify_center()
+            .size_full()
+            .when(is_empty, |el| el.child("No changes to review"))
+            .when(!is_empty, |el| el.child(self.editor.clone()))
+    }
+}
+
+fn render_diff_hunk_controls(
+    row: u32,
+    status: &DiffHunkStatus,
+    hunk_range: Range<editor::Anchor>,
+    is_created_file: bool,
+    line_height: Pixels,
+    assistant_diff: &Entity<AssistantDiff>,
+    cx: &mut App,
+) -> AnyElement {
+    let editor = assistant_diff.read(cx).editor.clone();
+    h_flex()
+        .h(line_height)
+        .mr_1()
+        .gap_1()
+        .px_0p5()
+        .pb_1()
+        .border_x_1()
+        .border_b_1()
+        .border_color(cx.theme().colors().border_variant)
+        .rounded_b_lg()
+        .bg(cx.theme().colors().editor_background)
+        .gap_1()
+        .occlude()
+        .shadow_md()
+        .children(if status.has_secondary_hunk() {
+            vec![
+                Button::new(("stage", row as u64), "Accept")
+                    .alpha(if status.is_pending() { 0.66 } else { 1.0 })
+                    // TODO: add tooltip
+                    // .tooltip({
+                    //     let focus_handle = editor.focus_handle(cx);
+                    //     move |window, cx| {
+                    //         Tooltip::for_action_in(
+                    //             "Stage Hunk",
+                    //             &::git::ToggleStaged,
+                    //             &focus_handle,
+                    //             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],
+                                    true,
+                                    window,
+                                    cx,
+                                );
+                            });
+                        }
+                    }),
+                Button::new("undo", "Undo")
+                    // TODO: add tooltip
+                    // .tooltip({
+                    //     let focus_handle = editor.focus_handle(cx);
+                    //     move |window, cx| {
+                    //         Tooltip::for_action_in("Undo Hunk", &::git::Undo, &focus_handle, window, cx)
+                    //     }
+                    // })
+                    .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.undo_hunks_in_ranges(vec![point..point], window, cx);
+                            // });
+                        }
+                    })
+                    .disabled(is_created_file),
+            ]
+        } else {
+            vec![Button::new(("review", row as u64), "Review")
+                .alpha(if status.is_pending() { 0.66 } else { 1.0 })
+                // TODO: add tooltip
+                // .tooltip({
+                //     let focus_handle = editor.focus_handle(cx);
+                //     move |window, cx| {
+                //         Tooltip::for_action_in(
+                //             "Review",
+                //             &::git::ToggleStaged,
+                //             &focus_handle,
+                //             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,
+                                window,
+                                cx,
+                            );
+                        });
+                    }
+                })]
+        })
+        .when(
+            !editor.read(cx).buffer().read(cx).all_diff_hunks_expanded(),
+            |el| {
+                el.child(
+                    IconButton::new(("next-hunk", row as u64), IconName::ArrowDown)
+                        .shape(IconButtonShape::Square)
+                        .icon_size(IconSize::Small)
+                        // .disabled(!has_multiple_hunks)
+                        // TODO: add tooltip
+                        // .tooltip({
+                        //     let focus_handle = editor.focus_handle(cx);
+                        //     move |window, cx| {
+                        //         Tooltip::for_action_in(
+                        //             "Next Hunk",
+                        //             &GoToHunk,
+                        //             &focus_handle,
+                        //             window,
+                        //             cx,
+                        //         )
+                        //     }
+                        // })
+                        .on_click({
+                            let _editor = editor.clone();
+                            move |_event, _window, _cx| {
+                                // TODO: wire this up
+                                // editor.update(cx, |editor, cx| {
+                                //     let snapshot = editor.snapshot(window, cx);
+                                //     let position =
+                                //         hunk_range.end.to_point(&snapshot.buffer_snapshot);
+                                //     editor.go_to_hunk_before_or_after_position(
+                                //         &snapshot,
+                                //         position,
+                                //         Direction::Next,
+                                //         window,
+                                //         cx,
+                                //     );
+                                //     editor.expand_selected_diff_hunks(cx);
+                                // });
+                            }
+                        }),
+                )
+                .child(
+                    IconButton::new(("prev-hunk", row as u64), IconName::ArrowUp)
+                        .shape(IconButtonShape::Square)
+                        .icon_size(IconSize::Small)
+                        // .disabled(!has_multiple_hunks)
+                        // TODO: add tooltip
+                        // .tooltip({
+                        //     let focus_handle = editor.focus_handle(cx);
+                        //     move |window, cx| {
+                        //         Tooltip::for_action_in(
+                        //             "Previous Hunk",
+                        //             &GoToPreviousHunk,
+                        //             &focus_handle,
+                        //             window,
+                        //             cx,
+                        //         )
+                        //     }
+                        // })
+                        .on_click({
+                            let _editor = editor.clone();
+                            move |_event, _window, _cx| {
+                                // TODO: wire this up
+                                // editor.update(cx, |editor, cx| {
+                                //     let snapshot = editor.snapshot(window, cx);
+                                //     let point =
+                                //         hunk_range.start.to_point(&snapshot.buffer_snapshot);
+                                //     editor.go_to_hunk_before_or_after_position(
+                                //         &snapshot,
+                                //         point,
+                                //         Direction::Prev,
+                                //         window,
+                                //         cx,
+                                //     );
+                                //     editor.expand_selected_diff_hunks(cx);
+                                // });
+                            }
+                        }),
+                )
+            },
+        )
+        .into_any_element()
+}

crates/assistant2/src/message_editor.rs 🔗

@@ -20,6 +20,7 @@ use ui::{
     prelude::*, ButtonLike, Disclosure, KeyBinding, PlatformStyle, PopoverMenu, PopoverMenuHandle,
     Tooltip,
 };
+use util::ResultExt;
 use vim_mode_setting::VimModeSetting;
 use workspace::notifications::{NotificationId, NotifyTaskExt};
 use workspace::{Toast, Workspace};
@@ -31,7 +32,7 @@ use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
 use crate::thread::{RequestKind, Thread};
 use crate::thread_store::ThreadStore;
 use crate::tool_selector::ToolSelector;
-use crate::{Chat, ChatMode, RemoveAllContext, ToggleContextPicker};
+use crate::{AssistantDiff, Chat, ChatMode, RemoveAllContext, ToggleContextPicker};
 
 pub struct MessageEditor {
     thread: Entity<Thread>,
@@ -313,6 +314,10 @@ impl MessageEditor {
         })
         .detach_and_notify_err(window, cx);
     }
+
+    fn handle_review_click(&self, window: &mut Window, cx: &mut Context<Self>) {
+        AssistantDiff::deploy(self.thread.clone(), self.workspace.clone(), window, cx).log_err();
+    }
 }
 
 impl Focusable for MessageEditor {
@@ -347,8 +352,9 @@ impl Render for MessageEditor {
             px(64.)
         };
 
-        let changed_buffers = self.thread.read(cx).scripting_changed_buffers(cx);
-        let changed_buffers_count = changed_buffers.len();
+        let action_log = self.thread.read(cx).action_log();
+        let unreviewed_buffers = action_log.read(cx).unreviewed_buffers();
+        let unreviewed_buffers_count = unreviewed_buffers.len();
 
         v_flex()
             .size_full()
@@ -410,7 +416,7 @@ impl Render for MessageEditor {
                     ),
                 )
             })
-            .when(changed_buffers_count > 0, |parent| {
+            .when(unreviewed_buffers_count > 0, |parent| {
                 parent.child(
                     v_flex()
                         .mx_2()
@@ -421,93 +427,130 @@ impl Render for MessageEditor {
                         .rounded_t_md()
                         .child(
                             h_flex()
-                                .gap_2()
                                 .p_2()
+                                .justify_between()
                                 .child(
-                                    Disclosure::new("edits-disclosure", self.edits_expanded)
-                                        .on_click(cx.listener(|this, _ev, _window, cx| {
-                                            this.edits_expanded = !this.edits_expanded;
-                                            cx.notify();
-                                        })),
-                                )
-                                .child(
-                                    Label::new("Edits")
-                                        .size(LabelSize::XSmall)
-                                        .color(Color::Muted),
+                                    h_flex()
+                                        .gap_2()
+                                        .child(
+                                            Disclosure::new(
+                                                "edits-disclosure",
+                                                self.edits_expanded,
+                                            )
+                                            .on_click(
+                                                cx.listener(|this, _ev, _window, cx| {
+                                                    this.edits_expanded = !this.edits_expanded;
+                                                    cx.notify();
+                                                }),
+                                            ),
+                                        )
+                                        .child(
+                                            Label::new("Edits")
+                                                .size(LabelSize::XSmall)
+                                                .color(Color::Muted),
+                                        )
+                                        .child(
+                                            Label::new("•")
+                                                .size(LabelSize::XSmall)
+                                                .color(Color::Muted),
+                                        )
+                                        .child(
+                                            Label::new(format!(
+                                                "{} {}",
+                                                unreviewed_buffers_count,
+                                                if unreviewed_buffers_count == 1 {
+                                                    "file"
+                                                } else {
+                                                    "files"
+                                                }
+                                            ))
+                                            .size(LabelSize::XSmall)
+                                            .color(Color::Muted),
+                                        ),
                                 )
-                                .child(Label::new("•").size(LabelSize::XSmall).color(Color::Muted))
                                 .child(
-                                    Label::new(format!(
-                                        "{} {}",
-                                        changed_buffers_count,
-                                        if changed_buffers_count == 1 {
-                                            "file"
-                                        } else {
-                                            "files"
-                                        }
-                                    ))
-                                    .size(LabelSize::XSmall)
-                                    .color(Color::Muted),
+                                    Button::new("review", "Review")
+                                        .label_size(LabelSize::XSmall)
+                                        .on_click(cx.listener(|this, _, window, cx| {
+                                            this.handle_review_click(window, cx)
+                                        })),
                                 ),
                         )
                         .when(self.edits_expanded, |parent| {
                             parent.child(
                                 v_flex().bg(cx.theme().colors().editor_background).children(
-                                    changed_buffers.enumerate().flat_map(|(index, buffer)| {
-                                        let file = buffer.read(cx).file()?;
-                                        let path = file.path();
-
-                                        let parent_label = path.parent().and_then(|parent| {
-                                            let parent_str = parent.to_string_lossy();
-
-                                            if parent_str.is_empty() {
-                                                None
-                                            } else {
-                                                Some(
-                                                    Label::new(format!(
-                                                        "{}{}",
-                                                        parent_str,
-                                                        std::path::MAIN_SEPARATOR_STR
-                                                    ))
-                                                    .color(Color::Muted)
-                                                    .size(LabelSize::Small),
-                                                )
-                                            }
-                                        });
-
-                                        let name_label = path.file_name().map(|name| {
-                                            Label::new(name.to_string_lossy().to_string())
-                                                .size(LabelSize::Small)
-                                        });
-
-                                        let file_icon = FileIcons::get_icon(&path, cx)
-                                            .map(Icon::from_path)
-                                            .unwrap_or_else(|| Icon::new(IconName::File));
-
-                                        let element = div()
-                                            .p_2()
-                                            .when(index + 1 < changed_buffers_count, |parent| {
-                                                parent
-                                                    .border_color(cx.theme().colors().border)
-                                                    .border_b_1()
-                                            })
-                                            .child(
-                                                h_flex()
-                                                    .gap_2()
-                                                    .child(file_icon)
-                                                    .child(
-                                                        // TODO: handle overflow
-                                                        h_flex()
-                                                            .children(parent_label)
-                                                            .children(name_label),
+                                    unreviewed_buffers.into_iter().enumerate().flat_map(
+                                        |(index, (buffer, tracked))| {
+                                            let file = buffer.read(cx).file()?;
+                                            let path = file.path();
+
+                                            let parent_label = path.parent().and_then(|parent| {
+                                                let parent_str = parent.to_string_lossy();
+
+                                                if parent_str.is_empty() {
+                                                    None
+                                                } else {
+                                                    Some(
+                                                        Label::new(format!(
+                                                            "{}{}",
+                                                            parent_str,
+                                                            std::path::MAIN_SEPARATOR_STR
+                                                        ))
+                                                        .color(Color::Muted)
+                                                        .size(LabelSize::Small),
                                                     )
-                                                    // TODO: show lines changed
-                                                    .child(Label::new("+").color(Color::Created))
-                                                    .child(Label::new("-").color(Color::Deleted)),
-                                            );
+                                                }
+                                            });
+
+                                            let name_label = path.file_name().map(|name| {
+                                                Label::new(name.to_string_lossy().to_string())
+                                                    .size(LabelSize::Small)
+                                            });
+
+                                            let file_icon = FileIcons::get_icon(&path, cx)
+                                                .map(Icon::from_path)
+                                                .unwrap_or_else(|| Icon::new(IconName::File));
+
+                                            let element = div()
+                                                .p_2()
+                                                .when(
+                                                    index + 1 < unreviewed_buffers_count,
+                                                    |parent| {
+                                                        parent
+                                                            .border_color(
+                                                                cx.theme().colors().border,
+                                                            )
+                                                            .border_b_1()
+                                                    },
+                                                )
+                                                .child(
+                                                    h_flex()
+                                                        .gap_2()
+                                                        .child(file_icon)
+                                                        .child(
+                                                            // TODO: handle overflow
+                                                            h_flex()
+                                                                .children(parent_label)
+                                                                .children(name_label),
+                                                        )
+                                                        // TODO: show lines changed
+                                                        .child(
+                                                            Label::new("+").color(Color::Created),
+                                                        )
+                                                        .child(
+                                                            Label::new("-").color(Color::Deleted),
+                                                        )
+                                                        .when(!tracked.needs_review(), |parent| {
+                                                            parent.child(
+                                                                Icon::new(IconName::Check)
+                                                                    .color(Color::Success),
+                                                            )
+                                                        }),
+                                                );
 
-                                        Some(element)
-                                    }),
+                                            Some(element)
+                                        },
+                                    ),
                                 ),
                             )
                         }),

crates/assistant2/src/thread.rs 🔗

@@ -1,5 +1,6 @@
 use std::fmt::Write as _;
 use std::io::Write;
+use std::ops::Range;
 use std::sync::Arc;
 
 use anyhow::{Context as _, Result};
@@ -975,6 +976,10 @@ impl Thread {
         })
     }
 
+    pub fn project(&self) -> &Entity<Project> {
+        &self.project
+    }
+
     /// Create a snapshot of the current project state including git information and unsaved buffers.
     fn project_snapshot(
         project: Entity<Project>,
@@ -1123,6 +1128,18 @@ impl Thread {
         Ok(String::from_utf8_lossy(&markdown).to_string())
     }
 
+    pub fn review_edits_in_range(
+        &mut self,
+        buffer: Entity<language::Buffer>,
+        buffer_range: Range<language::Anchor>,
+        accept: bool,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<()>> {
+        self.action_log.update(cx, |action_log, cx| {
+            action_log.review_edits_in_range(buffer, buffer_range, accept, cx)
+        })
+    }
+
     pub fn action_log(&self) -> &Entity<ActionLog> {
         &self.action_log
     }

crates/assistant_tool/Cargo.toml 🔗

@@ -13,6 +13,7 @@ path = "src/assistant_tool.rs"
 
 [dependencies]
 anyhow.workspace = true
+buffer_diff.workspace = true
 collections.workspace = true
 clock.workspace = true
 derive_more.workspace = true
@@ -23,3 +24,12 @@ parking_lot.workspace = true
 project.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+
+[dev-dependencies]
+buffer_diff = { workspace = true, features = ["test-support"] }
+collections = { workspace = true, features = ["test-support"] }
+clock = { workspace = true, features = ["test-support"] }
+gpui = { workspace = true, features = ["test-support"] }
+language = { workspace = true, features = ["test-support"] }
+language_model = { workspace = true, features = ["test-support"] }
+project = { workspace = true, features = ["test-support"] }

crates/assistant_tool/src/action_log.rs 🔗

@@ -0,0 +1,398 @@
+use anyhow::{anyhow, Result};
+use buffer_diff::BufferDiff;
+use collections::{BTreeMap, HashMap, HashSet};
+use gpui::{App, AppContext, Context, Entity, Task};
+use language::{Buffer, OffsetRangeExt, ToOffset};
+use std::{future::Future, ops::Range};
+
+/// Tracks actions performed by tools in a thread
+#[derive(Debug)]
+pub struct ActionLog {
+    /// Buffers that user manually added to the context, and whose content has
+    /// changed since the model last saw them.
+    stale_buffers_in_context: HashSet<Entity<Buffer>>,
+    /// Buffers that we want to notify the model about when they change.
+    tracked_buffers: BTreeMap<Entity<Buffer>, TrackedBuffer>,
+}
+
+#[derive(Debug, Clone)]
+pub struct TrackedBuffer {
+    buffer: Entity<Buffer>,
+    unreviewed_edit_ids: Vec<clock::Lamport>,
+    accepted_edit_ids: Vec<clock::Lamport>,
+    version: clock::Global,
+    diff: Entity<BufferDiff>,
+    secondary_diff: Entity<BufferDiff>,
+}
+
+impl TrackedBuffer {
+    pub fn needs_review(&self) -> bool {
+        !self.unreviewed_edit_ids.is_empty()
+    }
+
+    pub fn diff(&self) -> &Entity<BufferDiff> {
+        &self.diff
+    }
+
+    fn update_diff(&mut self, cx: &mut App) -> impl 'static + Future<Output = ()> {
+        let edits_to_undo = self
+            .unreviewed_edit_ids
+            .iter()
+            .chain(&self.accepted_edit_ids)
+            .map(|edit_id| (*edit_id, u32::MAX))
+            .collect::<HashMap<_, _>>();
+        let buffer_without_edits = self.buffer.update(cx, |buffer, cx| buffer.branch(cx));
+        buffer_without_edits.update(cx, |buffer, cx| {
+            buffer.undo_operations(edits_to_undo, cx);
+        });
+        let primary_diff_update = self.diff.update(cx, |diff, cx| {
+            diff.set_base_text(
+                buffer_without_edits,
+                self.buffer.read(cx).text_snapshot(),
+                cx,
+            )
+        });
+
+        let unreviewed_edits_to_undo = self
+            .unreviewed_edit_ids
+            .iter()
+            .map(|edit_id| (*edit_id, u32::MAX))
+            .collect::<HashMap<_, _>>();
+        let buffer_without_unreviewed_edits =
+            self.buffer.update(cx, |buffer, cx| buffer.branch(cx));
+        buffer_without_unreviewed_edits.update(cx, |buffer, cx| {
+            buffer.undo_operations(unreviewed_edits_to_undo, cx);
+        });
+        let secondary_diff_update = self.secondary_diff.update(cx, |diff, cx| {
+            diff.set_base_text(
+                buffer_without_unreviewed_edits.clone(),
+                self.buffer.read(cx).text_snapshot(),
+                cx,
+            )
+        });
+
+        async move {
+            _ = primary_diff_update.await;
+            _ = secondary_diff_update.await;
+        }
+    }
+}
+
+impl ActionLog {
+    /// Creates a new, empty action log.
+    pub fn new() -> Self {
+        Self {
+            stale_buffers_in_context: HashSet::default(),
+            tracked_buffers: BTreeMap::default(),
+        }
+    }
+
+    fn track_buffer(
+        &mut self,
+        buffer: Entity<Buffer>,
+        cx: &mut Context<Self>,
+    ) -> &mut TrackedBuffer {
+        let tracked_buffer = self
+            .tracked_buffers
+            .entry(buffer.clone())
+            .or_insert_with(|| {
+                let text_snapshot = buffer.read(cx).text_snapshot();
+                let unreviewed_diff = cx.new(|cx| BufferDiff::new(&text_snapshot, cx));
+                let diff = cx.new(|cx| {
+                    let mut diff = BufferDiff::new(&text_snapshot, cx);
+                    diff.set_secondary_diff(unreviewed_diff.clone());
+                    diff
+                });
+                TrackedBuffer {
+                    buffer: buffer.clone(),
+                    unreviewed_edit_ids: Vec::new(),
+                    accepted_edit_ids: Vec::new(),
+                    version: buffer.read(cx).version(),
+                    diff,
+                    secondary_diff: unreviewed_diff,
+                }
+            });
+        tracked_buffer.version = buffer.read(cx).version();
+        tracked_buffer
+    }
+
+    /// Track a buffer as read, so we can notify the model about user edits.
+    pub fn buffer_read(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
+        self.track_buffer(buffer, cx);
+    }
+
+    /// Mark a buffer as edited, so we can refresh it in the context
+    pub fn buffer_edited(
+        &mut self,
+        buffer: Entity<Buffer>,
+        edit_ids: Vec<clock::Lamport>,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<()>> {
+        self.stale_buffers_in_context.insert(buffer.clone());
+
+        let tracked_buffer = self.track_buffer(buffer.clone(), cx);
+        tracked_buffer
+            .unreviewed_edit_ids
+            .extend(edit_ids.iter().copied());
+        let update = tracked_buffer.update_diff(cx);
+        cx.spawn(async move |this, cx| {
+            update.await;
+            this.update(cx, |_this, cx| cx.notify())?;
+            Ok(())
+        })
+    }
+
+    /// Accepts edits in a given range within a buffer.
+    pub fn review_edits_in_range<T: ToOffset>(
+        &mut self,
+        buffer: Entity<Buffer>,
+        buffer_range: Range<T>,
+        accept: bool,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<()>> {
+        let Some(tracked_buffer) = self.tracked_buffers.get_mut(&buffer) else {
+            return Task::ready(Err(anyhow!("buffer not found")));
+        };
+
+        let buffer = buffer.read(cx);
+        let buffer_range = buffer_range.to_offset(buffer);
+
+        let source;
+        let destination;
+        if accept {
+            source = &mut tracked_buffer.unreviewed_edit_ids;
+            destination = &mut tracked_buffer.accepted_edit_ids;
+        } else {
+            source = &mut tracked_buffer.accepted_edit_ids;
+            destination = &mut tracked_buffer.unreviewed_edit_ids;
+        }
+
+        source.retain(|edit_id| {
+            for range in buffer.edited_ranges_for_edit_ids::<usize>([edit_id]) {
+                if buffer_range.end >= range.start && buffer_range.start <= range.end {
+                    destination.push(*edit_id);
+                    return false;
+                }
+            }
+            true
+        });
+
+        let update = tracked_buffer.update_diff(cx);
+        cx.spawn(async move |this, cx| {
+            update.await;
+            this.update(cx, |_this, cx| cx.notify())?;
+            Ok(())
+        })
+    }
+
+    /// Returns the set of buffers that contain changes that haven't been reviewed by the user.
+    pub fn unreviewed_buffers(&self) -> BTreeMap<Entity<Buffer>, TrackedBuffer> {
+        self.tracked_buffers
+            .iter()
+            .map(|(buffer, tracked)| (buffer.clone(), tracked.clone()))
+            .collect()
+    }
+
+    /// Iterate over buffers changed since last read or edited by the model
+    pub fn stale_buffers<'a>(&'a self, cx: &'a App) -> impl Iterator<Item = &'a Entity<Buffer>> {
+        self.tracked_buffers
+            .iter()
+            .filter(|(buffer, tracked)| tracked.version != buffer.read(cx).version)
+            .map(|(buffer, _)| buffer)
+    }
+
+    /// Takes and returns the set of buffers pending refresh, clearing internal state.
+    pub fn take_stale_buffers_in_context(&mut self) -> HashSet<Entity<Buffer>> {
+        std::mem::take(&mut self.stale_buffers_in_context)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use buffer_diff::DiffHunkStatusKind;
+    use gpui::TestAppContext;
+    use language::Point;
+
+    #[gpui::test]
+    async fn test_edit_review(cx: &mut TestAppContext) {
+        let action_log = cx.new(|_| ActionLog::new());
+        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl\nmno", cx));
+
+        let edit1 = buffer.update(cx, |buffer, cx| {
+            buffer
+                .edit([(Point::new(1, 1)..Point::new(1, 2), "E")], None, cx)
+                .unwrap()
+        });
+        let edit2 = buffer.update(cx, |buffer, cx| {
+            buffer
+                .edit([(Point::new(4, 2)..Point::new(4, 3), "O")], None, cx)
+                .unwrap()
+        });
+        assert_eq!(
+            buffer.read_with(cx, |buffer, _| buffer.text()),
+            "abc\ndEf\nghi\njkl\nmnO"
+        );
+
+        action_log
+            .update(cx, |log, cx| {
+                log.buffer_edited(buffer.clone(), vec![edit1, edit2], cx)
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            unreviewed_hunks(&action_log, cx),
+            vec![(
+                buffer.clone(),
+                vec![
+                    HunkStatus {
+                        range: Point::new(1, 0)..Point::new(2, 0),
+                        review_status: ReviewStatus::Unreviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    },
+                    HunkStatus {
+                        range: Point::new(4, 0)..Point::new(4, 3),
+                        review_status: ReviewStatus::Unreviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    }
+                ],
+            )]
+        );
+
+        action_log
+            .update(cx, |log, cx| {
+                log.review_edits_in_range(
+                    buffer.clone(),
+                    Point::new(3, 0)..Point::new(4, 3),
+                    true,
+                    cx,
+                )
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            unreviewed_hunks(&action_log, cx),
+            vec![(
+                buffer.clone(),
+                vec![
+                    HunkStatus {
+                        range: Point::new(1, 0)..Point::new(2, 0),
+                        review_status: ReviewStatus::Unreviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    },
+                    HunkStatus {
+                        range: Point::new(4, 0)..Point::new(4, 3),
+                        review_status: ReviewStatus::Reviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    }
+                ],
+            )]
+        );
+
+        action_log
+            .update(cx, |log, cx| {
+                log.review_edits_in_range(
+                    buffer.clone(),
+                    Point::new(3, 0)..Point::new(4, 3),
+                    false,
+                    cx,
+                )
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            unreviewed_hunks(&action_log, cx),
+            vec![(
+                buffer.clone(),
+                vec![
+                    HunkStatus {
+                        range: Point::new(1, 0)..Point::new(2, 0),
+                        review_status: ReviewStatus::Unreviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    },
+                    HunkStatus {
+                        range: Point::new(4, 0)..Point::new(4, 3),
+                        review_status: ReviewStatus::Unreviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    }
+                ],
+            )]
+        );
+
+        action_log
+            .update(cx, |log, cx| {
+                log.review_edits_in_range(
+                    buffer.clone(),
+                    Point::new(0, 0)..Point::new(4, 3),
+                    true,
+                    cx,
+                )
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            unreviewed_hunks(&action_log, cx),
+            vec![(
+                buffer.clone(),
+                vec![
+                    HunkStatus {
+                        range: Point::new(1, 0)..Point::new(2, 0),
+                        review_status: ReviewStatus::Reviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    },
+                    HunkStatus {
+                        range: Point::new(4, 0)..Point::new(4, 3),
+                        review_status: ReviewStatus::Reviewed,
+                        diff_status: DiffHunkStatusKind::Modified,
+                    }
+                ],
+            )]
+        );
+    }
+
+    #[derive(Debug, Clone, PartialEq, Eq)]
+    struct HunkStatus {
+        range: Range<Point>,
+        review_status: ReviewStatus,
+        diff_status: DiffHunkStatusKind,
+    }
+
+    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
+    enum ReviewStatus {
+        Unreviewed,
+        Reviewed,
+    }
+
+    fn unreviewed_hunks(
+        action_log: &Entity<ActionLog>,
+        cx: &TestAppContext,
+    ) -> Vec<(Entity<Buffer>, Vec<HunkStatus>)> {
+        cx.read(|cx| {
+            action_log
+                .read(cx)
+                .unreviewed_buffers()
+                .into_iter()
+                .map(|(buffer, tracked_buffer)| {
+                    let snapshot = buffer.read(cx).snapshot();
+                    (
+                        buffer,
+                        tracked_buffer
+                            .diff
+                            .read(cx)
+                            .hunks(&snapshot, cx)
+                            .map(|hunk| HunkStatus {
+                                review_status: if hunk.status().has_secondary_hunk() {
+                                    ReviewStatus::Unreviewed
+                                } else {
+                                    ReviewStatus::Reviewed
+                                },
+                                diff_status: hunk.status().kind,
+                                range: hunk.range,
+                            })
+                            .collect(),
+                    )
+                })
+                .collect()
+        })
+    }
+}

crates/assistant_tool/src/assistant_tool.rs 🔗

@@ -1,16 +1,14 @@
+mod action_log;
 mod tool_registry;
 mod tool_working_set;
 
-use std::sync::Arc;
-
 use anyhow::Result;
-use collections::{HashMap, HashSet};
-use gpui::Context;
 use gpui::{App, Entity, SharedString, Task};
-use language::Buffer;
 use language_model::LanguageModelRequestMessage;
 use project::Project;
+use std::sync::Arc;
 
+pub use crate::action_log::*;
 pub use crate::tool_registry::*;
 pub use crate::tool_working_set::*;
 
@@ -54,57 +52,3 @@ pub trait Tool: 'static + Send + Sync {
         cx: &mut App,
     ) -> Task<Result<String>>;
 }
-
-/// Tracks actions performed by tools in a thread
-#[derive(Debug)]
-pub struct ActionLog {
-    /// Buffers that user manually added to the context, and whose content has
-    /// changed since the model last saw them.
-    stale_buffers_in_context: HashSet<Entity<Buffer>>,
-    /// Buffers that we want to notify the model about when they change.
-    tracked_buffers: HashMap<Entity<Buffer>, TrackedBuffer>,
-}
-
-#[derive(Debug, Default)]
-struct TrackedBuffer {
-    version: clock::Global,
-}
-
-impl ActionLog {
-    /// Creates a new, empty action log.
-    pub fn new() -> Self {
-        Self {
-            stale_buffers_in_context: HashSet::default(),
-            tracked_buffers: HashMap::default(),
-        }
-    }
-
-    /// Track a buffer as read, so we can notify the model about user edits.
-    pub fn buffer_read(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
-        let tracked_buffer = self.tracked_buffers.entry(buffer.clone()).or_default();
-        tracked_buffer.version = buffer.read(cx).version();
-    }
-
-    /// Mark a buffer as edited, so we can refresh it in the context
-    pub fn buffer_edited(&mut self, buffers: HashSet<Entity<Buffer>>, cx: &mut Context<Self>) {
-        for buffer in &buffers {
-            let tracked_buffer = self.tracked_buffers.entry(buffer.clone()).or_default();
-            tracked_buffer.version = buffer.read(cx).version();
-        }
-
-        self.stale_buffers_in_context.extend(buffers);
-    }
-
-    /// Iterate over buffers changed since last read or edited by the model
-    pub fn stale_buffers<'a>(&'a self, cx: &'a App) -> impl Iterator<Item = &'a Entity<Buffer>> {
-        self.tracked_buffers
-            .iter()
-            .filter(|(buffer, tracked)| tracked.version != buffer.read(cx).version)
-            .map(|(buffer, _)| buffer)
-    }
-
-    /// Takes and returns the set of buffers pending refresh, clearing internal state.
-    pub fn take_stale_buffers_in_context(&mut self) -> HashSet<Entity<Buffer>> {
-        std::mem::take(&mut self.stale_buffers_in_context)
-    }
-}

crates/assistant_tools/src/edit_files_tool.rs 🔗

@@ -274,7 +274,17 @@ impl EditToolRequest {
                 self.bad_searches.push(invalid_replace);
             }
             DiffResult::Diff(diff) => {
-                let _clock = buffer.update(cx, |buffer, cx| buffer.apply_diff(diff, cx))?;
+                let edit_ids = buffer.update(cx, |buffer, cx| {
+                    buffer.finalize_last_transaction();
+                    buffer.apply_diff(diff, cx);
+                    let transaction = buffer.finalize_last_transaction();
+                    transaction.map_or(Vec::new(), |transaction| transaction.edit_ids.clone())
+                })?;
+                self.action_log
+                    .update(cx, |log, cx| {
+                        log.buffer_edited(buffer.clone(), edit_ids, cx)
+                    })?
+                    .await?;
 
                 write!(&mut self.output, "\n\n{}", source)?;
                 self.changed_buffers.insert(buffer);
@@ -322,10 +332,6 @@ impl EditToolRequest {
                 .await?;
         }
 
-        self.action_log
-            .update(cx, |log, cx| log.buffer_edited(self.changed_buffers, cx))
-            .log_err();
-
         let errors = self.parser.errors();
 
         if errors.is_empty() && self.bad_searches.is_empty() {

crates/editor/src/editor.rs 🔗

@@ -182,8 +182,8 @@ use theme::{
     ThemeColors, ThemeSettings,
 };
 use ui::{
-    h_flex, prelude::*, ButtonSize, ButtonStyle, Disclosure, IconButton, IconName, IconSize, Key,
-    Tooltip,
+    h_flex, prelude::*, ButtonSize, ButtonStyle, Disclosure, IconButton, IconButtonShape, IconName,
+    IconSize, Key, Tooltip,
 };
 use util::{maybe, post_inc, RangeExt, ResultExt, TryFutureExt};
 use workspace::{
@@ -221,6 +221,18 @@ pub(crate) const EDIT_PREDICTION_KEY_CONTEXT: &str = "edit_prediction";
 pub(crate) const EDIT_PREDICTION_CONFLICT_KEY_CONTEXT: &str = "edit_prediction_conflict";
 pub(crate) const MIN_LINE_NUMBER_DIGITS: u32 = 4;
 
+pub type RenderDiffHunkControlsFn = Arc<
+    dyn Fn(
+        u32,
+        &DiffHunkStatus,
+        Range<Anchor>,
+        bool,
+        Pixels,
+        &Entity<Editor>,
+        &mut App,
+    ) -> AnyElement,
+>;
+
 const COLUMNAR_SELECTION_MODIFIERS: Modifiers = Modifiers {
     alt: true,
     shift: true,
@@ -740,6 +752,7 @@ pub struct Editor {
     show_git_blame_inline_delay_task: Option<Task<()>>,
     git_blame_inline_tooltip: Option<WeakEntity<crate::commit_tooltip::CommitTooltip>>,
     git_blame_inline_enabled: bool,
+    render_diff_hunk_controls: RenderDiffHunkControlsFn,
     serialize_dirty_buffers: bool,
     show_selection_menu: Option<bool>,
     blame: Option<Entity<GitBlame>>,
@@ -1474,6 +1487,7 @@ impl Editor {
             show_git_blame_inline_delay_task: None,
             git_blame_inline_tooltip: None,
             git_blame_inline_enabled: ProjectSettings::get_global(cx).git.inline_blame_enabled(),
+            render_diff_hunk_controls: Arc::new(render_diff_hunk_controls),
             serialize_dirty_buffers: ProjectSettings::get_global(cx)
                 .session
                 .restore_unsaved_buffers,
@@ -14471,6 +14485,15 @@ impl Editor {
         self.stage_or_unstage_diff_hunks(stage, ranges, cx);
     }
 
+    pub fn set_render_diff_hunk_controls(
+        &mut self,
+        render_diff_hunk_controls: RenderDiffHunkControlsFn,
+        cx: &mut Context<Self>,
+    ) {
+        self.render_diff_hunk_controls = render_diff_hunk_controls;
+        cx.notify();
+    }
+
     pub fn stage_and_next(
         &mut self,
         _: &::git::StageAndNext,
@@ -19559,3 +19582,187 @@ impl From<Background> for LineHighlight {
         }
     }
 }
+
+fn render_diff_hunk_controls(
+    row: u32,
+    status: &DiffHunkStatus,
+    hunk_range: Range<Anchor>,
+    is_created_file: bool,
+    line_height: Pixels,
+    editor: &Entity<Editor>,
+    cx: &mut App,
+) -> AnyElement {
+    h_flex()
+        .h(line_height)
+        .mr_1()
+        .gap_1()
+        .px_0p5()
+        .pb_1()
+        .border_x_1()
+        .border_b_1()
+        .border_color(cx.theme().colors().border_variant)
+        .rounded_b_lg()
+        .bg(cx.theme().colors().editor_background)
+        .gap_1()
+        .occlude()
+        .shadow_md()
+        .child(if status.has_secondary_hunk() {
+            Button::new(("stage", row as u64), "Stage")
+                .alpha(if status.is_pending() { 0.66 } else { 1.0 })
+                .tooltip({
+                    let focus_handle = editor.focus_handle(cx);
+                    move |window, cx| {
+                        Tooltip::for_action_in(
+                            "Stage Hunk",
+                            &::git::ToggleStaged,
+                            &focus_handle,
+                            window,
+                            cx,
+                        )
+                    }
+                })
+                .on_click({
+                    let editor = editor.clone();
+                    move |_event, _window, cx| {
+                        editor.update(cx, |editor, cx| {
+                            editor.stage_or_unstage_diff_hunks(
+                                true,
+                                vec![hunk_range.start..hunk_range.start],
+                                cx,
+                            );
+                        });
+                    }
+                })
+        } else {
+            Button::new(("unstage", row as u64), "Unstage")
+                .alpha(if status.is_pending() { 0.66 } else { 1.0 })
+                .tooltip({
+                    let focus_handle = editor.focus_handle(cx);
+                    move |window, cx| {
+                        Tooltip::for_action_in(
+                            "Unstage Hunk",
+                            &::git::ToggleStaged,
+                            &focus_handle,
+                            window,
+                            cx,
+                        )
+                    }
+                })
+                .on_click({
+                    let editor = editor.clone();
+                    move |_event, _window, cx| {
+                        editor.update(cx, |editor, cx| {
+                            editor.stage_or_unstage_diff_hunks(
+                                false,
+                                vec![hunk_range.start..hunk_range.start],
+                                cx,
+                            );
+                        });
+                    }
+                })
+        })
+        .child(
+            Button::new("restore", "Restore")
+                .tooltip({
+                    let focus_handle = editor.focus_handle(cx);
+                    move |window, cx| {
+                        Tooltip::for_action_in(
+                            "Restore Hunk",
+                            &::git::Restore,
+                            &focus_handle,
+                            window,
+                            cx,
+                        )
+                    }
+                })
+                .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);
+                        });
+                    }
+                })
+                .disabled(is_created_file),
+        )
+        .when(
+            !editor.read(cx).buffer().read(cx).all_diff_hunks_expanded(),
+            |el| {
+                el.child(
+                    IconButton::new(("next-hunk", row as u64), IconName::ArrowDown)
+                        .shape(IconButtonShape::Square)
+                        .icon_size(IconSize::Small)
+                        // .disabled(!has_multiple_hunks)
+                        .tooltip({
+                            let focus_handle = editor.focus_handle(cx);
+                            move |window, cx| {
+                                Tooltip::for_action_in(
+                                    "Next Hunk",
+                                    &GoToHunk,
+                                    &focus_handle,
+                                    window,
+                                    cx,
+                                )
+                            }
+                        })
+                        .on_click({
+                            let editor = editor.clone();
+                            move |_event, window, cx| {
+                                editor.update(cx, |editor, cx| {
+                                    let snapshot = editor.snapshot(window, cx);
+                                    let position =
+                                        hunk_range.end.to_point(&snapshot.buffer_snapshot);
+                                    editor.go_to_hunk_before_or_after_position(
+                                        &snapshot,
+                                        position,
+                                        Direction::Next,
+                                        window,
+                                        cx,
+                                    );
+                                    editor.expand_selected_diff_hunks(cx);
+                                });
+                            }
+                        }),
+                )
+                .child(
+                    IconButton::new(("prev-hunk", row as u64), IconName::ArrowUp)
+                        .shape(IconButtonShape::Square)
+                        .icon_size(IconSize::Small)
+                        // .disabled(!has_multiple_hunks)
+                        .tooltip({
+                            let focus_handle = editor.focus_handle(cx);
+                            move |window, cx| {
+                                Tooltip::for_action_in(
+                                    "Previous Hunk",
+                                    &GoToPreviousHunk,
+                                    &focus_handle,
+                                    window,
+                                    cx,
+                                )
+                            }
+                        })
+                        .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.go_to_hunk_before_or_after_position(
+                                        &snapshot,
+                                        point,
+                                        Direction::Prev,
+                                        window,
+                                        cx,
+                                    );
+                                    editor.expand_selected_diff_hunks(cx);
+                                });
+                            }
+                        }),
+                )
+            },
+        )
+        .into_any_element()
+}

crates/editor/src/element.rs 🔗

@@ -18,12 +18,12 @@ use crate::{
     scroll::{axis_pair, scroll_amount::ScrollAmount, AxisPair},
     BlockId, ChunkReplacement, CursorShape, CustomBlockId, DisplayDiffHunk, DisplayPoint,
     DisplayRow, DocumentHighlightRead, DocumentHighlightWrite, EditDisplayMode, Editor, EditorMode,
-    EditorSettings, EditorSnapshot, EditorStyle, FocusedBlock, GoToHunk, GoToPreviousHunk,
-    GutterDimensions, HalfPageDown, HalfPageUp, HandleInput, HoveredCursor, InlayHintRefreshReason,
-    InlineCompletion, JumpData, LineDown, LineHighlight, LineUp, OpenExcerpts, PageDown, PageUp,
-    Point, RowExt, RowRangeExt, SelectPhase, SelectedTextHighlight, Selection, SoftWrap,
-    StickyHeaderExcerpt, ToPoint, ToggleFold, COLUMNAR_SELECTION_MODIFIERS, CURSORS_VISIBLE_FOR,
-    FILE_HEADER_HEIGHT, GIT_BLAME_MAX_AUTHOR_CHARS_DISPLAYED, MAX_LINE_LEN, MIN_LINE_NUMBER_DIGITS,
+    EditorSettings, EditorSnapshot, EditorStyle, FocusedBlock, GutterDimensions, HalfPageDown,
+    HalfPageUp, HandleInput, HoveredCursor, InlayHintRefreshReason, InlineCompletion, JumpData,
+    LineDown, LineHighlight, LineUp, OpenExcerpts, PageDown, PageUp, Point, RowExt, RowRangeExt,
+    SelectPhase, SelectedTextHighlight, Selection, SoftWrap, StickyHeaderExcerpt, ToPoint,
+    ToggleFold, COLUMNAR_SELECTION_MODIFIERS, CURSORS_VISIBLE_FOR, FILE_HEADER_HEIGHT,
+    GIT_BLAME_MAX_AUTHOR_CHARS_DISPLAYED, MAX_LINE_LEN, MIN_LINE_NUMBER_DIGITS,
     MULTI_BUFFER_EXCERPT_HEADER_HEIGHT,
 };
 use buffer_diff::{DiffHunkStatus, DiffHunkStatusKind};
@@ -42,7 +42,6 @@ use gpui::{
     ParentElement, Pixels, ScrollDelta, ScrollWheelEvent, ShapedLine, SharedString, Size,
     StatefulInteractiveElement, Style, Styled, Subscription, TextRun, TextStyleRefinement, Window,
 };
-use inline_completion::Direction;
 use itertools::Itertools;
 use language::{
     language_settings::{
@@ -75,10 +74,7 @@ use std::{
 use sum_tree::Bias;
 use text::BufferId;
 use theme::{ActiveTheme, Appearance, BufferLineHeight, PlayerColor};
-use ui::{
-    h_flex, prelude::*, ButtonLike, ContextMenu, IconButtonShape, KeyBinding, Tooltip,
-    POPOVER_Y_PADDING,
-};
+use ui::{h_flex, prelude::*, ButtonLike, ContextMenu, KeyBinding, Tooltip, POPOVER_Y_PADDING};
 use unicode_segmentation::UnicodeSegmentation;
 use util::{debug_panic, RangeExt, ResultExt};
 use workspace::{item::Item, notifications::NotifyTaskExt};
@@ -3986,6 +3982,7 @@ impl EditorElement {
         window: &mut Window,
         cx: &mut App,
     ) -> Vec<AnyElement> {
+        let render_diff_hunk_controls = editor.read(cx).render_diff_hunk_controls.clone();
         let point_for_position = position_map.point_for_position(window.mouse_position());
 
         let mut controls = vec![];
@@ -4028,7 +4025,7 @@ impl EditorElement {
                         + text_hitbox.bounds.top()
                         - scroll_pixel_position.y;
 
-                    let mut element = diff_hunk_controls(
+                    let mut element = render_diff_hunk_controls(
                         display_row_range.start.0,
                         status,
                         multi_buffer_range.clone(),
@@ -8931,187 +8928,3 @@ mod tests {
             .collect()
     }
 }
-
-fn diff_hunk_controls(
-    row: u32,
-    status: &DiffHunkStatus,
-    hunk_range: Range<Anchor>,
-    is_created_file: bool,
-    line_height: Pixels,
-    editor: &Entity<Editor>,
-    cx: &mut App,
-) -> AnyElement {
-    h_flex()
-        .h(line_height)
-        .mr_1()
-        .gap_1()
-        .px_0p5()
-        .pb_1()
-        .border_x_1()
-        .border_b_1()
-        .border_color(cx.theme().colors().border_variant)
-        .rounded_b_lg()
-        .bg(cx.theme().colors().editor_background)
-        .gap_1()
-        .occlude()
-        .shadow_md()
-        .child(if status.has_secondary_hunk() {
-            Button::new(("stage", row as u64), "Stage")
-                .alpha(if status.is_pending() { 0.66 } else { 1.0 })
-                .tooltip({
-                    let focus_handle = editor.focus_handle(cx);
-                    move |window, cx| {
-                        Tooltip::for_action_in(
-                            "Stage Hunk",
-                            &::git::ToggleStaged,
-                            &focus_handle,
-                            window,
-                            cx,
-                        )
-                    }
-                })
-                .on_click({
-                    let editor = editor.clone();
-                    move |_event, _window, cx| {
-                        editor.update(cx, |editor, cx| {
-                            editor.stage_or_unstage_diff_hunks(
-                                true,
-                                vec![hunk_range.start..hunk_range.start],
-                                cx,
-                            );
-                        });
-                    }
-                })
-        } else {
-            Button::new(("unstage", row as u64), "Unstage")
-                .alpha(if status.is_pending() { 0.66 } else { 1.0 })
-                .tooltip({
-                    let focus_handle = editor.focus_handle(cx);
-                    move |window, cx| {
-                        Tooltip::for_action_in(
-                            "Unstage Hunk",
-                            &::git::ToggleStaged,
-                            &focus_handle,
-                            window,
-                            cx,
-                        )
-                    }
-                })
-                .on_click({
-                    let editor = editor.clone();
-                    move |_event, _window, cx| {
-                        editor.update(cx, |editor, cx| {
-                            editor.stage_or_unstage_diff_hunks(
-                                false,
-                                vec![hunk_range.start..hunk_range.start],
-                                cx,
-                            );
-                        });
-                    }
-                })
-        })
-        .child(
-            Button::new("restore", "Restore")
-                .tooltip({
-                    let focus_handle = editor.focus_handle(cx);
-                    move |window, cx| {
-                        Tooltip::for_action_in(
-                            "Restore Hunk",
-                            &::git::Restore,
-                            &focus_handle,
-                            window,
-                            cx,
-                        )
-                    }
-                })
-                .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);
-                        });
-                    }
-                })
-                .disabled(is_created_file),
-        )
-        .when(
-            !editor.read(cx).buffer().read(cx).all_diff_hunks_expanded(),
-            |el| {
-                el.child(
-                    IconButton::new(("next-hunk", row as u64), IconName::ArrowDown)
-                        .shape(IconButtonShape::Square)
-                        .icon_size(IconSize::Small)
-                        // .disabled(!has_multiple_hunks)
-                        .tooltip({
-                            let focus_handle = editor.focus_handle(cx);
-                            move |window, cx| {
-                                Tooltip::for_action_in(
-                                    "Next Hunk",
-                                    &GoToHunk,
-                                    &focus_handle,
-                                    window,
-                                    cx,
-                                )
-                            }
-                        })
-                        .on_click({
-                            let editor = editor.clone();
-                            move |_event, window, cx| {
-                                editor.update(cx, |editor, cx| {
-                                    let snapshot = editor.snapshot(window, cx);
-                                    let position =
-                                        hunk_range.end.to_point(&snapshot.buffer_snapshot);
-                                    editor.go_to_hunk_before_or_after_position(
-                                        &snapshot,
-                                        position,
-                                        Direction::Next,
-                                        window,
-                                        cx,
-                                    );
-                                    editor.expand_selected_diff_hunks(cx);
-                                });
-                            }
-                        }),
-                )
-                .child(
-                    IconButton::new(("prev-hunk", row as u64), IconName::ArrowUp)
-                        .shape(IconButtonShape::Square)
-                        .icon_size(IconSize::Small)
-                        // .disabled(!has_multiple_hunks)
-                        .tooltip({
-                            let focus_handle = editor.focus_handle(cx);
-                            move |window, cx| {
-                                Tooltip::for_action_in(
-                                    "Previous Hunk",
-                                    &GoToPreviousHunk,
-                                    &focus_handle,
-                                    window,
-                                    cx,
-                                )
-                            }
-                        })
-                        .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.go_to_hunk_before_or_after_position(
-                                        &snapshot,
-                                        point,
-                                        Direction::Prev,
-                                        window,
-                                        cx,
-                                    );
-                                    editor.expand_selected_diff_hunks(cx);
-                                });
-                            }
-                        }),
-                )
-            },
-        )
-        .into_any_element()
-}

crates/text/src/text.rs 🔗

@@ -1498,9 +1498,9 @@ impl Buffer {
             .flat_map(|transaction| self.edited_ranges_for_transaction(transaction))
     }
 
-    pub fn edited_ranges_for_transaction<'a, D>(
+    pub fn edited_ranges_for_edit_ids<'a, D>(
         &'a self,
-        transaction: &'a Transaction,
+        edit_ids: impl IntoIterator<Item = &'a clock::Lamport>,
     ) -> impl 'a + Iterator<Item = Range<D>>
     where
         D: TextDimension,
@@ -1508,7 +1508,7 @@ impl Buffer {
         // get fragment ranges
         let mut cursor = self.fragments.cursor::<(Option<&Locator>, usize)>(&None);
         let offset_ranges = self
-            .fragment_ids_for_edits(transaction.edit_ids.iter())
+            .fragment_ids_for_edits(edit_ids.into_iter())
             .into_iter()
             .filter_map(move |fragment_id| {
                 cursor.seek_forward(&Some(fragment_id), Bias::Left, &None);
@@ -1547,6 +1547,16 @@ impl Buffer {
         })
     }
 
+    pub fn edited_ranges_for_transaction<'a, D>(
+        &'a self,
+        transaction: &'a Transaction,
+    ) -> impl 'a + Iterator<Item = Range<D>>
+    where
+        D: TextDimension,
+    {
+        self.edited_ranges_for_edit_ids(&transaction.edit_ids)
+    }
+
     pub fn subscribe(&mut self) -> Subscription {
         self.subscriptions.subscribe()
     }