message_editor.rs

   1use std::collections::BTreeMap;
   2use std::sync::Arc;
   3
   4use crate::assistant_model_selector::{AssistantModelSelector, ModelType};
   5use crate::context::{AgentContextKey, ContextCreasesAddon, ContextLoadResult, load_context};
   6use crate::tool_compatibility::{IncompatibleToolsState, IncompatibleToolsTooltip};
   7use crate::ui::{AgentPreview, AnimatedLabel, MaxModeTooltip};
   8use buffer_diff::BufferDiff;
   9use collections::{HashMap, HashSet};
  10use editor::actions::{MoveUp, Paste};
  11use editor::{
  12    AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorEvent,
  13    EditorMode, EditorStyle, MultiBuffer,
  14};
  15use feature_flags::{FeatureFlagAppExt, NewBillingFeatureFlag};
  16use file_icons::FileIcons;
  17use fs::Fs;
  18use futures::future::Shared;
  19use futures::{FutureExt as _, future};
  20use gpui::{
  21    Animation, AnimationExt, App, ClipboardEntry, Entity, EventEmitter, Focusable, Subscription,
  22    Task, TextStyle, WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,
  23};
  24use language::{Buffer, Language};
  25use language_model::{ConfiguredModel, LanguageModelRequestMessage, MessageContent};
  26use language_model_selector::ToggleModelSelector;
  27use multi_buffer;
  28use project::Project;
  29use prompt_store::PromptStore;
  30use settings::Settings;
  31use std::time::Duration;
  32use theme::ThemeSettings;
  33use ui::{Disclosure, KeyBinding, PopoverMenuHandle, Tooltip, prelude::*};
  34use util::ResultExt as _;
  35use workspace::{CollaboratorId, Workspace};
  36use zed_llm_client::CompletionMode;
  37
  38use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider, crease_for_mention};
  39use crate::context_store::ContextStore;
  40use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
  41use crate::profile_selector::ProfileSelector;
  42use crate::thread::{MessageCrease, Thread, TokenUsageRatio};
  43use crate::thread_store::ThreadStore;
  44use crate::{
  45    ActiveThread, AgentDiffPane, Chat, ExpandMessageEditor, Follow, NewThread, OpenAgentDiff,
  46    RemoveAllContext, ToggleContextPicker, ToggleProfileSelector, register_agent_preview,
  47};
  48
  49#[derive(RegisterComponent)]
  50pub struct MessageEditor {
  51    thread: Entity<Thread>,
  52    incompatible_tools_state: Entity<IncompatibleToolsState>,
  53    editor: Entity<Editor>,
  54    workspace: WeakEntity<Workspace>,
  55    project: Entity<Project>,
  56    context_store: Entity<ContextStore>,
  57    prompt_store: Option<Entity<PromptStore>>,
  58    context_strip: Entity<ContextStrip>,
  59    context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
  60    model_selector: Entity<AssistantModelSelector>,
  61    last_loaded_context: Option<ContextLoadResult>,
  62    load_context_task: Option<Shared<Task<()>>>,
  63    profile_selector: Entity<ProfileSelector>,
  64    edits_expanded: bool,
  65    editor_is_expanded: bool,
  66    last_estimated_token_count: Option<usize>,
  67    update_token_count_task: Option<Task<()>>,
  68    _subscriptions: Vec<Subscription>,
  69}
  70
  71const MAX_EDITOR_LINES: usize = 8;
  72
  73pub(crate) fn create_editor(
  74    workspace: WeakEntity<Workspace>,
  75    context_store: WeakEntity<ContextStore>,
  76    thread_store: WeakEntity<ThreadStore>,
  77    window: &mut Window,
  78    cx: &mut App,
  79) -> Entity<Editor> {
  80    let language = Language::new(
  81        language::LanguageConfig {
  82            completion_query_characters: HashSet::from_iter(['.', '-', '_', '@']),
  83            ..Default::default()
  84        },
  85        None,
  86    );
  87
  88    let editor = cx.new(|cx| {
  89        let buffer = cx.new(|cx| Buffer::local("", cx).with_language(Arc::new(language), cx));
  90        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
  91        let mut editor = Editor::new(
  92            editor::EditorMode::AutoHeight {
  93                max_lines: MAX_EDITOR_LINES,
  94            },
  95            buffer,
  96            None,
  97            window,
  98            cx,
  99        );
 100        editor.set_placeholder_text("Message the agent – @ to include context", cx);
 101        editor.set_show_indent_guides(false, cx);
 102        editor.set_soft_wrap();
 103        editor.set_context_menu_options(ContextMenuOptions {
 104            min_entries_visible: 12,
 105            max_entries_visible: 12,
 106            placement: Some(ContextMenuPlacement::Above),
 107        });
 108        editor.register_addon(ContextCreasesAddon::new());
 109        editor
 110    });
 111
 112    let editor_entity = editor.downgrade();
 113    editor.update(cx, |editor, _| {
 114        editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
 115            workspace,
 116            context_store,
 117            Some(thread_store),
 118            editor_entity,
 119            None,
 120        ))));
 121    });
 122    editor
 123}
 124
 125impl MessageEditor {
 126    pub fn new(
 127        fs: Arc<dyn Fs>,
 128        workspace: WeakEntity<Workspace>,
 129        context_store: Entity<ContextStore>,
 130        prompt_store: Option<Entity<PromptStore>>,
 131        thread_store: WeakEntity<ThreadStore>,
 132        thread: Entity<Thread>,
 133        window: &mut Window,
 134        cx: &mut Context<Self>,
 135    ) -> Self {
 136        let context_picker_menu_handle = PopoverMenuHandle::default();
 137        let model_selector_menu_handle = PopoverMenuHandle::default();
 138
 139        let editor = create_editor(
 140            workspace.clone(),
 141            context_store.downgrade(),
 142            thread_store.clone(),
 143            window,
 144            cx,
 145        );
 146
 147        let context_strip = cx.new(|cx| {
 148            ContextStrip::new(
 149                context_store.clone(),
 150                workspace.clone(),
 151                Some(thread_store.clone()),
 152                context_picker_menu_handle.clone(),
 153                SuggestContextKind::File,
 154                window,
 155                cx,
 156            )
 157        });
 158
 159        let incompatible_tools =
 160            cx.new(|cx| IncompatibleToolsState::new(thread.read(cx).tools().clone(), cx));
 161
 162        let subscriptions = vec![
 163            cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event),
 164            cx.subscribe(&editor, |this, _, event, cx| match event {
 165                EditorEvent::BufferEdited => this.handle_message_changed(cx),
 166                _ => {}
 167            }),
 168            cx.observe(&context_store, |this, _, cx| {
 169                // When context changes, reload it for token counting.
 170                let _ = this.reload_context(cx);
 171            }),
 172            cx.observe(&thread.read(cx).action_log().clone(), |_, _, cx| {
 173                cx.notify()
 174            }),
 175        ];
 176
 177        let model_selector = cx.new(|cx| {
 178            AssistantModelSelector::new(
 179                fs.clone(),
 180                model_selector_menu_handle,
 181                editor.focus_handle(cx),
 182                ModelType::Default(thread.clone()),
 183                window,
 184                cx,
 185            )
 186        });
 187
 188        Self {
 189            editor: editor.clone(),
 190            project: thread.read(cx).project().clone(),
 191            thread,
 192            incompatible_tools_state: incompatible_tools.clone(),
 193            workspace,
 194            context_store,
 195            prompt_store,
 196            context_strip,
 197            context_picker_menu_handle,
 198            load_context_task: None,
 199            last_loaded_context: None,
 200            model_selector,
 201            edits_expanded: false,
 202            editor_is_expanded: false,
 203            profile_selector: cx
 204                .new(|cx| ProfileSelector::new(fs, thread_store, editor.focus_handle(cx), cx)),
 205            last_estimated_token_count: None,
 206            update_token_count_task: None,
 207            _subscriptions: subscriptions,
 208        }
 209    }
 210
 211    pub fn context_store(&self) -> &Entity<ContextStore> {
 212        &self.context_store
 213    }
 214
 215    pub fn expand_message_editor(
 216        &mut self,
 217        _: &ExpandMessageEditor,
 218        _window: &mut Window,
 219        cx: &mut Context<Self>,
 220    ) {
 221        self.set_editor_is_expanded(!self.editor_is_expanded, cx);
 222    }
 223
 224    fn set_editor_is_expanded(&mut self, is_expanded: bool, cx: &mut Context<Self>) {
 225        self.editor_is_expanded = is_expanded;
 226        self.editor.update(cx, |editor, _| {
 227            if self.editor_is_expanded {
 228                editor.set_mode(EditorMode::Full {
 229                    scale_ui_elements_with_buffer_font_size: false,
 230                    show_active_line_background: false,
 231                    sized_by_content: false,
 232                })
 233            } else {
 234                editor.set_mode(EditorMode::AutoHeight {
 235                    max_lines: MAX_EDITOR_LINES,
 236                })
 237            }
 238        });
 239        cx.notify();
 240    }
 241
 242    fn toggle_context_picker(
 243        &mut self,
 244        _: &ToggleContextPicker,
 245        window: &mut Window,
 246        cx: &mut Context<Self>,
 247    ) {
 248        self.context_picker_menu_handle.toggle(window, cx);
 249    }
 250
 251    pub fn remove_all_context(
 252        &mut self,
 253        _: &RemoveAllContext,
 254        _window: &mut Window,
 255        cx: &mut Context<Self>,
 256    ) {
 257        self.context_store.update(cx, |store, _cx| store.clear());
 258        cx.notify();
 259    }
 260
 261    fn chat(&mut self, _: &Chat, window: &mut Window, cx: &mut Context<Self>) {
 262        if self.is_editor_empty(cx) {
 263            return;
 264        }
 265
 266        self.thread.update(cx, |thread, cx| {
 267            thread.cancel_editing(cx);
 268        });
 269
 270        if self.thread.read(cx).is_generating() {
 271            self.stop_current_and_send_new_message(window, cx);
 272            return;
 273        }
 274
 275        self.set_editor_is_expanded(false, cx);
 276        self.send_to_model(window, cx);
 277
 278        cx.notify();
 279    }
 280
 281    fn is_editor_empty(&self, cx: &App) -> bool {
 282        self.editor.read(cx).text(cx).trim().is_empty()
 283    }
 284
 285    fn send_to_model(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 286        let Some(ConfiguredModel { model, provider }) = self
 287            .thread
 288            .update(cx, |thread, cx| thread.get_or_init_configured_model(cx))
 289        else {
 290            return;
 291        };
 292
 293        if provider.must_accept_terms(cx) {
 294            cx.notify();
 295            return;
 296        }
 297
 298        let (user_message, user_message_creases) = self.editor.update(cx, |editor, cx| {
 299            let creases = extract_message_creases(editor, cx);
 300            let text = editor.text(cx);
 301            editor.clear(window, cx);
 302            (text, creases)
 303        });
 304
 305        self.last_estimated_token_count.take();
 306        cx.emit(MessageEditorEvent::EstimatedTokenCount);
 307
 308        let thread = self.thread.clone();
 309        let git_store = self.project.read(cx).git_store().clone();
 310        let checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
 311        let context_task = self.reload_context(cx);
 312        let window_handle = window.window_handle();
 313
 314        cx.spawn(async move |_this, cx| {
 315            let (checkpoint, loaded_context) = future::join(checkpoint, context_task).await;
 316            let loaded_context = loaded_context.unwrap_or_default();
 317
 318            thread
 319                .update(cx, |thread, cx| {
 320                    thread.insert_user_message(
 321                        user_message,
 322                        loaded_context,
 323                        checkpoint.ok(),
 324                        user_message_creases,
 325                        cx,
 326                    );
 327                })
 328                .log_err();
 329
 330            thread
 331                .update(cx, |thread, cx| {
 332                    thread.advance_prompt_id();
 333                    thread.send_to_model(model, Some(window_handle), cx);
 334                })
 335                .log_err();
 336        })
 337        .detach();
 338    }
 339
 340    fn stop_current_and_send_new_message(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 341        self.thread.update(cx, |thread, cx| {
 342            thread.cancel_editing(cx);
 343        });
 344
 345        let cancelled = self.thread.update(cx, |thread, cx| {
 346            thread.cancel_last_completion(Some(window.window_handle()), cx)
 347        });
 348
 349        if cancelled {
 350            self.set_editor_is_expanded(false, cx);
 351            self.send_to_model(window, cx);
 352        }
 353    }
 354
 355    fn handle_context_strip_event(
 356        &mut self,
 357        _context_strip: &Entity<ContextStrip>,
 358        event: &ContextStripEvent,
 359        window: &mut Window,
 360        cx: &mut Context<Self>,
 361    ) {
 362        match event {
 363            ContextStripEvent::PickerDismissed
 364            | ContextStripEvent::BlurredEmpty
 365            | ContextStripEvent::BlurredDown => {
 366                let editor_focus_handle = self.editor.focus_handle(cx);
 367                window.focus(&editor_focus_handle);
 368            }
 369            ContextStripEvent::BlurredUp => {}
 370        }
 371    }
 372
 373    fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
 374        if self.context_picker_menu_handle.is_deployed() {
 375            cx.propagate();
 376        } else {
 377            self.context_strip.focus_handle(cx).focus(window);
 378        }
 379    }
 380
 381    fn paste(&mut self, _: &Paste, _: &mut Window, cx: &mut Context<Self>) {
 382        let images = cx
 383            .read_from_clipboard()
 384            .map(|item| {
 385                item.into_entries()
 386                    .filter_map(|entry| {
 387                        if let ClipboardEntry::Image(image) = entry {
 388                            Some(image)
 389                        } else {
 390                            None
 391                        }
 392                    })
 393                    .collect::<Vec<_>>()
 394            })
 395            .unwrap_or_default();
 396
 397        if images.is_empty() {
 398            return;
 399        }
 400        cx.stop_propagation();
 401
 402        self.context_store.update(cx, |store, cx| {
 403            for image in images {
 404                store.add_image_instance(Arc::new(image), cx);
 405            }
 406        });
 407    }
 408
 409    fn handle_review_click(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 410        self.edits_expanded = true;
 411        AgentDiffPane::deploy(self.thread.clone(), self.workspace.clone(), window, cx).log_err();
 412        cx.notify();
 413    }
 414
 415    fn handle_file_click(
 416        &self,
 417        buffer: Entity<Buffer>,
 418        window: &mut Window,
 419        cx: &mut Context<Self>,
 420    ) {
 421        if let Ok(diff) =
 422            AgentDiffPane::deploy(self.thread.clone(), self.workspace.clone(), window, cx)
 423        {
 424            let path_key = multi_buffer::PathKey::for_buffer(&buffer, cx);
 425            diff.update(cx, |diff, cx| diff.move_to_path(path_key, window, cx));
 426        }
 427    }
 428
 429    fn render_max_mode_toggle(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
 430        if !cx.has_flag::<NewBillingFeatureFlag>() {
 431            return None;
 432        }
 433
 434        let thread = self.thread.read(cx);
 435        let model = thread.configured_model();
 436        if !model?.model.supports_max_mode() {
 437            return None;
 438        }
 439
 440        let active_completion_mode = thread.completion_mode();
 441
 442        Some(
 443            IconButton::new("max-mode", IconName::ZedMaxMode)
 444                .icon_size(IconSize::Small)
 445                .icon_color(Color::Muted)
 446                .toggle_state(active_completion_mode == CompletionMode::Max)
 447                .on_click(cx.listener(move |this, _event, _window, cx| {
 448                    this.thread.update(cx, |thread, _cx| {
 449                        thread.set_completion_mode(match active_completion_mode {
 450                            CompletionMode::Max => CompletionMode::Normal,
 451                            CompletionMode::Normal => CompletionMode::Max,
 452                        });
 453                    });
 454                }))
 455                .tooltip(|_, cx| cx.new(MaxModeTooltip::new).into())
 456                .into_any_element(),
 457        )
 458    }
 459
 460    fn render_follow_toggle(&self, cx: &mut Context<Self>) -> impl IntoElement {
 461        let following = self
 462            .workspace
 463            .read_with(cx, |workspace, _| {
 464                workspace.is_being_followed(CollaboratorId::Agent)
 465            })
 466            .unwrap_or(false);
 467
 468        IconButton::new("follow-agent", IconName::Crosshair)
 469            .icon_size(IconSize::Small)
 470            .icon_color(Color::Muted)
 471            .toggle_state(following)
 472            .selected_icon_color(Some(Color::Custom(cx.theme().players().agent().cursor)))
 473            .tooltip(move |window, cx| {
 474                if following {
 475                    Tooltip::for_action("Stop Following Agent", &Follow, window, cx)
 476                } else {
 477                    Tooltip::with_meta(
 478                        "Follow Agent",
 479                        Some(&Follow),
 480                        "Track the agent's location as it reads and edits files.",
 481                        window,
 482                        cx,
 483                    )
 484                }
 485            })
 486            .on_click(cx.listener(move |this, _, window, cx| {
 487                this.workspace
 488                    .update(cx, |workspace, cx| {
 489                        if following {
 490                            workspace.unfollow(CollaboratorId::Agent, window, cx);
 491                        } else {
 492                            workspace.follow(CollaboratorId::Agent, window, cx);
 493                        }
 494                    })
 495                    .ok();
 496            }))
 497    }
 498
 499    fn render_editor(
 500        &self,
 501        font_size: Rems,
 502        line_height: Pixels,
 503        window: &mut Window,
 504        cx: &mut Context<Self>,
 505    ) -> Div {
 506        let thread = self.thread.read(cx);
 507        let model = thread.configured_model();
 508
 509        let editor_bg_color = cx.theme().colors().editor_background;
 510        let is_generating = thread.is_generating();
 511        let focus_handle = self.editor.focus_handle(cx);
 512
 513        let is_model_selected = model.is_some();
 514        let is_editor_empty = self.is_editor_empty(cx);
 515
 516        let incompatible_tools = model
 517            .as_ref()
 518            .map(|model| {
 519                self.incompatible_tools_state.update(cx, |state, cx| {
 520                    state
 521                        .incompatible_tools(&model.model, cx)
 522                        .iter()
 523                        .cloned()
 524                        .collect::<Vec<_>>()
 525                })
 526            })
 527            .unwrap_or_default();
 528
 529        let is_editor_expanded = self.editor_is_expanded;
 530        let expand_icon = if is_editor_expanded {
 531            IconName::Minimize
 532        } else {
 533            IconName::Maximize
 534        };
 535
 536        v_flex()
 537            .key_context("MessageEditor")
 538            .on_action(cx.listener(Self::chat))
 539            .on_action(cx.listener(|this, _: &ToggleProfileSelector, window, cx| {
 540                this.profile_selector
 541                    .read(cx)
 542                    .menu_handle()
 543                    .toggle(window, cx);
 544            }))
 545            .on_action(cx.listener(|this, _: &ToggleModelSelector, window, cx| {
 546                this.model_selector
 547                    .update(cx, |model_selector, cx| model_selector.toggle(window, cx));
 548            }))
 549            .on_action(cx.listener(Self::toggle_context_picker))
 550            .on_action(cx.listener(Self::remove_all_context))
 551            .on_action(cx.listener(Self::move_up))
 552            .on_action(cx.listener(Self::expand_message_editor))
 553            .capture_action(cx.listener(Self::paste))
 554            .gap_2()
 555            .p_2()
 556            .bg(editor_bg_color)
 557            .border_t_1()
 558            .border_color(cx.theme().colors().border)
 559            .child(
 560                h_flex()
 561                    .items_start()
 562                    .justify_between()
 563                    .child(self.context_strip.clone())
 564                    .child(
 565                        h_flex()
 566                            .gap_1()
 567                            .when(focus_handle.is_focused(window), |this| {
 568                                this.child(
 569                                    IconButton::new("toggle-height", expand_icon)
 570                                        .icon_size(IconSize::XSmall)
 571                                        .icon_color(Color::Muted)
 572                                        .tooltip({
 573                                            let focus_handle = focus_handle.clone();
 574                                            move |window, cx| {
 575                                                let expand_label = if is_editor_expanded {
 576                                                    "Minimize Message Editor".to_string()
 577                                                } else {
 578                                                    "Expand Message Editor".to_string()
 579                                                };
 580
 581                                                Tooltip::for_action_in(
 582                                                    expand_label,
 583                                                    &ExpandMessageEditor,
 584                                                    &focus_handle,
 585                                                    window,
 586                                                    cx,
 587                                                )
 588                                            }
 589                                        })
 590                                        .on_click(cx.listener(|_, _, window, cx| {
 591                                            window
 592                                                .dispatch_action(Box::new(ExpandMessageEditor), cx);
 593                                        })),
 594                                )
 595                            }),
 596                    ),
 597            )
 598            .child(
 599                v_flex()
 600                    .size_full()
 601                    .gap_4()
 602                    .when(is_editor_expanded, |this| {
 603                        this.h(vh(0.8, window)).justify_between()
 604                    })
 605                    .child(
 606                        div()
 607                            .min_h_16()
 608                            .when(is_editor_expanded, |this| this.h_full())
 609                            .child({
 610                                let settings = ThemeSettings::get_global(cx);
 611
 612                                let text_style = TextStyle {
 613                                    color: cx.theme().colors().text,
 614                                    font_family: settings.buffer_font.family.clone(),
 615                                    font_fallbacks: settings.buffer_font.fallbacks.clone(),
 616                                    font_features: settings.buffer_font.features.clone(),
 617                                    font_size: font_size.into(),
 618                                    line_height: line_height.into(),
 619                                    ..Default::default()
 620                                };
 621
 622                                EditorElement::new(
 623                                    &self.editor,
 624                                    EditorStyle {
 625                                        background: editor_bg_color,
 626                                        local_player: cx.theme().players().local(),
 627                                        text: text_style,
 628                                        syntax: cx.theme().syntax().clone(),
 629                                        ..Default::default()
 630                                    },
 631                                )
 632                                .into_any()
 633                            }),
 634                    )
 635                    .child(
 636                        h_flex()
 637                            .flex_none()
 638                            .justify_between()
 639                            .child(
 640                                h_flex()
 641                                    .gap_1()
 642                                    .child(self.render_follow_toggle(cx))
 643                                    .children(self.render_max_mode_toggle(cx)),
 644                            )
 645                            .child(
 646                                h_flex()
 647                                    .gap_1()
 648                                    .when(!incompatible_tools.is_empty(), |this| {
 649                                        this.child(
 650                                            IconButton::new(
 651                                                "tools-incompatible-warning",
 652                                                IconName::Warning,
 653                                            )
 654                                            .icon_color(Color::Warning)
 655                                            .icon_size(IconSize::Small)
 656                                            .tooltip({
 657                                                move |_, cx| {
 658                                                    cx.new(|_| IncompatibleToolsTooltip {
 659                                                        incompatible_tools: incompatible_tools
 660                                                            .clone(),
 661                                                    })
 662                                                    .into()
 663                                                }
 664                                            }),
 665                                        )
 666                                    })
 667                                    .child(self.profile_selector.clone())
 668                                    .child(self.model_selector.clone())
 669                                    .map({
 670                                        let focus_handle = focus_handle.clone();
 671                                        move |parent| {
 672                                            if is_generating {
 673                                                parent
 674                                                    .when(is_editor_empty, |parent| {
 675                                                        parent.child(
 676                                                            IconButton::new(
 677                                                                "stop-generation",
 678                                                                IconName::StopFilled,
 679                                                            )
 680                                                            .icon_color(Color::Error)
 681                                                            .style(ButtonStyle::Tinted(
 682                                                                ui::TintColor::Error,
 683                                                            ))
 684                                                            .tooltip(move |window, cx| {
 685                                                                Tooltip::for_action(
 686                                                                    "Stop Generation",
 687                                                                    &editor::actions::Cancel,
 688                                                                    window,
 689                                                                    cx,
 690                                                                )
 691                                                            })
 692                                                            .on_click({
 693                                                                let focus_handle =
 694                                                                    focus_handle.clone();
 695                                                                move |_event, window, cx| {
 696                                                                    focus_handle.dispatch_action(
 697                                                                        &editor::actions::Cancel,
 698                                                                        window,
 699                                                                        cx,
 700                                                                    );
 701                                                                }
 702                                                            })
 703                                                            .with_animation(
 704                                                                "pulsating-label",
 705                                                                Animation::new(
 706                                                                    Duration::from_secs(2),
 707                                                                )
 708                                                                .repeat()
 709                                                                .with_easing(pulsating_between(
 710                                                                    0.4, 1.0,
 711                                                                )),
 712                                                                |icon_button, delta| {
 713                                                                    icon_button.alpha(delta)
 714                                                                },
 715                                                            ),
 716                                                        )
 717                                                    })
 718                                                    .when(!is_editor_empty, |parent| {
 719                                                        parent.child(
 720                                                            IconButton::new(
 721                                                                "send-message",
 722                                                                IconName::Send,
 723                                                            )
 724                                                            .icon_color(Color::Accent)
 725                                                            .style(ButtonStyle::Filled)
 726                                                            .disabled(!is_model_selected)
 727                                                            .on_click({
 728                                                                let focus_handle =
 729                                                                    focus_handle.clone();
 730                                                                move |_event, window, cx| {
 731                                                                    focus_handle.dispatch_action(
 732                                                                        &Chat, window, cx,
 733                                                                    );
 734                                                                }
 735                                                            })
 736                                                            .tooltip(move |window, cx| {
 737                                                                Tooltip::for_action(
 738                                                                    "Stop and Send New Message",
 739                                                                    &Chat,
 740                                                                    window,
 741                                                                    cx,
 742                                                                )
 743                                                            }),
 744                                                        )
 745                                                    })
 746                                            } else {
 747                                                parent.child(
 748                                                    IconButton::new("send-message", IconName::Send)
 749                                                        .icon_color(Color::Accent)
 750                                                        .style(ButtonStyle::Filled)
 751                                                        .disabled(
 752                                                            is_editor_empty || !is_model_selected,
 753                                                        )
 754                                                        .on_click({
 755                                                            let focus_handle = focus_handle.clone();
 756                                                            move |_event, window, cx| {
 757                                                                focus_handle.dispatch_action(
 758                                                                    &Chat, window, cx,
 759                                                                );
 760                                                            }
 761                                                        })
 762                                                        .when(
 763                                                            !is_editor_empty && is_model_selected,
 764                                                            |button| {
 765                                                                button.tooltip(move |window, cx| {
 766                                                                    Tooltip::for_action(
 767                                                                        "Send", &Chat, window, cx,
 768                                                                    )
 769                                                                })
 770                                                            },
 771                                                        )
 772                                                        .when(is_editor_empty, |button| {
 773                                                            button.tooltip(Tooltip::text(
 774                                                                "Type a message to submit",
 775                                                            ))
 776                                                        })
 777                                                        .when(!is_model_selected, |button| {
 778                                                            button.tooltip(Tooltip::text(
 779                                                                "Select a model to continue",
 780                                                            ))
 781                                                        }),
 782                                                )
 783                                            }
 784                                        }
 785                                    }),
 786                            ),
 787                    ),
 788            )
 789    }
 790
 791    fn render_changed_buffers(
 792        &self,
 793        changed_buffers: &BTreeMap<Entity<Buffer>, Entity<BufferDiff>>,
 794        window: &mut Window,
 795        cx: &mut Context<Self>,
 796    ) -> Div {
 797        let focus_handle = self.editor.focus_handle(cx);
 798
 799        let editor_bg_color = cx.theme().colors().editor_background;
 800        let border_color = cx.theme().colors().border;
 801        let active_color = cx.theme().colors().element_selected;
 802        let bg_edit_files_disclosure = editor_bg_color.blend(active_color.opacity(0.3));
 803
 804        let is_edit_changes_expanded = self.edits_expanded;
 805        let is_generating = self.thread.read(cx).is_generating();
 806
 807        v_flex()
 808            .mt_1()
 809            .mx_2()
 810            .bg(bg_edit_files_disclosure)
 811            .border_1()
 812            .border_b_0()
 813            .border_color(border_color)
 814            .rounded_t_md()
 815            .shadow(smallvec::smallvec![gpui::BoxShadow {
 816                color: gpui::black().opacity(0.15),
 817                offset: point(px(1.), px(-1.)),
 818                blur_radius: px(3.),
 819                spread_radius: px(0.),
 820            }])
 821            .child(
 822                h_flex()
 823                    .id("edits-container")
 824                    .cursor_pointer()
 825                    .p_1p5()
 826                    .justify_between()
 827                    .when(is_edit_changes_expanded, |this| {
 828                        this.border_b_1().border_color(border_color)
 829                    })
 830                    .on_click(
 831                        cx.listener(|this, _, window, cx| this.handle_review_click(window, cx)),
 832                    )
 833                    .child(
 834                        h_flex()
 835                            .gap_1()
 836                            .child(
 837                                Disclosure::new("edits-disclosure", is_edit_changes_expanded)
 838                                    .on_click(cx.listener(|this, _ev, _window, cx| {
 839                                        this.edits_expanded = !this.edits_expanded;
 840                                        cx.notify();
 841                                    })),
 842                            )
 843                            .map(|this| {
 844                                if is_generating {
 845                                    this.child(
 846                                        AnimatedLabel::new(format!(
 847                                            "Editing {} {}",
 848                                            changed_buffers.len(),
 849                                            if changed_buffers.len() == 1 {
 850                                                "file"
 851                                            } else {
 852                                                "files"
 853                                            }
 854                                        ))
 855                                        .size(LabelSize::Small),
 856                                    )
 857                                } else {
 858                                    this.child(
 859                                        Label::new("Edits")
 860                                            .size(LabelSize::Small)
 861                                            .color(Color::Muted),
 862                                    )
 863                                    .child(
 864                                        Label::new("").size(LabelSize::XSmall).color(Color::Muted),
 865                                    )
 866                                    .child(
 867                                        Label::new(format!(
 868                                            "{} {}",
 869                                            changed_buffers.len(),
 870                                            if changed_buffers.len() == 1 {
 871                                                "file"
 872                                            } else {
 873                                                "files"
 874                                            }
 875                                        ))
 876                                        .size(LabelSize::Small)
 877                                        .color(Color::Muted),
 878                                    )
 879                                }
 880                            }),
 881                    )
 882                    .child(
 883                        Button::new("review", "Review Changes")
 884                            .label_size(LabelSize::Small)
 885                            .key_binding(
 886                                KeyBinding::for_action_in(
 887                                    &OpenAgentDiff,
 888                                    &focus_handle,
 889                                    window,
 890                                    cx,
 891                                )
 892                                .map(|kb| kb.size(rems_from_px(12.))),
 893                            )
 894                            .on_click(cx.listener(|this, _, window, cx| {
 895                                this.handle_review_click(window, cx)
 896                            })),
 897                    ),
 898            )
 899            .when(is_edit_changes_expanded, |parent| {
 900                parent.child(
 901                    v_flex().children(changed_buffers.into_iter().enumerate().flat_map(
 902                        |(index, (buffer, _diff))| {
 903                            let file = buffer.read(cx).file()?;
 904                            let path = file.path();
 905
 906                            let parent_label = path.parent().and_then(|parent| {
 907                                let parent_str = parent.to_string_lossy();
 908
 909                                if parent_str.is_empty() {
 910                                    None
 911                                } else {
 912                                    Some(
 913                                        Label::new(format!(
 914                                            "/{}{}",
 915                                            parent_str,
 916                                            std::path::MAIN_SEPARATOR_STR
 917                                        ))
 918                                        .color(Color::Muted)
 919                                        .size(LabelSize::XSmall)
 920                                        .buffer_font(cx),
 921                                    )
 922                                }
 923                            });
 924
 925                            let name_label = path.file_name().map(|name| {
 926                                Label::new(name.to_string_lossy().to_string())
 927                                    .size(LabelSize::XSmall)
 928                                    .buffer_font(cx)
 929                            });
 930
 931                            let file_icon = FileIcons::get_icon(&path, cx)
 932                                .map(Icon::from_path)
 933                                .map(|icon| icon.color(Color::Muted).size(IconSize::Small))
 934                                .unwrap_or_else(|| {
 935                                    Icon::new(IconName::File)
 936                                        .color(Color::Muted)
 937                                        .size(IconSize::Small)
 938                                });
 939
 940                            let hover_color = cx
 941                                .theme()
 942                                .colors()
 943                                .element_background
 944                                .blend(cx.theme().colors().editor_foreground.opacity(0.025));
 945
 946                            let overlay_gradient = linear_gradient(
 947                                90.,
 948                                linear_color_stop(editor_bg_color, 1.),
 949                                linear_color_stop(editor_bg_color.opacity(0.2), 0.),
 950                            );
 951
 952                            let overlay_gradient_hover = linear_gradient(
 953                                90.,
 954                                linear_color_stop(hover_color, 1.),
 955                                linear_color_stop(hover_color.opacity(0.2), 0.),
 956                            );
 957
 958                            let element = h_flex()
 959                                .group("edited-code")
 960                                .id(("file-container", index))
 961                                .cursor_pointer()
 962                                .relative()
 963                                .py_1()
 964                                .pl_2()
 965                                .pr_1()
 966                                .gap_2()
 967                                .justify_between()
 968                                .bg(cx.theme().colors().editor_background)
 969                                .hover(|style| style.bg(hover_color))
 970                                .when(index < changed_buffers.len() - 1, |parent| {
 971                                    parent.border_color(border_color).border_b_1()
 972                                })
 973                                .child(
 974                                    h_flex()
 975                                        .id("file-name")
 976                                        .pr_8()
 977                                        .gap_1p5()
 978                                        .max_w_full()
 979                                        .overflow_x_scroll()
 980                                        .child(file_icon)
 981                                        .child(
 982                                            h_flex()
 983                                                .gap_0p5()
 984                                                .children(name_label)
 985                                                .children(parent_label),
 986                                        ), // TODO: Implement line diff
 987                                           // .child(Label::new("+").color(Color::Created))
 988                                           // .child(Label::new("-").color(Color::Deleted)),
 989                                )
 990                                .child(
 991                                    div().visible_on_hover("edited-code").child(
 992                                        Button::new("review", "Review")
 993                                            .label_size(LabelSize::Small)
 994                                            .on_click({
 995                                                let buffer = buffer.clone();
 996                                                cx.listener(move |this, _, window, cx| {
 997                                                    this.handle_file_click(
 998                                                        buffer.clone(),
 999                                                        window,
1000                                                        cx,
1001                                                    );
1002                                                })
1003                                            }),
1004                                    ),
1005                                )
1006                                .child(
1007                                    div()
1008                                        .id("gradient-overlay")
1009                                        .absolute()
1010                                        .h_5_6()
1011                                        .w_12()
1012                                        .bottom_0()
1013                                        .right(px(52.))
1014                                        .bg(overlay_gradient)
1015                                        .group_hover("edited-code", |style| {
1016                                            style.bg(overlay_gradient_hover)
1017                                        }),
1018                                )
1019                                .on_click({
1020                                    let buffer = buffer.clone();
1021                                    cx.listener(move |this, _, window, cx| {
1022                                        this.handle_file_click(buffer.clone(), window, cx);
1023                                    })
1024                                });
1025
1026                            Some(element)
1027                        },
1028                    )),
1029                )
1030            })
1031    }
1032
1033    fn render_token_limit_callout(
1034        &self,
1035        line_height: Pixels,
1036        token_usage_ratio: TokenUsageRatio,
1037        cx: &mut Context<Self>,
1038    ) -> Div {
1039        let heading = if token_usage_ratio == TokenUsageRatio::Exceeded {
1040            "Thread reached the token limit"
1041        } else {
1042            "Thread reaching the token limit soon"
1043        };
1044
1045        h_flex()
1046            .p_2()
1047            .gap_2()
1048            .flex_wrap()
1049            .justify_between()
1050            .bg(
1051                if token_usage_ratio == TokenUsageRatio::Exceeded {
1052                    cx.theme().status().error_background.opacity(0.1)
1053                } else {
1054                    cx.theme().status().warning_background.opacity(0.1)
1055                })
1056            .border_t_1()
1057            .border_color(cx.theme().colors().border)
1058            .child(
1059                h_flex()
1060                    .gap_2()
1061                    .items_start()
1062                    .child(
1063                        h_flex()
1064                            .h(line_height)
1065                            .justify_center()
1066                            .child(
1067                                if token_usage_ratio == TokenUsageRatio::Exceeded {
1068                                    Icon::new(IconName::X)
1069                                        .color(Color::Error)
1070                                        .size(IconSize::XSmall)
1071                                } else {
1072                                    Icon::new(IconName::Warning)
1073                                        .color(Color::Warning)
1074                                        .size(IconSize::XSmall)
1075                                }
1076                            ),
1077                    )
1078                    .child(
1079                        v_flex()
1080                            .mr_auto()
1081                            .child(Label::new(heading).size(LabelSize::Small))
1082                            .child(
1083                                Label::new(
1084                                    "Start a new thread from a summary to continue the conversation.",
1085                                )
1086                                .size(LabelSize::Small)
1087                                .color(Color::Muted),
1088                            ),
1089                    ),
1090            )
1091            .child(
1092                Button::new("new-thread", "Start New Thread")
1093                    .on_click(cx.listener(|this, _, window, cx| {
1094                        let from_thread_id = Some(this.thread.read(cx).id().clone());
1095
1096                        window.dispatch_action(Box::new(NewThread {
1097                            from_thread_id
1098                        }), cx);
1099                    }))
1100                    .icon(IconName::Plus)
1101                    .icon_position(IconPosition::Start)
1102                    .icon_size(IconSize::Small)
1103                    .style(ButtonStyle::Tinted(ui::TintColor::Accent))
1104                    .label_size(LabelSize::Small),
1105            )
1106    }
1107
1108    pub fn last_estimated_token_count(&self) -> Option<usize> {
1109        self.last_estimated_token_count
1110    }
1111
1112    pub fn is_waiting_to_update_token_count(&self) -> bool {
1113        self.update_token_count_task.is_some()
1114    }
1115
1116    fn reload_context(&mut self, cx: &mut Context<Self>) -> Task<Option<ContextLoadResult>> {
1117        let load_task = cx.spawn(async move |this, cx| {
1118            let Ok(load_task) = this.update(cx, |this, cx| {
1119                let new_context = this.context_store.read_with(cx, |context_store, cx| {
1120                    context_store.new_context_for_thread(this.thread.read(cx), None)
1121                });
1122                load_context(new_context, &this.project, &this.prompt_store, cx)
1123            }) else {
1124                return;
1125            };
1126            let result = load_task.await;
1127            this.update(cx, |this, cx| {
1128                this.last_loaded_context = Some(result);
1129                this.load_context_task = None;
1130                this.message_or_context_changed(false, cx);
1131            })
1132            .ok();
1133        });
1134        // Replace existing load task, if any, causing it to be cancelled.
1135        let load_task = load_task.shared();
1136        self.load_context_task = Some(load_task.clone());
1137        cx.spawn(async move |this, cx| {
1138            load_task.await;
1139            this.read_with(cx, |this, _cx| this.last_loaded_context.clone())
1140                .ok()
1141                .flatten()
1142        })
1143    }
1144
1145    fn handle_message_changed(&mut self, cx: &mut Context<Self>) {
1146        self.message_or_context_changed(true, cx);
1147    }
1148
1149    fn message_or_context_changed(&mut self, debounce: bool, cx: &mut Context<Self>) {
1150        cx.emit(MessageEditorEvent::Changed);
1151        self.update_token_count_task.take();
1152
1153        let Some(model) = self.thread.read(cx).configured_model() else {
1154            self.last_estimated_token_count.take();
1155            return;
1156        };
1157
1158        let editor = self.editor.clone();
1159
1160        self.update_token_count_task = Some(cx.spawn(async move |this, cx| {
1161            if debounce {
1162                cx.background_executor()
1163                    .timer(Duration::from_millis(200))
1164                    .await;
1165            }
1166
1167            let token_count = if let Some(task) = this
1168                .update(cx, |this, cx| {
1169                    let loaded_context = this
1170                        .last_loaded_context
1171                        .as_ref()
1172                        .map(|context_load_result| &context_load_result.loaded_context);
1173                    let message_text = editor.read(cx).text(cx);
1174
1175                    if message_text.is_empty()
1176                        && loaded_context.map_or(true, |loaded_context| loaded_context.is_empty())
1177                    {
1178                        return None;
1179                    }
1180
1181                    let mut request_message = LanguageModelRequestMessage {
1182                        role: language_model::Role::User,
1183                        content: Vec::new(),
1184                        cache: false,
1185                    };
1186
1187                    if let Some(loaded_context) = loaded_context {
1188                        loaded_context.add_to_request_message(&mut request_message);
1189                    }
1190
1191                    if !message_text.is_empty() {
1192                        request_message
1193                            .content
1194                            .push(MessageContent::Text(message_text));
1195                    }
1196
1197                    let request = language_model::LanguageModelRequest {
1198                        thread_id: None,
1199                        prompt_id: None,
1200                        mode: None,
1201                        messages: vec![request_message],
1202                        tools: vec![],
1203                        stop: vec![],
1204                        temperature: None,
1205                    };
1206
1207                    Some(model.model.count_tokens(request, cx))
1208                })
1209                .ok()
1210                .flatten()
1211            {
1212                task.await.log_err()
1213            } else {
1214                Some(0)
1215            };
1216
1217            this.update(cx, |this, cx| {
1218                if let Some(token_count) = token_count {
1219                    this.last_estimated_token_count = Some(token_count);
1220                    cx.emit(MessageEditorEvent::EstimatedTokenCount);
1221                }
1222                this.update_token_count_task.take();
1223            })
1224            .ok();
1225        }));
1226    }
1227}
1228
1229pub fn extract_message_creases(
1230    editor: &mut Editor,
1231    cx: &mut Context<'_, Editor>,
1232) -> Vec<MessageCrease> {
1233    let buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
1234    let mut contexts_by_crease_id = editor
1235        .addon_mut::<ContextCreasesAddon>()
1236        .map(std::mem::take)
1237        .unwrap_or_default()
1238        .into_inner()
1239        .into_iter()
1240        .flat_map(|(key, creases)| {
1241            let context = key.0;
1242            creases
1243                .into_iter()
1244                .map(move |(id, _)| (id, context.clone()))
1245        })
1246        .collect::<HashMap<_, _>>();
1247    // Filter the addon's list of creases based on what the editor reports,
1248    // since the addon might have removed creases in it.
1249    let creases = editor.display_map.update(cx, |display_map, cx| {
1250        display_map
1251            .snapshot(cx)
1252            .crease_snapshot
1253            .creases()
1254            .filter_map(|(id, crease)| {
1255                Some((
1256                    id,
1257                    (
1258                        crease.range().to_offset(&buffer_snapshot),
1259                        crease.metadata()?.clone(),
1260                    ),
1261                ))
1262            })
1263            .map(|(id, (range, metadata))| {
1264                let context = contexts_by_crease_id.remove(&id);
1265                MessageCrease {
1266                    range,
1267                    metadata,
1268                    context,
1269                }
1270            })
1271            .collect()
1272    });
1273    creases
1274}
1275
1276impl EventEmitter<MessageEditorEvent> for MessageEditor {}
1277
1278pub enum MessageEditorEvent {
1279    EstimatedTokenCount,
1280    Changed,
1281}
1282
1283impl Focusable for MessageEditor {
1284    fn focus_handle(&self, cx: &App) -> gpui::FocusHandle {
1285        self.editor.focus_handle(cx)
1286    }
1287}
1288
1289impl Render for MessageEditor {
1290    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1291        let thread = self.thread.read(cx);
1292        let token_usage_ratio = thread
1293            .total_token_usage()
1294            .map_or(TokenUsageRatio::Normal, |total_token_usage| {
1295                total_token_usage.ratio()
1296            });
1297
1298        let action_log = self.thread.read(cx).action_log();
1299        let changed_buffers = action_log.read(cx).changed_buffers(cx);
1300
1301        let font_size = TextSize::Small.rems(cx);
1302        let line_height = font_size.to_pixels(window.rem_size()) * 1.5;
1303
1304        v_flex()
1305            .size_full()
1306            .when(changed_buffers.len() > 0, |parent| {
1307                parent.child(self.render_changed_buffers(&changed_buffers, window, cx))
1308            })
1309            .child(self.render_editor(font_size, line_height, window, cx))
1310            .when(token_usage_ratio != TokenUsageRatio::Normal, |parent| {
1311                parent.child(self.render_token_limit_callout(line_height, token_usage_ratio, cx))
1312            })
1313    }
1314}
1315
1316pub fn insert_message_creases(
1317    editor: &mut Editor,
1318    message_creases: &[MessageCrease],
1319    context_store: &Entity<ContextStore>,
1320    window: &mut Window,
1321    cx: &mut Context<'_, Editor>,
1322) {
1323    let buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
1324    let creases = message_creases
1325        .iter()
1326        .map(|crease| {
1327            let start = buffer_snapshot.anchor_after(crease.range.start);
1328            let end = buffer_snapshot.anchor_before(crease.range.end);
1329            crease_for_mention(
1330                crease.metadata.label.clone(),
1331                crease.metadata.icon_path.clone(),
1332                start..end,
1333                cx.weak_entity(),
1334            )
1335        })
1336        .collect::<Vec<_>>();
1337    let ids = editor.insert_creases(creases.clone(), cx);
1338    editor.fold_creases(creases, false, window, cx);
1339    if let Some(addon) = editor.addon_mut::<ContextCreasesAddon>() {
1340        for (crease, id) in message_creases.iter().zip(ids) {
1341            if let Some(context) = crease.context.as_ref() {
1342                let key = AgentContextKey(context.clone());
1343                addon.add_creases(
1344                    context_store,
1345                    key,
1346                    vec![(id, crease.metadata.label.clone())],
1347                    cx,
1348                );
1349            }
1350        }
1351    }
1352}
1353impl Component for MessageEditor {
1354    fn scope() -> ComponentScope {
1355        ComponentScope::Agent
1356    }
1357}
1358
1359impl AgentPreview for MessageEditor {
1360    fn agent_preview(
1361        workspace: WeakEntity<Workspace>,
1362        active_thread: Entity<ActiveThread>,
1363        thread_store: WeakEntity<ThreadStore>,
1364        window: &mut Window,
1365        cx: &mut App,
1366    ) -> Option<AnyElement> {
1367        if let Some(workspace_entity) = workspace.upgrade() {
1368            let fs = workspace_entity.read(cx).app_state().fs.clone();
1369            let weak_project = workspace_entity.read(cx).project().clone().downgrade();
1370            let context_store = cx.new(|_cx| ContextStore::new(weak_project, None));
1371            let thread = active_thread.read(cx).thread().clone();
1372
1373            let example_message_editor = cx.new(|cx| {
1374                MessageEditor::new(
1375                    fs,
1376                    workspace,
1377                    context_store,
1378                    None,
1379                    thread_store,
1380                    thread,
1381                    window,
1382                    cx,
1383                )
1384            });
1385
1386            Some(
1387                v_flex()
1388                    .gap_4()
1389                    .children(vec![single_example(
1390                        "Default",
1391                        example_message_editor.clone().into_any_element(),
1392                    )])
1393                    .into_any_element(),
1394            )
1395        } else {
1396            None
1397        }
1398    }
1399}
1400
1401register_agent_preview!(MessageEditor);