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