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