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::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 window: &mut Window,
850 cx: &mut Context<PromptEditor<BufferCodegen>>,
851 ) -> PromptEditor<BufferCodegen> {
852 let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
853 let codegen_buffer = codegen.read(cx).buffer(cx).read(cx).as_singleton();
854 let mode = PromptEditorMode::Buffer {
855 id,
856 codegen,
857 gutter_dimensions,
858 };
859
860 let prompt_editor = cx.new(|cx| {
861 let mut editor = Editor::new(
862 EditorMode::AutoHeight {
863 max_lines: Self::MAX_LINES as usize,
864 },
865 prompt_buffer,
866 None,
867 window,
868 cx,
869 );
870 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
871 // Since the prompt editors for all inline assistants are linked,
872 // always show the cursor (even when it isn't focused) because
873 // typing in one will make what you typed appear in all of them.
874 editor.set_show_cursor_when_unfocused(true, cx);
875 editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
876 editor.register_addon(ContextCreasesAddon::new());
877 editor.set_context_menu_options(ContextMenuOptions {
878 min_entries_visible: 12,
879 max_entries_visible: 12,
880 placement: None,
881 });
882
883 editor
884 });
885
886 let prompt_editor_entity = prompt_editor.downgrade();
887 prompt_editor.update(cx, |editor, _| {
888 editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
889 workspace.clone(),
890 context_store.downgrade(),
891 thread_store.clone(),
892 prompt_editor_entity,
893 codegen_buffer.as_ref().map(Entity::downgrade),
894 ))));
895 });
896
897 let context_picker_menu_handle = PopoverMenuHandle::default();
898 let model_selector_menu_handle = PopoverMenuHandle::default();
899
900 let context_strip = cx.new(|cx| {
901 ContextStrip::new(
902 context_store.clone(),
903 workspace.clone(),
904 thread_store.clone(),
905 context_picker_menu_handle.clone(),
906 SuggestContextKind::Thread,
907 window,
908 cx,
909 )
910 });
911
912 let context_strip_subscription =
913 cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
914
915 let mut this: PromptEditor<BufferCodegen> = PromptEditor {
916 editor: prompt_editor.clone(),
917 context_store,
918 context_strip,
919 context_picker_menu_handle,
920 model_selector: cx.new(|cx| {
921 AssistantModelSelector::new(
922 fs,
923 model_selector_menu_handle,
924 prompt_editor.focus_handle(cx),
925 ModelType::InlineAssistant,
926 window,
927 cx,
928 )
929 }),
930 edited_since_done: false,
931 prompt_history,
932 prompt_history_ix: None,
933 pending_prompt: String::new(),
934 _codegen_subscription: codegen_subscription,
935 editor_subscriptions: Vec::new(),
936 _context_strip_subscription: context_strip_subscription,
937 show_rate_limit_notice: false,
938 mode,
939 _phantom: Default::default(),
940 };
941
942 this.subscribe_to_editor(window, cx);
943 this
944 }
945
946 fn handle_codegen_changed(
947 &mut self,
948 _: Entity<BufferCodegen>,
949 cx: &mut Context<PromptEditor<BufferCodegen>>,
950 ) {
951 match self.codegen_status(cx) {
952 CodegenStatus::Idle => {
953 self.editor
954 .update(cx, |editor, _| editor.set_read_only(false));
955 }
956 CodegenStatus::Pending => {
957 self.editor
958 .update(cx, |editor, _| editor.set_read_only(true));
959 }
960 CodegenStatus::Done => {
961 self.edited_since_done = false;
962 self.editor
963 .update(cx, |editor, _| editor.set_read_only(false));
964 }
965 CodegenStatus::Error(error) => {
966 if cx.has_flag::<ZedProFeatureFlag>()
967 && error.error_code() == proto::ErrorCode::RateLimitExceeded
968 && !dismissed_rate_limit_notice()
969 {
970 self.show_rate_limit_notice = true;
971 cx.notify();
972 }
973
974 self.edited_since_done = false;
975 self.editor
976 .update(cx, |editor, _| editor.set_read_only(false));
977 }
978 }
979 }
980
981 pub fn id(&self) -> InlineAssistId {
982 match &self.mode {
983 PromptEditorMode::Buffer { id, .. } => *id,
984 PromptEditorMode::Terminal { .. } => unreachable!(),
985 }
986 }
987
988 pub fn codegen(&self) -> &Entity<BufferCodegen> {
989 match &self.mode {
990 PromptEditorMode::Buffer { codegen, .. } => codegen,
991 PromptEditorMode::Terminal { .. } => unreachable!(),
992 }
993 }
994
995 pub fn gutter_dimensions(&self) -> &Arc<Mutex<GutterDimensions>> {
996 match &self.mode {
997 PromptEditorMode::Buffer {
998 gutter_dimensions, ..
999 } => gutter_dimensions,
1000 PromptEditorMode::Terminal { .. } => unreachable!(),
1001 }
1002 }
1003}
1004
1005#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1006pub struct TerminalInlineAssistId(pub usize);
1007
1008impl TerminalInlineAssistId {
1009 pub fn post_inc(&mut self) -> TerminalInlineAssistId {
1010 let id = *self;
1011 self.0 += 1;
1012 id
1013 }
1014}
1015
1016impl PromptEditor<TerminalCodegen> {
1017 pub fn new_terminal(
1018 id: TerminalInlineAssistId,
1019 prompt_history: VecDeque<String>,
1020 prompt_buffer: Entity<MultiBuffer>,
1021 codegen: Entity<TerminalCodegen>,
1022 fs: Arc<dyn Fs>,
1023 context_store: Entity<ContextStore>,
1024 workspace: WeakEntity<Workspace>,
1025 thread_store: Option<WeakEntity<ThreadStore>>,
1026 window: &mut Window,
1027 cx: &mut Context<Self>,
1028 ) -> Self {
1029 let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
1030 let mode = PromptEditorMode::Terminal {
1031 id,
1032 codegen,
1033 height_in_lines: 1,
1034 };
1035
1036 let prompt_editor = cx.new(|cx| {
1037 let mut editor = Editor::new(
1038 EditorMode::AutoHeight {
1039 max_lines: Self::MAX_LINES as usize,
1040 },
1041 prompt_buffer,
1042 None,
1043 window,
1044 cx,
1045 );
1046 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1047 editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
1048 editor.set_context_menu_options(ContextMenuOptions {
1049 min_entries_visible: 12,
1050 max_entries_visible: 12,
1051 placement: None,
1052 });
1053 editor
1054 });
1055
1056 let prompt_editor_entity = prompt_editor.downgrade();
1057 prompt_editor.update(cx, |editor, _| {
1058 editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
1059 workspace.clone(),
1060 context_store.downgrade(),
1061 thread_store.clone(),
1062 prompt_editor_entity,
1063 None,
1064 ))));
1065 });
1066
1067 let context_picker_menu_handle = PopoverMenuHandle::default();
1068 let model_selector_menu_handle = PopoverMenuHandle::default();
1069
1070 let context_strip = cx.new(|cx| {
1071 ContextStrip::new(
1072 context_store.clone(),
1073 workspace.clone(),
1074 thread_store.clone(),
1075 context_picker_menu_handle.clone(),
1076 SuggestContextKind::Thread,
1077 window,
1078 cx,
1079 )
1080 });
1081
1082 let context_strip_subscription =
1083 cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
1084
1085 let mut this = Self {
1086 editor: prompt_editor.clone(),
1087 context_store,
1088 context_strip,
1089 context_picker_menu_handle,
1090 model_selector: cx.new(|cx| {
1091 AssistantModelSelector::new(
1092 fs,
1093 model_selector_menu_handle.clone(),
1094 prompt_editor.focus_handle(cx),
1095 ModelType::InlineAssistant,
1096 window,
1097 cx,
1098 )
1099 }),
1100 edited_since_done: false,
1101 prompt_history,
1102 prompt_history_ix: None,
1103 pending_prompt: String::new(),
1104 _codegen_subscription: codegen_subscription,
1105 editor_subscriptions: Vec::new(),
1106 _context_strip_subscription: context_strip_subscription,
1107 mode,
1108 show_rate_limit_notice: false,
1109 _phantom: Default::default(),
1110 };
1111 this.count_lines(cx);
1112 this.subscribe_to_editor(window, cx);
1113 this
1114 }
1115
1116 fn count_lines(&mut self, cx: &mut Context<Self>) {
1117 let height_in_lines = cmp::max(
1118 2, // Make the editor at least two lines tall, to account for padding and buttons.
1119 cmp::min(
1120 self.editor
1121 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1122 Self::MAX_LINES as u32,
1123 ),
1124 ) as u8;
1125
1126 match &mut self.mode {
1127 PromptEditorMode::Terminal {
1128 height_in_lines: current_height,
1129 ..
1130 } => {
1131 if height_in_lines != *current_height {
1132 *current_height = height_in_lines;
1133 cx.emit(PromptEditorEvent::Resized { height_in_lines });
1134 }
1135 }
1136 PromptEditorMode::Buffer { .. } => unreachable!(),
1137 }
1138 }
1139
1140 fn handle_codegen_changed(&mut self, _: Entity<TerminalCodegen>, cx: &mut Context<Self>) {
1141 match &self.codegen().read(cx).status {
1142 CodegenStatus::Idle => {
1143 self.editor
1144 .update(cx, |editor, _| editor.set_read_only(false));
1145 }
1146 CodegenStatus::Pending => {
1147 self.editor
1148 .update(cx, |editor, _| editor.set_read_only(true));
1149 }
1150 CodegenStatus::Done | CodegenStatus::Error(_) => {
1151 self.edited_since_done = false;
1152 self.editor
1153 .update(cx, |editor, _| editor.set_read_only(false));
1154 }
1155 }
1156 }
1157
1158 pub fn codegen(&self) -> &Entity<TerminalCodegen> {
1159 match &self.mode {
1160 PromptEditorMode::Buffer { .. } => unreachable!(),
1161 PromptEditorMode::Terminal { codegen, .. } => codegen,
1162 }
1163 }
1164
1165 pub fn id(&self) -> TerminalInlineAssistId {
1166 match &self.mode {
1167 PromptEditorMode::Buffer { .. } => unreachable!(),
1168 PromptEditorMode::Terminal { id, .. } => *id,
1169 }
1170 }
1171}
1172
1173const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
1174
1175fn dismissed_rate_limit_notice() -> bool {
1176 db::kvp::KEY_VALUE_STORE
1177 .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
1178 .log_err()
1179 .map_or(false, |s| s.is_some())
1180}
1181
1182fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut App) {
1183 db::write_and_log(cx, move || async move {
1184 if is_dismissed {
1185 db::kvp::KEY_VALUE_STORE
1186 .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
1187 .await
1188 } else {
1189 db::kvp::KEY_VALUE_STORE
1190 .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
1191 .await
1192 }
1193 })
1194}
1195
1196pub enum CodegenStatus {
1197 Idle,
1198 Pending,
1199 Done,
1200 Error(anyhow::Error),
1201}
1202
1203/// This is just CodegenStatus without the anyhow::Error, which causes a lifetime issue for rendering the Cancel button.
1204#[derive(Copy, Clone)]
1205pub enum CancelButtonState {
1206 Idle,
1207 Pending,
1208 Done,
1209 Error,
1210}
1211
1212impl Into<CancelButtonState> for &CodegenStatus {
1213 fn into(self) -> CancelButtonState {
1214 match self {
1215 CodegenStatus::Idle => CancelButtonState::Idle,
1216 CodegenStatus::Pending => CancelButtonState::Pending,
1217 CodegenStatus::Done => CancelButtonState::Done,
1218 CodegenStatus::Error(_) => CancelButtonState::Error,
1219 }
1220 }
1221}
1222
1223#[derive(Copy, Clone)]
1224pub enum GenerationMode {
1225 Generate,
1226 Transform,
1227}
1228
1229impl GenerationMode {
1230 fn start_label(self) -> &'static str {
1231 match self {
1232 GenerationMode::Generate { .. } => "Generate",
1233 GenerationMode::Transform => "Transform",
1234 }
1235 }
1236 fn tooltip_interrupt(self) -> &'static str {
1237 match self {
1238 GenerationMode::Generate { .. } => "Interrupt Generation",
1239 GenerationMode::Transform => "Interrupt Transform",
1240 }
1241 }
1242
1243 fn tooltip_restart(self) -> &'static str {
1244 match self {
1245 GenerationMode::Generate { .. } => "Restart Generation",
1246 GenerationMode::Transform => "Restart Transform",
1247 }
1248 }
1249
1250 fn tooltip_accept(self) -> &'static str {
1251 match self {
1252 GenerationMode::Generate { .. } => "Accept Generation",
1253 GenerationMode::Transform => "Accept Transform",
1254 }
1255 }
1256}