Introduce "Keep All" and "Reject All" buttons when reviewing assistant edits (#27724)

Antonio Scandurra created 

Release Notes:

- N/A

Change summary

crates/assistant2/src/assistant.rs      |   6 
crates/assistant2/src/assistant_diff.rs | 170 ++++++++++++++++++++++----
crates/assistant2/src/thread.rs         |   7 +
crates/assistant_tool/src/action_log.rs |  22 +++
crates/editor/src/editor.rs             |   2 
crates/editor/src/element.rs            |   1 
crates/zed/src/zed.rs                   |   3 
7 files changed, 182 insertions(+), 29 deletions(-)

Detailed changes

crates/assistant2/src/assistant.rs 🔗 

@@ -40,7 +40,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;
+pub use assistant_diff::{AssistantDiff, AssistantDiffToolbar};
 
 actions!(
     assistant2,
@@ -66,7 +66,9 @@ actions!(
         AcceptSuggestedContext,
         OpenActiveThreadAsMarkdown,
         ToggleKeep,
-        Reject
+        Reject,
+        RejectAll,
+        KeepAll
     ]
 );

crates/assistant2/src/assistant_diff.rs 🔗 

@@ -7,10 +7,10 @@ use editor::{
     Direction, Editor, EditorEvent, MultiBuffer, ToPoint,
 };
 use gpui::{
-    prelude::*, AnyElement, AnyView, App, Entity, EventEmitter, FocusHandle, Focusable,
+    prelude::*, Action, AnyElement, AnyView, App, Entity, EventEmitter, FocusHandle, Focusable,
     SharedString, Subscription, Task, WeakEntity, Window,
 };
-use language::{Capability, DiskState, OffsetRangeExt};
+use language::{Capability, DiskState, OffsetRangeExt, Point};
 use multi_buffer::PathKey;
 use project::{Project, ProjectPath};
 use std::{
@@ -18,11 +18,12 @@ use std::{
     ops::Range,
     sync::Arc,
 };
-use ui::{prelude::*, IconButtonShape, Tooltip};
+use ui::{prelude::*, IconButtonShape, KeyBinding, Tooltip};
 use workspace::{
     item::{BreadcrumbText, ItemEvent, TabContentParams},
     searchable::SearchableItemHandle,
-    Item, ItemHandle, ItemNavHistory, ToolbarItemLocation, Workspace,
+    Item, ItemHandle, ItemNavHistory, ToolbarItemEvent, ToolbarItemLocation, ToolbarItemView,
+    Workspace,
 };
 
 pub struct AssistantDiff {
@@ -78,6 +79,7 @@ impl AssistantDiff {
                   is_created_file,
                   line_height,
                   _editor: &Entity<Editor>,
+                  window: &mut Window,
                   cx: &mut App| {
                 render_diff_hunk_controls(
                     row,
@@ -86,6 +88,7 @@ impl AssistantDiff {
                     is_created_file,
                     line_height,
                     &assistant_diff,
+                    window,
                     cx,
                 )
             }
@@ -253,6 +256,18 @@ impl AssistantDiff {
         })
     }
 
+    fn reject_all(&mut self, _: &crate::RejectAll, window: &mut Window, cx: &mut Context<Self>) {
+        self.editor.update(cx, |editor, cx| {
+            let max_point = editor.buffer().read(cx).read(cx).max_point();
+            editor.restore_hunks_in_ranges(vec![Point::zero()..max_point], window, cx)
+        })
+    }
+
+    fn keep_all(&mut self, _: &crate::KeepAll, _window: &mut Window, cx: &mut Context<Self>) {
+        self.thread
+            .update(cx, |thread, cx| thread.keep_all_edits(cx));
+    }
+
     fn review_diff_hunks(
         &mut self,
         hunk_ranges: Vec<Range<editor::Anchor>>,
@@ -465,6 +480,8 @@ impl Render for AssistantDiff {
             })
             .on_action(cx.listener(Self::toggle_keep))
             .on_action(cx.listener(Self::reject))
+            .on_action(cx.listener(Self::reject_all))
+            .on_action(cx.listener(Self::keep_all))
             .bg(cx.theme().colors().editor_background)
             .flex()
             .items_center()
@@ -482,6 +499,7 @@ fn render_diff_hunk_controls(
     is_created_file: bool,
     line_height: Pixels,
     assistant_diff: &Entity<AssistantDiff>,
+    window: &mut Window,
     cx: &mut App,
 ) -> AnyElement {
     let editor = assistant_diff.read(cx).editor.clone();
@@ -501,13 +519,19 @@ fn render_diff_hunk_controls(
         .shadow_md()
         .children(if status.has_secondary_hunk() {
             vec![
-                Button::new(("keep", row as u64), "Keep")
+                Button::new("reject", "Reject")
+                    .key_binding(KeyBinding::for_action_in(
+                        &crate::Reject,
+                        &editor.read(cx).focus_handle(cx),
+                        window,
+                        cx,
+                    ))
                     .tooltip({
                         let focus_handle = editor.focus_handle(cx);
                         move |window, cx| {
                             Tooltip::for_action_in(
-                                "Keep Hunk",
-                                &crate::ToggleKeep,
+                                "Reject Hunk",
+                                &crate::Reject,
                                 &focus_handle,
                                 window,
                                 cx,
@@ -515,24 +539,29 @@ fn render_diff_hunk_controls(
                         }
                     })
                     .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,
-                                );
+                        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("reject", "Reject")
+                    })
+                    .disabled(is_created_file),
+                Button::new(("keep", row as u64), "Keep")
+                    .key_binding(KeyBinding::for_action_in(
+                        &crate::ToggleKeep,
+                        &editor.read(cx).focus_handle(cx),
+                        window,
+                        cx,
+                    ))
                     .tooltip({
                         let focus_handle = editor.focus_handle(cx);
                         move |window, cx| {
                             Tooltip::for_action_in(
-                                "Reject Hunk",
-                                &crate::Reject,
+                                "Keep Hunk",
+                                &crate::ToggleKeep,
                                 &focus_handle,
                                 window,
                                 cx,
@@ -540,16 +569,17 @@ fn render_diff_hunk_controls(
                         }
                     })
                     .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);
+                        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,
+                                );
                             });
                         }
-                    })
-                    .disabled(is_created_file),
+                    }),
             ]
         } else {
             vec![Button::new(("review", row as u64), "Review")
@@ -663,3 +693,89 @@ impl editor::Addon for AssistantDiffAddon {
         key_context.add("assistant_diff");
     }
 }
+
+pub struct AssistantDiffToolbar {
+    assistant_diff: Option<WeakEntity<AssistantDiff>>,
+    _workspace: WeakEntity<Workspace>,
+}
+
+impl AssistantDiffToolbar {
+    pub fn new(workspace: &Workspace, _: &mut Context<Self>) -> Self {
+        Self {
+            assistant_diff: None,
+            _workspace: workspace.weak_handle(),
+        }
+    }
+
+    fn assistant_diff(&self, _: &App) -> Option<Entity<AssistantDiff>> {
+        self.assistant_diff.as_ref()?.upgrade()
+    }
+
+    fn dispatch_action(&self, action: &dyn Action, window: &mut Window, cx: &mut Context<Self>) {
+        if let Some(assistant_diff) = self.assistant_diff(cx) {
+            assistant_diff.focus_handle(cx).focus(window);
+        }
+        let action = action.boxed_clone();
+        cx.defer(move |cx| {
+            cx.dispatch_action(action.as_ref());
+        })
+    }
+}
+
+impl EventEmitter<ToolbarItemEvent> for AssistantDiffToolbar {}
+
+impl ToolbarItemView for AssistantDiffToolbar {
+    fn set_active_pane_item(
+        &mut self,
+        active_pane_item: Option<&dyn ItemHandle>,
+        _: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> ToolbarItemLocation {
+        self.assistant_diff = active_pane_item
+            .and_then(|item| item.act_as::<AssistantDiff>(cx))
+            .map(|entity| entity.downgrade());
+        if self.assistant_diff.is_some() {
+            ToolbarItemLocation::PrimaryRight
+        } else {
+            ToolbarItemLocation::Hidden
+        }
+    }
+
+    fn pane_focus_update(
+        &mut self,
+        _pane_focused: bool,
+        _window: &mut Window,
+        _cx: &mut Context<Self>,
+    ) {
+    }
+}
+
+impl Render for AssistantDiffToolbar {
+    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        if self.assistant_diff(cx).is_none() {
+            return div();
+        }
+
+        h_group_xl()
+            .my_neg_1()
+            .items_center()
+            .py_1()
+            .pl_2()
+            .pr_1()
+            .flex_wrap()
+            .justify_between()
+            .child(
+                h_group_sm()
+                    .child(
+                        Button::new("reject-all", "Reject All").on_click(cx.listener(
+                            |this, _, window, cx| {
+                                this.dispatch_action(&crate::RejectAll, window, cx)
+                            },
+                        )),
+                    )
+                    .child(Button::new("keep-all", "Keep All").on_click(cx.listener(
+                        |this, _, window, cx| this.dispatch_action(&crate::KeepAll, window, cx),
+                    ))),
+            )
+    }
+}

crates/assistant2/src/thread.rs 🔗 

@@ -1542,6 +1542,13 @@ impl Thread {
         });
     }
 
