1use crate::assistant_model_selector::{AssistantModelSelector, ModelType};
2use crate::buffer_codegen::BufferCodegen;
3use crate::context::ContextCreasesAddon;
4use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider};
5use crate::context_store::ContextStore;
6use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
7use crate::message_editor::{extract_message_creases, insert_message_creases};
8use crate::terminal_codegen::TerminalCodegen;
9use crate::thread_store::{TextThreadStore, ThreadStore};
10use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist};
11use crate::{RemoveAllContext, ToggleContextPicker};
12use client::ErrorExt;
13use collections::VecDeque;
14use editor::{
15 ContextMenuOptions, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle,
16 GutterDimensions, MultiBuffer,
17 actions::{MoveDown, MoveUp},
18};
19use feature_flags::{FeatureFlagAppExt as _, ZedProFeatureFlag};
20use fs::Fs;
21use gpui::{
22 AnyElement, App, ClickEvent, Context, CursorStyle, Entity, EventEmitter, FocusHandle,
23 Focusable, FontWeight, Subscription, TextStyle, WeakEntity, Window, anchored, deferred, point,
24};
25use language_model::{LanguageModel, LanguageModelRegistry};
26use language_model_selector::ToggleModelSelector;
27use parking_lot::Mutex;
28use settings::Settings;
29use std::cmp;
30use std::sync::Arc;
31use theme::ThemeSettings;
32use ui::utils::WithRemSize;
33use ui::{
34 CheckboxWithLabel, IconButtonShape, KeyBinding, Popover, PopoverMenuHandle, Tooltip, prelude::*,
35};
36use util::ResultExt;
37use workspace::Workspace;
38
39pub struct PromptEditor<T> {
40 pub editor: Entity<Editor>,
41 mode: PromptEditorMode,
42 context_store: Entity<ContextStore>,
43 context_strip: Entity<ContextStrip>,
44 context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
45 model_selector: Entity<AssistantModelSelector>,
46 edited_since_done: bool,
47 prompt_history: VecDeque<String>,
48 prompt_history_ix: Option<usize>,
49 pending_prompt: String,
50 _codegen_subscription: Subscription,
51 editor_subscriptions: Vec<Subscription>,
52 _context_strip_subscription: Subscription,
53 show_rate_limit_notice: bool,
54 _phantom: std::marker::PhantomData<T>,
55}
56
57impl<T: 'static> EventEmitter<PromptEditorEvent> for PromptEditor<T> {}
58
59impl<T: 'static> Render for PromptEditor<T> {
60 fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
61 let ui_font_size = ThemeSettings::get_global(cx).ui_font_size(cx);
62 let mut buttons = Vec::new();
63
64 let left_gutter_width = match &self.mode {
65 PromptEditorMode::Buffer {
66 id: _,
67 codegen,
68 gutter_dimensions,
69 } => {
70 let codegen = codegen.read(cx);
71
72 if codegen.alternative_count(cx) > 1 {
73 buttons.push(self.render_cycle_controls(&codegen, cx));
74 }
75
76 let gutter_dimensions = gutter_dimensions.lock();
77
78 gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0)
79 }
80 PromptEditorMode::Terminal { .. } => {
81 // Give the equivalent of the same left-padding that we're using on the right
82 Pixels::from(40.0)
83 }
84 };
85
86 let bottom_padding = match &self.mode {
87 PromptEditorMode::Buffer { .. } => Pixels::from(0.),
88 PromptEditorMode::Terminal { .. } => Pixels::from(8.0),
89 };
90
91 buttons.extend(self.render_buttons(window, cx));
92
93 v_flex()
94 .key_context("PromptEditor")
95 .bg(cx.theme().colors().editor_background)
96 .block_mouse_down()
97 .gap_0p5()
98 .border_y_1()
99 .border_color(cx.theme().status().info_border)
100 .size_full()
101 .pt_0p5()
102 .pb(bottom_padding)
103 .pr_6()
104 .child(
105 h_flex()
106 .items_start()
107 .cursor(CursorStyle::Arrow)
108 .on_action(cx.listener(Self::toggle_context_picker))
109 .on_action(cx.listener(|this, _: &ToggleModelSelector, window, cx| {
110 this.model_selector
111 .update(cx, |model_selector, cx| model_selector.toggle(window, cx));
112 }))
113 .on_action(cx.listener(Self::confirm))
114 .on_action(cx.listener(Self::cancel))
115 .on_action(cx.listener(Self::move_up))
116 .on_action(cx.listener(Self::move_down))
117 .on_action(cx.listener(Self::remove_all_context))
118 .capture_action(cx.listener(Self::cycle_prev))
119 .capture_action(cx.listener(Self::cycle_next))
120 .child(
121 WithRemSize::new(ui_font_size)
122 .flex()
123 .flex_row()
124 .flex_shrink_0()
125 .items_center()
126 .h_full()
127 .w(left_gutter_width)
128 .justify_center()
129 .gap_2()
130 .child(self.render_close_button(cx))
131 .map(|el| {
132 let CodegenStatus::Error(error) = self.codegen_status(cx) else {
133 return el;
134 };
135
136 let error_message = SharedString::from(error.to_string());
137 if error.error_code() == proto::ErrorCode::RateLimitExceeded
138 && cx.has_flag::<ZedProFeatureFlag>()
139 {
140 el.child(
141 v_flex()
142 .child(
143 IconButton::new(
144 "rate-limit-error",
145 IconName::XCircle,
146 )
147 .toggle_state(self.show_rate_limit_notice)
148 .shape(IconButtonShape::Square)
149 .icon_size(IconSize::Small)
150 .on_click(
151 cx.listener(Self::toggle_rate_limit_notice),
152 ),
153 )
154 .children(self.show_rate_limit_notice.then(|| {
155 deferred(
156 anchored()
157 .position_mode(
158 gpui::AnchoredPositionMode::Local,
159 )
160 .position(point(px(0.), px(24.)))
161 .anchor(gpui::Corner::TopLeft)
162 .child(self.render_rate_limit_notice(cx)),
163 )
164 })),
165 )
166 } else {
167 el.child(
168 div()
169 .id("error")
170 .tooltip(Tooltip::text(error_message))
171 .child(
172 Icon::new(IconName::XCircle)
173 .size(IconSize::Small)
174 .color(Color::Error),
175 ),
176 )
177 }
178 }),
179 )
180 .child(
181 h_flex()
182 .w_full()
183 .justify_between()
184 .child(div().flex_1().child(self.render_editor(window, cx)))
185 .child(
186 WithRemSize::new(ui_font_size)
187 .flex()
188 .flex_row()
189 .items_center()
190 .gap_1()
191 .children(buttons),
192 ),
193 ),
194 )
195 .child(
196 WithRemSize::new(ui_font_size)
197 .flex()
198 .flex_row()
199 .items_center()
200 .child(h_flex().flex_shrink_0().w(left_gutter_width))
201 .child(
202 h_flex()
203 .w_full()
204 .pl_1()
205 .items_start()
206 .justify_between()
207 .child(self.context_strip.clone())
208 .child(self.model_selector.clone()),
209 ),
210 )
211 }
212}
213
214impl<T: 'static> Focusable for PromptEditor<T> {
215 fn focus_handle(&self, cx: &App) -> FocusHandle {
216 self.editor.focus_handle(cx)
217 }
218}
219
220impl<T: 'static> PromptEditor<T> {
221 const MAX_LINES: u8 = 8;
222
223 fn codegen_status<'a>(&'a self, cx: &'a App) -> &'a CodegenStatus {
224 match &self.mode {
225 PromptEditorMode::Buffer { codegen, .. } => codegen.read(cx).status(cx),
226 PromptEditorMode::Terminal { codegen, .. } => &codegen.read(cx).status,
227 }
228 }
229
230 fn subscribe_to_editor(&mut self, window: &mut Window, cx: &mut Context<Self>) {
231 self.editor_subscriptions.clear();
232 self.editor_subscriptions.push(cx.subscribe_in(
233 &self.editor,
234 window,
235 Self::handle_prompt_editor_events,
236 ));
237 }
238
239 pub fn set_show_cursor_when_unfocused(
240 &mut self,
241 show_cursor_when_unfocused: bool,
242 cx: &mut Context<Self>,
243 ) {
244 self.editor.update(cx, |editor, cx| {
245 editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
246 });
247 }
248
249 pub fn unlink(&mut self, window: &mut Window, cx: &mut Context<Self>) {
250 let prompt = self.prompt(cx);
251 let existing_creases = self.editor.update(cx, extract_message_creases);
252
253 let focus = self.editor.focus_handle(cx).contains_focused(window, cx);
254 self.editor = cx.new(|cx| {
255 let mut editor = Editor::auto_height(Self::MAX_LINES as usize, window, cx);
256 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
257 editor.set_placeholder_text("Add a prompt…", cx);
258 editor.set_text(prompt, window, cx);
259 insert_message_creases(
260 &mut editor,
261 &existing_creases,
262 &self.context_store,
263 window,
264 cx,
265 );
266
267 if focus {
268 window.focus(&editor.focus_handle(cx));
269 }
270 editor
271 });
272 self.subscribe_to_editor(window, cx);
273 }
274
275 pub fn placeholder_text(mode: &PromptEditorMode, window: &mut Window, cx: &mut App) -> String {
276 let action = match mode {
277 PromptEditorMode::Buffer { codegen, .. } => {
278 if codegen.read(cx).is_insertion {
279 "Generate"
280 } else {
281 "Transform"
282 }
283 }
284 PromptEditorMode::Terminal { .. } => "Generate",
285 };
286
287 let assistant_panel_keybinding =
288 ui::text_for_action(&zed_actions::assistant::ToggleFocus, window, cx)
289 .map(|keybinding| format!("{keybinding} to chat ― "))
290 .unwrap_or_default();
291
292 format!("{action}… ({assistant_panel_keybinding}↓↑ for history)")
293 }
294
295 pub fn prompt(&self, cx: &App) -> String {
296 self.editor.read(cx).text(cx)
297 }
298
299 fn toggle_rate_limit_notice(
300 &mut self,
301 _: &ClickEvent,
302 window: &mut Window,
303 cx: &mut Context<Self>,
304 ) {
305 self.show_rate_limit_notice = !self.show_rate_limit_notice;
306 if self.show_rate_limit_notice {
307 window.focus(&self.editor.focus_handle(cx));
308 }
309 cx.notify();
310 }
311
312 fn handle_prompt_editor_events(
313 &mut self,
314 _: &Entity<Editor>,
315 event: &EditorEvent,
316 window: &mut Window,
317 cx: &mut Context<Self>,
318 ) {
319 match event {
320 EditorEvent::Edited { .. } => {
321 if let Some(workspace) = window.root::<Workspace>().flatten() {
322 workspace.update(cx, |workspace, cx| {
323 let is_via_ssh = workspace
324 .project()
325 .update(cx, |project, _| project.is_via_ssh());
326
327 workspace
328 .client()
329 .telemetry()
330 .log_edit_event("inline assist", is_via_ssh);
331 });
332 }
333 let prompt = self.editor.read(cx).text(cx);
334 if self
335 .prompt_history_ix
336 .map_or(true, |ix| self.prompt_history[ix] != prompt)
337 {
338 self.prompt_history_ix.take();
339 self.pending_prompt = prompt;
340 }
341
342 self.edited_since_done = true;
343 cx.notify();
344 }
345 EditorEvent::Blurred => {
346 if self.show_rate_limit_notice {
347 self.show_rate_limit_notice = false;
348 cx.notify();
349 }
350 }
351 _ => {}
352 }
353 }
354
355 fn toggle_context_picker(
356 &mut self,
357 _: &ToggleContextPicker,
358 window: &mut Window,
359 cx: &mut Context<Self>,
360 ) {
361 self.context_picker_menu_handle.toggle(window, cx);
362 }
363
364 pub fn remove_all_context(
365 &mut self,
366 _: &RemoveAllContext,
367 _window: &mut Window,
368 cx: &mut Context<Self>,
369 ) {
370 self.context_store.update(cx, |store, _cx| store.clear());
371 cx.notify();
372 }
373
374 fn cancel(
375 &mut self,
376 _: &editor::actions::Cancel,
377 _window: &mut Window,
378 cx: &mut Context<Self>,
379 ) {
380 match self.codegen_status(cx) {
381 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
382 cx.emit(PromptEditorEvent::CancelRequested);
383 }
384 CodegenStatus::Pending => {
385 cx.emit(PromptEditorEvent::StopRequested);
386 }
387 }
388 }
389
390 fn confirm(&mut self, _: &menu::Confirm, _window: &mut Window, cx: &mut Context<Self>) {
391 match self.codegen_status(cx) {
392 CodegenStatus::Idle => {
393 cx.emit(PromptEditorEvent::StartRequested);
394 }
395 CodegenStatus::Pending => {
396 cx.emit(PromptEditorEvent::DismissRequested);
397 }
398 CodegenStatus::Done => {
399 if self.edited_since_done {
400 cx.emit(PromptEditorEvent::StartRequested);
401 } else {
402 cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
403 }
404 }
405 CodegenStatus::Error(_) => {
406 cx.emit(PromptEditorEvent::StartRequested);
407 }
408 }
409 }
410
411 fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
412 if let Some(ix) = self.prompt_history_ix {
413 if ix > 0 {
414 self.prompt_history_ix = Some(ix - 1);
415 let prompt = self.prompt_history[ix - 1].as_str();
416 self.editor.update(cx, |editor, cx| {
417 editor.set_text(prompt, window, cx);
418 editor.move_to_beginning(&Default::default(), window, cx);
419 });
420 }
421 } else if !self.prompt_history.is_empty() {
422 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
423 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
424 self.editor.update(cx, |editor, cx| {
425 editor.set_text(prompt, window, cx);
426 editor.move_to_beginning(&Default::default(), window, cx);
427 });
428 }
429 }
430
431 fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {
432 if let Some(ix) = self.prompt_history_ix {
433 if ix < self.prompt_history.len() - 1 {
434 self.prompt_history_ix = Some(ix + 1);
435 let prompt = self.prompt_history[ix + 1].as_str();
436 self.editor.update(cx, |editor, cx| {
437 editor.set_text(prompt, window, cx);
438 editor.move_to_end(&Default::default(), window, cx)
439 });
440 } else {
441 self.prompt_history_ix = None;
442 let prompt = self.pending_prompt.as_str();
443 self.editor.update(cx, |editor, cx| {
444 editor.set_text(prompt, window, cx);
445 editor.move_to_end(&Default::default(), window, cx)
446 });
447 }
448 } else {
449 self.context_strip.focus_handle(cx).focus(window);
450 }
451 }
452
453 fn render_buttons(&self, _window: &mut Window, cx: &mut Context<Self>) -> Vec<AnyElement> {
454 let mode = match &self.mode {
455 PromptEditorMode::Buffer { codegen, .. } => {
456 let codegen = codegen.read(cx);
457 if codegen.is_insertion {
458 GenerationMode::Generate
459 } else {
460 GenerationMode::Transform
461 }
462 }
463 PromptEditorMode::Terminal { .. } => GenerationMode::Generate,
464 };
465
466 let codegen_status = self.codegen_status(cx);
467
468 match codegen_status {
469 CodegenStatus::Idle => {
470 vec![
471 Button::new("start", mode.start_label())
472 .label_size(LabelSize::Small)
473 .icon(IconName::Return)
474 .icon_size(IconSize::XSmall)
475 .icon_color(Color::Muted)
476 .on_click(
477 cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
478 )
479 .into_any_element(),
480 ]
481 }
482 CodegenStatus::Pending => vec![
483 IconButton::new("stop", IconName::Stop)
484 .icon_color(Color::Error)
485 .shape(IconButtonShape::Square)
486 .tooltip(move |window, cx| {
487 Tooltip::with_meta(
488 mode.tooltip_interrupt(),
489 Some(&menu::Cancel),
490 "Changes won't be discarded",
491 window,
492 cx,
493 )
494 })
495 .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
496 .into_any_element(),
497 ],
498 CodegenStatus::Done | CodegenStatus::Error(_) => {
499 let has_error = matches!(codegen_status, CodegenStatus::Error(_));
500 if has_error || self.edited_since_done {
501 vec![
502 IconButton::new("restart", IconName::RotateCw)
503 .icon_color(Color::Info)
504 .shape(IconButtonShape::Square)
505 .tooltip(move |window, cx| {
506 Tooltip::with_meta(
507 mode.tooltip_restart(),
508 Some(&menu::Confirm),
509 "Changes will be discarded",
510 window,
511 cx,
512 )
513 })
514 .on_click(cx.listener(|_, _, _, cx| {
515 cx.emit(PromptEditorEvent::StartRequested);
516 }))
517 .into_any_element(),
518 ]
519 } else {
520 let accept = IconButton::new("accept", IconName::Check)
521 .icon_color(Color::Info)
522 .shape(IconButtonShape::Square)
523 .tooltip(move |window, cx| {
524 Tooltip::for_action(mode.tooltip_accept(), &menu::Confirm, window, cx)
525 })
526 .on_click(cx.listener(|_, _, _, cx| {
527 cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
528 }))
529 .into_any_element();
530
531 match &self.mode {
532 PromptEditorMode::Terminal { .. } => vec![
533 accept,
534 IconButton::new("confirm", IconName::Play)
535 .icon_color(Color::Info)
536 .shape(IconButtonShape::Square)
537 .tooltip(|window, cx| {
538 Tooltip::for_action(
539 "Execute Generated Command",
540 &menu::SecondaryConfirm,
541 window,
542 cx,
543 )
544 })
545 .on_click(cx.listener(|_, _, _, cx| {
546 cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
547 }))
548 .into_any_element(),
549 ],
550 PromptEditorMode::Buffer { .. } => vec![accept],
551 }
552 }
553 }
554 }
555 }
556
557 fn cycle_prev(
558 &mut self,
559 _: &CyclePreviousInlineAssist,
560 _: &mut Window,
561 cx: &mut Context<Self>,
562 ) {
563 match &self.mode {
564 PromptEditorMode::Buffer { codegen, .. } => {
565 codegen.update(cx, |codegen, cx| codegen.cycle_prev(cx));
566 }
567 PromptEditorMode::Terminal { .. } => {
568 // no cycle buttons in terminal mode
569 }
570 }
571 }
572
573 fn cycle_next(&mut self, _: &CycleNextInlineAssist, _: &mut Window, cx: &mut Context<Self>) {
574 match &self.mode {
575 PromptEditorMode::Buffer { codegen, .. } => {
576 codegen.update(cx, |codegen, cx| codegen.cycle_next(cx));
577 }
578 PromptEditorMode::Terminal { .. } => {
579 // no cycle buttons in terminal mode
580 }
581 }
582 }
583
584 fn render_close_button(&self, cx: &mut Context<Self>) -> AnyElement {
585 IconButton::new("cancel", IconName::Close)
586 .icon_color(Color::Muted)
587 .shape(IconButtonShape::Square)
588 .tooltip(Tooltip::text("Close Assistant"))
589 .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
590 .into_any_element()
591 }
592
593 fn render_cycle_controls(&self, codegen: &BufferCodegen, cx: &Context<Self>) -> AnyElement {
594 let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
595
596 let model_registry = LanguageModelRegistry::read_global(cx);
597 let default_model = model_registry.default_model().map(|default| default.model);
598 let alternative_models = model_registry.inline_alternative_models();
599
600 let get_model_name = |index: usize| -> String {
601 let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
602
603 match index {
604 0 => default_model.as_ref().map_or_else(String::new, name),
605 index if index <= alternative_models.len() => alternative_models
606 .get(index - 1)
607 .map_or_else(String::new, name),
608 _ => String::new(),
609 }
610 };
611
612 let total_models = alternative_models.len() + 1;
613
614 if total_models <= 1 {
615 return div().into_any_element();
616 }
617
618 let current_index = codegen.active_alternative;
619 let prev_index = (current_index + total_models - 1) % total_models;
620 let next_index = (current_index + 1) % total_models;
621
622 let prev_model_name = get_model_name(prev_index);
623 let next_model_name = get_model_name(next_index);
624
625 h_flex()
626 .child(
627 IconButton::new("previous", IconName::ChevronLeft)
628 .icon_color(Color::Muted)
629 .disabled(disabled || current_index == 0)
630 .shape(IconButtonShape::Square)
631 .tooltip({
632 let focus_handle = self.editor.focus_handle(cx);
633 move |window, cx| {
634 cx.new(|cx| {
635 let mut tooltip = Tooltip::new("Previous Alternative").key_binding(
636 KeyBinding::for_action_in(
637 &CyclePreviousInlineAssist,
638 &focus_handle,
639 window,
640 cx,
641 ),
642 );
643 if !disabled && current_index != 0 {
644 tooltip = tooltip.meta(prev_model_name.clone());
645 }
646 tooltip
647 })
648 .into()
649 }
650 })
651 .on_click(cx.listener(|this, _, window, cx| {
652 this.cycle_prev(&CyclePreviousInlineAssist, window, cx);
653 })),
654 )
655 .child(
656 Label::new(format!(
657 "{}/{}",
658 codegen.active_alternative + 1,
659 codegen.alternative_count(cx)
660 ))
661 .size(LabelSize::Small)
662 .color(if disabled {
663 Color::Disabled
664 } else {
665 Color::Muted
666 }),
667 )
668 .child(
669 IconButton::new("next", IconName::ChevronRight)
670 .icon_color(Color::Muted)
671 .disabled(disabled || current_index == total_models - 1)
672 .shape(IconButtonShape::Square)
673 .tooltip({
674 let focus_handle = self.editor.focus_handle(cx);
675 move |window, cx| {
676 cx.new(|cx| {
677 let mut tooltip = Tooltip::new("Next Alternative").key_binding(
678 KeyBinding::for_action_in(
679 &CycleNextInlineAssist,
680 &focus_handle,
681 window,
682 cx,
683 ),
684 );
685 if !disabled && current_index != total_models - 1 {
686 tooltip = tooltip.meta(next_model_name.clone());
687 }
688 tooltip
689 })
690 .into()
691 }
692 })
693 .on_click(cx.listener(|this, _, window, cx| {
694 this.cycle_next(&CycleNextInlineAssist, window, cx)
695 })),
696 )
697 .into_any_element()
698 }
699
700 fn render_rate_limit_notice(&self, cx: &mut Context<Self>) -> impl IntoElement {
701 Popover::new().child(
702 v_flex()
703 .occlude()
704 .p_2()
705 .child(
706 Label::new("Out of Tokens")
707 .size(LabelSize::Small)
708 .weight(FontWeight::BOLD),
709 )
710 .child(Label::new(
711 "Try Zed Pro for higher limits, a wider range of models, and more.",
712 ))
713 .child(
714 h_flex()
715 .justify_between()
716 .child(CheckboxWithLabel::new(
717 "dont-show-again",
718 Label::new("Don't show again"),
719 if dismissed_rate_limit_notice() {
720 ui::ToggleState::Selected
721 } else {
722 ui::ToggleState::Unselected
723 },
724 |selection, _, cx| {
725 let is_dismissed = match selection {
726 ui::ToggleState::Unselected => false,
727 ui::ToggleState::Indeterminate => return,
728 ui::ToggleState::Selected => true,
729 };
730
731 set_rate_limit_notice_dismissed(is_dismissed, cx)
732 },
733 ))
734 .child(
735 h_flex()
736 .gap_2()
737 .child(
738 Button::new("dismiss", "Dismiss")
739 .style(ButtonStyle::Transparent)
740 .on_click(cx.listener(Self::toggle_rate_limit_notice)),
741 )
742 .child(Button::new("more-info", "More Info").on_click(
743 |_event, window, cx| {
744 window.dispatch_action(
745 Box::new(zed_actions::OpenAccountSettings),
746 cx,
747 )
748 },
749 )),
750 ),
751 ),
752 )
753 }
754
755 fn render_editor(&mut self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
756 let font_size = TextSize::Default.rems(cx);
757 let line_height = font_size.to_pixels(window.rem_size()) * 1.3;
758
759 div()
760 .key_context("InlineAssistEditor")
761 .size_full()
762 .p_2()
763 .pl_1()
764 .bg(cx.theme().colors().editor_background)
765 .child({
766 let settings = ThemeSettings::get_global(cx);
767 let text_style = TextStyle {
768 color: cx.theme().colors().editor_foreground,
769 font_family: settings.buffer_font.family.clone(),
770 font_features: settings.buffer_font.features.clone(),
771 font_size: font_size.into(),
772 line_height: line_height.into(),
773 ..Default::default()
774 };
775
776 EditorElement::new(
777 &self.editor,
778 EditorStyle {
779 background: cx.theme().colors().editor_background,
780 local_player: cx.theme().players().local(),
781 text: text_style,
782 ..Default::default()
783 },
784 )
785 })
786 .into_any_element()
787 }
788
789 fn handle_context_strip_event(
790 &mut self,
791 _context_strip: &Entity<ContextStrip>,
792 event: &ContextStripEvent,
793 window: &mut Window,
794 cx: &mut Context<Self>,
795 ) {
796 match event {
797 ContextStripEvent::PickerDismissed
798 | ContextStripEvent::BlurredEmpty
799 | ContextStripEvent::BlurredUp => self.editor.focus_handle(cx).focus(window),
800 ContextStripEvent::BlurredDown => {}
801 }
802 }
803}
804
805pub enum PromptEditorMode {
806 Buffer {
807 id: InlineAssistId,
808 codegen: Entity<BufferCodegen>,
809 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
810 },
811 Terminal {
812 id: TerminalInlineAssistId,
813 codegen: Entity<TerminalCodegen>,
814 height_in_lines: u8,
815 },
816}
817
818pub enum PromptEditorEvent {
819 StartRequested,
820 StopRequested,
821 ConfirmRequested { execute: bool },
822 CancelRequested,
823 DismissRequested,
824 Resized { height_in_lines: u8 },
825}
826
827#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
828pub struct InlineAssistId(pub usize);
829
830impl InlineAssistId {
831 pub fn post_inc(&mut self) -> InlineAssistId {
832 let id = *self;
833 self.0 += 1;
834 id
835 }
836}
837
838impl PromptEditor<BufferCodegen> {
839 pub fn new_buffer(
840 id: InlineAssistId,
841 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
842 prompt_history: VecDeque<String>,
843 prompt_buffer: Entity<MultiBuffer>,
844 codegen: Entity<BufferCodegen>,
845 fs: Arc<dyn Fs>,
846 context_store: Entity<ContextStore>,
847 workspace: WeakEntity<Workspace>,
848 thread_store: Option<WeakEntity<ThreadStore>>,
849 text_thread_store: Option<WeakEntity<TextThreadStore>>,
850 window: &mut Window,
851 cx: &mut Context<PromptEditor<BufferCodegen>>,
852 ) -> PromptEditor<BufferCodegen> {
853 let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
854 let codegen_buffer = codegen.read(cx).buffer(cx).read(cx).as_singleton();
855 let mode = PromptEditorMode::Buffer {
856 id,
857 codegen,
858 gutter_dimensions,
859 };
860
861 let prompt_editor = cx.new(|cx| {
862 let mut editor = Editor::new(
863 EditorMode::AutoHeight {
864 max_lines: Self::MAX_LINES as usize,
865 },
866 prompt_buffer,
867 None,
868 window,
869 cx,
870 );
871 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
872 // Since the prompt editors for all inline assistants are linked,
873 // always show the cursor (even when it isn't focused) because
874 // typing in one will make what you typed appear in all of them.
875 editor.set_show_cursor_when_unfocused(true, cx);
876 editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
877 editor.register_addon(ContextCreasesAddon::new());
878 editor.set_context_menu_options(ContextMenuOptions {
879 min_entries_visible: 12,
880 max_entries_visible: 12,
881 placement: None,
882 });
883
884 editor
885 });
886
887 let prompt_editor_entity = prompt_editor.downgrade();
888 prompt_editor.update(cx, |editor, _| {
889 editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
890 workspace.clone(),
891 context_store.downgrade(),
892 thread_store.clone(),
893 text_thread_store.clone(),
894 prompt_editor_entity,
895 codegen_buffer.as_ref().map(Entity::downgrade),
896 ))));
897 });
898
899 let context_picker_menu_handle = PopoverMenuHandle::default();
900 let model_selector_menu_handle = PopoverMenuHandle::default();
901
902 let context_strip = cx.new(|cx| {
903 ContextStrip::new(
904 context_store.clone(),
905 workspace.clone(),
906 thread_store.clone(),
907 text_thread_store.clone(),
908 context_picker_menu_handle.clone(),
909 SuggestContextKind::Thread,
910 window,
911 cx,
912 )
913 });
914
915 let context_strip_subscription =
916 cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
917
918 let mut this: PromptEditor<BufferCodegen> = PromptEditor {
919 editor: prompt_editor.clone(),
920 context_store,
921 context_strip,
922 context_picker_menu_handle,
923 model_selector: cx.new(|cx| {
924 AssistantModelSelector::new(
925 fs,
926 model_selector_menu_handle,
927 prompt_editor.focus_handle(cx),
928 ModelType::InlineAssistant,
929 window,
930 cx,
931 )
932 }),
933 edited_since_done: false,
934 prompt_history,
935 prompt_history_ix: None,
936 pending_prompt: String::new(),
937 _codegen_subscription: codegen_subscription,
938 editor_subscriptions: Vec::new(),
939 _context_strip_subscription: context_strip_subscription,
940 show_rate_limit_notice: false,
941 mode,
942 _phantom: Default::default(),
943 };
944
945 this.subscribe_to_editor(window, cx);
946 this
947 }
948
949 fn handle_codegen_changed(
950 &mut self,
951 _: Entity<BufferCodegen>,
952 cx: &mut Context<PromptEditor<BufferCodegen>>,
953 ) {
954 match self.codegen_status(cx) {
955 CodegenStatus::Idle => {
956 self.editor
957 .update(cx, |editor, _| editor.set_read_only(false));
958 }
959 CodegenStatus::Pending => {
960 self.editor
961 .update(cx, |editor, _| editor.set_read_only(true));
962 }
963 CodegenStatus::Done => {
964 self.edited_since_done = false;
965 self.editor
966 .update(cx, |editor, _| editor.set_read_only(false));
967 }
968 CodegenStatus::Error(error) => {
969 if cx.has_flag::<ZedProFeatureFlag>()
970 && error.error_code() == proto::ErrorCode::RateLimitExceeded
971 && !dismissed_rate_limit_notice()
972 {
973 self.show_rate_limit_notice = true;
974 cx.notify();
975 }
976
977 self.edited_since_done = false;
978 self.editor
979 .update(cx, |editor, _| editor.set_read_only(false));
980 }
981 }
982 }
983
984 pub fn id(&self) -> InlineAssistId {
985 match &self.mode {
986 PromptEditorMode::Buffer { id, .. } => *id,
987 PromptEditorMode::Terminal { .. } => unreachable!(),
988 }
989 }
990
991 pub fn codegen(&self) -> &Entity<BufferCodegen> {
992 match &self.mode {
993 PromptEditorMode::Buffer { codegen, .. } => codegen,
994 PromptEditorMode::Terminal { .. } => unreachable!(),
995 }
996 }
997
998 pub fn gutter_dimensions(&self) -> &Arc<Mutex<GutterDimensions>> {
999 match &self.mode {
1000 PromptEditorMode::Buffer {
1001 gutter_dimensions, ..
1002 } => gutter_dimensions,
1003 PromptEditorMode::Terminal { .. } => unreachable!(),
1004 }
1005 }
1006}
1007
1008#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1009pub struct TerminalInlineAssistId(pub usize);
1010
1011impl TerminalInlineAssistId {
1012 pub fn post_inc(&mut self) -> TerminalInlineAssistId {
1013 let id = *self;
1014 self.0 += 1;
1015 id
1016 }
1017}
1018
1019impl PromptEditor<TerminalCodegen> {
1020 pub fn new_terminal(
1021 id: TerminalInlineAssistId,
1022 prompt_history: VecDeque<String>,
1023 prompt_buffer: Entity<MultiBuffer>,
1024 codegen: Entity<TerminalCodegen>,
1025 fs: Arc<dyn Fs>,
1026 context_store: Entity<ContextStore>,
1027 workspace: WeakEntity<Workspace>,
1028 thread_store: Option<WeakEntity<ThreadStore>>,
1029 text_thread_store: Option<WeakEntity<TextThreadStore>>,
1030 window: &mut Window,
1031 cx: &mut Context<Self>,
1032 ) -> Self {
1033 let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
1034 let mode = PromptEditorMode::Terminal {
1035 id,
1036 codegen,
1037 height_in_lines: 1,
1038 };
1039
1040 let prompt_editor = cx.new(|cx| {
1041 let mut editor = Editor::new(
1042 EditorMode::AutoHeight {
1043 max_lines: Self::MAX_LINES as usize,
1044 },
1045 prompt_buffer,
1046 None,
1047 window,
1048 cx,
1049 );
1050 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1051 editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
1052 editor.set_context_menu_options(ContextMenuOptions {
1053 min_entries_visible: 12,
1054 max_entries_visible: 12,
1055 placement: None,
1056 });
1057 editor
1058 });
1059
1060 let prompt_editor_entity = prompt_editor.downgrade();
1061 prompt_editor.update(cx, |editor, _| {
1062 editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
1063 workspace.clone(),
1064 context_store.downgrade(),
1065 thread_store.clone(),
1066 text_thread_store.clone(),
1067 prompt_editor_entity,
1068 None,
1069 ))));
1070 });
1071
1072 let context_picker_menu_handle = PopoverMenuHandle::default();
1073 let model_selector_menu_handle = PopoverMenuHandle::default();
1074
1075 let context_strip = cx.new(|cx| {
1076 ContextStrip::new(
1077 context_store.clone(),
1078 workspace.clone(),
1079 thread_store.clone(),
1080 text_thread_store.clone(),
1081 context_picker_menu_handle.clone(),
1082 SuggestContextKind::Thread,
1083 window,
1084 cx,
1085 )
1086 });
1087
1088 let context_strip_subscription =
1089 cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
1090
1091 let mut this = Self {
1092 editor: prompt_editor.clone(),
1093 context_store,
1094 context_strip,
1095 context_picker_menu_handle,
1096 model_selector: cx.new(|cx| {
1097 AssistantModelSelector::new(
1098 fs,
1099 model_selector_menu_handle.clone(),
1100 prompt_editor.focus_handle(cx),
1101 ModelType::InlineAssistant,
1102 window,
1103 cx,
1104 )
1105 }),
1106 edited_since_done: false,
1107 prompt_history,
1108 prompt_history_ix: None,
1109 pending_prompt: String::new(),
1110 _codegen_subscription: codegen_subscription,
1111 editor_subscriptions: Vec::new(),
1112 _context_strip_subscription: context_strip_subscription,
1113 mode,
1114 show_rate_limit_notice: false,
1115 _phantom: Default::default(),
1116 };
1117 this.count_lines(cx);
1118 this.subscribe_to_editor(window, cx);
1119 this
1120 }
1121
1122 fn count_lines(&mut self, cx: &mut Context<Self>) {
1123 let height_in_lines = cmp::max(
1124 2, // Make the editor at least two lines tall, to account for padding and buttons.
1125 cmp::min(
1126 self.editor
1127 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1128 Self::MAX_LINES as u32,
1129 ),
1130 ) as u8;
1131
1132 match &mut self.mode {
1133 PromptEditorMode::Terminal {
1134 height_in_lines: current_height,
1135 ..
1136 } => {
1137 if height_in_lines != *current_height {
1138 *current_height = height_in_lines;
1139 cx.emit(PromptEditorEvent::Resized { height_in_lines });
1140 }
1141 }
1142 PromptEditorMode::Buffer { .. } => unreachable!(),
1143 }
1144 }
1145
1146 fn handle_codegen_changed(&mut self, _: Entity<TerminalCodegen>, cx: &mut Context<Self>) {
1147 match &self.codegen().read(cx).status {
1148 CodegenStatus::Idle => {
1149 self.editor
1150 .update(cx, |editor, _| editor.set_read_only(false));
1151 }
1152 CodegenStatus::Pending => {
1153 self.editor
1154 .update(cx, |editor, _| editor.set_read_only(true));
1155 }
1156 CodegenStatus::Done | CodegenStatus::Error(_) => {
1157 self.edited_since_done = false;
1158 self.editor
1159 .update(cx, |editor, _| editor.set_read_only(false));
1160 }
1161 }
1162 }
1163
1164 pub fn codegen(&self) -> &Entity<TerminalCodegen> {
1165 match &self.mode {
1166 PromptEditorMode::Buffer { .. } => unreachable!(),
1167 PromptEditorMode::Terminal { codegen, .. } => codegen,
1168 }
1169 }
1170
1171 pub fn id(&self) -> TerminalInlineAssistId {
1172 match &self.mode {
1173 PromptEditorMode::Buffer { .. } => unreachable!(),
1174 PromptEditorMode::Terminal { id, .. } => *id,
1175 }
1176 }
1177}
1178
1179const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
1180
1181fn dismissed_rate_limit_notice() -> bool {
1182 db::kvp::KEY_VALUE_STORE
1183 .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
1184 .log_err()
1185 .map_or(false, |s| s.is_some())
1186}
1187
1188fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut App) {
1189 db::write_and_log(cx, move || async move {
1190 if is_dismissed {
1191 db::kvp::KEY_VALUE_STORE
1192 .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
1193 .await
1194 } else {
1195 db::kvp::KEY_VALUE_STORE
1196 .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
1197 .await
1198 }
1199 })
1200}
1201
1202pub enum CodegenStatus {
1203 Idle,
1204 Pending,
1205 Done,
1206 Error(anyhow::Error),
1207}
1208
1209/// This is just CodegenStatus without the anyhow::Error, which causes a lifetime issue for rendering the Cancel button.
1210#[derive(Copy, Clone)]
1211pub enum CancelButtonState {
1212 Idle,
1213 Pending,
1214 Done,
1215 Error,
1216}
1217
1218impl Into<CancelButtonState> for &CodegenStatus {
1219 fn into(self) -> CancelButtonState {
1220 match self {
1221 CodegenStatus::Idle => CancelButtonState::Idle,
1222 CodegenStatus::Pending => CancelButtonState::Pending,
1223 CodegenStatus::Done => CancelButtonState::Done,
1224 CodegenStatus::Error(_) => CancelButtonState::Error,
1225 }
1226 }
1227}
1228
1229#[derive(Copy, Clone)]
1230pub enum GenerationMode {
1231 Generate,
1232 Transform,
1233}
1234
1235impl GenerationMode {
1236 fn start_label(self) -> &'static str {
1237 match self {
1238 GenerationMode::Generate { .. } => "Generate",
1239 GenerationMode::Transform => "Transform",
1240 }
1241 }
1242 fn tooltip_interrupt(self) -> &'static str {
1243 match self {
1244 GenerationMode::Generate { .. } => "Interrupt Generation",
1245 GenerationMode::Transform => "Interrupt Transform",
1246 }
1247 }
1248
1249 fn tooltip_restart(self) -> &'static str {
1250 match self {
1251 GenerationMode::Generate { .. } => "Restart Generation",
1252 GenerationMode::Transform => "Restart Transform",
1253 }
1254 }
1255
1256 fn tooltip_accept(self) -> &'static str {
1257 match self {
1258 GenerationMode::Generate { .. } => "Accept Generation",
1259 GenerationMode::Transform => "Accept Transform",
1260 }
1261 }
1262}