agent_ui: Add the ability to undo "reject all" (#48462)

Danilo Leal and Bennet Bo Fenner created

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 <bennetbo@gmx.de>

Change summary

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 
crates/agent_ui/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(-)

Detailed changes

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",

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",

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",

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<Buffer>,
+    pub edits_to_restore: Vec<(Range<Anchor>, 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<PerBufferUndo>,
+}
+
 /// 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<Entity<Buffer>, TrackedBuffer>,
     /// The project this action log is associated with
     project: Entity<Project>,
+    /// Stores undo information for the most recent reject operation
+    last_reject_undo: Option<LastRejectUndo>,
 }
 
 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<Range<impl language::ToPoint>>,
         telemetry: Option<ActionLogTelemetry>,
         cx: &mut Context<Self>,
-    ) -> Task<Result<()>> {
+    ) -> (Task<Result<()>>, Option<PerBufferUndo>) {
         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<PerBufferUndo> = 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::<String>();
+
+                            // 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::<String>();
+                            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<ActionLogTelemetry>,
         cx: &mut Context<Self>,
     ) -> 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<Self>) -> 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<Buffer>, Entity<BufferDiff>> {
         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<Point>,
         diff_status: DiffHunkStatusKind,

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);

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<Self>,
+    ) {
+        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))

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>) {
         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<AcpThread>,
+    workspace: WeakEntity<Workspace>,
     window: &mut Window,
     cx: &mut Context<Editor>,
 ) {
@@ -327,7 +337,15 @@ fn reject_edits_in_selection(
         .selections
         .disjoint_anchor_ranges()
         .collect::<Vec<_>>();
-    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<AcpThread>,
     ranges: Vec<Range<editor::Anchor>>,
+    workspace: WeakEntity<Workspace>,
     window: &mut Window,
     cx: &mut Context<Editor>,
 ) {
@@ -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<AcpThread>) -> editor::RenderDiffHunkControlsFn {
+fn diff_hunk_controls(
+    thread: &Entity<AcpThread>,
+    workspace: WeakEntity<Workspace>,
+) -> editor::RenderDiffHunkControlsFn {
     let thread = thread.clone();
 
     Arc::new(
@@ -681,6 +723,7 @@ fn diff_hunk_controls(thread: &Entity<AcpThread>) -> editor::RenderDiffHunkContr
                     line_height,
                     &thread,
                     editor,
+                    workspace.clone(),
                     cx,
                 )
             }
@@ -696,6 +739,7 @@ fn render_diff_hunk_controls(
     line_height: Pixels,
     thread: &Entity<AcpThread>,
     editor: &Entity<Editor>,
+    workspace: WeakEntity<Workspace>,
     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<T: Action>(
         workspace: &mut Workspace,
-        review: impl Fn(&Entity<Editor>, &Entity<AcpThread>, &mut Window, &mut App) -> PostReviewState
+        review: impl Fn(
+            &Entity<Editor>,
+            &Entity<AcpThread>,
+            &WeakEntity<Workspace>,
+            &mut Window,
+            &mut App,
+        ) -> PostReviewState
         + 'static,
         this: &Entity<AgentDiff>,
     ) {
@@ -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<Editor>,
         thread: &Entity<AcpThread>,
+        _workspace: &WeakEntity<Workspace>,
         window: &mut Window,
         cx: &mut App,
     ) -> PostReviewState {
@@ -1606,6 +1661,7 @@ impl AgentDiff {
     fn reject_all(
         editor: &Entity<Editor>,
         thread: &Entity<AcpThread>,
+        workspace: &WeakEntity<Workspace>,
         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<Editor>,
         thread: &Entity<AcpThread>,
+        _workspace: &WeakEntity<Workspace>,
         window: &mut Window,
         cx: &mut App,
     ) -> PostReviewState {
@@ -1639,12 +1697,13 @@ impl AgentDiff {
     fn reject(
         editor: &Entity<Editor>,
         thread: &Entity<AcpThread>,
+        workspace: &WeakEntity<Workspace>,
         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<Editor>, &Entity<AcpThread>, &mut Window, &mut App) -> PostReviewState,
+        review: impl Fn(
+            &Entity<Editor>,
+            &Entity<AcpThread>,
+            &WeakEntity<Workspace>,
+            &mut Window,
+            &mut App,
+        ) -> PostReviewState,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) -> Option<Task<Result<()>>> {
@@ -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);

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.

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::*;

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<ActionLog>,
+    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);
+}