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