1use crate::{
2 humanize_token_count, prompts::PromptBuilder, AssistantPanel, AssistantPanelEvent,
3 ModelSelector, DEFAULT_CONTEXT_LINES,
4};
5use anyhow::{Context as _, Result};
6use client::telemetry::Telemetry;
7use collections::{HashMap, VecDeque};
8use editor::{
9 actions::{MoveDown, MoveUp, SelectAll},
10 Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, MultiBuffer,
11};
12use fs::Fs;
13use futures::{channel::mpsc, SinkExt, StreamExt};
14use gpui::{
15 AppContext, Context, EventEmitter, FocusHandle, FocusableView, Global, Model, ModelContext,
16 Subscription, Task, TextStyle, UpdateGlobal, View, WeakView,
17};
18use language::Buffer;
19use language_model::{
20 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, Role,
21};
22use settings::Settings;
23use std::{
24 cmp,
25 sync::Arc,
26 time::{Duration, Instant},
27};
28use terminal::Terminal;
29use terminal_view::TerminalView;
30use theme::ThemeSettings;
31use ui::{prelude::*, IconButtonShape, Tooltip};
32use util::ResultExt;
33use workspace::{notifications::NotificationId, Toast, Workspace};
34
35pub fn init(
36 fs: Arc<dyn Fs>,
37 prompt_builder: Arc<PromptBuilder>,
38 telemetry: Arc<Telemetry>,
39 cx: &mut AppContext,
40) {
41 cx.set_global(TerminalInlineAssistant::new(fs, prompt_builder, telemetry));
42}
43
44const PROMPT_HISTORY_MAX_LEN: usize = 20;
45
46#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
47struct TerminalInlineAssistId(usize);
48
49impl TerminalInlineAssistId {
50 fn post_inc(&mut self) -> TerminalInlineAssistId {
51 let id = *self;
52 self.0 += 1;
53 id
54 }
55}
56
57pub struct TerminalInlineAssistant {
58 next_assist_id: TerminalInlineAssistId,
59 assists: HashMap<TerminalInlineAssistId, TerminalInlineAssist>,
60 prompt_history: VecDeque<String>,
61 telemetry: Option<Arc<Telemetry>>,
62 fs: Arc<dyn Fs>,
63 prompt_builder: Arc<PromptBuilder>,
64}
65
66impl Global for TerminalInlineAssistant {}
67
68impl TerminalInlineAssistant {
69 pub fn new(
70 fs: Arc<dyn Fs>,
71 prompt_builder: Arc<PromptBuilder>,
72 telemetry: Arc<Telemetry>,
73 ) -> Self {
74 Self {
75 next_assist_id: TerminalInlineAssistId::default(),
76 assists: HashMap::default(),
77 prompt_history: VecDeque::default(),
78 telemetry: Some(telemetry),
79 fs,
80 prompt_builder,
81 }
82 }
83
84 pub fn assist(
85 &mut self,
86 terminal_view: &View<TerminalView>,
87 workspace: Option<WeakView<Workspace>>,
88 assistant_panel: Option<&View<AssistantPanel>>,
89 initial_prompt: Option<String>,
90 cx: &mut WindowContext,
91 ) {
92 let terminal = terminal_view.read(cx).terminal().clone();
93 let assist_id = self.next_assist_id.post_inc();
94 let prompt_buffer =
95 cx.new_model(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx));
96 let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
97 let codegen = cx.new_model(|_| Codegen::new(terminal, self.telemetry.clone()));
98
99 let prompt_editor = cx.new_view(|cx| {
100 PromptEditor::new(
101 assist_id,
102 self.prompt_history.clone(),
103 prompt_buffer.clone(),
104 codegen,
105 assistant_panel,
106 workspace.clone(),
107 self.fs.clone(),
108 cx,
109 )
110 });
111 let prompt_editor_render = prompt_editor.clone();
112 let block = terminal_view::BlockProperties {
113 height: 2,
114 render: Box::new(move |_| prompt_editor_render.clone().into_any_element()),
115 };
116 terminal_view.update(cx, |terminal_view, cx| {
117 terminal_view.set_block_below_cursor(block, cx);
118 });
119
120 let terminal_assistant = TerminalInlineAssist::new(
121 assist_id,
122 terminal_view,
123 assistant_panel.is_some(),
124 prompt_editor,
125 workspace.clone(),
126 cx,
127 );
128
129 self.assists.insert(assist_id, terminal_assistant);
130
131 self.focus_assist(assist_id, cx);
132 }
133
134 fn focus_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
135 let assist = &self.assists[&assist_id];
136 if let Some(prompt_editor) = assist.prompt_editor.as_ref() {
137 prompt_editor.update(cx, |this, cx| {
138 this.editor.update(cx, |editor, cx| {
139 editor.focus(cx);
140 editor.select_all(&SelectAll, cx);
141 });
142 });
143 }
144 }
145
146 fn handle_prompt_editor_event(
147 &mut self,
148 prompt_editor: View<PromptEditor>,
149 event: &PromptEditorEvent,
150 cx: &mut WindowContext,
151 ) {
152 let assist_id = prompt_editor.read(cx).id;
153 match event {
154 PromptEditorEvent::StartRequested => {
155 self.start_assist(assist_id, cx);
156 }
157 PromptEditorEvent::StopRequested => {
158 self.stop_assist(assist_id, cx);
159 }
160 PromptEditorEvent::ConfirmRequested { execute } => {
161 self.finish_assist(assist_id, false, *execute, cx);
162 }
163 PromptEditorEvent::CancelRequested => {
164 self.finish_assist(assist_id, true, false, cx);
165 }
166 PromptEditorEvent::DismissRequested => {
167 self.dismiss_assist(assist_id, cx);
168 }
169 PromptEditorEvent::Resized { height_in_lines } => {
170 self.insert_prompt_editor_into_terminal(assist_id, *height_in_lines, cx);
171 }
172 }
173 }
174
175 fn start_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
176 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
177 assist
178 } else {
179 return;
180 };
181
182 let Some(user_prompt) = assist
183 .prompt_editor
184 .as_ref()
185 .map(|editor| editor.read(cx).prompt(cx))
186 else {
187 return;
188 };
189
190 self.prompt_history.retain(|prompt| *prompt != user_prompt);
191 self.prompt_history.push_back(user_prompt.clone());
192 if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
193 self.prompt_history.pop_front();
194 }
195
196 assist
197 .terminal
198 .update(cx, |terminal, cx| {
199 terminal
200 .terminal()
201 .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
202 })
203 .log_err();
204
205 let codegen = assist.codegen.clone();
206 let Some(request) = self.request_for_inline_assist(assist_id, cx).log_err() else {
207 return;
208 };
209
210 codegen.update(cx, |codegen, cx| codegen.start(request, cx));
211 }
212
213 fn stop_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
214 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
215 assist
216 } else {
217 return;
218 };
219
220 assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
221 }
222
223 fn request_for_inline_assist(
224 &self,
225 assist_id: TerminalInlineAssistId,
226 cx: &mut WindowContext,
227 ) -> Result<LanguageModelRequest> {
228 let assist = self.assists.get(&assist_id).context("invalid assist")?;
229
230 let shell = std::env::var("SHELL").ok();
231 let (latest_output, working_directory) = assist
232 .terminal
233 .update(cx, |terminal, cx| {
234 let terminal = terminal.model().read(cx);
235 let latest_output = terminal.last_n_non_empty_lines(DEFAULT_CONTEXT_LINES);
236 let working_directory = terminal
237 .working_directory()
238 .map(|path| path.to_string_lossy().to_string());
239 (latest_output, working_directory)
240 })
241 .ok()
242 .unwrap_or_default();
243
244 let context_request = if assist.include_context {
245 assist.workspace.as_ref().and_then(|workspace| {
246 let workspace = workspace.upgrade()?.read(cx);
247 let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
248 Some(
249 assistant_panel
250 .read(cx)
251 .active_context(cx)?
252 .read(cx)
253 .to_completion_request(cx),
254 )
255 })
256 } else {
257 None
258 };
259
260 let prompt = self.prompt_builder.generate_terminal_assistant_prompt(
261 &assist
262 .prompt_editor
263 .clone()
264 .context("invalid assist")?
265 .read(cx)
266 .prompt(cx),
267 shell.as_deref(),
268 working_directory.as_deref(),
269 &latest_output,
270 )?;
271
272 let mut messages = Vec::new();
273 if let Some(context_request) = context_request {
274 messages = context_request.messages;
275 }
276
277 messages.push(LanguageModelRequestMessage {
278 role: Role::User,
279 content: vec![prompt.into()],
280 cache: false,
281 });
282
283 Ok(LanguageModelRequest {
284 messages,
285 tools: Vec::new(),
286 stop: Vec::new(),
287 temperature: 1.0,
288 })
289 }
290
291 fn finish_assist(
292 &mut self,
293 assist_id: TerminalInlineAssistId,
294 undo: bool,
295 execute: bool,
296 cx: &mut WindowContext,
297 ) {
298 self.dismiss_assist(assist_id, cx);
299
300 if let Some(assist) = self.assists.remove(&assist_id) {
301 assist
302 .terminal
303 .update(cx, |this, cx| {
304 this.clear_block_below_cursor(cx);
305 this.focus_handle(cx).focus(cx);
306 })
307 .log_err();
308 assist.codegen.update(cx, |codegen, cx| {
309 if undo {
310 codegen.undo(cx);
311 } else if execute {
312 codegen.complete(cx);
313 }
314 });
315 }
316 }
317
318 fn dismiss_assist(
319 &mut self,
320 assist_id: TerminalInlineAssistId,
321 cx: &mut WindowContext,
322 ) -> bool {
323 let Some(assist) = self.assists.get_mut(&assist_id) else {
324 return false;
325 };
326 if assist.prompt_editor.is_none() {
327 return false;
328 }
329 assist.prompt_editor = None;
330 assist
331 .terminal
332 .update(cx, |this, cx| {
333 this.clear_block_below_cursor(cx);
334 this.focus_handle(cx).focus(cx);
335 })
336 .is_ok()
337 }
338
339 fn insert_prompt_editor_into_terminal(
340 &mut self,
341 assist_id: TerminalInlineAssistId,
342 height: u8,
343 cx: &mut WindowContext,
344 ) {
345 if let Some(assist) = self.assists.get_mut(&assist_id) {
346 if let Some(prompt_editor) = assist.prompt_editor.as_ref().cloned() {
347 assist
348 .terminal
349 .update(cx, |terminal, cx| {
350 terminal.clear_block_below_cursor(cx);
351 let block = terminal_view::BlockProperties {
352 height,
353 render: Box::new(move |_| prompt_editor.clone().into_any_element()),
354 };
355 terminal.set_block_below_cursor(block, cx);
356 })
357 .log_err();
358 }
359 }
360 }
361}
362
363struct TerminalInlineAssist {
364 terminal: WeakView<TerminalView>,
365 prompt_editor: Option<View<PromptEditor>>,
366 codegen: Model<Codegen>,
367 workspace: Option<WeakView<Workspace>>,
368 include_context: bool,
369 _subscriptions: Vec<Subscription>,
370}
371
372impl TerminalInlineAssist {
373 pub fn new(
374 assist_id: TerminalInlineAssistId,
375 terminal: &View<TerminalView>,
376 include_context: bool,
377 prompt_editor: View<PromptEditor>,
378 workspace: Option<WeakView<Workspace>>,
379 cx: &mut WindowContext,
380 ) -> Self {
381 let codegen = prompt_editor.read(cx).codegen.clone();
382 Self {
383 terminal: terminal.downgrade(),
384 prompt_editor: Some(prompt_editor.clone()),
385 codegen: codegen.clone(),
386 workspace: workspace.clone(),
387 include_context,
388 _subscriptions: vec![
389 cx.subscribe(&prompt_editor, |prompt_editor, event, cx| {
390 TerminalInlineAssistant::update_global(cx, |this, cx| {
391 this.handle_prompt_editor_event(prompt_editor, event, cx)
392 })
393 }),
394 cx.subscribe(&codegen, move |codegen, event, cx| {
395 TerminalInlineAssistant::update_global(cx, |this, cx| match event {
396 CodegenEvent::Finished => {
397 let assist = if let Some(assist) = this.assists.get(&assist_id) {
398 assist
399 } else {
400 return;
401 };
402
403 if let CodegenStatus::Error(error) = &codegen.read(cx).status {
404 if assist.prompt_editor.is_none() {
405 if let Some(workspace) = assist
406 .workspace
407 .as_ref()
408 .and_then(|workspace| workspace.upgrade())
409 {
410 let error =
411 format!("Terminal inline assistant error: {}", error);
412 workspace.update(cx, |workspace, cx| {
413 struct InlineAssistantError;
414
415 let id =
416 NotificationId::identified::<InlineAssistantError>(
417 assist_id.0,
418 );
419
420 workspace.show_toast(Toast::new(id, error), cx);
421 })
422 }
423 }
424 }
425
426 if assist.prompt_editor.is_none() {
427 this.finish_assist(assist_id, false, false, cx);
428 }
429 }
430 })
431 }),
432 ],
433 }
434 }
435}
436
437enum PromptEditorEvent {
438 StartRequested,
439 StopRequested,
440 ConfirmRequested { execute: bool },
441 CancelRequested,
442 DismissRequested,
443 Resized { height_in_lines: u8 },
444}
445
446struct PromptEditor {
447 id: TerminalInlineAssistId,
448 fs: Arc<dyn Fs>,
449 height_in_lines: u8,
450 editor: View<Editor>,
451 edited_since_done: bool,
452 prompt_history: VecDeque<String>,
453 prompt_history_ix: Option<usize>,
454 pending_prompt: String,
455 codegen: Model<Codegen>,
456 _codegen_subscription: Subscription,
457 editor_subscriptions: Vec<Subscription>,
458 pending_token_count: Task<Result<()>>,
459 token_count: Option<usize>,
460 _token_count_subscriptions: Vec<Subscription>,
461 workspace: Option<WeakView<Workspace>>,
462}
463
464impl EventEmitter<PromptEditorEvent> for PromptEditor {}
465
466impl Render for PromptEditor {
467 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
468 let status = &self.codegen.read(cx).status;
469 let buttons = match status {
470 CodegenStatus::Idle => {
471 vec![
472 IconButton::new("cancel", IconName::Close)
473 .icon_color(Color::Muted)
474 .shape(IconButtonShape::Square)
475 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
476 .on_click(
477 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
478 ),
479 IconButton::new("start", IconName::SparkleAlt)
480 .icon_color(Color::Muted)
481 .shape(IconButtonShape::Square)
482 .tooltip(|cx| Tooltip::for_action("Generate", &menu::Confirm, cx))
483 .on_click(
484 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
485 ),
486 ]
487 }
488 CodegenStatus::Pending => {
489 vec![
490 IconButton::new("cancel", IconName::Close)
491 .icon_color(Color::Muted)
492 .shape(IconButtonShape::Square)
493 .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
494 .on_click(
495 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
496 ),
497 IconButton::new("stop", IconName::Stop)
498 .icon_color(Color::Error)
499 .shape(IconButtonShape::Square)
500 .tooltip(|cx| {
501 Tooltip::with_meta(
502 "Interrupt Generation",
503 Some(&menu::Cancel),
504 "Changes won't be discarded",
505 cx,
506 )
507 })
508 .on_click(
509 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
510 ),
511 ]
512 }
513 CodegenStatus::Error(_) | CodegenStatus::Done => {
514 let cancel = IconButton::new("cancel", IconName::Close)
515 .icon_color(Color::Muted)
516 .shape(IconButtonShape::Square)
517 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
518 .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)));
519
520 let has_error = matches!(status, CodegenStatus::Error(_));
521 if has_error || self.edited_since_done {
522 vec![
523 cancel,
524 IconButton::new("restart", IconName::RotateCw)
525 .icon_color(Color::Info)
526 .shape(IconButtonShape::Square)
527 .tooltip(|cx| {
528 Tooltip::with_meta(
529 "Restart Generation",
530 Some(&menu::Confirm),
531 "Changes will be discarded",
532 cx,
533 )
534 })
535 .on_click(cx.listener(|_, _, cx| {
536 cx.emit(PromptEditorEvent::StartRequested);
537 })),
538 ]
539 } else {
540 vec![
541 cancel,
542 IconButton::new("accept", IconName::Check)
543 .icon_color(Color::Info)
544 .shape(IconButtonShape::Square)
545 .tooltip(|cx| {
546 Tooltip::for_action("Accept Generated Command", &menu::Confirm, cx)
547 })
548 .on_click(cx.listener(|_, _, cx| {
549 cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
550 })),
551 IconButton::new("confirm", IconName::Play)
552 .icon_color(Color::Info)
553 .shape(IconButtonShape::Square)
554 .tooltip(|cx| {
555 Tooltip::for_action(
556 "Execute Generated Command",
557 &menu::SecondaryConfirm,
558 cx,
559 )
560 })
561 .on_click(cx.listener(|_, _, cx| {
562 cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
563 })),
564 ]
565 }
566 }
567 };
568
569 h_flex()
570 .bg(cx.theme().colors().editor_background)
571 .border_y_1()
572 .border_color(cx.theme().status().info_border)
573 .py_1p5()
574 .h_full()
575 .w_full()
576 .on_action(cx.listener(Self::confirm))
577 .on_action(cx.listener(Self::secondary_confirm))
578 .on_action(cx.listener(Self::cancel))
579 .on_action(cx.listener(Self::move_up))
580 .on_action(cx.listener(Self::move_down))
581 .child(
582 h_flex()
583 .w_12()
584 .justify_center()
585 .gap_2()
586 .child(ModelSelector::new(
587 self.fs.clone(),
588 IconButton::new("context", IconName::SettingsAlt)
589 .shape(IconButtonShape::Square)
590 .icon_size(IconSize::Small)
591 .icon_color(Color::Muted)
592 .tooltip(move |cx| {
593 Tooltip::with_meta(
594 format!(
595 "Using {}",
596 LanguageModelRegistry::read_global(cx)
597 .active_model()
598 .map(|model| model.name().0)
599 .unwrap_or_else(|| "No model selected".into()),
600 ),
601 None,
602 "Change Model",
603 cx,
604 )
605 }),
606 ))
607 .children(
608 if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
609 let error_message = SharedString::from(error.to_string());
610 Some(
611 div()
612 .id("error")
613 .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
614 .child(
615 Icon::new(IconName::XCircle)
616 .size(IconSize::Small)
617 .color(Color::Error),
618 ),
619 )
620 } else {
621 None
622 },
623 ),
624 )
625 .child(div().flex_1().child(self.render_prompt_editor(cx)))
626 .child(
627 h_flex()
628 .gap_1()
629 .pr_4()
630 .children(self.render_token_count(cx))
631 .children(buttons),
632 )
633 }
634}
635
636impl FocusableView for PromptEditor {
637 fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
638 self.editor.focus_handle(cx)
639 }
640}
641
642impl PromptEditor {
643 const MAX_LINES: u8 = 8;
644
645 #[allow(clippy::too_many_arguments)]
646 fn new(
647 id: TerminalInlineAssistId,
648 prompt_history: VecDeque<String>,
649 prompt_buffer: Model<MultiBuffer>,
650 codegen: Model<Codegen>,
651 assistant_panel: Option<&View<AssistantPanel>>,
652 workspace: Option<WeakView<Workspace>>,
653 fs: Arc<dyn Fs>,
654 cx: &mut ViewContext<Self>,
655 ) -> Self {
656 let prompt_editor = cx.new_view(|cx| {
657 let mut editor = Editor::new(
658 EditorMode::AutoHeight {
659 max_lines: Self::MAX_LINES as usize,
660 },
661 prompt_buffer,
662 None,
663 false,
664 cx,
665 );
666 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
667 editor.set_placeholder_text("Add a prompt…", cx);
668 editor
669 });
670
671 let mut token_count_subscriptions = Vec::new();
672 if let Some(assistant_panel) = assistant_panel {
673 token_count_subscriptions
674 .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
675 }
676
677 let mut this = Self {
678 id,
679 height_in_lines: 1,
680 editor: prompt_editor,
681 edited_since_done: false,
682 prompt_history,
683 prompt_history_ix: None,
684 pending_prompt: String::new(),
685 _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
686 editor_subscriptions: Vec::new(),
687 codegen,
688 fs,
689 pending_token_count: Task::ready(Ok(())),
690 token_count: None,
691 _token_count_subscriptions: token_count_subscriptions,
692 workspace,
693 };
694 this.count_lines(cx);
695 this.count_tokens(cx);
696 this.subscribe_to_editor(cx);
697 this
698 }
699
700 fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
701 self.editor_subscriptions.clear();
702 self.editor_subscriptions
703 .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
704 self.editor_subscriptions
705 .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
706 }
707
708 fn prompt(&self, cx: &AppContext) -> String {
709 self.editor.read(cx).text(cx)
710 }
711
712 fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
713 let height_in_lines = cmp::max(
714 2, // Make the editor at least two lines tall, to account for padding and buttons.
715 cmp::min(
716 self.editor
717 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
718 Self::MAX_LINES as u32,
719 ),
720 ) as u8;
721
722 if height_in_lines != self.height_in_lines {
723 self.height_in_lines = height_in_lines;
724 cx.emit(PromptEditorEvent::Resized { height_in_lines });
725 }
726 }
727
728 fn handle_assistant_panel_event(
729 &mut self,
730 _: View<AssistantPanel>,
731 event: &AssistantPanelEvent,
732 cx: &mut ViewContext<Self>,
733 ) {
734 let AssistantPanelEvent::ContextEdited { .. } = event;
735 self.count_tokens(cx);
736 }
737
738 fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
739 let assist_id = self.id;
740 let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
741 return;
742 };
743 self.pending_token_count = cx.spawn(|this, mut cx| async move {
744 cx.background_executor().timer(Duration::from_secs(1)).await;
745 let request =
746 cx.update_global(|inline_assistant: &mut TerminalInlineAssistant, cx| {
747 inline_assistant.request_for_inline_assist(assist_id, cx)
748 })??;
749
750 let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
751 this.update(&mut cx, |this, cx| {
752 this.token_count = Some(token_count);
753 cx.notify();
754 })
755 })
756 }
757
758 fn handle_prompt_editor_changed(&mut self, _: View<Editor>, cx: &mut ViewContext<Self>) {
759 self.count_lines(cx);
760 }
761
762 fn handle_prompt_editor_events(
763 &mut self,
764 _: View<Editor>,
765 event: &EditorEvent,
766 cx: &mut ViewContext<Self>,
767 ) {
768 match event {
769 EditorEvent::Edited { .. } => {
770 let prompt = self.editor.read(cx).text(cx);
771 if self
772 .prompt_history_ix
773 .map_or(true, |ix| self.prompt_history[ix] != prompt)
774 {
775 self.prompt_history_ix.take();
776 self.pending_prompt = prompt;
777 }
778
779 self.edited_since_done = true;
780 cx.notify();
781 }
782 EditorEvent::BufferEdited => {
783 self.count_tokens(cx);
784 }
785 _ => {}
786 }
787 }
788
789 fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
790 match &self.codegen.read(cx).status {
791 CodegenStatus::Idle => {
792 self.editor
793 .update(cx, |editor, _| editor.set_read_only(false));
794 }
795 CodegenStatus::Pending => {
796 self.editor
797 .update(cx, |editor, _| editor.set_read_only(true));
798 }
799 CodegenStatus::Done | CodegenStatus::Error(_) => {
800 self.edited_since_done = false;
801 self.editor
802 .update(cx, |editor, _| editor.set_read_only(false));
803 }
804 }
805 }
806
807 fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
808 match &self.codegen.read(cx).status {
809 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
810 cx.emit(PromptEditorEvent::CancelRequested);
811 }
812 CodegenStatus::Pending => {
813 cx.emit(PromptEditorEvent::StopRequested);
814 }
815 }
816 }
817
818 fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
819 match &self.codegen.read(cx).status {
820 CodegenStatus::Idle => {
821 if !self.editor.read(cx).text(cx).trim().is_empty() {
822 cx.emit(PromptEditorEvent::StartRequested);
823 }
824 }
825 CodegenStatus::Pending => {
826 cx.emit(PromptEditorEvent::DismissRequested);
827 }
828 CodegenStatus::Done => {
829 if self.edited_since_done {
830 cx.emit(PromptEditorEvent::StartRequested);
831 } else {
832 cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
833 }
834 }
835 CodegenStatus::Error(_) => {
836 cx.emit(PromptEditorEvent::StartRequested);
837 }
838 }
839 }
840
841 fn secondary_confirm(&mut self, _: &menu::SecondaryConfirm, cx: &mut ViewContext<Self>) {
842 if matches!(self.codegen.read(cx).status, CodegenStatus::Done) {
843 cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
844 }
845 }
846
847 fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
848 if let Some(ix) = self.prompt_history_ix {
849 if ix > 0 {
850 self.prompt_history_ix = Some(ix - 1);
851 let prompt = self.prompt_history[ix - 1].as_str();
852 self.editor.update(cx, |editor, cx| {
853 editor.set_text(prompt, cx);
854 editor.move_to_beginning(&Default::default(), cx);
855 });
856 }
857 } else if !self.prompt_history.is_empty() {
858 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
859 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
860 self.editor.update(cx, |editor, cx| {
861 editor.set_text(prompt, cx);
862 editor.move_to_beginning(&Default::default(), cx);
863 });
864 }
865 }
866
867 fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
868 if let Some(ix) = self.prompt_history_ix {
869 if ix < self.prompt_history.len() - 1 {
870 self.prompt_history_ix = Some(ix + 1);
871 let prompt = self.prompt_history[ix + 1].as_str();
872 self.editor.update(cx, |editor, cx| {
873 editor.set_text(prompt, cx);
874 editor.move_to_end(&Default::default(), cx)
875 });
876 } else {
877 self.prompt_history_ix = None;
878 let prompt = self.pending_prompt.as_str();
879 self.editor.update(cx, |editor, cx| {
880 editor.set_text(prompt, cx);
881 editor.move_to_end(&Default::default(), cx)
882 });
883 }
884 }
885 }
886
887 fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
888 let model = LanguageModelRegistry::read_global(cx).active_model()?;
889 let token_count = self.token_count?;
890 let max_token_count = model.max_token_count();
891
892 let remaining_tokens = max_token_count as isize - token_count as isize;
893 let token_count_color = if remaining_tokens <= 0 {
894 Color::Error
895 } else if token_count as f32 / max_token_count as f32 >= 0.8 {
896 Color::Warning
897 } else {
898 Color::Muted
899 };
900
901 let mut token_count = h_flex()
902 .id("token_count")
903 .gap_0p5()
904 .child(
905 Label::new(humanize_token_count(token_count))
906 .size(LabelSize::Small)
907 .color(token_count_color),
908 )
909 .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
910 .child(
911 Label::new(humanize_token_count(max_token_count))
912 .size(LabelSize::Small)
913 .color(Color::Muted),
914 );
915 if let Some(workspace) = self.workspace.clone() {
916 token_count = token_count
917 .tooltip(|cx| {
918 Tooltip::with_meta(
919 "Tokens Used by Inline Assistant",
920 None,
921 "Click to Open Assistant Panel",
922 cx,
923 )
924 })
925 .cursor_pointer()
926 .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
927 .on_click(move |_, cx| {
928 cx.stop_propagation();
929 workspace
930 .update(cx, |workspace, cx| {
931 workspace.focus_panel::<AssistantPanel>(cx)
932 })
933 .ok();
934 });
935 } else {
936 token_count = token_count
937 .cursor_default()
938 .tooltip(|cx| Tooltip::text("Tokens Used by Inline Assistant", cx));
939 }
940
941 Some(token_count)
942 }
943
944 fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
945 let settings = ThemeSettings::get_global(cx);
946 let text_style = TextStyle {
947 color: if self.editor.read(cx).read_only(cx) {
948 cx.theme().colors().text_disabled
949 } else {
950 cx.theme().colors().text
951 },
952 font_family: settings.ui_font.family.clone(),
953 font_features: settings.ui_font.features.clone(),
954 font_fallbacks: settings.ui_font.fallbacks.clone(),
955 font_size: rems(0.875).into(),
956 font_weight: settings.ui_font.weight,
957 line_height: relative(1.3),
958 ..Default::default()
959 };
960 EditorElement::new(
961 &self.editor,
962 EditorStyle {
963 background: cx.theme().colors().editor_background,
964 local_player: cx.theme().players().local(),
965 text: text_style,
966 ..Default::default()
967 },
968 )
969 }
970}
971
972#[derive(Debug)]
973pub enum CodegenEvent {
974 Finished,
975}
976
977impl EventEmitter<CodegenEvent> for Codegen {}
978
979const CLEAR_INPUT: &str = "\x15";
980const CARRIAGE_RETURN: &str = "\x0d";
981
982struct TerminalTransaction {
983 terminal: Model<Terminal>,
984}
985
986impl TerminalTransaction {
987 pub fn start(terminal: Model<Terminal>) -> Self {
988 Self { terminal }
989 }
990
991 pub fn push(&mut self, hunk: String, cx: &mut AppContext) {
992 // Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal
993 let input = Self::sanitize_input(hunk);
994 self.terminal
995 .update(cx, |terminal, _| terminal.input(input));
996 }
997
998 pub fn undo(&self, cx: &mut AppContext) {
999 self.terminal
1000 .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
1001 }
1002
1003 pub fn complete(&self, cx: &mut AppContext) {
1004 self.terminal.update(cx, |terminal, _| {
1005 terminal.input(CARRIAGE_RETURN.to_string())
1006 });
1007 }
1008
1009 fn sanitize_input(input: String) -> String {
1010 input.replace(['\r', '\n'], "")
1011 }
1012}
1013
1014pub struct Codegen {
1015 status: CodegenStatus,
1016 telemetry: Option<Arc<Telemetry>>,
1017 terminal: Model<Terminal>,
1018 generation: Task<()>,
1019 transaction: Option<TerminalTransaction>,
1020}
1021
1022impl Codegen {
1023 pub fn new(terminal: Model<Terminal>, telemetry: Option<Arc<Telemetry>>) -> Self {
1024 Self {
1025 terminal,
1026 telemetry,
1027 status: CodegenStatus::Idle,
1028 generation: Task::ready(()),
1029 transaction: None,
1030 }
1031 }
1032
1033 pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut ModelContext<Self>) {
1034 let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
1035 return;
1036 };
1037
1038 let telemetry = self.telemetry.clone();
1039 self.status = CodegenStatus::Pending;
1040 self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
1041 self.generation = cx.spawn(|this, mut cx| async move {
1042 let model_telemetry_id = model.telemetry_id();
1043 let response = model.stream_completion_text(prompt, &cx).await;
1044 let generate = async {
1045 let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1046
1047 let task = cx.background_executor().spawn(async move {
1048 let mut response_latency = None;
1049 let request_start = Instant::now();
1050 let task = async {
1051 let mut chunks = response?;
1052 while let Some(chunk) = chunks.next().await {
1053 if response_latency.is_none() {
1054 response_latency = Some(request_start.elapsed());
1055 }
1056 let chunk = chunk?;
1057 hunks_tx.send(chunk).await?;
1058 }
1059
1060 anyhow::Ok(())
1061 };
1062
1063 let result = task.await;
1064
1065 let error_message = result.as_ref().err().map(|error| error.to_string());
1066 if let Some(telemetry) = telemetry {
1067 telemetry.report_assistant_event(
1068 None,
1069 telemetry_events::AssistantKind::Inline,
1070 model_telemetry_id,
1071 response_latency,
1072 error_message,
1073 );
1074 }
1075
1076 result?;
1077 anyhow::Ok(())
1078 });
1079
1080 while let Some(hunk) = hunks_rx.next().await {
1081 this.update(&mut cx, |this, cx| {
1082 if let Some(transaction) = &mut this.transaction {
1083 transaction.push(hunk, cx);
1084 cx.notify();
1085 }
1086 })?;
1087 }
1088
1089 task.await?;
1090 anyhow::Ok(())
1091 };
1092
1093 let result = generate.await;
1094
1095 this.update(&mut cx, |this, cx| {
1096 if let Err(error) = result {
1097 this.status = CodegenStatus::Error(error);
1098 } else {
1099 this.status = CodegenStatus::Done;
1100 }
1101 cx.emit(CodegenEvent::Finished);
1102 cx.notify();
1103 })
1104 .ok();
1105 });
1106 cx.notify();
1107 }
1108
1109 pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
1110 self.status = CodegenStatus::Done;
1111 self.generation = Task::ready(());
1112 cx.emit(CodegenEvent::Finished);
1113 cx.notify();
1114 }
1115
1116 pub fn complete(&mut self, cx: &mut ModelContext<Self>) {
1117 if let Some(transaction) = self.transaction.take() {
1118 transaction.complete(cx);
1119 }
1120 }
1121
1122 pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
1123 if let Some(transaction) = self.transaction.take() {
1124 transaction.undo(cx);
1125 }
1126 }
1127}
1128
1129enum CodegenStatus {
1130 Idle,
1131 Pending,
1132 Done,
1133 Error(anyhow::Error),
1134}