From 52fc14f36d92cfb3b9af0726ceeba01ad58236ba Mon Sep 17 00:00:00 2001 From: Danilo Leal <67129314+danilo-leal@users.noreply.github.com> Date: Wed, 18 Feb 2026 15:49:36 -0300 Subject: [PATCH] agent_ui: Add the ability to undo "reject all" (#48462) Closes https://github.com/zed-industries/zed/issues/39294 This PR implements the ability to undo the "reject all" action from the agent panel (and the other places where this button is surfaced). Effectively, this allows to recuperate the edits an agent has introduced in case you either accidentally clicked the button or regretted the decision to reject all. https://github.com/user-attachments/assets/6f048b95-dd0a-4a45-8b4f-cd8f99d45cb3 Release Notes: - Agent: Introduced the ability to undo the "reject all" action from agent-made changes. --------- Co-authored-by: Bennet Bo Fenner --- assets/keymaps/default-linux.json | 1 + assets/keymaps/default-macos.json | 1 + assets/keymaps/default-windows.json | 1 + crates/action_log/src/action_log.rs | 291 ++++++++++++++++-- crates/agent_ui/src/acp/thread_view.rs | 2 +- .../src/acp/thread_view/active_thread.rs | 25 ++ crates/agent_ui/src/agent_diff.rs | 89 +++++- crates/agent_ui/src/agent_ui.rs | 2 + crates/agent_ui/src/ui.rs | 2 + crates/agent_ui/src/ui/undo_reject_toast.rs | 25 ++ 10 files changed, 397 insertions(+), 42 deletions(-) create mode 100644 crates/agent_ui/src/ui/undo_reject_toast.rs diff --git a/assets/keymaps/default-linux.json b/assets/keymaps/default-linux.json index 1be1de0230a74c29e96f44d54a4405dfa4c0b29d..e569b4b5c9f1f71624942b547d2e719cc52fe86a 100644 --- a/assets/keymaps/default-linux.json +++ b/assets/keymaps/default-linux.json @@ -315,6 +315,7 @@ "ctrl-shift-d": "git::Diff", "shift-alt-y": "agent::KeepAll", "shift-alt-z": "agent::RejectAll", + "shift-alt-u": "agent::UndoLastReject", "ctrl-enter": "agent::ChatWithFollow", "ctrl-shift-enter": "agent::SendImmediately", "ctrl-shift-alt-enter": "agent::SendNextQueuedMessage", diff --git a/assets/keymaps/default-macos.json b/assets/keymaps/default-macos.json index 2d121bce142c109af36480e6d11a455ce7fb848a..5a478542d6b0ccd5cb358b0ebb684c1c5b0c794c 100644 --- a/assets/keymaps/default-macos.json +++ b/assets/keymaps/default-macos.json @@ -361,6 +361,7 @@ "shift-ctrl-d": "git::Diff", "shift-alt-y": "agent::KeepAll", "shift-alt-z": "agent::RejectAll", + "shift-alt-u": "agent::UndoLastReject", "cmd-enter": "agent::ChatWithFollow", "cmd-shift-enter": "agent::SendImmediately", "cmd-shift-alt-enter": "agent::SendNextQueuedMessage", diff --git a/assets/keymaps/default-windows.json b/assets/keymaps/default-windows.json index 273e733b0cdef263ae5d2ee5d4004ac312f49f4b..52dfce74efd2a0240734c44472bec19b81978f44 100644 --- a/assets/keymaps/default-windows.json +++ b/assets/keymaps/default-windows.json @@ -317,6 +317,7 @@ "ctrl-shift-d": "git::Diff", "shift-alt-y": "agent::KeepAll", "shift-alt-z": "agent::RejectAll", + "shift-alt-u": "agent::UndoLastReject", "ctrl-enter": "agent::ChatWithFollow", "ctrl-shift-enter": "agent::SendImmediately", "ctrl-shift-alt-enter": "agent::SendNextQueuedMessage", diff --git a/crates/action_log/src/action_log.rs b/crates/action_log/src/action_log.rs index 4f1f3da02704acb142570780957a8da06b9eccde..1157d8d6f881ecb33df8104dd4be04bd9d846b5e 100644 --- a/crates/action_log/src/action_log.rs +++ b/crates/action_log/src/action_log.rs @@ -6,18 +6,50 @@ use futures::{FutureExt, StreamExt, channel::mpsc}; use gpui::{ App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity, }; -use language::{Anchor, Buffer, BufferEvent, Point, ToPoint}; +use language::{Anchor, Buffer, BufferEvent, Point, ToOffset, ToPoint}; use project::{Project, ProjectItem, lsp_store::OpenLspBufferHandle}; use std::{cmp, ops::Range, sync::Arc}; use text::{Edit, Patch, Rope}; use util::{RangeExt, ResultExt as _}; +/// Stores undo information for a single buffer's rejected edits +#[derive(Clone)] +pub struct PerBufferUndo { + pub buffer: WeakEntity, + pub edits_to_restore: Vec<(Range, String)>, + pub status: UndoBufferStatus, +} + +/// Tracks the buffer status for undo purposes +#[derive(Clone, Debug)] +pub enum UndoBufferStatus { + Modified, + /// Buffer was created by the agent. + /// - `had_existing_content: true` - Agent overwrote an existing file. On reject, the + /// original content was restored. Undo is supported: we restore the agent's content. + /// - `had_existing_content: false` - Agent created a new file that didn't exist before. + /// On reject, the file was deleted. Undo is NOT currently supported (would require + /// recreating the file). Future TODO. + Created { + had_existing_content: bool, + }, +} + +/// Stores undo information for the most recent reject operation +#[derive(Clone)] +pub struct LastRejectUndo { + /// Per-buffer undo information + pub buffers: Vec, +} + /// Tracks actions performed by tools in a thread pub struct ActionLog { /// Buffers that we want to notify the model about when they change. tracked_buffers: BTreeMap, TrackedBuffer>, /// The project this action log is associated with project: Entity, + /// Stores undo information for the most recent reject operation + last_reject_undo: Option, } impl ActionLog { @@ -26,6 +58,7 @@ impl ActionLog { Self { tracked_buffers: BTreeMap::default(), project, + last_reject_undo: None, } } @@ -577,17 +610,21 @@ impl ActionLog { buffer_ranges: Vec>, telemetry: Option, cx: &mut Context, - ) -> Task> { + ) -> (Task>, Option) { let Some(tracked_buffer) = self.tracked_buffers.get_mut(&buffer) else { - return Task::ready(Ok(())); + return (Task::ready(Ok(())), None); }; let mut metrics = ActionLogMetrics::for_buffer(buffer.read(cx)); + let mut undo_info: Option = None; let task = match &tracked_buffer.status { TrackedBufferStatus::Created { existing_file_content, } => { let task = if let Some(existing_file_content) = existing_file_content { + // Capture the agent's content before restoring existing file content + let agent_content = buffer.read(cx).text(); + buffer.update(cx, |buffer, cx| { buffer.start_transaction(); buffer.set_text("", cx); @@ -596,6 +633,15 @@ impl ActionLog { } buffer.end_transaction(cx); }); + + undo_info = Some(PerBufferUndo { + buffer: buffer.downgrade(), + edits_to_restore: vec![(Anchor::MIN..Anchor::MAX, agent_content)], + status: UndoBufferStatus::Created { + had_existing_content: true, + }, + }); + self.project .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx)) } else { @@ -654,7 +700,7 @@ impl ActionLog { save } TrackedBufferStatus::Modified => { - buffer.update(cx, |buffer, cx| { + let edits_to_restore = buffer.update(cx, |buffer, cx| { let mut buffer_row_ranges = buffer_ranges .into_iter() .map(|range| { @@ -663,6 +709,7 @@ impl ActionLog { .peekable(); let mut edits_to_revert = Vec::new(); + let mut edits_for_undo = Vec::new(); for edit in tracked_buffer.unreviewed_edits.edits() { let new_range = tracked_buffer .snapshot @@ -699,12 +746,30 @@ impl ActionLog { .diff_base .chunks_in_range(old_range) .collect::(); + + // Capture the agent's text before we revert it (for undo) + let new_range_offset = + new_range.start.to_offset(buffer)..new_range.end.to_offset(buffer); + let agent_text = + buffer.text_for_range(new_range_offset).collect::(); + edits_for_undo.push((new_range.clone(), agent_text)); + edits_to_revert.push((new_range, old_text)); } } buffer.edit(edits_to_revert, None, cx); + edits_for_undo }); + + if !edits_to_restore.is_empty() { + undo_info = Some(PerBufferUndo { + buffer: buffer.downgrade(), + edits_to_restore, + status: UndoBufferStatus::Modified, + }); + } + self.project .update(cx, |project, cx| project.save_buffer(buffer, cx)) } @@ -712,7 +777,7 @@ impl ActionLog { if let Some(telemetry) = telemetry { telemetry_report_rejected_edits(&telemetry, metrics); } - task + (task, undo_info) } pub fn keep_all_edits( @@ -748,16 +813,34 @@ impl ActionLog { telemetry: Option, cx: &mut Context, ) -> Task<()> { - let futures = self.changed_buffers(cx).into_keys().map(|buffer| { + // Clear any previous undo state before starting a new reject operation + self.last_reject_undo = None; + + let mut undo_buffers = Vec::new(); + let mut futures = Vec::new(); + + for buffer in self.changed_buffers(cx).into_keys() { let buffer_ranges = vec![Anchor::min_max_range_for_buffer( buffer.read(cx).remote_id(), )]; - let reject = self.reject_edits_in_ranges(buffer, buffer_ranges, telemetry.clone(), cx); + let (reject_task, undo_info) = + self.reject_edits_in_ranges(buffer, buffer_ranges, telemetry.clone(), cx); - async move { - reject.await.log_err(); + if let Some(undo) = undo_info { + undo_buffers.push(undo); } - }); + + futures.push(async move { + reject_task.await.log_err(); + }); + } + + // Store the undo information if we have any + if !undo_buffers.is_empty() { + self.last_reject_undo = Some(LastRejectUndo { + buffers: undo_buffers, + }); + } let task = futures::future::join_all(futures); cx.background_spawn(async move { @@ -765,6 +848,63 @@ impl ActionLog { }) } + pub fn has_pending_undo(&self) -> bool { + self.last_reject_undo.is_some() + } + + pub fn set_last_reject_undo(&mut self, undo: LastRejectUndo) { + self.last_reject_undo = Some(undo); + } + + /// Undoes the most recent reject operation, restoring the rejected agent changes. + /// This is a best-effort operation: if buffers have been closed or modified externally, + /// those buffers will be skipped. + pub fn undo_last_reject(&mut self, cx: &mut Context) -> Task<()> { + let Some(undo) = self.last_reject_undo.take() else { + return Task::ready(()); + }; + + let mut save_tasks = Vec::with_capacity(undo.buffers.len()); + + for per_buffer_undo in undo.buffers { + // Skip if the buffer entity has been deallocated + let Some(buffer) = per_buffer_undo.buffer.upgrade() else { + continue; + }; + + buffer.update(cx, |buffer, cx| { + let mut valid_edits = Vec::new(); + + for (anchor_range, text_to_restore) in per_buffer_undo.edits_to_restore { + if anchor_range.start.buffer_id == Some(buffer.remote_id()) + && anchor_range.end.buffer_id == Some(buffer.remote_id()) + { + valid_edits.push((anchor_range, text_to_restore)); + } + } + + if !valid_edits.is_empty() { + buffer.edit(valid_edits, None, cx); + } + }); + + if !self.tracked_buffers.contains_key(&buffer) { + self.buffer_edited(buffer.clone(), cx); + } + + let save = self + .project + .update(cx, |project, cx| project.save_buffer(buffer, cx)); + save_tasks.push(save); + } + + cx.notify(); + + cx.background_spawn(async move { + futures::future::join_all(save_tasks).await; + }) + } + /// Returns the set of buffers that contain edits that haven't been reviewed by the user. pub fn changed_buffers(&self, cx: &App) -> BTreeMap, Entity> { self.tracked_buffers @@ -1417,7 +1557,8 @@ mod tests { action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges(buffer.clone(), vec![2..5], None, cx) + let (task, _) = log.reject_edits_in_ranges(buffer.clone(), vec![2..5], None, cx); + task }) .await .unwrap(); @@ -1497,7 +1638,8 @@ mod tests { action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges(buffer.clone(), vec![2..5], None, cx) + let (task, _) = log.reject_edits_in_ranges(buffer.clone(), vec![2..5], None, cx); + task }) .await .unwrap(); @@ -1677,12 +1819,13 @@ mod tests { // If the rejected range doesn't overlap with any hunk, we ignore it. action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges( + let (task, _) = log.reject_edits_in_ranges( buffer.clone(), vec![Point::new(4, 0)..Point::new(4, 0)], None, cx, - ) + ); + task }) .await .unwrap(); @@ -1712,12 +1855,13 @@ mod tests { action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges( + let (task, _) = log.reject_edits_in_ranges( buffer.clone(), vec![Point::new(0, 0)..Point::new(1, 0)], None, cx, - ) + ); + task }) .await .unwrap(); @@ -1740,12 +1884,13 @@ mod tests { action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges( + let (task, _) = log.reject_edits_in_ranges( buffer.clone(), vec![Point::new(4, 0)..Point::new(4, 0)], None, cx, - ) + ); + task }) .await .unwrap(); @@ -1818,8 +1963,9 @@ mod tests { let range_2 = buffer.read(cx).anchor_before(Point::new(5, 0)) ..buffer.read(cx).anchor_before(Point::new(5, 3)); - log.reject_edits_in_ranges(buffer.clone(), vec![range_1, range_2], None, cx) - .detach(); + let (task, _) = + log.reject_edits_in_ranges(buffer.clone(), vec![range_1, range_2], None, cx); + task.detach(); assert_eq!( buffer.read_with(cx, |buffer, _| buffer.text()), "abc\ndef\nghi\njkl\nmno" @@ -1876,12 +2022,13 @@ mod tests { action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges( + let (task, _) = log.reject_edits_in_ranges( buffer.clone(), vec![Point::new(0, 0)..Point::new(0, 0)], None, cx, - ) + ); + task }) .await .unwrap(); @@ -1932,12 +2079,13 @@ mod tests { action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges( + let (task, _) = log.reject_edits_in_ranges( buffer.clone(), vec![Point::new(0, 0)..Point::new(0, 11)], None, cx, - ) + ); + task }) .await .unwrap(); @@ -1995,12 +2143,13 @@ mod tests { // Reject all action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges( + let (task, _) = log.reject_edits_in_ranges( buffer.clone(), vec![Point::new(0, 0)..Point::new(100, 0)], None, cx, - ) + ); + task }) .await .unwrap(); @@ -2068,14 +2217,15 @@ mod tests { // User rejects the hunk action_log .update(cx, |log, cx| { - log.reject_edits_in_ranges( + let (task, _) = log.reject_edits_in_ranges( buffer.clone(), vec![Anchor::min_max_range_for_buffer( buffer.read(cx).remote_id(), )], None, cx, - ) + ); + task }) .await .unwrap(); @@ -2186,7 +2336,9 @@ mod tests { .update(cx, |log, cx| { let range = buffer.read(cx).random_byte_range(0, &mut rng); log::info!("rejecting edits in range {:?}", range); - log.reject_edits_in_ranges(buffer.clone(), vec![range], None, cx) + let (task, _) = + log.reject_edits_in_ranges(buffer.clone(), vec![range], None, cx); + task }) .await .unwrap(); @@ -2403,7 +2555,86 @@ mod tests { assert_eq!(unreviewed_hunks(&action_log, cx), vec![]); } - #[derive(Debug, Clone, PartialEq, Eq)] + #[gpui::test] + async fn test_undo_last_reject(cx: &mut TestAppContext) { + init_test(cx); + + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + path!("/dir"), + json!({ + "file1": "abc\ndef\nghi" + }), + ) + .await; + let project = Project::test(fs.clone(), [path!("/dir").as_ref()], cx).await; + let action_log = cx.new(|_| ActionLog::new(project.clone())); + let file_path = project + .read_with(cx, |project, cx| project.find_project_path("dir/file1", cx)) + .unwrap(); + + let buffer = project + .update(cx, |project, cx| project.open_buffer(file_path, cx)) + .await + .unwrap(); + + // Track the buffer and make an agent edit + cx.update(|cx| { + action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx)); + buffer.update(cx, |buffer, cx| { + buffer + .edit( + [(Point::new(1, 0)..Point::new(1, 3), "AGENT_EDIT")], + None, + cx, + ) + .unwrap() + }); + action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx)); + }); + cx.run_until_parked(); + + // Verify the agent edit is there + assert_eq!( + buffer.read_with(cx, |buffer, _| buffer.text()), + "abc\nAGENT_EDIT\nghi" + ); + assert!(!unreviewed_hunks(&action_log, cx).is_empty()); + + // Reject all edits + action_log + .update(cx, |log, cx| log.reject_all_edits(None, cx)) + .await; + cx.run_until_parked(); + + // Verify the buffer is back to original + assert_eq!( + buffer.read_with(cx, |buffer, _| buffer.text()), + "abc\ndef\nghi" + ); + assert!(unreviewed_hunks(&action_log, cx).is_empty()); + + // Verify undo state is available + assert!(action_log.read_with(cx, |log, _| log.has_pending_undo())); + + // Undo the reject + action_log + .update(cx, |log, cx| log.undo_last_reject(cx)) + .await; + + cx.run_until_parked(); + + // Verify the agent edit is restored + assert_eq!( + buffer.read_with(cx, |buffer, _| buffer.text()), + "abc\nAGENT_EDIT\nghi" + ); + + // Verify undo state is cleared + assert!(!action_log.read_with(cx, |log, _| log.has_pending_undo())); + } + + #[derive(Debug, PartialEq)] struct HunkStatus { range: Range, diff_status: DiffHunkStatusKind, diff --git a/crates/agent_ui/src/acp/thread_view.rs b/crates/agent_ui/src/acp/thread_view.rs index 41e0548efe4cf2191337e4d1c962b0a764bb2316..0393205c3bb00d2a955e3a76fc037ade842e3358 100644 --- a/crates/agent_ui/src/acp/thread_view.rs +++ b/crates/agent_ui/src/acp/thread_view.rs @@ -76,7 +76,7 @@ use crate::{ ExpandMessageEditor, ExternalAgentInitialContent, Follow, KeepAll, NewThread, OpenAddContextMenu, OpenAgentDiff, OpenHistory, RejectAll, RejectOnce, RemoveFirstQueuedMessage, SelectPermissionGranularity, SendImmediately, SendNextQueuedMessage, - ToggleProfileSelector, ToggleThinkingEffortMenu, ToggleThinkingMode, + ToggleProfileSelector, ToggleThinkingEffortMenu, ToggleThinkingMode, UndoLastReject, }; const STOPWATCH_THRESHOLD: Duration = Duration::from_secs(30); diff --git a/crates/agent_ui/src/acp/thread_view/active_thread.rs b/crates/agent_ui/src/acp/thread_view/active_thread.rs index 306994e5db0da15ff29cad9e6e8a7ee24f583305..904a42bfbecac8dc75de85d0e387edbec43174ef 100644 --- a/crates/agent_ui/src/acp/thread_view/active_thread.rs +++ b/crates/agent_ui/src/acp/thread_view/active_thread.rs @@ -1390,11 +1390,34 @@ impl AcpThreadView { let thread = &self.thread; let telemetry = ActionLogTelemetry::from(thread.read(cx)); let action_log = thread.read(cx).action_log().clone(); + let has_changes = action_log.read(cx).changed_buffers(cx).len() > 0; + action_log .update(cx, |action_log, cx| { action_log.reject_all_edits(Some(telemetry), cx) }) .detach(); + + if has_changes { + if let Some(workspace) = self.workspace.upgrade() { + workspace.update(cx, |workspace, cx| { + crate::ui::show_undo_reject_toast(workspace, action_log, cx); + }); + } + } + } + + pub fn undo_last_reject( + &mut self, + _: &UndoLastReject, + _window: &mut Window, + cx: &mut Context, + ) { + let thread = &self.thread; + let action_log = thread.read(cx).action_log().clone(); + action_log + .update(cx, |action_log, cx| action_log.undo_last_reject(cx)) + .detach() } pub fn open_edited_buffer( @@ -1946,6 +1969,7 @@ impl AcpThreadView { Some(telemetry.clone()), cx, ) + .0 .detach_and_log_err(cx); }) } @@ -7312,6 +7336,7 @@ impl Render for AcpThreadView { })) .on_action(cx.listener(Self::keep_all)) .on_action(cx.listener(Self::reject_all)) + .on_action(cx.listener(Self::undo_last_reject)) .on_action(cx.listener(Self::allow_always)) .on_action(cx.listener(Self::allow_once)) .on_action(cx.listener(Self::reject_once)) diff --git a/crates/agent_ui/src/agent_diff.rs b/crates/agent_ui/src/agent_diff.rs index d950d0b75cdf1cdc7aebeaf8083054f6881ec2e4..7a7220e5497946ccb705febfe4fcc127055225a7 100644 --- a/crates/agent_ui/src/agent_diff.rs +++ b/crates/agent_ui/src/agent_diff.rs @@ -1,6 +1,6 @@ use crate::{Keep, KeepAll, OpenAgentDiff, Reject, RejectAll}; use acp_thread::{AcpThread, AcpThreadEvent}; -use action_log::ActionLogTelemetry; +use action_log::{ActionLogTelemetry, LastRejectUndo}; use agent_settings::AgentSettings; use anyhow::Result; use buffer_diff::DiffHunkStatus; @@ -96,7 +96,8 @@ impl AgentDiffPane { 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(diff_hunk_controls(&thread), cx); + editor + .set_render_diff_hunk_controls(diff_hunk_controls(&thread, workspace.clone()), cx); editor.register_addon(AgentDiffAddon); editor }); @@ -274,7 +275,14 @@ impl AgentDiffPane { fn reject(&mut self, _: &Reject, window: &mut Window, cx: &mut Context) { self.editor.update(cx, |editor, cx| { let snapshot = editor.buffer().read(cx).snapshot(cx); - reject_edits_in_selection(editor, &snapshot, &self.thread, window, cx); + reject_edits_in_selection( + editor, + &snapshot, + &self.thread, + self.workspace.clone(), + window, + cx, + ); }); } @@ -286,6 +294,7 @@ impl AgentDiffPane { &snapshot, &self.thread, vec![editor::Anchor::min()..editor::Anchor::max()], + self.workspace.clone(), window, cx, ); @@ -320,6 +329,7 @@ fn reject_edits_in_selection( editor: &mut Editor, buffer_snapshot: &MultiBufferSnapshot, thread: &Entity, + workspace: WeakEntity, window: &mut Window, cx: &mut Context, ) { @@ -327,7 +337,15 @@ fn reject_edits_in_selection( .selections .disjoint_anchor_ranges() .collect::>(); - reject_edits_in_ranges(editor, buffer_snapshot, thread, ranges, window, cx) + reject_edits_in_ranges( + editor, + buffer_snapshot, + thread, + ranges, + workspace, + window, + cx, + ) } fn keep_edits_in_ranges( @@ -367,6 +385,7 @@ fn reject_edits_in_ranges( buffer_snapshot: &MultiBufferSnapshot, thread: &Entity, ranges: Vec>, + workspace: WeakEntity, window: &mut Window, cx: &mut Context, ) { @@ -391,13 +410,33 @@ fn reject_edits_in_ranges( let action_log = thread.read(cx).action_log().clone(); let telemetry = ActionLogTelemetry::from(thread.read(cx)); + let mut undo_buffers = Vec::new(); + for (buffer, ranges) in ranges_by_buffer { action_log .update(cx, |action_log, cx| { - action_log.reject_edits_in_ranges(buffer, ranges, Some(telemetry.clone()), cx) + let (task, undo_info) = + action_log.reject_edits_in_ranges(buffer, ranges, Some(telemetry.clone()), cx); + undo_buffers.extend(undo_info); + task }) .detach_and_log_err(cx); } + if !undo_buffers.is_empty() { + action_log.update(cx, |action_log, _cx| { + action_log.set_last_reject_undo(LastRejectUndo { + buffers: undo_buffers, + }); + }); + + if let Some(workspace) = workspace.upgrade() { + cx.defer(move |cx| { + workspace.update(cx, |workspace, cx| { + crate::ui::show_undo_reject_toast(workspace, action_log, cx); + }); + }); + } + } } fn update_editor_selection( @@ -667,7 +706,10 @@ impl Render for AgentDiffPane { } } -fn diff_hunk_controls(thread: &Entity) -> editor::RenderDiffHunkControlsFn { +fn diff_hunk_controls( + thread: &Entity, + workspace: WeakEntity, +) -> editor::RenderDiffHunkControlsFn { let thread = thread.clone(); Arc::new( @@ -681,6 +723,7 @@ fn diff_hunk_controls(thread: &Entity) -> editor::RenderDiffHunkContr line_height, &thread, editor, + workspace.clone(), cx, ) } @@ -696,6 +739,7 @@ fn render_diff_hunk_controls( line_height: Pixels, thread: &Entity, editor: &Entity, + workspace: WeakEntity, cx: &mut App, ) -> AnyElement { let editor = editor.clone(); @@ -732,6 +776,7 @@ fn render_diff_hunk_controls( &snapshot, &thread, vec![hunk_range.start..hunk_range.start], + workspace.clone(), window, cx, ); @@ -1304,7 +1349,13 @@ impl AgentDiff { fn register_review_action( workspace: &mut Workspace, - review: impl Fn(&Entity, &Entity, &mut Window, &mut App) -> PostReviewState + review: impl Fn( + &Entity, + &Entity, + &WeakEntity, + &mut Window, + &mut App, + ) -> PostReviewState + 'static, this: &Entity, ) { @@ -1500,7 +1551,10 @@ impl AgentDiff { if previous_state.is_none() { editor.update(cx, |editor, cx| { editor.start_temporary_diff_override(); - editor.set_render_diff_hunk_controls(diff_hunk_controls(&thread), cx); + editor.set_render_diff_hunk_controls( + diff_hunk_controls(&thread, workspace.clone()), + cx, + ); editor.set_expand_all_diff_hunks(cx); editor.register_addon(EditorAgentDiffAddon); }); @@ -1586,6 +1640,7 @@ impl AgentDiff { fn keep_all( editor: &Entity, thread: &Entity, + _workspace: &WeakEntity, window: &mut Window, cx: &mut App, ) -> PostReviewState { @@ -1606,6 +1661,7 @@ impl AgentDiff { fn reject_all( editor: &Entity, thread: &Entity, + workspace: &WeakEntity, window: &mut Window, cx: &mut App, ) -> PostReviewState { @@ -1616,6 +1672,7 @@ impl AgentDiff { &snapshot, thread, vec![editor::Anchor::min()..editor::Anchor::max()], + workspace.clone(), window, cx, ); @@ -1626,6 +1683,7 @@ impl AgentDiff { fn keep( editor: &Entity, thread: &Entity, + _workspace: &WeakEntity, window: &mut Window, cx: &mut App, ) -> PostReviewState { @@ -1639,12 +1697,13 @@ impl AgentDiff { fn reject( editor: &Entity, thread: &Entity, + workspace: &WeakEntity, window: &mut Window, cx: &mut App, ) -> PostReviewState { editor.update(cx, |editor, cx| { let snapshot = editor.buffer().read(cx).snapshot(cx); - reject_edits_in_selection(editor, &snapshot, thread, window, cx); + reject_edits_in_selection(editor, &snapshot, thread, workspace.clone(), window, cx); Self::post_review_state(&snapshot) }) } @@ -1661,7 +1720,13 @@ impl AgentDiff { fn review_in_active_editor( &mut self, workspace: &mut Workspace, - review: impl Fn(&Entity, &Entity, &mut Window, &mut App) -> PostReviewState, + review: impl Fn( + &Entity, + &Entity, + &WeakEntity, + &mut Window, + &mut App, + ) -> PostReviewState, window: &mut Window, cx: &mut Context, ) -> Option>> { @@ -1680,7 +1745,9 @@ impl AgentDiff { let thread = thread.upgrade()?; - if let PostReviewState::AllReviewed = review(&editor, &thread, window, cx) + let review_result = review(&editor, &thread, &workspace.weak_handle(), window, cx); + + if matches!(review_result, PostReviewState::AllReviewed) && let Some(curr_buffer) = editor.read(cx).buffer().read(cx).as_singleton() { let changed_buffers = thread.read(cx).action_log().read(cx).changed_buffers(cx); diff --git a/crates/agent_ui/src/agent_ui.rs b/crates/agent_ui/src/agent_ui.rs index a517ea866bc5c5ed18e807d2bcfd4cf0bcd77532..9c06ba95ba2117b01d0851b88339b2b7db5e358b 100644 --- a/crates/agent_ui/src/agent_ui.rs +++ b/crates/agent_ui/src/agent_ui.rs @@ -111,6 +111,8 @@ actions!( Reject, /// Rejects all suggestions or changes. RejectAll, + /// Undoes the most recent reject operation, restoring the rejected changes. + UndoLastReject, /// Keeps all suggestions or changes. KeepAll, /// Allow this operation only this time. diff --git a/crates/agent_ui/src/ui.rs b/crates/agent_ui/src/ui.rs index 4df34592cfe73db514c25368c933637e8c18f720..16732951ce67d76ca8d65259e309c4b81df30c3b 100644 --- a/crates/agent_ui/src/ui.rs +++ b/crates/agent_ui/src/ui.rs @@ -5,6 +5,7 @@ mod end_trial_upsell; mod hold_for_default; mod mention_crease; mod model_selector_components; +mod undo_reject_toast; pub use acp_onboarding_modal::*; pub use agent_notification::*; @@ -13,3 +14,4 @@ pub use end_trial_upsell::*; pub use hold_for_default::*; pub use mention_crease::*; pub use model_selector_components::*; +pub use undo_reject_toast::*; diff --git a/crates/agent_ui/src/ui/undo_reject_toast.rs b/crates/agent_ui/src/ui/undo_reject_toast.rs new file mode 100644 index 0000000000000000000000000000000000000000..90c8a9c7ea98edd56ca935eddb36206a83bcc4bc --- /dev/null +++ b/crates/agent_ui/src/ui/undo_reject_toast.rs @@ -0,0 +1,25 @@ +use action_log::ActionLog; +use gpui::{App, Entity}; +use notifications::status_toast::{StatusToast, ToastIcon}; +use ui::prelude::*; +use workspace::Workspace; + +pub fn show_undo_reject_toast( + workspace: &mut Workspace, + action_log: Entity, + cx: &mut App, +) { + let action_log_weak = action_log.downgrade(); + let status_toast = StatusToast::new("Agent Changes Rejected", cx, move |this, _cx| { + this.icon(ToastIcon::new(IconName::Undo).color(Color::Muted)) + .action("Undo", move |_window, cx| { + if let Some(action_log) = action_log_weak.upgrade() { + action_log + .update(cx, |action_log, cx| action_log.undo_last_reject(cx)) + .detach(); + } + }) + .dismiss_button(true) + }); + workspace.toggle_status_toast(status_toast, cx); +}