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