agent_diff.rs

   1use crate::{Keep, Reject, Thread, ThreadEvent};
   2use anyhow::Result;
   3use buffer_diff::DiffHunkStatus;
   4use collections::HashSet;
   5use editor::{
   6    AnchorRangeExt, Direction, Editor, EditorEvent, MultiBuffer, ToPoint,
   7    actions::{GoToHunk, GoToPreviousHunk},
   8    scroll::Autoscroll,
   9};
  10use gpui::{
  11    Action, AnyElement, AnyView, App, Entity, EventEmitter, FocusHandle, Focusable, SharedString,
  12    Subscription, Task, WeakEntity, Window, prelude::*,
  13};
  14use language::{Capability, DiskState, OffsetRangeExt, Point};
  15use multi_buffer::PathKey;
  16use project::{Project, ProjectPath};
  17use std::{
  18    any::{Any, TypeId},
  19    ops::Range,
  20    sync::Arc,
  21};
  22use ui::{IconButtonShape, KeyBinding, Tooltip, prelude::*};
  23use workspace::{
  24    Item, ItemHandle, ItemNavHistory, ToolbarItemEvent, ToolbarItemLocation, ToolbarItemView,
  25    Workspace,
  26    item::{BreadcrumbText, ItemEvent, TabContentParams},
  27    searchable::SearchableItemHandle,
  28};
  29use zed_actions::assistant::ToggleFocus;
  30
  31pub struct AgentDiff {
  32    multibuffer: Entity<MultiBuffer>,
  33    editor: Entity<Editor>,
  34    thread: Entity<Thread>,
  35    focus_handle: FocusHandle,
  36    workspace: WeakEntity<Workspace>,
  37    title: SharedString,
  38    _subscriptions: Vec<Subscription>,
  39}
  40
  41impl AgentDiff {
  42    pub fn deploy(
  43        thread: Entity<Thread>,
  44        workspace: WeakEntity<Workspace>,
  45        window: &mut Window,
  46        cx: &mut App,
  47    ) -> Result<Entity<Self>> {
  48        let existing_diff = workspace.update(cx, |workspace, cx| {
  49            workspace
  50                .items_of_type::<AgentDiff>(cx)
  51                .find(|diff| diff.read(cx).thread == thread)
  52        })?;
  53        if let Some(existing_diff) = existing_diff {
  54            workspace.update(cx, |workspace, cx| {
  55                workspace.activate_item(&existing_diff, true, true, window, cx);
  56            })?;
  57            Ok(existing_diff)
  58        } else {
  59            let agent_diff =
  60                cx.new(|cx| AgentDiff::new(thread.clone(), workspace.clone(), window, cx));
  61            workspace.update(cx, |workspace, cx| {
  62                workspace.add_item_to_center(Box::new(agent_diff.clone()), window, cx);
  63            })?;
  64            Ok(agent_diff)
  65        }
  66    }
  67
  68    pub fn new(
  69        thread: Entity<Thread>,
  70        workspace: WeakEntity<Workspace>,
  71        window: &mut Window,
  72        cx: &mut Context<Self>,
  73    ) -> Self {
  74        let focus_handle = cx.focus_handle();
  75        let multibuffer = cx.new(|_| MultiBuffer::new(Capability::ReadWrite));
  76
  77        let project = thread.read(cx).project().clone();
  78        let render_diff_hunk_controls = Arc::new({
  79            let agent_diff = cx.entity();
  80            move |row,
  81                  status: &DiffHunkStatus,
  82                  hunk_range,
  83                  is_created_file,
  84                  line_height,
  85                  editor: &Entity<Editor>,
  86                  window: &mut Window,
  87                  cx: &mut App| {
  88                render_diff_hunk_controls(
  89                    row,
  90                    status,
  91                    hunk_range,
  92                    is_created_file,
  93                    line_height,
  94                    &agent_diff,
  95                    editor,
  96                    window,
  97                    cx,
  98                )
  99            }
 100        });
 101        let editor = cx.new(|cx| {
 102            let mut editor =
 103                Editor::for_multibuffer(multibuffer.clone(), Some(project.clone()), window, cx);
 104            editor.disable_inline_diagnostics();
 105            editor.set_expand_all_diff_hunks(cx);
 106            editor.set_render_diff_hunk_controls(render_diff_hunk_controls, cx);
 107            editor.register_addon(AgentDiffAddon);
 108            editor
 109        });
 110
 111        let action_log = thread.read(cx).action_log().clone();
 112        let mut this = Self {
 113            _subscriptions: vec![
 114                cx.observe_in(&action_log, window, |this, _action_log, window, cx| {
 115                    this.update_excerpts(window, cx)
 116                }),
 117                cx.subscribe(&thread, |this, _thread, event, cx| {
 118                    this.handle_thread_event(event, cx)
 119                }),
 120            ],
 121            title: SharedString::default(),
 122            multibuffer,
 123            editor,
 124            thread,
 125            focus_handle,
 126            workspace,
 127        };
 128        this.update_excerpts(window, cx);
 129        this.update_title(cx);
 130        this
 131    }
 132
 133    fn update_excerpts(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 134        let thread = self.thread.read(cx);
 135        let changed_buffers = thread.action_log().read(cx).changed_buffers(cx);
 136        let mut paths_to_delete = self.multibuffer.read(cx).paths().collect::<HashSet<_>>();
 137
 138        for (buffer, diff_handle) in changed_buffers {
 139            if buffer.read(cx).file().is_none() {
 140                continue;
 141            }
 142
 143            let path_key = PathKey::for_buffer(&buffer, cx);
 144            paths_to_delete.remove(&path_key);
 145
 146            let snapshot = buffer.read(cx).snapshot();
 147            let diff = diff_handle.read(cx);
 148
 149            let diff_hunk_ranges = diff
 150                .hunks_intersecting_range(
 151                    language::Anchor::MIN..language::Anchor::MAX,
 152                    &snapshot,
 153                    cx,
 154                )
 155                .map(|diff_hunk| diff_hunk.buffer_range.to_point(&snapshot))
 156                .collect::<Vec<_>>();
 157
 158            let (was_empty, is_excerpt_newly_added) =
 159                self.multibuffer.update(cx, |multibuffer, cx| {
 160                    let was_empty = multibuffer.is_empty();
 161                    let (_, is_excerpt_newly_added) = multibuffer.set_excerpts_for_path(
 162                        path_key.clone(),
 163                        buffer.clone(),
 164                        diff_hunk_ranges,
 165                        editor::DEFAULT_MULTIBUFFER_CONTEXT,
 166                        cx,
 167                    );
 168                    multibuffer.add_diff(diff_handle, cx);
 169                    (was_empty, is_excerpt_newly_added)
 170                });
 171
 172            self.editor.update(cx, |editor, cx| {
 173                if was_empty {
 174                    let first_hunk = editor
 175                        .diff_hunks_in_ranges(
 176                            &[editor::Anchor::min()..editor::Anchor::max()],
 177                            &self.multibuffer.read(cx).read(cx),
 178                        )
 179                        .next();
 180
 181                    if let Some(first_hunk) = first_hunk {
 182                        let first_hunk_start = first_hunk.multi_buffer_range().start;
 183                        editor.change_selections(
 184                            Some(Autoscroll::fit()),
 185                            window,
 186                            cx,
 187                            |selections| {
 188                                selections
 189                                    .select_anchor_ranges([first_hunk_start..first_hunk_start]);
 190                            },
 191                        )
 192                    }
 193                }
 194
 195                if is_excerpt_newly_added
 196                    && buffer
 197                        .read(cx)
 198                        .file()
 199                        .map_or(false, |file| file.disk_state() == DiskState::Deleted)
 200                {
 201                    editor.fold_buffer(snapshot.text.remote_id(), cx)
 202                }
 203            });
 204        }
 205
 206        self.multibuffer.update(cx, |multibuffer, cx| {
 207            for path in paths_to_delete {
 208                multibuffer.remove_excerpts_for_path(path, cx);
 209            }
 210        });
 211
 212        if self.multibuffer.read(cx).is_empty()
 213            && self
 214                .editor
 215                .read(cx)
 216                .focus_handle(cx)
 217                .contains_focused(window, cx)
 218        {
 219            self.focus_handle.focus(window);
 220        } else if self.focus_handle.is_focused(window) && !self.multibuffer.read(cx).is_empty() {
 221            self.editor.update(cx, |editor, cx| {
 222                editor.focus_handle(cx).focus(window);
 223            });
 224        }
 225    }
 226
 227    fn update_title(&mut self, cx: &mut Context<Self>) {
 228        let new_title = self
 229            .thread
 230            .read(cx)
 231            .summary()
 232            .unwrap_or("Assistant Changes".into());
 233        if new_title != self.title {
 234            self.title = new_title;
 235            cx.emit(EditorEvent::TitleChanged);
 236        }
 237    }
 238
 239    fn handle_thread_event(&mut self, event: &ThreadEvent, cx: &mut Context<Self>) {
 240        match event {
 241            ThreadEvent::SummaryChanged => self.update_title(cx),
 242            _ => {}
 243        }
 244    }
 245
 246    pub fn move_to_path(&mut self, path_key: PathKey, window: &mut Window, cx: &mut Context<Self>) {
 247        if let Some(position) = self.multibuffer.read(cx).location_for_path(&path_key, cx) {
 248            self.editor.update(cx, |editor, cx| {
 249                let first_hunk = editor
 250                    .diff_hunks_in_ranges(
 251                        &[position..editor::Anchor::max()],
 252                        &self.multibuffer.read(cx).read(cx),
 253                    )
 254                    .next();
 255
 256                if let Some(first_hunk) = first_hunk {
 257                    let first_hunk_start = first_hunk.multi_buffer_range().start;
 258                    editor.change_selections(Some(Autoscroll::fit()), window, cx, |selections| {
 259                        selections.select_anchor_ranges([first_hunk_start..first_hunk_start]);
 260                    })
 261                }
 262            });
 263        }
 264    }
 265
 266    fn keep(&mut self, _: &crate::Keep, window: &mut Window, cx: &mut Context<Self>) {
 267        let ranges = self
 268            .editor
 269            .read(cx)
 270            .selections
 271            .disjoint_anchor_ranges()
 272            .collect::<Vec<_>>();
 273        self.keep_edits_in_ranges(ranges, window, cx);
 274    }
 275
 276    fn reject(&mut self, _: &crate::Reject, window: &mut Window, cx: &mut Context<Self>) {
 277        let ranges = self
 278            .editor
 279            .read(cx)
 280            .selections
 281            .disjoint_anchor_ranges()
 282            .collect::<Vec<_>>();
 283        self.reject_edits_in_ranges(ranges, window, cx);
 284    }
 285
 286    fn reject_all(&mut self, _: &crate::RejectAll, window: &mut Window, cx: &mut Context<Self>) {
 287        self.reject_edits_in_ranges(
 288            vec![editor::Anchor::min()..editor::Anchor::max()],
 289            window,
 290            cx,
 291        );
 292    }
 293
 294    fn keep_all(&mut self, _: &crate::KeepAll, _window: &mut Window, cx: &mut Context<Self>) {
 295        self.thread
 296            .update(cx, |thread, cx| thread.keep_all_edits(cx));
 297    }
 298
 299    fn keep_edits_in_ranges(
 300        &mut self,
 301        ranges: Vec<Range<editor::Anchor>>,
 302        window: &mut Window,
 303        cx: &mut Context<Self>,
 304    ) {
 305        let snapshot = self.multibuffer.read(cx).snapshot(cx);
 306        let diff_hunks_in_ranges = self
 307            .editor
 308            .read(cx)
 309            .diff_hunks_in_ranges(&ranges, &snapshot)
 310            .collect::<Vec<_>>();
 311        let newest_cursor = self.editor.update(cx, |editor, cx| {
 312            editor.selections.newest::<Point>(cx).head()
 313        });
 314        if diff_hunks_in_ranges.iter().any(|hunk| {
 315            hunk.row_range
 316                .contains(&multi_buffer::MultiBufferRow(newest_cursor.row))
 317        }) {
 318            self.update_selection(&diff_hunks_in_ranges, window, cx);
 319        }
 320
 321        for hunk in &diff_hunks_in_ranges {
 322            let buffer = self.multibuffer.read(cx).buffer(hunk.buffer_id);
 323            if let Some(buffer) = buffer {
 324                self.thread.update(cx, |thread, cx| {
 325                    thread.keep_edits_in_range(buffer, hunk.buffer_range.clone(), cx)
 326                });
 327            }
 328        }
 329    }
 330
 331    fn reject_edits_in_ranges(
 332        &mut self,
 333        ranges: Vec<Range<editor::Anchor>>,
 334        window: &mut Window,
 335        cx: &mut Context<Self>,
 336    ) {
 337        let snapshot = self.multibuffer.read(cx).snapshot(cx);
 338        let diff_hunks_in_ranges = self
 339            .editor
 340            .read(cx)
 341            .diff_hunks_in_ranges(&ranges, &snapshot)
 342            .collect::<Vec<_>>();
 343        let newest_cursor = self.editor.update(cx, |editor, cx| {
 344            editor.selections.newest::<Point>(cx).head()
 345        });
 346        if diff_hunks_in_ranges.iter().any(|hunk| {
 347            hunk.row_range
 348                .contains(&multi_buffer::MultiBufferRow(newest_cursor.row))
 349        }) {
 350            self.update_selection(&diff_hunks_in_ranges, window, cx);
 351        }
 352
 353        let point_ranges = ranges
 354            .into_iter()
 355            .map(|range| range.to_point(&snapshot))
 356            .collect();
 357        self.editor.update(cx, |editor, cx| {
 358            editor.restore_hunks_in_ranges(point_ranges, window, cx)
 359        });
 360    }
 361
 362    fn update_selection(
 363        &mut self,
 364        diff_hunks: &[multi_buffer::MultiBufferDiffHunk],
 365        window: &mut Window,
 366        cx: &mut Context<Self>,
 367    ) {
 368        let snapshot = self.multibuffer.read(cx).snapshot(cx);
 369        let target_hunk = diff_hunks
 370            .last()
 371            .and_then(|last_kept_hunk| {
 372                let last_kept_hunk_end = last_kept_hunk.multi_buffer_range().end;
 373                self.editor
 374                    .read(cx)
 375                    .diff_hunks_in_ranges(&[last_kept_hunk_end..editor::Anchor::max()], &snapshot)
 376                    .skip(1)
 377                    .next()
 378            })
 379            .or_else(|| {
 380                let first_kept_hunk = diff_hunks.first()?;
 381                let first_kept_hunk_start = first_kept_hunk.multi_buffer_range().start;
 382                self.editor
 383                    .read(cx)
 384                    .diff_hunks_in_ranges(
 385                        &[editor::Anchor::min()..first_kept_hunk_start],
 386                        &snapshot,
 387                    )
 388                    .next()
 389            });
 390
 391        if let Some(target_hunk) = target_hunk {
 392            self.editor.update(cx, |editor, cx| {
 393                editor.change_selections(Some(Autoscroll::fit()), window, cx, |selections| {
 394                    let next_hunk_start = target_hunk.multi_buffer_range().start;
 395                    selections.select_anchor_ranges([next_hunk_start..next_hunk_start]);
 396                })
 397            });
 398        }
 399    }
 400}
 401
 402impl EventEmitter<EditorEvent> for AgentDiff {}
 403
 404impl Focusable for AgentDiff {
 405    fn focus_handle(&self, cx: &App) -> FocusHandle {
 406        if self.multibuffer.read(cx).is_empty() {
 407            self.focus_handle.clone()
 408        } else {
 409            self.editor.focus_handle(cx)
 410        }
 411    }
 412}
 413
 414impl Item for AgentDiff {
 415    type Event = EditorEvent;
 416
 417    fn tab_icon(&self, _window: &Window, _cx: &App) -> Option<Icon> {
 418        Some(Icon::new(IconName::ZedAssistant).color(Color::Muted))
 419    }
 420
 421    fn to_item_events(event: &EditorEvent, f: impl FnMut(ItemEvent)) {
 422        Editor::to_item_events(event, f)
 423    }
 424
 425    fn deactivated(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 426        self.editor
 427            .update(cx, |editor, cx| editor.deactivated(window, cx));
 428    }
 429
 430    fn navigate(
 431        &mut self,
 432        data: Box<dyn Any>,
 433        window: &mut Window,
 434        cx: &mut Context<Self>,
 435    ) -> bool {
 436        self.editor
 437            .update(cx, |editor, cx| editor.navigate(data, window, cx))
 438    }
 439
 440    fn tab_tooltip_text(&self, _: &App) -> Option<SharedString> {
 441        Some("Agent Diff".into())
 442    }
 443
 444    fn tab_content(&self, params: TabContentParams, _window: &Window, cx: &App) -> AnyElement {
 445        let summary = self
 446            .thread
 447            .read(cx)
 448            .summary()
 449            .unwrap_or("Assistant Changes".into());
 450        Label::new(format!("Review: {}", summary))
 451            .color(if params.selected {
 452                Color::Default
 453            } else {
 454                Color::Muted
 455            })
 456            .into_any_element()
 457    }
 458
 459    fn telemetry_event_text(&self) -> Option<&'static str> {
 460        Some("Assistant Diff Opened")
 461    }
 462
 463    fn as_searchable(&self, _: &Entity<Self>) -> Option<Box<dyn SearchableItemHandle>> {
 464        Some(Box::new(self.editor.clone()))
 465    }
 466
 467    fn for_each_project_item(
 468        &self,
 469        cx: &App,
 470        f: &mut dyn FnMut(gpui::EntityId, &dyn project::ProjectItem),
 471    ) {
 472        self.editor.for_each_project_item(cx, f)
 473    }
 474
 475    fn is_singleton(&self, _: &App) -> bool {
 476        false
 477    }
 478
 479    fn set_nav_history(
 480        &mut self,
 481        nav_history: ItemNavHistory,
 482        _: &mut Window,
 483        cx: &mut Context<Self>,
 484    ) {
 485        self.editor.update(cx, |editor, _| {
 486            editor.set_nav_history(Some(nav_history));
 487        });
 488    }
 489
 490    fn clone_on_split(
 491        &self,
 492        _workspace_id: Option<workspace::WorkspaceId>,
 493        window: &mut Window,
 494        cx: &mut Context<Self>,
 495    ) -> Option<Entity<Self>>
 496    where
 497        Self: Sized,
 498    {
 499        Some(cx.new(|cx| Self::new(self.thread.clone(), self.workspace.clone(), window, cx)))
 500    }
 501
 502    fn is_dirty(&self, cx: &App) -> bool {
 503        self.multibuffer.read(cx).is_dirty(cx)
 504    }
 505
 506    fn has_conflict(&self, cx: &App) -> bool {
 507        self.multibuffer.read(cx).has_conflict(cx)
 508    }
 509
 510    fn can_save(&self, _: &App) -> bool {
 511        true
 512    }
 513
 514    fn save(
 515        &mut self,
 516        format: bool,
 517        project: Entity<Project>,
 518        window: &mut Window,
 519        cx: &mut Context<Self>,
 520    ) -> Task<Result<()>> {
 521        self.editor.save(format, project, window, cx)
 522    }
 523
 524    fn save_as(
 525        &mut self,
 526        _: Entity<Project>,
 527        _: ProjectPath,
 528        _window: &mut Window,
 529        _: &mut Context<Self>,
 530    ) -> Task<Result<()>> {
 531        unreachable!()
 532    }
 533
 534    fn reload(
 535        &mut self,
 536        project: Entity<Project>,
 537        window: &mut Window,
 538        cx: &mut Context<Self>,
 539    ) -> Task<Result<()>> {
 540        self.editor.reload(project, window, cx)
 541    }
 542
 543    fn act_as_type<'a>(
 544        &'a self,
 545        type_id: TypeId,
 546        self_handle: &'a Entity<Self>,
 547        _: &'a App,
 548    ) -> Option<AnyView> {
 549        if type_id == TypeId::of::<Self>() {
 550            Some(self_handle.to_any())
 551        } else if type_id == TypeId::of::<Editor>() {
 552            Some(self.editor.to_any())
 553        } else {
 554            None
 555        }
 556    }
 557
 558    fn breadcrumb_location(&self, _: &App) -> ToolbarItemLocation {
 559        ToolbarItemLocation::PrimaryLeft
 560    }
 561
 562    fn breadcrumbs(&self, theme: &theme::Theme, cx: &App) -> Option<Vec<BreadcrumbText>> {
 563        self.editor.breadcrumbs(theme, cx)
 564    }
 565
 566    fn added_to_workspace(
 567        &mut self,
 568        workspace: &mut Workspace,
 569        window: &mut Window,
 570        cx: &mut Context<Self>,
 571    ) {
 572        self.editor.update(cx, |editor, cx| {
 573            editor.added_to_workspace(workspace, window, cx)
 574        });
 575    }
 576}
 577
 578impl Render for AgentDiff {
 579    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 580        let is_empty = self.multibuffer.read(cx).is_empty();
 581        let focus_handle = &self.focus_handle;
 582
 583        div()
 584            .track_focus(focus_handle)
 585            .key_context(if is_empty { "EmptyPane" } else { "AgentDiff" })
 586            .on_action(cx.listener(Self::keep))
 587            .on_action(cx.listener(Self::reject))
 588            .on_action(cx.listener(Self::reject_all))
 589            .on_action(cx.listener(Self::keep_all))
 590            .bg(cx.theme().colors().editor_background)
 591            .flex()
 592            .items_center()
 593            .justify_center()
 594            .size_full()
 595            .when(is_empty, |el| {
 596                el.child(
 597                    v_flex()
 598                        .items_center()
 599                        .gap_2()
 600                        .child("No changes to review")
 601                        .child(
 602                            Button::new("continue-iterating", "Continue Iterating")
 603                                .style(ButtonStyle::Filled)
 604                                .icon(IconName::ForwardArrow)
 605                                .icon_position(IconPosition::Start)
 606                                .icon_size(IconSize::Small)
 607                                .icon_color(Color::Muted)
 608                                .full_width()
 609                                .key_binding(KeyBinding::for_action_in(
 610                                    &ToggleFocus,
 611                                    &focus_handle.clone(),
 612                                    window,
 613                                    cx,
 614                                ))
 615                                .on_click(|_event, window, cx| {
 616                                    window.dispatch_action(ToggleFocus.boxed_clone(), cx)
 617                                }),
 618                        ),
 619                )
 620            })
 621            .when(!is_empty, |el| el.child(self.editor.clone()))
 622    }
 623}
 624
 625fn render_diff_hunk_controls(
 626    row: u32,
 627    _status: &DiffHunkStatus,
 628    hunk_range: Range<editor::Anchor>,
 629    is_created_file: bool,
 630    line_height: Pixels,
 631    agent_diff: &Entity<AgentDiff>,
 632    editor: &Entity<Editor>,
 633    window: &mut Window,
 634    cx: &mut App,
 635) -> AnyElement {
 636    let editor = editor.clone();
 637    h_flex()
 638        .h(line_height)
 639        .mr_0p5()
 640        .gap_1()
 641        .px_0p5()
 642        .pb_1()
 643        .border_x_1()
 644        .border_b_1()
 645        .border_color(cx.theme().colors().border)
 646        .rounded_b_md()
 647        .bg(cx.theme().colors().editor_background)
 648        .gap_1()
 649        .occlude()
 650        .shadow_md()
 651        .children(vec![
 652            Button::new(("reject", row as u64), "Reject")
 653                .disabled(is_created_file)
 654                .key_binding(
 655                    KeyBinding::for_action_in(
 656                        &Reject,
 657                        &editor.read(cx).focus_handle(cx),
 658                        window,
 659                        cx,
 660                    )
 661                    .map(|kb| kb.size(rems_from_px(12.))),
 662                )
 663                .on_click({
 664                    let agent_diff = agent_diff.clone();
 665                    move |_event, window, cx| {
 666                        agent_diff.update(cx, |diff, cx| {
 667                            diff.reject_edits_in_ranges(
 668                                vec![hunk_range.start..hunk_range.start],
 669                                window,
 670                                cx,
 671                            );
 672                        });
 673                    }
 674                }),
 675            Button::new(("keep", row as u64), "Keep")
 676                .key_binding(
 677                    KeyBinding::for_action_in(&Keep, &editor.read(cx).focus_handle(cx), window, cx)
 678                        .map(|kb| kb.size(rems_from_px(12.))),
 679                )
 680                .on_click({
 681                    let agent_diff = agent_diff.clone();
 682                    move |_event, window, cx| {
 683                        agent_diff.update(cx, |diff, cx| {
 684                            diff.keep_edits_in_ranges(
 685                                vec![hunk_range.start..hunk_range.start],
 686                                window,
 687                                cx,
 688                            );
 689                        });
 690                    }
 691                }),
 692        ])
 693        .when(
 694            !editor.read(cx).buffer().read(cx).all_diff_hunks_expanded(),
 695            |el| {
 696                el.child(
 697                    IconButton::new(("next-hunk", row as u64), IconName::ArrowDown)
 698                        .shape(IconButtonShape::Square)
 699                        .icon_size(IconSize::Small)
 700                        // .disabled(!has_multiple_hunks)
 701                        .tooltip({
 702                            let focus_handle = editor.focus_handle(cx);
 703                            move |window, cx| {
 704                                Tooltip::for_action_in(
 705                                    "Next Hunk",
 706                                    &GoToHunk,
 707                                    &focus_handle,
 708                                    window,
 709                                    cx,
 710                                )
 711                            }
 712                        })
 713                        .on_click({
 714                            let editor = editor.clone();
 715                            move |_event, window, cx| {
 716                                editor.update(cx, |editor, cx| {
 717                                    let snapshot = editor.snapshot(window, cx);
 718                                    let position =
 719                                        hunk_range.end.to_point(&snapshot.buffer_snapshot);
 720                                    editor.go_to_hunk_before_or_after_position(
 721                                        &snapshot,
 722                                        position,
 723                                        Direction::Next,
 724                                        window,
 725                                        cx,
 726                                    );
 727                                    editor.expand_selected_diff_hunks(cx);
 728                                });
 729                            }
 730                        }),
 731                )
 732                .child(
 733                    IconButton::new(("prev-hunk", row as u64), IconName::ArrowUp)
 734                        .shape(IconButtonShape::Square)
 735                        .icon_size(IconSize::Small)
 736                        // .disabled(!has_multiple_hunks)
 737                        .tooltip({
 738                            let focus_handle = editor.focus_handle(cx);
 739                            move |window, cx| {
 740                                Tooltip::for_action_in(
 741                                    "Previous Hunk",
 742                                    &GoToPreviousHunk,
 743                                    &focus_handle,
 744                                    window,
 745                                    cx,
 746                                )
 747                            }
 748                        })
 749                        .on_click({
 750                            let editor = editor.clone();
 751                            move |_event, window, cx| {
 752                                editor.update(cx, |editor, cx| {
 753                                    let snapshot = editor.snapshot(window, cx);
 754                                    let point =
 755                                        hunk_range.start.to_point(&snapshot.buffer_snapshot);
 756                                    editor.go_to_hunk_before_or_after_position(
 757                                        &snapshot,
 758                                        point,
 759                                        Direction::Prev,
 760                                        window,
 761                                        cx,
 762                                    );
 763                                    editor.expand_selected_diff_hunks(cx);
 764                                });
 765                            }
 766                        }),
 767                )
 768            },
 769        )
 770        .into_any_element()
 771}
 772
 773struct AgentDiffAddon;
 774
 775impl editor::Addon for AgentDiffAddon {
 776    fn to_any(&self) -> &dyn std::any::Any {
 777        self
 778    }
 779
 780    fn extend_key_context(&self, key_context: &mut gpui::KeyContext, _: &App) {
 781        key_context.add("agent_diff");
 782    }
 783}
 784
 785pub struct AgentDiffToolbar {
 786    agent_diff: Option<WeakEntity<AgentDiff>>,
 787    _workspace: WeakEntity<Workspace>,
 788}
 789
 790impl AgentDiffToolbar {
 791    pub fn new(workspace: &Workspace, _: &mut Context<Self>) -> Self {
 792        Self {
 793            agent_diff: None,
 794            _workspace: workspace.weak_handle(),
 795        }
 796    }
 797
 798    fn agent_diff(&self, _: &App) -> Option<Entity<AgentDiff>> {
 799        self.agent_diff.as_ref()?.upgrade()
 800    }
 801
 802    fn dispatch_action(&self, action: &dyn Action, window: &mut Window, cx: &mut Context<Self>) {
 803        if let Some(agent_diff) = self.agent_diff(cx) {
 804            agent_diff.focus_handle(cx).focus(window);
 805        }
 806        let action = action.boxed_clone();
 807        cx.defer(move |cx| {
 808            cx.dispatch_action(action.as_ref());
 809        })
 810    }
 811}
 812
 813impl EventEmitter<ToolbarItemEvent> for AgentDiffToolbar {}
 814
 815impl ToolbarItemView for AgentDiffToolbar {
 816    fn set_active_pane_item(
 817        &mut self,
 818        active_pane_item: Option<&dyn ItemHandle>,
 819        _: &mut Window,
 820        cx: &mut Context<Self>,
 821    ) -> ToolbarItemLocation {
 822        self.agent_diff = active_pane_item
 823            .and_then(|item| item.act_as::<AgentDiff>(cx))
 824            .map(|entity| entity.downgrade());
 825        if self.agent_diff.is_some() {
 826            ToolbarItemLocation::PrimaryRight
 827        } else {
 828            ToolbarItemLocation::Hidden
 829        }
 830    }
 831
 832    fn pane_focus_update(
 833        &mut self,
 834        _pane_focused: bool,
 835        _window: &mut Window,
 836        _cx: &mut Context<Self>,
 837    ) {
 838    }
 839}
 840
 841impl Render for AgentDiffToolbar {
 842    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 843        let agent_diff = match self.agent_diff(cx) {
 844            Some(ad) => ad,
 845            None => return div(),
 846        };
 847
 848        let is_empty = agent_diff.read(cx).multibuffer.read(cx).is_empty();
 849
 850        if is_empty {
 851            return div();
 852        }
 853
 854        h_group_xl()
 855            .my_neg_1()
 856            .items_center()
 857            .p_1()
 858            .flex_wrap()
 859            .justify_between()
 860            .child(
 861                h_group_sm()
 862                    .child(
 863                        Button::new("reject-all", "Reject All").on_click(cx.listener(
 864                            |this, _, window, cx| {
 865                                this.dispatch_action(&crate::RejectAll, window, cx)
 866                            },
 867                        )),
 868                    )
 869                    .child(Button::new("keep-all", "Keep All").on_click(cx.listener(
 870                        |this, _, window, cx| this.dispatch_action(&crate::KeepAll, window, cx),
 871                    ))),
 872            )
 873    }
 874}
 875
 876#[cfg(test)]
 877mod tests {
 878    use super::*;
 879    use crate::{ThreadStore, thread_store};
 880    use assistant_settings::AssistantSettings;
 881    use context_server::ContextServerSettings;
 882    use editor::EditorSettings;
 883    use gpui::TestAppContext;
 884    use project::{FakeFs, Project};
 885    use prompt_store::PromptBuilder;
 886    use serde_json::json;
 887    use settings::{Settings, SettingsStore};
 888    use std::sync::Arc;
 889    use theme::ThemeSettings;
 890    use util::path;
 891
 892    #[gpui::test]
 893    async fn test_agent_diff(cx: &mut TestAppContext) {
 894        cx.update(|cx| {
 895            let settings_store = SettingsStore::test(cx);
 896            cx.set_global(settings_store);
 897            language::init(cx);
 898            Project::init_settings(cx);
 899            AssistantSettings::register(cx);
 900            thread_store::init(cx);
 901            workspace::init_settings(cx);
 902            ThemeSettings::register(cx);
 903            ContextServerSettings::register(cx);
 904            EditorSettings::register(cx);
 905        });
 906
 907        let fs = FakeFs::new(cx.executor());
 908        fs.insert_tree(
 909            path!("/test"),
 910            json!({"file1": "abc\ndef\nghi\njkl\nmno\npqr\nstu\nvwx\nyz"}),
 911        )
 912        .await;
 913        let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
 914        let buffer_path = project
 915            .read_with(cx, |project, cx| {
 916                project.find_project_path("test/file1", cx)
 917            })
 918            .unwrap();
 919
 920        let thread_store = cx.update(|cx| {
 921            ThreadStore::new(
 922                project.clone(),
 923                Arc::default(),
 924                Arc::new(PromptBuilder::new(None).unwrap()),
 925                cx,
 926            )
 927            .unwrap()
 928        });
 929        let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
 930        let action_log = thread.read_with(cx, |thread, _| thread.action_log().clone());
 931
 932        let (workspace, cx) =
 933            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
 934        let agent_diff = cx.new_window_entity(|window, cx| {
 935            AgentDiff::new(thread.clone(), workspace.downgrade(), window, cx)
 936        });
 937        let editor = agent_diff.read_with(cx, |diff, _cx| diff.editor.clone());
 938
 939        let buffer = project
 940            .update(cx, |project, cx| project.open_buffer(buffer_path, cx))
 941            .await
 942            .unwrap();
 943        cx.update(|_, cx| {
 944            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
 945            buffer.update(cx, |buffer, cx| {
 946                buffer
 947                    .edit(
 948                        [
 949                            (Point::new(1, 1)..Point::new(1, 2), "E"),
 950                            (Point::new(3, 2)..Point::new(3, 3), "L"),
 951                            (Point::new(5, 0)..Point::new(5, 1), "P"),
 952                            (Point::new(7, 1)..Point::new(7, 2), "W"),
 953                        ],
 954                        None,
 955                        cx,
 956                    )
 957                    .unwrap()
 958            });
 959            action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
 960        });
 961        cx.run_until_parked();
 962
 963        // When opening the assistant diff, the cursor is positioned on the first hunk.
 964        assert_eq!(
 965            editor.read_with(cx, |editor, cx| editor.text(cx)),
 966            "abc\ndef\ndEf\nghi\njkl\njkL\nmno\npqr\nPqr\nstu\nvwx\nvWx\nyz"
 967        );
 968        assert_eq!(
 969            editor
 970                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
 971                .range(),
 972            Point::new(1, 0)..Point::new(1, 0)
 973        );
 974
 975        // After keeping a hunk, the cursor should be positioned on the second hunk.
 976        agent_diff.update_in(cx, |diff, window, cx| diff.keep(&crate::Keep, window, cx));
 977        cx.run_until_parked();
 978        assert_eq!(
 979            editor.read_with(cx, |editor, cx| editor.text(cx)),
 980            "abc\ndEf\nghi\njkl\njkL\nmno\npqr\nPqr\nstu\nvwx\nvWx\nyz"
 981        );
 982        assert_eq!(
 983            editor
 984                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
 985                .range(),
 986            Point::new(3, 0)..Point::new(3, 0)
 987        );
 988
 989        // Restoring a hunk also moves the cursor to the next hunk, possibly cycling if it's at the end.
 990        editor.update_in(cx, |editor, window, cx| {
 991            editor.change_selections(None, window, cx, |selections| {
 992                selections.select_ranges([Point::new(10, 0)..Point::new(10, 0)])
 993            });
 994        });
 995        agent_diff.update_in(cx, |diff, window, cx| {
 996            diff.reject(&crate::Reject, window, cx)
 997        });
 998        cx.run_until_parked();
 999        assert_eq!(
1000            editor.read_with(cx, |editor, cx| editor.text(cx)),
1001            "abc\ndEf\nghi\njkl\njkL\nmno\npqr\nPqr\nstu\nvwx\nyz"
1002        );
1003        assert_eq!(
1004            editor
1005                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
1006                .range(),
1007            Point::new(3, 0)..Point::new(3, 0)
1008        );
1009
1010        // Keeping a range that doesn't intersect the current selection doesn't move it.
1011        agent_diff.update_in(cx, |diff, window, cx| {
1012            let position = editor
1013                .read(cx)
1014                .buffer()
1015                .read(cx)
1016                .read(cx)
1017                .anchor_before(Point::new(7, 0));
1018            diff.keep_edits_in_ranges(vec![position..position], window, cx)
1019        });
1020        cx.run_until_parked();
1021        assert_eq!(
1022            editor.read_with(cx, |editor, cx| editor.text(cx)),
1023            "abc\ndEf\nghi\njkl\njkL\nmno\nPqr\nstu\nvwx\nyz"
1024        );
1025        assert_eq!(
1026            editor
1027                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
1028                .range(),
1029            Point::new(3, 0)..Point::new(3, 0)
1030        );
1031    }
1032}