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