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