1use std::cmp;
2use std::mem;
3use std::ops::Range;
4use std::rc::Rc;
5use std::sync::Arc;
6
7use crate::context::load_context;
8use crate::mention_set::MentionSet;
9use crate::{
10 AgentPanel,
11 buffer_codegen::{BufferCodegen, CodegenAlternative, CodegenEvent},
12 inline_prompt_editor::{CodegenStatus, InlineAssistId, PromptEditor, PromptEditorEvent},
13 terminal_inline_assistant::TerminalInlineAssistant,
14};
15use agent::HistoryStore;
16use agent_settings::AgentSettings;
17use anyhow::{Context as _, Result};
18use client::telemetry::Telemetry;
19use collections::{HashMap, HashSet, VecDeque, hash_map};
20use editor::EditorSnapshot;
21use editor::MultiBufferOffset;
22use editor::RowExt;
23use editor::SelectionEffects;
24use editor::scroll::ScrollOffset;
25use editor::{
26 Anchor, AnchorRangeExt, CodeActionProvider, Editor, EditorEvent, ExcerptId, ExcerptRange,
27 MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint,
28 actions::SelectAll,
29 display_map::{
30 BlockContext, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId, EditorMargins,
31 RenderBlock, ToDisplayPoint,
32 },
33};
34use fs::Fs;
35use futures::{FutureExt, channel::mpsc};
36use gpui::{
37 App, Context, Entity, Focusable, Global, HighlightStyle, Subscription, Task, UpdateGlobal,
38 WeakEntity, Window, point,
39};
40use language::{Buffer, Point, Selection, TransactionId};
41use language_model::{
42 ConfigurationError, ConfiguredModel, LanguageModelRegistry, report_assistant_event,
43};
44use multi_buffer::MultiBufferRow;
45use parking_lot::Mutex;
46use project::{CodeAction, DisableAiSettings, LspAction, Project, ProjectTransaction};
47use prompt_store::{PromptBuilder, PromptStore};
48use settings::{Settings, SettingsStore};
49use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
50use terminal_view::{TerminalView, terminal_panel::TerminalPanel};
51use text::{OffsetRangeExt, ToPoint as _};
52use ui::prelude::*;
53use util::{RangeExt, ResultExt, maybe};
54use workspace::{ItemHandle, Toast, Workspace, dock::Panel, notifications::NotificationId};
55use zed_actions::agent::OpenSettings;
56
57pub fn init(
58 fs: Arc<dyn Fs>,
59 prompt_builder: Arc<PromptBuilder>,
60 telemetry: Arc<Telemetry>,
61 cx: &mut App,
62) {
63 cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));
64
65 cx.observe_global::<SettingsStore>(|cx| {
66 if DisableAiSettings::get_global(cx).disable_ai {
67 // Hide any active inline assist UI when AI is disabled
68 InlineAssistant::update_global(cx, |assistant, cx| {
69 assistant.cancel_all_active_completions(cx);
70 });
71 }
72 })
73 .detach();
74
75 cx.observe_new(|_workspace: &mut Workspace, window, cx| {
76 let Some(window) = window else {
77 return;
78 };
79 let workspace = cx.entity();
80 InlineAssistant::update_global(cx, |inline_assistant, cx| {
81 inline_assistant.register_workspace(&workspace, window, cx)
82 });
83 })
84 .detach();
85}
86
87const PROMPT_HISTORY_MAX_LEN: usize = 20;
88
89enum InlineAssistTarget {
90 Editor(Entity<Editor>),
91 Terminal(Entity<TerminalView>),
92}
93
94pub struct InlineAssistant {
95 next_assist_id: InlineAssistId,
96 next_assist_group_id: InlineAssistGroupId,
97 assists: HashMap<InlineAssistId, InlineAssist>,
98 assists_by_editor: HashMap<WeakEntity<Editor>, EditorInlineAssists>,
99 assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
100 confirmed_assists: HashMap<InlineAssistId, Entity<CodegenAlternative>>,
101 prompt_history: VecDeque<String>,
102 prompt_builder: Arc<PromptBuilder>,
103 telemetry: Arc<Telemetry>,
104 fs: Arc<dyn Fs>,
105 _inline_assistant_completions: Option<mpsc::UnboundedSender<anyhow::Result<InlineAssistId>>>,
106}
107
108impl Global for InlineAssistant {}
109
110impl InlineAssistant {
111 pub fn new(
112 fs: Arc<dyn Fs>,
113 prompt_builder: Arc<PromptBuilder>,
114 telemetry: Arc<Telemetry>,
115 ) -> Self {
116 Self {
117 next_assist_id: InlineAssistId::default(),
118 next_assist_group_id: InlineAssistGroupId::default(),
119 assists: HashMap::default(),
120 assists_by_editor: HashMap::default(),
121 assist_groups: HashMap::default(),
122 confirmed_assists: HashMap::default(),
123 prompt_history: VecDeque::default(),
124 prompt_builder,
125 telemetry,
126 fs,
127 _inline_assistant_completions: None,
128 }
129 }
130
131 #[cfg(any(test, feature = "test-support"))]
132 pub fn set_completion_receiver(
133 &mut self,
134 sender: mpsc::UnboundedSender<anyhow::Result<InlineAssistId>>,
135 ) {
136 self._inline_assistant_completions = Some(sender);
137 }
138
139 pub fn register_workspace(
140 &mut self,
141 workspace: &Entity<Workspace>,
142 window: &mut Window,
143 cx: &mut App,
144 ) {
145 window
146 .subscribe(workspace, cx, |workspace, event, window, cx| {
147 Self::update_global(cx, |this, cx| {
148 this.handle_workspace_event(workspace, event, window, cx)
149 });
150 })
151 .detach();
152
153 let workspace = workspace.downgrade();
154 cx.observe_global::<SettingsStore>(move |cx| {
155 let Some(workspace) = workspace.upgrade() else {
156 return;
157 };
158 let Some(terminal_panel) = workspace.read(cx).panel::<TerminalPanel>(cx) else {
159 return;
160 };
161 let enabled = AgentSettings::get_global(cx).enabled(cx);
162 terminal_panel.update(cx, |terminal_panel, cx| {
163 terminal_panel.set_assistant_enabled(enabled, cx)
164 });
165 })
166 .detach();
167 }
168
169 /// Hides all active inline assists when AI is disabled
170 pub fn cancel_all_active_completions(&mut self, cx: &mut App) {
171 // Cancel all active completions in editors
172 for (editor_handle, _) in self.assists_by_editor.iter() {
173 if let Some(editor) = editor_handle.upgrade() {
174 let windows = cx.windows();
175 if !windows.is_empty() {
176 let window = windows[0];
177 let _ = window.update(cx, |_, window, cx| {
178 editor.update(cx, |editor, cx| {
179 if editor.has_active_edit_prediction() {
180 editor.cancel(&Default::default(), window, cx);
181 }
182 });
183 });
184 }
185 }
186 }
187 }
188
189 fn handle_workspace_event(
190 &mut self,
191 workspace: Entity<Workspace>,
192 event: &workspace::Event,
193 window: &mut Window,
194 cx: &mut App,
195 ) {
196 match event {
197 workspace::Event::UserSavedItem { item, .. } => {
198 // When the user manually saves an editor, automatically accepts all finished transformations.
199 if let Some(editor) = item.upgrade().and_then(|item| item.act_as::<Editor>(cx))
200 && let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade())
201 {
202 for assist_id in editor_assists.assist_ids.clone() {
203 let assist = &self.assists[&assist_id];
204 if let CodegenStatus::Done = assist.codegen.read(cx).status(cx) {
205 self.finish_assist(assist_id, false, window, cx)
206 }
207 }
208 }
209 }
210 workspace::Event::ItemAdded { item } => {
211 self.register_workspace_item(&workspace, item.as_ref(), window, cx);
212 }
213 _ => (),
214 }
215 }
216
217 fn register_workspace_item(
218 &mut self,
219 workspace: &Entity<Workspace>,
220 item: &dyn ItemHandle,
221 window: &mut Window,
222 cx: &mut App,
223 ) {
224 let is_ai_enabled = !DisableAiSettings::get_global(cx).disable_ai;
225
226 if let Some(editor) = item.act_as::<Editor>(cx) {
227 editor.update(cx, |editor, cx| {
228 if is_ai_enabled {
229 editor.add_code_action_provider(
230 Rc::new(AssistantCodeActionProvider {
231 editor: cx.entity().downgrade(),
232 workspace: workspace.downgrade(),
233 }),
234 window,
235 cx,
236 );
237
238 if DisableAiSettings::get_global(cx).disable_ai {
239 // Cancel any active edit predictions
240 if editor.has_active_edit_prediction() {
241 editor.cancel(&Default::default(), window, cx);
242 }
243 }
244 } else {
245 editor.remove_code_action_provider(
246 ASSISTANT_CODE_ACTION_PROVIDER_ID.into(),
247 window,
248 cx,
249 );
250 }
251 });
252 }
253 }
254
255 pub fn inline_assist(
256 workspace: &mut Workspace,
257 action: &zed_actions::assistant::InlineAssist,
258 window: &mut Window,
259 cx: &mut Context<Workspace>,
260 ) {
261 if !AgentSettings::get_global(cx).enabled(cx) {
262 return;
263 }
264
265 let Some(inline_assist_target) = Self::resolve_inline_assist_target(
266 workspace,
267 workspace.panel::<AgentPanel>(cx),
268 window,
269 cx,
270 ) else {
271 return;
272 };
273
274 let configuration_error = || {
275 let model_registry = LanguageModelRegistry::read_global(cx);
276 model_registry.configuration_error(model_registry.inline_assistant_model(), cx)
277 };
278
279 let Some(agent_panel) = workspace.panel::<AgentPanel>(cx) else {
280 return;
281 };
282 let agent_panel = agent_panel.read(cx);
283
284 let prompt_store = agent_panel.prompt_store().as_ref().cloned();
285 let thread_store = agent_panel.thread_store().clone();
286
287 let handle_assist =
288 |window: &mut Window, cx: &mut Context<Workspace>| match inline_assist_target {
289 InlineAssistTarget::Editor(active_editor) => {
290 InlineAssistant::update_global(cx, |assistant, cx| {
291 assistant.assist(
292 &active_editor,
293 cx.entity().downgrade(),
294 workspace.project().downgrade(),
295 thread_store,
296 prompt_store,
297 action.prompt.clone(),
298 window,
299 cx,
300 );
301 })
302 }
303 InlineAssistTarget::Terminal(active_terminal) => {
304 TerminalInlineAssistant::update_global(cx, |assistant, cx| {
305 assistant.assist(
306 &active_terminal,
307 cx.entity().downgrade(),
308 workspace.project().downgrade(),
309 thread_store,
310 prompt_store,
311 action.prompt.clone(),
312 window,
313 cx,
314 );
315 });
316 }
317 };
318
319 if let Some(error) = configuration_error() {
320 if let ConfigurationError::ProviderNotAuthenticated(provider) = error {
321 cx.spawn(async move |_, cx| {
322 cx.update(|cx| provider.authenticate(cx))?.await?;
323 anyhow::Ok(())
324 })
325 .detach_and_log_err(cx);
326
327 if configuration_error().is_none() {
328 handle_assist(window, cx);
329 }
330 } else {
331 cx.spawn_in(window, async move |_, cx| {
332 let answer = cx
333 .prompt(
334 gpui::PromptLevel::Warning,
335 &error.to_string(),
336 None,
337 &["Configure", "Cancel"],
338 )
339 .await
340 .ok();
341 if let Some(answer) = answer
342 && answer == 0
343 {
344 cx.update(|window, cx| window.dispatch_action(Box::new(OpenSettings), cx))
345 .ok();
346 }
347 anyhow::Ok(())
348 })
349 .detach_and_log_err(cx);
350 }
351 } else {
352 handle_assist(window, cx);
353 }
354 }
355
356 fn codegen_ranges(
357 &mut self,
358 editor: &Entity<Editor>,
359 snapshot: &EditorSnapshot,
360 window: &mut Window,
361 cx: &mut App,
362 ) -> Option<(Vec<Range<Anchor>>, Selection<Point>)> {
363 let (initial_selections, newest_selection) = editor.update(cx, |editor, _| {
364 (
365 editor.selections.all::<Point>(&snapshot.display_snapshot),
366 editor
367 .selections
368 .newest::<Point>(&snapshot.display_snapshot),
369 )
370 });
371
372 // Check if there is already an inline assistant that contains the
373 // newest selection, if there is, focus it
374 if let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) {
375 for assist_id in &editor_assists.assist_ids {
376 let assist = &self.assists[assist_id];
377 let range = assist.range.to_point(&snapshot.buffer_snapshot());
378 if range.start.row <= newest_selection.start.row
379 && newest_selection.end.row <= range.end.row
380 {
381 self.focus_assist(*assist_id, window, cx);
382 return None;
383 }
384 }
385 }
386
387 let mut selections = Vec::<Selection<Point>>::new();
388 let mut newest_selection = None;
389 for mut selection in initial_selections {
390 if selection.end > selection.start {
391 selection.start.column = 0;
392 // If the selection ends at the start of the line, we don't want to include it.
393 if selection.end.column == 0 {
394 selection.end.row -= 1;
395 }
396 selection.end.column = snapshot
397 .buffer_snapshot()
398 .line_len(MultiBufferRow(selection.end.row));
399 } else if let Some(fold) =
400 snapshot.crease_for_buffer_row(MultiBufferRow(selection.end.row))
401 {
402 selection.start = fold.range().start;
403 selection.end = fold.range().end;
404 if MultiBufferRow(selection.end.row) < snapshot.buffer_snapshot().max_row() {
405 let chars = snapshot
406 .buffer_snapshot()
407 .chars_at(Point::new(selection.end.row + 1, 0));
408
409 for c in chars {
410 if c == '\n' {
411 break;
412 }
413 if c.is_whitespace() {
414 continue;
415 }
416 if snapshot
417 .language_at(selection.end)
418 .is_some_and(|language| language.config().brackets.is_closing_brace(c))
419 {
420 selection.end.row += 1;
421 selection.end.column = snapshot
422 .buffer_snapshot()
423 .line_len(MultiBufferRow(selection.end.row));
424 }
425 }
426 }
427 }
428
429 if let Some(prev_selection) = selections.last_mut()
430 && selection.start <= prev_selection.end
431 {
432 prev_selection.end = selection.end;
433 continue;
434 }
435
436 let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
437 if selection.id > latest_selection.id {
438 *latest_selection = selection.clone();
439 }
440 selections.push(selection);
441 }
442 let snapshot = &snapshot.buffer_snapshot();
443 let newest_selection = newest_selection.unwrap();
444
445 let mut codegen_ranges = Vec::new();
446 for (buffer, buffer_range, excerpt_id) in
447 snapshot.ranges_to_buffer_ranges(selections.iter().map(|selection| {
448 snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
449 }))
450 {
451 let anchor_range = Anchor::range_in_buffer(
452 excerpt_id,
453 buffer.anchor_before(buffer_range.start)..buffer.anchor_after(buffer_range.end),
454 );
455
456 codegen_ranges.push(anchor_range);
457
458 if let Some(model) = LanguageModelRegistry::read_global(cx).inline_assistant_model() {
459 self.telemetry.report_assistant_event(AssistantEventData {
460 conversation_id: None,
461 kind: AssistantKind::Inline,
462 phase: AssistantPhase::Invoked,
463 message_id: None,
464 model: model.model.telemetry_id(),
465 model_provider: model.provider.id().to_string(),
466 response_latency: None,
467 error_message: None,
468 language_name: buffer.language().map(|language| language.name().to_proto()),
469 });
470 }
471 }
472
473 Some((codegen_ranges, newest_selection))
474 }
475
476 fn batch_assist(
477 &mut self,
478 editor: &Entity<Editor>,
479 workspace: WeakEntity<Workspace>,
480 project: WeakEntity<Project>,
481 thread_store: Entity<HistoryStore>,
482 prompt_store: Option<Entity<PromptStore>>,
483 initial_prompt: Option<String>,
484 window: &mut Window,
485 codegen_ranges: &[Range<Anchor>],
486 newest_selection: Option<Selection<Point>>,
487 initial_transaction_id: Option<TransactionId>,
488 cx: &mut App,
489 ) -> Option<InlineAssistId> {
490 let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
491
492 let assist_group_id = self.next_assist_group_id.post_inc();
493 let prompt_buffer = cx.new(|cx| {
494 MultiBuffer::singleton(
495 cx.new(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx)),
496 cx,
497 )
498 });
499
500 let mut assists = Vec::new();
501 let mut assist_to_focus = None;
502
503 for range in codegen_ranges {
504 let assist_id = self.next_assist_id.post_inc();
505 let codegen = cx.new(|cx| {
506 BufferCodegen::new(
507 editor.read(cx).buffer().clone(),
508 range.clone(),
509 initial_transaction_id,
510 self.telemetry.clone(),
511 self.prompt_builder.clone(),
512 cx,
513 )
514 });
515
516 let editor_margins = Arc::new(Mutex::new(EditorMargins::default()));
517 let prompt_editor = cx.new(|cx| {
518 PromptEditor::new_buffer(
519 assist_id,
520 editor_margins,
521 self.prompt_history.clone(),
522 prompt_buffer.clone(),
523 codegen.clone(),
524 self.fs.clone(),
525 thread_store.clone(),
526 prompt_store.clone(),
527 project.clone(),
528 workspace.clone(),
529 window,
530 cx,
531 )
532 });
533
534 if let Some(newest_selection) = newest_selection.as_ref()
535 && assist_to_focus.is_none()
536 {
537 let focus_assist = if newest_selection.reversed {
538 range.start.to_point(&snapshot) == newest_selection.start
539 } else {
540 range.end.to_point(&snapshot) == newest_selection.end
541 };
542 if focus_assist {
543 assist_to_focus = Some(assist_id);
544 }
545 }
546
547 let [prompt_block_id, end_block_id] =
548 self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
549
550 assists.push((
551 assist_id,
552 range.clone(),
553 prompt_editor,
554 prompt_block_id,
555 end_block_id,
556 ));
557 }
558
559 let editor_assists = self
560 .assists_by_editor
561 .entry(editor.downgrade())
562 .or_insert_with(|| EditorInlineAssists::new(editor, window, cx));
563
564 let assist_to_focus = if let Some(focus_id) = assist_to_focus {
565 Some(focus_id)
566 } else if assists.len() >= 1 {
567 Some(assists[0].0)
568 } else {
569 None
570 };
571
572 let mut assist_group = InlineAssistGroup::new();
573 for (assist_id, range, prompt_editor, prompt_block_id, end_block_id) in assists {
574 let codegen = prompt_editor.read(cx).codegen().clone();
575
576 self.assists.insert(
577 assist_id,
578 InlineAssist::new(
579 assist_id,
580 assist_group_id,
581 editor,
582 &prompt_editor,
583 prompt_block_id,
584 end_block_id,
585 range,
586 codegen,
587 workspace.clone(),
588 window,
589 cx,
590 ),
591 );
592 assist_group.assist_ids.push(assist_id);
593 editor_assists.assist_ids.push(assist_id);
594 }
595
596 self.assist_groups.insert(assist_group_id, assist_group);
597
598 assist_to_focus
599 }
600
601 pub fn assist(
602 &mut self,
603 editor: &Entity<Editor>,
604 workspace: WeakEntity<Workspace>,
605 project: WeakEntity<Project>,
606 thread_store: Entity<HistoryStore>,
607 prompt_store: Option<Entity<PromptStore>>,
608 initial_prompt: Option<String>,
609 window: &mut Window,
610 cx: &mut App,
611 ) -> Option<InlineAssistId> {
612 let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
613
614 let Some((codegen_ranges, newest_selection)) =
615 self.codegen_ranges(editor, &snapshot, window, cx)
616 else {
617 return None;
618 };
619
620 let assist_to_focus = self.batch_assist(
621 editor,
622 workspace,
623 project,
624 thread_store,
625 prompt_store,
626 initial_prompt,
627 window,
628 &codegen_ranges,
629 Some(newest_selection),
630 None,
631 cx,
632 );
633
634 if let Some(assist_id) = assist_to_focus {
635 self.focus_assist(assist_id, window, cx);
636 }
637
638 assist_to_focus
639 }
640
641 pub fn suggest_assist(
642 &mut self,
643 editor: &Entity<Editor>,
644 mut range: Range<Anchor>,
645 initial_prompt: String,
646 initial_transaction_id: Option<TransactionId>,
647 focus: bool,
648 workspace: Entity<Workspace>,
649 thread_store: Entity<HistoryStore>,
650 prompt_store: Option<Entity<PromptStore>>,
651 window: &mut Window,
652 cx: &mut App,
653 ) -> InlineAssistId {
654 let buffer = editor.read(cx).buffer().clone();
655 {
656 let snapshot = buffer.read(cx).read(cx);
657 range.start = range.start.bias_left(&snapshot);
658 range.end = range.end.bias_right(&snapshot);
659 }
660
661 let project = workspace.read(cx).project().downgrade();
662
663 let assist_id = self
664 .batch_assist(
665 editor,
666 workspace.downgrade(),
667 project,
668 thread_store,
669 prompt_store,
670 Some(initial_prompt),
671 window,
672 &[range],
673 None,
674 initial_transaction_id,
675 cx,
676 )
677 .expect("batch_assist returns an id if there's only one range");
678
679 if focus {
680 self.focus_assist(assist_id, window, cx);
681 }
682
683 assist_id
684 }
685
686 fn insert_assist_blocks(
687 &self,
688 editor: &Entity<Editor>,
689 range: &Range<Anchor>,
690 prompt_editor: &Entity<PromptEditor<BufferCodegen>>,
691 cx: &mut App,
692 ) -> [CustomBlockId; 2] {
693 let prompt_editor_height = prompt_editor.update(cx, |prompt_editor, cx| {
694 prompt_editor
695 .editor
696 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1 + 2)
697 });
698 let assist_blocks = vec![
699 BlockProperties {
700 style: BlockStyle::Sticky,
701 placement: BlockPlacement::Above(range.start),
702 height: Some(prompt_editor_height),
703 render: build_assist_editor_renderer(prompt_editor),
704 priority: 0,
705 },
706 BlockProperties {
707 style: BlockStyle::Sticky,
708 placement: BlockPlacement::Below(range.end),
709 height: None,
710 render: Arc::new(|cx| {
711 v_flex()
712 .h_full()
713 .w_full()
714 .border_t_1()
715 .border_color(cx.theme().status().info_border)
716 .into_any_element()
717 }),
718 priority: 0,
719 },
720 ];
721
722 editor.update(cx, |editor, cx| {
723 let block_ids = editor.insert_blocks(assist_blocks, None, cx);
724 [block_ids[0], block_ids[1]]
725 })
726 }
727
728 fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut App) {
729 let assist = &self.assists[&assist_id];
730 let Some(decorations) = assist.decorations.as_ref() else {
731 return;
732 };
733 let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
734 let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
735
736 assist_group.active_assist_id = Some(assist_id);
737 if assist_group.linked {
738 for assist_id in &assist_group.assist_ids {
739 if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
740 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
741 prompt_editor.set_show_cursor_when_unfocused(true, cx)
742 });
743 }
744 }
745 }
746
747 assist
748 .editor
749 .update(cx, |editor, cx| {
750 let scroll_top = editor.scroll_position(cx).y;
751 let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
752 editor_assists.scroll_lock = editor
753 .row_for_block(decorations.prompt_block_id, cx)
754 .map(|row| row.as_f64())
755 .filter(|prompt_row| (scroll_top..scroll_bottom).contains(&prompt_row))
756 .map(|prompt_row| InlineAssistScrollLock {
757 assist_id,
758 distance_from_top: prompt_row - scroll_top,
759 });
760 })
761 .ok();
762 }
763
764 fn handle_prompt_editor_focus_out(&mut self, assist_id: InlineAssistId, cx: &mut App) {
765 let assist = &self.assists[&assist_id];
766 let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
767 if assist_group.active_assist_id == Some(assist_id) {
768 assist_group.active_assist_id = None;
769 if assist_group.linked {
770 for assist_id in &assist_group.assist_ids {
771 if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
772 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
773 prompt_editor.set_show_cursor_when_unfocused(false, cx)
774 });
775 }
776 }
777 }
778 }
779 }
780
781 fn handle_prompt_editor_event(
782 &mut self,
783 prompt_editor: Entity<PromptEditor<BufferCodegen>>,
784 event: &PromptEditorEvent,
785 window: &mut Window,
786 cx: &mut App,
787 ) {
788 let assist_id = prompt_editor.read(cx).id();
789 match event {
790 PromptEditorEvent::StartRequested => {
791 self.start_assist(assist_id, window, cx);
792 }
793 PromptEditorEvent::StopRequested => {
794 self.stop_assist(assist_id, cx);
795 }
796 PromptEditorEvent::ConfirmRequested { execute: _ } => {
797 self.finish_assist(assist_id, false, window, cx);
798 }
799 PromptEditorEvent::CancelRequested => {
800 self.finish_assist(assist_id, true, window, cx);
801 }
802 PromptEditorEvent::Resized { .. } => {
803 // This only matters for the terminal inline assistant
804 }
805 }
806 }
807
808 fn handle_editor_newline(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
809 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
810 return;
811 };
812
813 if editor.read(cx).selections.count() == 1 {
814 let (selection, buffer) = editor.update(cx, |editor, cx| {
815 (
816 editor
817 .selections
818 .newest::<MultiBufferOffset>(&editor.display_snapshot(cx)),
819 editor.buffer().read(cx).snapshot(cx),
820 )
821 });
822 for assist_id in &editor_assists.assist_ids {
823 let assist = &self.assists[assist_id];
824 let assist_range = assist.range.to_offset(&buffer);
825 if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
826 {
827 if matches!(assist.codegen.read(cx).status(cx), CodegenStatus::Pending) {
828 self.dismiss_assist(*assist_id, window, cx);
829 } else {
830 self.finish_assist(*assist_id, false, window, cx);
831 }
832
833 return;
834 }
835 }
836 }
837
838 cx.propagate();
839 }
840
841 fn handle_editor_cancel(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
842 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
843 return;
844 };
845
846 if editor.read(cx).selections.count() == 1 {
847 let (selection, buffer) = editor.update(cx, |editor, cx| {
848 (
849 editor
850 .selections
851 .newest::<MultiBufferOffset>(&editor.display_snapshot(cx)),
852 editor.buffer().read(cx).snapshot(cx),
853 )
854 });
855 let mut closest_assist_fallback = None;
856 for assist_id in &editor_assists.assist_ids {
857 let assist = &self.assists[assist_id];
858 let assist_range = assist.range.to_offset(&buffer);
859 if assist.decorations.is_some() {
860 if assist_range.contains(&selection.start)
861 && assist_range.contains(&selection.end)
862 {
863 self.focus_assist(*assist_id, window, cx);
864 return;
865 } else {
866 let distance_from_selection = assist_range
867 .start
868 .0
869 .abs_diff(selection.start.0)
870 .min(assist_range.start.0.abs_diff(selection.end.0))
871 + assist_range
872 .end
873 .0
874 .abs_diff(selection.start.0)
875 .min(assist_range.end.0.abs_diff(selection.end.0));
876 match closest_assist_fallback {
877 Some((_, old_distance)) => {
878 if distance_from_selection < old_distance {
879 closest_assist_fallback =
880 Some((assist_id, distance_from_selection));
881 }
882 }
883 None => {
884 closest_assist_fallback = Some((assist_id, distance_from_selection))
885 }
886 }
887 }
888 }
889 }
890
891 if let Some((&assist_id, _)) = closest_assist_fallback {
892 self.focus_assist(assist_id, window, cx);
893 }
894 }
895
896 cx.propagate();
897 }
898
899 fn handle_editor_release(
900 &mut self,
901 editor: WeakEntity<Editor>,
902 window: &mut Window,
903 cx: &mut App,
904 ) {
905 if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor) {
906 for assist_id in editor_assists.assist_ids.clone() {
907 self.finish_assist(assist_id, true, window, cx);
908 }
909 }
910 }
911
912 fn handle_editor_change(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
913 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
914 return;
915 };
916 let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
917 return;
918 };
919 let assist = &self.assists[&scroll_lock.assist_id];
920 let Some(decorations) = assist.decorations.as_ref() else {
921 return;
922 };
923
924 editor.update(cx, |editor, cx| {
925 let scroll_position = editor.scroll_position(cx);
926 let target_scroll_top = editor
927 .row_for_block(decorations.prompt_block_id, cx)?
928 .as_f64()
929 - scroll_lock.distance_from_top;
930 if target_scroll_top != scroll_position.y {
931 editor.set_scroll_position(point(scroll_position.x, target_scroll_top), window, cx);
932 }
933 Some(())
934 });
935 }
936
937 fn handle_editor_event(
938 &mut self,
939 editor: Entity<Editor>,
940 event: &EditorEvent,
941 window: &mut Window,
942 cx: &mut App,
943 ) {
944 let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
945 return;
946 };
947
948 match event {
949 EditorEvent::Edited { transaction_id } => {
950 let buffer = editor.read(cx).buffer().read(cx);
951 let edited_ranges =
952 buffer.edited_ranges_for_transaction::<MultiBufferOffset>(*transaction_id, cx);
953 let snapshot = buffer.snapshot(cx);
954
955 for assist_id in editor_assists.assist_ids.clone() {
956 let assist = &self.assists[&assist_id];
957 if matches!(
958 assist.codegen.read(cx).status(cx),
959 CodegenStatus::Error(_) | CodegenStatus::Done
960 ) {
961 let assist_range = assist.range.to_offset(&snapshot);
962 if edited_ranges
963 .iter()
964 .any(|range| range.overlaps(&assist_range))
965 {
966 self.finish_assist(assist_id, false, window, cx);
967 }
968 }
969 }
970 }
971 EditorEvent::ScrollPositionChanged { .. } => {
972 if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
973 let assist = &self.assists[&scroll_lock.assist_id];
974 if let Some(decorations) = assist.decorations.as_ref() {
975 let distance_from_top = editor.update(cx, |editor, cx| {
976 let scroll_top = editor.scroll_position(cx).y;
977 let prompt_row = editor
978 .row_for_block(decorations.prompt_block_id, cx)?
979 .0 as ScrollOffset;
980 Some(prompt_row - scroll_top)
981 });
982
983 if distance_from_top.is_none_or(|distance_from_top| {
984 distance_from_top != scroll_lock.distance_from_top
985 }) {
986 editor_assists.scroll_lock = None;
987 }
988 }
989 }
990 }
991 EditorEvent::SelectionsChanged { .. } => {
992 for assist_id in editor_assists.assist_ids.clone() {
993 let assist = &self.assists[&assist_id];
994 if let Some(decorations) = assist.decorations.as_ref()
995 && decorations
996 .prompt_editor
997 .focus_handle(cx)
998 .is_focused(window)
999 {
1000 return;
1001 }
1002 }
1003
1004 editor_assists.scroll_lock = None;
1005 }
1006 _ => {}
1007 }
1008 }
1009
1010 pub fn finish_assist(
1011 &mut self,
1012 assist_id: InlineAssistId,
1013 undo: bool,
1014 window: &mut Window,
1015 cx: &mut App,
1016 ) {
1017 if let Some(assist) = self.assists.get(&assist_id) {
1018 let assist_group_id = assist.group_id;
1019 if self.assist_groups[&assist_group_id].linked {
1020 for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1021 self.finish_assist(assist_id, undo, window, cx);
1022 }
1023 return;
1024 }
1025 }
1026
1027 self.dismiss_assist(assist_id, window, cx);
1028
1029 if let Some(assist) = self.assists.remove(&assist_id) {
1030 if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
1031 {
1032 entry.get_mut().assist_ids.retain(|id| *id != assist_id);
1033 if entry.get().assist_ids.is_empty() {
1034 entry.remove();
1035 }
1036 }
1037
1038 if let hash_map::Entry::Occupied(mut entry) =
1039 self.assists_by_editor.entry(assist.editor.clone())
1040 {
1041 entry.get_mut().assist_ids.retain(|id| *id != assist_id);
1042 if entry.get().assist_ids.is_empty() {
1043 entry.remove();
1044 if let Some(editor) = assist.editor.upgrade() {
1045 self.update_editor_highlights(&editor, cx);
1046 }
1047 } else {
1048 entry.get_mut().highlight_updates.send(()).ok();
1049 }
1050 }
1051
1052 let active_alternative = assist.codegen.read(cx).active_alternative().clone();
1053 let message_id = active_alternative.read(cx).message_id.clone();
1054
1055 if let Some(model) = LanguageModelRegistry::read_global(cx).inline_assistant_model() {
1056 let language_name = assist.editor.upgrade().and_then(|editor| {
1057 let multibuffer = editor.read(cx).buffer().read(cx);
1058 let snapshot = multibuffer.snapshot(cx);
1059 let ranges = snapshot.range_to_buffer_ranges(assist.range.clone());
1060 ranges
1061 .first()
1062 .and_then(|(buffer, _, _)| buffer.language())
1063 .map(|language| language.name())
1064 });
1065 report_assistant_event(
1066 AssistantEventData {
1067 conversation_id: None,
1068 kind: AssistantKind::Inline,
1069 message_id,
1070 phase: if undo {
1071 AssistantPhase::Rejected
1072 } else {
1073 AssistantPhase::Accepted
1074 },
1075 model: model.model.telemetry_id(),
1076 model_provider: model.model.provider_id().to_string(),
1077 response_latency: None,
1078 error_message: None,
1079 language_name: language_name.map(|name| name.to_proto()),
1080 },
1081 Some(self.telemetry.clone()),
1082 cx.http_client(),
1083 model.model.api_key(cx),
1084 cx.background_executor(),
1085 );
1086 }
1087
1088 if undo {
1089 assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
1090 } else {
1091 self.confirmed_assists.insert(assist_id, active_alternative);
1092 }
1093 }
1094 }
1095
1096 fn dismiss_assist(
1097 &mut self,
1098 assist_id: InlineAssistId,
1099 window: &mut Window,
1100 cx: &mut App,
1101 ) -> bool {
1102 let Some(assist) = self.assists.get_mut(&assist_id) else {
1103 return false;
1104 };
1105 let Some(editor) = assist.editor.upgrade() else {
1106 return false;
1107 };
1108 let Some(decorations) = assist.decorations.take() else {
1109 return false;
1110 };
1111
1112 editor.update(cx, |editor, cx| {
1113 let mut to_remove = decorations.removed_line_block_ids;
1114 to_remove.insert(decorations.prompt_block_id);
1115 to_remove.insert(decorations.end_block_id);
1116 editor.remove_blocks(to_remove, None, cx);
1117 });
1118
1119 if decorations
1120 .prompt_editor
1121 .focus_handle(cx)
1122 .contains_focused(window, cx)
1123 {
1124 self.focus_next_assist(assist_id, window, cx);
1125 }
1126
1127 if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
1128 if editor_assists
1129 .scroll_lock
1130 .as_ref()
1131 .is_some_and(|lock| lock.assist_id == assist_id)
1132 {
1133 editor_assists.scroll_lock = None;
1134 }
1135 editor_assists.highlight_updates.send(()).ok();
1136 }
1137
1138 true
1139 }
1140
1141 fn focus_next_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1142 let Some(assist) = self.assists.get(&assist_id) else {
1143 return;
1144 };
1145
1146 let assist_group = &self.assist_groups[&assist.group_id];
1147 let assist_ix = assist_group
1148 .assist_ids
1149 .iter()
1150 .position(|id| *id == assist_id)
1151 .unwrap();
1152 let assist_ids = assist_group
1153 .assist_ids
1154 .iter()
1155 .skip(assist_ix + 1)
1156 .chain(assist_group.assist_ids.iter().take(assist_ix));
1157
1158 for assist_id in assist_ids {
1159 let assist = &self.assists[assist_id];
1160 if assist.decorations.is_some() {
1161 self.focus_assist(*assist_id, window, cx);
1162 return;
1163 }
1164 }
1165
1166 assist
1167 .editor
1168 .update(cx, |editor, cx| window.focus(&editor.focus_handle(cx)))
1169 .ok();
1170 }
1171
1172 fn focus_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1173 let Some(assist) = self.assists.get(&assist_id) else {
1174 return;
1175 };
1176
1177 if let Some(decorations) = assist.decorations.as_ref() {
1178 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
1179 prompt_editor.editor.update(cx, |editor, cx| {
1180 window.focus(&editor.focus_handle(cx));
1181 editor.select_all(&SelectAll, window, cx);
1182 })
1183 });
1184 }
1185
1186 self.scroll_to_assist(assist_id, window, cx);
1187 }
1188
1189 pub fn scroll_to_assist(
1190 &mut self,
1191 assist_id: InlineAssistId,
1192 window: &mut Window,
1193 cx: &mut App,
1194 ) {
1195 let Some(assist) = self.assists.get(&assist_id) else {
1196 return;
1197 };
1198 let Some(editor) = assist.editor.upgrade() else {
1199 return;
1200 };
1201
1202 let position = assist.range.start;
1203 editor.update(cx, |editor, cx| {
1204 editor.change_selections(SelectionEffects::no_scroll(), window, cx, |selections| {
1205 selections.select_anchor_ranges([position..position])
1206 });
1207
1208 let mut scroll_target_range = None;
1209 if let Some(decorations) = assist.decorations.as_ref() {
1210 scroll_target_range = maybe!({
1211 let top = editor.row_for_block(decorations.prompt_block_id, cx)?.0 as f64;
1212 let bottom = editor.row_for_block(decorations.end_block_id, cx)?.0 as f64;
1213 Some((top, bottom))
1214 });
1215 if scroll_target_range.is_none() {
1216 log::error!("bug: failed to find blocks for scrolling to inline assist");
1217 }
1218 }
1219 let scroll_target_range = scroll_target_range.unwrap_or_else(|| {
1220 let snapshot = editor.snapshot(window, cx);
1221 let start_row = assist
1222 .range
1223 .start
1224 .to_display_point(&snapshot.display_snapshot)
1225 .row();
1226 let top = start_row.0 as ScrollOffset;
1227 let bottom = top + 1.0;
1228 (top, bottom)
1229 });
1230 let mut scroll_target_top = scroll_target_range.0;
1231 let mut scroll_target_bottom = scroll_target_range.1;
1232
1233 scroll_target_top -= editor.vertical_scroll_margin() as ScrollOffset;
1234 scroll_target_bottom += editor.vertical_scroll_margin() as ScrollOffset;
1235
1236 let height_in_lines = editor.visible_line_count().unwrap_or(0.);
1237 let scroll_top = editor.scroll_position(cx).y;
1238 let scroll_bottom = scroll_top + height_in_lines;
1239
1240 if scroll_target_top < scroll_top {
1241 editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1242 } else if scroll_target_bottom > scroll_bottom {
1243 if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
1244 editor.set_scroll_position(
1245 point(0., scroll_target_bottom - height_in_lines),
1246 window,
1247 cx,
1248 );
1249 } else {
1250 editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1251 }
1252 }
1253 });
1254 }
1255
1256 fn unlink_assist_group(
1257 &mut self,
1258 assist_group_id: InlineAssistGroupId,
1259 window: &mut Window,
1260 cx: &mut App,
1261 ) -> Vec<InlineAssistId> {
1262 let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
1263 assist_group.linked = false;
1264
1265 for assist_id in &assist_group.assist_ids {
1266 let assist = self.assists.get_mut(assist_id).unwrap();
1267 if let Some(editor_decorations) = assist.decorations.as_ref() {
1268 editor_decorations
1269 .prompt_editor
1270 .update(cx, |prompt_editor, cx| prompt_editor.unlink(window, cx));
1271 }
1272 }
1273 assist_group.assist_ids.clone()
1274 }
1275
1276 pub fn start_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1277 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1278 assist
1279 } else {
1280 return;
1281 };
1282
1283 let assist_group_id = assist.group_id;
1284 if self.assist_groups[&assist_group_id].linked {
1285 for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1286 self.start_assist(assist_id, window, cx);
1287 }
1288 return;
1289 }
1290
1291 let Some((user_prompt, mention_set)) = assist.user_prompt(cx).zip(assist.mention_set(cx))
1292 else {
1293 return;
1294 };
1295
1296 self.prompt_history.retain(|prompt| *prompt != user_prompt);
1297 self.prompt_history.push_back(user_prompt.clone());
1298 if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
1299 self.prompt_history.pop_front();
1300 }
1301
1302 let Some(ConfiguredModel { model, .. }) =
1303 LanguageModelRegistry::read_global(cx).inline_assistant_model()
1304 else {
1305 return;
1306 };
1307
1308 let context_task = load_context(&mention_set, cx).shared();
1309 assist
1310 .codegen
1311 .update(cx, |codegen, cx| {
1312 codegen.start(model, user_prompt, context_task, cx)
1313 })
1314 .log_err();
1315 }
1316
1317 pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut App) {
1318 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1319 assist
1320 } else {
1321 return;
1322 };
1323
1324 assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
1325 }
1326
1327 fn update_editor_highlights(&self, editor: &Entity<Editor>, cx: &mut App) {
1328 let mut gutter_pending_ranges = Vec::new();
1329 let mut gutter_transformed_ranges = Vec::new();
1330 let mut foreground_ranges = Vec::new();
1331 let mut inserted_row_ranges = Vec::new();
1332 let empty_assist_ids = Vec::new();
1333 let assist_ids = self
1334 .assists_by_editor
1335 .get(&editor.downgrade())
1336 .map_or(&empty_assist_ids, |editor_assists| {
1337 &editor_assists.assist_ids
1338 });
1339
1340 for assist_id in assist_ids {
1341 if let Some(assist) = self.assists.get(assist_id) {
1342 let codegen = assist.codegen.read(cx);
1343 let buffer = codegen.buffer(cx).read(cx).read(cx);
1344 foreground_ranges.extend(codegen.last_equal_ranges(cx).iter().cloned());
1345
1346 let pending_range =
1347 codegen.edit_position(cx).unwrap_or(assist.range.start)..assist.range.end;
1348 if pending_range.end.to_offset(&buffer) > pending_range.start.to_offset(&buffer) {
1349 gutter_pending_ranges.push(pending_range);
1350 }
1351
1352 if let Some(edit_position) = codegen.edit_position(cx) {
1353 let edited_range = assist.range.start..edit_position;
1354 if edited_range.end.to_offset(&buffer) > edited_range.start.to_offset(&buffer) {
1355 gutter_transformed_ranges.push(edited_range);
1356 }
1357 }
1358
1359 if assist.decorations.is_some() {
1360 inserted_row_ranges
1361 .extend(codegen.diff(cx).inserted_row_ranges.iter().cloned());
1362 }
1363 }
1364 }
1365
1366 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
1367 merge_ranges(&mut foreground_ranges, &snapshot);
1368 merge_ranges(&mut gutter_pending_ranges, &snapshot);
1369 merge_ranges(&mut gutter_transformed_ranges, &snapshot);
1370 editor.update(cx, |editor, cx| {
1371 enum GutterPendingRange {}
1372 if gutter_pending_ranges.is_empty() {
1373 editor.clear_gutter_highlights::<GutterPendingRange>(cx);
1374 } else {
1375 editor.highlight_gutter::<GutterPendingRange>(
1376 gutter_pending_ranges,
1377 |cx| cx.theme().status().info_background,
1378 cx,
1379 )
1380 }
1381
1382 enum GutterTransformedRange {}
1383 if gutter_transformed_ranges.is_empty() {
1384 editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
1385 } else {
1386 editor.highlight_gutter::<GutterTransformedRange>(
1387 gutter_transformed_ranges,
1388 |cx| cx.theme().status().info,
1389 cx,
1390 )
1391 }
1392
1393 if foreground_ranges.is_empty() {
1394 editor.clear_highlights::<InlineAssist>(cx);
1395 } else {
1396 editor.highlight_text::<InlineAssist>(
1397 foreground_ranges,
1398 HighlightStyle {
1399 fade_out: Some(0.6),
1400 ..Default::default()
1401 },
1402 cx,
1403 );
1404 }
1405
1406 editor.clear_row_highlights::<InlineAssist>();
1407 for row_range in inserted_row_ranges {
1408 editor.highlight_rows::<InlineAssist>(
1409 row_range,
1410 cx.theme().status().info_background,
1411 Default::default(),
1412 cx,
1413 );
1414 }
1415 });
1416 }
1417
1418 fn update_editor_blocks(
1419 &mut self,
1420 editor: &Entity<Editor>,
1421 assist_id: InlineAssistId,
1422 window: &mut Window,
1423 cx: &mut App,
1424 ) {
1425 let Some(assist) = self.assists.get_mut(&assist_id) else {
1426 return;
1427 };
1428 let Some(decorations) = assist.decorations.as_mut() else {
1429 return;
1430 };
1431
1432 let codegen = assist.codegen.read(cx);
1433 let old_snapshot = codegen.snapshot(cx);
1434 let old_buffer = codegen.old_buffer(cx);
1435 let deleted_row_ranges = codegen.diff(cx).deleted_row_ranges.clone();
1436
1437 editor.update(cx, |editor, cx| {
1438 let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
1439 editor.remove_blocks(old_blocks, None, cx);
1440
1441 let mut new_blocks = Vec::new();
1442 for (new_row, old_row_range) in deleted_row_ranges {
1443 let (_, buffer_start) = old_snapshot
1444 .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
1445 .unwrap();
1446 let (_, buffer_end) = old_snapshot
1447 .point_to_buffer_offset(Point::new(
1448 *old_row_range.end(),
1449 old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
1450 ))
1451 .unwrap();
1452
1453 let deleted_lines_editor = cx.new(|cx| {
1454 let multi_buffer =
1455 cx.new(|_| MultiBuffer::without_headers(language::Capability::ReadOnly));
1456 multi_buffer.update(cx, |multi_buffer, cx| {
1457 multi_buffer.push_excerpts(
1458 old_buffer.clone(),
1459 // todo(lw): buffer_start and buffer_end might come from different snapshots!
1460 Some(ExcerptRange::new(buffer_start..buffer_end)),
1461 cx,
1462 );
1463 });
1464
1465 enum DeletedLines {}
1466 let mut editor = Editor::for_multibuffer(multi_buffer, None, window, cx);
1467 editor.disable_scrollbars_and_minimap(window, cx);
1468 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1469 editor.set_show_wrap_guides(false, cx);
1470 editor.set_show_gutter(false, cx);
1471 editor.set_offset_content(false, cx);
1472 editor.scroll_manager.set_forbid_vertical_scroll(true);
1473 editor.set_read_only(true);
1474 editor.set_show_edit_predictions(Some(false), window, cx);
1475 editor.highlight_rows::<DeletedLines>(
1476 Anchor::min()..Anchor::max(),
1477 cx.theme().status().deleted_background,
1478 Default::default(),
1479 cx,
1480 );
1481 editor
1482 });
1483
1484 let height =
1485 deleted_lines_editor.update(cx, |editor, cx| editor.max_point(cx).row().0 + 1);
1486 new_blocks.push(BlockProperties {
1487 placement: BlockPlacement::Above(new_row),
1488 height: Some(height),
1489 style: BlockStyle::Flex,
1490 render: Arc::new(move |cx| {
1491 div()
1492 .block_mouse_except_scroll()
1493 .bg(cx.theme().status().deleted_background)
1494 .size_full()
1495 .h(height as f32 * cx.window.line_height())
1496 .pl(cx.margins.gutter.full_width())
1497 .child(deleted_lines_editor.clone())
1498 .into_any_element()
1499 }),
1500 priority: 0,
1501 });
1502 }
1503
1504 decorations.removed_line_block_ids = editor
1505 .insert_blocks(new_blocks, None, cx)
1506 .into_iter()
1507 .collect();
1508 })
1509 }
1510
1511 fn resolve_inline_assist_target(
1512 workspace: &mut Workspace,
1513 agent_panel: Option<Entity<AgentPanel>>,
1514 window: &mut Window,
1515 cx: &mut App,
1516 ) -> Option<InlineAssistTarget> {
1517 if let Some(terminal_panel) = workspace.panel::<TerminalPanel>(cx)
1518 && terminal_panel
1519 .read(cx)
1520 .focus_handle(cx)
1521 .contains_focused(window, cx)
1522 && let Some(terminal_view) = terminal_panel.read(cx).pane().and_then(|pane| {
1523 pane.read(cx)
1524 .active_item()
1525 .and_then(|t| t.downcast::<TerminalView>())
1526 })
1527 {
1528 return Some(InlineAssistTarget::Terminal(terminal_view));
1529 }
1530
1531 let text_thread_editor = agent_panel
1532 .and_then(|panel| panel.read(cx).active_text_thread_editor())
1533 .and_then(|editor| {
1534 let editor = &editor.read(cx).editor().clone();
1535 if editor.read(cx).is_focused(window) {
1536 Some(editor.clone())
1537 } else {
1538 None
1539 }
1540 });
1541
1542 if let Some(text_thread_editor) = text_thread_editor {
1543 Some(InlineAssistTarget::Editor(text_thread_editor))
1544 } else if let Some(workspace_editor) = workspace
1545 .active_item(cx)
1546 .and_then(|item| item.act_as::<Editor>(cx))
1547 {
1548 Some(InlineAssistTarget::Editor(workspace_editor))
1549 } else {
1550 workspace
1551 .active_item(cx)
1552 .and_then(|item| item.act_as::<TerminalView>(cx))
1553 .map(InlineAssistTarget::Terminal)
1554 }
1555 }
1556}
1557
1558struct EditorInlineAssists {
1559 assist_ids: Vec<InlineAssistId>,
1560 scroll_lock: Option<InlineAssistScrollLock>,
1561 highlight_updates: watch::Sender<()>,
1562 _update_highlights: Task<Result<()>>,
1563 _subscriptions: Vec<gpui::Subscription>,
1564}
1565
1566struct InlineAssistScrollLock {
1567 assist_id: InlineAssistId,
1568 distance_from_top: ScrollOffset,
1569}
1570
1571impl EditorInlineAssists {
1572 fn new(editor: &Entity<Editor>, window: &mut Window, cx: &mut App) -> Self {
1573 let (highlight_updates_tx, mut highlight_updates_rx) = watch::channel(());
1574 Self {
1575 assist_ids: Vec::new(),
1576 scroll_lock: None,
1577 highlight_updates: highlight_updates_tx,
1578 _update_highlights: cx.spawn({
1579 let editor = editor.downgrade();
1580 async move |cx| {
1581 while let Ok(()) = highlight_updates_rx.changed().await {
1582 let editor = editor.upgrade().context("editor was dropped")?;
1583 cx.update_global(|assistant: &mut InlineAssistant, cx| {
1584 assistant.update_editor_highlights(&editor, cx);
1585 })?;
1586 }
1587 Ok(())
1588 }
1589 }),
1590 _subscriptions: vec![
1591 cx.observe_release_in(editor, window, {
1592 let editor = editor.downgrade();
1593 |_, window, cx| {
1594 InlineAssistant::update_global(cx, |this, cx| {
1595 this.handle_editor_release(editor, window, cx);
1596 })
1597 }
1598 }),
1599 window.observe(editor, cx, move |editor, window, cx| {
1600 InlineAssistant::update_global(cx, |this, cx| {
1601 this.handle_editor_change(editor, window, cx)
1602 })
1603 }),
1604 window.subscribe(editor, cx, move |editor, event, window, cx| {
1605 InlineAssistant::update_global(cx, |this, cx| {
1606 this.handle_editor_event(editor, event, window, cx)
1607 })
1608 }),
1609 editor.update(cx, |editor, cx| {
1610 let editor_handle = cx.entity().downgrade();
1611 editor.register_action(move |_: &editor::actions::Newline, window, cx| {
1612 InlineAssistant::update_global(cx, |this, cx| {
1613 if let Some(editor) = editor_handle.upgrade() {
1614 this.handle_editor_newline(editor, window, cx)
1615 }
1616 })
1617 })
1618 }),
1619 editor.update(cx, |editor, cx| {
1620 let editor_handle = cx.entity().downgrade();
1621 editor.register_action(move |_: &editor::actions::Cancel, window, cx| {
1622 InlineAssistant::update_global(cx, |this, cx| {
1623 if let Some(editor) = editor_handle.upgrade() {
1624 this.handle_editor_cancel(editor, window, cx)
1625 }
1626 })
1627 })
1628 }),
1629 ],
1630 }
1631 }
1632}
1633
1634struct InlineAssistGroup {
1635 assist_ids: Vec<InlineAssistId>,
1636 linked: bool,
1637 active_assist_id: Option<InlineAssistId>,
1638}
1639
1640impl InlineAssistGroup {
1641 fn new() -> Self {
1642 Self {
1643 assist_ids: Vec::new(),
1644 linked: true,
1645 active_assist_id: None,
1646 }
1647 }
1648}
1649
1650fn build_assist_editor_renderer(editor: &Entity<PromptEditor<BufferCodegen>>) -> RenderBlock {
1651 let editor = editor.clone();
1652
1653 Arc::new(move |cx: &mut BlockContext| {
1654 let editor_margins = editor.read(cx).editor_margins();
1655
1656 *editor_margins.lock() = *cx.margins;
1657 editor.clone().into_any_element()
1658 })
1659}
1660
1661#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1662struct InlineAssistGroupId(usize);
1663
1664impl InlineAssistGroupId {
1665 fn post_inc(&mut self) -> InlineAssistGroupId {
1666 let id = *self;
1667 self.0 += 1;
1668 id
1669 }
1670}
1671
1672pub struct InlineAssist {
1673 group_id: InlineAssistGroupId,
1674 range: Range<Anchor>,
1675 editor: WeakEntity<Editor>,
1676 decorations: Option<InlineAssistDecorations>,
1677 codegen: Entity<BufferCodegen>,
1678 _subscriptions: Vec<Subscription>,
1679 workspace: WeakEntity<Workspace>,
1680}
1681
1682impl InlineAssist {
1683 fn new(
1684 assist_id: InlineAssistId,
1685 group_id: InlineAssistGroupId,
1686 editor: &Entity<Editor>,
1687 prompt_editor: &Entity<PromptEditor<BufferCodegen>>,
1688 prompt_block_id: CustomBlockId,
1689 end_block_id: CustomBlockId,
1690 range: Range<Anchor>,
1691 codegen: Entity<BufferCodegen>,
1692 workspace: WeakEntity<Workspace>,
1693 window: &mut Window,
1694 cx: &mut App,
1695 ) -> Self {
1696 let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
1697 InlineAssist {
1698 group_id,
1699 editor: editor.downgrade(),
1700 decorations: Some(InlineAssistDecorations {
1701 prompt_block_id,
1702 prompt_editor: prompt_editor.clone(),
1703 removed_line_block_ids: HashSet::default(),
1704 end_block_id,
1705 }),
1706 range,
1707 codegen: codegen.clone(),
1708 workspace,
1709 _subscriptions: vec![
1710 window.on_focus_in(&prompt_editor_focus_handle, cx, move |_, cx| {
1711 InlineAssistant::update_global(cx, |this, cx| {
1712 this.handle_prompt_editor_focus_in(assist_id, cx)
1713 })
1714 }),
1715 window.on_focus_out(&prompt_editor_focus_handle, cx, move |_, _, cx| {
1716 InlineAssistant::update_global(cx, |this, cx| {
1717 this.handle_prompt_editor_focus_out(assist_id, cx)
1718 })
1719 }),
1720 window.subscribe(prompt_editor, cx, |prompt_editor, event, window, cx| {
1721 InlineAssistant::update_global(cx, |this, cx| {
1722 this.handle_prompt_editor_event(prompt_editor, event, window, cx)
1723 })
1724 }),
1725 window.observe(&codegen, cx, {
1726 let editor = editor.downgrade();
1727 move |_, window, cx| {
1728 if let Some(editor) = editor.upgrade() {
1729 InlineAssistant::update_global(cx, |this, cx| {
1730 if let Some(editor_assists) =
1731 this.assists_by_editor.get_mut(&editor.downgrade())
1732 {
1733 editor_assists.highlight_updates.send(()).ok();
1734 }
1735
1736 this.update_editor_blocks(&editor, assist_id, window, cx);
1737 })
1738 }
1739 }
1740 }),
1741 window.subscribe(&codegen, cx, move |codegen, event, window, cx| {
1742 InlineAssistant::update_global(cx, |this, cx| match event {
1743 CodegenEvent::Undone => this.finish_assist(assist_id, false, window, cx),
1744 CodegenEvent::Finished => {
1745 let assist = if let Some(assist) = this.assists.get(&assist_id) {
1746 assist
1747 } else {
1748 return;
1749 };
1750
1751 if let CodegenStatus::Error(error) = codegen.read(cx).status(cx)
1752 && assist.decorations.is_none()
1753 && let Some(workspace) = assist.workspace.upgrade()
1754 {
1755 #[cfg(any(test, feature = "test-support"))]
1756 if let Some(sender) = &mut this._inline_assistant_completions {
1757 sender
1758 .unbounded_send(Err(anyhow::anyhow!(
1759 "Inline assistant error: {}",
1760 error
1761 )))
1762 .ok();
1763 }
1764
1765 let error = format!("Inline assistant error: {}", error);
1766 workspace.update(cx, |workspace, cx| {
1767 struct InlineAssistantError;
1768
1769 let id = NotificationId::composite::<InlineAssistantError>(
1770 assist_id.0,
1771 );
1772
1773 workspace.show_toast(Toast::new(id, error), cx);
1774 })
1775 } else {
1776 #[cfg(any(test, feature = "test-support"))]
1777 if let Some(sender) = &mut this._inline_assistant_completions {
1778 sender.unbounded_send(Ok(assist_id)).ok();
1779 }
1780 }
1781
1782 if assist.decorations.is_none() {
1783 this.finish_assist(assist_id, false, window, cx);
1784 }
1785 }
1786 })
1787 }),
1788 ],
1789 }
1790 }
1791
1792 fn user_prompt(&self, cx: &App) -> Option<String> {
1793 let decorations = self.decorations.as_ref()?;
1794 Some(decorations.prompt_editor.read(cx).prompt(cx))
1795 }
1796
1797 fn mention_set(&self, cx: &App) -> Option<Entity<MentionSet>> {
1798 let decorations = self.decorations.as_ref()?;
1799 Some(decorations.prompt_editor.read(cx).mention_set().clone())
1800 }
1801}
1802
1803struct InlineAssistDecorations {
1804 prompt_block_id: CustomBlockId,
1805 prompt_editor: Entity<PromptEditor<BufferCodegen>>,
1806 removed_line_block_ids: HashSet<CustomBlockId>,
1807 end_block_id: CustomBlockId,
1808}
1809
1810struct AssistantCodeActionProvider {
1811 editor: WeakEntity<Editor>,
1812 workspace: WeakEntity<Workspace>,
1813}
1814
1815const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant";
1816
1817impl CodeActionProvider for AssistantCodeActionProvider {
1818 fn id(&self) -> Arc<str> {
1819 ASSISTANT_CODE_ACTION_PROVIDER_ID.into()
1820 }
1821
1822 fn code_actions(
1823 &self,
1824 buffer: &Entity<Buffer>,
1825 range: Range<text::Anchor>,
1826 _: &mut Window,
1827 cx: &mut App,
1828 ) -> Task<Result<Vec<CodeAction>>> {
1829 if !AgentSettings::get_global(cx).enabled(cx) {
1830 return Task::ready(Ok(Vec::new()));
1831 }
1832
1833 let snapshot = buffer.read(cx).snapshot();
1834 let mut range = range.to_point(&snapshot);
1835
1836 // Expand the range to line boundaries.
1837 range.start.column = 0;
1838 range.end.column = snapshot.line_len(range.end.row);
1839
1840 let mut has_diagnostics = false;
1841 for diagnostic in snapshot.diagnostics_in_range::<_, Point>(range.clone(), false) {
1842 range.start = cmp::min(range.start, diagnostic.range.start);
1843 range.end = cmp::max(range.end, diagnostic.range.end);
1844 has_diagnostics = true;
1845 }
1846 if has_diagnostics {
1847 let symbols_containing_start = snapshot.symbols_containing(range.start, None);
1848 if let Some(symbol) = symbols_containing_start.last() {
1849 range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1850 range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1851 }
1852 let symbols_containing_end = snapshot.symbols_containing(range.end, None);
1853 if let Some(symbol) = symbols_containing_end.last() {
1854 range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1855 range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1856 }
1857
1858 Task::ready(Ok(vec![CodeAction {
1859 server_id: language::LanguageServerId(0),
1860 range: snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end),
1861 lsp_action: LspAction::Action(Box::new(lsp::CodeAction {
1862 title: "Fix with Assistant".into(),
1863 ..Default::default()
1864 })),
1865 resolved: true,
1866 }]))
1867 } else {
1868 Task::ready(Ok(Vec::new()))
1869 }
1870 }
1871
1872 fn apply_code_action(
1873 &self,
1874 buffer: Entity<Buffer>,
1875 action: CodeAction,
1876 excerpt_id: ExcerptId,
1877 _push_to_history: bool,
1878 window: &mut Window,
1879 cx: &mut App,
1880 ) -> Task<Result<ProjectTransaction>> {
1881 let editor = self.editor.clone();
1882 let workspace = self.workspace.clone();
1883 let prompt_store = PromptStore::global(cx);
1884 window.spawn(cx, async move |cx| {
1885 let workspace = workspace.upgrade().context("workspace was released")?;
1886 let thread_store = cx.update(|_window, cx| {
1887 anyhow::Ok(
1888 workspace
1889 .read(cx)
1890 .panel::<AgentPanel>(cx)
1891 .context("missing agent panel")?
1892 .read(cx)
1893 .thread_store()
1894 .clone(),
1895 )
1896 })??;
1897 let editor = editor.upgrade().context("editor was released")?;
1898 let range = editor
1899 .update(cx, |editor, cx| {
1900 editor.buffer().update(cx, |multibuffer, cx| {
1901 let buffer = buffer.read(cx);
1902 let multibuffer_snapshot = multibuffer.read(cx);
1903
1904 let old_context_range =
1905 multibuffer_snapshot.context_range_for_excerpt(excerpt_id)?;
1906 let mut new_context_range = old_context_range.clone();
1907 if action
1908 .range
1909 .start
1910 .cmp(&old_context_range.start, buffer)
1911 .is_lt()
1912 {
1913 new_context_range.start = action.range.start;
1914 }
1915 if action.range.end.cmp(&old_context_range.end, buffer).is_gt() {
1916 new_context_range.end = action.range.end;
1917 }
1918 drop(multibuffer_snapshot);
1919
1920 if new_context_range != old_context_range {
1921 multibuffer.resize_excerpt(excerpt_id, new_context_range, cx);
1922 }
1923
1924 let multibuffer_snapshot = multibuffer.read(cx);
1925 multibuffer_snapshot.anchor_range_in_excerpt(excerpt_id, action.range)
1926 })
1927 })?
1928 .context("invalid range")?;
1929
1930 let prompt_store = prompt_store.await.ok();
1931 cx.update_global(|assistant: &mut InlineAssistant, window, cx| {
1932 let assist_id = assistant.suggest_assist(
1933 &editor,
1934 range,
1935 "Fix Diagnostics".into(),
1936 None,
1937 true,
1938 workspace,
1939 thread_store,
1940 prompt_store,
1941 window,
1942 cx,
1943 );
1944 assistant.start_assist(assist_id, window, cx);
1945 })?;
1946
1947 Ok(ProjectTransaction::default())
1948 })
1949 }
1950}
1951
1952fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
1953 ranges.sort_unstable_by(|a, b| {
1954 a.start
1955 .cmp(&b.start, buffer)
1956 .then_with(|| b.end.cmp(&a.end, buffer))
1957 });
1958
1959 let mut ix = 0;
1960 while ix + 1 < ranges.len() {
1961 let b = ranges[ix + 1].clone();
1962 let a = &mut ranges[ix];
1963 if a.end.cmp(&b.start, buffer).is_gt() {
1964 if a.end.cmp(&b.end, buffer).is_lt() {
1965 a.end = b.end;
1966 }
1967 ranges.remove(ix + 1);
1968 } else {
1969 ix += 1;
1970 }
1971 }
1972}
1973
1974#[cfg(any(test, feature = "test-support"))]
1975pub mod test {
1976 use std::sync::Arc;
1977
1978 use agent::HistoryStore;
1979 use assistant_text_thread::TextThreadStore;
1980 use client::{Client, UserStore};
1981 use editor::{Editor, MultiBuffer, MultiBufferOffset};
1982 use fs::FakeFs;
1983 use futures::channel::mpsc;
1984 use gpui::{AppContext, TestAppContext, UpdateGlobal as _};
1985 use language::Buffer;
1986 use language_model::LanguageModelRegistry;
1987 use project::Project;
1988 use prompt_store::PromptBuilder;
1989 use smol::stream::StreamExt as _;
1990 use util::test::marked_text_ranges;
1991 use workspace::Workspace;
1992
1993 use crate::InlineAssistant;
1994
1995 pub fn run_inline_assistant_test<SetupF, TestF>(
1996 base_buffer: String,
1997 prompt: String,
1998 setup: SetupF,
1999 test: TestF,
2000 cx: &mut TestAppContext,
2001 ) -> String
2002 where
2003 SetupF: FnOnce(&mut gpui::VisualTestContext),
2004 TestF: FnOnce(&mut gpui::VisualTestContext),
2005 {
2006 let fs = FakeFs::new(cx.executor());
2007 let app_state = cx.update(|cx| workspace::AppState::test(cx));
2008 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2009 let http = Arc::new(reqwest_client::ReqwestClient::user_agent("agent tests").unwrap());
2010 let client = cx.update(|cx| {
2011 cx.set_http_client(http);
2012 Client::production(cx)
2013 });
2014 let mut inline_assistant =
2015 InlineAssistant::new(fs.clone(), prompt_builder, client.telemetry().clone());
2016
2017 let (tx, mut completion_rx) = mpsc::unbounded();
2018 inline_assistant.set_completion_receiver(tx);
2019
2020 // Initialize settings and client
2021 cx.update(|cx| {
2022 gpui_tokio::init(cx);
2023 settings::init(cx);
2024 client::init(&client, cx);
2025 workspace::init(app_state.clone(), cx);
2026 let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
2027 language_model::init(client.clone(), cx);
2028 language_models::init(user_store, client.clone(), cx);
2029
2030 cx.set_global(inline_assistant);
2031 });
2032
2033 let project = cx
2034 .executor()
2035 .block_test(async { Project::test(fs.clone(), [], cx).await });
2036
2037 // Create workspace with window
2038 let (workspace, cx) = cx.add_window_view(|window, cx| {
2039 window.activate_window();
2040 Workspace::new(None, project.clone(), app_state.clone(), window, cx)
2041 });
2042
2043 setup(cx);
2044
2045 let (_editor, buffer) = cx.update(|window, cx| {
2046 let buffer = cx.new(|cx| Buffer::local("", cx));
2047 let multibuffer = cx.new(|cx| MultiBuffer::singleton(buffer.clone(), cx));
2048 let editor = cx.new(|cx| Editor::for_multibuffer(multibuffer, None, window, cx));
2049 editor.update(cx, |editor, cx| {
2050 let (unmarked_text, selection_ranges) = marked_text_ranges(&base_buffer, true);
2051 editor.set_text(unmarked_text, window, cx);
2052 editor.change_selections(Default::default(), window, cx, |s| {
2053 s.select_ranges(
2054 selection_ranges.into_iter().map(|range| {
2055 MultiBufferOffset(range.start)..MultiBufferOffset(range.end)
2056 }),
2057 )
2058 })
2059 });
2060
2061 let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
2062 let history_store = cx.new(|cx| HistoryStore::new(text_thread_store, cx));
2063
2064 // Add editor to workspace
2065 workspace.update(cx, |workspace, cx| {
2066 workspace.add_item_to_active_pane(Box::new(editor.clone()), None, true, window, cx);
2067 });
2068
2069 // Call assist method
2070 InlineAssistant::update_global(cx, |inline_assistant, cx| {
2071 let assist_id = inline_assistant
2072 .assist(
2073 &editor,
2074 workspace.downgrade(),
2075 project.downgrade(),
2076 history_store, // thread_store
2077 None, // prompt_store
2078 Some(prompt),
2079 window,
2080 cx,
2081 )
2082 .unwrap();
2083
2084 inline_assistant.start_assist(assist_id, window, cx);
2085 });
2086
2087 (editor, buffer)
2088 });
2089
2090 cx.run_until_parked();
2091
2092 test(cx);
2093
2094 cx.executor()
2095 .block_test(async { completion_rx.next().await });
2096
2097 buffer.read_with(cx, |buffer, _| buffer.text())
2098 }
2099
2100 #[allow(unused)]
2101 pub fn test_inline_assistant(
2102 base_buffer: &'static str,
2103 llm_output: &'static str,
2104 cx: &mut TestAppContext,
2105 ) -> String {
2106 run_inline_assistant_test(
2107 base_buffer.to_string(),
2108 "Prompt doesn't matter because we're using a fake model".to_string(),
2109 |cx| {
2110 cx.update(|_, cx| LanguageModelRegistry::test(cx));
2111 },
2112 |cx| {
2113 let fake_model = cx.update(|_, cx| {
2114 LanguageModelRegistry::global(cx)
2115 .update(cx, |registry, _| registry.fake_model())
2116 });
2117 let fake = fake_model.as_fake();
2118
2119 // let fake = fake_model;
2120 fake.send_last_completion_stream_text_chunk(llm_output.to_string());
2121 fake.end_last_completion_stream();
2122
2123 // Run again to process the model's response
2124 cx.run_until_parked();
2125 },
2126 cx,
2127 )
2128 }
2129}