+    /// 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());
+    }
+
     pub fn action_log(&self) -> &Entity<ActionLog> {
         &self.action_log
     }

crates/assistant_tool/src/action_log.rs 🔗 

@@ -353,6 +353,28 @@ impl ActionLog {
         tracked_buffer.schedule_diff_update();
     }
 
+    /// Keep all edits across all buffers.
+    /// This is a more performant alternative to calling review_edits_in_range for each buffer.
+    pub fn keep_all_edits(&mut self) {
+        // Process all tracked buffers
+        for (_, tracked_buffer) in self.tracked_buffers.iter_mut() {
+            match &mut tracked_buffer.change {
+                Change::Deleted { reviewed, .. } => {
+                    *reviewed = true;
+                }
+                Change::Edited {
+                    unreviewed_edit_ids,
+                    accepted_edit_ids,
+                    ..
+                } => {
+                    accepted_edit_ids.extend(unreviewed_edit_ids.drain());
+                }
+            }
+
+            tracked_buffer.schedule_diff_update();
+        }
+    }
+
     /// Returns the set of buffers that contain changes that haven't been reviewed by the user.
     pub fn changed_buffers(&self, cx: &App) -> BTreeMap<Entity<Buffer>, ChangedBuffer> {
         self.tracked_buffers

crates/editor/src/editor.rs 🔗 

@@ -228,6 +228,7 @@ pub type RenderDiffHunkControlsFn = Arc<
         bool,
         Pixels,
         &Entity<Editor>,
+        &mut Window,
         &mut App,
     ) -> AnyElement,
 >;
@@ -20011,6 +20012,7 @@ fn render_diff_hunk_controls(
     is_created_file: bool,
     line_height: Pixels,
     editor: &Entity<Editor>,
+    _window: &mut Window,
     cx: &mut App,
 ) -> AnyElement {
     h_flex()

crates/zed/src/zed.rs 🔗 

@@ -11,6 +11,7 @@ pub(crate) mod windows_only_instance;
 use anyhow::Context as _;
 pub use app_menus::*;
 use assets::Assets;
+use assistant2::AssistantDiffToolbar;
 use assistant_context_editor::AssistantPanelDelegate;
 use breadcrumbs::Breadcrumbs;
 use client::{zed_urls, ZED_URL_SCHEME};
@@ -939,6 +940,8 @@ fn initialize_pane(
             toolbar.add_item(migration_banner, window, cx);
             let project_diff_toolbar = cx.new(|cx| ProjectDiffToolbar::new(workspace, cx));
             toolbar.add_item(project_diff_toolbar, window, cx);
+            let assistant_diff_toolbar = cx.new(|cx| AssistantDiffToolbar::new(workspace, cx));
+            toolbar.add_item(assistant_diff_toolbar, window, cx);
         })
     });
 }