1use crate::{
2 assistant_settings::AssistantSettings, humanize_token_count, prompts::PromptBuilder,
3 AssistantPanel, AssistantPanelEvent, CharOperation, CycleNextInlineAssist,
4 CyclePreviousInlineAssist, LineDiff, LineOperation, RequestType, StreamingDiff,
5};
6use anyhow::{anyhow, Context as _, Result};
7use client::{telemetry::Telemetry, ErrorExt};
8use collections::{hash_map, HashMap, HashSet, VecDeque};
9use editor::{
10 actions::{MoveDown, MoveUp, SelectAll},
11 display_map::{
12 BlockContext, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,
13 ToDisplayPoint,
14 },
15 Anchor, AnchorRangeExt, CodeActionProvider, Editor, EditorElement, EditorEvent, EditorMode,
16 EditorStyle, ExcerptId, ExcerptRange, GutterDimensions, MultiBuffer, MultiBufferSnapshot,
17 ToOffset as _, ToPoint,
18};
19use feature_flags::{
20 Assistant2FeatureFlag, FeatureFlagAppExt as _, FeatureFlagViewExt as _, ZedPro,
21};
22use fs::Fs;
23use futures::{
24 channel::mpsc,
25 future::{BoxFuture, LocalBoxFuture},
26 join, SinkExt, Stream, StreamExt,
27};
28use gpui::{
29 anchored, deferred, point, AnyElement, AppContext, ClickEvent, CursorStyle, EventEmitter,
30 FocusHandle, FocusableView, FontWeight, Global, HighlightStyle, Model, ModelContext,
31 Subscription, Task, TextStyle, UpdateGlobal, View, ViewContext, WeakView, WindowContext,
32};
33use language::{Buffer, IndentKind, Point, Selection, TransactionId};
34use language_model::{
35 LanguageModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
36 LanguageModelTextStream, Role,
37};
38use language_model_selector::{LanguageModelSelector, LanguageModelSelectorPopoverMenu};
39use language_models::report_assistant_event;
40use multi_buffer::MultiBufferRow;
41use parking_lot::Mutex;
42use project::{CodeAction, ProjectTransaction};
43use rope::Rope;
44use settings::{update_settings_file, Settings, SettingsStore};
45use smol::future::FutureExt;
46use std::{
47 cmp,
48 future::{self, Future},
49 iter, mem,
50 ops::{Range, RangeInclusive},
51 pin::Pin,
52 rc::Rc,
53 sync::Arc,
54 task::{self, Poll},
55 time::{Duration, Instant},
56};
57use telemetry_events::{AssistantEvent, AssistantKind, AssistantPhase};
58use terminal_view::terminal_panel::TerminalPanel;
59use text::{OffsetRangeExt, ToPoint as _};
60use theme::ThemeSettings;
61use ui::{
62 prelude::*, text_for_action, CheckboxWithLabel, IconButtonShape, KeyBinding, Popover, Tooltip,
63};
64use util::{RangeExt, ResultExt};
65use workspace::{notifications::NotificationId, ItemHandle, Toast, Workspace};
66
67pub fn init(
68 fs: Arc<dyn Fs>,
69 prompt_builder: Arc<PromptBuilder>,
70 telemetry: Arc<Telemetry>,
71 cx: &mut AppContext,
72) {
73 cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));
74 cx.observe_new_views(|_, cx| {
75 cx.observe_flag::<Assistant2FeatureFlag, _>({
76 |is_assistant2_enabled, _view, cx| {
77 if is_assistant2_enabled {
78 // Assistant2 enabled, nothing to do for Assistant1.
79 } else {
80 let workspace = cx.view().clone();
81 InlineAssistant::update_global(cx, |inline_assistant, cx| {
82 inline_assistant.register_workspace(&workspace, cx)
83 })
84 }
85 }
86 })
87 .detach();
88 })
89 .detach();
90}
91
92const PROMPT_HISTORY_MAX_LEN: usize = 20;
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<WeakView<Editor>, EditorInlineAssists>,
99 assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
100 confirmed_assists: HashMap<InlineAssistId, Model<CodegenAlternative>>,
101 prompt_history: VecDeque<String>,
102 prompt_builder: Arc<PromptBuilder>,
103 telemetry: Arc<Telemetry>,
104 fs: Arc<dyn Fs>,
105}
106
107impl Global for InlineAssistant {}
108
109impl InlineAssistant {
110 pub fn new(
111 fs: Arc<dyn Fs>,
112 prompt_builder: Arc<PromptBuilder>,
113 telemetry: Arc<Telemetry>,
114 ) -> Self {
115 Self {
116 next_assist_id: InlineAssistId::default(),
117 next_assist_group_id: InlineAssistGroupId::default(),
118 assists: HashMap::default(),
119 assists_by_editor: HashMap::default(),
120 assist_groups: HashMap::default(),
121 confirmed_assists: HashMap::default(),
122 prompt_history: VecDeque::default(),
123 prompt_builder,
124 telemetry,
125 fs,
126 }
127 }
128
129 pub fn register_workspace(&mut self, workspace: &View<Workspace>, cx: &mut WindowContext) {
130 cx.subscribe(workspace, |workspace, event, cx| {
131 Self::update_global(cx, |this, cx| {
132 this.handle_workspace_event(workspace, event, cx)
133 });
134 })
135 .detach();
136
137 let workspace = workspace.downgrade();
138 cx.observe_global::<SettingsStore>(move |cx| {
139 let Some(workspace) = workspace.upgrade() else {
140 return;
141 };
142 let Some(terminal_panel) = workspace.read(cx).panel::<TerminalPanel>(cx) else {
143 return;
144 };
145 let enabled = AssistantSettings::get_global(cx).enabled;
146 terminal_panel.update(cx, |terminal_panel, cx| {
147 terminal_panel.set_assistant_enabled(enabled, cx)
148 });
149 })
150 .detach();
151 }
152
153 fn handle_workspace_event(
154 &mut self,
155 workspace: View<Workspace>,
156 event: &workspace::Event,
157 cx: &mut WindowContext,
158 ) {
159 match event {
160 workspace::Event::UserSavedItem { item, .. } => {
161 // When the user manually saves an editor, automatically accepts all finished transformations.
162 if let Some(editor) = item.upgrade().and_then(|item| item.act_as::<Editor>(cx)) {
163 if let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) {
164 for assist_id in editor_assists.assist_ids.clone() {
165 let assist = &self.assists[&assist_id];
166 if let CodegenStatus::Done = assist.codegen.read(cx).status(cx) {
167 self.finish_assist(assist_id, false, cx)
168 }
169 }
170 }
171 }
172 }
173 workspace::Event::ItemAdded { item } => {
174 self.register_workspace_item(&workspace, item.as_ref(), cx);
175 }
176 _ => (),
177 }
178 }
179
180 fn register_workspace_item(
181 &mut self,
182 workspace: &View<Workspace>,
183 item: &dyn ItemHandle,
184 cx: &mut WindowContext,
185 ) {
186 if let Some(editor) = item.act_as::<Editor>(cx) {
187 editor.update(cx, |editor, cx| {
188 editor.push_code_action_provider(
189 Rc::new(AssistantCodeActionProvider {
190 editor: cx.view().downgrade(),
191 workspace: workspace.downgrade(),
192 }),
193 cx,
194 );
195 });
196 }
197 }
198
199 pub fn assist(
200 &mut self,
201 editor: &View<Editor>,
202 workspace: Option<WeakView<Workspace>>,
203 assistant_panel: Option<&View<AssistantPanel>>,
204 initial_prompt: Option<String>,
205 cx: &mut WindowContext,
206 ) {
207 let (snapshot, initial_selections) = editor.update(cx, |editor, cx| {
208 (
209 editor.buffer().read(cx).snapshot(cx),
210 editor.selections.all::<Point>(cx),
211 )
212 });
213
214 let mut selections = Vec::<Selection<Point>>::new();
215 let mut newest_selection = None;
216 for mut selection in initial_selections {
217 if selection.end > selection.start {
218 selection.start.column = 0;
219 // If the selection ends at the start of the line, we don't want to include it.
220 if selection.end.column == 0 {
221 selection.end.row -= 1;
222 }
223 selection.end.column = snapshot.line_len(MultiBufferRow(selection.end.row));
224 }
225
226 if let Some(prev_selection) = selections.last_mut() {
227 if selection.start <= prev_selection.end {
228 prev_selection.end = selection.end;
229 continue;
230 }
231 }
232
233 let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
234 if selection.id > latest_selection.id {
235 *latest_selection = selection.clone();
236 }
237 selections.push(selection);
238 }
239 let newest_selection = newest_selection.unwrap();
240
241 let mut codegen_ranges = Vec::new();
242 for (excerpt_id, buffer, buffer_range) in
243 snapshot.excerpts_in_ranges(selections.iter().map(|selection| {
244 snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
245 }))
246 {
247 let start = Anchor {
248 buffer_id: Some(buffer.remote_id()),
249 excerpt_id,
250 text_anchor: buffer.anchor_before(buffer_range.start),
251 };
252 let end = Anchor {
253 buffer_id: Some(buffer.remote_id()),
254 excerpt_id,
255 text_anchor: buffer.anchor_after(buffer_range.end),
256 };
257 codegen_ranges.push(start..end);
258
259 if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
260 self.telemetry.report_assistant_event(AssistantEvent {
261 conversation_id: None,
262 kind: AssistantKind::Inline,
263 phase: AssistantPhase::Invoked,
264 message_id: None,
265 model: model.telemetry_id(),
266 model_provider: model.provider_id().to_string(),
267 response_latency: None,
268 error_message: None,
269 language_name: buffer.language().map(|language| language.name().to_proto()),
270 });
271 }
272 }
273
274 let assist_group_id = self.next_assist_group_id.post_inc();
275 let prompt_buffer =
276 cx.new_model(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx));
277 let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
278
279 let mut assists = Vec::new();
280 let mut assist_to_focus = None;
281 for range in codegen_ranges {
282 let assist_id = self.next_assist_id.post_inc();
283 let codegen = cx.new_model(|cx| {
284 Codegen::new(
285 editor.read(cx).buffer().clone(),
286 range.clone(),
287 None,
288 self.telemetry.clone(),
289 self.prompt_builder.clone(),
290 cx,
291 )
292 });
293
294 let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
295 let prompt_editor = cx.new_view(|cx| {
296 PromptEditor::new(
297 assist_id,
298 gutter_dimensions.clone(),
299 self.prompt_history.clone(),
300 prompt_buffer.clone(),
301 codegen.clone(),
302 editor,
303 assistant_panel,
304 workspace.clone(),
305 self.fs.clone(),
306 cx,
307 )
308 });
309
310 if assist_to_focus.is_none() {
311 let focus_assist = if newest_selection.reversed {
312 range.start.to_point(&snapshot) == newest_selection.start
313 } else {
314 range.end.to_point(&snapshot) == newest_selection.end
315 };
316 if focus_assist {
317 assist_to_focus = Some(assist_id);
318 }
319 }
320
321 let [prompt_block_id, end_block_id] =
322 self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
323
324 assists.push((
325 assist_id,
326 range,
327 prompt_editor,
328 prompt_block_id,
329 end_block_id,
330 ));
331 }
332
333 let editor_assists = self
334 .assists_by_editor
335 .entry(editor.downgrade())
336 .or_insert_with(|| EditorInlineAssists::new(&editor, cx));
337 let mut assist_group = InlineAssistGroup::new();
338 for (assist_id, range, prompt_editor, prompt_block_id, end_block_id) in assists {
339 self.assists.insert(
340 assist_id,
341 InlineAssist::new(
342 assist_id,
343 assist_group_id,
344 assistant_panel.is_some(),
345 editor,
346 &prompt_editor,
347 prompt_block_id,
348 end_block_id,
349 range,
350 prompt_editor.read(cx).codegen.clone(),
351 workspace.clone(),
352 cx,
353 ),
354 );
355 assist_group.assist_ids.push(assist_id);
356 editor_assists.assist_ids.push(assist_id);
357 }
358 self.assist_groups.insert(assist_group_id, assist_group);
359
360 if let Some(assist_id) = assist_to_focus {
361 self.focus_assist(assist_id, cx);
362 }
363 }
364
365 #[allow(clippy::too_many_arguments)]
366 pub fn suggest_assist(
367 &mut self,
368 editor: &View<Editor>,
369 mut range: Range<Anchor>,
370 initial_prompt: String,
371 initial_transaction_id: Option<TransactionId>,
372 focus: bool,
373 workspace: Option<WeakView<Workspace>>,
374 assistant_panel: Option<&View<AssistantPanel>>,
375 cx: &mut WindowContext,
376 ) -> InlineAssistId {
377 let assist_group_id = self.next_assist_group_id.post_inc();
378 let prompt_buffer = cx.new_model(|cx| Buffer::local(&initial_prompt, cx));
379 let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
380
381 let assist_id = self.next_assist_id.post_inc();
382
383 let buffer = editor.read(cx).buffer().clone();
384 {
385 let snapshot = buffer.read(cx).read(cx);
386 range.start = range.start.bias_left(&snapshot);
387 range.end = range.end.bias_right(&snapshot);
388 }
389
390 let codegen = cx.new_model(|cx| {
391 Codegen::new(
392 editor.read(cx).buffer().clone(),
393 range.clone(),
394 initial_transaction_id,
395 self.telemetry.clone(),
396 self.prompt_builder.clone(),
397 cx,
398 )
399 });
400
401 let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
402 let prompt_editor = cx.new_view(|cx| {
403 PromptEditor::new(
404 assist_id,
405 gutter_dimensions.clone(),
406 self.prompt_history.clone(),
407 prompt_buffer.clone(),
408 codegen.clone(),
409 editor,
410 assistant_panel,
411 workspace.clone(),
412 self.fs.clone(),
413 cx,
414 )
415 });
416
417 let [prompt_block_id, end_block_id] =
418 self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
419
420 let editor_assists = self
421 .assists_by_editor
422 .entry(editor.downgrade())
423 .or_insert_with(|| EditorInlineAssists::new(&editor, cx));
424
425 let mut assist_group = InlineAssistGroup::new();
426 self.assists.insert(
427 assist_id,
428 InlineAssist::new(
429 assist_id,
430 assist_group_id,
431 assistant_panel.is_some(),
432 editor,
433 &prompt_editor,
434 prompt_block_id,
435 end_block_id,
436 range,
437 prompt_editor.read(cx).codegen.clone(),
438 workspace.clone(),
439 cx,
440 ),
441 );
442 assist_group.assist_ids.push(assist_id);
443 editor_assists.assist_ids.push(assist_id);
444 self.assist_groups.insert(assist_group_id, assist_group);
445
446 if focus {
447 self.focus_assist(assist_id, cx);
448 }
449
450 assist_id
451 }
452
453 fn insert_assist_blocks(
454 &self,
455 editor: &View<Editor>,
456 range: &Range<Anchor>,
457 prompt_editor: &View<PromptEditor>,
458 cx: &mut WindowContext,
459 ) -> [CustomBlockId; 2] {
460 let prompt_editor_height = prompt_editor.update(cx, |prompt_editor, cx| {
461 prompt_editor
462 .editor
463 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1 + 2)
464 });
465 let assist_blocks = vec![
466 BlockProperties {
467 style: BlockStyle::Sticky,
468 placement: BlockPlacement::Above(range.start),
469 height: prompt_editor_height,
470 render: build_assist_editor_renderer(prompt_editor),
471 priority: 0,
472 },
473 BlockProperties {
474 style: BlockStyle::Sticky,
475 placement: BlockPlacement::Below(range.end),
476 height: 0,
477 render: Arc::new(|cx| {
478 v_flex()
479 .h_full()
480 .w_full()
481 .border_t_1()
482 .border_color(cx.theme().status().info_border)
483 .into_any_element()
484 }),
485 priority: 0,
486 },
487 ];
488
489 editor.update(cx, |editor, cx| {
490 let block_ids = editor.insert_blocks(assist_blocks, None, cx);
491 [block_ids[0], block_ids[1]]
492 })
493 }
494
495 fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
496 let assist = &self.assists[&assist_id];
497 let Some(decorations) = assist.decorations.as_ref() else {
498 return;
499 };
500 let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
501 let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
502
503 assist_group.active_assist_id = Some(assist_id);
504 if assist_group.linked {
505 for assist_id in &assist_group.assist_ids {
506 if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
507 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
508 prompt_editor.set_show_cursor_when_unfocused(true, cx)
509 });
510 }
511 }
512 }
513
514 assist
515 .editor
516 .update(cx, |editor, cx| {
517 let scroll_top = editor.scroll_position(cx).y;
518 let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
519 let prompt_row = editor
520 .row_for_block(decorations.prompt_block_id, cx)
521 .unwrap()
522 .0 as f32;
523
524 if (scroll_top..scroll_bottom).contains(&prompt_row) {
525 editor_assists.scroll_lock = Some(InlineAssistScrollLock {
526 assist_id,
527 distance_from_top: prompt_row - scroll_top,
528 });
529 } else {
530 editor_assists.scroll_lock = None;
531 }
532 })
533 .ok();
534 }
535
536 fn handle_prompt_editor_focus_out(
537 &mut self,
538 assist_id: InlineAssistId,
539 cx: &mut WindowContext,
540 ) {
541 let assist = &self.assists[&assist_id];
542 let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
543 if assist_group.active_assist_id == Some(assist_id) {
544 assist_group.active_assist_id = None;
545 if assist_group.linked {
546 for assist_id in &assist_group.assist_ids {
547 if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
548 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
549 prompt_editor.set_show_cursor_when_unfocused(false, cx)
550 });
551 }
552 }
553 }
554 }
555 }
556
557 fn handle_prompt_editor_event(
558 &mut self,
559 prompt_editor: View<PromptEditor>,
560 event: &PromptEditorEvent,
561 cx: &mut WindowContext,
562 ) {
563 let assist_id = prompt_editor.read(cx).id;
564 match event {
565 PromptEditorEvent::StartRequested => {
566 self.start_assist(assist_id, cx);
567 }
568 PromptEditorEvent::StopRequested => {
569 self.stop_assist(assist_id, cx);
570 }
571 PromptEditorEvent::ConfirmRequested => {
572 self.finish_assist(assist_id, false, cx);
573 }
574 PromptEditorEvent::CancelRequested => {
575 self.finish_assist(assist_id, true, cx);
576 }
577 PromptEditorEvent::DismissRequested => {
578 self.dismiss_assist(assist_id, cx);
579 }
580 }
581 }
582
583 fn handle_editor_newline(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
584 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
585 return;
586 };
587
588 if editor.read(cx).selections.count() == 1 {
589 let (selection, buffer) = editor.update(cx, |editor, cx| {
590 (
591 editor.selections.newest::<usize>(cx),
592 editor.buffer().read(cx).snapshot(cx),
593 )
594 });
595 for assist_id in &editor_assists.assist_ids {
596 let assist = &self.assists[assist_id];
597 let assist_range = assist.range.to_offset(&buffer);
598 if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
599 {
600 if matches!(assist.codegen.read(cx).status(cx), CodegenStatus::Pending) {
601 self.dismiss_assist(*assist_id, cx);
602 } else {
603 self.finish_assist(*assist_id, false, cx);
604 }
605
606 return;
607 }
608 }
609 }
610
611 cx.propagate();
612 }
613
614 fn handle_editor_cancel(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
615 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
616 return;
617 };
618
619 if editor.read(cx).selections.count() == 1 {
620 let (selection, buffer) = editor.update(cx, |editor, cx| {
621 (
622 editor.selections.newest::<usize>(cx),
623 editor.buffer().read(cx).snapshot(cx),
624 )
625 });
626 let mut closest_assist_fallback = None;
627 for assist_id in &editor_assists.assist_ids {
628 let assist = &self.assists[assist_id];
629 let assist_range = assist.range.to_offset(&buffer);
630 if assist.decorations.is_some() {
631 if assist_range.contains(&selection.start)
632 && assist_range.contains(&selection.end)
633 {
634 self.focus_assist(*assist_id, cx);
635 return;
636 } else {
637 let distance_from_selection = assist_range
638 .start
639 .abs_diff(selection.start)
640 .min(assist_range.start.abs_diff(selection.end))
641 + assist_range
642 .end
643 .abs_diff(selection.start)
644 .min(assist_range.end.abs_diff(selection.end));
645 match closest_assist_fallback {
646 Some((_, old_distance)) => {
647 if distance_from_selection < old_distance {
648 closest_assist_fallback =
649 Some((assist_id, distance_from_selection));
650 }
651 }
652 None => {
653 closest_assist_fallback = Some((assist_id, distance_from_selection))
654 }
655 }
656 }
657 }
658 }
659
660 if let Some((&assist_id, _)) = closest_assist_fallback {
661 self.focus_assist(assist_id, cx);
662 }
663 }
664
665 cx.propagate();
666 }
667
668 fn handle_editor_release(&mut self, editor: WeakView<Editor>, cx: &mut WindowContext) {
669 if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor) {
670 for assist_id in editor_assists.assist_ids.clone() {
671 self.finish_assist(assist_id, true, cx);
672 }
673 }
674 }
675
676 fn handle_editor_change(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
677 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
678 return;
679 };
680 let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
681 return;
682 };
683 let assist = &self.assists[&scroll_lock.assist_id];
684 let Some(decorations) = assist.decorations.as_ref() else {
685 return;
686 };
687
688 editor.update(cx, |editor, cx| {
689 let scroll_position = editor.scroll_position(cx);
690 let target_scroll_top = editor
691 .row_for_block(decorations.prompt_block_id, cx)
692 .unwrap()
693 .0 as f32
694 - scroll_lock.distance_from_top;
695 if target_scroll_top != scroll_position.y {
696 editor.set_scroll_position(point(scroll_position.x, target_scroll_top), cx);
697 }
698 });
699 }
700
701 fn handle_editor_event(
702 &mut self,
703 editor: View<Editor>,
704 event: &EditorEvent,
705 cx: &mut WindowContext,
706 ) {
707 let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
708 return;
709 };
710
711 match event {
712 EditorEvent::Edited { transaction_id } => {
713 let buffer = editor.read(cx).buffer().read(cx);
714 let edited_ranges =
715 buffer.edited_ranges_for_transaction::<usize>(*transaction_id, cx);
716 let snapshot = buffer.snapshot(cx);
717
718 for assist_id in editor_assists.assist_ids.clone() {
719 let assist = &self.assists[&assist_id];
720 if matches!(
721 assist.codegen.read(cx).status(cx),
722 CodegenStatus::Error(_) | CodegenStatus::Done
723 ) {
724 let assist_range = assist.range.to_offset(&snapshot);
725 if edited_ranges
726 .iter()
727 .any(|range| range.overlaps(&assist_range))
728 {
729 self.finish_assist(assist_id, false, cx);
730 }
731 }
732 }
733 }
734 EditorEvent::ScrollPositionChanged { .. } => {
735 if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
736 let assist = &self.assists[&scroll_lock.assist_id];
737 if let Some(decorations) = assist.decorations.as_ref() {
738 let distance_from_top = editor.update(cx, |editor, cx| {
739 let scroll_top = editor.scroll_position(cx).y;
740 let prompt_row = editor
741 .row_for_block(decorations.prompt_block_id, cx)
742 .unwrap()
743 .0 as f32;
744 prompt_row - scroll_top
745 });
746
747 if distance_from_top != scroll_lock.distance_from_top {
748 editor_assists.scroll_lock = None;
749 }
750 }
751 }
752 }
753 EditorEvent::SelectionsChanged { .. } => {
754 for assist_id in editor_assists.assist_ids.clone() {
755 let assist = &self.assists[&assist_id];
756 if let Some(decorations) = assist.decorations.as_ref() {
757 if decorations.prompt_editor.focus_handle(cx).is_focused(cx) {
758 return;
759 }
760 }
761 }
762
763 editor_assists.scroll_lock = None;
764 }
765 _ => {}
766 }
767 }
768
769 pub fn finish_assist(&mut self, assist_id: InlineAssistId, undo: bool, cx: &mut WindowContext) {
770 if let Some(assist) = self.assists.get(&assist_id) {
771 let assist_group_id = assist.group_id;
772 if self.assist_groups[&assist_group_id].linked {
773 for assist_id in self.unlink_assist_group(assist_group_id, cx) {
774 self.finish_assist(assist_id, undo, cx);
775 }
776 return;
777 }
778 }
779
780 self.dismiss_assist(assist_id, cx);
781
782 if let Some(assist) = self.assists.remove(&assist_id) {
783 if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
784 {
785 entry.get_mut().assist_ids.retain(|id| *id != assist_id);
786 if entry.get().assist_ids.is_empty() {
787 entry.remove();
788 }
789 }
790
791 if let hash_map::Entry::Occupied(mut entry) =
792 self.assists_by_editor.entry(assist.editor.clone())
793 {
794 entry.get_mut().assist_ids.retain(|id| *id != assist_id);
795 if entry.get().assist_ids.is_empty() {
796 entry.remove();
797 if let Some(editor) = assist.editor.upgrade() {
798 self.update_editor_highlights(&editor, cx);
799 }
800 } else {
801 entry.get().highlight_updates.send(()).ok();
802 }
803 }
804
805 let active_alternative = assist.codegen.read(cx).active_alternative().clone();
806 let message_id = active_alternative.read(cx).message_id.clone();
807
808 if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
809 let language_name = assist.editor.upgrade().and_then(|editor| {
810 let multibuffer = editor.read(cx).buffer().read(cx);
811 let multibuffer_snapshot = multibuffer.snapshot(cx);
812 let ranges = multibuffer_snapshot.range_to_buffer_ranges(assist.range.clone());
813 ranges
814 .first()
815 .and_then(|(excerpt, _)| excerpt.buffer().language())
816 .map(|language| language.name())
817 });
818 report_assistant_event(
819 AssistantEvent {
820 conversation_id: None,
821 kind: AssistantKind::Inline,
822 message_id,
823 phase: if undo {
824 AssistantPhase::Rejected
825 } else {
826 AssistantPhase::Accepted
827 },
828 model: model.telemetry_id(),
829 model_provider: model.provider_id().to_string(),
830 response_latency: None,
831 error_message: None,
832 language_name: language_name.map(|name| name.to_proto()),
833 },
834 Some(self.telemetry.clone()),
835 cx.http_client(),
836 model.api_key(cx),
837 cx.background_executor(),
838 );
839 }
840
841 if undo {
842 assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
843 } else {
844 self.confirmed_assists.insert(assist_id, active_alternative);
845 }
846 }
847 }
848
849 fn dismiss_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) -> bool {
850 let Some(assist) = self.assists.get_mut(&assist_id) else {
851 return false;
852 };
853 let Some(editor) = assist.editor.upgrade() else {
854 return false;
855 };
856 let Some(decorations) = assist.decorations.take() else {
857 return false;
858 };
859
860 editor.update(cx, |editor, cx| {
861 let mut to_remove = decorations.removed_line_block_ids;
862 to_remove.insert(decorations.prompt_block_id);
863 to_remove.insert(decorations.end_block_id);
864 editor.remove_blocks(to_remove, None, cx);
865 });
866
867 if decorations
868 .prompt_editor
869 .focus_handle(cx)
870 .contains_focused(cx)
871 {
872 self.focus_next_assist(assist_id, cx);
873 }
874
875 if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
876 if editor_assists
877 .scroll_lock
878 .as_ref()
879 .map_or(false, |lock| lock.assist_id == assist_id)
880 {
881 editor_assists.scroll_lock = None;
882 }
883 editor_assists.highlight_updates.send(()).ok();
884 }
885
886 true
887 }
888
889 fn focus_next_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
890 let Some(assist) = self.assists.get(&assist_id) else {
891 return;
892 };
893
894 let assist_group = &self.assist_groups[&assist.group_id];
895 let assist_ix = assist_group
896 .assist_ids
897 .iter()
898 .position(|id| *id == assist_id)
899 .unwrap();
900 let assist_ids = assist_group
901 .assist_ids
902 .iter()
903 .skip(assist_ix + 1)
904 .chain(assist_group.assist_ids.iter().take(assist_ix));
905
906 for assist_id in assist_ids {
907 let assist = &self.assists[assist_id];
908 if assist.decorations.is_some() {
909 self.focus_assist(*assist_id, cx);
910 return;
911 }
912 }
913
914 assist.editor.update(cx, |editor, cx| editor.focus(cx)).ok();
915 }
916
917 fn focus_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
918 let Some(assist) = self.assists.get(&assist_id) else {
919 return;
920 };
921
922 if let Some(decorations) = assist.decorations.as_ref() {
923 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
924 prompt_editor.editor.update(cx, |editor, cx| {
925 editor.focus(cx);
926 editor.select_all(&SelectAll, cx);
927 })
928 });
929 }
930
931 self.scroll_to_assist(assist_id, cx);
932 }
933
934 pub fn scroll_to_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
935 let Some(assist) = self.assists.get(&assist_id) else {
936 return;
937 };
938 let Some(editor) = assist.editor.upgrade() else {
939 return;
940 };
941
942 let position = assist.range.start;
943 editor.update(cx, |editor, cx| {
944 editor.change_selections(None, cx, |selections| {
945 selections.select_anchor_ranges([position..position])
946 });
947
948 let mut scroll_target_top;
949 let mut scroll_target_bottom;
950 if let Some(decorations) = assist.decorations.as_ref() {
951 scroll_target_top = editor
952 .row_for_block(decorations.prompt_block_id, cx)
953 .unwrap()
954 .0 as f32;
955 scroll_target_bottom = editor
956 .row_for_block(decorations.end_block_id, cx)
957 .unwrap()
958 .0 as f32;
959 } else {
960 let snapshot = editor.snapshot(cx);
961 let start_row = assist
962 .range
963 .start
964 .to_display_point(&snapshot.display_snapshot)
965 .row();
966 scroll_target_top = start_row.0 as f32;
967 scroll_target_bottom = scroll_target_top + 1.;
968 }
969 scroll_target_top -= editor.vertical_scroll_margin() as f32;
970 scroll_target_bottom += editor.vertical_scroll_margin() as f32;
971
972 let height_in_lines = editor.visible_line_count().unwrap_or(0.);
973 let scroll_top = editor.scroll_position(cx).y;
974 let scroll_bottom = scroll_top + height_in_lines;
975
976 if scroll_target_top < scroll_top {
977 editor.set_scroll_position(point(0., scroll_target_top), cx);
978 } else if scroll_target_bottom > scroll_bottom {
979 if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
980 editor
981 .set_scroll_position(point(0., scroll_target_bottom - height_in_lines), cx);
982 } else {
983 editor.set_scroll_position(point(0., scroll_target_top), cx);
984 }
985 }
986 });
987 }
988
989 fn unlink_assist_group(
990 &mut self,
991 assist_group_id: InlineAssistGroupId,
992 cx: &mut WindowContext,
993 ) -> Vec<InlineAssistId> {
994 let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
995 assist_group.linked = false;
996 for assist_id in &assist_group.assist_ids {
997 let assist = self.assists.get_mut(assist_id).unwrap();
998 if let Some(editor_decorations) = assist.decorations.as_ref() {
999 editor_decorations
1000 .prompt_editor
1001 .update(cx, |prompt_editor, cx| prompt_editor.unlink(cx));
1002 }
1003 }
1004 assist_group.assist_ids.clone()
1005 }
1006
1007 pub fn start_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
1008 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1009 assist
1010 } else {
1011 return;
1012 };
1013
1014 let assist_group_id = assist.group_id;
1015 if self.assist_groups[&assist_group_id].linked {
1016 for assist_id in self.unlink_assist_group(assist_group_id, cx) {
1017 self.start_assist(assist_id, cx);
1018 }
1019 return;
1020 }
1021
1022 let Some(user_prompt) = assist.user_prompt(cx) else {
1023 return;
1024 };
1025
1026 self.prompt_history.retain(|prompt| *prompt != user_prompt);
1027 self.prompt_history.push_back(user_prompt.clone());
1028 if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
1029 self.prompt_history.pop_front();
1030 }
1031
1032 let assistant_panel_context = assist.assistant_panel_context(cx);
1033
1034 assist
1035 .codegen
1036 .update(cx, |codegen, cx| {
1037 codegen.start(user_prompt, assistant_panel_context, cx)
1038 })
1039 .log_err();
1040 }
1041
1042 pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
1043 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1044 assist
1045 } else {
1046 return;
1047 };
1048
1049 assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
1050 }
1051
1052 fn update_editor_highlights(&self, editor: &View<Editor>, cx: &mut WindowContext) {
1053 let mut gutter_pending_ranges = Vec::new();
1054 let mut gutter_transformed_ranges = Vec::new();
1055 let mut foreground_ranges = Vec::new();
1056 let mut inserted_row_ranges = Vec::new();
1057 let empty_assist_ids = Vec::new();
1058 let assist_ids = self
1059 .assists_by_editor
1060 .get(&editor.downgrade())
1061 .map_or(&empty_assist_ids, |editor_assists| {
1062 &editor_assists.assist_ids
1063 });
1064
1065 for assist_id in assist_ids {
1066 if let Some(assist) = self.assists.get(assist_id) {
1067 let codegen = assist.codegen.read(cx);
1068 let buffer = codegen.buffer(cx).read(cx).read(cx);
1069 foreground_ranges.extend(codegen.last_equal_ranges(cx).iter().cloned());
1070
1071 let pending_range =
1072 codegen.edit_position(cx).unwrap_or(assist.range.start)..assist.range.end;
1073 if pending_range.end.to_offset(&buffer) > pending_range.start.to_offset(&buffer) {
1074 gutter_pending_ranges.push(pending_range);
1075 }
1076
1077 if let Some(edit_position) = codegen.edit_position(cx) {
1078 let edited_range = assist.range.start..edit_position;
1079 if edited_range.end.to_offset(&buffer) > edited_range.start.to_offset(&buffer) {
1080 gutter_transformed_ranges.push(edited_range);
1081 }
1082 }
1083
1084 if assist.decorations.is_some() {
1085 inserted_row_ranges
1086 .extend(codegen.diff(cx).inserted_row_ranges.iter().cloned());
1087 }
1088 }
1089 }
1090
1091 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
1092 merge_ranges(&mut foreground_ranges, &snapshot);
1093 merge_ranges(&mut gutter_pending_ranges, &snapshot);
1094 merge_ranges(&mut gutter_transformed_ranges, &snapshot);
1095 editor.update(cx, |editor, cx| {
1096 enum GutterPendingRange {}
1097 if gutter_pending_ranges.is_empty() {
1098 editor.clear_gutter_highlights::<GutterPendingRange>(cx);
1099 } else {
1100 editor.highlight_gutter::<GutterPendingRange>(
1101 &gutter_pending_ranges,
1102 |cx| cx.theme().status().info_background,
1103 cx,
1104 )
1105 }
1106
1107 enum GutterTransformedRange {}
1108 if gutter_transformed_ranges.is_empty() {
1109 editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
1110 } else {
1111 editor.highlight_gutter::<GutterTransformedRange>(
1112 &gutter_transformed_ranges,
1113 |cx| cx.theme().status().info,
1114 cx,
1115 )
1116 }
1117
1118 if foreground_ranges.is_empty() {
1119 editor.clear_highlights::<InlineAssist>(cx);
1120 } else {
1121 editor.highlight_text::<InlineAssist>(
1122 foreground_ranges,
1123 HighlightStyle {
1124 fade_out: Some(0.6),
1125 ..Default::default()
1126 },
1127 cx,
1128 );
1129 }
1130
1131 editor.clear_row_highlights::<InlineAssist>();
1132 for row_range in inserted_row_ranges {
1133 editor.highlight_rows::<InlineAssist>(
1134 row_range,
1135 cx.theme().status().info_background,
1136 false,
1137 cx,
1138 );
1139 }
1140 });
1141 }
1142
1143 fn update_editor_blocks(
1144 &mut self,
1145 editor: &View<Editor>,
1146 assist_id: InlineAssistId,
1147 cx: &mut WindowContext,
1148 ) {
1149 let Some(assist) = self.assists.get_mut(&assist_id) else {
1150 return;
1151 };
1152 let Some(decorations) = assist.decorations.as_mut() else {
1153 return;
1154 };
1155
1156 let codegen = assist.codegen.read(cx);
1157 let old_snapshot = codegen.snapshot(cx);
1158 let old_buffer = codegen.old_buffer(cx);
1159 let deleted_row_ranges = codegen.diff(cx).deleted_row_ranges.clone();
1160
1161 editor.update(cx, |editor, cx| {
1162 let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
1163 editor.remove_blocks(old_blocks, None, cx);
1164
1165 let mut new_blocks = Vec::new();
1166 for (new_row, old_row_range) in deleted_row_ranges {
1167 let (_, buffer_start) = old_snapshot
1168 .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
1169 .unwrap();
1170 let (_, buffer_end) = old_snapshot
1171 .point_to_buffer_offset(Point::new(
1172 *old_row_range.end(),
1173 old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
1174 ))
1175 .unwrap();
1176
1177 let deleted_lines_editor = cx.new_view(|cx| {
1178 let multi_buffer = cx.new_model(|_| {
1179 MultiBuffer::without_headers(language::Capability::ReadOnly)
1180 });
1181 multi_buffer.update(cx, |multi_buffer, cx| {
1182 multi_buffer.push_excerpts(
1183 old_buffer.clone(),
1184 Some(ExcerptRange {
1185 context: buffer_start..buffer_end,
1186 primary: None,
1187 }),
1188 cx,
1189 );
1190 });
1191
1192 enum DeletedLines {}
1193 let mut editor = Editor::for_multibuffer(multi_buffer, None, true, cx);
1194 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1195 editor.set_show_wrap_guides(false, cx);
1196 editor.set_show_gutter(false, cx);
1197 editor.scroll_manager.set_forbid_vertical_scroll(true);
1198 editor.set_read_only(true);
1199 editor.set_show_inline_completions(Some(false), cx);
1200 editor.highlight_rows::<DeletedLines>(
1201 Anchor::min()..Anchor::max(),
1202 cx.theme().status().deleted_background,
1203 false,
1204 cx,
1205 );
1206 editor
1207 });
1208
1209 let height =
1210 deleted_lines_editor.update(cx, |editor, cx| editor.max_point(cx).row().0 + 1);
1211 new_blocks.push(BlockProperties {
1212 placement: BlockPlacement::Above(new_row),
1213 height,
1214 style: BlockStyle::Flex,
1215 render: Arc::new(move |cx| {
1216 div()
1217 .block_mouse_down()
1218 .bg(cx.theme().status().deleted_background)
1219 .size_full()
1220 .h(height as f32 * cx.line_height())
1221 .pl(cx.gutter_dimensions.full_width())
1222 .child(deleted_lines_editor.clone())
1223 .into_any_element()
1224 }),
1225 priority: 0,
1226 });
1227 }
1228
1229 decorations.removed_line_block_ids = editor
1230 .insert_blocks(new_blocks, None, cx)
1231 .into_iter()
1232 .collect();
1233 })
1234 }
1235}
1236
1237struct EditorInlineAssists {
1238 assist_ids: Vec<InlineAssistId>,
1239 scroll_lock: Option<InlineAssistScrollLock>,
1240 highlight_updates: async_watch::Sender<()>,
1241 _update_highlights: Task<Result<()>>,
1242 _subscriptions: Vec<gpui::Subscription>,
1243}
1244
1245struct InlineAssistScrollLock {
1246 assist_id: InlineAssistId,
1247 distance_from_top: f32,
1248}
1249
1250impl EditorInlineAssists {
1251 #[allow(clippy::too_many_arguments)]
1252 fn new(editor: &View<Editor>, cx: &mut WindowContext) -> Self {
1253 let (highlight_updates_tx, mut highlight_updates_rx) = async_watch::channel(());
1254 Self {
1255 assist_ids: Vec::new(),
1256 scroll_lock: None,
1257 highlight_updates: highlight_updates_tx,
1258 _update_highlights: cx.spawn(|mut cx| {
1259 let editor = editor.downgrade();
1260 async move {
1261 while let Ok(()) = highlight_updates_rx.changed().await {
1262 let editor = editor.upgrade().context("editor was dropped")?;
1263 cx.update_global(|assistant: &mut InlineAssistant, cx| {
1264 assistant.update_editor_highlights(&editor, cx);
1265 })?;
1266 }
1267 Ok(())
1268 }
1269 }),
1270 _subscriptions: vec![
1271 cx.observe_release(editor, {
1272 let editor = editor.downgrade();
1273 |_, cx| {
1274 InlineAssistant::update_global(cx, |this, cx| {
1275 this.handle_editor_release(editor, cx);
1276 })
1277 }
1278 }),
1279 cx.observe(editor, move |editor, cx| {
1280 InlineAssistant::update_global(cx, |this, cx| {
1281 this.handle_editor_change(editor, cx)
1282 })
1283 }),
1284 cx.subscribe(editor, move |editor, event, cx| {
1285 InlineAssistant::update_global(cx, |this, cx| {
1286 this.handle_editor_event(editor, event, cx)
1287 })
1288 }),
1289 editor.update(cx, |editor, cx| {
1290 let editor_handle = cx.view().downgrade();
1291 editor.register_action(
1292 move |_: &editor::actions::Newline, cx: &mut WindowContext| {
1293 InlineAssistant::update_global(cx, |this, cx| {
1294 if let Some(editor) = editor_handle.upgrade() {
1295 this.handle_editor_newline(editor, cx)
1296 }
1297 })
1298 },
1299 )
1300 }),
1301 editor.update(cx, |editor, cx| {
1302 let editor_handle = cx.view().downgrade();
1303 editor.register_action(
1304 move |_: &editor::actions::Cancel, cx: &mut WindowContext| {
1305 InlineAssistant::update_global(cx, |this, cx| {
1306 if let Some(editor) = editor_handle.upgrade() {
1307 this.handle_editor_cancel(editor, cx)
1308 }
1309 })
1310 },
1311 )
1312 }),
1313 ],
1314 }
1315 }
1316}
1317
1318struct InlineAssistGroup {
1319 assist_ids: Vec<InlineAssistId>,
1320 linked: bool,
1321 active_assist_id: Option<InlineAssistId>,
1322}
1323
1324impl InlineAssistGroup {
1325 fn new() -> Self {
1326 Self {
1327 assist_ids: Vec::new(),
1328 linked: true,
1329 active_assist_id: None,
1330 }
1331 }
1332}
1333
1334fn build_assist_editor_renderer(editor: &View<PromptEditor>) -> RenderBlock {
1335 let editor = editor.clone();
1336 Arc::new(move |cx: &mut BlockContext| {
1337 *editor.read(cx).gutter_dimensions.lock() = *cx.gutter_dimensions;
1338 editor.clone().into_any_element()
1339 })
1340}
1341
1342#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1343pub struct InlineAssistId(usize);
1344
1345impl InlineAssistId {
1346 fn post_inc(&mut self) -> InlineAssistId {
1347 let id = *self;
1348 self.0 += 1;
1349 id
1350 }
1351}
1352
1353#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1354struct InlineAssistGroupId(usize);
1355
1356impl InlineAssistGroupId {
1357 fn post_inc(&mut self) -> InlineAssistGroupId {
1358 let id = *self;
1359 self.0 += 1;
1360 id
1361 }
1362}
1363
1364enum PromptEditorEvent {
1365 StartRequested,
1366 StopRequested,
1367 ConfirmRequested,
1368 CancelRequested,
1369 DismissRequested,
1370}
1371
1372struct PromptEditor {
1373 id: InlineAssistId,
1374 editor: View<Editor>,
1375 language_model_selector: View<LanguageModelSelector>,
1376 edited_since_done: bool,
1377 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1378 prompt_history: VecDeque<String>,
1379 prompt_history_ix: Option<usize>,
1380 pending_prompt: String,
1381 codegen: Model<Codegen>,
1382 _codegen_subscription: Subscription,
1383 editor_subscriptions: Vec<Subscription>,
1384 pending_token_count: Task<Result<()>>,
1385 token_counts: Option<TokenCounts>,
1386 _token_count_subscriptions: Vec<Subscription>,
1387 workspace: Option<WeakView<Workspace>>,
1388 show_rate_limit_notice: bool,
1389}
1390
1391#[derive(Copy, Clone)]
1392pub struct TokenCounts {
1393 total: usize,
1394 assistant_panel: usize,
1395}
1396
1397impl EventEmitter<PromptEditorEvent> for PromptEditor {}
1398
1399impl Render for PromptEditor {
1400 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1401 let gutter_dimensions = *self.gutter_dimensions.lock();
1402 let codegen = self.codegen.read(cx);
1403
1404 let mut buttons = Vec::new();
1405 if codegen.alternative_count(cx) > 1 {
1406 buttons.push(self.render_cycle_controls(cx));
1407 }
1408
1409 let status = codegen.status(cx);
1410 buttons.extend(match status {
1411 CodegenStatus::Idle => {
1412 vec![
1413 IconButton::new("cancel", IconName::Close)
1414 .icon_color(Color::Muted)
1415 .shape(IconButtonShape::Square)
1416 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1417 .on_click(
1418 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1419 )
1420 .into_any_element(),
1421 IconButton::new("start", IconName::SparkleAlt)
1422 .icon_color(Color::Muted)
1423 .shape(IconButtonShape::Square)
1424 .tooltip(|cx| Tooltip::for_action("Transform", &menu::Confirm, cx))
1425 .on_click(
1426 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
1427 )
1428 .into_any_element(),
1429 ]
1430 }
1431 CodegenStatus::Pending => {
1432 vec![
1433 IconButton::new("cancel", IconName::Close)
1434 .icon_color(Color::Muted)
1435 .shape(IconButtonShape::Square)
1436 .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
1437 .on_click(
1438 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1439 )
1440 .into_any_element(),
1441 IconButton::new("stop", IconName::Stop)
1442 .icon_color(Color::Error)
1443 .shape(IconButtonShape::Square)
1444 .tooltip(|cx| {
1445 Tooltip::with_meta(
1446 "Interrupt Transformation",
1447 Some(&menu::Cancel),
1448 "Changes won't be discarded",
1449 cx,
1450 )
1451 })
1452 .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
1453 .into_any_element(),
1454 ]
1455 }
1456 CodegenStatus::Error(_) | CodegenStatus::Done => {
1457 let must_rerun =
1458 self.edited_since_done || matches!(status, CodegenStatus::Error(_));
1459 // when accept button isn't visible, then restart maps to confirm
1460 // when accept button is visible, then restart must be mapped to an alternate keyboard shortcut
1461 let restart_key: &dyn gpui::Action = if must_rerun {
1462 &menu::Confirm
1463 } else {
1464 &menu::Restart
1465 };
1466 vec![
1467 IconButton::new("cancel", IconName::Close)
1468 .icon_color(Color::Muted)
1469 .shape(IconButtonShape::Square)
1470 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1471 .on_click(
1472 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1473 )
1474 .into_any_element(),
1475 IconButton::new("restart", IconName::RotateCw)
1476 .icon_color(Color::Muted)
1477 .shape(IconButtonShape::Square)
1478 .tooltip(|cx| {
1479 Tooltip::with_meta(
1480 "Regenerate Transformation",
1481 Some(restart_key),
1482 "Current change will be discarded",
1483 cx,
1484 )
1485 })
1486 .on_click(cx.listener(|_, _, cx| {
1487 cx.emit(PromptEditorEvent::StartRequested);
1488 }))
1489 .into_any_element(),
1490 if !must_rerun {
1491 IconButton::new("confirm", IconName::Check)
1492 .icon_color(Color::Info)
1493 .shape(IconButtonShape::Square)
1494 .tooltip(|cx| Tooltip::for_action("Confirm Assist", &menu::Confirm, cx))
1495 .on_click(cx.listener(|_, _, cx| {
1496 cx.emit(PromptEditorEvent::ConfirmRequested);
1497 }))
1498 .into_any_element()
1499 } else {
1500 div().into_any_element()
1501 },
1502 ]
1503 }
1504 });
1505
1506 h_flex()
1507 .key_context("PromptEditor")
1508 .bg(cx.theme().colors().editor_background)
1509 .block_mouse_down()
1510 .cursor(CursorStyle::Arrow)
1511 .border_y_1()
1512 .border_color(cx.theme().status().info_border)
1513 .size_full()
1514 .py(cx.line_height() / 2.5)
1515 .on_action(cx.listener(Self::confirm))
1516 .on_action(cx.listener(Self::cancel))
1517 .on_action(cx.listener(Self::restart))
1518 .on_action(cx.listener(Self::move_up))
1519 .on_action(cx.listener(Self::move_down))
1520 .capture_action(cx.listener(Self::cycle_prev))
1521 .capture_action(cx.listener(Self::cycle_next))
1522 .child(
1523 h_flex()
1524 .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
1525 .justify_center()
1526 .gap_2()
1527 .child(LanguageModelSelectorPopoverMenu::new(
1528 self.language_model_selector.clone(),
1529 IconButton::new("context", IconName::SettingsAlt)
1530 .shape(IconButtonShape::Square)
1531 .icon_size(IconSize::Small)
1532 .icon_color(Color::Muted)
1533 .tooltip(move |cx| {
1534 Tooltip::with_meta(
1535 format!(
1536 "Using {}",
1537 LanguageModelRegistry::read_global(cx)
1538 .active_model()
1539 .map(|model| model.name().0)
1540 .unwrap_or_else(|| "No model selected".into()),
1541 ),
1542 None,
1543 "Change Model",
1544 cx,
1545 )
1546 }),
1547 ))
1548 .map(|el| {
1549 let CodegenStatus::Error(error) = self.codegen.read(cx).status(cx) else {
1550 return el;
1551 };
1552
1553 let error_message = SharedString::from(error.to_string());
1554 if error.error_code() == proto::ErrorCode::RateLimitExceeded
1555 && cx.has_flag::<ZedPro>()
1556 {
1557 el.child(
1558 v_flex()
1559 .child(
1560 IconButton::new("rate-limit-error", IconName::XCircle)
1561 .toggle_state(self.show_rate_limit_notice)
1562 .shape(IconButtonShape::Square)
1563 .icon_size(IconSize::Small)
1564 .on_click(cx.listener(Self::toggle_rate_limit_notice)),
1565 )
1566 .children(self.show_rate_limit_notice.then(|| {
1567 deferred(
1568 anchored()
1569 .position_mode(gpui::AnchoredPositionMode::Local)
1570 .position(point(px(0.), px(24.)))
1571 .anchor(gpui::Corner::TopLeft)
1572 .child(self.render_rate_limit_notice(cx)),
1573 )
1574 })),
1575 )
1576 } else {
1577 el.child(
1578 div()
1579 .id("error")
1580 .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
1581 .child(
1582 Icon::new(IconName::XCircle)
1583 .size(IconSize::Small)
1584 .color(Color::Error),
1585 ),
1586 )
1587 }
1588 }),
1589 )
1590 .child(div().flex_1().child(self.render_prompt_editor(cx)))
1591 .child(
1592 h_flex()
1593 .gap_2()
1594 .pr_6()
1595 .children(self.render_token_count(cx))
1596 .children(buttons),
1597 )
1598 }
1599}
1600
1601impl FocusableView for PromptEditor {
1602 fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
1603 self.editor.focus_handle(cx)
1604 }
1605}
1606
1607impl PromptEditor {
1608 const MAX_LINES: u8 = 8;
1609
1610 #[allow(clippy::too_many_arguments)]
1611 fn new(
1612 id: InlineAssistId,
1613 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1614 prompt_history: VecDeque<String>,
1615 prompt_buffer: Model<MultiBuffer>,
1616 codegen: Model<Codegen>,
1617 parent_editor: &View<Editor>,
1618 assistant_panel: Option<&View<AssistantPanel>>,
1619 workspace: Option<WeakView<Workspace>>,
1620 fs: Arc<dyn Fs>,
1621 cx: &mut ViewContext<Self>,
1622 ) -> Self {
1623 let prompt_editor = cx.new_view(|cx| {
1624 let mut editor = Editor::new(
1625 EditorMode::AutoHeight {
1626 max_lines: Self::MAX_LINES as usize,
1627 },
1628 prompt_buffer,
1629 None,
1630 false,
1631 cx,
1632 );
1633 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1634 // Since the prompt editors for all inline assistants are linked,
1635 // always show the cursor (even when it isn't focused) because
1636 // typing in one will make what you typed appear in all of them.
1637 editor.set_show_cursor_when_unfocused(true, cx);
1638 editor.set_placeholder_text(Self::placeholder_text(codegen.read(cx), cx), cx);
1639 editor
1640 });
1641
1642 let mut token_count_subscriptions = Vec::new();
1643 token_count_subscriptions
1644 .push(cx.subscribe(parent_editor, Self::handle_parent_editor_event));
1645 if let Some(assistant_panel) = assistant_panel {
1646 token_count_subscriptions
1647 .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
1648 }
1649
1650 let mut this = Self {
1651 id,
1652 editor: prompt_editor,
1653 language_model_selector: cx.new_view(|cx| {
1654 let fs = fs.clone();
1655 LanguageModelSelector::new(
1656 move |model, cx| {
1657 update_settings_file::<AssistantSettings>(
1658 fs.clone(),
1659 cx,
1660 move |settings, _| settings.set_model(model.clone()),
1661 );
1662 },
1663 cx,
1664 )
1665 }),
1666 edited_since_done: false,
1667 gutter_dimensions,
1668 prompt_history,
1669 prompt_history_ix: None,
1670 pending_prompt: String::new(),
1671 _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
1672 editor_subscriptions: Vec::new(),
1673 codegen,
1674 pending_token_count: Task::ready(Ok(())),
1675 token_counts: None,
1676 _token_count_subscriptions: token_count_subscriptions,
1677 workspace,
1678 show_rate_limit_notice: false,
1679 };
1680 this.count_tokens(cx);
1681 this.subscribe_to_editor(cx);
1682 this
1683 }
1684
1685 fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
1686 self.editor_subscriptions.clear();
1687 self.editor_subscriptions
1688 .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
1689 }
1690
1691 fn set_show_cursor_when_unfocused(
1692 &mut self,
1693 show_cursor_when_unfocused: bool,
1694 cx: &mut ViewContext<Self>,
1695 ) {
1696 self.editor.update(cx, |editor, cx| {
1697 editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
1698 });
1699 }
1700
1701 fn unlink(&mut self, cx: &mut ViewContext<Self>) {
1702 let prompt = self.prompt(cx);
1703 let focus = self.editor.focus_handle(cx).contains_focused(cx);
1704 self.editor = cx.new_view(|cx| {
1705 let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
1706 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1707 editor.set_placeholder_text(Self::placeholder_text(self.codegen.read(cx), cx), cx);
1708 editor.set_placeholder_text("Add a prompt…", cx);
1709 editor.set_text(prompt, cx);
1710 if focus {
1711 editor.focus(cx);
1712 }
1713 editor
1714 });
1715 self.subscribe_to_editor(cx);
1716 }
1717
1718 fn placeholder_text(codegen: &Codegen, cx: &WindowContext) -> String {
1719 let context_keybinding = text_for_action(&crate::ToggleFocus, cx)
1720 .map(|keybinding| format!(" • {keybinding} for context"))
1721 .unwrap_or_default();
1722
1723 let action = if codegen.is_insertion {
1724 "Generate"
1725 } else {
1726 "Transform"
1727 };
1728
1729 format!("{action}…{context_keybinding} • ↓↑ for history")
1730 }
1731
1732 fn prompt(&self, cx: &AppContext) -> String {
1733 self.editor.read(cx).text(cx)
1734 }
1735
1736 fn toggle_rate_limit_notice(&mut self, _: &ClickEvent, cx: &mut ViewContext<Self>) {
1737 self.show_rate_limit_notice = !self.show_rate_limit_notice;
1738 if self.show_rate_limit_notice {
1739 cx.focus_view(&self.editor);
1740 }
1741 cx.notify();
1742 }
1743
1744 fn handle_parent_editor_event(
1745 &mut self,
1746 _: View<Editor>,
1747 event: &EditorEvent,
1748 cx: &mut ViewContext<Self>,
1749 ) {
1750 if let EditorEvent::BufferEdited { .. } = event {
1751 self.count_tokens(cx);
1752 }
1753 }
1754
1755 fn handle_assistant_panel_event(
1756 &mut self,
1757 _: View<AssistantPanel>,
1758 event: &AssistantPanelEvent,
1759 cx: &mut ViewContext<Self>,
1760 ) {
1761 let AssistantPanelEvent::ContextEdited { .. } = event;
1762 self.count_tokens(cx);
1763 }
1764
1765 fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
1766 let assist_id = self.id;
1767 self.pending_token_count = cx.spawn(|this, mut cx| async move {
1768 cx.background_executor().timer(Duration::from_secs(1)).await;
1769 let token_count = cx
1770 .update_global(|inline_assistant: &mut InlineAssistant, cx| {
1771 let assist = inline_assistant
1772 .assists
1773 .get(&assist_id)
1774 .context("assist not found")?;
1775 anyhow::Ok(assist.count_tokens(cx))
1776 })??
1777 .await?;
1778
1779 this.update(&mut cx, |this, cx| {
1780 this.token_counts = Some(token_count);
1781 cx.notify();
1782 })
1783 })
1784 }
1785
1786 fn handle_prompt_editor_events(
1787 &mut self,
1788 _: View<Editor>,
1789 event: &EditorEvent,
1790 cx: &mut ViewContext<Self>,
1791 ) {
1792 match event {
1793 EditorEvent::Edited { .. } => {
1794 if let Some(workspace) = cx.window_handle().downcast::<Workspace>() {
1795 workspace
1796 .update(cx, |workspace, cx| {
1797 let is_via_ssh = workspace
1798 .project()
1799 .update(cx, |project, _| project.is_via_ssh());
1800
1801 workspace
1802 .client()
1803 .telemetry()
1804 .log_edit_event("inline assist", is_via_ssh);
1805 })
1806 .log_err();
1807 }
1808 let prompt = self.editor.read(cx).text(cx);
1809 if self
1810 .prompt_history_ix
1811 .map_or(true, |ix| self.prompt_history[ix] != prompt)
1812 {
1813 self.prompt_history_ix.take();
1814 self.pending_prompt = prompt;
1815 }
1816
1817 self.edited_since_done = true;
1818 cx.notify();
1819 }
1820 EditorEvent::BufferEdited => {
1821 self.count_tokens(cx);
1822 }
1823 EditorEvent::Blurred => {
1824 if self.show_rate_limit_notice {
1825 self.show_rate_limit_notice = false;
1826 cx.notify();
1827 }
1828 }
1829 _ => {}
1830 }
1831 }
1832
1833 fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
1834 match self.codegen.read(cx).status(cx) {
1835 CodegenStatus::Idle => {
1836 self.editor
1837 .update(cx, |editor, _| editor.set_read_only(false));
1838 }
1839 CodegenStatus::Pending => {
1840 self.editor
1841 .update(cx, |editor, _| editor.set_read_only(true));
1842 }
1843 CodegenStatus::Done => {
1844 self.edited_since_done = false;
1845 self.editor
1846 .update(cx, |editor, _| editor.set_read_only(false));
1847 }
1848 CodegenStatus::Error(error) => {
1849 if cx.has_flag::<ZedPro>()
1850 && error.error_code() == proto::ErrorCode::RateLimitExceeded
1851 && !dismissed_rate_limit_notice()
1852 {
1853 self.show_rate_limit_notice = true;
1854 cx.notify();
1855 }
1856
1857 self.edited_since_done = false;
1858 self.editor
1859 .update(cx, |editor, _| editor.set_read_only(false));
1860 }
1861 }
1862 }
1863
1864 fn restart(&mut self, _: &menu::Restart, cx: &mut ViewContext<Self>) {
1865 cx.emit(PromptEditorEvent::StartRequested);
1866 }
1867
1868 fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
1869 match self.codegen.read(cx).status(cx) {
1870 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1871 cx.emit(PromptEditorEvent::CancelRequested);
1872 }
1873 CodegenStatus::Pending => {
1874 cx.emit(PromptEditorEvent::StopRequested);
1875 }
1876 }
1877 }
1878
1879 fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
1880 match self.codegen.read(cx).status(cx) {
1881 CodegenStatus::Idle => {
1882 cx.emit(PromptEditorEvent::StartRequested);
1883 }
1884 CodegenStatus::Pending => {
1885 cx.emit(PromptEditorEvent::DismissRequested);
1886 }
1887 CodegenStatus::Done => {
1888 if self.edited_since_done {
1889 cx.emit(PromptEditorEvent::StartRequested);
1890 } else {
1891 cx.emit(PromptEditorEvent::ConfirmRequested);
1892 }
1893 }
1894 CodegenStatus::Error(_) => {
1895 cx.emit(PromptEditorEvent::StartRequested);
1896 }
1897 }
1898 }
1899
1900 fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
1901 if let Some(ix) = self.prompt_history_ix {
1902 if ix > 0 {
1903 self.prompt_history_ix = Some(ix - 1);
1904 let prompt = self.prompt_history[ix - 1].as_str();
1905 self.editor.update(cx, |editor, cx| {
1906 editor.set_text(prompt, cx);
1907 editor.move_to_beginning(&Default::default(), cx);
1908 });
1909 }
1910 } else if !self.prompt_history.is_empty() {
1911 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1912 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1913 self.editor.update(cx, |editor, cx| {
1914 editor.set_text(prompt, cx);
1915 editor.move_to_beginning(&Default::default(), cx);
1916 });
1917 }
1918 }
1919
1920 fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
1921 if let Some(ix) = self.prompt_history_ix {
1922 if ix < self.prompt_history.len() - 1 {
1923 self.prompt_history_ix = Some(ix + 1);
1924 let prompt = self.prompt_history[ix + 1].as_str();
1925 self.editor.update(cx, |editor, cx| {
1926 editor.set_text(prompt, cx);
1927 editor.move_to_end(&Default::default(), cx)
1928 });
1929 } else {
1930 self.prompt_history_ix = None;
1931 let prompt = self.pending_prompt.as_str();
1932 self.editor.update(cx, |editor, cx| {
1933 editor.set_text(prompt, cx);
1934 editor.move_to_end(&Default::default(), cx)
1935 });
1936 }
1937 }
1938 }
1939
1940 fn cycle_prev(&mut self, _: &CyclePreviousInlineAssist, cx: &mut ViewContext<Self>) {
1941 self.codegen
1942 .update(cx, |codegen, cx| codegen.cycle_prev(cx));
1943 }
1944
1945 fn cycle_next(&mut self, _: &CycleNextInlineAssist, cx: &mut ViewContext<Self>) {
1946 self.codegen
1947 .update(cx, |codegen, cx| codegen.cycle_next(cx));
1948 }
1949
1950 fn render_cycle_controls(&self, cx: &ViewContext<Self>) -> AnyElement {
1951 let codegen = self.codegen.read(cx);
1952 let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
1953
1954 let model_registry = LanguageModelRegistry::read_global(cx);
1955 let default_model = model_registry.active_model();
1956 let alternative_models = model_registry.inline_alternative_models();
1957
1958 let get_model_name = |index: usize| -> String {
1959 let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
1960
1961 match index {
1962 0 => default_model.as_ref().map_or_else(String::new, name),
1963 index if index <= alternative_models.len() => alternative_models
1964 .get(index - 1)
1965 .map_or_else(String::new, name),
1966 _ => String::new(),
1967 }
1968 };
1969
1970 let total_models = alternative_models.len() + 1;
1971
1972 if total_models <= 1 {
1973 return div().into_any_element();
1974 }
1975
1976 let current_index = codegen.active_alternative;
1977 let prev_index = (current_index + total_models - 1) % total_models;
1978 let next_index = (current_index + 1) % total_models;
1979
1980 let prev_model_name = get_model_name(prev_index);
1981 let next_model_name = get_model_name(next_index);
1982
1983 h_flex()
1984 .child(
1985 IconButton::new("previous", IconName::ChevronLeft)
1986 .icon_color(Color::Muted)
1987 .disabled(disabled || current_index == 0)
1988 .shape(IconButtonShape::Square)
1989 .tooltip({
1990 let focus_handle = self.editor.focus_handle(cx);
1991 move |cx| {
1992 cx.new_view(|cx| {
1993 let mut tooltip = Tooltip::new("Previous Alternative").key_binding(
1994 KeyBinding::for_action_in(
1995 &CyclePreviousInlineAssist,
1996 &focus_handle,
1997 cx,
1998 ),
1999 );
2000 if !disabled && current_index != 0 {
2001 tooltip = tooltip.meta(prev_model_name.clone());
2002 }
2003 tooltip
2004 })
2005 .into()
2006 }
2007 })
2008 .on_click(cx.listener(|this, _, cx| {
2009 this.codegen
2010 .update(cx, |codegen, cx| codegen.cycle_prev(cx))
2011 })),
2012 )
2013 .child(
2014 Label::new(format!(
2015 "{}/{}",
2016 codegen.active_alternative + 1,
2017 codegen.alternative_count(cx)
2018 ))
2019 .size(LabelSize::Small)
2020 .color(if disabled {
2021 Color::Disabled
2022 } else {
2023 Color::Muted
2024 }),
2025 )
2026 .child(
2027 IconButton::new("next", IconName::ChevronRight)
2028 .icon_color(Color::Muted)
2029 .disabled(disabled || current_index == total_models - 1)
2030 .shape(IconButtonShape::Square)
2031 .tooltip({
2032 let focus_handle = self.editor.focus_handle(cx);
2033 move |cx| {
2034 cx.new_view(|cx| {
2035 let mut tooltip = Tooltip::new("Next Alternative").key_binding(
2036 KeyBinding::for_action_in(
2037 &CycleNextInlineAssist,
2038 &focus_handle,
2039 cx,
2040 ),
2041 );
2042 if !disabled && current_index != total_models - 1 {
2043 tooltip = tooltip.meta(next_model_name.clone());
2044 }
2045 tooltip
2046 })
2047 .into()
2048 }
2049 })
2050 .on_click(cx.listener(|this, _, cx| {
2051 this.codegen
2052 .update(cx, |codegen, cx| codegen.cycle_next(cx))
2053 })),
2054 )
2055 .into_any_element()
2056 }
2057
2058 fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
2059 let model = LanguageModelRegistry::read_global(cx).active_model()?;
2060 let token_counts = self.token_counts?;
2061 let max_token_count = model.max_token_count();
2062
2063 let remaining_tokens = max_token_count as isize - token_counts.total as isize;
2064 let token_count_color = if remaining_tokens <= 0 {
2065 Color::Error
2066 } else if token_counts.total as f32 / max_token_count as f32 >= 0.8 {
2067 Color::Warning
2068 } else {
2069 Color::Muted
2070 };
2071
2072 let mut token_count = h_flex()
2073 .id("token_count")
2074 .gap_0p5()
2075 .child(
2076 Label::new(humanize_token_count(token_counts.total))
2077 .size(LabelSize::Small)
2078 .color(token_count_color),
2079 )
2080 .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
2081 .child(
2082 Label::new(humanize_token_count(max_token_count))
2083 .size(LabelSize::Small)
2084 .color(Color::Muted),
2085 );
2086 if let Some(workspace) = self.workspace.clone() {
2087 token_count = token_count
2088 .tooltip(move |cx| {
2089 Tooltip::with_meta(
2090 format!(
2091 "Tokens Used ({} from the Assistant Panel)",
2092 humanize_token_count(token_counts.assistant_panel)
2093 ),
2094 None,
2095 "Click to open the Assistant Panel",
2096 cx,
2097 )
2098 })
2099 .cursor_pointer()
2100 .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
2101 .on_click(move |_, cx| {
2102 cx.stop_propagation();
2103 workspace
2104 .update(cx, |workspace, cx| {
2105 workspace.focus_panel::<AssistantPanel>(cx)
2106 })
2107 .ok();
2108 });
2109 } else {
2110 token_count = token_count
2111 .cursor_default()
2112 .tooltip(|cx| Tooltip::text("Tokens used", cx));
2113 }
2114
2115 Some(token_count)
2116 }
2117
2118 fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
2119 let settings = ThemeSettings::get_global(cx);
2120 let text_style = TextStyle {
2121 color: if self.editor.read(cx).read_only(cx) {
2122 cx.theme().colors().text_disabled
2123 } else {
2124 cx.theme().colors().text
2125 },
2126 font_family: settings.buffer_font.family.clone(),
2127 font_fallbacks: settings.buffer_font.fallbacks.clone(),
2128 font_size: settings.buffer_font_size.into(),
2129 font_weight: settings.buffer_font.weight,
2130 line_height: relative(settings.buffer_line_height.value()),
2131 ..Default::default()
2132 };
2133 EditorElement::new(
2134 &self.editor,
2135 EditorStyle {
2136 background: cx.theme().colors().editor_background,
2137 local_player: cx.theme().players().local(),
2138 text: text_style,
2139 ..Default::default()
2140 },
2141 )
2142 }
2143
2144 fn render_rate_limit_notice(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
2145 Popover::new().child(
2146 v_flex()
2147 .occlude()
2148 .p_2()
2149 .child(
2150 Label::new("Out of Tokens")
2151 .size(LabelSize::Small)
2152 .weight(FontWeight::BOLD),
2153 )
2154 .child(Label::new(
2155 "Try Zed Pro for higher limits, a wider range of models, and more.",
2156 ))
2157 .child(
2158 h_flex()
2159 .justify_between()
2160 .child(CheckboxWithLabel::new(
2161 "dont-show-again",
2162 Label::new("Don't show again"),
2163 if dismissed_rate_limit_notice() {
2164 ui::ToggleState::Selected
2165 } else {
2166 ui::ToggleState::Unselected
2167 },
2168 |selection, cx| {
2169 let is_dismissed = match selection {
2170 ui::ToggleState::Unselected => false,
2171 ui::ToggleState::Indeterminate => return,
2172 ui::ToggleState::Selected => true,
2173 };
2174
2175 set_rate_limit_notice_dismissed(is_dismissed, cx)
2176 },
2177 ))
2178 .child(
2179 h_flex()
2180 .gap_2()
2181 .child(
2182 Button::new("dismiss", "Dismiss")
2183 .style(ButtonStyle::Transparent)
2184 .on_click(cx.listener(Self::toggle_rate_limit_notice)),
2185 )
2186 .child(Button::new("more-info", "More Info").on_click(
2187 |_event, cx| {
2188 cx.dispatch_action(Box::new(
2189 zed_actions::OpenAccountSettings,
2190 ))
2191 },
2192 )),
2193 ),
2194 ),
2195 )
2196 }
2197}
2198
2199const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
2200
2201fn dismissed_rate_limit_notice() -> bool {
2202 db::kvp::KEY_VALUE_STORE
2203 .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
2204 .log_err()
2205 .map_or(false, |s| s.is_some())
2206}
2207
2208fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut AppContext) {
2209 db::write_and_log(cx, move || async move {
2210 if is_dismissed {
2211 db::kvp::KEY_VALUE_STORE
2212 .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
2213 .await
2214 } else {
2215 db::kvp::KEY_VALUE_STORE
2216 .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
2217 .await
2218 }
2219 })
2220}
2221
2222struct InlineAssist {
2223 group_id: InlineAssistGroupId,
2224 range: Range<Anchor>,
2225 editor: WeakView<Editor>,
2226 decorations: Option<InlineAssistDecorations>,
2227 codegen: Model<Codegen>,
2228 _subscriptions: Vec<Subscription>,
2229 workspace: Option<WeakView<Workspace>>,
2230 include_context: bool,
2231}
2232
2233impl InlineAssist {
2234 #[allow(clippy::too_many_arguments)]
2235 fn new(
2236 assist_id: InlineAssistId,
2237 group_id: InlineAssistGroupId,
2238 include_context: bool,
2239 editor: &View<Editor>,
2240 prompt_editor: &View<PromptEditor>,
2241 prompt_block_id: CustomBlockId,
2242 end_block_id: CustomBlockId,
2243 range: Range<Anchor>,
2244 codegen: Model<Codegen>,
2245 workspace: Option<WeakView<Workspace>>,
2246 cx: &mut WindowContext,
2247 ) -> Self {
2248 let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
2249 InlineAssist {
2250 group_id,
2251 include_context,
2252 editor: editor.downgrade(),
2253 decorations: Some(InlineAssistDecorations {
2254 prompt_block_id,
2255 prompt_editor: prompt_editor.clone(),
2256 removed_line_block_ids: HashSet::default(),
2257 end_block_id,
2258 }),
2259 range,
2260 codegen: codegen.clone(),
2261 workspace: workspace.clone(),
2262 _subscriptions: vec![
2263 cx.on_focus_in(&prompt_editor_focus_handle, move |cx| {
2264 InlineAssistant::update_global(cx, |this, cx| {
2265 this.handle_prompt_editor_focus_in(assist_id, cx)
2266 })
2267 }),
2268 cx.on_focus_out(&prompt_editor_focus_handle, move |_, cx| {
2269 InlineAssistant::update_global(cx, |this, cx| {
2270 this.handle_prompt_editor_focus_out(assist_id, cx)
2271 })
2272 }),
2273 cx.subscribe(prompt_editor, |prompt_editor, event, cx| {
2274 InlineAssistant::update_global(cx, |this, cx| {
2275 this.handle_prompt_editor_event(prompt_editor, event, cx)
2276 })
2277 }),
2278 cx.observe(&codegen, {
2279 let editor = editor.downgrade();
2280 move |_, cx| {
2281 if let Some(editor) = editor.upgrade() {
2282 InlineAssistant::update_global(cx, |this, cx| {
2283 if let Some(editor_assists) =
2284 this.assists_by_editor.get(&editor.downgrade())
2285 {
2286 editor_assists.highlight_updates.send(()).ok();
2287 }
2288
2289 this.update_editor_blocks(&editor, assist_id, cx);
2290 })
2291 }
2292 }
2293 }),
2294 cx.subscribe(&codegen, move |codegen, event, cx| {
2295 InlineAssistant::update_global(cx, |this, cx| match event {
2296 CodegenEvent::Undone => this.finish_assist(assist_id, false, cx),
2297 CodegenEvent::Finished => {
2298 let assist = if let Some(assist) = this.assists.get(&assist_id) {
2299 assist
2300 } else {
2301 return;
2302 };
2303
2304 if let CodegenStatus::Error(error) = codegen.read(cx).status(cx) {
2305 if assist.decorations.is_none() {
2306 if let Some(workspace) = assist
2307 .workspace
2308 .as_ref()
2309 .and_then(|workspace| workspace.upgrade())
2310 {
2311 let error = format!("Inline assistant error: {}", error);
2312 workspace.update(cx, |workspace, cx| {
2313 struct InlineAssistantError;
2314
2315 let id =
2316 NotificationId::composite::<InlineAssistantError>(
2317 assist_id.0,
2318 );
2319
2320 workspace.show_toast(Toast::new(id, error), cx);
2321 })
2322 }
2323 }
2324 }
2325
2326 if assist.decorations.is_none() {
2327 this.finish_assist(assist_id, false, cx);
2328 }
2329 }
2330 })
2331 }),
2332 ],
2333 }
2334 }
2335
2336 fn user_prompt(&self, cx: &AppContext) -> Option<String> {
2337 let decorations = self.decorations.as_ref()?;
2338 Some(decorations.prompt_editor.read(cx).prompt(cx))
2339 }
2340
2341 fn assistant_panel_context(&self, cx: &WindowContext) -> Option<LanguageModelRequest> {
2342 if self.include_context {
2343 let workspace = self.workspace.as_ref()?;
2344 let workspace = workspace.upgrade()?.read(cx);
2345 let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
2346 Some(
2347 assistant_panel
2348 .read(cx)
2349 .active_context(cx)?
2350 .read(cx)
2351 .to_completion_request(RequestType::Chat, cx),
2352 )
2353 } else {
2354 None
2355 }
2356 }
2357
2358 pub fn count_tokens(&self, cx: &WindowContext) -> BoxFuture<'static, Result<TokenCounts>> {
2359 let Some(user_prompt) = self.user_prompt(cx) else {
2360 return future::ready(Err(anyhow!("no user prompt"))).boxed();
2361 };
2362 let assistant_panel_context = self.assistant_panel_context(cx);
2363 self.codegen
2364 .read(cx)
2365 .count_tokens(user_prompt, assistant_panel_context, cx)
2366 }
2367}
2368
2369struct InlineAssistDecorations {
2370 prompt_block_id: CustomBlockId,
2371 prompt_editor: View<PromptEditor>,
2372 removed_line_block_ids: HashSet<CustomBlockId>,
2373 end_block_id: CustomBlockId,
2374}
2375
2376#[derive(Copy, Clone, Debug)]
2377pub enum CodegenEvent {
2378 Finished,
2379 Undone,
2380}
2381
2382pub struct Codegen {
2383 alternatives: Vec<Model<CodegenAlternative>>,
2384 active_alternative: usize,
2385 seen_alternatives: HashSet<usize>,
2386 subscriptions: Vec<Subscription>,
2387 buffer: Model<MultiBuffer>,
2388 range: Range<Anchor>,
2389 initial_transaction_id: Option<TransactionId>,
2390 telemetry: Arc<Telemetry>,
2391 builder: Arc<PromptBuilder>,
2392 is_insertion: bool,
2393}
2394
2395impl Codegen {
2396 pub fn new(
2397 buffer: Model<MultiBuffer>,
2398 range: Range<Anchor>,
2399 initial_transaction_id: Option<TransactionId>,
2400 telemetry: Arc<Telemetry>,
2401 builder: Arc<PromptBuilder>,
2402 cx: &mut ModelContext<Self>,
2403 ) -> Self {
2404 let codegen = cx.new_model(|cx| {
2405 CodegenAlternative::new(
2406 buffer.clone(),
2407 range.clone(),
2408 false,
2409 Some(telemetry.clone()),
2410 builder.clone(),
2411 cx,
2412 )
2413 });
2414 let mut this = Self {
2415 is_insertion: range.to_offset(&buffer.read(cx).snapshot(cx)).is_empty(),
2416 alternatives: vec![codegen],
2417 active_alternative: 0,
2418 seen_alternatives: HashSet::default(),
2419 subscriptions: Vec::new(),
2420 buffer,
2421 range,
2422 initial_transaction_id,
2423 telemetry,
2424 builder,
2425 };
2426 this.activate(0, cx);
2427 this
2428 }
2429
2430 fn subscribe_to_alternative(&mut self, cx: &mut ModelContext<Self>) {
2431 let codegen = self.active_alternative().clone();
2432 self.subscriptions.clear();
2433 self.subscriptions
2434 .push(cx.observe(&codegen, |_, _, cx| cx.notify()));
2435 self.subscriptions
2436 .push(cx.subscribe(&codegen, |_, _, event, cx| cx.emit(*event)));
2437 }
2438
2439 fn active_alternative(&self) -> &Model<CodegenAlternative> {
2440 &self.alternatives[self.active_alternative]
2441 }
2442
2443 fn status<'a>(&self, cx: &'a AppContext) -> &'a CodegenStatus {
2444 &self.active_alternative().read(cx).status
2445 }
2446
2447 fn alternative_count(&self, cx: &AppContext) -> usize {
2448 LanguageModelRegistry::read_global(cx)
2449 .inline_alternative_models()
2450 .len()
2451 + 1
2452 }
2453
2454 pub fn cycle_prev(&mut self, cx: &mut ModelContext<Self>) {
2455 let next_active_ix = if self.active_alternative == 0 {
2456 self.alternatives.len() - 1
2457 } else {
2458 self.active_alternative - 1
2459 };
2460 self.activate(next_active_ix, cx);
2461 }
2462
2463 pub fn cycle_next(&mut self, cx: &mut ModelContext<Self>) {
2464 let next_active_ix = (self.active_alternative + 1) % self.alternatives.len();
2465 self.activate(next_active_ix, cx);
2466 }
2467
2468 fn activate(&mut self, index: usize, cx: &mut ModelContext<Self>) {
2469 self.active_alternative()
2470 .update(cx, |codegen, cx| codegen.set_active(false, cx));
2471 self.seen_alternatives.insert(index);
2472 self.active_alternative = index;
2473 self.active_alternative()
2474 .update(cx, |codegen, cx| codegen.set_active(true, cx));
2475 self.subscribe_to_alternative(cx);
2476 cx.notify();
2477 }
2478
2479 pub fn start(
2480 &mut self,
2481 user_prompt: String,
2482 assistant_panel_context: Option<LanguageModelRequest>,
2483 cx: &mut ModelContext<Self>,
2484 ) -> Result<()> {
2485 let alternative_models = LanguageModelRegistry::read_global(cx)
2486 .inline_alternative_models()
2487 .to_vec();
2488
2489 self.active_alternative()
2490 .update(cx, |alternative, cx| alternative.undo(cx));
2491 self.activate(0, cx);
2492 self.alternatives.truncate(1);
2493
2494 for _ in 0..alternative_models.len() {
2495 self.alternatives.push(cx.new_model(|cx| {
2496 CodegenAlternative::new(
2497 self.buffer.clone(),
2498 self.range.clone(),
2499 false,
2500 Some(self.telemetry.clone()),
2501 self.builder.clone(),
2502 cx,
2503 )
2504 }));
2505 }
2506
2507 let primary_model = LanguageModelRegistry::read_global(cx)
2508 .active_model()
2509 .context("no active model")?;
2510
2511 for (model, alternative) in iter::once(primary_model)
2512 .chain(alternative_models)
2513 .zip(&self.alternatives)
2514 {
2515 alternative.update(cx, |alternative, cx| {
2516 alternative.start(
2517 user_prompt.clone(),
2518 assistant_panel_context.clone(),
2519 model.clone(),
2520 cx,
2521 )
2522 })?;
2523 }
2524
2525 Ok(())
2526 }
2527
2528 pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
2529 for codegen in &self.alternatives {
2530 codegen.update(cx, |codegen, cx| codegen.stop(cx));
2531 }
2532 }
2533
2534 pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
2535 self.active_alternative()
2536 .update(cx, |codegen, cx| codegen.undo(cx));
2537
2538 self.buffer.update(cx, |buffer, cx| {
2539 if let Some(transaction_id) = self.initial_transaction_id.take() {
2540 buffer.undo_transaction(transaction_id, cx);
2541 buffer.refresh_preview(cx);
2542 }
2543 });
2544 }
2545
2546 pub fn count_tokens(
2547 &self,
2548 user_prompt: String,
2549 assistant_panel_context: Option<LanguageModelRequest>,
2550 cx: &AppContext,
2551 ) -> BoxFuture<'static, Result<TokenCounts>> {
2552 self.active_alternative()
2553 .read(cx)
2554 .count_tokens(user_prompt, assistant_panel_context, cx)
2555 }
2556
2557 pub fn buffer(&self, cx: &AppContext) -> Model<MultiBuffer> {
2558 self.active_alternative().read(cx).buffer.clone()
2559 }
2560
2561 pub fn old_buffer(&self, cx: &AppContext) -> Model<Buffer> {
2562 self.active_alternative().read(cx).old_buffer.clone()
2563 }
2564
2565 pub fn snapshot(&self, cx: &AppContext) -> MultiBufferSnapshot {
2566 self.active_alternative().read(cx).snapshot.clone()
2567 }
2568
2569 pub fn edit_position(&self, cx: &AppContext) -> Option<Anchor> {
2570 self.active_alternative().read(cx).edit_position
2571 }
2572
2573 fn diff<'a>(&self, cx: &'a AppContext) -> &'a Diff {
2574 &self.active_alternative().read(cx).diff
2575 }
2576
2577 pub fn last_equal_ranges<'a>(&self, cx: &'a AppContext) -> &'a [Range<Anchor>] {
2578 self.active_alternative().read(cx).last_equal_ranges()
2579 }
2580}
2581
2582impl EventEmitter<CodegenEvent> for Codegen {}
2583
2584pub struct CodegenAlternative {
2585 buffer: Model<MultiBuffer>,
2586 old_buffer: Model<Buffer>,
2587 snapshot: MultiBufferSnapshot,
2588 edit_position: Option<Anchor>,
2589 range: Range<Anchor>,
2590 last_equal_ranges: Vec<Range<Anchor>>,
2591 transformation_transaction_id: Option<TransactionId>,
2592 status: CodegenStatus,
2593 generation: Task<()>,
2594 diff: Diff,
2595 telemetry: Option<Arc<Telemetry>>,
2596 _subscription: gpui::Subscription,
2597 builder: Arc<PromptBuilder>,
2598 active: bool,
2599 edits: Vec<(Range<Anchor>, String)>,
2600 line_operations: Vec<LineOperation>,
2601 request: Option<LanguageModelRequest>,
2602 elapsed_time: Option<f64>,
2603 completion: Option<String>,
2604 message_id: Option<String>,
2605}
2606
2607enum CodegenStatus {
2608 Idle,
2609 Pending,
2610 Done,
2611 Error(anyhow::Error),
2612}
2613
2614#[derive(Default)]
2615struct Diff {
2616 deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
2617 inserted_row_ranges: Vec<Range<Anchor>>,
2618}
2619
2620impl Diff {
2621 fn is_empty(&self) -> bool {
2622 self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
2623 }
2624}
2625
2626impl EventEmitter<CodegenEvent> for CodegenAlternative {}
2627
2628impl CodegenAlternative {
2629 pub fn new(
2630 multi_buffer: Model<MultiBuffer>,
2631 range: Range<Anchor>,
2632 active: bool,
2633 telemetry: Option<Arc<Telemetry>>,
2634 builder: Arc<PromptBuilder>,
2635 cx: &mut ModelContext<Self>,
2636 ) -> Self {
2637 let snapshot = multi_buffer.read(cx).snapshot(cx);
2638
2639 let (old_excerpt, _) = snapshot
2640 .range_to_buffer_ranges(range.clone())
2641 .pop()
2642 .unwrap();
2643 let old_buffer = cx.new_model(|cx| {
2644 let text = old_excerpt.buffer().as_rope().clone();
2645 let line_ending = old_excerpt.buffer().line_ending();
2646 let language = old_excerpt.buffer().language().cloned();
2647 let language_registry = multi_buffer
2648 .read(cx)
2649 .buffer(old_excerpt.buffer_id())
2650 .unwrap()
2651 .read(cx)
2652 .language_registry();
2653
2654 let mut buffer = Buffer::local_normalized(text, line_ending, cx);
2655 buffer.set_language(language, cx);
2656 if let Some(language_registry) = language_registry {
2657 buffer.set_language_registry(language_registry)
2658 }
2659 buffer
2660 });
2661
2662 Self {
2663 buffer: multi_buffer.clone(),
2664 old_buffer,
2665 edit_position: None,
2666 message_id: None,
2667 snapshot,
2668 last_equal_ranges: Default::default(),
2669 transformation_transaction_id: None,
2670 status: CodegenStatus::Idle,
2671 generation: Task::ready(()),
2672 diff: Diff::default(),
2673 telemetry,
2674 _subscription: cx.subscribe(&multi_buffer, Self::handle_buffer_event),
2675 builder,
2676 active,
2677 edits: Vec::new(),
2678 line_operations: Vec::new(),
2679 range,
2680 request: None,
2681 elapsed_time: None,
2682 completion: None,
2683 }
2684 }
2685
2686 fn set_active(&mut self, active: bool, cx: &mut ModelContext<Self>) {
2687 if active != self.active {
2688 self.active = active;
2689
2690 if self.active {
2691 let edits = self.edits.clone();
2692 self.apply_edits(edits, cx);
2693 if matches!(self.status, CodegenStatus::Pending) {
2694 let line_operations = self.line_operations.clone();
2695 self.reapply_line_based_diff(line_operations, cx);
2696 } else {
2697 self.reapply_batch_diff(cx).detach();
2698 }
2699 } else if let Some(transaction_id) = self.transformation_transaction_id.take() {
2700 self.buffer.update(cx, |buffer, cx| {
2701 buffer.undo_transaction(transaction_id, cx);
2702 buffer.forget_transaction(transaction_id, cx);
2703 });
2704 }
2705 }
2706 }
2707
2708 fn handle_buffer_event(
2709 &mut self,
2710 _buffer: Model<MultiBuffer>,
2711 event: &multi_buffer::Event,
2712 cx: &mut ModelContext<Self>,
2713 ) {
2714 if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
2715 if self.transformation_transaction_id == Some(*transaction_id) {
2716 self.transformation_transaction_id = None;
2717 self.generation = Task::ready(());
2718 cx.emit(CodegenEvent::Undone);
2719 }
2720 }
2721 }
2722
2723 pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
2724 &self.last_equal_ranges
2725 }
2726
2727 pub fn count_tokens(
2728 &self,
2729 user_prompt: String,
2730 assistant_panel_context: Option<LanguageModelRequest>,
2731 cx: &AppContext,
2732 ) -> BoxFuture<'static, Result<TokenCounts>> {
2733 if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
2734 let request = self.build_request(user_prompt, assistant_panel_context.clone(), cx);
2735 match request {
2736 Ok(request) => {
2737 let total_count = model.count_tokens(request.clone(), cx);
2738 let assistant_panel_count = assistant_panel_context
2739 .map(|context| model.count_tokens(context, cx))
2740 .unwrap_or_else(|| future::ready(Ok(0)).boxed());
2741
2742 async move {
2743 Ok(TokenCounts {
2744 total: total_count.await?,
2745 assistant_panel: assistant_panel_count.await?,
2746 })
2747 }
2748 .boxed()
2749 }
2750 Err(error) => futures::future::ready(Err(error)).boxed(),
2751 }
2752 } else {
2753 future::ready(Err(anyhow!("no active model"))).boxed()
2754 }
2755 }
2756
2757 pub fn start(
2758 &mut self,
2759 user_prompt: String,
2760 assistant_panel_context: Option<LanguageModelRequest>,
2761 model: Arc<dyn LanguageModel>,
2762 cx: &mut ModelContext<Self>,
2763 ) -> Result<()> {
2764 if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {
2765 self.buffer.update(cx, |buffer, cx| {
2766 buffer.undo_transaction(transformation_transaction_id, cx);
2767 });
2768 }
2769
2770 self.edit_position = Some(self.range.start.bias_right(&self.snapshot));
2771
2772 let api_key = model.api_key(cx);
2773 let telemetry_id = model.telemetry_id();
2774 let provider_id = model.provider_id();
2775 let stream: LocalBoxFuture<Result<LanguageModelTextStream>> =
2776 if user_prompt.trim().to_lowercase() == "delete" {
2777 async { Ok(LanguageModelTextStream::default()) }.boxed_local()
2778 } else {
2779 let request = self.build_request(user_prompt, assistant_panel_context, cx)?;
2780 self.request = Some(request.clone());
2781
2782 cx.spawn(|_, cx| async move { model.stream_completion_text(request, &cx).await })
2783 .boxed_local()
2784 };
2785 self.handle_stream(telemetry_id, provider_id.to_string(), api_key, stream, cx);
2786 Ok(())
2787 }
2788
2789 fn build_request(
2790 &self,
2791 user_prompt: String,
2792 assistant_panel_context: Option<LanguageModelRequest>,
2793 cx: &AppContext,
2794 ) -> Result<LanguageModelRequest> {
2795 let buffer = self.buffer.read(cx).snapshot(cx);
2796 let language = buffer.language_at(self.range.start);
2797 let language_name = if let Some(language) = language.as_ref() {
2798 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
2799 None
2800 } else {
2801 Some(language.name())
2802 }
2803 } else {
2804 None
2805 };
2806
2807 let language_name = language_name.as_ref();
2808 let start = buffer.point_to_buffer_offset(self.range.start);
2809 let end = buffer.point_to_buffer_offset(self.range.end);
2810 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
2811 let (start_buffer, start_buffer_offset) = start;
2812 let (end_buffer, end_buffer_offset) = end;
2813 if start_buffer.remote_id() == end_buffer.remote_id() {
2814 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
2815 } else {
2816 return Err(anyhow::anyhow!("invalid transformation range"));
2817 }
2818 } else {
2819 return Err(anyhow::anyhow!("invalid transformation range"));
2820 };
2821
2822 let prompt = self
2823 .builder
2824 .generate_inline_transformation_prompt(user_prompt, language_name, buffer, range)
2825 .map_err(|e| anyhow::anyhow!("Failed to generate content prompt: {}", e))?;
2826
2827 let mut messages = Vec::new();
2828 if let Some(context_request) = assistant_panel_context {
2829 messages = context_request.messages;
2830 }
2831
2832 messages.push(LanguageModelRequestMessage {
2833 role: Role::User,
2834 content: vec![prompt.into()],
2835 cache: false,
2836 });
2837
2838 Ok(LanguageModelRequest {
2839 messages,
2840 tools: Vec::new(),
2841 stop: Vec::new(),
2842 temperature: None,
2843 })
2844 }
2845
2846 pub fn handle_stream(
2847 &mut self,
2848 model_telemetry_id: String,
2849 model_provider_id: String,
2850 model_api_key: Option<String>,
2851 stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,
2852 cx: &mut ModelContext<Self>,
2853 ) {
2854 let start_time = Instant::now();
2855 let snapshot = self.snapshot.clone();
2856 let selected_text = snapshot
2857 .text_for_range(self.range.start..self.range.end)
2858 .collect::<Rope>();
2859
2860 let selection_start = self.range.start.to_point(&snapshot);
2861
2862 // Start with the indentation of the first line in the selection
2863 let mut suggested_line_indent = snapshot
2864 .suggested_indents(selection_start.row..=selection_start.row, cx)
2865 .into_values()
2866 .next()
2867 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
2868
2869 // If the first line in the selection does not have indentation, check the following lines
2870 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
2871 for row in selection_start.row..=self.range.end.to_point(&snapshot).row {
2872 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
2873 // Prefer tabs if a line in the selection uses tabs as indentation
2874 if line_indent.kind == IndentKind::Tab {
2875 suggested_line_indent.kind = IndentKind::Tab;
2876 break;
2877 }
2878 }
2879 }
2880
2881 let http_client = cx.http_client().clone();
2882 let telemetry = self.telemetry.clone();
2883 let language_name = {
2884 let multibuffer = self.buffer.read(cx);
2885 let snapshot = multibuffer.snapshot(cx);
2886 let ranges = snapshot.range_to_buffer_ranges(self.range.clone());
2887 ranges
2888 .first()
2889 .and_then(|(excerpt, _)| excerpt.buffer().language())
2890 .map(|language| language.name())
2891 };
2892
2893 self.diff = Diff::default();
2894 self.status = CodegenStatus::Pending;
2895 let mut edit_start = self.range.start.to_offset(&snapshot);
2896 let completion = Arc::new(Mutex::new(String::new()));
2897 let completion_clone = completion.clone();
2898
2899 self.generation = cx.spawn(|codegen, mut cx| {
2900 async move {
2901 let stream = stream.await;
2902 let message_id = stream
2903 .as_ref()
2904 .ok()
2905 .and_then(|stream| stream.message_id.clone());
2906 let generate = async {
2907 let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
2908 let executor = cx.background_executor().clone();
2909 let message_id = message_id.clone();
2910 let line_based_stream_diff: Task<anyhow::Result<()>> =
2911 cx.background_executor().spawn(async move {
2912 let mut response_latency = None;
2913 let request_start = Instant::now();
2914 let diff = async {
2915 let chunks = StripInvalidSpans::new(stream?.stream);
2916 futures::pin_mut!(chunks);
2917 let mut diff = StreamingDiff::new(selected_text.to_string());
2918 let mut line_diff = LineDiff::default();
2919
2920 let mut new_text = String::new();
2921 let mut base_indent = None;
2922 let mut line_indent = None;
2923 let mut first_line = true;
2924
2925 while let Some(chunk) = chunks.next().await {
2926 if response_latency.is_none() {
2927 response_latency = Some(request_start.elapsed());
2928 }
2929 let chunk = chunk?;
2930 completion_clone.lock().push_str(&chunk);
2931
2932 let mut lines = chunk.split('\n').peekable();
2933 while let Some(line) = lines.next() {
2934 new_text.push_str(line);
2935 if line_indent.is_none() {
2936 if let Some(non_whitespace_ch_ix) =
2937 new_text.find(|ch: char| !ch.is_whitespace())
2938 {
2939 line_indent = Some(non_whitespace_ch_ix);
2940 base_indent = base_indent.or(line_indent);
2941
2942 let line_indent = line_indent.unwrap();
2943 let base_indent = base_indent.unwrap();
2944 let indent_delta =
2945 line_indent as i32 - base_indent as i32;
2946 let mut corrected_indent_len = cmp::max(
2947 0,
2948 suggested_line_indent.len as i32 + indent_delta,
2949 )
2950 as usize;
2951 if first_line {
2952 corrected_indent_len = corrected_indent_len
2953 .saturating_sub(
2954 selection_start.column as usize,
2955 );
2956 }
2957
2958 let indent_char = suggested_line_indent.char();
2959 let mut indent_buffer = [0; 4];
2960 let indent_str =
2961 indent_char.encode_utf8(&mut indent_buffer);
2962 new_text.replace_range(
2963 ..line_indent,
2964 &indent_str.repeat(corrected_indent_len),
2965 );
2966 }
2967 }
2968
2969 if line_indent.is_some() {
2970 let char_ops = diff.push_new(&new_text);
2971 line_diff
2972 .push_char_operations(&char_ops, &selected_text);
2973 diff_tx
2974 .send((char_ops, line_diff.line_operations()))
2975 .await?;
2976 new_text.clear();
2977 }
2978
2979 if lines.peek().is_some() {
2980 let char_ops = diff.push_new("\n");
2981 line_diff
2982 .push_char_operations(&char_ops, &selected_text);
2983 diff_tx
2984 .send((char_ops, line_diff.line_operations()))
2985 .await?;
2986 if line_indent.is_none() {
2987 // Don't write out the leading indentation in empty lines on the next line
2988 // This is the case where the above if statement didn't clear the buffer
2989 new_text.clear();
2990 }
2991 line_indent = None;
2992 first_line = false;
2993 }
2994 }
2995 }
2996
2997 let mut char_ops = diff.push_new(&new_text);
2998 char_ops.extend(diff.finish());
2999 line_diff.push_char_operations(&char_ops, &selected_text);
3000 line_diff.finish(&selected_text);
3001 diff_tx
3002 .send((char_ops, line_diff.line_operations()))
3003 .await?;
3004
3005 anyhow::Ok(())
3006 };
3007
3008 let result = diff.await;
3009
3010 let error_message =
3011 result.as_ref().err().map(|error| error.to_string());
3012 report_assistant_event(
3013 AssistantEvent {
3014 conversation_id: None,
3015 message_id,
3016 kind: AssistantKind::Inline,
3017 phase: AssistantPhase::Response,
3018 model: model_telemetry_id,
3019 model_provider: model_provider_id.to_string(),
3020 response_latency,
3021 error_message,
3022 language_name: language_name.map(|name| name.to_proto()),
3023 },
3024 telemetry,
3025 http_client,
3026 model_api_key,
3027 &executor,
3028 );
3029
3030 result?;
3031 Ok(())
3032 });
3033
3034 while let Some((char_ops, line_ops)) = diff_rx.next().await {
3035 codegen.update(&mut cx, |codegen, cx| {
3036 codegen.last_equal_ranges.clear();
3037
3038 let edits = char_ops
3039 .into_iter()
3040 .filter_map(|operation| match operation {
3041 CharOperation::Insert { text } => {
3042 let edit_start = snapshot.anchor_after(edit_start);
3043 Some((edit_start..edit_start, text))
3044 }
3045 CharOperation::Delete { bytes } => {
3046 let edit_end = edit_start + bytes;
3047 let edit_range = snapshot.anchor_after(edit_start)
3048 ..snapshot.anchor_before(edit_end);
3049 edit_start = edit_end;
3050 Some((edit_range, String::new()))
3051 }
3052 CharOperation::Keep { bytes } => {
3053 let edit_end = edit_start + bytes;
3054 let edit_range = snapshot.anchor_after(edit_start)
3055 ..snapshot.anchor_before(edit_end);
3056 edit_start = edit_end;
3057 codegen.last_equal_ranges.push(edit_range);
3058 None
3059 }
3060 })
3061 .collect::<Vec<_>>();
3062
3063 if codegen.active {
3064 codegen.apply_edits(edits.iter().cloned(), cx);
3065 codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
3066 }
3067 codegen.edits.extend(edits);
3068 codegen.line_operations = line_ops;
3069 codegen.edit_position = Some(snapshot.anchor_after(edit_start));
3070
3071 cx.notify();
3072 })?;
3073 }
3074
3075 // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
3076 // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
3077 // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
3078 let batch_diff_task =
3079 codegen.update(&mut cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
3080 let (line_based_stream_diff, ()) =
3081 join!(line_based_stream_diff, batch_diff_task);
3082 line_based_stream_diff?;
3083
3084 anyhow::Ok(())
3085 };
3086
3087 let result = generate.await;
3088 let elapsed_time = start_time.elapsed().as_secs_f64();
3089
3090 codegen
3091 .update(&mut cx, |this, cx| {
3092 this.message_id = message_id;
3093 this.last_equal_ranges.clear();
3094 if let Err(error) = result {
3095 this.status = CodegenStatus::Error(error);
3096 } else {
3097 this.status = CodegenStatus::Done;
3098 }
3099 this.elapsed_time = Some(elapsed_time);
3100 this.completion = Some(completion.lock().clone());
3101 cx.emit(CodegenEvent::Finished);
3102 cx.notify();
3103 })
3104 .ok();
3105 }
3106 });
3107 cx.notify();
3108 }
3109
3110 pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
3111 self.last_equal_ranges.clear();
3112 if self.diff.is_empty() {
3113 self.status = CodegenStatus::Idle;
3114 } else {
3115 self.status = CodegenStatus::Done;
3116 }
3117 self.generation = Task::ready(());
3118 cx.emit(CodegenEvent::Finished);
3119 cx.notify();
3120 }
3121
3122 pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
3123 self.buffer.update(cx, |buffer, cx| {
3124 if let Some(transaction_id) = self.transformation_transaction_id.take() {
3125 buffer.undo_transaction(transaction_id, cx);
3126 buffer.refresh_preview(cx);
3127 }
3128 });
3129 }
3130
3131 fn apply_edits(
3132 &mut self,
3133 edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
3134 cx: &mut ModelContext<CodegenAlternative>,
3135 ) {
3136 let transaction = self.buffer.update(cx, |buffer, cx| {
3137 // Avoid grouping assistant edits with user edits.
3138 buffer.finalize_last_transaction(cx);
3139 buffer.start_transaction(cx);
3140 buffer.edit(edits, None, cx);
3141 buffer.end_transaction(cx)
3142 });
3143
3144 if let Some(transaction) = transaction {
3145 if let Some(first_transaction) = self.transformation_transaction_id {
3146 // Group all assistant edits into the first transaction.
3147 self.buffer.update(cx, |buffer, cx| {
3148 buffer.merge_transactions(transaction, first_transaction, cx)
3149 });
3150 } else {
3151 self.transformation_transaction_id = Some(transaction);
3152 self.buffer
3153 .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
3154 }
3155 }
3156 }
3157
3158 fn reapply_line_based_diff(
3159 &mut self,
3160 line_operations: impl IntoIterator<Item = LineOperation>,
3161 cx: &mut ModelContext<Self>,
3162 ) {
3163 let old_snapshot = self.snapshot.clone();
3164 let old_range = self.range.to_point(&old_snapshot);
3165 let new_snapshot = self.buffer.read(cx).snapshot(cx);
3166 let new_range = self.range.to_point(&new_snapshot);
3167
3168 let mut old_row = old_range.start.row;
3169 let mut new_row = new_range.start.row;
3170
3171 self.diff.deleted_row_ranges.clear();
3172 self.diff.inserted_row_ranges.clear();
3173 for operation in line_operations {
3174 match operation {
3175 LineOperation::Keep { lines } => {
3176 old_row += lines;
3177 new_row += lines;
3178 }
3179 LineOperation::Delete { lines } => {
3180 let old_end_row = old_row + lines - 1;
3181 let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
3182
3183 if let Some((_, last_deleted_row_range)) =
3184 self.diff.deleted_row_ranges.last_mut()
3185 {
3186 if *last_deleted_row_range.end() + 1 == old_row {
3187 *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
3188 } else {
3189 self.diff
3190 .deleted_row_ranges
3191 .push((new_row, old_row..=old_end_row));
3192 }
3193 } else {
3194 self.diff
3195 .deleted_row_ranges
3196 .push((new_row, old_row..=old_end_row));
3197 }
3198
3199 old_row += lines;
3200 }
3201 LineOperation::Insert { lines } => {
3202 let new_end_row = new_row + lines - 1;
3203 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
3204 let end = new_snapshot.anchor_before(Point::new(
3205 new_end_row,
3206 new_snapshot.line_len(MultiBufferRow(new_end_row)),
3207 ));
3208 self.diff.inserted_row_ranges.push(start..end);
3209 new_row += lines;
3210 }
3211 }
3212
3213 cx.notify();
3214 }
3215 }
3216
3217 fn reapply_batch_diff(&mut self, cx: &mut ModelContext<Self>) -> Task<()> {
3218 let old_snapshot = self.snapshot.clone();
3219 let old_range = self.range.to_point(&old_snapshot);
3220 let new_snapshot = self.buffer.read(cx).snapshot(cx);
3221 let new_range = self.range.to_point(&new_snapshot);
3222
3223 cx.spawn(|codegen, mut cx| async move {
3224 let (deleted_row_ranges, inserted_row_ranges) = cx
3225 .background_executor()
3226 .spawn(async move {
3227 let old_text = old_snapshot
3228 .text_for_range(
3229 Point::new(old_range.start.row, 0)
3230 ..Point::new(
3231 old_range.end.row,
3232 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
3233 ),
3234 )
3235 .collect::<String>();
3236 let new_text = new_snapshot
3237 .text_for_range(
3238 Point::new(new_range.start.row, 0)
3239 ..Point::new(
3240 new_range.end.row,
3241 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
3242 ),
3243 )
3244 .collect::<String>();
3245
3246 let mut old_row = old_range.start.row;
3247 let mut new_row = new_range.start.row;
3248 let batch_diff =
3249 similar::TextDiff::from_lines(old_text.as_str(), new_text.as_str());
3250
3251 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
3252 let mut inserted_row_ranges = Vec::new();
3253 for change in batch_diff.iter_all_changes() {
3254 let line_count = change.value().lines().count() as u32;
3255 match change.tag() {
3256 similar::ChangeTag::Equal => {
3257 old_row += line_count;
3258 new_row += line_count;
3259 }
3260 similar::ChangeTag::Delete => {
3261 let old_end_row = old_row + line_count - 1;
3262 let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
3263
3264 if let Some((_, last_deleted_row_range)) =
3265 deleted_row_ranges.last_mut()
3266 {
3267 if *last_deleted_row_range.end() + 1 == old_row {
3268 *last_deleted_row_range =
3269 *last_deleted_row_range.start()..=old_end_row;
3270 } else {
3271 deleted_row_ranges.push((new_row, old_row..=old_end_row));
3272 }
3273 } else {
3274 deleted_row_ranges.push((new_row, old_row..=old_end_row));
3275 }
3276
3277 old_row += line_count;
3278 }
3279 similar::ChangeTag::Insert => {
3280 let new_end_row = new_row + line_count - 1;
3281 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
3282 let end = new_snapshot.anchor_before(Point::new(
3283 new_end_row,
3284 new_snapshot.line_len(MultiBufferRow(new_end_row)),
3285 ));
3286 inserted_row_ranges.push(start..end);
3287 new_row += line_count;
3288 }
3289 }
3290 }
3291
3292 (deleted_row_ranges, inserted_row_ranges)
3293 })
3294 .await;
3295
3296 codegen
3297 .update(&mut cx, |codegen, cx| {
3298 codegen.diff.deleted_row_ranges = deleted_row_ranges;
3299 codegen.diff.inserted_row_ranges = inserted_row_ranges;
3300 cx.notify();
3301 })
3302 .ok();
3303 })
3304 }
3305}
3306
3307struct StripInvalidSpans<T> {
3308 stream: T,
3309 stream_done: bool,
3310 buffer: String,
3311 first_line: bool,
3312 line_end: bool,
3313 starts_with_code_block: bool,
3314}
3315
3316impl<T> StripInvalidSpans<T>
3317where
3318 T: Stream<Item = Result<String>>,
3319{
3320 fn new(stream: T) -> Self {
3321 Self {
3322 stream,
3323 stream_done: false,
3324 buffer: String::new(),
3325 first_line: true,
3326 line_end: false,
3327 starts_with_code_block: false,
3328 }
3329 }
3330}
3331
3332impl<T> Stream for StripInvalidSpans<T>
3333where
3334 T: Stream<Item = Result<String>>,
3335{
3336 type Item = Result<String>;
3337
3338 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
3339 const CODE_BLOCK_DELIMITER: &str = "```";
3340 const CURSOR_SPAN: &str = "<|CURSOR|>";
3341
3342 let this = unsafe { self.get_unchecked_mut() };
3343 loop {
3344 if !this.stream_done {
3345 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
3346 match stream.as_mut().poll_next(cx) {
3347 Poll::Ready(Some(Ok(chunk))) => {
3348 this.buffer.push_str(&chunk);
3349 }
3350 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
3351 Poll::Ready(None) => {
3352 this.stream_done = true;
3353 }
3354 Poll::Pending => return Poll::Pending,
3355 }
3356 }
3357
3358 let mut chunk = String::new();
3359 let mut consumed = 0;
3360 if !this.buffer.is_empty() {
3361 let mut lines = this.buffer.split('\n').enumerate().peekable();
3362 while let Some((line_ix, line)) = lines.next() {
3363 if line_ix > 0 {
3364 this.first_line = false;
3365 }
3366
3367 if this.first_line {
3368 let trimmed_line = line.trim();
3369 if lines.peek().is_some() {
3370 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
3371 consumed += line.len() + 1;
3372 this.starts_with_code_block = true;
3373 continue;
3374 }
3375 } else if trimmed_line.is_empty()
3376 || prefixes(CODE_BLOCK_DELIMITER)
3377 .any(|prefix| trimmed_line.starts_with(prefix))
3378 {
3379 break;
3380 }
3381 }
3382
3383 let line_without_cursor = line.replace(CURSOR_SPAN, "");
3384 if lines.peek().is_some() {
3385 if this.line_end {
3386 chunk.push('\n');
3387 }
3388
3389 chunk.push_str(&line_without_cursor);
3390 this.line_end = true;
3391 consumed += line.len() + 1;
3392 } else if this.stream_done {
3393 if !this.starts_with_code_block
3394 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
3395 {
3396 if this.line_end {
3397 chunk.push('\n');
3398 }
3399
3400 chunk.push_str(&line);
3401 }
3402
3403 consumed += line.len();
3404 } else {
3405 let trimmed_line = line.trim();
3406 if trimmed_line.is_empty()
3407 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
3408 || prefixes(CODE_BLOCK_DELIMITER)
3409 .any(|prefix| trimmed_line.ends_with(prefix))
3410 {
3411 break;
3412 } else {
3413 if this.line_end {
3414 chunk.push('\n');
3415 this.line_end = false;
3416 }
3417
3418 chunk.push_str(&line_without_cursor);
3419 consumed += line.len();
3420 }
3421 }
3422 }
3423 }
3424
3425 this.buffer = this.buffer.split_off(consumed);
3426 if !chunk.is_empty() {
3427 return Poll::Ready(Some(Ok(chunk)));
3428 } else if this.stream_done {
3429 return Poll::Ready(None);
3430 }
3431 }
3432 }
3433}
3434
3435struct AssistantCodeActionProvider {
3436 editor: WeakView<Editor>,
3437 workspace: WeakView<Workspace>,
3438}
3439
3440impl CodeActionProvider for AssistantCodeActionProvider {
3441 fn code_actions(
3442 &self,
3443 buffer: &Model<Buffer>,
3444 range: Range<text::Anchor>,
3445 cx: &mut WindowContext,
3446 ) -> Task<Result<Vec<CodeAction>>> {
3447 if !AssistantSettings::get_global(cx).enabled {
3448 return Task::ready(Ok(Vec::new()));
3449 }
3450
3451 let snapshot = buffer.read(cx).snapshot();
3452 let mut range = range.to_point(&snapshot);
3453
3454 // Expand the range to line boundaries.
3455 range.start.column = 0;
3456 range.end.column = snapshot.line_len(range.end.row);
3457
3458 let mut has_diagnostics = false;
3459 for diagnostic in snapshot.diagnostics_in_range::<_, Point>(range.clone(), false) {
3460 range.start = cmp::min(range.start, diagnostic.range.start);
3461 range.end = cmp::max(range.end, diagnostic.range.end);
3462 has_diagnostics = true;
3463 }
3464 if has_diagnostics {
3465 if let Some(symbols_containing_start) = snapshot.symbols_containing(range.start, None) {
3466 if let Some(symbol) = symbols_containing_start.last() {
3467 range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
3468 range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
3469 }
3470 }
3471
3472 if let Some(symbols_containing_end) = snapshot.symbols_containing(range.end, None) {
3473 if let Some(symbol) = symbols_containing_end.last() {
3474 range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
3475 range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
3476 }
3477 }
3478
3479 Task::ready(Ok(vec![CodeAction {
3480 server_id: language::LanguageServerId(0),
3481 range: snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end),
3482 lsp_action: lsp::CodeAction {
3483 title: "Fix with Assistant".into(),
3484 ..Default::default()
3485 },
3486 }]))
3487 } else {
3488 Task::ready(Ok(Vec::new()))
3489 }
3490 }
3491
3492 fn apply_code_action(
3493 &self,
3494 buffer: Model<Buffer>,
3495 action: CodeAction,
3496 excerpt_id: ExcerptId,
3497 _push_to_history: bool,
3498 cx: &mut WindowContext,
3499 ) -> Task<Result<ProjectTransaction>> {
3500 let editor = self.editor.clone();
3501 let workspace = self.workspace.clone();
3502 cx.spawn(|mut cx| async move {
3503 let editor = editor.upgrade().context("editor was released")?;
3504 let range = editor
3505 .update(&mut cx, |editor, cx| {
3506 editor.buffer().update(cx, |multibuffer, cx| {
3507 let buffer = buffer.read(cx);
3508 let multibuffer_snapshot = multibuffer.read(cx);
3509
3510 let old_context_range =
3511 multibuffer_snapshot.context_range_for_excerpt(excerpt_id)?;
3512 let mut new_context_range = old_context_range.clone();
3513 if action
3514 .range
3515 .start
3516 .cmp(&old_context_range.start, buffer)
3517 .is_lt()
3518 {
3519 new_context_range.start = action.range.start;
3520 }
3521 if action.range.end.cmp(&old_context_range.end, buffer).is_gt() {
3522 new_context_range.end = action.range.end;
3523 }
3524 drop(multibuffer_snapshot);
3525
3526 if new_context_range != old_context_range {
3527 multibuffer.resize_excerpt(excerpt_id, new_context_range, cx);
3528 }
3529
3530 let multibuffer_snapshot = multibuffer.read(cx);
3531 Some(
3532 multibuffer_snapshot
3533 .anchor_in_excerpt(excerpt_id, action.range.start)?
3534 ..multibuffer_snapshot
3535 .anchor_in_excerpt(excerpt_id, action.range.end)?,
3536 )
3537 })
3538 })?
3539 .context("invalid range")?;
3540 let assistant_panel = workspace.update(&mut cx, |workspace, cx| {
3541 workspace
3542 .panel::<AssistantPanel>(cx)
3543 .context("assistant panel was released")
3544 })??;
3545
3546 cx.update_global(|assistant: &mut InlineAssistant, cx| {
3547 let assist_id = assistant.suggest_assist(
3548 &editor,
3549 range,
3550 "Fix Diagnostics".into(),
3551 None,
3552 true,
3553 Some(workspace),
3554 Some(&assistant_panel),
3555 cx,
3556 );
3557 assistant.start_assist(assist_id, cx);
3558 })?;
3559
3560 Ok(ProjectTransaction::default())
3561 })
3562 }
3563}
3564
3565fn prefixes(text: &str) -> impl Iterator<Item = &str> {
3566 (0..text.len() - 1).map(|ix| &text[..ix + 1])
3567}
3568
3569fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
3570 ranges.sort_unstable_by(|a, b| {
3571 a.start
3572 .cmp(&b.start, buffer)
3573 .then_with(|| b.end.cmp(&a.end, buffer))
3574 });
3575
3576 let mut ix = 0;
3577 while ix + 1 < ranges.len() {
3578 let b = ranges[ix + 1].clone();
3579 let a = &mut ranges[ix];
3580 if a.end.cmp(&b.start, buffer).is_gt() {
3581 if a.end.cmp(&b.end, buffer).is_lt() {
3582 a.end = b.end;
3583 }
3584 ranges.remove(ix + 1);
3585 } else {
3586 ix += 1;
3587 }
3588 }
3589}
3590
3591#[cfg(test)]
3592mod tests {
3593 use super::*;
3594 use futures::stream::{self};
3595 use gpui::{Context, TestAppContext};
3596 use indoc::indoc;
3597 use language::{
3598 language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
3599 Point,
3600 };
3601 use language_model::LanguageModelRegistry;
3602 use rand::prelude::*;
3603 use serde::Serialize;
3604 use settings::SettingsStore;
3605 use std::{future, sync::Arc};
3606
3607 #[derive(Serialize)]
3608 pub struct DummyCompletionRequest {
3609 pub name: String,
3610 }
3611
3612 #[gpui::test(iterations = 10)]
3613 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
3614 cx.set_global(cx.update(SettingsStore::test));
3615 cx.update(language_model::LanguageModelRegistry::test);
3616 cx.update(language_settings::init);
3617
3618 let text = indoc! {"
3619 fn main() {
3620 let x = 0;
3621 for _ in 0..10 {
3622 x += 1;
3623 }
3624 }
3625 "};
3626 let buffer =
3627 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3628 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
3629 let range = buffer.read_with(cx, |buffer, cx| {
3630 let snapshot = buffer.snapshot(cx);
3631 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
3632 });
3633 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3634 let codegen = cx.new_model(|cx| {
3635 CodegenAlternative::new(
3636 buffer.clone(),
3637 range.clone(),
3638 true,
3639 None,
3640 prompt_builder,
3641 cx,
3642 )
3643 });
3644
3645 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3646
3647 let mut new_text = concat!(
3648 " let mut x = 0;\n",
3649 " while x < 10 {\n",
3650 " x += 1;\n",
3651 " }",
3652 );
3653 while !new_text.is_empty() {
3654 let max_len = cmp::min(new_text.len(), 10);
3655 let len = rng.gen_range(1..=max_len);
3656 let (chunk, suffix) = new_text.split_at(len);
3657 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
3658 new_text = suffix;
3659 cx.background_executor.run_until_parked();
3660 }
3661 drop(chunks_tx);
3662 cx.background_executor.run_until_parked();
3663
3664 assert_eq!(
3665 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3666 indoc! {"
3667 fn main() {
3668 let mut x = 0;
3669 while x < 10 {
3670 x += 1;
3671 }
3672 }
3673 "}
3674 );
3675 }
3676
3677 #[gpui::test(iterations = 10)]
3678 async fn test_autoindent_when_generating_past_indentation(
3679 cx: &mut TestAppContext,
3680 mut rng: StdRng,
3681 ) {
3682 cx.set_global(cx.update(SettingsStore::test));
3683 cx.update(language_settings::init);
3684
3685 let text = indoc! {"
3686 fn main() {
3687 le
3688 }
3689 "};
3690 let buffer =
3691 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3692 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
3693 let range = buffer.read_with(cx, |buffer, cx| {
3694 let snapshot = buffer.snapshot(cx);
3695 snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
3696 });
3697 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3698 let codegen = cx.new_model(|cx| {
3699 CodegenAlternative::new(
3700 buffer.clone(),
3701 range.clone(),
3702 true,
3703 None,
3704 prompt_builder,
3705 cx,
3706 )
3707 });
3708
3709 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3710
3711 cx.background_executor.run_until_parked();
3712
3713 let mut new_text = concat!(
3714 "t mut x = 0;\n",
3715 "while x < 10 {\n",
3716 " x += 1;\n",
3717 "}", //
3718 );
3719 while !new_text.is_empty() {
3720 let max_len = cmp::min(new_text.len(), 10);
3721 let len = rng.gen_range(1..=max_len);
3722 let (chunk, suffix) = new_text.split_at(len);
3723 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
3724 new_text = suffix;
3725 cx.background_executor.run_until_parked();
3726 }
3727 drop(chunks_tx);
3728 cx.background_executor.run_until_parked();
3729
3730 assert_eq!(
3731 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3732 indoc! {"
3733 fn main() {
3734 let mut x = 0;
3735 while x < 10 {
3736 x += 1;
3737 }
3738 }
3739 "}
3740 );
3741 }
3742
3743 #[gpui::test(iterations = 10)]
3744 async fn test_autoindent_when_generating_before_indentation(
3745 cx: &mut TestAppContext,
3746 mut rng: StdRng,
3747 ) {
3748 cx.update(LanguageModelRegistry::test);
3749 cx.set_global(cx.update(SettingsStore::test));
3750 cx.update(language_settings::init);
3751
3752 let text = concat!(
3753 "fn main() {\n",
3754 " \n",
3755 "}\n" //
3756 );
3757 let buffer =
3758 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3759 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
3760 let range = buffer.read_with(cx, |buffer, cx| {
3761 let snapshot = buffer.snapshot(cx);
3762 snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
3763 });
3764 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3765 let codegen = cx.new_model(|cx| {
3766 CodegenAlternative::new(
3767 buffer.clone(),
3768 range.clone(),
3769 true,
3770 None,
3771 prompt_builder,
3772 cx,
3773 )
3774 });
3775
3776 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3777
3778 cx.background_executor.run_until_parked();
3779
3780 let mut new_text = concat!(
3781 "let mut x = 0;\n",
3782 "while x < 10 {\n",
3783 " x += 1;\n",
3784 "}", //
3785 );
3786 while !new_text.is_empty() {
3787 let max_len = cmp::min(new_text.len(), 10);
3788 let len = rng.gen_range(1..=max_len);
3789 let (chunk, suffix) = new_text.split_at(len);
3790 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
3791 new_text = suffix;
3792 cx.background_executor.run_until_parked();
3793 }
3794 drop(chunks_tx);
3795 cx.background_executor.run_until_parked();
3796
3797 assert_eq!(
3798 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3799 indoc! {"
3800 fn main() {
3801 let mut x = 0;
3802 while x < 10 {
3803 x += 1;
3804 }
3805 }
3806 "}
3807 );
3808 }
3809
3810 #[gpui::test(iterations = 10)]
3811 async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
3812 cx.update(LanguageModelRegistry::test);
3813 cx.set_global(cx.update(SettingsStore::test));
3814 cx.update(language_settings::init);
3815
3816 let text = indoc! {"
3817 func main() {
3818 \tx := 0
3819 \tfor i := 0; i < 10; i++ {
3820 \t\tx++
3821 \t}
3822 }
3823 "};
3824 let buffer = cx.new_model(|cx| Buffer::local(text, cx));
3825 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
3826 let range = buffer.read_with(cx, |buffer, cx| {
3827 let snapshot = buffer.snapshot(cx);
3828 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
3829 });
3830 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3831 let codegen = cx.new_model(|cx| {
3832 CodegenAlternative::new(
3833 buffer.clone(),
3834 range.clone(),
3835 true,
3836 None,
3837 prompt_builder,
3838 cx,
3839 )
3840 });
3841
3842 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3843 let new_text = concat!(
3844 "func main() {\n",
3845 "\tx := 0\n",
3846 "\tfor x < 10 {\n",
3847 "\t\tx++\n",
3848 "\t}", //
3849 );
3850 chunks_tx.unbounded_send(new_text.to_string()).unwrap();
3851 drop(chunks_tx);
3852 cx.background_executor.run_until_parked();
3853
3854 assert_eq!(
3855 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3856 indoc! {"
3857 func main() {
3858 \tx := 0
3859 \tfor x < 10 {
3860 \t\tx++
3861 \t}
3862 }
3863 "}
3864 );
3865 }
3866
3867 #[gpui::test]
3868 async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
3869 cx.update(LanguageModelRegistry::test);
3870 cx.set_global(cx.update(SettingsStore::test));
3871 cx.update(language_settings::init);
3872
3873 let text = indoc! {"
3874 fn main() {
3875 let x = 0;
3876 }
3877 "};
3878 let buffer =
3879 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3880 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
3881 let range = buffer.read_with(cx, |buffer, cx| {
3882 let snapshot = buffer.snapshot(cx);
3883 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
3884 });
3885 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3886 let codegen = cx.new_model(|cx| {
3887 CodegenAlternative::new(
3888 buffer.clone(),
3889 range.clone(),
3890 false,
3891 None,
3892 prompt_builder,
3893 cx,
3894 )
3895 });
3896
3897 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3898 chunks_tx
3899 .unbounded_send("let mut x = 0;\nx += 1;".to_string())
3900 .unwrap();
3901 drop(chunks_tx);
3902 cx.run_until_parked();
3903
3904 // The codegen is inactive, so the buffer doesn't get modified.
3905 assert_eq!(
3906 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3907 text
3908 );
3909
3910 // Activating the codegen applies the changes.
3911 codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
3912 assert_eq!(
3913 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3914 indoc! {"
3915 fn main() {
3916 let mut x = 0;
3917 x += 1;
3918 }
3919 "}
3920 );
3921
3922 // Deactivating the codegen undoes the changes.
3923 codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
3924 cx.run_until_parked();
3925 assert_eq!(
3926 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3927 text
3928 );
3929 }
3930
3931 #[gpui::test]
3932 async fn test_strip_invalid_spans_from_codeblock() {
3933 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
3934 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
3935 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
3936 assert_chunks(
3937 "```html\n```js\nLorem ipsum dolor\n```\n```",
3938 "```js\nLorem ipsum dolor\n```",
3939 )
3940 .await;
3941 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
3942 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
3943 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
3944 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
3945
3946 async fn assert_chunks(text: &str, expected_text: &str) {
3947 for chunk_size in 1..=text.len() {
3948 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
3949 .map(|chunk| chunk.unwrap())
3950 .collect::<String>()
3951 .await;
3952 assert_eq!(
3953 actual_text, expected_text,
3954 "failed to strip invalid spans, chunk size: {}",
3955 chunk_size
3956 );
3957 }
3958 }
3959
3960 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
3961 stream::iter(
3962 text.chars()
3963 .collect::<Vec<_>>()
3964 .chunks(size)
3965 .map(|chunk| Ok(chunk.iter().collect::<String>()))
3966 .collect::<Vec<_>>(),
3967 )
3968 }
3969 }
3970
3971 fn simulate_response_stream(
3972 codegen: Model<CodegenAlternative>,
3973 cx: &mut TestAppContext,
3974 ) -> mpsc::UnboundedSender<String> {
3975 let (chunks_tx, chunks_rx) = mpsc::unbounded();
3976 codegen.update(cx, |codegen, cx| {
3977 codegen.handle_stream(
3978 String::new(),
3979 String::new(),
3980 None,
3981 future::ready(Ok(LanguageModelTextStream {
3982 message_id: None,
3983 stream: chunks_rx.map(Ok).boxed(),
3984 })),
3985 cx,
3986 );
3987 });
3988 chunks_tx
3989 }
3990
3991 fn rust_lang() -> Language {
3992 Language::new(
3993 LanguageConfig {
3994 name: "Rust".into(),
3995 matcher: LanguageMatcher {
3996 path_suffixes: vec!["rs".to_string()],
3997 ..Default::default()
3998 },
3999 ..Default::default()
4000 },
4001 Some(tree_sitter_rust::LANGUAGE.into()),
4002 )
4003 .with_indents_query(
4004 r#"
4005 (call_expression) @indent
4006 (field_expression) @indent
4007 (_ "(" ")" @end) @indent
4008 (_ "{" "}" @end) @indent
4009 "#,
4010 )
4011 .unwrap()
4012 }
4013}