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 .pl(cx.gutter_dimensions.full_width())
1018 .child(deleted_lines_editor.clone())
1019 .into_any_element()
1020 }),
1021 disposition: BlockDisposition::Above,
1022 });
1023 }
1024
1025 decorations.removed_line_block_ids = editor
1026 .insert_blocks(new_blocks, None, cx)
1027 .into_iter()
1028 .collect();
1029 })
1030 }
1031}
1032
1033struct EditorInlineAssists {
1034 assist_ids: Vec<InlineAssistId>,
1035 scroll_lock: Option<InlineAssistScrollLock>,
1036 highlight_updates: async_watch::Sender<()>,
1037 _update_highlights: Task<Result<()>>,
1038 _subscriptions: Vec<gpui::Subscription>,
1039}
1040
1041struct InlineAssistScrollLock {
1042 assist_id: InlineAssistId,
1043 distance_from_top: f32,
1044}
1045
1046impl EditorInlineAssists {
1047 #[allow(clippy::too_many_arguments)]
1048 fn new(editor: &View<Editor>, cx: &mut WindowContext) -> Self {
1049 let (highlight_updates_tx, mut highlight_updates_rx) = async_watch::channel(());
1050 Self {
1051 assist_ids: Vec::new(),
1052 scroll_lock: None,
1053 highlight_updates: highlight_updates_tx,
1054 _update_highlights: cx.spawn(|mut cx| {
1055 let editor = editor.downgrade();
1056 async move {
1057 while let Ok(()) = highlight_updates_rx.changed().await {
1058 let editor = editor.upgrade().context("editor was dropped")?;
1059 cx.update_global(|assistant: &mut InlineAssistant, cx| {
1060 assistant.update_editor_highlights(&editor, cx);
1061 })?;
1062 }
1063 Ok(())
1064 }
1065 }),
1066 _subscriptions: vec![
1067 cx.observe_release(editor, {
1068 let editor = editor.downgrade();
1069 |_, cx| {
1070 InlineAssistant::update_global(cx, |this, cx| {
1071 this.handle_editor_release(editor, cx);
1072 })
1073 }
1074 }),
1075 cx.observe(editor, move |editor, cx| {
1076 InlineAssistant::update_global(cx, |this, cx| {
1077 this.handle_editor_change(editor, cx)
1078 })
1079 }),
1080 cx.subscribe(editor, move |editor, event, cx| {
1081 InlineAssistant::update_global(cx, |this, cx| {
1082 this.handle_editor_event(editor, event, cx)
1083 })
1084 }),
1085 editor.update(cx, |editor, cx| {
1086 let editor_handle = cx.view().downgrade();
1087 editor.register_action(
1088 move |_: &editor::actions::Newline, cx: &mut WindowContext| {
1089 InlineAssistant::update_global(cx, |this, cx| {
1090 if let Some(editor) = editor_handle.upgrade() {
1091 this.handle_editor_newline(editor, cx)
1092 }
1093 })
1094 },
1095 )
1096 }),
1097 editor.update(cx, |editor, cx| {
1098 let editor_handle = cx.view().downgrade();
1099 editor.register_action(
1100 move |_: &editor::actions::Cancel, cx: &mut WindowContext| {
1101 InlineAssistant::update_global(cx, |this, cx| {
1102 if let Some(editor) = editor_handle.upgrade() {
1103 this.handle_editor_cancel(editor, cx)
1104 }
1105 })
1106 },
1107 )
1108 }),
1109 ],
1110 }
1111 }
1112}
1113
1114struct InlineAssistGroup {
1115 assist_ids: Vec<InlineAssistId>,
1116 linked: bool,
1117 active_assist_id: Option<InlineAssistId>,
1118}
1119
1120impl InlineAssistGroup {
1121 fn new() -> Self {
1122 Self {
1123 assist_ids: Vec::new(),
1124 linked: true,
1125 active_assist_id: None,
1126 }
1127 }
1128}
1129
1130fn build_assist_editor_renderer(editor: &View<PromptEditor>) -> RenderBlock {
1131 let editor = editor.clone();
1132 Box::new(move |cx: &mut BlockContext| {
1133 *editor.read(cx).gutter_dimensions.lock() = *cx.gutter_dimensions;
1134 editor.clone().into_any_element()
1135 })
1136}
1137
1138#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1139pub enum InitialInsertion {
1140 NewlineBefore,
1141 NewlineAfter,
1142}
1143
1144#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1145pub struct InlineAssistId(usize);
1146
1147impl InlineAssistId {
1148 fn post_inc(&mut self) -> InlineAssistId {
1149 let id = *self;
1150 self.0 += 1;
1151 id
1152 }
1153}
1154
1155#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1156struct InlineAssistGroupId(usize);
1157
1158impl InlineAssistGroupId {
1159 fn post_inc(&mut self) -> InlineAssistGroupId {
1160 let id = *self;
1161 self.0 += 1;
1162 id
1163 }
1164}
1165
1166enum PromptEditorEvent {
1167 StartRequested,
1168 StopRequested,
1169 ConfirmRequested,
1170 CancelRequested,
1171 DismissRequested,
1172}
1173
1174struct PromptEditor {
1175 id: InlineAssistId,
1176 fs: Arc<dyn Fs>,
1177 editor: View<Editor>,
1178 edited_since_done: bool,
1179 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1180 prompt_history: VecDeque<String>,
1181 prompt_history_ix: Option<usize>,
1182 pending_prompt: String,
1183 codegen: Model<Codegen>,
1184 _codegen_subscription: Subscription,
1185 editor_subscriptions: Vec<Subscription>,
1186 pending_token_count: Task<Result<()>>,
1187 token_count: Option<usize>,
1188 _token_count_subscriptions: Vec<Subscription>,
1189 workspace: Option<WeakView<Workspace>>,
1190}
1191
1192impl EventEmitter<PromptEditorEvent> for PromptEditor {}
1193
1194impl Render for PromptEditor {
1195 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1196 let gutter_dimensions = *self.gutter_dimensions.lock();
1197 let status = &self.codegen.read(cx).status;
1198 let buttons = match status {
1199 CodegenStatus::Idle => {
1200 vec![
1201 IconButton::new("cancel", IconName::Close)
1202 .icon_color(Color::Muted)
1203 .shape(IconButtonShape::Square)
1204 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1205 .on_click(
1206 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1207 ),
1208 IconButton::new("start", IconName::SparkleAlt)
1209 .icon_color(Color::Muted)
1210 .shape(IconButtonShape::Square)
1211 .tooltip(|cx| Tooltip::for_action("Transform", &menu::Confirm, cx))
1212 .on_click(
1213 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
1214 ),
1215 ]
1216 }
1217 CodegenStatus::Pending => {
1218 vec![
1219 IconButton::new("cancel", IconName::Close)
1220 .icon_color(Color::Muted)
1221 .shape(IconButtonShape::Square)
1222 .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
1223 .on_click(
1224 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1225 ),
1226 IconButton::new("stop", IconName::Stop)
1227 .icon_color(Color::Error)
1228 .shape(IconButtonShape::Square)
1229 .tooltip(|cx| {
1230 Tooltip::with_meta(
1231 "Interrupt Transformation",
1232 Some(&menu::Cancel),
1233 "Changes won't be discarded",
1234 cx,
1235 )
1236 })
1237 .on_click(
1238 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
1239 ),
1240 ]
1241 }
1242 CodegenStatus::Error(_) | CodegenStatus::Done => {
1243 vec![
1244 IconButton::new("cancel", IconName::Close)
1245 .icon_color(Color::Muted)
1246 .shape(IconButtonShape::Square)
1247 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1248 .on_click(
1249 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1250 ),
1251 if self.edited_since_done || matches!(status, CodegenStatus::Error(_)) {
1252 IconButton::new("restart", IconName::RotateCw)
1253 .icon_color(Color::Info)
1254 .shape(IconButtonShape::Square)
1255 .tooltip(|cx| {
1256 Tooltip::with_meta(
1257 "Restart Transformation",
1258 Some(&menu::Confirm),
1259 "Changes will be discarded",
1260 cx,
1261 )
1262 })
1263 .on_click(cx.listener(|_, _, cx| {
1264 cx.emit(PromptEditorEvent::StartRequested);
1265 }))
1266 } else {
1267 IconButton::new("confirm", IconName::Check)
1268 .icon_color(Color::Info)
1269 .shape(IconButtonShape::Square)
1270 .tooltip(|cx| Tooltip::for_action("Confirm Assist", &menu::Confirm, cx))
1271 .on_click(cx.listener(|_, _, cx| {
1272 cx.emit(PromptEditorEvent::ConfirmRequested);
1273 }))
1274 },
1275 ]
1276 }
1277 };
1278
1279 h_flex()
1280 .bg(cx.theme().colors().editor_background)
1281 .border_y_1()
1282 .border_color(cx.theme().status().info_border)
1283 .size_full()
1284 .py(cx.line_height() / 2.)
1285 .on_action(cx.listener(Self::confirm))
1286 .on_action(cx.listener(Self::cancel))
1287 .on_action(cx.listener(Self::move_up))
1288 .on_action(cx.listener(Self::move_down))
1289 .child(
1290 h_flex()
1291 .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
1292 .justify_center()
1293 .gap_2()
1294 .child(
1295 ModelSelector::new(
1296 self.fs.clone(),
1297 IconButton::new("context", IconName::SlidersAlt)
1298 .shape(IconButtonShape::Square)
1299 .icon_size(IconSize::Small)
1300 .icon_color(Color::Muted)
1301 .tooltip(move |cx| {
1302 Tooltip::with_meta(
1303 format!(
1304 "Using {}",
1305 LanguageModelRegistry::read_global(cx)
1306 .active_model()
1307 .map(|model| model.name().0)
1308 .unwrap_or_else(|| "No model selected".into()),
1309 ),
1310 None,
1311 "Change Model",
1312 cx,
1313 )
1314 }),
1315 )
1316 .with_info_text(
1317 "Inline edits use context\n\
1318 from the currently selected\n\
1319 assistant panel tab.",
1320 ),
1321 )
1322 .children(
1323 if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
1324 let error_message = SharedString::from(error.to_string());
1325 Some(
1326 div()
1327 .id("error")
1328 .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
1329 .child(
1330 Icon::new(IconName::XCircle)
1331 .size(IconSize::Small)
1332 .color(Color::Error),
1333 ),
1334 )
1335 } else {
1336 None
1337 },
1338 ),
1339 )
1340 .child(div().flex_1().child(self.render_prompt_editor(cx)))
1341 .child(
1342 h_flex()
1343 .gap_2()
1344 .pr_6()
1345 .children(self.render_token_count(cx))
1346 .children(buttons),
1347 )
1348 }
1349}
1350
1351impl FocusableView for PromptEditor {
1352 fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
1353 self.editor.focus_handle(cx)
1354 }
1355}
1356
1357impl PromptEditor {
1358 const MAX_LINES: u8 = 8;
1359
1360 #[allow(clippy::too_many_arguments)]
1361 fn new(
1362 id: InlineAssistId,
1363 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1364 prompt_history: VecDeque<String>,
1365 prompt_buffer: Model<MultiBuffer>,
1366 codegen: Model<Codegen>,
1367 parent_editor: &View<Editor>,
1368 assistant_panel: Option<&View<AssistantPanel>>,
1369 workspace: Option<WeakView<Workspace>>,
1370 fs: Arc<dyn Fs>,
1371 cx: &mut ViewContext<Self>,
1372 ) -> Self {
1373 let prompt_editor = cx.new_view(|cx| {
1374 let mut editor = Editor::new(
1375 EditorMode::AutoHeight {
1376 max_lines: Self::MAX_LINES as usize,
1377 },
1378 prompt_buffer,
1379 None,
1380 false,
1381 cx,
1382 );
1383 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1384 // Since the prompt editors for all inline assistants are linked,
1385 // always show the cursor (even when it isn't focused) because
1386 // typing in one will make what you typed appear in all of them.
1387 editor.set_show_cursor_when_unfocused(true, cx);
1388 editor.set_placeholder_text("Add a prompt…", cx);
1389 editor
1390 });
1391
1392 let mut token_count_subscriptions = Vec::new();
1393 token_count_subscriptions
1394 .push(cx.subscribe(parent_editor, Self::handle_parent_editor_event));
1395 if let Some(assistant_panel) = assistant_panel {
1396 token_count_subscriptions
1397 .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
1398 }
1399
1400 let mut this = Self {
1401 id,
1402 editor: prompt_editor,
1403 edited_since_done: false,
1404 gutter_dimensions,
1405 prompt_history,
1406 prompt_history_ix: None,
1407 pending_prompt: String::new(),
1408 _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
1409 editor_subscriptions: Vec::new(),
1410 codegen,
1411 fs,
1412 pending_token_count: Task::ready(Ok(())),
1413 token_count: None,
1414 _token_count_subscriptions: token_count_subscriptions,
1415 workspace,
1416 };
1417 this.count_tokens(cx);
1418 this.subscribe_to_editor(cx);
1419 this
1420 }
1421
1422 fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
1423 self.editor_subscriptions.clear();
1424 self.editor_subscriptions
1425 .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
1426 }
1427
1428 fn set_show_cursor_when_unfocused(
1429 &mut self,
1430 show_cursor_when_unfocused: bool,
1431 cx: &mut ViewContext<Self>,
1432 ) {
1433 self.editor.update(cx, |editor, cx| {
1434 editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
1435 });
1436 }
1437
1438 fn unlink(&mut self, cx: &mut ViewContext<Self>) {
1439 let prompt = self.prompt(cx);
1440 let focus = self.editor.focus_handle(cx).contains_focused(cx);
1441 self.editor = cx.new_view(|cx| {
1442 let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
1443 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1444 editor.set_placeholder_text("Add a prompt…", cx);
1445 editor.set_text(prompt, cx);
1446 if focus {
1447 editor.focus(cx);
1448 }
1449 editor
1450 });
1451 self.subscribe_to_editor(cx);
1452 }
1453
1454 fn prompt(&self, cx: &AppContext) -> String {
1455 self.editor.read(cx).text(cx)
1456 }
1457
1458 fn handle_parent_editor_event(
1459 &mut self,
1460 _: View<Editor>,
1461 event: &EditorEvent,
1462 cx: &mut ViewContext<Self>,
1463 ) {
1464 if let EditorEvent::BufferEdited { .. } = event {
1465 self.count_tokens(cx);
1466 }
1467 }
1468
1469 fn handle_assistant_panel_event(
1470 &mut self,
1471 _: View<AssistantPanel>,
1472 event: &AssistantPanelEvent,
1473 cx: &mut ViewContext<Self>,
1474 ) {
1475 let AssistantPanelEvent::ContextEdited { .. } = event;
1476 self.count_tokens(cx);
1477 }
1478
1479 fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
1480 let assist_id = self.id;
1481 self.pending_token_count = cx.spawn(|this, mut cx| async move {
1482 cx.background_executor().timer(Duration::from_secs(1)).await;
1483 let token_count = cx
1484 .update_global(|inline_assistant: &mut InlineAssistant, cx| {
1485 let assist = inline_assistant
1486 .assists
1487 .get(&assist_id)
1488 .context("assist not found")?;
1489 anyhow::Ok(assist.count_tokens(cx))
1490 })??
1491 .await?;
1492
1493 this.update(&mut cx, |this, cx| {
1494 this.token_count = Some(token_count);
1495 cx.notify();
1496 })
1497 })
1498 }
1499
1500 fn handle_prompt_editor_events(
1501 &mut self,
1502 _: View<Editor>,
1503 event: &EditorEvent,
1504 cx: &mut ViewContext<Self>,
1505 ) {
1506 match event {
1507 EditorEvent::Edited { .. } => {
1508 let prompt = self.editor.read(cx).text(cx);
1509 if self
1510 .prompt_history_ix
1511 .map_or(true, |ix| self.prompt_history[ix] != prompt)
1512 {
1513 self.prompt_history_ix.take();
1514 self.pending_prompt = prompt;
1515 }
1516
1517 self.edited_since_done = true;
1518 cx.notify();
1519 }
1520 EditorEvent::BufferEdited => {
1521 self.count_tokens(cx);
1522 }
1523 _ => {}
1524 }
1525 }
1526
1527 fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
1528 match &self.codegen.read(cx).status {
1529 CodegenStatus::Idle => {
1530 self.editor
1531 .update(cx, |editor, _| editor.set_read_only(false));
1532 }
1533 CodegenStatus::Pending => {
1534 self.editor
1535 .update(cx, |editor, _| editor.set_read_only(true));
1536 }
1537 CodegenStatus::Done | CodegenStatus::Error(_) => {
1538 self.edited_since_done = false;
1539 self.editor
1540 .update(cx, |editor, _| editor.set_read_only(false));
1541 }
1542 }
1543 }
1544
1545 fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
1546 match &self.codegen.read(cx).status {
1547 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1548 cx.emit(PromptEditorEvent::CancelRequested);
1549 }
1550 CodegenStatus::Pending => {
1551 cx.emit(PromptEditorEvent::StopRequested);
1552 }
1553 }
1554 }
1555
1556 fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
1557 match &self.codegen.read(cx).status {
1558 CodegenStatus::Idle => {
1559 cx.emit(PromptEditorEvent::StartRequested);
1560 }
1561 CodegenStatus::Pending => {
1562 cx.emit(PromptEditorEvent::DismissRequested);
1563 }
1564 CodegenStatus::Done | CodegenStatus::Error(_) => {
1565 if self.edited_since_done {
1566 cx.emit(PromptEditorEvent::StartRequested);
1567 } else {
1568 cx.emit(PromptEditorEvent::ConfirmRequested);
1569 }
1570 }
1571 }
1572 }
1573
1574 fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
1575 if let Some(ix) = self.prompt_history_ix {
1576 if ix > 0 {
1577 self.prompt_history_ix = Some(ix - 1);
1578 let prompt = self.prompt_history[ix - 1].as_str();
1579 self.editor.update(cx, |editor, cx| {
1580 editor.set_text(prompt, cx);
1581 editor.move_to_beginning(&Default::default(), cx);
1582 });
1583 }
1584 } else if !self.prompt_history.is_empty() {
1585 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1586 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1587 self.editor.update(cx, |editor, cx| {
1588 editor.set_text(prompt, cx);
1589 editor.move_to_beginning(&Default::default(), cx);
1590 });
1591 }
1592 }
1593
1594 fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
1595 if let Some(ix) = self.prompt_history_ix {
1596 if ix < self.prompt_history.len() - 1 {
1597 self.prompt_history_ix = Some(ix + 1);
1598 let prompt = self.prompt_history[ix + 1].as_str();
1599 self.editor.update(cx, |editor, cx| {
1600 editor.set_text(prompt, cx);
1601 editor.move_to_end(&Default::default(), cx)
1602 });
1603 } else {
1604 self.prompt_history_ix = None;
1605 let prompt = self.pending_prompt.as_str();
1606 self.editor.update(cx, |editor, cx| {
1607 editor.set_text(prompt, cx);
1608 editor.move_to_end(&Default::default(), cx)
1609 });
1610 }
1611 }
1612 }
1613
1614 fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
1615 let model = LanguageModelRegistry::read_global(cx).active_model()?;
1616 let token_count = self.token_count?;
1617 let max_token_count = model.max_token_count();
1618
1619 let remaining_tokens = max_token_count as isize - token_count as isize;
1620 let token_count_color = if remaining_tokens <= 0 {
1621 Color::Error
1622 } else if token_count as f32 / max_token_count as f32 >= 0.8 {
1623 Color::Warning
1624 } else {
1625 Color::Muted
1626 };
1627
1628 let mut token_count = h_flex()
1629 .id("token_count")
1630 .gap_0p5()
1631 .child(
1632 Label::new(humanize_token_count(token_count))
1633 .size(LabelSize::Small)
1634 .color(token_count_color),
1635 )
1636 .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
1637 .child(
1638 Label::new(humanize_token_count(max_token_count))
1639 .size(LabelSize::Small)
1640 .color(Color::Muted),
1641 );
1642 if let Some(workspace) = self.workspace.clone() {
1643 token_count = token_count
1644 .tooltip(|cx| {
1645 Tooltip::with_meta(
1646 "Tokens Used by Inline Assistant",
1647 None,
1648 "Click to Open Assistant Panel",
1649 cx,
1650 )
1651 })
1652 .cursor_pointer()
1653 .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
1654 .on_click(move |_, cx| {
1655 cx.stop_propagation();
1656 workspace
1657 .update(cx, |workspace, cx| {
1658 workspace.focus_panel::<AssistantPanel>(cx)
1659 })
1660 .ok();
1661 });
1662 } else {
1663 token_count = token_count
1664 .cursor_default()
1665 .tooltip(|cx| Tooltip::text("Tokens Used by Inline Assistant", cx));
1666 }
1667
1668 Some(token_count)
1669 }
1670
1671 fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1672 let settings = ThemeSettings::get_global(cx);
1673 let text_style = TextStyle {
1674 color: if self.editor.read(cx).read_only(cx) {
1675 cx.theme().colors().text_disabled
1676 } else {
1677 cx.theme().colors().text
1678 },
1679 font_family: settings.ui_font.family.clone(),
1680 font_features: settings.ui_font.features.clone(),
1681 font_fallbacks: settings.ui_font.fallbacks.clone(),
1682 font_size: rems(0.875).into(),
1683 font_weight: settings.ui_font.weight,
1684 line_height: relative(1.3),
1685 ..Default::default()
1686 };
1687 EditorElement::new(
1688 &self.editor,
1689 EditorStyle {
1690 background: cx.theme().colors().editor_background,
1691 local_player: cx.theme().players().local(),
1692 text: text_style,
1693 ..Default::default()
1694 },
1695 )
1696 }
1697}
1698
1699struct InlineAssist {
1700 group_id: InlineAssistGroupId,
1701 range: Range<Anchor>,
1702 editor: WeakView<Editor>,
1703 decorations: Option<InlineAssistDecorations>,
1704 codegen: Model<Codegen>,
1705 _subscriptions: Vec<Subscription>,
1706 workspace: Option<WeakView<Workspace>>,
1707 include_context: bool,
1708}
1709
1710impl InlineAssist {
1711 #[allow(clippy::too_many_arguments)]
1712 fn new(
1713 assist_id: InlineAssistId,
1714 group_id: InlineAssistGroupId,
1715 include_context: bool,
1716 editor: &View<Editor>,
1717 prompt_editor: &View<PromptEditor>,
1718 prompt_block_id: CustomBlockId,
1719 end_block_id: CustomBlockId,
1720 range: Range<Anchor>,
1721 codegen: Model<Codegen>,
1722 workspace: Option<WeakView<Workspace>>,
1723 cx: &mut WindowContext,
1724 ) -> Self {
1725 let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
1726 InlineAssist {
1727 group_id,
1728 include_context,
1729 editor: editor.downgrade(),
1730 decorations: Some(InlineAssistDecorations {
1731 prompt_block_id,
1732 prompt_editor: prompt_editor.clone(),
1733 removed_line_block_ids: HashSet::default(),
1734 end_block_id,
1735 }),
1736 range,
1737 codegen: codegen.clone(),
1738 workspace: workspace.clone(),
1739 _subscriptions: vec![
1740 cx.on_focus_in(&prompt_editor_focus_handle, move |cx| {
1741 InlineAssistant::update_global(cx, |this, cx| {
1742 this.handle_prompt_editor_focus_in(assist_id, cx)
1743 })
1744 }),
1745 cx.on_focus_out(&prompt_editor_focus_handle, move |_, cx| {
1746 InlineAssistant::update_global(cx, |this, cx| {
1747 this.handle_prompt_editor_focus_out(assist_id, cx)
1748 })
1749 }),
1750 cx.subscribe(prompt_editor, |prompt_editor, event, cx| {
1751 InlineAssistant::update_global(cx, |this, cx| {
1752 this.handle_prompt_editor_event(prompt_editor, event, cx)
1753 })
1754 }),
1755 cx.observe(&codegen, {
1756 let editor = editor.downgrade();
1757 move |_, cx| {
1758 if let Some(editor) = editor.upgrade() {
1759 InlineAssistant::update_global(cx, |this, cx| {
1760 if let Some(editor_assists) =
1761 this.assists_by_editor.get(&editor.downgrade())
1762 {
1763 editor_assists.highlight_updates.send(()).ok();
1764 }
1765
1766 this.update_editor_blocks(&editor, assist_id, cx);
1767 })
1768 }
1769 }
1770 }),
1771 cx.subscribe(&codegen, move |codegen, event, cx| {
1772 InlineAssistant::update_global(cx, |this, cx| match event {
1773 CodegenEvent::Undone => this.finish_assist(assist_id, false, cx),
1774 CodegenEvent::Finished => {
1775 let assist = if let Some(assist) = this.assists.get(&assist_id) {
1776 assist
1777 } else {
1778 return;
1779 };
1780
1781 if let CodegenStatus::Error(error) = &codegen.read(cx).status {
1782 if assist.decorations.is_none() {
1783 if let Some(workspace) = assist
1784 .workspace
1785 .as_ref()
1786 .and_then(|workspace| workspace.upgrade())
1787 {
1788 let error = format!("Inline assistant error: {}", error);
1789 workspace.update(cx, |workspace, cx| {
1790 struct InlineAssistantError;
1791
1792 let id =
1793 NotificationId::identified::<InlineAssistantError>(
1794 assist_id.0,
1795 );
1796
1797 workspace.show_toast(Toast::new(id, error), cx);
1798 })
1799 }
1800 }
1801 }
1802
1803 if assist.decorations.is_none() {
1804 this.finish_assist(assist_id, false, cx);
1805 }
1806 }
1807 })
1808 }),
1809 ],
1810 }
1811 }
1812
1813 fn user_prompt(&self, cx: &AppContext) -> Option<String> {
1814 let decorations = self.decorations.as_ref()?;
1815 Some(decorations.prompt_editor.read(cx).prompt(cx))
1816 }
1817
1818 fn assistant_panel_context(&self, cx: &WindowContext) -> Option<LanguageModelRequest> {
1819 if self.include_context {
1820 let workspace = self.workspace.as_ref()?;
1821 let workspace = workspace.upgrade()?.read(cx);
1822 let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
1823 Some(
1824 assistant_panel
1825 .read(cx)
1826 .active_context(cx)?
1827 .read(cx)
1828 .to_completion_request(cx),
1829 )
1830 } else {
1831 None
1832 }
1833 }
1834
1835 pub fn count_tokens(&self, cx: &WindowContext) -> BoxFuture<'static, Result<usize>> {
1836 let Some(user_prompt) = self.user_prompt(cx) else {
1837 return future::ready(Err(anyhow!("no user prompt"))).boxed();
1838 };
1839 let assistant_panel_context = self.assistant_panel_context(cx);
1840 self.codegen.read(cx).count_tokens(
1841 self.range.clone(),
1842 user_prompt,
1843 assistant_panel_context,
1844 cx,
1845 )
1846 }
1847}
1848
1849struct InlineAssistDecorations {
1850 prompt_block_id: CustomBlockId,
1851 prompt_editor: View<PromptEditor>,
1852 removed_line_block_ids: HashSet<CustomBlockId>,
1853 end_block_id: CustomBlockId,
1854}
1855
1856#[derive(Debug)]
1857pub enum CodegenEvent {
1858 Finished,
1859 Undone,
1860}
1861
1862pub struct Codegen {
1863 buffer: Model<MultiBuffer>,
1864 old_buffer: Model<Buffer>,
1865 snapshot: MultiBufferSnapshot,
1866 edit_position: Option<Anchor>,
1867 last_equal_ranges: Vec<Range<Anchor>>,
1868 transaction_id: Option<TransactionId>,
1869 status: CodegenStatus,
1870 generation: Task<()>,
1871 diff: Diff,
1872 telemetry: Option<Arc<Telemetry>>,
1873 _subscription: gpui::Subscription,
1874 initial_insertion: Option<InitialInsertion>,
1875}
1876
1877enum CodegenStatus {
1878 Idle,
1879 Pending,
1880 Done,
1881 Error(anyhow::Error),
1882}
1883
1884#[derive(Default)]
1885struct Diff {
1886 task: Option<Task<()>>,
1887 should_update: bool,
1888 deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1889 inserted_row_ranges: Vec<RangeInclusive<Anchor>>,
1890}
1891
1892impl EventEmitter<CodegenEvent> for Codegen {}
1893
1894impl Codegen {
1895 pub fn new(
1896 buffer: Model<MultiBuffer>,
1897 range: Range<Anchor>,
1898 initial_insertion: Option<InitialInsertion>,
1899 telemetry: Option<Arc<Telemetry>>,
1900 cx: &mut ModelContext<Self>,
1901 ) -> Self {
1902 let snapshot = buffer.read(cx).snapshot(cx);
1903
1904 let (old_buffer, _, _) = buffer
1905 .read(cx)
1906 .range_to_buffer_ranges(range.clone(), cx)
1907 .pop()
1908 .unwrap();
1909 let old_buffer = cx.new_model(|cx| {
1910 let old_buffer = old_buffer.read(cx);
1911 let text = old_buffer.as_rope().clone();
1912 let line_ending = old_buffer.line_ending();
1913 let language = old_buffer.language().cloned();
1914 let language_registry = old_buffer.language_registry();
1915
1916 let mut buffer = Buffer::local_normalized(text, line_ending, cx);
1917 buffer.set_language(language, cx);
1918 if let Some(language_registry) = language_registry {
1919 buffer.set_language_registry(language_registry)
1920 }
1921 buffer
1922 });
1923
1924 Self {
1925 buffer: buffer.clone(),
1926 old_buffer,
1927 edit_position: None,
1928 snapshot,
1929 last_equal_ranges: Default::default(),
1930 transaction_id: None,
1931 status: CodegenStatus::Idle,
1932 generation: Task::ready(()),
1933 diff: Diff::default(),
1934 telemetry,
1935 _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
1936 initial_insertion,
1937 }
1938 }
1939
1940 fn handle_buffer_event(
1941 &mut self,
1942 _buffer: Model<MultiBuffer>,
1943 event: &multi_buffer::Event,
1944 cx: &mut ModelContext<Self>,
1945 ) {
1946 if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
1947 if self.transaction_id == Some(*transaction_id) {
1948 self.transaction_id = None;
1949 self.generation = Task::ready(());
1950 cx.emit(CodegenEvent::Undone);
1951 }
1952 }
1953 }
1954
1955 pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
1956 &self.last_equal_ranges
1957 }
1958
1959 pub fn count_tokens(
1960 &self,
1961 edit_range: Range<Anchor>,
1962 user_prompt: String,
1963 assistant_panel_context: Option<LanguageModelRequest>,
1964 cx: &AppContext,
1965 ) -> BoxFuture<'static, Result<usize>> {
1966 if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
1967 let request = self.build_request(user_prompt, assistant_panel_context, edit_range, cx);
1968 model.count_tokens(request, cx)
1969 } else {
1970 future::ready(Err(anyhow!("no active model"))).boxed()
1971 }
1972 }
1973
1974 pub fn start(
1975 &mut self,
1976 mut edit_range: Range<Anchor>,
1977 user_prompt: String,
1978 assistant_panel_context: Option<LanguageModelRequest>,
1979 cx: &mut ModelContext<Self>,
1980 ) -> Result<()> {
1981 let model = LanguageModelRegistry::read_global(cx)
1982 .active_model()
1983 .context("no active model")?;
1984
1985 self.undo(cx);
1986
1987 // Handle initial insertion
1988 self.transaction_id = if let Some(initial_insertion) = self.initial_insertion {
1989 self.buffer.update(cx, |buffer, cx| {
1990 buffer.start_transaction(cx);
1991 let offset = edit_range.start.to_offset(&self.snapshot);
1992 let edit_position;
1993 match initial_insertion {
1994 InitialInsertion::NewlineBefore => {
1995 buffer.edit([(offset..offset, "\n\n")], None, cx);
1996 self.snapshot = buffer.snapshot(cx);
1997 edit_position = self.snapshot.anchor_after(offset + 1);
1998 }
1999 InitialInsertion::NewlineAfter => {
2000 buffer.edit([(offset..offset, "\n")], None, cx);
2001 self.snapshot = buffer.snapshot(cx);
2002 edit_position = self.snapshot.anchor_after(offset);
2003 }
2004 }
2005 self.edit_position = Some(edit_position);
2006 edit_range = edit_position.bias_left(&self.snapshot)..edit_position;
2007 buffer.end_transaction(cx)
2008 })
2009 } else {
2010 self.edit_position = Some(edit_range.start.bias_right(&self.snapshot));
2011 None
2012 };
2013
2014 let telemetry_id = model.telemetry_id();
2015 let chunks: LocalBoxFuture<Result<BoxStream<Result<String>>>> = if user_prompt
2016 .trim()
2017 .to_lowercase()
2018 == "delete"
2019 {
2020 async { Ok(stream::empty().boxed()) }.boxed_local()
2021 } else {
2022 let request =
2023 self.build_request(user_prompt, assistant_panel_context, edit_range.clone(), cx);
2024 let chunks =
2025 cx.spawn(|_, cx| async move { model.stream_completion(request, &cx).await });
2026 async move { Ok(chunks.await?.boxed()) }.boxed_local()
2027 };
2028 self.handle_stream(telemetry_id, edit_range, chunks, cx);
2029 Ok(())
2030 }
2031
2032 fn build_request(
2033 &self,
2034 user_prompt: String,
2035 assistant_panel_context: Option<LanguageModelRequest>,
2036 edit_range: Range<Anchor>,
2037 cx: &AppContext,
2038 ) -> LanguageModelRequest {
2039 let buffer = self.buffer.read(cx).snapshot(cx);
2040 let language = buffer.language_at(edit_range.start);
2041 let language_name = if let Some(language) = language.as_ref() {
2042 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
2043 None
2044 } else {
2045 Some(language.name())
2046 }
2047 } else {
2048 None
2049 };
2050
2051 // Higher Temperature increases the randomness of model outputs.
2052 // If Markdown or No Language is Known, increase the randomness for more creative output
2053 // If Code, decrease temperature to get more deterministic outputs
2054 let temperature = if let Some(language) = language_name.clone() {
2055 if language.as_ref() == "Markdown" {
2056 1.0
2057 } else {
2058 0.5
2059 }
2060 } else {
2061 1.0
2062 };
2063
2064 let language_name = language_name.as_deref();
2065 let start = buffer.point_to_buffer_offset(edit_range.start);
2066 let end = buffer.point_to_buffer_offset(edit_range.end);
2067 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
2068 let (start_buffer, start_buffer_offset) = start;
2069 let (end_buffer, end_buffer_offset) = end;
2070 if start_buffer.remote_id() == end_buffer.remote_id() {
2071 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
2072 } else {
2073 panic!("invalid transformation range");
2074 }
2075 } else {
2076 panic!("invalid transformation range");
2077 };
2078 let prompt = generate_content_prompt(user_prompt, language_name, buffer, range);
2079
2080 let mut messages = Vec::new();
2081 if let Some(context_request) = assistant_panel_context {
2082 messages = context_request.messages;
2083 }
2084
2085 messages.push(LanguageModelRequestMessage {
2086 role: Role::User,
2087 content: prompt,
2088 });
2089
2090 LanguageModelRequest {
2091 messages,
2092 stop: vec!["|END|>".to_string()],
2093 temperature,
2094 }
2095 }
2096
2097 pub fn handle_stream(
2098 &mut self,
2099 model_telemetry_id: String,
2100 edit_range: Range<Anchor>,
2101 stream: impl 'static + Future<Output = Result<BoxStream<'static, Result<String>>>>,
2102 cx: &mut ModelContext<Self>,
2103 ) {
2104 let snapshot = self.snapshot.clone();
2105 let selected_text = snapshot
2106 .text_for_range(edit_range.start..edit_range.end)
2107 .collect::<Rope>();
2108
2109 let selection_start = edit_range.start.to_point(&snapshot);
2110
2111 // Start with the indentation of the first line in the selection
2112 let mut suggested_line_indent = snapshot
2113 .suggested_indents(selection_start.row..=selection_start.row, cx)
2114 .into_values()
2115 .next()
2116 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
2117
2118 // If the first line in the selection does not have indentation, check the following lines
2119 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
2120 for row in selection_start.row..=edit_range.end.to_point(&snapshot).row {
2121 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
2122 // Prefer tabs if a line in the selection uses tabs as indentation
2123 if line_indent.kind == IndentKind::Tab {
2124 suggested_line_indent.kind = IndentKind::Tab;
2125 break;
2126 }
2127 }
2128 }
2129
2130 let telemetry = self.telemetry.clone();
2131 self.diff = Diff::default();
2132 self.status = CodegenStatus::Pending;
2133 let mut edit_start = edit_range.start.to_offset(&snapshot);
2134 self.generation = cx.spawn(|this, mut cx| {
2135 async move {
2136 let chunks = stream.await;
2137 let generate = async {
2138 let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
2139 let diff: Task<anyhow::Result<()>> =
2140 cx.background_executor().spawn(async move {
2141 let mut response_latency = None;
2142 let request_start = Instant::now();
2143 let diff = async {
2144 let chunks = StripInvalidSpans::new(chunks?);
2145 futures::pin_mut!(chunks);
2146 let mut diff = StreamingDiff::new(selected_text.to_string());
2147
2148 let mut new_text = String::new();
2149 let mut base_indent = None;
2150 let mut line_indent = None;
2151 let mut first_line = true;
2152
2153 while let Some(chunk) = chunks.next().await {
2154 if response_latency.is_none() {
2155 response_latency = Some(request_start.elapsed());
2156 }
2157 let chunk = chunk?;
2158
2159 let mut lines = chunk.split('\n').peekable();
2160 while let Some(line) = lines.next() {
2161 new_text.push_str(line);
2162 if line_indent.is_none() {
2163 if let Some(non_whitespace_ch_ix) =
2164 new_text.find(|ch: char| !ch.is_whitespace())
2165 {
2166 line_indent = Some(non_whitespace_ch_ix);
2167 base_indent = base_indent.or(line_indent);
2168
2169 let line_indent = line_indent.unwrap();
2170 let base_indent = base_indent.unwrap();
2171 let indent_delta =
2172 line_indent as i32 - base_indent as i32;
2173 let mut corrected_indent_len = cmp::max(
2174 0,
2175 suggested_line_indent.len as i32 + indent_delta,
2176 )
2177 as usize;
2178 if first_line {
2179 corrected_indent_len = corrected_indent_len
2180 .saturating_sub(
2181 selection_start.column as usize,
2182 );
2183 }
2184
2185 let indent_char = suggested_line_indent.char();
2186 let mut indent_buffer = [0; 4];
2187 let indent_str =
2188 indent_char.encode_utf8(&mut indent_buffer);
2189 new_text.replace_range(
2190 ..line_indent,
2191 &indent_str.repeat(corrected_indent_len),
2192 );
2193 }
2194 }
2195
2196 if line_indent.is_some() {
2197 hunks_tx.send(diff.push_new(&new_text)).await?;
2198 new_text.clear();
2199 }
2200
2201 if lines.peek().is_some() {
2202 hunks_tx.send(diff.push_new("\n")).await?;
2203 if line_indent.is_none() {
2204 // Don't write out the leading indentation in empty lines on the next line
2205 // This is the case where the above if statement didn't clear the buffer
2206 new_text.clear();
2207 }
2208 line_indent = None;
2209 first_line = false;
2210 }
2211 }
2212 }
2213 hunks_tx.send(diff.push_new(&new_text)).await?;
2214 hunks_tx.send(diff.finish()).await?;
2215
2216 anyhow::Ok(())
2217 };
2218
2219 let result = diff.await;
2220
2221 let error_message =
2222 result.as_ref().err().map(|error| error.to_string());
2223 if let Some(telemetry) = telemetry {
2224 telemetry.report_assistant_event(
2225 None,
2226 telemetry_events::AssistantKind::Inline,
2227 model_telemetry_id,
2228 response_latency,
2229 error_message,
2230 );
2231 }
2232
2233 result?;
2234 Ok(())
2235 });
2236
2237 while let Some(hunks) = hunks_rx.next().await {
2238 this.update(&mut cx, |this, cx| {
2239 this.last_equal_ranges.clear();
2240
2241 let transaction = this.buffer.update(cx, |buffer, cx| {
2242 // Avoid grouping assistant edits with user edits.
2243 buffer.finalize_last_transaction(cx);
2244
2245 buffer.start_transaction(cx);
2246 buffer.edit(
2247 hunks.into_iter().filter_map(|hunk| match hunk {
2248 Hunk::Insert { text } => {
2249 let edit_start = snapshot.anchor_after(edit_start);
2250 Some((edit_start..edit_start, text))
2251 }
2252 Hunk::Remove { len } => {
2253 let edit_end = edit_start + len;
2254 let edit_range = snapshot.anchor_after(edit_start)
2255 ..snapshot.anchor_before(edit_end);
2256 edit_start = edit_end;
2257 Some((edit_range, String::new()))
2258 }
2259 Hunk::Keep { len } => {
2260 let edit_end = edit_start + len;
2261 let edit_range = snapshot.anchor_after(edit_start)
2262 ..snapshot.anchor_before(edit_end);
2263 edit_start = edit_end;
2264 this.last_equal_ranges.push(edit_range);
2265 None
2266 }
2267 }),
2268 None,
2269 cx,
2270 );
2271 this.edit_position = Some(snapshot.anchor_after(edit_start));
2272
2273 buffer.end_transaction(cx)
2274 });
2275
2276 if let Some(transaction) = transaction {
2277 if let Some(first_transaction) = this.transaction_id {
2278 // Group all assistant edits into the first transaction.
2279 this.buffer.update(cx, |buffer, cx| {
2280 buffer.merge_transactions(
2281 transaction,
2282 first_transaction,
2283 cx,
2284 )
2285 });
2286 } else {
2287 this.transaction_id = Some(transaction);
2288 this.buffer.update(cx, |buffer, cx| {
2289 buffer.finalize_last_transaction(cx)
2290 });
2291 }
2292 }
2293
2294 this.update_diff(edit_range.clone(), cx);
2295 cx.notify();
2296 })?;
2297 }
2298
2299 diff.await?;
2300
2301 anyhow::Ok(())
2302 };
2303
2304 let result = generate.await;
2305 this.update(&mut cx, |this, cx| {
2306 this.last_equal_ranges.clear();
2307 if let Err(error) = result {
2308 this.status = CodegenStatus::Error(error);
2309 } else {
2310 this.status = CodegenStatus::Done;
2311 }
2312 cx.emit(CodegenEvent::Finished);
2313 cx.notify();
2314 })
2315 .ok();
2316 }
2317 });
2318 cx.notify();
2319 }
2320
2321 pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
2322 self.last_equal_ranges.clear();
2323 self.status = CodegenStatus::Done;
2324 self.generation = Task::ready(());
2325 cx.emit(CodegenEvent::Finished);
2326 cx.notify();
2327 }
2328
2329 pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
2330 if let Some(transaction_id) = self.transaction_id.take() {
2331 self.buffer
2332 .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
2333 }
2334 }
2335
2336 fn update_diff(&mut self, edit_range: Range<Anchor>, cx: &mut ModelContext<Self>) {
2337 if self.diff.task.is_some() {
2338 self.diff.should_update = true;
2339 } else {
2340 self.diff.should_update = false;
2341
2342 let old_snapshot = self.snapshot.clone();
2343 let old_range = edit_range.to_point(&old_snapshot);
2344 let new_snapshot = self.buffer.read(cx).snapshot(cx);
2345 let new_range = edit_range.to_point(&new_snapshot);
2346
2347 self.diff.task = Some(cx.spawn(|this, mut cx| async move {
2348 let (deleted_row_ranges, inserted_row_ranges) = cx
2349 .background_executor()
2350 .spawn(async move {
2351 let old_text = old_snapshot
2352 .text_for_range(
2353 Point::new(old_range.start.row, 0)
2354 ..Point::new(
2355 old_range.end.row,
2356 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
2357 ),
2358 )
2359 .collect::<String>();
2360 let new_text = new_snapshot
2361 .text_for_range(
2362 Point::new(new_range.start.row, 0)
2363 ..Point::new(
2364 new_range.end.row,
2365 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
2366 ),
2367 )
2368 .collect::<String>();
2369
2370 let mut old_row = old_range.start.row;
2371 let mut new_row = new_range.start.row;
2372 let diff = TextDiff::from_lines(old_text.as_str(), new_text.as_str());
2373
2374 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
2375 let mut inserted_row_ranges = Vec::new();
2376 for change in diff.iter_all_changes() {
2377 let line_count = change.value().lines().count() as u32;
2378 match change.tag() {
2379 similar::ChangeTag::Equal => {
2380 old_row += line_count;
2381 new_row += line_count;
2382 }
2383 similar::ChangeTag::Delete => {
2384 let old_end_row = old_row + line_count - 1;
2385 let new_row =
2386 new_snapshot.anchor_before(Point::new(new_row, 0));
2387
2388 if let Some((_, last_deleted_row_range)) =
2389 deleted_row_ranges.last_mut()
2390 {
2391 if *last_deleted_row_range.end() + 1 == old_row {
2392 *last_deleted_row_range =
2393 *last_deleted_row_range.start()..=old_end_row;
2394 } else {
2395 deleted_row_ranges
2396 .push((new_row, old_row..=old_end_row));
2397 }
2398 } else {
2399 deleted_row_ranges.push((new_row, old_row..=old_end_row));
2400 }
2401
2402 old_row += line_count;
2403 }
2404 similar::ChangeTag::Insert => {
2405 let new_end_row = new_row + line_count - 1;
2406 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
2407 let end = new_snapshot.anchor_before(Point::new(
2408 new_end_row,
2409 new_snapshot.line_len(MultiBufferRow(new_end_row)),
2410 ));
2411 inserted_row_ranges.push(start..=end);
2412 new_row += line_count;
2413 }
2414 }
2415 }
2416
2417 (deleted_row_ranges, inserted_row_ranges)
2418 })
2419 .await;
2420
2421 this.update(&mut cx, |this, cx| {
2422 this.diff.deleted_row_ranges = deleted_row_ranges;
2423 this.diff.inserted_row_ranges = inserted_row_ranges;
2424 this.diff.task = None;
2425 if this.diff.should_update {
2426 this.update_diff(edit_range, cx);
2427 }
2428 cx.notify();
2429 })
2430 .ok();
2431 }));
2432 }
2433 }
2434}
2435
2436struct StripInvalidSpans<T> {
2437 stream: T,
2438 stream_done: bool,
2439 buffer: String,
2440 first_line: bool,
2441 line_end: bool,
2442 starts_with_code_block: bool,
2443}
2444
2445impl<T> StripInvalidSpans<T>
2446where
2447 T: Stream<Item = Result<String>>,
2448{
2449 fn new(stream: T) -> Self {
2450 Self {
2451 stream,
2452 stream_done: false,
2453 buffer: String::new(),
2454 first_line: true,
2455 line_end: false,
2456 starts_with_code_block: false,
2457 }
2458 }
2459}
2460
2461impl<T> Stream for StripInvalidSpans<T>
2462where
2463 T: Stream<Item = Result<String>>,
2464{
2465 type Item = Result<String>;
2466
2467 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
2468 const CODE_BLOCK_DELIMITER: &str = "```";
2469 const CURSOR_SPAN: &str = "<|CURSOR|>";
2470
2471 let this = unsafe { self.get_unchecked_mut() };
2472 loop {
2473 if !this.stream_done {
2474 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
2475 match stream.as_mut().poll_next(cx) {
2476 Poll::Ready(Some(Ok(chunk))) => {
2477 this.buffer.push_str(&chunk);
2478 }
2479 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
2480 Poll::Ready(None) => {
2481 this.stream_done = true;
2482 }
2483 Poll::Pending => return Poll::Pending,
2484 }
2485 }
2486
2487 let mut chunk = String::new();
2488 let mut consumed = 0;
2489 if !this.buffer.is_empty() {
2490 let mut lines = this.buffer.split('\n').enumerate().peekable();
2491 while let Some((line_ix, line)) = lines.next() {
2492 if line_ix > 0 {
2493 this.first_line = false;
2494 }
2495
2496 if this.first_line {
2497 let trimmed_line = line.trim();
2498 if lines.peek().is_some() {
2499 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
2500 consumed += line.len() + 1;
2501 this.starts_with_code_block = true;
2502 continue;
2503 }
2504 } else if trimmed_line.is_empty()
2505 || prefixes(CODE_BLOCK_DELIMITER)
2506 .any(|prefix| trimmed_line.starts_with(prefix))
2507 {
2508 break;
2509 }
2510 }
2511
2512 let line_without_cursor = line.replace(CURSOR_SPAN, "");
2513 if lines.peek().is_some() {
2514 if this.line_end {
2515 chunk.push('\n');
2516 }
2517
2518 chunk.push_str(&line_without_cursor);
2519 this.line_end = true;
2520 consumed += line.len() + 1;
2521 } else if this.stream_done {
2522 if !this.starts_with_code_block
2523 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
2524 {
2525 if this.line_end {
2526 chunk.push('\n');
2527 }
2528
2529 chunk.push_str(&line);
2530 }
2531
2532 consumed += line.len();
2533 } else {
2534 let trimmed_line = line.trim();
2535 if trimmed_line.is_empty()
2536 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
2537 || prefixes(CODE_BLOCK_DELIMITER)
2538 .any(|prefix| trimmed_line.ends_with(prefix))
2539 {
2540 break;
2541 } else {
2542 if this.line_end {
2543 chunk.push('\n');
2544 this.line_end = false;
2545 }
2546
2547 chunk.push_str(&line_without_cursor);
2548 consumed += line.len();
2549 }
2550 }
2551 }
2552 }
2553
2554 this.buffer = this.buffer.split_off(consumed);
2555 if !chunk.is_empty() {
2556 return Poll::Ready(Some(Ok(chunk)));
2557 } else if this.stream_done {
2558 return Poll::Ready(None);
2559 }
2560 }
2561 }
2562}
2563
2564fn prefixes(text: &str) -> impl Iterator<Item = &str> {
2565 (0..text.len() - 1).map(|ix| &text[..ix + 1])
2566}
2567
2568fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
2569 ranges.sort_unstable_by(|a, b| {
2570 a.start
2571 .cmp(&b.start, buffer)
2572 .then_with(|| b.end.cmp(&a.end, buffer))
2573 });
2574
2575 let mut ix = 0;
2576 while ix + 1 < ranges.len() {
2577 let b = ranges[ix + 1].clone();
2578 let a = &mut ranges[ix];
2579 if a.end.cmp(&b.start, buffer).is_gt() {
2580 if a.end.cmp(&b.end, buffer).is_lt() {
2581 a.end = b.end;
2582 }
2583 ranges.remove(ix + 1);
2584 } else {
2585 ix += 1;
2586 }
2587 }
2588}
2589
2590#[cfg(test)]
2591mod tests {
2592 use super::*;
2593 use futures::stream::{self};
2594 use gpui::{Context, TestAppContext};
2595 use indoc::indoc;
2596 use language::{
2597 language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
2598 Point,
2599 };
2600 use language_model::LanguageModelRegistry;
2601 use rand::prelude::*;
2602 use serde::Serialize;
2603 use settings::SettingsStore;
2604 use std::{future, sync::Arc};
2605
2606 #[derive(Serialize)]
2607 pub struct DummyCompletionRequest {
2608 pub name: String,
2609 }
2610
2611 #[gpui::test(iterations = 10)]
2612 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
2613 cx.set_global(cx.update(SettingsStore::test));
2614 cx.update(language_model::LanguageModelRegistry::test);
2615 cx.update(language_settings::init);
2616
2617 let text = indoc! {"
2618 fn main() {
2619 let x = 0;
2620 for _ in 0..10 {
2621 x += 1;
2622 }
2623 }
2624 "};
2625 let buffer =
2626 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2627 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2628 let range = buffer.read_with(cx, |buffer, cx| {
2629 let snapshot = buffer.snapshot(cx);
2630 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
2631 });
2632 let codegen =
2633 cx.new_model(|cx| Codegen::new(buffer.clone(), range.clone(), None, None, cx));
2634
2635 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2636 codegen.update(cx, |codegen, cx| {
2637 codegen.handle_stream(
2638 String::new(),
2639 range,
2640 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2641 cx,
2642 )
2643 });
2644
2645 let mut new_text = concat!(
2646 " let mut x = 0;\n",
2647 " while x < 10 {\n",
2648 " x += 1;\n",
2649 " }",
2650 );
2651 while !new_text.is_empty() {
2652 let max_len = cmp::min(new_text.len(), 10);
2653 let len = rng.gen_range(1..=max_len);
2654 let (chunk, suffix) = new_text.split_at(len);
2655 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
2656 new_text = suffix;
2657 cx.background_executor.run_until_parked();
2658 }
2659 drop(chunks_tx);
2660 cx.background_executor.run_until_parked();
2661
2662 assert_eq!(
2663 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2664 indoc! {"
2665 fn main() {
2666 let mut x = 0;
2667 while x < 10 {
2668 x += 1;
2669 }
2670 }
2671 "}
2672 );
2673 }
2674
2675 #[gpui::test(iterations = 10)]
2676 async fn test_autoindent_when_generating_past_indentation(
2677 cx: &mut TestAppContext,
2678 mut rng: StdRng,
2679 ) {
2680 cx.set_global(cx.update(SettingsStore::test));
2681 cx.update(language_settings::init);
2682
2683 let text = indoc! {"
2684 fn main() {
2685 le
2686 }
2687 "};
2688 let buffer =
2689 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2690 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2691 let range = buffer.read_with(cx, |buffer, cx| {
2692 let snapshot = buffer.snapshot(cx);
2693 snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
2694 });
2695 let codegen =
2696 cx.new_model(|cx| Codegen::new(buffer.clone(), range.clone(), None, None, cx));
2697
2698 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2699 codegen.update(cx, |codegen, cx| {
2700 codegen.handle_stream(
2701 String::new(),
2702 range.clone(),
2703 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2704 cx,
2705 )
2706 });
2707
2708 cx.background_executor.run_until_parked();
2709
2710 let mut new_text = concat!(
2711 "t mut x = 0;\n",
2712 "while x < 10 {\n",
2713 " x += 1;\n",
2714 "}", //
2715 );
2716 while !new_text.is_empty() {
2717 let max_len = cmp::min(new_text.len(), 10);
2718 let len = rng.gen_range(1..=max_len);
2719 let (chunk, suffix) = new_text.split_at(len);
2720 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
2721 new_text = suffix;
2722 cx.background_executor.run_until_parked();
2723 }
2724 drop(chunks_tx);
2725 cx.background_executor.run_until_parked();
2726
2727 assert_eq!(
2728 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2729 indoc! {"
2730 fn main() {
2731 let mut x = 0;
2732 while x < 10 {
2733 x += 1;
2734 }
2735 }
2736 "}
2737 );
2738 }
2739
2740 #[gpui::test(iterations = 10)]
2741 async fn test_autoindent_when_generating_before_indentation(
2742 cx: &mut TestAppContext,
2743 mut rng: StdRng,
2744 ) {
2745 cx.update(LanguageModelRegistry::test);
2746 cx.set_global(cx.update(SettingsStore::test));
2747 cx.update(language_settings::init);
2748
2749 let text = concat!(
2750 "fn main() {\n",
2751 " \n",
2752 "}\n" //
2753 );
2754 let buffer =
2755 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2756 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2757 let range = buffer.read_with(cx, |buffer, cx| {
2758 let snapshot = buffer.snapshot(cx);
2759 snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
2760 });
2761 let codegen =
2762 cx.new_model(|cx| Codegen::new(buffer.clone(), range.clone(), None, None, cx));
2763
2764 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2765 codegen.update(cx, |codegen, cx| {
2766 codegen.handle_stream(
2767 String::new(),
2768 range.clone(),
2769 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2770 cx,
2771 )
2772 });
2773
2774 cx.background_executor.run_until_parked();
2775
2776 let mut new_text = concat!(
2777 "let mut x = 0;\n",
2778 "while x < 10 {\n",
2779 " x += 1;\n",
2780 "}", //
2781 );
2782 while !new_text.is_empty() {
2783 let max_len = cmp::min(new_text.len(), 10);
2784 let len = rng.gen_range(1..=max_len);
2785 let (chunk, suffix) = new_text.split_at(len);
2786 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
2787 new_text = suffix;
2788 cx.background_executor.run_until_parked();
2789 }
2790 drop(chunks_tx);
2791 cx.background_executor.run_until_parked();
2792
2793 assert_eq!(
2794 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2795 indoc! {"
2796 fn main() {
2797 let mut x = 0;
2798 while x < 10 {
2799 x += 1;
2800 }
2801 }
2802 "}
2803 );
2804 }
2805
2806 #[gpui::test(iterations = 10)]
2807 async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
2808 cx.update(LanguageModelRegistry::test);
2809 cx.set_global(cx.update(SettingsStore::test));
2810 cx.update(language_settings::init);
2811
2812 let text = indoc! {"
2813 func main() {
2814 \tx := 0
2815 \tfor i := 0; i < 10; i++ {
2816 \t\tx++
2817 \t}
2818 }
2819 "};
2820 let buffer = cx.new_model(|cx| Buffer::local(text, cx));
2821 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2822 let range = buffer.read_with(cx, |buffer, cx| {
2823 let snapshot = buffer.snapshot(cx);
2824 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
2825 });
2826 let codegen =
2827 cx.new_model(|cx| Codegen::new(buffer.clone(), range.clone(), None, None, cx));
2828
2829 let (chunks_tx, chunks_rx) = mpsc::unbounded();
2830 codegen.update(cx, |codegen, cx| {
2831 codegen.handle_stream(
2832 String::new(),
2833 range.clone(),
2834 future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
2835 cx,
2836 )
2837 });
2838
2839 let new_text = concat!(
2840 "func main() {\n",
2841 "\tx := 0\n",
2842 "\tfor x < 10 {\n",
2843 "\t\tx++\n",
2844 "\t}", //
2845 );
2846 chunks_tx.unbounded_send(new_text.to_string()).unwrap();
2847 drop(chunks_tx);
2848 cx.background_executor.run_until_parked();
2849
2850 assert_eq!(
2851 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2852 indoc! {"
2853 func main() {
2854 \tx := 0
2855 \tfor x < 10 {
2856 \t\tx++
2857 \t}
2858 }
2859 "}
2860 );
2861 }
2862
2863 #[gpui::test]
2864 async fn test_strip_invalid_spans_from_codeblock() {
2865 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
2866 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
2867 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
2868 assert_chunks(
2869 "```html\n```js\nLorem ipsum dolor\n```\n```",
2870 "```js\nLorem ipsum dolor\n```",
2871 )
2872 .await;
2873 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
2874 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
2875 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
2876 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
2877
2878 async fn assert_chunks(text: &str, expected_text: &str) {
2879 for chunk_size in 1..=text.len() {
2880 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
2881 .map(|chunk| chunk.unwrap())
2882 .collect::<String>()
2883 .await;
2884 assert_eq!(
2885 actual_text, expected_text,
2886 "failed to strip invalid spans, chunk size: {}",
2887 chunk_size
2888 );
2889 }
2890 }
2891
2892 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
2893 stream::iter(
2894 text.chars()
2895 .collect::<Vec<_>>()
2896 .chunks(size)
2897 .map(|chunk| Ok(chunk.iter().collect::<String>()))
2898 .collect::<Vec<_>>(),
2899 )
2900 }
2901 }
2902
2903 fn rust_lang() -> Language {
2904 Language::new(
2905 LanguageConfig {
2906 name: "Rust".into(),
2907 matcher: LanguageMatcher {
2908 path_suffixes: vec!["rs".to_string()],
2909 ..Default::default()
2910 },
2911 ..Default::default()
2912 },
2913 Some(tree_sitter_rust::language()),
2914 )
2915 .with_indents_query(
2916 r#"
2917 (call_expression) @indent
2918 (field_expression) @indent
2919 (_ "(" ")" @end) @indent
2920 (_ "{" "}" @end) @indent
2921 "#,
2922 )
2923 .unwrap()
2924 }
2925}