agent_diff.rs

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