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