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