1use crate::{
2 humanize_token_count, prompts::generate_content_prompt, AssistantPanel, AssistantPanelEvent,
3 Hunk, LanguageModelCompletionProvider, ModelSelector, StreamingDiff,
4};
5use anyhow::{anyhow, Context as _, Result};
6use client::telemetry::Telemetry;
7use collections::{hash_map, HashMap, HashSet, VecDeque};
8use editor::{
9 actions::{MoveDown, MoveUp, SelectAll},
10 display_map::{
11 BlockContext, BlockDisposition, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,
12 ToDisplayPoint,
13 },
14 Anchor, AnchorRangeExt, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle,
15 ExcerptRange, GutterDimensions, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint,
16};
17use fs::Fs;
18use futures::{
19 channel::mpsc,
20 future::LocalBoxFuture,
21 stream::{self, BoxStream},
22 SinkExt, Stream, StreamExt,
23};
24use gpui::{
25 point, AppContext, EventEmitter, FocusHandle, FocusableView, Global, HighlightStyle, Model,
26 ModelContext, Subscription, Task, TextStyle, UpdateGlobal, View, ViewContext, WeakView,
27 WindowContext,
28};
29use language::{Buffer, IndentKind, Point, Selection, TransactionId};
30use language_model::{LanguageModelRequest, LanguageModelRequestMessage, Role};
31use multi_buffer::MultiBufferRow;
32use parking_lot::Mutex;
33use rope::Rope;
34use settings::Settings;
35use similar::TextDiff;
36use smol::future::FutureExt;
37use std::{
38 cmp,
39 future::Future,
40 mem,
41 ops::{Range, RangeInclusive},
42 pin::Pin,
43 sync::Arc,
44 task::{self, Poll},
45 time::{Duration, Instant},
46};
47use theme::ThemeSettings;
48use ui::{prelude::*, IconButtonShape, Tooltip};
49use util::RangeExt;
50use workspace::{notifications::NotificationId, Toast, Workspace};
51
52pub fn init(fs: Arc<dyn Fs>, telemetry: Arc<Telemetry>, cx: &mut AppContext) {
53 cx.set_global(InlineAssistant::new(fs, telemetry));
54}
55
56const PROMPT_HISTORY_MAX_LEN: usize = 20;
57
58pub struct InlineAssistant {
59 next_assist_id: InlineAssistId,
60 next_assist_group_id: InlineAssistGroupId,
61 assists: HashMap<InlineAssistId, InlineAssist>,
62 assists_by_editor: HashMap<WeakView<Editor>, EditorInlineAssists>,
63 assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
64 prompt_history: VecDeque<String>,
65 telemetry: Option<Arc<Telemetry>>,
66 fs: Arc<dyn Fs>,
67}
68
69impl Global for InlineAssistant {}
70
71impl InlineAssistant {
72 pub fn new(fs: Arc<dyn Fs>, telemetry: Arc<Telemetry>) -> Self {
73 Self {
74 next_assist_id: InlineAssistId::default(),
75 next_assist_group_id: InlineAssistGroupId::default(),
76 assists: HashMap::default(),
77 assists_by_editor: HashMap::default(),
78 assist_groups: HashMap::default(),
79 prompt_history: VecDeque::default(),
80 telemetry: Some(telemetry),
81 fs,
82 }
83 }
84
85 pub fn assist(
86 &mut self,
87 editor: &View<Editor>,
88 workspace: Option<WeakView<Workspace>>,
89 assistant_panel: Option<&View<AssistantPanel>>,
90 initial_prompt: Option<String>,
91 cx: &mut WindowContext,
92 ) {
93 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
94
95 let mut selections = Vec::<Selection<Point>>::new();
96 let mut newest_selection = None;
97 for mut selection in editor.read(cx).selections.all::<Point>(cx) {
98 if selection.end > selection.start {
99 selection.start.column = 0;
100 // If the selection ends at the start of the line, we don't want to include it.
101 if selection.end.column == 0 {
102 selection.end.row -= 1;
103 }
104 selection.end.column = snapshot.line_len(MultiBufferRow(selection.end.row));
105 }
106
107 if let Some(prev_selection) = selections.last_mut() {
108 if selection.start <= prev_selection.end {
109 prev_selection.end = selection.end;
110 continue;
111 }
112 }
113
114 let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
115 if selection.id > latest_selection.id {
116 *latest_selection = selection.clone();
117 }
118 selections.push(selection);
119 }
120 let newest_selection = newest_selection.unwrap();
121
122 let mut codegen_ranges = Vec::new();
123 for (excerpt_id, buffer, buffer_range) in
124 snapshot.excerpts_in_ranges(selections.iter().map(|selection| {
125 snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
126 }))
127 {
128 let start = Anchor {
129 buffer_id: Some(buffer.remote_id()),
130 excerpt_id,
131 text_anchor: buffer.anchor_before(buffer_range.start),
132 };
133 let end = Anchor {
134 buffer_id: Some(buffer.remote_id()),
135 excerpt_id,
136 text_anchor: buffer.anchor_after(buffer_range.end),
137 };
138 codegen_ranges.push(start..end);
139 }
140
141 let assist_group_id = self.next_assist_group_id.post_inc();
142 let prompt_buffer =
143 cx.new_model(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx));
144 let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
145
146 let mut assists = Vec::new();
147 let mut assist_to_focus = None;
148 for range in codegen_ranges {
149 let assist_id = self.next_assist_id.post_inc();
150 let codegen = cx.new_model(|cx| {
151 Codegen::new(
152 editor.read(cx).buffer().clone(),
153 range.clone(),
154 None,
155 self.telemetry.clone(),
156 cx,
157 )
158 });
159
160 let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
161 let prompt_editor = cx.new_view(|cx| {
162 PromptEditor::new(
163 assist_id,
164 gutter_dimensions.clone(),
165 self.prompt_history.clone(),
166 prompt_buffer.clone(),
167 codegen.clone(),
168 editor,
169 assistant_panel,
170 workspace.clone(),
171 self.fs.clone(),
172 cx,
173 )
174 });
175
176 if assist_to_focus.is_none() {
177 let focus_assist = if newest_selection.reversed {
178 range.start.to_point(&snapshot) == newest_selection.start
179 } else {
180 range.end.to_point(&snapshot) == newest_selection.end
181 };
182 if focus_assist {
183 assist_to_focus = Some(assist_id);
184 }
185 }
186
187 let [prompt_block_id, end_block_id] =
188 self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
189
190 assists.push((assist_id, prompt_editor, prompt_block_id, end_block_id));
191 }
192
193 let editor_assists = self
194 .assists_by_editor
195 .entry(editor.downgrade())
196 .or_insert_with(|| EditorInlineAssists::new(&editor, cx));
197 let mut assist_group = InlineAssistGroup::new();
198 for (assist_id, prompt_editor, prompt_block_id, end_block_id) in assists {
199 self.assists.insert(
200 assist_id,
201 InlineAssist::new(
202 assist_id,
203 assist_group_id,
204 assistant_panel.is_some(),
205 editor,
206 &prompt_editor,
207 prompt_block_id,
208 end_block_id,
209 prompt_editor.read(cx).codegen.clone(),
210 workspace.clone(),
211 cx,
212 ),
213 );
214 assist_group.assist_ids.push(assist_id);
215 editor_assists.assist_ids.push(assist_id);
216 }
217 self.assist_groups.insert(assist_group_id, assist_group);
218
219 if let Some(assist_id) = assist_to_focus {
220 self.focus_assist(assist_id, cx);
221 }
222 }
223
224 #[allow(clippy::too_many_arguments)]
225 pub fn suggest_assist(
226 &mut self,
227 editor: &View<Editor>,
228 mut range: Range<Anchor>,
229 initial_prompt: String,
230 initial_insertion: Option<String>,
231 workspace: Option<WeakView<Workspace>>,
232 assistant_panel: Option<&View<AssistantPanel>>,
233 cx: &mut WindowContext,
234 ) -> InlineAssistId {
235 let assist_group_id = self.next_assist_group_id.post_inc();
236 let prompt_buffer = cx.new_model(|cx| Buffer::local(&initial_prompt, cx));
237 let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
238
239 let assist_id = self.next_assist_id.post_inc();
240
241 let buffer = editor.read(cx).buffer().clone();
242 let prepend_transaction_id = initial_insertion.and_then(|initial_insertion| {
243 buffer.update(cx, |buffer, cx| {
244 buffer.start_transaction(cx);
245 buffer.edit([(range.start..range.start, initial_insertion)], None, cx);
246 buffer.end_transaction(cx)
247 })
248 });
249
250 range.start = range.start.bias_left(&buffer.read(cx).read(cx));
251 range.end = range.end.bias_right(&buffer.read(cx).read(cx));
252
253 let codegen = cx.new_model(|cx| {
254 Codegen::new(
255 editor.read(cx).buffer().clone(),
256 range.clone(),
257 prepend_transaction_id,
258 self.telemetry.clone(),
259 cx,
260 )
261 });
262
263 let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
264 let prompt_editor = cx.new_view(|cx| {
265 PromptEditor::new(
266 assist_id,
267 gutter_dimensions.clone(),
268 self.prompt_history.clone(),
269 prompt_buffer.clone(),
270 codegen.clone(),
271 editor,
272 assistant_panel,
273 workspace.clone(),
274 self.fs.clone(),
275 cx,
276 )
277 });
278
279 let [prompt_block_id, end_block_id] =
280 self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
281
282 let editor_assists = self
283 .assists_by_editor
284 .entry(editor.downgrade())
285 .or_insert_with(|| EditorInlineAssists::new(&editor, cx));
286
287 let mut assist_group = InlineAssistGroup::new();
288 self.assists.insert(
289 assist_id,
290 InlineAssist::new(
291 assist_id,
292 assist_group_id,
293 assistant_panel.is_some(),
294 editor,
295 &prompt_editor,
296 prompt_block_id,
297 end_block_id,
298 prompt_editor.read(cx).codegen.clone(),
299 workspace.clone(),
300 cx,
301 ),
302 );
303 assist_group.assist_ids.push(assist_id);
304 editor_assists.assist_ids.push(assist_id);
305 self.assist_groups.insert(assist_group_id, assist_group);
306 assist_id
307 }
308
309 fn insert_assist_blocks(
310 &self,
311 editor: &View<Editor>,
312 range: &Range<Anchor>,
313 prompt_editor: &View<PromptEditor>,
314 cx: &mut WindowContext,
315 ) -> [CustomBlockId; 2] {
316 let assist_blocks = vec![
317 BlockProperties {
318 style: BlockStyle::Sticky,
319 position: range.start,
320 height: prompt_editor.read(cx).height_in_lines,
321 render: build_assist_editor_renderer(prompt_editor),
322 disposition: BlockDisposition::Above,
323 },
324 BlockProperties {
325 style: BlockStyle::Sticky,
326 position: range.end,
327 height: 1,
328 render: Box::new(|cx| {
329 v_flex()
330 .h_full()
331 .w_full()
332 .border_t_1()
333 .border_color(cx.theme().status().info_border)
334 .into_any_element()
335 }),
336 disposition: BlockDisposition::Below,
337 },
338 ];
339
340 editor.update(cx, |editor, cx| {
341 let block_ids = editor.insert_blocks(assist_blocks, None, cx);
342 [block_ids[0], block_ids[1]]
343 })
344 }
345
346 fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
347 let assist = &self.assists[&assist_id];
348 let Some(decorations) = assist.decorations.as_ref() else {
349 return;
350 };
351 let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
352 let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
353
354 assist_group.active_assist_id = Some(assist_id);
355 if assist_group.linked {
356 for assist_id in &assist_group.assist_ids {
357 if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
358 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
359 prompt_editor.set_show_cursor_when_unfocused(true, cx)
360 });
361 }
362 }
363 }
364
365 assist
366 .editor
367 .update(cx, |editor, cx| {
368 let scroll_top = editor.scroll_position(cx).y;
369 let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
370 let prompt_row = editor
371 .row_for_block(decorations.prompt_block_id, cx)
372 .unwrap()
373 .0 as f32;
374
375 if (scroll_top..scroll_bottom).contains(&prompt_row) {
376 editor_assists.scroll_lock = Some(InlineAssistScrollLock {
377 assist_id,
378 distance_from_top: prompt_row - scroll_top,
379 });
380 } else {
381 editor_assists.scroll_lock = None;
382 }
383 })
384 .ok();
385 }
386
387 fn handle_prompt_editor_focus_out(
388 &mut self,
389 assist_id: InlineAssistId,
390 cx: &mut WindowContext,
391 ) {
392 let assist = &self.assists[&assist_id];
393 let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
394 if assist_group.active_assist_id == Some(assist_id) {
395 assist_group.active_assist_id = None;
396 if assist_group.linked {
397 for assist_id in &assist_group.assist_ids {
398 if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
399 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
400 prompt_editor.set_show_cursor_when_unfocused(false, cx)
401 });
402 }
403 }
404 }
405 }
406 }
407
408 fn handle_prompt_editor_event(
409 &mut self,
410 prompt_editor: View<PromptEditor>,
411 event: &PromptEditorEvent,
412 cx: &mut WindowContext,
413 ) {
414 let assist_id = prompt_editor.read(cx).id;
415 match event {
416 PromptEditorEvent::StartRequested => {
417 self.start_assist(assist_id, cx);
418 }
419 PromptEditorEvent::StopRequested => {
420 self.stop_assist(assist_id, cx);
421 }
422 PromptEditorEvent::ConfirmRequested => {
423 self.finish_assist(assist_id, false, cx);
424 }
425 PromptEditorEvent::CancelRequested => {
426 self.finish_assist(assist_id, true, cx);
427 }
428 PromptEditorEvent::DismissRequested => {
429 self.dismiss_assist(assist_id, cx);
430 }
431 PromptEditorEvent::Resized { height_in_lines } => {
432 self.resize_assist(assist_id, *height_in_lines, cx);
433 }
434 }
435 }
436
437 fn handle_editor_newline(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
438 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
439 return;
440 };
441
442 let editor = editor.read(cx);
443 if editor.selections.count() == 1 {
444 let selection = editor.selections.newest::<usize>(cx);
445 let buffer = editor.buffer().read(cx).snapshot(cx);
446 for assist_id in &editor_assists.assist_ids {
447 let assist = &self.assists[assist_id];
448 let assist_range = assist.codegen.read(cx).range.to_offset(&buffer);
449 if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
450 {
451 if matches!(assist.codegen.read(cx).status, CodegenStatus::Pending) {
452 self.dismiss_assist(*assist_id, cx);
453 } else {
454 self.finish_assist(*assist_id, false, cx);
455 }
456
457 return;
458 }
459 }
460 }
461
462 cx.propagate();
463 }
464
465 fn handle_editor_cancel(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
466 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
467 return;
468 };
469
470 let editor = editor.read(cx);
471 if editor.selections.count() == 1 {
472 let selection = editor.selections.newest::<usize>(cx);
473 let buffer = editor.buffer().read(cx).snapshot(cx);
474 for assist_id in &editor_assists.assist_ids {
475 let assist = &self.assists[assist_id];
476 let assist_range = assist.codegen.read(cx).range.to_offset(&buffer);
477 if assist.decorations.is_some()
478 && assist_range.contains(&selection.start)
479 && assist_range.contains(&selection.end)
480 {
481 self.focus_assist(*assist_id, cx);
482 return;
483 }
484 }
485 }
486
487 cx.propagate();
488 }
489
490 fn handle_editor_release(&mut self, editor: WeakView<Editor>, cx: &mut WindowContext) {
491 if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor) {
492 for assist_id in editor_assists.assist_ids.clone() {
493 self.finish_assist(assist_id, true, cx);
494 }
495 }
496 }
497
498 fn handle_editor_change(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
499 let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
500 return;
501 };
502 let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
503 return;
504 };
505 let assist = &self.assists[&scroll_lock.assist_id];
506 let Some(decorations) = assist.decorations.as_ref() else {
507 return;
508 };
509
510 editor.update(cx, |editor, cx| {
511 let scroll_position = editor.scroll_position(cx);
512 let target_scroll_top = editor
513 .row_for_block(decorations.prompt_block_id, cx)
514 .unwrap()
515 .0 as f32
516 - scroll_lock.distance_from_top;
517 if target_scroll_top != scroll_position.y {
518 editor.set_scroll_position(point(scroll_position.x, target_scroll_top), cx);
519 }
520 });
521 }
522
523 fn handle_editor_event(
524 &mut self,
525 editor: View<Editor>,
526 event: &EditorEvent,
527 cx: &mut WindowContext,
528 ) {
529 let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
530 return;
531 };
532
533 match event {
534 EditorEvent::Saved => {
535 for assist_id in editor_assists.assist_ids.clone() {
536 let assist = &self.assists[&assist_id];
537 if let CodegenStatus::Done = &assist.codegen.read(cx).status {
538 self.finish_assist(assist_id, false, cx)
539 }
540 }
541 }
542 EditorEvent::Edited { transaction_id } => {
543 let buffer = editor.read(cx).buffer().read(cx);
544 let edited_ranges =
545 buffer.edited_ranges_for_transaction::<usize>(*transaction_id, cx);
546 let snapshot = buffer.snapshot(cx);
547
548 for assist_id in editor_assists.assist_ids.clone() {
549 let assist = &self.assists[&assist_id];
550 if matches!(
551 assist.codegen.read(cx).status,
552 CodegenStatus::Error(_) | CodegenStatus::Done
553 ) {
554 let assist_range = assist.codegen.read(cx).range.to_offset(&snapshot);
555 if edited_ranges
556 .iter()
557 .any(|range| range.overlaps(&assist_range))
558 {
559 self.finish_assist(assist_id, false, cx);
560 }
561 }
562 }
563 }
564 EditorEvent::ScrollPositionChanged { .. } => {
565 if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
566 let assist = &self.assists[&scroll_lock.assist_id];
567 if let Some(decorations) = assist.decorations.as_ref() {
568 let distance_from_top = editor.update(cx, |editor, cx| {
569 let scroll_top = editor.scroll_position(cx).y;
570 let prompt_row = editor
571 .row_for_block(decorations.prompt_block_id, cx)
572 .unwrap()
573 .0 as f32;
574 prompt_row - scroll_top
575 });
576
577 if distance_from_top != scroll_lock.distance_from_top {
578 editor_assists.scroll_lock = None;
579 }
580 }
581 }
582 }
583 EditorEvent::SelectionsChanged { .. } => {
584 for assist_id in editor_assists.assist_ids.clone() {
585 let assist = &self.assists[&assist_id];
586 if let Some(decorations) = assist.decorations.as_ref() {
587 if decorations.prompt_editor.focus_handle(cx).is_focused(cx) {
588 return;
589 }
590 }
591 }
592
593 editor_assists.scroll_lock = None;
594 }
595 _ => {}
596 }
597 }
598
599 fn finish_assist(&mut self, assist_id: InlineAssistId, undo: bool, cx: &mut WindowContext) {
600 if let Some(assist) = self.assists.get(&assist_id) {
601 let assist_group_id = assist.group_id;
602 if self.assist_groups[&assist_group_id].linked {
603 for assist_id in self.unlink_assist_group(assist_group_id, cx) {
604 self.finish_assist(assist_id, undo, cx);
605 }
606 return;
607 }
608 }
609
610 self.dismiss_assist(assist_id, cx);
611
612 if let Some(assist) = self.assists.remove(&assist_id) {
613 if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
614 {
615 entry.get_mut().assist_ids.retain(|id| *id != assist_id);
616 if entry.get().assist_ids.is_empty() {
617 entry.remove();
618 }
619 }
620
621 if let hash_map::Entry::Occupied(mut entry) =
622 self.assists_by_editor.entry(assist.editor.clone())
623 {
624 entry.get_mut().assist_ids.retain(|id| *id != assist_id);
625 if entry.get().assist_ids.is_empty() {
626 entry.remove();
627 if let Some(editor) = assist.editor.upgrade() {
628 self.update_editor_highlights(&editor, cx);
629 }
630 } else {
631 entry.get().highlight_updates.send(()).ok();
632 }
633 }
634
635 if undo {
636 assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
637 }
638 }
639 }
640
641 fn dismiss_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) -> bool {
642 let Some(assist) = self.assists.get_mut(&assist_id) else {
643 return false;
644 };
645 let Some(editor) = assist.editor.upgrade() else {
646 return false;
647 };
648 let Some(decorations) = assist.decorations.take() else {
649 return false;
650 };
651
652 editor.update(cx, |editor, cx| {
653 let mut to_remove = decorations.removed_line_block_ids;
654 to_remove.insert(decorations.prompt_block_id);
655 to_remove.insert(decorations.end_block_id);
656 editor.remove_blocks(to_remove, None, cx);
657 });
658
659 if decorations
660 .prompt_editor
661 .focus_handle(cx)
662 .contains_focused(cx)
663 {
664 self.focus_next_assist(assist_id, cx);
665 }
666
667 if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
668 if editor_assists
669 .scroll_lock
670 .as_ref()
671 .map_or(false, |lock| lock.assist_id == assist_id)
672 {
673 editor_assists.scroll_lock = None;
674 }
675 editor_assists.highlight_updates.send(()).ok();
676 }
677
678 true
679 }
680
681 fn focus_next_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
682 let Some(assist) = self.assists.get(&assist_id) else {
683 return;
684 };
685
686 let assist_group = &self.assist_groups[&assist.group_id];
687 let assist_ix = assist_group
688 .assist_ids
689 .iter()
690 .position(|id| *id == assist_id)
691 .unwrap();
692 let assist_ids = assist_group
693 .assist_ids
694 .iter()
695 .skip(assist_ix + 1)
696 .chain(assist_group.assist_ids.iter().take(assist_ix));
697
698 for assist_id in assist_ids {
699 let assist = &self.assists[assist_id];
700 if assist.decorations.is_some() {
701 self.focus_assist(*assist_id, cx);
702 return;
703 }
704 }
705
706 assist.editor.update(cx, |editor, cx| editor.focus(cx)).ok();
707 }
708
709 fn focus_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
710 let assist = &self.assists[&assist_id];
711 let Some(editor) = assist.editor.upgrade() else {
712 return;
713 };
714
715 if let Some(decorations) = assist.decorations.as_ref() {
716 decorations.prompt_editor.update(cx, |prompt_editor, cx| {
717 prompt_editor.editor.update(cx, |editor, cx| {
718 editor.focus(cx);
719 editor.select_all(&SelectAll, cx);
720 })
721 });
722 }
723
724 let position = assist.codegen.read(cx).range.start;
725 editor.update(cx, |editor, cx| {
726 editor.change_selections(None, cx, |selections| {
727 selections.select_anchor_ranges([position..position])
728 });
729
730 let mut scroll_target_top;
731 let mut scroll_target_bottom;
732 if let Some(decorations) = assist.decorations.as_ref() {
733 scroll_target_top = editor
734 .row_for_block(decorations.prompt_block_id, cx)
735 .unwrap()
736 .0 as f32;
737 scroll_target_bottom = editor
738 .row_for_block(decorations.end_block_id, cx)
739 .unwrap()
740 .0 as f32;
741 } else {
742 let snapshot = editor.snapshot(cx);
743 let codegen = assist.codegen.read(cx);
744 let start_row = codegen
745 .range
746 .start
747 .to_display_point(&snapshot.display_snapshot)
748 .row();
749 scroll_target_top = start_row.0 as f32;
750 scroll_target_bottom = scroll_target_top + 1.;
751 }
752 scroll_target_top -= editor.vertical_scroll_margin() as f32;
753 scroll_target_bottom += editor.vertical_scroll_margin() as f32;
754
755 let height_in_lines = editor.visible_line_count().unwrap_or(0.);
756 let scroll_top = editor.scroll_position(cx).y;
757 let scroll_bottom = scroll_top + height_in_lines;
758
759 if scroll_target_top < scroll_top {
760 editor.set_scroll_position(point(0., scroll_target_top), cx);
761 } else if scroll_target_bottom > scroll_bottom {
762 if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
763 editor
764 .set_scroll_position(point(0., scroll_target_bottom - height_in_lines), cx);
765 } else {
766 editor.set_scroll_position(point(0., scroll_target_top), cx);
767 }
768 }
769 });
770 }
771
772 fn resize_assist(
773 &mut self,
774 assist_id: InlineAssistId,
775 height_in_lines: u8,
776 cx: &mut WindowContext,
777 ) {
778 if let Some(assist) = self.assists.get_mut(&assist_id) {
779 if let Some(editor) = assist.editor.upgrade() {
780 if let Some(decorations) = assist.decorations.as_ref() {
781 let mut new_blocks = HashMap::default();
782 new_blocks.insert(
783 decorations.prompt_block_id,
784 (
785 Some(height_in_lines),
786 build_assist_editor_renderer(&decorations.prompt_editor),
787 ),
788 );
789 editor.update(cx, |editor, cx| {
790 editor
791 .display_map
792 .update(cx, |map, cx| map.replace_blocks(new_blocks, cx))
793 });
794 }
795 }
796 }
797 }
798
799 fn unlink_assist_group(
800 &mut self,
801 assist_group_id: InlineAssistGroupId,
802 cx: &mut WindowContext,
803 ) -> Vec<InlineAssistId> {
804 let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
805 assist_group.linked = false;
806 for assist_id in &assist_group.assist_ids {
807 let assist = self.assists.get_mut(assist_id).unwrap();
808 if let Some(editor_decorations) = assist.decorations.as_ref() {
809 editor_decorations
810 .prompt_editor
811 .update(cx, |prompt_editor, cx| prompt_editor.unlink(cx));
812 }
813 }
814 assist_group.assist_ids.clone()
815 }
816
817 pub fn start_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
818 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
819 assist
820 } else {
821 return;
822 };
823
824 let assist_group_id = assist.group_id;
825 if self.assist_groups[&assist_group_id].linked {
826 for assist_id in self.unlink_assist_group(assist_group_id, cx) {
827 self.start_assist(assist_id, cx);
828 }
829 return;
830 }
831
832 let Some(user_prompt) = assist
833 .decorations
834 .as_ref()
835 .map(|decorations| decorations.prompt_editor.read(cx).prompt(cx))
836 else {
837 return;
838 };
839
840 self.prompt_history.retain(|prompt| *prompt != user_prompt);
841 self.prompt_history.push_back(user_prompt.clone());
842 if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
843 self.prompt_history.pop_front();
844 }
845
846 let codegen = assist.codegen.clone();
847 let telemetry_id = LanguageModelCompletionProvider::read_global(cx)
848 .active_model()
849 .map(|m| m.telemetry_id())
850 .unwrap_or_default();
851 let chunks: LocalBoxFuture<Result<BoxStream<Result<String>>>> =
852 if user_prompt.trim().to_lowercase() == "delete" {
853 async { Ok(stream::empty().boxed()) }.boxed_local()
854 } else {
855 let request = self.request_for_inline_assist(assist_id, cx);
856 let mut cx = cx.to_async();
857 async move {
858 let request = request.await?;
859 let chunks = cx
860 .update(|cx| {
861 LanguageModelCompletionProvider::read_global(cx)
862 .stream_completion(request, cx)
863 })?
864 .await?;
865 Ok(chunks.boxed())
866 }
867 .boxed_local()
868 };
869 codegen.update(cx, |codegen, cx| {
870 codegen.start(telemetry_id, chunks, cx);
871 });
872 }
873
874 fn request_for_inline_assist(
875 &self,
876 assist_id: InlineAssistId,
877 cx: &mut WindowContext,
878 ) -> Task<Result<LanguageModelRequest>> {
879 cx.spawn(|mut cx| async move {
880 let (user_prompt, context_request, project_name, buffer, range) =
881 cx.read_global(|this: &InlineAssistant, cx: &WindowContext| {
882 let assist = this.assists.get(&assist_id).context("invalid assist")?;
883 let decorations = assist.decorations.as_ref().context("invalid assist")?;
884 let editor = assist.editor.upgrade().context("invalid assist")?;
885 let user_prompt = decorations.prompt_editor.read(cx).prompt(cx);
886 let context_request = if assist.include_context {
887 assist.workspace.as_ref().and_then(|workspace| {
888 let workspace = workspace.upgrade()?.read(cx);
889 let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
890 Some(
891 assistant_panel
892 .read(cx)
893 .active_context(cx)?
894 .read(cx)
895 .to_completion_request(cx),
896 )
897 })
898 } else {
899 None
900 };
901 let project_name = assist.workspace.as_ref().and_then(|workspace| {
902 let workspace = workspace.upgrade()?;
903 Some(
904 workspace
905 .read(cx)
906 .project()
907 .read(cx)
908 .worktree_root_names(cx)
909 .collect::<Vec<&str>>()
910 .join("/"),
911 )
912 });
913 let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
914 let range = assist.codegen.read(cx).range.clone();
915 anyhow::Ok((user_prompt, context_request, project_name, buffer, range))
916 })??;
917
918 let language = buffer.language_at(range.start);
919 let language_name = if let Some(language) = language.as_ref() {
920 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
921 None
922 } else {
923 Some(language.name())
924 }
925 } else {
926 None
927 };
928
929 // Higher Temperature increases the randomness of model outputs.
930 // If Markdown or No Language is Known, increase the randomness for more creative output
931 // If Code, decrease temperature to get more deterministic outputs
932 let temperature = if let Some(language) = language_name.clone() {
933 if language.as_ref() == "Markdown" {
934 1.0
935 } else {
936 0.5
937 }
938 } else {
939 1.0
940 };
941
942 let prompt = cx
943 .background_executor()
944 .spawn(async move {
945 let language_name = language_name.as_deref();
946 let start = buffer.point_to_buffer_offset(range.start);
947 let end = buffer.point_to_buffer_offset(range.end);
948 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
949 let (start_buffer, start_buffer_offset) = start;
950 let (end_buffer, end_buffer_offset) = end;
951 if start_buffer.remote_id() == end_buffer.remote_id() {
952 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
953 } else {
954 return Err(anyhow!("invalid transformation range"));
955 }
956 } else {
957 return Err(anyhow!("invalid transformation range"));
958 };
959 generate_content_prompt(user_prompt, language_name, buffer, range, project_name)
960 })
961 .await?;
962
963 let mut messages = Vec::new();
964 if let Some(context_request) = context_request {
965 messages = context_request.messages;
966 }
967
968 messages.push(LanguageModelRequestMessage {
969 role: Role::User,
970 content: prompt,
971 });
972
973 Ok(LanguageModelRequest {
974 messages,
975 stop: vec!["|END|>".to_string()],
976 temperature,
977 })
978 })
979 }
980
981 pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
982 let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
983 assist
984 } else {
985 return;
986 };
987
988 assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
989 }
990
991 fn update_editor_highlights(&self, editor: &View<Editor>, cx: &mut WindowContext) {
992 let mut gutter_pending_ranges = Vec::new();
993 let mut gutter_transformed_ranges = Vec::new();
994 let mut foreground_ranges = Vec::new();
995 let mut inserted_row_ranges = Vec::new();
996 let empty_assist_ids = Vec::new();
997 let assist_ids = self
998 .assists_by_editor
999 .get(&editor.downgrade())
1000 .map_or(&empty_assist_ids, |editor_assists| {
1001 &editor_assists.assist_ids
1002 });
1003
1004 for assist_id in assist_ids {
1005 if let Some(assist) = self.assists.get(assist_id) {
1006 let codegen = assist.codegen.read(cx);
1007 foreground_ranges.extend(codegen.last_equal_ranges().iter().cloned());
1008
1009 if codegen.edit_position != codegen.range.end {
1010 gutter_pending_ranges.push(codegen.edit_position..codegen.range.end);
1011 }
1012
1013 if codegen.range.start != codegen.edit_position {
1014 gutter_transformed_ranges.push(codegen.range.start..codegen.edit_position);
1015 }
1016
1017 if assist.decorations.is_some() {
1018 inserted_row_ranges.extend(codegen.diff.inserted_row_ranges.iter().cloned());
1019 }
1020 }
1021 }
1022
1023 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
1024 merge_ranges(&mut foreground_ranges, &snapshot);
1025 merge_ranges(&mut gutter_pending_ranges, &snapshot);
1026 merge_ranges(&mut gutter_transformed_ranges, &snapshot);
1027 editor.update(cx, |editor, cx| {
1028 enum GutterPendingRange {}
1029 if gutter_pending_ranges.is_empty() {
1030 editor.clear_gutter_highlights::<GutterPendingRange>(cx);
1031 } else {
1032 editor.highlight_gutter::<GutterPendingRange>(
1033 &gutter_pending_ranges,
1034 |cx| cx.theme().status().info_background,
1035 cx,
1036 )
1037 }
1038
1039 enum GutterTransformedRange {}
1040 if gutter_transformed_ranges.is_empty() {
1041 editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
1042 } else {
1043 editor.highlight_gutter::<GutterTransformedRange>(
1044 &gutter_transformed_ranges,
1045 |cx| cx.theme().status().info,
1046 cx,
1047 )
1048 }
1049
1050 if foreground_ranges.is_empty() {
1051 editor.clear_highlights::<InlineAssist>(cx);
1052 } else {
1053 editor.highlight_text::<InlineAssist>(
1054 foreground_ranges,
1055 HighlightStyle {
1056 fade_out: Some(0.6),
1057 ..Default::default()
1058 },
1059 cx,
1060 );
1061 }
1062
1063 editor.clear_row_highlights::<InlineAssist>();
1064 for row_range in inserted_row_ranges {
1065 editor.highlight_rows::<InlineAssist>(
1066 row_range,
1067 Some(cx.theme().status().info_background),
1068 false,
1069 cx,
1070 );
1071 }
1072 });
1073 }
1074
1075 fn update_editor_blocks(
1076 &mut self,
1077 editor: &View<Editor>,
1078 assist_id: InlineAssistId,
1079 cx: &mut WindowContext,
1080 ) {
1081 let Some(assist) = self.assists.get_mut(&assist_id) else {
1082 return;
1083 };
1084 let Some(decorations) = assist.decorations.as_mut() else {
1085 return;
1086 };
1087
1088 let codegen = assist.codegen.read(cx);
1089 let old_snapshot = codegen.snapshot.clone();
1090 let old_buffer = codegen.old_buffer.clone();
1091 let deleted_row_ranges = codegen.diff.deleted_row_ranges.clone();
1092
1093 editor.update(cx, |editor, cx| {
1094 let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
1095 editor.remove_blocks(old_blocks, None, cx);
1096
1097 let mut new_blocks = Vec::new();
1098 for (new_row, old_row_range) in deleted_row_ranges {
1099 let (_, buffer_start) = old_snapshot
1100 .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
1101 .unwrap();
1102 let (_, buffer_end) = old_snapshot
1103 .point_to_buffer_offset(Point::new(
1104 *old_row_range.end(),
1105 old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
1106 ))
1107 .unwrap();
1108
1109 let deleted_lines_editor = cx.new_view(|cx| {
1110 let multi_buffer = cx.new_model(|_| {
1111 MultiBuffer::without_headers(0, language::Capability::ReadOnly)
1112 });
1113 multi_buffer.update(cx, |multi_buffer, cx| {
1114 multi_buffer.push_excerpts(
1115 old_buffer.clone(),
1116 Some(ExcerptRange {
1117 context: buffer_start..buffer_end,
1118 primary: None,
1119 }),
1120 cx,
1121 );
1122 });
1123
1124 enum DeletedLines {}
1125 let mut editor = Editor::for_multibuffer(multi_buffer, None, true, cx);
1126 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1127 editor.set_show_wrap_guides(false, cx);
1128 editor.set_show_gutter(false, cx);
1129 editor.scroll_manager.set_forbid_vertical_scroll(true);
1130 editor.set_read_only(true);
1131 editor.highlight_rows::<DeletedLines>(
1132 Anchor::min()..=Anchor::max(),
1133 Some(cx.theme().status().deleted_background),
1134 false,
1135 cx,
1136 );
1137 editor
1138 });
1139
1140 let height = deleted_lines_editor
1141 .update(cx, |editor, cx| editor.max_point(cx).row().0 as u8 + 1);
1142 new_blocks.push(BlockProperties {
1143 position: new_row,
1144 height,
1145 style: BlockStyle::Flex,
1146 render: Box::new(move |cx| {
1147 div()
1148 .bg(cx.theme().status().deleted_background)
1149 .size_full()
1150 .pl(cx.gutter_dimensions.full_width())
1151 .child(deleted_lines_editor.clone())
1152 .into_any_element()
1153 }),
1154 disposition: BlockDisposition::Above,
1155 });
1156 }
1157
1158 decorations.removed_line_block_ids = editor
1159 .insert_blocks(new_blocks, None, cx)
1160 .into_iter()
1161 .collect();
1162 })
1163 }
1164}
1165
1166struct EditorInlineAssists {
1167 assist_ids: Vec<InlineAssistId>,
1168 scroll_lock: Option<InlineAssistScrollLock>,
1169 highlight_updates: async_watch::Sender<()>,
1170 _update_highlights: Task<Result<()>>,
1171 _subscriptions: Vec<gpui::Subscription>,
1172}
1173
1174struct InlineAssistScrollLock {
1175 assist_id: InlineAssistId,
1176 distance_from_top: f32,
1177}
1178
1179impl EditorInlineAssists {
1180 #[allow(clippy::too_many_arguments)]
1181 fn new(editor: &View<Editor>, cx: &mut WindowContext) -> Self {
1182 let (highlight_updates_tx, mut highlight_updates_rx) = async_watch::channel(());
1183 Self {
1184 assist_ids: Vec::new(),
1185 scroll_lock: None,
1186 highlight_updates: highlight_updates_tx,
1187 _update_highlights: cx.spawn(|mut cx| {
1188 let editor = editor.downgrade();
1189 async move {
1190 while let Ok(()) = highlight_updates_rx.changed().await {
1191 let editor = editor.upgrade().context("editor was dropped")?;
1192 cx.update_global(|assistant: &mut InlineAssistant, cx| {
1193 assistant.update_editor_highlights(&editor, cx);
1194 })?;
1195 }
1196 Ok(())
1197 }
1198 }),
1199 _subscriptions: vec![
1200 cx.observe_release(editor, {
1201 let editor = editor.downgrade();
1202 |_, cx| {
1203 InlineAssistant::update_global(cx, |this, cx| {
1204 this.handle_editor_release(editor, cx);
1205 })
1206 }
1207 }),
1208 cx.observe(editor, move |editor, cx| {
1209 InlineAssistant::update_global(cx, |this, cx| {
1210 this.handle_editor_change(editor, cx)
1211 })
1212 }),
1213 cx.subscribe(editor, move |editor, event, cx| {
1214 InlineAssistant::update_global(cx, |this, cx| {
1215 this.handle_editor_event(editor, event, cx)
1216 })
1217 }),
1218 editor.update(cx, |editor, cx| {
1219 let editor_handle = cx.view().downgrade();
1220 editor.register_action(
1221 move |_: &editor::actions::Newline, cx: &mut WindowContext| {
1222 InlineAssistant::update_global(cx, |this, cx| {
1223 if let Some(editor) = editor_handle.upgrade() {
1224 this.handle_editor_newline(editor, cx)
1225 }
1226 })
1227 },
1228 )
1229 }),
1230 editor.update(cx, |editor, cx| {
1231 let editor_handle = cx.view().downgrade();
1232 editor.register_action(
1233 move |_: &editor::actions::Cancel, cx: &mut WindowContext| {
1234 InlineAssistant::update_global(cx, |this, cx| {
1235 if let Some(editor) = editor_handle.upgrade() {
1236 this.handle_editor_cancel(editor, cx)
1237 }
1238 })
1239 },
1240 )
1241 }),
1242 ],
1243 }
1244 }
1245}
1246
1247struct InlineAssistGroup {
1248 assist_ids: Vec<InlineAssistId>,
1249 linked: bool,
1250 active_assist_id: Option<InlineAssistId>,
1251}
1252
1253impl InlineAssistGroup {
1254 fn new() -> Self {
1255 Self {
1256 assist_ids: Vec::new(),
1257 linked: true,
1258 active_assist_id: None,
1259 }
1260 }
1261}
1262
1263fn build_assist_editor_renderer(editor: &View<PromptEditor>) -> RenderBlock {
1264 let editor = editor.clone();
1265 Box::new(move |cx: &mut BlockContext| {
1266 *editor.read(cx).gutter_dimensions.lock() = *cx.gutter_dimensions;
1267 editor.clone().into_any_element()
1268 })
1269}
1270
1271#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1272pub struct InlineAssistId(usize);
1273
1274impl InlineAssistId {
1275 fn post_inc(&mut self) -> InlineAssistId {
1276 let id = *self;
1277 self.0 += 1;
1278 id
1279 }
1280}
1281
1282#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1283struct InlineAssistGroupId(usize);
1284
1285impl InlineAssistGroupId {
1286 fn post_inc(&mut self) -> InlineAssistGroupId {
1287 let id = *self;
1288 self.0 += 1;
1289 id
1290 }
1291}
1292
1293enum PromptEditorEvent {
1294 StartRequested,
1295 StopRequested,
1296 ConfirmRequested,
1297 CancelRequested,
1298 DismissRequested,
1299 Resized { height_in_lines: u8 },
1300}
1301
1302struct PromptEditor {
1303 id: InlineAssistId,
1304 fs: Arc<dyn Fs>,
1305 height_in_lines: u8,
1306 editor: View<Editor>,
1307 edited_since_done: bool,
1308 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1309 prompt_history: VecDeque<String>,
1310 prompt_history_ix: Option<usize>,
1311 pending_prompt: String,
1312 codegen: Model<Codegen>,
1313 _codegen_subscription: Subscription,
1314 editor_subscriptions: Vec<Subscription>,
1315 pending_token_count: Task<Result<()>>,
1316 token_count: Option<usize>,
1317 _token_count_subscriptions: Vec<Subscription>,
1318 workspace: Option<WeakView<Workspace>>,
1319}
1320
1321impl EventEmitter<PromptEditorEvent> for PromptEditor {}
1322
1323impl Render for PromptEditor {
1324 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1325 let gutter_dimensions = *self.gutter_dimensions.lock();
1326 let buttons = match &self.codegen.read(cx).status {
1327 CodegenStatus::Idle => {
1328 vec![
1329 IconButton::new("cancel", IconName::Close)
1330 .icon_color(Color::Muted)
1331 .shape(IconButtonShape::Square)
1332 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1333 .on_click(
1334 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1335 ),
1336 IconButton::new("start", IconName::SparkleAlt)
1337 .icon_color(Color::Muted)
1338 .shape(IconButtonShape::Square)
1339 .tooltip(|cx| Tooltip::for_action("Transform", &menu::Confirm, cx))
1340 .on_click(
1341 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
1342 ),
1343 ]
1344 }
1345 CodegenStatus::Pending => {
1346 vec![
1347 IconButton::new("cancel", IconName::Close)
1348 .icon_color(Color::Muted)
1349 .shape(IconButtonShape::Square)
1350 .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
1351 .on_click(
1352 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1353 ),
1354 IconButton::new("stop", IconName::Stop)
1355 .icon_color(Color::Error)
1356 .shape(IconButtonShape::Square)
1357 .tooltip(|cx| {
1358 Tooltip::with_meta(
1359 "Interrupt Transformation",
1360 Some(&menu::Cancel),
1361 "Changes won't be discarded",
1362 cx,
1363 )
1364 })
1365 .on_click(
1366 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
1367 ),
1368 ]
1369 }
1370 CodegenStatus::Error(_) | CodegenStatus::Done => {
1371 vec![
1372 IconButton::new("cancel", IconName::Close)
1373 .icon_color(Color::Muted)
1374 .shape(IconButtonShape::Square)
1375 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1376 .on_click(
1377 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1378 ),
1379 if self.edited_since_done {
1380 IconButton::new("restart", IconName::RotateCw)
1381 .icon_color(Color::Info)
1382 .shape(IconButtonShape::Square)
1383 .tooltip(|cx| {
1384 Tooltip::with_meta(
1385 "Restart Transformation",
1386 Some(&menu::Confirm),
1387 "Changes will be discarded",
1388 cx,
1389 )
1390 })
1391 .on_click(cx.listener(|_, _, cx| {
1392 cx.emit(PromptEditorEvent::StartRequested);
1393 }))
1394 } else {
1395 IconButton::new("confirm", IconName::Check)
1396 .icon_color(Color::Info)
1397 .shape(IconButtonShape::Square)
1398 .tooltip(|cx| Tooltip::for_action("Confirm Assist", &menu::Confirm, cx))
1399 .on_click(cx.listener(|_, _, cx| {
1400 cx.emit(PromptEditorEvent::ConfirmRequested);
1401 }))
1402 },
1403 ]
1404 }
1405 };
1406
1407 h_flex()
1408 .bg(cx.theme().colors().editor_background)
1409 .border_y_1()
1410 .border_color(cx.theme().status().info_border)
1411 .py_1p5()
1412 .h_full()
1413 .w_full()
1414 .on_action(cx.listener(Self::confirm))
1415 .on_action(cx.listener(Self::cancel))
1416 .on_action(cx.listener(Self::move_up))
1417 .on_action(cx.listener(Self::move_down))
1418 .child(
1419 h_flex()
1420 .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
1421 .justify_center()
1422 .gap_2()
1423 .child(
1424 ModelSelector::new(
1425 self.fs.clone(),
1426 IconButton::new("context", IconName::SlidersAlt)
1427 .shape(IconButtonShape::Square)
1428 .icon_size(IconSize::Small)
1429 .icon_color(Color::Muted)
1430 .tooltip(move |cx| {
1431 Tooltip::with_meta(
1432 format!(
1433 "Using {}",
1434 LanguageModelCompletionProvider::read_global(cx)
1435 .active_model()
1436 .map(|model| model.name().0)
1437 .unwrap_or_else(|| "No model selected".into()),
1438 ),
1439 None,
1440 "Change Model",
1441 cx,
1442 )
1443 }),
1444 )
1445 .with_info_text(
1446 "Inline edits use context\n\
1447 from the currently selected\n\
1448 assistant panel tab.",
1449 ),
1450 )
1451 .children(
1452 if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
1453 let error_message = SharedString::from(error.to_string());
1454 Some(
1455 div()
1456 .id("error")
1457 .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
1458 .child(
1459 Icon::new(IconName::XCircle)
1460 .size(IconSize::Small)
1461 .color(Color::Error),
1462 ),
1463 )
1464 } else {
1465 None
1466 },
1467 ),
1468 )
1469 .child(div().flex_1().child(self.render_prompt_editor(cx)))
1470 .child(
1471 h_flex()
1472 .gap_2()
1473 .pr_6()
1474 .children(self.render_token_count(cx))
1475 .children(buttons),
1476 )
1477 }
1478}
1479
1480impl FocusableView for PromptEditor {
1481 fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
1482 self.editor.focus_handle(cx)
1483 }
1484}
1485
1486impl PromptEditor {
1487 const MAX_LINES: u8 = 8;
1488
1489 #[allow(clippy::too_many_arguments)]
1490 fn new(
1491 id: InlineAssistId,
1492 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1493 prompt_history: VecDeque<String>,
1494 prompt_buffer: Model<MultiBuffer>,
1495 codegen: Model<Codegen>,
1496 parent_editor: &View<Editor>,
1497 assistant_panel: Option<&View<AssistantPanel>>,
1498 workspace: Option<WeakView<Workspace>>,
1499 fs: Arc<dyn Fs>,
1500 cx: &mut ViewContext<Self>,
1501 ) -> Self {
1502 let prompt_editor = cx.new_view(|cx| {
1503 let mut editor = Editor::new(
1504 EditorMode::AutoHeight {
1505 max_lines: Self::MAX_LINES as usize,
1506 },
1507 prompt_buffer,
1508 None,
1509 false,
1510 cx,
1511 );
1512 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1513 // Since the prompt editors for all inline assistants are linked,
1514 // always show the cursor (even when it isn't focused) because
1515 // typing in one will make what you typed appear in all of them.
1516 editor.set_show_cursor_when_unfocused(true, cx);
1517 editor.set_placeholder_text("Add a prompt…", cx);
1518 editor
1519 });
1520
1521 let mut token_count_subscriptions = Vec::new();
1522 token_count_subscriptions
1523 .push(cx.subscribe(parent_editor, Self::handle_parent_editor_event));
1524 if let Some(assistant_panel) = assistant_panel {
1525 token_count_subscriptions
1526 .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
1527 }
1528
1529 let mut this = Self {
1530 id,
1531 height_in_lines: 1,
1532 editor: prompt_editor,
1533 edited_since_done: false,
1534 gutter_dimensions,
1535 prompt_history,
1536 prompt_history_ix: None,
1537 pending_prompt: String::new(),
1538 _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
1539 editor_subscriptions: Vec::new(),
1540 codegen,
1541 fs,
1542 pending_token_count: Task::ready(Ok(())),
1543 token_count: None,
1544 _token_count_subscriptions: token_count_subscriptions,
1545 workspace,
1546 };
1547 this.count_lines(cx);
1548 this.count_tokens(cx);
1549 this.subscribe_to_editor(cx);
1550 this
1551 }
1552
1553 fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
1554 self.editor_subscriptions.clear();
1555 self.editor_subscriptions
1556 .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
1557 self.editor_subscriptions
1558 .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
1559 }
1560
1561 fn set_show_cursor_when_unfocused(
1562 &mut self,
1563 show_cursor_when_unfocused: bool,
1564 cx: &mut ViewContext<Self>,
1565 ) {
1566 self.editor.update(cx, |editor, cx| {
1567 editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
1568 });
1569 }
1570
1571 fn unlink(&mut self, cx: &mut ViewContext<Self>) {
1572 let prompt = self.prompt(cx);
1573 let focus = self.editor.focus_handle(cx).contains_focused(cx);
1574 self.editor = cx.new_view(|cx| {
1575 let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
1576 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1577 editor.set_placeholder_text("Add a prompt…", cx);
1578 editor.set_text(prompt, cx);
1579 if focus {
1580 editor.focus(cx);
1581 }
1582 editor
1583 });
1584 self.subscribe_to_editor(cx);
1585 }
1586
1587 fn prompt(&self, cx: &AppContext) -> String {
1588 self.editor.read(cx).text(cx)
1589 }
1590
1591 fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
1592 let height_in_lines = cmp::max(
1593 2, // Make the editor at least two lines tall, to account for padding and buttons.
1594 cmp::min(
1595 self.editor
1596 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1597 Self::MAX_LINES as u32,
1598 ),
1599 ) as u8;
1600
1601 if height_in_lines != self.height_in_lines {
1602 self.height_in_lines = height_in_lines;
1603 cx.emit(PromptEditorEvent::Resized { height_in_lines });
1604 }
1605 }
1606
1607 fn handle_parent_editor_event(
1608 &mut self,
1609 _: View<Editor>,
1610 event: &EditorEvent,
1611 cx: &mut ViewContext<Self>,
1612 ) {
1613 if let EditorEvent::BufferEdited { .. } = event {
1614 self.count_tokens(cx);
1615 }
1616 }
1617
1618 fn handle_assistant_panel_event(
1619 &mut self,
1620 _: View<AssistantPanel>,
1621 event: &AssistantPanelEvent,
1622 cx: &mut ViewContext<Self>,
1623 ) {
1624 let AssistantPanelEvent::ContextEdited { .. } = event;
1625 self.count_tokens(cx);
1626 }
1627
1628 fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
1629 let assist_id = self.id;
1630 self.pending_token_count = cx.spawn(|this, mut cx| async move {
1631 cx.background_executor().timer(Duration::from_secs(1)).await;
1632 let request = cx
1633 .update_global(|inline_assistant: &mut InlineAssistant, cx| {
1634 inline_assistant.request_for_inline_assist(assist_id, cx)
1635 })?
1636 .await?;
1637
1638 let token_count = cx
1639 .update(|cx| {
1640 LanguageModelCompletionProvider::read_global(cx).count_tokens(request, cx)
1641 })?
1642 .await?;
1643 this.update(&mut cx, |this, cx| {
1644 this.token_count = Some(token_count);
1645 cx.notify();
1646 })
1647 })
1648 }
1649
1650 fn handle_prompt_editor_changed(&mut self, _: View<Editor>, cx: &mut ViewContext<Self>) {
1651 self.count_lines(cx);
1652 }
1653
1654 fn handle_prompt_editor_events(
1655 &mut self,
1656 _: View<Editor>,
1657 event: &EditorEvent,
1658 cx: &mut ViewContext<Self>,
1659 ) {
1660 match event {
1661 EditorEvent::Edited { .. } => {
1662 let prompt = self.editor.read(cx).text(cx);
1663 if self
1664 .prompt_history_ix
1665 .map_or(true, |ix| self.prompt_history[ix] != prompt)
1666 {
1667 self.prompt_history_ix.take();
1668 self.pending_prompt = prompt;
1669 }
1670
1671 self.edited_since_done = true;
1672 cx.notify();
1673 }
1674 EditorEvent::BufferEdited => {
1675 self.count_tokens(cx);
1676 }
1677 _ => {}
1678 }
1679 }
1680
1681 fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
1682 match &self.codegen.read(cx).status {
1683 CodegenStatus::Idle => {
1684 self.editor
1685 .update(cx, |editor, _| editor.set_read_only(false));
1686 }
1687 CodegenStatus::Pending => {
1688 self.editor
1689 .update(cx, |editor, _| editor.set_read_only(true));
1690 }
1691 CodegenStatus::Done | CodegenStatus::Error(_) => {
1692 self.edited_since_done = false;
1693 self.editor
1694 .update(cx, |editor, _| editor.set_read_only(false));
1695 }
1696 }
1697 }
1698
1699 fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
1700 match &self.codegen.read(cx).status {
1701 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1702 cx.emit(PromptEditorEvent::CancelRequested);
1703 }
1704 CodegenStatus::Pending => {
1705 cx.emit(PromptEditorEvent::StopRequested);
1706 }
1707 }
1708 }
1709
1710 fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
1711 match &self.codegen.read(cx).status {
1712 CodegenStatus::Idle => {
1713 cx.emit(PromptEditorEvent::StartRequested);
1714 }
1715 CodegenStatus::Pending => {
1716 cx.emit(PromptEditorEvent::DismissRequested);
1717 }
1718 CodegenStatus::Done | CodegenStatus::Error(_) => {
1719 if self.edited_since_done {
1720 cx.emit(PromptEditorEvent::StartRequested);
1721 } else {
1722 cx.emit(PromptEditorEvent::ConfirmRequested);
1723 }
1724 }
1725 }
1726 }
1727
1728 fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
1729 if let Some(ix) = self.prompt_history_ix {
1730 if ix > 0 {
1731 self.prompt_history_ix = Some(ix - 1);
1732 let prompt = self.prompt_history[ix - 1].as_str();
1733 self.editor.update(cx, |editor, cx| {
1734 editor.set_text(prompt, cx);
1735 editor.move_to_beginning(&Default::default(), cx);
1736 });
1737 }
1738 } else if !self.prompt_history.is_empty() {
1739 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1740 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1741 self.editor.update(cx, |editor, cx| {
1742 editor.set_text(prompt, cx);
1743 editor.move_to_beginning(&Default::default(), cx);
1744 });
1745 }
1746 }
1747
1748 fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
1749 if let Some(ix) = self.prompt_history_ix {
1750 if ix < self.prompt_history.len() - 1 {
1751 self.prompt_history_ix = Some(ix + 1);
1752 let prompt = self.prompt_history[ix + 1].as_str();
1753 self.editor.update(cx, |editor, cx| {
1754 editor.set_text(prompt, cx);
1755 editor.move_to_end(&Default::default(), cx)
1756 });
1757 } else {
1758 self.prompt_history_ix = None;
1759 let prompt = self.pending_prompt.as_str();
1760 self.editor.update(cx, |editor, cx| {
1761 editor.set_text(prompt, cx);
1762 editor.move_to_end(&Default::default(), cx)
1763 });
1764 }
1765 }
1766 }
1767
1768 fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
1769 let model = LanguageModelCompletionProvider::read_global(cx).active_model()?;
1770 let token_count = self.token_count?;
1771 let max_token_count = model.max_token_count();
1772
1773 let remaining_tokens = max_token_count as isize - token_count as isize;
1774 let token_count_color = if remaining_tokens <= 0 {
1775 Color::Error
1776 } else if token_count as f32 / max_token_count as f32 >= 0.8 {
1777 Color::Warning
1778 } else {
1779 Color::Muted
1780 };
1781
1782 let mut token_count = h_flex()
1783 .id("token_count")
1784 .gap_0p5()
1785 .child(
1786 Label::new(humanize_token_count(token_count))
1787 .size(LabelSize::Small)
1788 .color(token_count_color),
1789 )
1790 .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
1791 .child(
1792 Label::new(humanize_token_count(max_token_count))
1793 .size(LabelSize::Small)
1794 .color(Color::Muted),
1795 );
1796 if let Some(workspace) = self.workspace.clone() {
1797 token_count = token_count
1798 .tooltip(|cx| {
1799 Tooltip::with_meta(
1800 "Tokens Used by Inline Assistant",
1801 None,
1802 "Click to Open Assistant Panel",
1803 cx,
1804 )
1805 })
1806 .cursor_pointer()
1807 .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
1808 .on_click(move |_, cx| {
1809 cx.stop_propagation();
1810 workspace
1811 .update(cx, |workspace, cx| {
1812 workspace.focus_panel::<AssistantPanel>(cx)
1813 })
1814 .ok();
1815 });
1816 } else {
1817 token_count = token_count
1818 .cursor_default()
1819 .tooltip(|cx| Tooltip::text("Tokens Used by Inline Assistant", cx));
1820 }
1821
1822 Some(token_count)
1823 }
1824
1825 fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1826 let settings = ThemeSettings::get_global(cx);
1827 let text_style = TextStyle {
1828 color: if self.editor.read(cx).read_only(cx) {
1829 cx.theme().colors().text_disabled
1830 } else {
1831 cx.theme().colors().text
1832 },
1833 font_family: settings.ui_font.family.clone(),
1834 font_features: settings.ui_font.features.clone(),
1835 font_size: rems(0.875).into(),
1836 font_weight: settings.ui_font.weight,
1837 line_height: relative(1.3),
1838 ..Default::default()
1839 };
1840 EditorElement::new(
1841 &self.editor,
1842 EditorStyle {
1843 background: cx.theme().colors().editor_background,
1844 local_player: cx.theme().players().local(),
1845 text: text_style,
1846 ..Default::default()
1847 },
1848 )
1849 }
1850}
1851
1852struct InlineAssist {
1853 group_id: InlineAssistGroupId,
1854 editor: WeakView<Editor>,
1855 decorations: Option<InlineAssistDecorations>,
1856 codegen: Model<Codegen>,
1857 _subscriptions: Vec<Subscription>,
1858 workspace: Option<WeakView<Workspace>>,
1859 include_context: bool,
1860}
1861
1862impl InlineAssist {
1863 #[allow(clippy::too_many_arguments)]
1864 fn new(
1865 assist_id: InlineAssistId,
1866 group_id: InlineAssistGroupId,
1867 include_context: bool,
1868 editor: &View<Editor>,
1869 prompt_editor: &View<PromptEditor>,
1870 prompt_block_id: CustomBlockId,
1871 end_block_id: CustomBlockId,
1872 codegen: Model<Codegen>,
1873 workspace: Option<WeakView<Workspace>>,
1874 cx: &mut WindowContext,
1875 ) -> Self {
1876 let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
1877 InlineAssist {
1878 group_id,
1879 include_context,
1880 editor: editor.downgrade(),
1881 decorations: Some(InlineAssistDecorations {
1882 prompt_block_id,
1883 prompt_editor: prompt_editor.clone(),
1884 removed_line_block_ids: HashSet::default(),
1885 end_block_id,
1886 }),
1887 codegen: codegen.clone(),
1888 workspace: workspace.clone(),
1889 _subscriptions: vec![
1890 cx.on_focus_in(&prompt_editor_focus_handle, move |cx| {
1891 InlineAssistant::update_global(cx, |this, cx| {
1892 this.handle_prompt_editor_focus_in(assist_id, cx)
1893 })
1894 }),
1895 cx.on_focus_out(&prompt_editor_focus_handle, move |_, cx| {
1896 InlineAssistant::update_global(cx, |this, cx| {
1897 this.handle_prompt_editor_focus_out(assist_id, cx)
1898 })
1899 }),
1900 cx.subscribe(prompt_editor, |prompt_editor, event, cx| {
1901 InlineAssistant::update_global(cx, |this, cx| {
1902 this.handle_prompt_editor_event(prompt_editor, event, cx)
1903 })
1904 }),
1905 cx.observe(&codegen, {
1906 let editor = editor.downgrade();
1907 move |_, cx| {
1908 if let Some(editor) = editor.upgrade() {
1909 InlineAssistant::update_global(cx, |this, cx| {
1910 if let Some(editor_assists) =
1911 this.assists_by_editor.get(&editor.downgrade())
1912 {
1913 editor_assists.highlight_updates.send(()).ok();
1914 }
1915
1916 this.update_editor_blocks(&editor, assist_id, cx);
1917 })
1918 }
1919 }
1920 }),
1921 cx.subscribe(&codegen, move |codegen, event, cx| {
1922 InlineAssistant::update_global(cx, |this, cx| match event {
1923 CodegenEvent::Undone => this.finish_assist(assist_id, false, cx),
1924 CodegenEvent::Finished => {
1925 let assist = if let Some(assist) = this.assists.get(&assist_id) {
1926 assist
1927 } else {
1928 return;
1929 };
1930
1931 if let CodegenStatus::Error(error) = &codegen.read(cx).status {
1932 if assist.decorations.is_none() {
1933 if let Some(workspace) = assist
1934 .workspace
1935 .as_ref()
1936 .and_then(|workspace| workspace.upgrade())
1937 {
1938 let error = format!("Inline assistant error: {}", error);
1939 workspace.update(cx, |workspace, cx| {
1940 struct InlineAssistantError;
1941
1942 let id =
1943 NotificationId::identified::<InlineAssistantError>(
1944 assist_id.0,
1945 );
1946
1947 workspace.show_toast(Toast::new(id, error), cx);
1948 })
1949 }
1950 }
1951 }
1952
1953 if assist.decorations.is_none() {
1954 this.finish_assist(assist_id, false, cx);
1955 }
1956 }
1957 })
1958 }),
1959 ],
1960 }
1961 }
1962}
1963
1964struct InlineAssistDecorations {
1965 prompt_block_id: CustomBlockId,
1966 prompt_editor: View<PromptEditor>,
1967 removed_line_block_ids: HashSet<CustomBlockId>,
1968 end_block_id: CustomBlockId,
1969}
1970
1971#[derive(Debug)]
1972pub enum CodegenEvent {
1973 Finished,
1974 Undone,
1975}
1976
1977pub struct Codegen {
1978 buffer: Model<MultiBuffer>,
1979 old_buffer: Model<Buffer>,
1980 snapshot: MultiBufferSnapshot,
1981 range: Range<Anchor>,
1982 edit_position: Anchor,
1983 last_equal_ranges: Vec<Range<Anchor>>,
1984 prepend_transaction_id: Option<TransactionId>,
1985 generation_transaction_id: Option<TransactionId>,
1986 status: CodegenStatus,
1987 generation: Task<()>,
1988 diff: Diff,
1989 telemetry: Option<Arc<Telemetry>>,
1990 _subscription: gpui::Subscription,
1991}
1992
1993enum CodegenStatus {
1994 Idle,
1995 Pending,
1996 Done,
1997 Error(anyhow::Error),
1998}
1999
2000#[derive(Default)]
2001struct Diff {
2002 task: Option<Task<()>>,
2003 should_update: bool,
2004 deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
2005 inserted_row_ranges: Vec<RangeInclusive<Anchor>>,
2006}
2007
2008impl EventEmitter<CodegenEvent> for Codegen {}
2009
2010impl Codegen {
2011 pub fn new(
2012 buffer: Model<MultiBuffer>,
2013 range: Range<Anchor>,
2014 prepend_transaction_id: Option<TransactionId>,
2015 telemetry: Option<Arc<Telemetry>>,
2016 cx: &mut ModelContext<Self>,
2017 ) -> Self {
2018 let snapshot = buffer.read(cx).snapshot(cx);
2019
2020 let (old_buffer, _, _) = buffer
2021 .read(cx)
2022 .range_to_buffer_ranges(range.clone(), cx)
2023 .pop()
2024 .unwrap();
2025 let old_buffer = cx.new_model(|cx| {
2026 let old_buffer = old_buffer.read(cx);
2027 let text = old_buffer.as_rope().clone();
2028 let line_ending = old_buffer.line_ending();
2029 let language = old_buffer.language().cloned();
2030 let language_registry = old_buffer.language_registry();
2031
2032 let mut buffer = Buffer::local_normalized(text, line_ending, cx);
2033 buffer.set_language(language, cx);
2034 if let Some(language_registry) = language_registry {
2035 buffer.set_language_registry(language_registry)
2036 }
2037 buffer
2038 });
2039
2040 Self {
2041 buffer: buffer.clone(),
2042 old_buffer,
2043 edit_position: range.start,
2044 range,
2045 snapshot,
2046 last_equal_ranges: Default::default(),
2047 prepend_transaction_id,
2048 generation_transaction_id: None,
2049 status: CodegenStatus::Idle,
2050 generation: Task::ready(()),
2051 diff: Diff::default(),
2052 telemetry,
2053 _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
2054 }
2055 }
2056
2057 fn handle_buffer_event(
2058 &mut self,
2059 _buffer: Model<MultiBuffer>,
2060 event: &multi_buffer::Event,
2061 cx: &mut ModelContext<Self>,
2062 ) {
2063 if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
2064 if self.generation_transaction_id == Some(*transaction_id) {
2065 self.generation_transaction_id = None;
2066 self.generation = Task::ready(());
2067 cx.emit(CodegenEvent::Undone);
2068 } else if self.prepend_transaction_id == Some(*transaction_id) {
2069 self.prepend_transaction_id = None;
2070 self.generation_transaction_id = None;
2071 self.generation = Task::ready(());
2072 cx.emit(CodegenEvent::Undone);
2073 }
2074 }
2075 }
2076
2077 pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
2078 &self.last_equal_ranges
2079 }
2080
2081 pub fn start(
2082 &mut self,
2083 telemetry_id: String,
2084 stream: impl 'static + Future<Output = Result<BoxStream<'static, Result<String>>>>,
2085 cx: &mut ModelContext<Self>,
2086 ) {
2087 let range = self.range.clone();
2088 let snapshot = self.snapshot.clone();
2089 let selected_text = snapshot
2090 .text_for_range(range.start..range.end)
2091 .collect::<Rope>();
2092
2093 let selection_start = range.start.to_point(&snapshot);
2094
2095 // Start with the indentation of the first line in the selection
2096 let mut suggested_line_indent = snapshot
2097 .suggested_indents(selection_start.row..=selection_start.row, cx)
2098 .into_values()
2099 .next()
2100 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
2101
2102 // If the first line in the selection does not have indentation, check the following lines
2103 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
2104 for row in selection_start.row..=range.end.to_point(&snapshot).row {
2105 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
2106 // Prefer tabs if a line in the selection uses tabs as indentation
2107 if line_indent.kind == IndentKind::Tab {
2108 suggested_line_indent.kind = IndentKind::Tab;
2109 break;
2110 }
2111 }
2112 }
2113
2114 let telemetry = self.telemetry.clone();
2115 self.edit_position = range.start;
2116 self.diff = Diff::default();
2117 self.status = CodegenStatus::Pending;
2118 if let Some(transaction_id) = self.generation_transaction_id.take() {
2119 self.buffer
2120 .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
2121 }
2122 self.generation = cx.spawn(|this, mut cx| {
2123 async move {
2124 let chunks = stream.await;
2125 let generate = async {
2126 let mut edit_start = range.start.to_offset(&snapshot);
2127
2128 let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
2129 let diff: Task<anyhow::Result<()>> =
2130 cx.background_executor().spawn(async move {
2131 let mut response_latency = None;
2132 let request_start = Instant::now();
2133 let diff = async {
2134 let chunks = StripInvalidSpans::new(chunks?);
2135 futures::pin_mut!(chunks);
2136 let mut diff = StreamingDiff::new(selected_text.to_string());
2137
2138 let mut new_text = String::new();
2139 let mut base_indent = None;
2140 let mut line_indent = None;
2141 let mut first_line = true;
2142
2143 while let Some(chunk) = chunks.next().await {
2144 if response_latency.is_none() {
2145 response_latency = Some(request_start.elapsed());
2146 }
2147 let chunk = chunk?;
2148
2149 let mut lines = chunk.split('\n').peekable();
2150 while let Some(line) = lines.next() {
2151 new_text.push_str(line);
2152 if line_indent.is_none() {
2153 if let Some(non_whitespace_ch_ix) =
2154 new_text.find(|ch: char| !ch.is_whitespace())
2155 {
2156 line_indent = Some(non_whitespace_ch_ix);
2157 base_indent = base_indent.or(line_indent);
2158
2159 let line_indent = line_indent.unwrap();
2160 let base_indent = base_indent.unwrap();
2161 let indent_delta =
2162 line_indent as i32 - base_indent as i32;
2163 let mut corrected_indent_len = cmp::max(
2164 0,
2165 suggested_line_indent.len as i32 + indent_delta,
2166 )
2167 as usize;
2168 if first_line {
2169 corrected_indent_len = corrected_indent_len
2170 .saturating_sub(
2171 selection_start.column as usize,
2172 );
2173 }
2174
2175 let indent_char = suggested_line_indent.char();
2176 let mut indent_buffer = [0; 4];
2177 let indent_str =
2178 indent_char.encode_utf8(&mut indent_buffer);
2179 new_text.replace_range(
2180 ..line_indent,
2181 &indent_str.repeat(corrected_indent_len),
2182 );
2183 }
2184 }
2185
2186 if line_indent.is_some() {
2187 hunks_tx.send(diff.push_new(&new_text)).await?;
2188 new_text.clear();
2189 }
2190
2191 if lines.peek().is_some() {
2192 hunks_tx.send(diff.push_new("\n")).await?;
2193 if line_indent.is_none() {
2194 // Don't write out the leading indentation in empty lines on the next line
2195 // This is the case where the above if statement didn't clear the buffer
2196 new_text.clear();
2197 }
2198 line_indent = None;
2199 first_line = false;
2200 }
2201 }
2202 }
2203 hunks_tx.send(diff.push_new(&new_text)).await?;
2204 hunks_tx.send(diff.finish()).await?;
2205
2206 anyhow::Ok(())
2207 };
2208
2209 let result = diff.await;
2210
2211 let error_message =
2212 result.as_ref().err().map(|error| error.to_string());
2213 if let Some(telemetry) = telemetry {
2214 telemetry.report_assistant_event(
2215 None,
2216 telemetry_events::AssistantKind::Inline,
2217 telemetry_id,
2218 response_latency,
2219 error_message,
2220 );
2221 }
2222
2223 result?;
2224 Ok(())
2225 });
2226
2227 while let Some(hunks) = hunks_rx.next().await {
2228 this.update(&mut cx, |this, cx| {
2229 this.last_equal_ranges.clear();
2230
2231 let transaction = this.buffer.update(cx, |buffer, cx| {
2232 // Avoid grouping assistant edits with user edits.
2233 buffer.finalize_last_transaction(cx);
2234
2235 buffer.start_transaction(cx);
2236 buffer.edit(
2237 hunks.into_iter().filter_map(|hunk| match hunk {
2238 Hunk::Insert { text } => {
2239 let edit_start = snapshot.anchor_after(edit_start);
2240 Some((edit_start..edit_start, text))
2241 }
2242 Hunk::Remove { len } => {
2243 let edit_end = edit_start + len;
2244 let edit_range = snapshot.anchor_after(edit_start)
2245 ..snapshot.anchor_before(edit_end);
2246 edit_start = edit_end;
2247 Some((edit_range, String::new()))
2248 }
2249 Hunk::Keep { len } => {
2250 let edit_end = edit_start + len;
2251 let edit_range = snapshot.anchor_after(edit_start)
2252 ..snapshot.anchor_before(edit_end);
2253 edit_start = edit_end;
2254 this.last_equal_ranges.push(edit_range);
2255 None
2256 }
2257 }),
2258 None,
2259 cx,
2260 );
2261 this.edit_position = snapshot.anchor_after(edit_start);
2262
2263 buffer.end_transaction(cx)
2264 });
2265
2266 if let Some(transaction) = transaction {
2267 if let Some(first_transaction) = this.generation_transaction_id {
2268 // Group all assistant edits into the first transaction.
2269 this.buffer.update(cx, |buffer, cx| {
2270 buffer.merge_transactions(
2271 transaction,
2272 first_transaction,
2273 cx,
2274 )
2275 });
2276 } else {
2277 this.generation_transaction_id = Some(transaction);
2278 this.buffer.update(cx, |buffer, cx| {
2279 buffer.finalize_last_transaction(cx)
2280 });
2281 }
2282 }
2283
2284 this.update_diff(cx);
2285 cx.notify();
2286 })?;
2287 }
2288
2289 diff.await?;
2290
2291 anyhow::Ok(())
2292 };
2293
2294 let result = generate.await;
2295 this.update(&mut cx, |this, cx| {
2296 this.last_equal_ranges.clear();
2297 if let Err(error) = result {
2298 this.status = CodegenStatus::Error(error);
2299 } else {
2300 this.status = CodegenStatus::Done;
2301 }
2302 cx.emit(CodegenEvent::Finished);
2303 cx.notify();
2304 })
2305 .ok();
2306 }
2307 });
2308 cx.notify();
2309 }
2310
2311 pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
2312 self.last_equal_ranges.clear();
2313 self.status = CodegenStatus::Done;
2314 self.generation = Task::ready(());
2315 cx.emit(CodegenEvent::Finished);
2316 cx.notify();
2317 }
2318
2319 pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
2320 if let Some(transaction_id) = self.prepend_transaction_id.take() {
2321 self.buffer
2322 .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
2323 }
2324
2325 if let Some(transaction_id) = self.generation_transaction_id.take() {
2326 self.buffer
2327 .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
2328 }
2329 }
2330
2331 fn update_diff(&mut self, cx: &mut ModelContext<Self>) {
2332 if self.diff.task.is_some() {
2333 self.diff.should_update = true;
2334 } else {
2335 self.diff.should_update = false;
2336
2337 let old_snapshot = self.snapshot.clone();
2338 let old_range = self.range.to_point(&old_snapshot);
2339 let new_snapshot = self.buffer.read(cx).snapshot(cx);
2340 let new_range = self.range.to_point(&new_snapshot);
2341
2342 self.diff.task = Some(cx.spawn(|this, mut cx| async move {
2343 let (deleted_row_ranges, inserted_row_ranges) = cx
2344 .background_executor()
2345 .spawn(async move {
2346 let old_text = old_snapshot
2347 .text_for_range(
2348 Point::new(old_range.start.row, 0)
2349 ..Point::new(
2350 old_range.end.row,
2351 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
2352 ),
2353 )
2354 .collect::<String>();
2355 let new_text = new_snapshot
2356 .text_for_range(
2357 Point::new(new_range.start.row, 0)
2358 ..Point::new(
2359 new_range.end.row,
2360 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
2361 ),
2362 )
2363 .collect::<String>();
2364
2365 let mut old_row = old_range.start.row;
2366 let mut new_row = new_range.start.row;
2367 let diff = TextDiff::from_lines(old_text.as_str(), new_text.as_str());
2368
2369 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
2370 let mut inserted_row_ranges = Vec::new();
2371 for change in diff.iter_all_changes() {
2372 let line_count = change.value().lines().count() as u32;
2373 match change.tag() {
2374 similar::ChangeTag::Equal => {
2375 old_row += line_count;
2376 new_row += line_count;
2377 }
2378 similar::ChangeTag::Delete => {
2379 let old_end_row = old_row + line_count - 1;
2380 let new_row =
2381 new_snapshot.anchor_before(Point::new(new_row, 0));
2382
2383 if let Some((_, last_deleted_row_range)) =
2384 deleted_row_ranges.last_mut()
2385 {
2386 if *last_deleted_row_range.end() + 1 == old_row {
2387 *last_deleted_row_range =
2388 *last_deleted_row_range.start()..=old_end_row;
2389 } else {
2390 deleted_row_ranges
2391 .push((new_row, old_row..=old_end_row));
2392 }
2393 } else {
2394 deleted_row_ranges.push((new_row, old_row..=old_end_row));
2395 }
2396
2397 old_row += line_count;
2398 }
2399 similar::ChangeTag::Insert => {
2400 let new_end_row = new_row + line_count - 1;
2401 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
2402 let end = new_snapshot.anchor_before(Point::new(
2403 new_end_row,
2404 new_snapshot.line_len(MultiBufferRow(new_end_row)),
2405 ));
2406 inserted_row_ranges.push(start..=end);
2407 new_row += line_count;
2408 }
2409 }
2410 }
2411
2412 (deleted_row_ranges, inserted_row_ranges)
2413 })
2414 .await;
2415
2416 this.update(&mut cx, |this, cx| {
2417 this.diff.deleted_row_ranges = deleted_row_ranges;
2418 this.diff.inserted_row_ranges = inserted_row_ranges;
2419 this.diff.task = None;
2420 if this.diff.should_update {
2421 this.update_diff(cx);
2422 }
2423 cx.notify();
2424 })
2425 .ok();
2426 }));
2427 }
2428 }
2429}
2430
2431struct StripInvalidSpans<T> {
2432 stream: T,
2433 stream_done: bool,
2434 buffer: String,
2435 first_line: bool,
2436 line_end: bool,
2437 starts_with_code_block: bool,
2438}
2439
2440impl<T> StripInvalidSpans<T>
2441where
2442 T: Stream<Item = Result<String>>,
2443{
2444 fn new(stream: T) -> Self {
2445 Self {
2446 stream,
2447 stream_done: false,
2448 buffer: String::new(),
2449 first_line: true,
2450 line_end: false,
2451 starts_with_code_block: false,
2452 }
2453 }
2454}
2455
2456impl<T> Stream for StripInvalidSpans<T>
2457where
2458 T: Stream<Item = Result<String>>,
2459{
2460 type Item = Result<String>;
2461
2462 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
2463 const CODE_BLOCK_DELIMITER: &str = "```";
2464 const CURSOR_SPAN: &str = "<|CURSOR|>";
2465
2466 let this = unsafe { self.get_unchecked_mut() };
2467 loop {
2468 if !this.stream_done {
2469 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
2470 match stream.as_mut().poll_next(cx) {
2471 Poll::Ready(Some(Ok(chunk))) => {
2472 this.buffer.push_str(&chunk);
2473 }
2474 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
2475 Poll::Ready(None) => {
2476 this.stream_done = true;
2477 }
2478 Poll::Pending => return Poll::Pending,
2479 }
2480 }
2481
2482 let mut chunk = String::new();
2483 let mut consumed = 0;
2484 if !this.buffer.is_empty() {
2485 let mut lines = this.buffer.split('\n').enumerate().peekable();
2486 while let Some((line_ix, line)) = lines.next() {
2487 if line_ix > 0 {
2488 this.first_line = false;
2489 }
2490
2491 if this.first_line {
2492 let trimmed_line = line.trim();
2493 if lines.peek().is_some() {
2494 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
2495 consumed += line.len() + 1;
2496 this.starts_with_code_block = true;
2497 continue;
2498 }
2499 } else if trimmed_line.is_empty()
2500 || prefixes(CODE_BLOCK_DELIMITER)
2501 .any(|prefix| trimmed_line.starts_with(prefix))
2502 {
2503 break;
2504 }
2505 }
2506
2507 let line_without_cursor = line.replace(CURSOR_SPAN, "");
2508 if lines.peek().is_some() {
2509 if this.line_end {
2510 chunk.push('\n');
2511 }
2512
2513 chunk.push_str(&line_without_cursor);
2514 this.line_end = true;
2515 consumed += line.len() + 1;
2516 } else if this.stream_done {
2517 if !this.starts_with_code_block
2518 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
2519 {
2520 if this.line_end {
2521 chunk.push('\n');
2522 }
2523
2524 chunk.push_str(&line);
2525 }
2526
2527 consumed += line.len();
2528 } else {
2529 let trimmed_line = line.trim();
2530 if trimmed_line.is_empty()
2531 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
2532 || prefixes(CODE_BLOCK_DELIMITER)
2533 .any(|prefix| trimmed_line.ends_with(prefix))
2534 {
2535 break;
2536 } else {
2537 if this.line_end {
2538 chunk.push('\n');
2539 this.line_end = false;
2540 }
2541
2542 chunk.push_str(&line_without_cursor);
2543 consumed += line.len();
2544 }
2545 }
2546 }
2547 }
2548
2549 this.buffer = this.buffer.split_off(consumed);
2550 if !chunk.is_empty() {
2551 return Poll::Ready(Some(Ok(chunk)));
2552 } else if this.stream_done {
2553 return Poll::Ready(None);
2554 }
2555 }
2556 }
2557}
2558
2559fn prefixes(text: &str) -> impl Iterator<Item = &str> {
2560 (0..text.len() - 1).map(|ix| &text[..ix + 1])
2561}
2562
2563fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
2564 ranges.sort_unstable_by(|a, b| {
2565 a.start
2566 .cmp(&b.start, buffer)
2567 .then_with(|| b.end.cmp(&a.end, buffer))
2568 });
2569
2570 let mut ix = 0;
2571 while ix + 1 < ranges.len() {
2572 let b = ranges[ix + 1].clone();
2573 let a = &mut ranges[ix];
2574 if a.end.cmp(&b.start, buffer).is_gt() {
2575 if a.end.cmp(&b.end, buffer).is_lt() {
2576 a.end = b.end;
2577 }
2578 ranges.remove(ix + 1);
2579 } else {
2580 ix += 1;
2581 }
2582 }
2583}
2584
2585#[cfg(test)]
2586mod tests {
2587 use super::*;
2588 use futures::stream::{self};
2589 use gpui::{Context, TestAppContext};
2590 use indoc::indoc;
2591 use language::{
2592 language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
2593 Point,
2594 };
2595 use language_model::LanguageModelRegistry;
2596 use rand::prelude::*;
2597 use serde::Serialize;
2598 use settings::SettingsStore;
2599 use std::{future, sync::Arc};
2600
2601 #[derive(Serialize)]
2602 pub struct DummyCompletionRequest {
2603 pub name: String,
2604 }
2605
2606 #[gpui::test(iterations = 10)]
2607 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
2608 cx.set_global(cx.update(SettingsStore::test));
2609 cx.update(language_model::LanguageModelRegistry::test);
2610 cx.update(completion::LanguageModelCompletionProvider::test);
2611 cx.update(language_settings::init);
2612
2613 let text = indoc! {"
2614 fn main() {
2615 let x = 0;
2616 for _ in 0..10 {
2617 x += 1;
2618 }
2619 }
2620 "};
2621 let buffer =
2622 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2623 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2624 let range = buffer.read_with(cx, |buffer, cx| {
2625 let snapshot = buffer.snapshot(cx);
2626 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
2627 });
2628 let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, None, cx));
2629
2630 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2631 codegen.update(cx, |codegen, cx| {
2632 codegen.start(
2633 String::new(),
2634 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2635 cx,
2636 )
2637 });
2638
2639 let mut new_text = concat!(
2640 " let mut x = 0;\n",
2641 " while x < 10 {\n",
2642 " x += 1;\n",
2643 " }",
2644 );
2645 while !new_text.is_empty() {
2646 let max_len = cmp::min(new_text.len(), 10);
2647 let len = rng.gen_range(1..=max_len);
2648 let (chunk, suffix) = new_text.split_at(len);
2649 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
2650 new_text = suffix;
2651 cx.background_executor.run_until_parked();
2652 }
2653 drop(chunks_tx);
2654 cx.background_executor.run_until_parked();
2655
2656 assert_eq!(
2657 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2658 indoc! {"
2659 fn main() {
2660 let mut x = 0;
2661 while x < 10 {
2662 x += 1;
2663 }
2664 }
2665 "}
2666 );
2667 }
2668
2669 #[gpui::test(iterations = 10)]
2670 async fn test_autoindent_when_generating_past_indentation(
2671 cx: &mut TestAppContext,
2672 mut rng: StdRng,
2673 ) {
2674 cx.set_global(cx.update(SettingsStore::test));
2675 cx.update(language_settings::init);
2676
2677 let text = indoc! {"
2678 fn main() {
2679 le
2680 }
2681 "};
2682 let buffer =
2683 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2684 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2685 let range = buffer.read_with(cx, |buffer, cx| {
2686 let snapshot = buffer.snapshot(cx);
2687 snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
2688 });
2689 let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, None, cx));
2690
2691 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2692 codegen.update(cx, |codegen, cx| {
2693 codegen.start(
2694 String::new(),
2695 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2696 cx,
2697 )
2698 });
2699
2700 cx.background_executor.run_until_parked();
2701
2702 let mut new_text = concat!(
2703 "t mut x = 0;\n",
2704 "while x < 10 {\n",
2705 " x += 1;\n",
2706 "}", //
2707 );
2708 while !new_text.is_empty() {
2709 let max_len = cmp::min(new_text.len(), 10);
2710 let len = rng.gen_range(1..=max_len);
2711 let (chunk, suffix) = new_text.split_at(len);
2712 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
2713 new_text = suffix;
2714 cx.background_executor.run_until_parked();
2715 }
2716 drop(chunks_tx);
2717 cx.background_executor.run_until_parked();
2718
2719 assert_eq!(
2720 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2721 indoc! {"
2722 fn main() {
2723 let mut x = 0;
2724 while x < 10 {
2725 x += 1;
2726 }
2727 }
2728 "}
2729 );
2730 }
2731
2732 #[gpui::test(iterations = 10)]
2733 async fn test_autoindent_when_generating_before_indentation(
2734 cx: &mut TestAppContext,
2735 mut rng: StdRng,
2736 ) {
2737 cx.update(LanguageModelRegistry::test);
2738 cx.update(completion::LanguageModelCompletionProvider::test);
2739 cx.set_global(cx.update(SettingsStore::test));
2740 cx.update(language_settings::init);
2741
2742 let text = concat!(
2743 "fn main() {\n",
2744 " \n",
2745 "}\n" //
2746 );
2747 let buffer =
2748 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2749 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2750 let range = buffer.read_with(cx, |buffer, cx| {
2751 let snapshot = buffer.snapshot(cx);
2752 snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
2753 });
2754 let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, None, cx));
2755
2756 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2757 codegen.update(cx, |codegen, cx| {
2758 codegen.start(
2759 String::new(),
2760 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2761 cx,
2762 )
2763 });
2764
2765 cx.background_executor.run_until_parked();
2766
2767 let mut new_text = concat!(
2768 "let mut x = 0;\n",
2769 "while x < 10 {\n",
2770 " x += 1;\n",
2771 "}", //
2772 );
2773 while !new_text.is_empty() {
2774 let max_len = cmp::min(new_text.len(), 10);
2775 let len = rng.gen_range(1..=max_len);
2776 let (chunk, suffix) = new_text.split_at(len);
2777 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
2778 new_text = suffix;
2779 cx.background_executor.run_until_parked();
2780 }
2781 drop(chunks_tx);
2782 cx.background_executor.run_until_parked();
2783
2784 assert_eq!(
2785 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2786 indoc! {"
2787 fn main() {
2788 let mut x = 0;
2789 while x < 10 {
2790 x += 1;
2791 }
2792 }
2793 "}
2794 );
2795 }
2796
2797 #[gpui::test(iterations = 10)]
2798 async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
2799 cx.update(LanguageModelRegistry::test);
2800 cx.update(completion::LanguageModelCompletionProvider::test);
2801 cx.set_global(cx.update(SettingsStore::test));
2802 cx.update(language_settings::init);
2803
2804 let text = indoc! {"
2805 func main() {
2806 \tx := 0
2807 \tfor i := 0; i < 10; i++ {
2808 \t\tx++
2809 \t}
2810 }
2811 "};
2812 let buffer = cx.new_model(|cx| Buffer::local(text, cx));
2813 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2814 let range = buffer.read_with(cx, |buffer, cx| {
2815 let snapshot = buffer.snapshot(cx);
2816 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
2817 });
2818 let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, None, cx));
2819
2820 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2821 codegen.update(cx, |codegen, cx| {
2822 codegen.start(
2823 String::new(),
2824 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2825 cx,
2826 )
2827 });
2828
2829 let new_text = concat!(
2830 "func main() {\n",
2831 "\tx := 0\n",
2832 "\tfor x < 10 {\n",
2833 "\t\tx++\n",
2834 "\t}", //
2835 );
2836 chunks_tx.unbounded_send(new_text.to_string()).unwrap();
2837 drop(chunks_tx);
2838 cx.background_executor.run_until_parked();
2839
2840 assert_eq!(
2841 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2842 indoc! {"
2843 func main() {
2844 \tx := 0
2845 \tfor x < 10 {
2846 \t\tx++
2847 \t}
2848 }
2849 "}
2850 );
2851 }
2852
2853 #[gpui::test]
2854 async fn test_strip_invalid_spans_from_codeblock() {
2855 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
2856 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
2857 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
2858 assert_chunks(
2859 "```html\n```js\nLorem ipsum dolor\n```\n```",
2860 "```js\nLorem ipsum dolor\n```",
2861 )
2862 .await;
2863 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
2864 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
2865 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
2866 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
2867
2868 async fn assert_chunks(text: &str, expected_text: &str) {
2869 for chunk_size in 1..=text.len() {
2870 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
2871 .map(|chunk| chunk.unwrap())
2872 .collect::<String>()
2873 .await;
2874 assert_eq!(
2875 actual_text, expected_text,
2876 "failed to strip invalid spans, chunk size: {}",
2877 chunk_size
2878 );
2879 }
2880 }
2881
2882 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
2883 stream::iter(
2884 text.chars()
2885 .collect::<Vec<_>>()
2886 .chunks(size)
2887 .map(|chunk| Ok(chunk.iter().collect::<String>()))
2888 .collect::<Vec<_>>(),
2889 )
2890 }
2891 }
2892
2893 fn rust_lang() -> Language {
2894 Language::new(
2895 LanguageConfig {
2896 name: "Rust".into(),
2897 matcher: LanguageMatcher {
2898 path_suffixes: vec!["rs".to_string()],
2899 ..Default::default()
2900 },
2901 ..Default::default()
2902 },
2903 Some(tree_sitter_rust::language()),
2904 )
2905 .with_indents_query(
2906 r#"
2907 (call_expression) @indent
2908 (field_expression) @indent
2909 (_ "(" ")" @end) @indent
2910 (_ "{" "}" @end) @indent
2911 "#,
2912 )
2913 .unwrap()
2914 }
2915}