message_editor.rs

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