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