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