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
330 .project()
331 .update(cx, |project, _| project.is_via_ssh());
332
333 workspace
334 .client()
335 .telemetry()
336 .log_edit_event("inline assist", is_via_ssh);
337 });
338 }
339 let prompt = self.editor.read(cx).text(cx);
340 if self
341 .prompt_history_ix
342 .map_or(true, |ix| self.prompt_history[ix] != prompt)
343 {
344 self.prompt_history_ix.take();
345 self.pending_prompt = prompt;
346 }
347
348 self.edited_since_done = true;
349 cx.notify();
350 }
351 EditorEvent::Blurred => {
352 if self.show_rate_limit_notice {
353 self.show_rate_limit_notice = false;
354 cx.notify();
355 }
356 }
357 _ => {}
358 }
359 }
360
361 fn toggle_context_picker(
362 &mut self,
363 _: &ToggleContextPicker,
364 window: &mut Window,
365 cx: &mut Context<Self>,
366 ) {
367 self.context_picker_menu_handle.toggle(window, cx);
368 }
369
370 pub fn remove_all_context(
371 &mut self,
372 _: &RemoveAllContext,
373 _window: &mut Window,
374 cx: &mut Context<Self>,
375 ) {
376 self.context_store.update(cx, |store, _cx| store.clear());
377 cx.notify();
378 }
379
380 fn cancel(
381 &mut self,
382 _: &editor::actions::Cancel,
383 _window: &mut Window,
384 cx: &mut Context<Self>,
385 ) {
386 match self.codegen_status(cx) {
387 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
388 cx.emit(PromptEditorEvent::CancelRequested);
389 }
390 CodegenStatus::Pending => {
391 cx.emit(PromptEditorEvent::StopRequested);
392 }
393 }
394 }
395
396 fn confirm(&mut self, _: &menu::Confirm, _window: &mut Window, cx: &mut Context<Self>) {
397 match self.codegen_status(cx) {
398 CodegenStatus::Idle => {
399 cx.emit(PromptEditorEvent::StartRequested);
400 }
401 CodegenStatus::Pending => {
402 cx.emit(PromptEditorEvent::DismissRequested);
403 }
404 CodegenStatus::Done => {
405 if self.edited_since_done {
406 cx.emit(PromptEditorEvent::StartRequested);
407 } else {
408 cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
409 }
410 }
411 CodegenStatus::Error(_) => {
412 cx.emit(PromptEditorEvent::StartRequested);
413 }
414 }
415 }
416
417 fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
418 if let Some(ix) = self.prompt_history_ix {
419 if ix > 0 {
420 self.prompt_history_ix = Some(ix - 1);
421 let prompt = self.prompt_history[ix - 1].as_str();
422 self.editor.update(cx, |editor, cx| {
423 editor.set_text(prompt, window, cx);
424 editor.move_to_beginning(&Default::default(), window, cx);
425 });
426 }
427 } else if !self.prompt_history.is_empty() {
428 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
429 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
430 self.editor.update(cx, |editor, cx| {
431 editor.set_text(prompt, window, cx);
432 editor.move_to_beginning(&Default::default(), window, cx);
433 });
434 }
435 }
436
437 fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {
438 if let Some(ix) = self.prompt_history_ix {
439 if ix < self.prompt_history.len() - 1 {
440 self.prompt_history_ix = Some(ix + 1);
441 let prompt = self.prompt_history[ix + 1].as_str();
442 self.editor.update(cx, |editor, cx| {
443 editor.set_text(prompt, window, cx);
444 editor.move_to_end(&Default::default(), window, cx)
445 });
446 } else {
447 self.prompt_history_ix = None;
448 let prompt = self.pending_prompt.as_str();
449 self.editor.update(cx, |editor, cx| {
450 editor.set_text(prompt, window, cx);
451 editor.move_to_end(&Default::default(), window, cx)
452 });
453 }
454 } else if self.context_strip.read(cx).has_context_items(cx) {
455 self.context_strip.focus_handle(cx).focus(window);
456 }
457 }
458
459 fn render_buttons(&self, _window: &mut Window, cx: &mut Context<Self>) -> Vec<AnyElement> {
460 let mode = match &self.mode {
461 PromptEditorMode::Buffer { codegen, .. } => {
462 let codegen = codegen.read(cx);
463 if codegen.is_insertion {
464 GenerationMode::Generate
465 } else {
466 GenerationMode::Transform
467 }
468 }
469 PromptEditorMode::Terminal { .. } => GenerationMode::Generate,
470 };
471
472 let codegen_status = self.codegen_status(cx);
473
474 match codegen_status {
475 CodegenStatus::Idle => {
476 vec![
477 Button::new("start", mode.start_label())
478 .label_size(LabelSize::Small)
479 .icon(IconName::Return)
480 .icon_size(IconSize::XSmall)
481 .icon_color(Color::Muted)
482 .on_click(
483 cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
484 )
485 .into_any_element(),
486 ]
487 }
488 CodegenStatus::Pending => vec![
489 IconButton::new("stop", IconName::Stop)
490 .icon_color(Color::Error)
491 .shape(IconButtonShape::Square)
492 .tooltip(move |window, cx| {
493 Tooltip::with_meta(
494 mode.tooltip_interrupt(),
495 Some(&menu::Cancel),
496 "Changes won't be discarded",
497 window,
498 cx,
499 )
500 })
501 .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
502 .into_any_element(),
503 ],
504 CodegenStatus::Done | CodegenStatus::Error(_) => {
505 let has_error = matches!(codegen_status, CodegenStatus::Error(_));
506 if has_error || self.edited_since_done {
507 vec![
508 IconButton::new("restart", IconName::RotateCw)
509 .icon_color(Color::Info)
510 .shape(IconButtonShape::Square)
511 .tooltip(move |window, cx| {
512 Tooltip::with_meta(
513 mode.tooltip_restart(),
514 Some(&menu::Confirm),
515 "Changes will be discarded",
516 window,
517 cx,
518 )
519 })
520 .on_click(cx.listener(|_, _, _, cx| {
521 cx.emit(PromptEditorEvent::StartRequested);
522 }))
523 .into_any_element(),
524 ]
525 } else {
526 let accept = IconButton::new("accept", IconName::Check)
527 .icon_color(Color::Info)
528 .shape(IconButtonShape::Square)
529 .tooltip(move |window, cx| {
530 Tooltip::for_action(mode.tooltip_accept(), &menu::Confirm, window, cx)
531 })
532 .on_click(cx.listener(|_, _, _, cx| {
533 cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
534 }))
535 .into_any_element();
536
537 match &self.mode {
538 PromptEditorMode::Terminal { .. } => vec![
539 accept,
540 IconButton::new("confirm", IconName::Play)
541 .icon_color(Color::Info)
542 .shape(IconButtonShape::Square)
543 .tooltip(|window, cx| {
544 Tooltip::for_action(
545 "Execute Generated Command",
546 &menu::SecondaryConfirm,
547 window,
548 cx,
549 )
550 })
551 .on_click(cx.listener(|_, _, _, cx| {
552 cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
553 }))
554 .into_any_element(),
555 ],
556 PromptEditorMode::Buffer { .. } => vec![accept],
557 }
558 }
559 }
560 }
561 }
562
563 fn cycle_prev(
564 &mut self,
565 _: &CyclePreviousInlineAssist,
566 _: &mut Window,
567 cx: &mut Context<Self>,
568 ) {
569 match &self.mode {
570 PromptEditorMode::Buffer { codegen, .. } => {
571 codegen.update(cx, |codegen, cx| codegen.cycle_prev(cx));
572 }
573 PromptEditorMode::Terminal { .. } => {
574 // no cycle buttons in terminal mode
575 }
576 }
577 }
578
579 fn cycle_next(&mut self, _: &CycleNextInlineAssist, _: &mut Window, cx: &mut Context<Self>) {
580 match &self.mode {
581 PromptEditorMode::Buffer { codegen, .. } => {
582 codegen.update(cx, |codegen, cx| codegen.cycle_next(cx));
583 }
584 PromptEditorMode::Terminal { .. } => {
585 // no cycle buttons in terminal mode
586 }
587 }
588 }
589
590 fn render_close_button(&self, cx: &mut Context<Self>) -> AnyElement {
591 IconButton::new("cancel", IconName::Close)
592 .icon_color(Color::Muted)
593 .shape(IconButtonShape::Square)
594 .tooltip(Tooltip::text("Close Assistant"))
595 .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
596 .into_any_element()
597 }
598
599 fn render_cycle_controls(&self, codegen: &BufferCodegen, cx: &Context<Self>) -> AnyElement {
600 let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
601
602 let model_registry = LanguageModelRegistry::read_global(cx);
603 let default_model = model_registry.default_model().map(|default| default.model);
604 let alternative_models = model_registry.inline_alternative_models();
605
606 let get_model_name = |index: usize| -> String {
607 let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
608
609 match index {
610 0 => default_model.as_ref().map_or_else(String::new, name),
611 index if index <= alternative_models.len() => alternative_models
612 .get(index - 1)
613 .map_or_else(String::new, name),
614 _ => String::new(),
615 }
616 };
617
618 let total_models = alternative_models.len() + 1;
619
620 if total_models <= 1 {
621 return div().into_any_element();
622 }
623
624 let current_index = codegen.active_alternative;
625 let prev_index = (current_index + total_models - 1) % total_models;
626 let next_index = (current_index + 1) % total_models;
627
628 let prev_model_name = get_model_name(prev_index);
629 let next_model_name = get_model_name(next_index);
630
631 h_flex()
632 .child(
633 IconButton::new("previous", IconName::ChevronLeft)
634 .icon_color(Color::Muted)
635 .disabled(disabled || current_index == 0)
636 .shape(IconButtonShape::Square)
637 .tooltip({
638 let focus_handle = self.editor.focus_handle(cx);
639 move |window, cx| {
640 cx.new(|cx| {
641 let mut tooltip = Tooltip::new("Previous Alternative").key_binding(
642 KeyBinding::for_action_in(
643 &CyclePreviousInlineAssist,
644 &focus_handle,
645 window,
646 cx,
647 ),
648 );
649 if !disabled && current_index != 0 {
650 tooltip = tooltip.meta(prev_model_name.clone());
651 }
652 tooltip
653 })
654 .into()
655 }
656 })
657 .on_click(cx.listener(|this, _, window, cx| {
658 this.cycle_prev(&CyclePreviousInlineAssist, window, cx);
659 })),
660 )
661 .child(
662 Label::new(format!(
663 "{}/{}",
664 codegen.active_alternative + 1,
665 codegen.alternative_count(cx)
666 ))
667 .size(LabelSize::Small)
668 .color(if disabled {
669 Color::Disabled
670 } else {
671 Color::Muted
672 }),
673 )
674 .child(
675 IconButton::new("next", IconName::ChevronRight)
676 .icon_color(Color::Muted)
677 .disabled(disabled || current_index == total_models - 1)
678 .shape(IconButtonShape::Square)
679 .tooltip({
680 let focus_handle = self.editor.focus_handle(cx);
681 move |window, cx| {
682 cx.new(|cx| {
683 let mut tooltip = Tooltip::new("Next Alternative").key_binding(
684 KeyBinding::for_action_in(
685 &CycleNextInlineAssist,
686 &focus_handle,
687 window,
688 cx,
689 ),
690 );
691 if !disabled && current_index != total_models - 1 {
692 tooltip = tooltip.meta(next_model_name.clone());
693 }
694 tooltip
695 })
696 .into()
697 }
698 })
699 .on_click(cx.listener(|this, _, window, cx| {
700 this.cycle_next(&CycleNextInlineAssist, window, cx)
701 })),
702 )
703 .into_any_element()
704 }
705
706 fn render_rate_limit_notice(&self, cx: &mut Context<Self>) -> impl IntoElement {
707 Popover::new().child(
708 v_flex()
709 .occlude()
710 .p_2()
711 .child(
712 Label::new("Out of Tokens")
713 .size(LabelSize::Small)
714 .weight(FontWeight::BOLD),
715 )
716 .child(Label::new(
717 "Try Zed Pro for higher limits, a wider range of models, and more.",
718 ))
719 .child(
720 h_flex()
721 .justify_between()
722 .child(CheckboxWithLabel::new(
723 "dont-show-again",
724 Label::new("Don't show again"),
725 if RateLimitNotice::dismissed() {
726 ui::ToggleState::Selected
727 } else {
728 ui::ToggleState::Unselected
729 },
730 |selection, _, cx| {
731 let is_dismissed = match selection {
732 ui::ToggleState::Unselected => false,
733 ui::ToggleState::Indeterminate => return,
734 ui::ToggleState::Selected => true,
735 };
736
737 RateLimitNotice::set_dismissed(is_dismissed, cx);
738 },
739 ))
740 .child(
741 h_flex()
742 .gap_2()
743 .child(
744 Button::new("dismiss", "Dismiss")
745 .style(ButtonStyle::Transparent)
746 .on_click(cx.listener(Self::toggle_rate_limit_notice)),
747 )
748 .child(Button::new("more-info", "More Info").on_click(
749 |_event, window, cx| {
750 window.dispatch_action(
751 Box::new(zed_actions::OpenAccountSettings),
752 cx,
753 )
754 },
755 )),
756 ),
757 ),
758 )
759 }
760
761 fn render_editor(&mut self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
762 let font_size = TextSize::Default.rems(cx);
763 let line_height = font_size.to_pixels(window.rem_size()) * 1.3;
764
765 div()
766 .key_context("InlineAssistEditor")
767 .size_full()
768 .p_2()
769 .pl_1()
770 .bg(cx.theme().colors().editor_background)
771 .child({
772 let settings = ThemeSettings::get_global(cx);
773 let text_style = TextStyle {
774 color: cx.theme().colors().editor_foreground,
775 font_family: settings.buffer_font.family.clone(),
776 font_features: settings.buffer_font.features.clone(),
777 font_size: font_size.into(),
778 line_height: line_height.into(),
779 ..Default::default()
780 };
781
782 EditorElement::new(
783 &self.editor,
784 EditorStyle {
785 background: cx.theme().colors().editor_background,
786 local_player: cx.theme().players().local(),
787 text: text_style,
788 ..Default::default()
789 },
790 )
791 })
792 .into_any_element()
793 }
794
795 fn handle_context_strip_event(
796 &mut self,
797 _context_strip: &Entity<ContextStrip>,
798 event: &ContextStripEvent,
799 window: &mut Window,
800 cx: &mut Context<Self>,
801 ) {
802 match event {
803 ContextStripEvent::PickerDismissed
804 | ContextStripEvent::BlurredEmpty
805 | ContextStripEvent::BlurredUp => self.editor.focus_handle(cx).focus(window),
806 ContextStripEvent::BlurredDown => {}
807 }
808 }
809}
810
811pub enum PromptEditorMode {
812 Buffer {
813 id: InlineAssistId,
814 codegen: Entity<BufferCodegen>,
815 editor_margins: Arc<Mutex<EditorMargins>>,
816 },
817 Terminal {
818 id: TerminalInlineAssistId,
819 codegen: Entity<TerminalCodegen>,
820 height_in_lines: u8,
821 },
822}
823
824pub enum PromptEditorEvent {
825 StartRequested,
826 StopRequested,
827 ConfirmRequested { execute: bool },
828 CancelRequested,
829 DismissRequested,
830 Resized { height_in_lines: u8 },
831}
832
833#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
834pub struct InlineAssistId(pub usize);
835
836impl InlineAssistId {
837 pub fn post_inc(&mut self) -> InlineAssistId {
838 let id = *self;
839 self.0 += 1;
840 id
841 }
842}
843
844impl PromptEditor<BufferCodegen> {
845 pub fn new_buffer(
846 id: InlineAssistId,
847 editor_margins: Arc<Mutex<EditorMargins>>,
848 prompt_history: VecDeque<String>,
849 prompt_buffer: Entity<MultiBuffer>,
850 codegen: Entity<BufferCodegen>,
851 fs: Arc<dyn Fs>,
852 context_store: Entity<ContextStore>,
853 workspace: WeakEntity<Workspace>,
854 thread_store: Option<WeakEntity<ThreadStore>>,
855 text_thread_store: Option<WeakEntity<TextThreadStore>>,
856 window: &mut Window,
857 cx: &mut Context<PromptEditor<BufferCodegen>>,
858 ) -> PromptEditor<BufferCodegen> {
859 let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
860 let codegen_buffer = codegen.read(cx).buffer(cx).read(cx).as_singleton();
861 let mode = PromptEditorMode::Buffer {
862 id,
863 codegen,
864 editor_margins,
865 };
866
867 let prompt_editor = cx.new(|cx| {
868 let mut editor = Editor::new(
869 EditorMode::AutoHeight {
870 max_lines: Self::MAX_LINES as usize,
871 },
872 prompt_buffer,
873 None,
874 window,
875 cx,
876 );
877 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
878 // Since the prompt editors for all inline assistants are linked,
879 // always show the cursor (even when it isn't focused) because
880 // typing in one will make what you typed appear in all of them.
881 editor.set_show_cursor_when_unfocused(true, cx);
882 editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
883 editor.register_addon(ContextCreasesAddon::new());
884 editor.set_context_menu_options(ContextMenuOptions {
885 min_entries_visible: 12,
886 max_entries_visible: 12,
887 placement: None,
888 });
889
890 editor
891 });
892
893 let prompt_editor_entity = prompt_editor.downgrade();
894 prompt_editor.update(cx, |editor, _| {
895 editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
896 workspace.clone(),
897 context_store.downgrade(),
898 thread_store.clone(),
899 text_thread_store.clone(),
900 prompt_editor_entity,
901 codegen_buffer.as_ref().map(Entity::downgrade),
902 ))));
903 });
904
905 let context_picker_menu_handle = PopoverMenuHandle::default();
906 let model_selector_menu_handle = PopoverMenuHandle::default();
907
908 let context_strip = cx.new(|cx| {
909 ContextStrip::new(
910 context_store.clone(),
911 workspace.clone(),
912 thread_store.clone(),
913 text_thread_store.clone(),
914 context_picker_menu_handle.clone(),
915 SuggestContextKind::Thread,
916 window,
917 cx,
918 )
919 });
920
921 let context_strip_subscription =
922 cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
923
924 let mut this: PromptEditor<BufferCodegen> = PromptEditor {
925 editor: prompt_editor.clone(),
926 context_store,
927 context_strip,
928 context_picker_menu_handle,
929 model_selector: cx.new(|cx| {
930 AgentModelSelector::new(
931 fs,
932 model_selector_menu_handle,
933 prompt_editor.focus_handle(cx),
934 ModelType::InlineAssistant,
935 window,
936 cx,
937 )
938 }),
939 edited_since_done: false,
940 prompt_history,
941 prompt_history_ix: None,
942 pending_prompt: String::new(),
943 _codegen_subscription: codegen_subscription,
944 editor_subscriptions: Vec::new(),
945 _context_strip_subscription: context_strip_subscription,
946 show_rate_limit_notice: false,
947 mode,
948 _phantom: Default::default(),
949 };
950
951 this.subscribe_to_editor(window, cx);
952 this
953 }
954
955 fn handle_codegen_changed(
956 &mut self,
957 _: Entity<BufferCodegen>,
958 cx: &mut Context<PromptEditor<BufferCodegen>>,
959 ) {
960 match self.codegen_status(cx) {
961 CodegenStatus::Idle => {
962 self.editor
963 .update(cx, |editor, _| editor.set_read_only(false));
964 }
965 CodegenStatus::Pending => {
966 self.editor
967 .update(cx, |editor, _| editor.set_read_only(true));
968 }
969 CodegenStatus::Done => {
970 self.edited_since_done = false;
971 self.editor
972 .update(cx, |editor, _| editor.set_read_only(false));
973 }
974 CodegenStatus::Error(error) => {
975 if cx.has_flag::<ZedProFeatureFlag>()
976 && error.error_code() == proto::ErrorCode::RateLimitExceeded
977 && !RateLimitNotice::dismissed()
978 {
979 self.show_rate_limit_notice = true;
980 cx.notify();
981 }
982
983 self.edited_since_done = false;
984 self.editor
985 .update(cx, |editor, _| editor.set_read_only(false));
986 }
987 }
988 }
989
990 pub fn id(&self) -> InlineAssistId {
991 match &self.mode {
992 PromptEditorMode::Buffer { id, .. } => *id,
993 PromptEditorMode::Terminal { .. } => unreachable!(),
994 }
995 }
996
997 pub fn codegen(&self) -> &Entity<BufferCodegen> {
998 match &self.mode {
999 PromptEditorMode::Buffer { codegen, .. } => codegen,
1000 PromptEditorMode::Terminal { .. } => unreachable!(),
1001 }
1002 }
1003
1004 pub fn editor_margins(&self) -> &Arc<Mutex<EditorMargins>> {
1005 match &self.mode {
1006 PromptEditorMode::Buffer { editor_margins, .. } => editor_margins,
1007 PromptEditorMode::Terminal { .. } => unreachable!(),
1008 }
1009 }
1010}
1011
1012#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1013pub struct TerminalInlineAssistId(pub usize);
1014
1015impl TerminalInlineAssistId {
1016 pub fn post_inc(&mut self) -> TerminalInlineAssistId {
1017 let id = *self;
1018 self.0 += 1;
1019 id
1020 }
1021}
1022
1023impl PromptEditor<TerminalCodegen> {
1024 pub fn new_terminal(
1025 id: TerminalInlineAssistId,
1026 prompt_history: VecDeque<String>,
1027 prompt_buffer: Entity<MultiBuffer>,
1028 codegen: Entity<TerminalCodegen>,
1029 fs: Arc<dyn Fs>,
1030 context_store: Entity<ContextStore>,
1031 workspace: WeakEntity<Workspace>,
1032 thread_store: Option<WeakEntity<ThreadStore>>,
1033 text_thread_store: Option<WeakEntity<TextThreadStore>>,
1034 window: &mut Window,
1035 cx: &mut Context<Self>,
1036 ) -> Self {
1037 let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
1038 let mode = PromptEditorMode::Terminal {
1039 id,
1040 codegen,
1041 height_in_lines: 1,
1042 };
1043
1044 let prompt_editor = cx.new(|cx| {
1045 let mut editor = Editor::new(
1046 EditorMode::AutoHeight {
1047 max_lines: Self::MAX_LINES as usize,
1048 },
1049 prompt_buffer,
1050 None,
1051 window,
1052 cx,
1053 );
1054 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1055 editor.set_placeholder_text(Self::placeholder_text(&mode, window, cx), cx);
1056 editor.set_context_menu_options(ContextMenuOptions {
1057 min_entries_visible: 12,
1058 max_entries_visible: 12,
1059 placement: None,
1060 });
1061 editor
1062 });
1063
1064 let prompt_editor_entity = prompt_editor.downgrade();
1065 prompt_editor.update(cx, |editor, _| {
1066 editor.set_completion_provider(Some(Box::new(ContextPickerCompletionProvider::new(
1067 workspace.clone(),
1068 context_store.downgrade(),
1069 thread_store.clone(),
1070 text_thread_store.clone(),
1071 prompt_editor_entity,
1072 None,
1073 ))));
1074 });
1075
1076 let context_picker_menu_handle = PopoverMenuHandle::default();
1077 let model_selector_menu_handle = PopoverMenuHandle::default();
1078
1079 let context_strip = cx.new(|cx| {
1080 ContextStrip::new(
1081 context_store.clone(),
1082 workspace.clone(),
1083 thread_store.clone(),
1084 text_thread_store.clone(),
1085 context_picker_menu_handle.clone(),
1086 SuggestContextKind::Thread,
1087 window,
1088 cx,
1089 )
1090 });
1091
1092 let context_strip_subscription =
1093 cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
1094
1095 let mut this = Self {
1096 editor: prompt_editor.clone(),
1097 context_store,
1098 context_strip,
1099 context_picker_menu_handle,
1100 model_selector: cx.new(|cx| {
1101 AgentModelSelector::new(
1102 fs,
1103 model_selector_menu_handle.clone(),
1104 prompt_editor.focus_handle(cx),
1105 ModelType::InlineAssistant,
1106 window,
1107 cx,
1108 )
1109 }),
1110 edited_since_done: false,
1111 prompt_history,
1112 prompt_history_ix: None,
1113 pending_prompt: String::new(),
1114 _codegen_subscription: codegen_subscription,
1115 editor_subscriptions: Vec::new(),
1116 _context_strip_subscription: context_strip_subscription,
1117 mode,
1118 show_rate_limit_notice: false,
1119 _phantom: Default::default(),
1120 };
1121 this.count_lines(cx);
1122 this.subscribe_to_editor(window, cx);
1123 this
1124 }
1125
1126 fn count_lines(&mut self, cx: &mut Context<Self>) {
1127 let height_in_lines = cmp::max(
1128 2, // Make the editor at least two lines tall, to account for padding and buttons.
1129 cmp::min(
1130 self.editor
1131 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1132 Self::MAX_LINES as u32,
1133 ),
1134 ) as u8;
1135
1136 match &mut self.mode {
1137 PromptEditorMode::Terminal {
1138 height_in_lines: current_height,
1139 ..
1140 } => {
1141 if height_in_lines != *current_height {
1142 *current_height = height_in_lines;
1143 cx.emit(PromptEditorEvent::Resized { height_in_lines });
1144 }
1145 }
1146 PromptEditorMode::Buffer { .. } => unreachable!(),
1147 }
1148 }
1149
1150 fn handle_codegen_changed(&mut self, _: Entity<TerminalCodegen>, cx: &mut Context<Self>) {
1151 match &self.codegen().read(cx).status {
1152 CodegenStatus::Idle => {
1153 self.editor
1154 .update(cx, |editor, _| editor.set_read_only(false));
1155 }
1156 CodegenStatus::Pending => {
1157 self.editor
1158 .update(cx, |editor, _| editor.set_read_only(true));
1159 }
1160 CodegenStatus::Done | CodegenStatus::Error(_) => {
1161 self.edited_since_done = false;
1162 self.editor
1163 .update(cx, |editor, _| editor.set_read_only(false));
1164 }
1165 }
1166 }
1167
1168 pub fn codegen(&self) -> &Entity<TerminalCodegen> {
1169 match &self.mode {
1170 PromptEditorMode::Buffer { .. } => unreachable!(),
1171 PromptEditorMode::Terminal { codegen, .. } => codegen,
1172 }
1173 }
1174
1175 pub fn id(&self) -> TerminalInlineAssistId {
1176 match &self.mode {
1177 PromptEditorMode::Buffer { .. } => unreachable!(),
1178 PromptEditorMode::Terminal { id, .. } => *id,
1179 }
1180 }
1181}
1182
1183struct RateLimitNotice;
1184
1185impl Dismissable for RateLimitNotice {
1186 const KEY: &'static str = "dismissed-rate-limit-notice";
1187}
1188
1189pub enum CodegenStatus {
1190 Idle,
1191 Pending,
1192 Done,
1193 Error(anyhow::Error),
1194}
1195
1196/// This is just CodegenStatus without the anyhow::Error, which causes a lifetime issue for rendering the Cancel button.
1197#[derive(Copy, Clone)]
1198pub enum CancelButtonState {
1199 Idle,
1200 Pending,
1201 Done,
1202 Error,
1203}
1204
1205impl Into<CancelButtonState> for &CodegenStatus {
1206 fn into(self) -> CancelButtonState {
1207 match self {
1208 CodegenStatus::Idle => CancelButtonState::Idle,
1209 CodegenStatus::Pending => CancelButtonState::Pending,
1210 CodegenStatus::Done => CancelButtonState::Done,
1211 CodegenStatus::Error(_) => CancelButtonState::Error,
1212 }
1213 }
1214}
1215
1216#[derive(Copy, Clone)]
1217pub enum GenerationMode {
1218 Generate,
1219 Transform,
1220}
1221
1222impl GenerationMode {
1223 fn start_label(self) -> &'static str {
1224 match self {
1225 GenerationMode::Generate { .. } => "Generate",
1226 GenerationMode::Transform => "Transform",
1227 }
1228 }
1229 fn tooltip_interrupt(self) -> &'static str {
1230 match self {
1231 GenerationMode::Generate { .. } => "Interrupt Generation",
1232 GenerationMode::Transform => "Interrupt Transform",
1233 }
1234 }
1235
1236 fn tooltip_restart(self) -> &'static str {
1237 match self {
1238 GenerationMode::Generate { .. } => "Restart Generation",
1239 GenerationMode::Transform => "Restart Transform",
1240 }
1241 }
1242
1243 fn tooltip_accept(self) -> &'static str {
1244 match self {
1245 GenerationMode::Generate { .. } => "Accept Generation",
1246 GenerationMode::Transform => "Accept Transform",
1247 }
1248 }
1249}