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