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    fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
 602        "Agent Diff".into()
 603    }
 604}
 605
 606impl Render for AgentDiff {
 607    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 608        let is_empty = self.multibuffer.read(cx).is_empty();
 609        let focus_handle = &self.focus_handle;
 610
 611        div()
 612            .track_focus(focus_handle)
 613            .key_context(if is_empty { "EmptyPane" } else { "AgentDiff" })
 614            .on_action(cx.listener(Self::keep))
 615            .on_action(cx.listener(Self::reject))
 616            .on_action(cx.listener(Self::reject_all))
 617            .on_action(cx.listener(Self::keep_all))
 618            .bg(cx.theme().colors().editor_background)
 619            .flex()
 620            .items_center()
 621            .justify_center()
 622            .size_full()
 623            .when(is_empty, |el| {
 624                el.child(
 625                    v_flex()
 626                        .items_center()
 627                        .gap_2()
 628                        .child("No changes to review")
 629                        .child(
 630                            Button::new("continue-iterating", "Continue Iterating")
 631                                .style(ButtonStyle::Filled)
 632                                .icon(IconName::ForwardArrow)
 633                                .icon_position(IconPosition::Start)
 634                                .icon_size(IconSize::Small)
 635                                .icon_color(Color::Muted)
 636                                .full_width()
 637                                .key_binding(KeyBinding::for_action_in(
 638                                    &ToggleFocus,
 639                                    &focus_handle.clone(),
 640                                    window,
 641                                    cx,
 642                                ))
 643                                .on_click(|_event, window, cx| {
 644                                    window.dispatch_action(ToggleFocus.boxed_clone(), cx)
 645                                }),
 646                        ),
 647                )
 648            })
 649            .when(!is_empty, |el| el.child(self.editor.clone()))
 650    }
 651}
 652
 653fn render_diff_hunk_controls(
 654    row: u32,
 655    _status: &DiffHunkStatus,
 656    hunk_range: Range<editor::Anchor>,
 657    is_created_file: bool,
 658    line_height: Pixels,
 659    agent_diff: &Entity<AgentDiff>,
 660    editor: &Entity<Editor>,
 661    window: &mut Window,
 662    cx: &mut App,
 663) -> AnyElement {
 664    let editor = editor.clone();
 665
 666    if agent_diff.read(cx).thread.read(cx).is_generating() {
 667        return Empty.into_any();
 668    }
 669
 670    h_flex()
 671        .h(line_height)
 672        .mr_0p5()
 673        .gap_1()
 674        .px_0p5()
 675        .pb_1()
 676        .border_x_1()
 677        .border_b_1()
 678        .border_color(cx.theme().colors().border)
 679        .rounded_b_md()
 680        .bg(cx.theme().colors().editor_background)
 681        .gap_1()
 682        .occlude()
 683        .shadow_md()
 684        .children(vec![
 685            Button::new(("reject", row as u64), "Reject")
 686                .disabled(is_created_file)
 687                .key_binding(
 688                    KeyBinding::for_action_in(
 689                        &Reject,
 690                        &editor.read(cx).focus_handle(cx),
 691                        window,
 692                        cx,
 693                    )
 694                    .map(|kb| kb.size(rems_from_px(12.))),
 695                )
 696                .on_click({
 697                    let agent_diff = agent_diff.clone();
 698                    move |_event, window, cx| {
 699                        agent_diff.update(cx, |diff, cx| {
 700                            diff.reject_edits_in_ranges(
 701                                vec![hunk_range.start..hunk_range.start],
 702                                window,
 703                                cx,
 704                            );
 705                        });
 706                    }
 707                }),
 708            Button::new(("keep", row as u64), "Keep")
 709                .key_binding(
 710                    KeyBinding::for_action_in(&Keep, &editor.read(cx).focus_handle(cx), window, cx)
 711                        .map(|kb| kb.size(rems_from_px(12.))),
 712                )
 713                .on_click({
 714                    let agent_diff = agent_diff.clone();
 715                    move |_event, window, cx| {
 716                        agent_diff.update(cx, |diff, cx| {
 717                            diff.keep_edits_in_ranges(
 718                                vec![hunk_range.start..hunk_range.start],
 719                                window,
 720                                cx,
 721                            );
 722                        });
 723                    }
 724                }),
 725        ])
 726        .when(
 727            !editor.read(cx).buffer().read(cx).all_diff_hunks_expanded(),
 728            |el| {
 729                el.child(
 730                    IconButton::new(("next-hunk", row as u64), IconName::ArrowDown)
 731                        .shape(IconButtonShape::Square)
 732                        .icon_size(IconSize::Small)
 733                        // .disabled(!has_multiple_hunks)
 734                        .tooltip({
 735                            let focus_handle = editor.focus_handle(cx);
 736                            move |window, cx| {
 737                                Tooltip::for_action_in(
 738                                    "Next Hunk",
 739                                    &GoToHunk,
 740                                    &focus_handle,
 741                                    window,
 742                                    cx,
 743                                )
 744                            }
 745                        })
 746                        .on_click({
 747                            let editor = editor.clone();
 748                            move |_event, window, cx| {
 749                                editor.update(cx, |editor, cx| {
 750                                    let snapshot = editor.snapshot(window, cx);
 751                                    let position =
 752                                        hunk_range.end.to_point(&snapshot.buffer_snapshot);
 753                                    editor.go_to_hunk_before_or_after_position(
 754                                        &snapshot,
 755                                        position,
 756                                        Direction::Next,
 757                                        window,
 758                                        cx,
 759                                    );
 760                                    editor.expand_selected_diff_hunks(cx);
 761                                });
 762                            }
 763                        }),
 764                )
 765                .child(
 766                    IconButton::new(("prev-hunk", row as u64), IconName::ArrowUp)
 767                        .shape(IconButtonShape::Square)
 768                        .icon_size(IconSize::Small)
 769                        // .disabled(!has_multiple_hunks)
 770                        .tooltip({
 771                            let focus_handle = editor.focus_handle(cx);
 772                            move |window, cx| {
 773                                Tooltip::for_action_in(
 774                                    "Previous Hunk",
 775                                    &GoToPreviousHunk,
 776                                    &focus_handle,
 777                                    window,
 778                                    cx,
 779                                )
 780                            }
 781                        })
 782                        .on_click({
 783                            let editor = editor.clone();
 784                            move |_event, window, cx| {
 785                                editor.update(cx, |editor, cx| {
 786                                    let snapshot = editor.snapshot(window, cx);
 787                                    let point =
 788                                        hunk_range.start.to_point(&snapshot.buffer_snapshot);
 789                                    editor.go_to_hunk_before_or_after_position(
 790                                        &snapshot,
 791                                        point,
 792                                        Direction::Prev,
 793                                        window,
 794                                        cx,
 795                                    );
 796                                    editor.expand_selected_diff_hunks(cx);
 797                                });
 798                            }
 799                        }),
 800                )
 801            },
 802        )
 803        .into_any_element()
 804}
 805
 806struct AgentDiffAddon;
 807
 808impl editor::Addon for AgentDiffAddon {
 809    fn to_any(&self) -> &dyn std::any::Any {
 810        self
 811    }
 812
 813    fn extend_key_context(&self, key_context: &mut gpui::KeyContext, _: &App) {
 814        key_context.add("agent_diff");
 815    }
 816}
 817
 818pub struct AgentDiffToolbar {
 819    agent_diff: Option<WeakEntity<AgentDiff>>,
 820}
 821
 822impl AgentDiffToolbar {
 823    pub fn new() -> Self {
 824        Self { agent_diff: None }
 825    }
 826
 827    fn agent_diff(&self, _: &App) -> Option<Entity<AgentDiff>> {
 828        self.agent_diff.as_ref()?.upgrade()
 829    }
 830
 831    fn dispatch_action(&self, action: &dyn Action, window: &mut Window, cx: &mut Context<Self>) {
 832        if let Some(agent_diff) = self.agent_diff(cx) {
 833            agent_diff.focus_handle(cx).focus(window);
 834        }
 835        let action = action.boxed_clone();
 836        cx.defer(move |cx| {
 837            cx.dispatch_action(action.as_ref());
 838        })
 839    }
 840}
 841
 842impl EventEmitter<ToolbarItemEvent> for AgentDiffToolbar {}
 843
 844impl ToolbarItemView for AgentDiffToolbar {
 845    fn set_active_pane_item(
 846        &mut self,
 847        active_pane_item: Option<&dyn ItemHandle>,
 848        _: &mut Window,
 849        cx: &mut Context<Self>,
 850    ) -> ToolbarItemLocation {
 851        self.agent_diff = active_pane_item
 852            .and_then(|item| item.act_as::<AgentDiff>(cx))
 853            .map(|entity| entity.downgrade());
 854        if self.agent_diff.is_some() {
 855            ToolbarItemLocation::PrimaryRight
 856        } else {
 857            ToolbarItemLocation::Hidden
 858        }
 859    }
 860
 861    fn pane_focus_update(
 862        &mut self,
 863        _pane_focused: bool,
 864        _window: &mut Window,
 865        _cx: &mut Context<Self>,
 866    ) {
 867    }
 868}
 869
 870impl Render for AgentDiffToolbar {
 871    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 872        let agent_diff = match self.agent_diff(cx) {
 873            Some(ad) => ad,
 874            None => return div(),
 875        };
 876
 877        let is_generating = agent_diff.read(cx).thread.read(cx).is_generating();
 878        if is_generating {
 879            return div()
 880                .w(rems(6.5625)) // Arbitrary 105px size—so the label doesn't dance around
 881                .child(AnimatedLabel::new("Generating"));
 882        }
 883
 884        let is_empty = agent_diff.read(cx).multibuffer.read(cx).is_empty();
 885        if is_empty {
 886            return div();
 887        }
 888
 889        let focus_handle = agent_diff.focus_handle(cx);
 890
 891        h_group_xl()
 892            .my_neg_1()
 893            .items_center()
 894            .p_1()
 895            .flex_wrap()
 896            .justify_between()
 897            .child(
 898                h_group_sm()
 899                    .child(
 900                        Button::new("reject-all", "Reject All")
 901                            .key_binding({
 902                                KeyBinding::for_action_in(&RejectAll, &focus_handle, window, cx)
 903                                    .map(|kb| kb.size(rems_from_px(12.)))
 904                            })
 905                            .on_click(cx.listener(|this, _, window, cx| {
 906                                this.dispatch_action(&RejectAll, window, cx)
 907                            })),
 908                    )
 909                    .child(
 910                        Button::new("keep-all", "Keep All")
 911                            .key_binding({
 912                                KeyBinding::for_action_in(&KeepAll, &focus_handle, window, cx)
 913                                    .map(|kb| kb.size(rems_from_px(12.)))
 914                            })
 915                            .on_click(cx.listener(|this, _, window, cx| {
 916                                this.dispatch_action(&KeepAll, window, cx)
 917                            })),
 918                    ),
 919            )
 920    }
 921}
 922
 923#[cfg(test)]
 924mod tests {
 925    use super::*;
 926    use crate::{ThreadStore, thread_store};
 927    use assistant_settings::AssistantSettings;
 928    use assistant_tool::ToolWorkingSet;
 929    use context_server::ContextServerSettings;
 930    use editor::EditorSettings;
 931    use gpui::TestAppContext;
 932    use project::{FakeFs, Project};
 933    use prompt_store::PromptBuilder;
 934    use serde_json::json;
 935    use settings::{Settings, SettingsStore};
 936    use std::sync::Arc;
 937    use theme::ThemeSettings;
 938    use util::path;
 939
 940    #[gpui::test]
 941    async fn test_agent_diff(cx: &mut TestAppContext) {
 942        cx.update(|cx| {
 943            let settings_store = SettingsStore::test(cx);
 944            cx.set_global(settings_store);
 945            language::init(cx);
 946            Project::init_settings(cx);
 947            AssistantSettings::register(cx);
 948            prompt_store::init(cx);
 949            thread_store::init(cx);
 950            workspace::init_settings(cx);
 951            ThemeSettings::register(cx);
 952            ContextServerSettings::register(cx);
 953            EditorSettings::register(cx);
 954        });
 955
 956        let fs = FakeFs::new(cx.executor());
 957        fs.insert_tree(
 958            path!("/test"),
 959            json!({"file1": "abc\ndef\nghi\njkl\nmno\npqr\nstu\nvwx\nyz"}),
 960        )
 961        .await;
 962        let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
 963        let buffer_path = project
 964            .read_with(cx, |project, cx| {
 965                project.find_project_path("test/file1", cx)
 966            })
 967            .unwrap();
 968
 969        let prompt_store = None;
 970        let thread_store = cx
 971            .update(|cx| {
 972                ThreadStore::load(
 973                    project.clone(),
 974                    cx.new(|_| ToolWorkingSet::default()),
 975                    prompt_store,
 976                    Arc::new(PromptBuilder::new(None).unwrap()),
 977                    cx,
 978                )
 979            })
 980            .await
 981            .unwrap();
 982        let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
 983        let action_log = thread.read_with(cx, |thread, _| thread.action_log().clone());
 984
 985        let (workspace, cx) =
 986            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
 987        let agent_diff = cx.new_window_entity(|window, cx| {
 988            AgentDiff::new(thread.clone(), workspace.downgrade(), window, cx)
 989        });
 990        let editor = agent_diff.read_with(cx, |diff, _cx| diff.editor.clone());
 991
 992        let buffer = project
 993            .update(cx, |project, cx| project.open_buffer(buffer_path, cx))
 994            .await
 995            .unwrap();
 996        cx.update(|_, cx| {
 997            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
 998            buffer.update(cx, |buffer, cx| {
 999                buffer
1000                    .edit(
1001                        [
1002                            (Point::new(1, 1)..Point::new(1, 2), "E"),
1003                            (Point::new(3, 2)..Point::new(3, 3), "L"),
1004                            (Point::new(5, 0)..Point::new(5, 1), "P"),
1005                            (Point::new(7, 1)..Point::new(7, 2), "W"),
1006                        ],
1007                        None,
1008                        cx,
1009                    )
1010                    .unwrap()
1011            });
1012            action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
1013        });
1014        cx.run_until_parked();
1015
1016        // When opening the assistant diff, the cursor is positioned on the first hunk.
1017        assert_eq!(
1018            editor.read_with(cx, |editor, cx| editor.text(cx)),
1019            "abc\ndef\ndEf\nghi\njkl\njkL\nmno\npqr\nPqr\nstu\nvwx\nvWx\nyz"
1020        );
1021        assert_eq!(
1022            editor
1023                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
1024                .range(),
1025            Point::new(1, 0)..Point::new(1, 0)
1026        );
1027
1028        // After keeping a hunk, the cursor should be positioned on the second hunk.
1029        agent_diff.update_in(cx, |diff, window, cx| diff.keep(&crate::Keep, window, cx));
1030        cx.run_until_parked();
1031        assert_eq!(
1032            editor.read_with(cx, |editor, cx| editor.text(cx)),
1033            "abc\ndEf\nghi\njkl\njkL\nmno\npqr\nPqr\nstu\nvwx\nvWx\nyz"
1034        );
1035        assert_eq!(
1036            editor
1037                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
1038                .range(),
1039            Point::new(3, 0)..Point::new(3, 0)
1040        );
1041
1042        // Rejecting a hunk also moves the cursor to the next hunk, possibly cycling if it's at the end.
1043        editor.update_in(cx, |editor, window, cx| {
1044            editor.change_selections(None, window, cx, |selections| {
1045                selections.select_ranges([Point::new(10, 0)..Point::new(10, 0)])
1046            });
1047        });
1048        agent_diff.update_in(cx, |diff, window, cx| {
1049            diff.reject(&crate::Reject, window, cx)
1050        });
1051        cx.run_until_parked();
1052        assert_eq!(
1053            editor.read_with(cx, |editor, cx| editor.text(cx)),
1054            "abc\ndEf\nghi\njkl\njkL\nmno\npqr\nPqr\nstu\nvwx\nyz"
1055        );
1056        assert_eq!(
1057            editor
1058                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
1059                .range(),
1060            Point::new(3, 0)..Point::new(3, 0)
1061        );
1062
1063        // Keeping a range that doesn't intersect the current selection doesn't move it.
1064        agent_diff.update_in(cx, |diff, window, cx| {
1065            let position = editor
1066                .read(cx)
1067                .buffer()
1068                .read(cx)
1069                .read(cx)
1070                .anchor_before(Point::new(7, 0));
1071            diff.keep_edits_in_ranges(vec![position..position], window, cx)
1072        });
1073        cx.run_until_parked();
1074        assert_eq!(
1075            editor.read_with(cx, |editor, cx| editor.text(cx)),
1076            "abc\ndEf\nghi\njkl\njkL\nmno\nPqr\nstu\nvwx\nyz"
1077        );
1078        assert_eq!(
1079            editor
1080                .update(cx, |editor, cx| editor.selections.newest::<Point>(cx))
1081                .range(),
1082            Point::new(3, 0)..Point::new(3, 0)
1083        );
1084    }
1085}