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