message_editor.rs

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