1use crate::{
2 prompts::generate_content_prompt, AssistantPanel, CompletionProvider, Hunk,
3 LanguageModelRequest, LanguageModelRequestMessage, Role, StreamingDiff,
4};
5use anyhow::Result;
6use client::telemetry::Telemetry;
7use collections::{hash_map, HashMap, HashSet, VecDeque};
8use editor::{
9 actions::{MoveDown, MoveUp},
10 display_map::{
11 BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,
12 },
13 scroll::{Autoscroll, AutoscrollStrategy},
14 Anchor, AnchorRangeExt, Editor, EditorElement, EditorEvent, EditorStyle, ExcerptRange,
15 GutterDimensions, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint,
16};
17use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
18use gpui::{
19 AppContext, EventEmitter, FocusHandle, FocusableView, FontStyle, FontWeight, Global,
20 HighlightStyle, Model, ModelContext, Subscription, Task, TextStyle, UpdateGlobal, View,
21 ViewContext, WeakView, WhiteSpace, WindowContext,
22};
23use language::{Buffer, Point, TransactionId};
24use multi_buffer::MultiBufferRow;
25use parking_lot::Mutex;
26use rope::Rope;
27use settings::Settings;
28use similar::TextDiff;
29use std::{
30 cmp, future, mem,
31 ops::{Range, RangeInclusive},
32 sync::Arc,
33 time::Instant,
34};
35use theme::ThemeSettings;
36use ui::{prelude::*, Tooltip};
37use util::RangeExt;
38use workspace::{notifications::NotificationId, Toast, Workspace};
39
40pub fn init(telemetry: Arc<Telemetry>, cx: &mut AppContext) {
41 cx.set_global(InlineAssistant::new(telemetry));
42}
43
44const PROMPT_HISTORY_MAX_LEN: usize = 20;
45
46pub struct InlineAssistant {
47 next_assist_id: InlineAssistId,
48 pending_assists: HashMap<InlineAssistId, PendingInlineAssist>,
49 pending_assist_ids_by_editor: HashMap<WeakView<Editor>, Vec<InlineAssistId>>,
50 prompt_history: VecDeque<String>,
51 telemetry: Option<Arc<Telemetry>>,
52}
53
54impl Global for InlineAssistant {}
55
56impl InlineAssistant {
57 pub fn new(telemetry: Arc<Telemetry>) -> Self {
58 Self {
59 next_assist_id: InlineAssistId::default(),
60 pending_assists: HashMap::default(),
61 pending_assist_ids_by_editor: HashMap::default(),
62 prompt_history: VecDeque::default(),
63 telemetry: Some(telemetry),
64 }
65 }
66
67 pub fn assist(
68 &mut self,
69 editor: &View<Editor>,
70 workspace: Option<WeakView<Workspace>>,
71 include_context: bool,
72 cx: &mut WindowContext,
73 ) {
74 let selection = editor.read(cx).selections.newest_anchor().clone();
75 if selection.start.excerpt_id != selection.end.excerpt_id {
76 return;
77 }
78 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
79
80 // Extend the selection to the start and the end of the line.
81 let mut point_selection = selection.map(|selection| selection.to_point(&snapshot));
82 if point_selection.end > point_selection.start {
83 point_selection.start.column = 0;
84 // If the selection ends at the start of the line, we don't want to include it.
85 if point_selection.end.column == 0 {
86 point_selection.end.row -= 1;
87 }
88 point_selection.end.column = snapshot.line_len(MultiBufferRow(point_selection.end.row));
89 }
90
91 let codegen_kind = if point_selection.start == point_selection.end {
92 CodegenKind::Generate {
93 position: snapshot.anchor_after(point_selection.start),
94 }
95 } else {
96 CodegenKind::Transform {
97 range: snapshot.anchor_before(point_selection.start)
98 ..snapshot.anchor_after(point_selection.end),
99 }
100 };
101
102 let assist_id = self.next_assist_id.post_inc();
103 let codegen = cx.new_model(|cx| {
104 Codegen::new(
105 editor.read(cx).buffer().clone(),
106 codegen_kind,
107 self.telemetry.clone(),
108 cx,
109 )
110 });
111
112 let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
113 let prompt_editor = cx.new_view(|cx| {
114 InlineAssistEditor::new(
115 assist_id,
116 gutter_dimensions.clone(),
117 self.prompt_history.clone(),
118 codegen.clone(),
119 workspace.clone(),
120 cx,
121 )
122 });
123 let (prompt_block_id, end_block_id) = editor.update(cx, |editor, cx| {
124 let start_anchor = snapshot.anchor_before(point_selection.start);
125 let end_anchor = snapshot.anchor_after(point_selection.end);
126 editor.change_selections(Some(Autoscroll::newest()), cx, |selections| {
127 selections.select_anchor_ranges([start_anchor..start_anchor])
128 });
129 let block_ids = editor.insert_blocks(
130 [
131 BlockProperties {
132 style: BlockStyle::Sticky,
133 position: start_anchor,
134 height: prompt_editor.read(cx).height_in_lines,
135 render: build_inline_assist_editor_renderer(
136 &prompt_editor,
137 gutter_dimensions,
138 ),
139 disposition: BlockDisposition::Above,
140 },
141 BlockProperties {
142 style: BlockStyle::Sticky,
143 position: end_anchor,
144 height: 1,
145 render: Box::new(|cx| {
146 v_flex()
147 .h_full()
148 .w_full()
149 .border_t_1()
150 .border_color(cx.theme().status().info_border)
151 .into_any_element()
152 }),
153 disposition: BlockDisposition::Below,
154 },
155 ],
156 Some(Autoscroll::Strategy(AutoscrollStrategy::Newest)),
157 cx,
158 );
159 (block_ids[0], block_ids[1])
160 });
161
162 self.pending_assists.insert(
163 assist_id,
164 PendingInlineAssist {
165 include_context,
166 editor: editor.downgrade(),
167 editor_decorations: Some(PendingInlineAssistDecorations {
168 prompt_block_id,
169 prompt_editor: prompt_editor.clone(),
170 removed_line_block_ids: HashSet::default(),
171 end_block_id,
172 }),
173 codegen: codegen.clone(),
174 workspace,
175 _subscriptions: vec![
176 cx.subscribe(&prompt_editor, |inline_assist_editor, event, cx| {
177 InlineAssistant::update_global(cx, |this, cx| {
178 this.handle_inline_assistant_editor_event(
179 inline_assist_editor,
180 event,
181 cx,
182 )
183 })
184 }),
185 editor.update(cx, |editor, _cx| {
186 editor.register_action(
187 move |_: &editor::actions::Newline, cx: &mut WindowContext| {
188 InlineAssistant::update_global(cx, |this, cx| {
189 this.handle_editor_action(assist_id, false, cx)
190 })
191 },
192 )
193 }),
194 editor.update(cx, |editor, _cx| {
195 editor.register_action(
196 move |_: &editor::actions::Cancel, cx: &mut WindowContext| {
197 InlineAssistant::update_global(cx, |this, cx| {
198 this.handle_editor_action(assist_id, true, cx)
199 })
200 },
201 )
202 }),
203 cx.subscribe(editor, move |editor, event, cx| {
204 InlineAssistant::update_global(cx, |this, cx| {
205 this.handle_editor_event(assist_id, editor, event, cx)
206 })
207 }),
208 cx.observe(&codegen, {
209 let editor = editor.downgrade();
210 move |_, cx| {
211 if let Some(editor) = editor.upgrade() {
212 InlineAssistant::update_global(cx, |this, cx| {
213 this.update_editor_highlights(&editor, cx);
214 this.update_editor_blocks(&editor, assist_id, cx);
215 })
216 }
217 }
218 }),
219 cx.subscribe(&codegen, move |codegen, event, cx| {
220 InlineAssistant::update_global(cx, |this, cx| match event {
221 CodegenEvent::Undone => this.finish_inline_assist(assist_id, false, cx),
222 CodegenEvent::Finished => {
223 let pending_assist = if let Some(pending_assist) =
224 this.pending_assists.get(&assist_id)
225 {
226 pending_assist
227 } else {
228 return;
229 };
230
231 if let CodegenStatus::Error(error) = &codegen.read(cx).status {
232 if pending_assist.editor_decorations.is_none() {
233 if let Some(workspace) = pending_assist
234 .workspace
235 .as_ref()
236 .and_then(|workspace| workspace.upgrade())
237 {
238 let error =
239 format!("Inline assistant error: {}", error);
240 workspace.update(cx, |workspace, cx| {
241 struct InlineAssistantError;
242
243 let id = NotificationId::identified::<
244 InlineAssistantError,
245 >(
246 assist_id.0
247 );
248
249 workspace.show_toast(Toast::new(id, error), cx);
250 })
251 }
252 }
253 }
254
255 if pending_assist.editor_decorations.is_none() {
256 this.finish_inline_assist(assist_id, false, cx);
257 }
258 }
259 })
260 }),
261 ],
262 },
263 );
264
265 self.pending_assist_ids_by_editor
266 .entry(editor.downgrade())
267 .or_default()
268 .push(assist_id);
269 self.update_editor_highlights(editor, cx);
270 }
271
272 fn handle_inline_assistant_editor_event(
273 &mut self,
274 inline_assist_editor: View<InlineAssistEditor>,
275 event: &InlineAssistEditorEvent,
276 cx: &mut WindowContext,
277 ) {
278 let assist_id = inline_assist_editor.read(cx).id;
279 match event {
280 InlineAssistEditorEvent::Started => {
281 self.start_inline_assist(assist_id, cx);
282 }
283 InlineAssistEditorEvent::Stopped => {
284 self.stop_inline_assist(assist_id, cx);
285 }
286 InlineAssistEditorEvent::Confirmed => {
287 self.finish_inline_assist(assist_id, false, cx);
288 }
289 InlineAssistEditorEvent::Canceled => {
290 self.finish_inline_assist(assist_id, true, cx);
291 }
292 InlineAssistEditorEvent::Dismissed => {
293 self.dismiss_inline_assist(assist_id, cx);
294 }
295 InlineAssistEditorEvent::Resized { height_in_lines } => {
296 self.resize_inline_assist(assist_id, *height_in_lines, cx);
297 }
298 }
299 }
300
301 fn handle_editor_action(
302 &mut self,
303 assist_id: InlineAssistId,
304 undo: bool,
305 cx: &mut WindowContext,
306 ) {
307 let Some(assist) = self.pending_assists.get(&assist_id) else {
308 return;
309 };
310 let Some(editor) = assist.editor.upgrade() else {
311 return;
312 };
313
314 let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
315 let assist_range = assist.codegen.read(cx).range().to_offset(&buffer);
316 let editor = editor.read(cx);
317 if editor.selections.count() == 1 {
318 let selection = editor.selections.newest::<usize>(cx);
319 if assist_range.contains(&selection.start) && assist_range.contains(&selection.end) {
320 if undo {
321 self.finish_inline_assist(assist_id, true, cx);
322 } else if matches!(assist.codegen.read(cx).status, CodegenStatus::Pending) {
323 self.dismiss_inline_assist(assist_id, cx);
324 } else {
325 self.finish_inline_assist(assist_id, false, cx);
326 }
327
328 return;
329 }
330 }
331
332 cx.propagate();
333 }
334
335 fn handle_editor_event(
336 &mut self,
337 assist_id: InlineAssistId,
338 editor: View<Editor>,
339 event: &EditorEvent,
340 cx: &mut WindowContext,
341 ) {
342 let Some(assist) = self.pending_assists.get(&assist_id) else {
343 return;
344 };
345
346 match event {
347 EditorEvent::SelectionsChanged { local } if *local => {
348 if let Some(decorations) = assist.editor_decorations.as_ref() {
349 if decorations
350 .prompt_editor
351 .focus_handle(cx)
352 .contains_focused(cx)
353 {
354 cx.focus_view(&editor);
355 }
356 }
357 }
358 EditorEvent::Saved => {
359 if let CodegenStatus::Done = &assist.codegen.read(cx).status {
360 self.finish_inline_assist(assist_id, false, cx)
361 }
362 }
363 EditorEvent::Edited { transaction_id }
364 if matches!(
365 assist.codegen.read(cx).status,
366 CodegenStatus::Error(_) | CodegenStatus::Done
367 ) =>
368 {
369 let buffer = editor.read(cx).buffer().read(cx);
370 let edited_ranges =
371 buffer.edited_ranges_for_transaction::<usize>(*transaction_id, cx);
372 let assist_range = assist.codegen.read(cx).range().to_offset(&buffer.read(cx));
373 if edited_ranges
374 .iter()
375 .any(|range| range.overlaps(&assist_range))
376 {
377 self.finish_inline_assist(assist_id, false, cx);
378 }
379 }
380 _ => {}
381 }
382 }
383
384 fn finish_inline_assist(
385 &mut self,
386 assist_id: InlineAssistId,
387 undo: bool,
388 cx: &mut WindowContext,
389 ) {
390 self.dismiss_inline_assist(assist_id, cx);
391
392 if let Some(pending_assist) = self.pending_assists.remove(&assist_id) {
393 if let hash_map::Entry::Occupied(mut entry) = self
394 .pending_assist_ids_by_editor
395 .entry(pending_assist.editor.clone())
396 {
397 entry.get_mut().retain(|id| *id != assist_id);
398 if entry.get().is_empty() {
399 entry.remove();
400 }
401 }
402
403 if let Some(editor) = pending_assist.editor.upgrade() {
404 self.update_editor_highlights(&editor, cx);
405
406 if undo {
407 pending_assist
408 .codegen
409 .update(cx, |codegen, cx| codegen.undo(cx));
410 }
411 }
412 }
413 }
414
415 fn dismiss_inline_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) -> bool {
416 let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) else {
417 return false;
418 };
419 let Some(editor) = pending_assist.editor.upgrade() else {
420 return false;
421 };
422 let Some(decorations) = pending_assist.editor_decorations.take() else {
423 return false;
424 };
425
426 editor.update(cx, |editor, cx| {
427 let mut to_remove = decorations.removed_line_block_ids;
428 to_remove.insert(decorations.prompt_block_id);
429 to_remove.insert(decorations.end_block_id);
430 editor.remove_blocks(to_remove, None, cx);
431 if decorations
432 .prompt_editor
433 .focus_handle(cx)
434 .contains_focused(cx)
435 {
436 editor.focus(cx);
437 }
438 });
439
440 self.update_editor_highlights(&editor, cx);
441 true
442 }
443
444 fn resize_inline_assist(
445 &mut self,
446 assist_id: InlineAssistId,
447 height_in_lines: u8,
448 cx: &mut WindowContext,
449 ) {
450 if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) {
451 if let Some(editor) = pending_assist.editor.upgrade() {
452 if let Some(decorations) = pending_assist.editor_decorations.as_ref() {
453 let gutter_dimensions =
454 decorations.prompt_editor.read(cx).gutter_dimensions.clone();
455 let mut new_blocks = HashMap::default();
456 new_blocks.insert(
457 decorations.prompt_block_id,
458 (
459 Some(height_in_lines),
460 build_inline_assist_editor_renderer(
461 &decorations.prompt_editor,
462 gutter_dimensions,
463 ),
464 ),
465 );
466 editor.update(cx, |editor, cx| {
467 editor
468 .display_map
469 .update(cx, |map, cx| map.replace_blocks(new_blocks, cx))
470 });
471 }
472 }
473 }
474 }
475
476 fn start_inline_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
477 let pending_assist = if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id)
478 {
479 pending_assist
480 } else {
481 return;
482 };
483
484 pending_assist
485 .codegen
486 .update(cx, |codegen, cx| codegen.undo(cx));
487
488 let Some(user_prompt) = pending_assist
489 .editor_decorations
490 .as_ref()
491 .map(|decorations| decorations.prompt_editor.read(cx).prompt(cx))
492 else {
493 return;
494 };
495
496 let context = if pending_assist.include_context {
497 pending_assist.workspace.as_ref().and_then(|workspace| {
498 let workspace = workspace.upgrade()?.read(cx);
499 let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
500 assistant_panel.read(cx).active_context(cx)
501 })
502 } else {
503 None
504 };
505
506 let editor = if let Some(editor) = pending_assist.editor.upgrade() {
507 editor
508 } else {
509 return;
510 };
511
512 let project_name = pending_assist.workspace.as_ref().and_then(|workspace| {
513 let workspace = workspace.upgrade()?;
514 Some(
515 workspace
516 .read(cx)
517 .project()
518 .read(cx)
519 .worktree_root_names(cx)
520 .collect::<Vec<&str>>()
521 .join("/"),
522 )
523 });
524
525 self.prompt_history.retain(|prompt| *prompt != user_prompt);
526 self.prompt_history.push_back(user_prompt.clone());
527 if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
528 self.prompt_history.pop_front();
529 }
530
531 let codegen = pending_assist.codegen.clone();
532 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
533 let range = codegen.read(cx).range();
534 let start = snapshot.point_to_buffer_offset(range.start);
535 let end = snapshot.point_to_buffer_offset(range.end);
536 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
537 let (start_buffer, start_buffer_offset) = start;
538 let (end_buffer, end_buffer_offset) = end;
539 if start_buffer.remote_id() == end_buffer.remote_id() {
540 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
541 } else {
542 self.finish_inline_assist(assist_id, false, cx);
543 return;
544 }
545 } else {
546 self.finish_inline_assist(assist_id, false, cx);
547 return;
548 };
549
550 let language = buffer.language_at(range.start);
551 let language_name = if let Some(language) = language.as_ref() {
552 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
553 None
554 } else {
555 Some(language.name())
556 }
557 } else {
558 None
559 };
560
561 // Higher Temperature increases the randomness of model outputs.
562 // If Markdown or No Language is Known, increase the randomness for more creative output
563 // If Code, decrease temperature to get more deterministic outputs
564 let temperature = if let Some(language) = language_name.clone() {
565 if language.as_ref() == "Markdown" {
566 1.0
567 } else {
568 0.5
569 }
570 } else {
571 1.0
572 };
573
574 let prompt = cx.background_executor().spawn(async move {
575 let language_name = language_name.as_deref();
576 generate_content_prompt(user_prompt, language_name, buffer, range, project_name)
577 });
578
579 let mut messages = Vec::new();
580 if let Some(context) = context {
581 let request = context.read(cx).to_completion_request(cx);
582 messages = request.messages;
583 }
584 let model = CompletionProvider::global(cx).model();
585
586 cx.spawn(|mut cx| async move {
587 let prompt = prompt.await?;
588
589 messages.push(LanguageModelRequestMessage {
590 role: Role::User,
591 content: prompt,
592 });
593
594 let request = LanguageModelRequest {
595 model,
596 messages,
597 stop: vec!["|END|>".to_string()],
598 temperature,
599 };
600
601 codegen.update(&mut cx, |codegen, cx| codegen.start(request, cx))?;
602 anyhow::Ok(())
603 })
604 .detach_and_log_err(cx);
605 }
606
607 fn stop_inline_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
608 let pending_assist = if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id)
609 {
610 pending_assist
611 } else {
612 return;
613 };
614
615 pending_assist
616 .codegen
617 .update(cx, |codegen, cx| codegen.stop(cx));
618 }
619
620 fn update_editor_highlights(&self, editor: &View<Editor>, cx: &mut WindowContext) {
621 let mut gutter_pending_ranges = Vec::new();
622 let mut gutter_transformed_ranges = Vec::new();
623 let mut foreground_ranges = Vec::new();
624 let mut inserted_row_ranges = Vec::new();
625 let empty_assist_ids = Vec::new();
626 let assist_ids = self
627 .pending_assist_ids_by_editor
628 .get(&editor.downgrade())
629 .unwrap_or(&empty_assist_ids);
630
631 for assist_id in assist_ids {
632 if let Some(pending_assist) = self.pending_assists.get(assist_id) {
633 let codegen = pending_assist.codegen.read(cx);
634 foreground_ranges.extend(codegen.last_equal_ranges().iter().cloned());
635
636 if codegen.edit_position != codegen.range().end {
637 gutter_pending_ranges.push(codegen.edit_position..codegen.range().end);
638 }
639
640 if codegen.range().start != codegen.edit_position {
641 gutter_transformed_ranges.push(codegen.range().start..codegen.edit_position);
642 }
643
644 if pending_assist.editor_decorations.is_some() {
645 inserted_row_ranges.extend(codegen.diff.inserted_row_ranges.iter().cloned());
646 }
647 }
648 }
649
650 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
651 merge_ranges(&mut foreground_ranges, &snapshot);
652 merge_ranges(&mut gutter_pending_ranges, &snapshot);
653 merge_ranges(&mut gutter_transformed_ranges, &snapshot);
654 editor.update(cx, |editor, cx| {
655 enum GutterPendingRange {}
656 if gutter_pending_ranges.is_empty() {
657 editor.clear_gutter_highlights::<GutterPendingRange>(cx);
658 } else {
659 editor.highlight_gutter::<GutterPendingRange>(
660 &gutter_pending_ranges,
661 |cx| cx.theme().status().info_background,
662 cx,
663 )
664 }
665
666 enum GutterTransformedRange {}
667 if gutter_transformed_ranges.is_empty() {
668 editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
669 } else {
670 editor.highlight_gutter::<GutterTransformedRange>(
671 &gutter_transformed_ranges,
672 |cx| cx.theme().status().info,
673 cx,
674 )
675 }
676
677 if foreground_ranges.is_empty() {
678 editor.clear_highlights::<PendingInlineAssist>(cx);
679 } else {
680 editor.highlight_text::<PendingInlineAssist>(
681 foreground_ranges,
682 HighlightStyle {
683 fade_out: Some(0.6),
684 ..Default::default()
685 },
686 cx,
687 );
688 }
689
690 editor.clear_row_highlights::<PendingInlineAssist>();
691 for row_range in inserted_row_ranges {
692 editor.highlight_rows::<PendingInlineAssist>(
693 row_range,
694 Some(cx.theme().status().info_background),
695 false,
696 cx,
697 );
698 }
699 });
700 }
701
702 fn update_editor_blocks(
703 &mut self,
704 editor: &View<Editor>,
705 assist_id: InlineAssistId,
706 cx: &mut WindowContext,
707 ) {
708 let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) else {
709 return;
710 };
711 let Some(decorations) = pending_assist.editor_decorations.as_mut() else {
712 return;
713 };
714
715 let codegen = pending_assist.codegen.read(cx);
716 let old_snapshot = codegen.snapshot.clone();
717 let old_buffer = codegen.old_buffer.clone();
718 let deleted_row_ranges = codegen.diff.deleted_row_ranges.clone();
719
720 editor.update(cx, |editor, cx| {
721 let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
722 editor.remove_blocks(old_blocks, None, cx);
723
724 let mut new_blocks = Vec::new();
725 for (new_row, old_row_range) in deleted_row_ranges {
726 let (_, buffer_start) = old_snapshot
727 .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
728 .unwrap();
729 let (_, buffer_end) = old_snapshot
730 .point_to_buffer_offset(Point::new(
731 *old_row_range.end(),
732 old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
733 ))
734 .unwrap();
735
736 let deleted_lines_editor = cx.new_view(|cx| {
737 let multi_buffer = cx.new_model(|_| {
738 MultiBuffer::without_headers(0, language::Capability::ReadOnly)
739 });
740 multi_buffer.update(cx, |multi_buffer, cx| {
741 multi_buffer.push_excerpts(
742 old_buffer.clone(),
743 Some(ExcerptRange {
744 context: buffer_start..buffer_end,
745 primary: None,
746 }),
747 cx,
748 );
749 });
750
751 enum DeletedLines {}
752 let mut editor = Editor::for_multibuffer(multi_buffer, None, true, cx);
753 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
754 editor.set_show_wrap_guides(false, cx);
755 editor.set_show_gutter(false, cx);
756 editor.scroll_manager.set_forbid_vertical_scroll(true);
757 editor.set_read_only(true);
758 editor.highlight_rows::<DeletedLines>(
759 Anchor::min()..=Anchor::max(),
760 Some(cx.theme().status().deleted_background),
761 false,
762 cx,
763 );
764 editor
765 });
766
767 let height = deleted_lines_editor
768 .update(cx, |editor, cx| editor.max_point(cx).row().0 as u8 + 1);
769 new_blocks.push(BlockProperties {
770 position: new_row,
771 height,
772 style: BlockStyle::Flex,
773 render: Box::new(move |cx| {
774 div()
775 .bg(cx.theme().status().deleted_background)
776 .size_full()
777 .pl(cx.gutter_dimensions.full_width())
778 .child(deleted_lines_editor.clone())
779 .into_any_element()
780 }),
781 disposition: BlockDisposition::Above,
782 });
783 }
784
785 decorations.removed_line_block_ids = editor
786 .insert_blocks(new_blocks, None, cx)
787 .into_iter()
788 .collect();
789 })
790 }
791}
792
793fn build_inline_assist_editor_renderer(
794 editor: &View<InlineAssistEditor>,
795 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
796) -> RenderBlock {
797 let editor = editor.clone();
798 Box::new(move |cx: &mut BlockContext| {
799 *gutter_dimensions.lock() = *cx.gutter_dimensions;
800 editor.clone().into_any_element()
801 })
802}
803
804#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
805struct InlineAssistId(usize);
806
807impl InlineAssistId {
808 fn post_inc(&mut self) -> InlineAssistId {
809 let id = *self;
810 self.0 += 1;
811 id
812 }
813}
814
815enum InlineAssistEditorEvent {
816 Started,
817 Stopped,
818 Confirmed,
819 Canceled,
820 Dismissed,
821 Resized { height_in_lines: u8 },
822}
823
824struct InlineAssistEditor {
825 id: InlineAssistId,
826 height_in_lines: u8,
827 prompt_editor: View<Editor>,
828 edited_since_done: bool,
829 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
830 prompt_history: VecDeque<String>,
831 prompt_history_ix: Option<usize>,
832 pending_prompt: String,
833 codegen: Model<Codegen>,
834 workspace: Option<WeakView<Workspace>>,
835 _subscriptions: Vec<Subscription>,
836}
837
838impl EventEmitter<InlineAssistEditorEvent> for InlineAssistEditor {}
839
840impl Render for InlineAssistEditor {
841 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
842 let gutter_dimensions = *self.gutter_dimensions.lock();
843
844 let buttons = match &self.codegen.read(cx).status {
845 CodegenStatus::Idle => {
846 vec![
847 IconButton::new("start", IconName::Sparkle)
848 .icon_color(Color::Muted)
849 .size(ButtonSize::None)
850 .icon_size(IconSize::XSmall)
851 .tooltip(|cx| Tooltip::for_action("Transform", &menu::Confirm, cx))
852 .on_click(
853 cx.listener(|_, _, cx| cx.emit(InlineAssistEditorEvent::Started)),
854 ),
855 IconButton::new("cancel", IconName::Close)
856 .icon_color(Color::Muted)
857 .size(ButtonSize::None)
858 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
859 .on_click(
860 cx.listener(|_, _, cx| cx.emit(InlineAssistEditorEvent::Canceled)),
861 ),
862 ]
863 }
864 CodegenStatus::Pending => {
865 vec![
866 IconButton::new("stop", IconName::Stop)
867 .icon_color(Color::Error)
868 .size(ButtonSize::None)
869 .icon_size(IconSize::XSmall)
870 .tooltip(|cx| {
871 Tooltip::with_meta(
872 "Interrupt Transformation",
873 Some(&menu::Cancel),
874 "Changes won't be discarded",
875 cx,
876 )
877 })
878 .on_click(
879 cx.listener(|_, _, cx| cx.emit(InlineAssistEditorEvent::Stopped)),
880 ),
881 IconButton::new("cancel", IconName::Close)
882 .icon_color(Color::Muted)
883 .size(ButtonSize::None)
884 .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
885 .on_click(
886 cx.listener(|_, _, cx| cx.emit(InlineAssistEditorEvent::Canceled)),
887 ),
888 ]
889 }
890 CodegenStatus::Error(_) | CodegenStatus::Done => {
891 vec![
892 if self.edited_since_done {
893 IconButton::new("restart", IconName::RotateCw)
894 .icon_color(Color::Info)
895 .icon_size(IconSize::XSmall)
896 .size(ButtonSize::None)
897 .tooltip(|cx| {
898 Tooltip::with_meta(
899 "Restart Transformation",
900 Some(&menu::Confirm),
901 "Changes will be discarded",
902 cx,
903 )
904 })
905 .on_click(cx.listener(|_, _, cx| {
906 cx.emit(InlineAssistEditorEvent::Started);
907 }))
908 } else {
909 IconButton::new("confirm", IconName::Check)
910 .icon_color(Color::Info)
911 .size(ButtonSize::None)
912 .tooltip(|cx| Tooltip::for_action("Confirm Assist", &menu::Confirm, cx))
913 .on_click(cx.listener(|_, _, cx| {
914 cx.emit(InlineAssistEditorEvent::Confirmed);
915 }))
916 },
917 IconButton::new("cancel", IconName::Close)
918 .icon_color(Color::Muted)
919 .size(ButtonSize::None)
920 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
921 .on_click(
922 cx.listener(|_, _, cx| cx.emit(InlineAssistEditorEvent::Canceled)),
923 ),
924 ]
925 }
926 };
927
928 v_flex().h_full().w_full().justify_end().child(
929 h_flex()
930 .bg(cx.theme().colors().editor_background)
931 .border_y_1()
932 .border_color(cx.theme().status().info_border)
933 .py_1p5()
934 .w_full()
935 .on_action(cx.listener(Self::confirm))
936 .on_action(cx.listener(Self::cancel))
937 .on_action(cx.listener(Self::move_up))
938 .on_action(cx.listener(Self::move_down))
939 .child(
940 h_flex()
941 .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
942 // .pr(gutter_dimensions.fold_area_width())
943 .justify_center()
944 .gap_2()
945 .children(self.workspace.clone().map(|workspace| {
946 IconButton::new("context", IconName::Context)
947 .size(ButtonSize::None)
948 .icon_size(IconSize::XSmall)
949 .icon_color(Color::Muted)
950 .on_click({
951 let workspace = workspace.clone();
952 cx.listener(move |_, _, cx| {
953 workspace
954 .update(cx, |workspace, cx| {
955 workspace.focus_panel::<AssistantPanel>(cx);
956 })
957 .ok();
958 })
959 })
960 .tooltip(move |cx| {
961 let token_count = workspace.upgrade().and_then(|workspace| {
962 let panel =
963 workspace.read(cx).panel::<AssistantPanel>(cx)?;
964 let context = panel.read(cx).active_context(cx)?;
965 context.read(cx).token_count()
966 });
967 if let Some(token_count) = token_count {
968 Tooltip::with_meta(
969 format!(
970 "{} Additional Context Tokens from Assistant",
971 token_count
972 ),
973 Some(&crate::ToggleFocus),
974 "Click to open…",
975 cx,
976 )
977 } else {
978 Tooltip::for_action(
979 "Toggle Assistant Panel",
980 &crate::ToggleFocus,
981 cx,
982 )
983 }
984 })
985 }))
986 .children(
987 if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
988 let error_message = SharedString::from(error.to_string());
989 Some(
990 div()
991 .id("error")
992 .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
993 .child(
994 Icon::new(IconName::XCircle)
995 .size(IconSize::Small)
996 .color(Color::Error),
997 ),
998 )
999 } else {
1000 None
1001 },
1002 ),
1003 )
1004 .child(div().flex_1().child(self.render_prompt_editor(cx)))
1005 .child(h_flex().gap_2().pr_4().children(buttons)),
1006 )
1007 }
1008}
1009
1010impl FocusableView for InlineAssistEditor {
1011 fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
1012 self.prompt_editor.focus_handle(cx)
1013 }
1014}
1015
1016impl InlineAssistEditor {
1017 const MAX_LINES: u8 = 8;
1018
1019 #[allow(clippy::too_many_arguments)]
1020 fn new(
1021 id: InlineAssistId,
1022 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1023 prompt_history: VecDeque<String>,
1024 codegen: Model<Codegen>,
1025 workspace: Option<WeakView<Workspace>>,
1026 cx: &mut ViewContext<Self>,
1027 ) -> Self {
1028 let prompt_editor = cx.new_view(|cx| {
1029 let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
1030 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1031 editor.set_placeholder_text("Add a prompt…", cx);
1032 editor
1033 });
1034 cx.focus_view(&prompt_editor);
1035
1036 let subscriptions = vec![
1037 cx.observe(&codegen, Self::handle_codegen_changed),
1038 cx.observe(&prompt_editor, Self::handle_prompt_editor_changed),
1039 cx.subscribe(&prompt_editor, Self::handle_prompt_editor_events),
1040 ];
1041
1042 let mut this = Self {
1043 id,
1044 height_in_lines: 1,
1045 prompt_editor,
1046 edited_since_done: false,
1047 gutter_dimensions,
1048 prompt_history,
1049 prompt_history_ix: None,
1050 pending_prompt: String::new(),
1051 codegen,
1052 workspace,
1053 _subscriptions: subscriptions,
1054 };
1055 this.count_lines(cx);
1056 this
1057 }
1058
1059 fn prompt(&self, cx: &AppContext) -> String {
1060 self.prompt_editor.read(cx).text(cx)
1061 }
1062
1063 fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
1064 let height_in_lines = cmp::max(
1065 2, // Make the editor at least two lines tall, to account for padding and buttons.
1066 cmp::min(
1067 self.prompt_editor
1068 .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1069 Self::MAX_LINES as u32,
1070 ),
1071 ) as u8;
1072
1073 if height_in_lines != self.height_in_lines {
1074 self.height_in_lines = height_in_lines;
1075 cx.emit(InlineAssistEditorEvent::Resized { height_in_lines });
1076 }
1077 }
1078
1079 fn handle_prompt_editor_changed(&mut self, _: View<Editor>, cx: &mut ViewContext<Self>) {
1080 self.count_lines(cx);
1081 }
1082
1083 fn handle_prompt_editor_events(
1084 &mut self,
1085 _: View<Editor>,
1086 event: &EditorEvent,
1087 cx: &mut ViewContext<Self>,
1088 ) {
1089 match event {
1090 EditorEvent::Edited { .. } => {
1091 let prompt = self.prompt_editor.read(cx).text(cx);
1092 if self
1093 .prompt_history_ix
1094 .map_or(true, |ix| self.prompt_history[ix] != prompt)
1095 {
1096 self.prompt_history_ix.take();
1097 self.pending_prompt = prompt;
1098 }
1099
1100 self.edited_since_done = true;
1101 cx.notify();
1102 }
1103 EditorEvent::Blurred => {
1104 if let CodegenStatus::Idle = &self.codegen.read(cx).status {
1105 let assistant_panel_is_focused = self
1106 .workspace
1107 .as_ref()
1108 .and_then(|workspace| {
1109 let panel =
1110 workspace.upgrade()?.read(cx).panel::<AssistantPanel>(cx)?;
1111 Some(panel.focus_handle(cx).contains_focused(cx))
1112 })
1113 .unwrap_or(false);
1114
1115 if !assistant_panel_is_focused {
1116 cx.emit(InlineAssistEditorEvent::Canceled);
1117 }
1118 }
1119 }
1120 _ => {}
1121 }
1122 }
1123
1124 fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
1125 match &self.codegen.read(cx).status {
1126 CodegenStatus::Idle => {
1127 self.prompt_editor
1128 .update(cx, |editor, _| editor.set_read_only(false));
1129 }
1130 CodegenStatus::Pending => {
1131 self.prompt_editor
1132 .update(cx, |editor, _| editor.set_read_only(true));
1133 }
1134 CodegenStatus::Done | CodegenStatus::Error(_) => {
1135 self.edited_since_done = false;
1136 self.prompt_editor
1137 .update(cx, |editor, _| editor.set_read_only(false));
1138 }
1139 }
1140 }
1141
1142 fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
1143 match &self.codegen.read(cx).status {
1144 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1145 cx.emit(InlineAssistEditorEvent::Canceled);
1146 }
1147 CodegenStatus::Pending => {
1148 cx.emit(InlineAssistEditorEvent::Stopped);
1149 }
1150 }
1151 }
1152
1153 fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
1154 match &self.codegen.read(cx).status {
1155 CodegenStatus::Idle => {
1156 cx.emit(InlineAssistEditorEvent::Started);
1157 }
1158 CodegenStatus::Pending => {
1159 cx.emit(InlineAssistEditorEvent::Dismissed);
1160 }
1161 CodegenStatus::Done | CodegenStatus::Error(_) => {
1162 if self.edited_since_done {
1163 cx.emit(InlineAssistEditorEvent::Started);
1164 } else {
1165 cx.emit(InlineAssistEditorEvent::Confirmed);
1166 }
1167 }
1168 }
1169 }
1170
1171 fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
1172 if let Some(ix) = self.prompt_history_ix {
1173 if ix > 0 {
1174 self.prompt_history_ix = Some(ix - 1);
1175 let prompt = self.prompt_history[ix - 1].as_str();
1176 self.prompt_editor.update(cx, |editor, cx| {
1177 editor.set_text(prompt, cx);
1178 editor.move_to_beginning(&Default::default(), cx);
1179 });
1180 }
1181 } else if !self.prompt_history.is_empty() {
1182 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1183 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1184 self.prompt_editor.update(cx, |editor, cx| {
1185 editor.set_text(prompt, cx);
1186 editor.move_to_beginning(&Default::default(), cx);
1187 });
1188 }
1189 }
1190
1191 fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
1192 if let Some(ix) = self.prompt_history_ix {
1193 if ix < self.prompt_history.len() - 1 {
1194 self.prompt_history_ix = Some(ix + 1);
1195 let prompt = self.prompt_history[ix + 1].as_str();
1196 self.prompt_editor.update(cx, |editor, cx| {
1197 editor.set_text(prompt, cx);
1198 editor.move_to_end(&Default::default(), cx)
1199 });
1200 } else {
1201 self.prompt_history_ix = None;
1202 let prompt = self.pending_prompt.as_str();
1203 self.prompt_editor.update(cx, |editor, cx| {
1204 editor.set_text(prompt, cx);
1205 editor.move_to_end(&Default::default(), cx)
1206 });
1207 }
1208 }
1209 }
1210
1211 fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1212 let settings = ThemeSettings::get_global(cx);
1213 let text_style = TextStyle {
1214 color: if self.prompt_editor.read(cx).read_only(cx) {
1215 cx.theme().colors().text_disabled
1216 } else {
1217 cx.theme().colors().text
1218 },
1219 font_family: settings.ui_font.family.clone(),
1220 font_features: settings.ui_font.features.clone(),
1221 font_size: rems(0.875).into(),
1222 font_weight: FontWeight::NORMAL,
1223 font_style: FontStyle::Normal,
1224 line_height: relative(1.3),
1225 background_color: None,
1226 underline: None,
1227 strikethrough: None,
1228 white_space: WhiteSpace::Normal,
1229 };
1230 EditorElement::new(
1231 &self.prompt_editor,
1232 EditorStyle {
1233 background: cx.theme().colors().editor_background,
1234 local_player: cx.theme().players().local(),
1235 text: text_style,
1236 ..Default::default()
1237 },
1238 )
1239 }
1240}
1241
1242struct PendingInlineAssist {
1243 editor: WeakView<Editor>,
1244 editor_decorations: Option<PendingInlineAssistDecorations>,
1245 codegen: Model<Codegen>,
1246 _subscriptions: Vec<Subscription>,
1247 workspace: Option<WeakView<Workspace>>,
1248 include_context: bool,
1249}
1250
1251struct PendingInlineAssistDecorations {
1252 prompt_block_id: BlockId,
1253 prompt_editor: View<InlineAssistEditor>,
1254 removed_line_block_ids: HashSet<BlockId>,
1255 end_block_id: BlockId,
1256}
1257
1258#[derive(Debug)]
1259pub enum CodegenEvent {
1260 Finished,
1261 Undone,
1262}
1263
1264#[derive(Clone)]
1265pub enum CodegenKind {
1266 Transform { range: Range<Anchor> },
1267 Generate { position: Anchor },
1268}
1269
1270impl CodegenKind {
1271 fn range(&self, snapshot: &MultiBufferSnapshot) -> Range<Anchor> {
1272 match self {
1273 CodegenKind::Transform { range } => range.clone(),
1274 CodegenKind::Generate { position } => position.bias_left(snapshot)..*position,
1275 }
1276 }
1277}
1278
1279pub struct Codegen {
1280 buffer: Model<MultiBuffer>,
1281 old_buffer: Model<Buffer>,
1282 snapshot: MultiBufferSnapshot,
1283 kind: CodegenKind,
1284 edit_position: Anchor,
1285 last_equal_ranges: Vec<Range<Anchor>>,
1286 transaction_id: Option<TransactionId>,
1287 status: CodegenStatus,
1288 generation: Task<()>,
1289 diff: Diff,
1290 telemetry: Option<Arc<Telemetry>>,
1291 _subscription: gpui::Subscription,
1292}
1293
1294enum CodegenStatus {
1295 Idle,
1296 Pending,
1297 Done,
1298 Error(anyhow::Error),
1299}
1300
1301#[derive(Default)]
1302struct Diff {
1303 task: Option<Task<()>>,
1304 should_update: bool,
1305 deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1306 inserted_row_ranges: Vec<RangeInclusive<Anchor>>,
1307}
1308
1309impl EventEmitter<CodegenEvent> for Codegen {}
1310
1311impl Codegen {
1312 pub fn new(
1313 buffer: Model<MultiBuffer>,
1314 kind: CodegenKind,
1315 telemetry: Option<Arc<Telemetry>>,
1316 cx: &mut ModelContext<Self>,
1317 ) -> Self {
1318 let snapshot = buffer.read(cx).snapshot(cx);
1319
1320 let (old_buffer, _, _) = buffer
1321 .read(cx)
1322 .range_to_buffer_ranges(kind.range(&snapshot), cx)
1323 .pop()
1324 .unwrap();
1325 let old_buffer = cx.new_model(|cx| {
1326 let old_buffer = old_buffer.read(cx);
1327 let text = old_buffer.as_rope().clone();
1328 let line_ending = old_buffer.line_ending();
1329 let language = old_buffer.language().cloned();
1330 let language_registry = old_buffer.language_registry();
1331
1332 let mut buffer = Buffer::local_normalized(text, line_ending, cx);
1333 buffer.set_language(language, cx);
1334 if let Some(language_registry) = language_registry {
1335 buffer.set_language_registry(language_registry)
1336 }
1337 buffer
1338 });
1339
1340 Self {
1341 buffer: buffer.clone(),
1342 old_buffer,
1343 edit_position: kind.range(&snapshot).start,
1344 snapshot,
1345 kind,
1346 last_equal_ranges: Default::default(),
1347 transaction_id: Default::default(),
1348 status: CodegenStatus::Idle,
1349 generation: Task::ready(()),
1350 diff: Diff::default(),
1351 telemetry,
1352 _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
1353 }
1354 }
1355
1356 fn handle_buffer_event(
1357 &mut self,
1358 _buffer: Model<MultiBuffer>,
1359 event: &multi_buffer::Event,
1360 cx: &mut ModelContext<Self>,
1361 ) {
1362 if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
1363 if self.transaction_id == Some(*transaction_id) {
1364 self.transaction_id = None;
1365 self.generation = Task::ready(());
1366 cx.emit(CodegenEvent::Undone);
1367 }
1368 }
1369 }
1370
1371 pub fn range(&self) -> Range<Anchor> {
1372 self.kind.range(&self.snapshot)
1373 }
1374
1375 pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
1376 &self.last_equal_ranges
1377 }
1378
1379 pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut ModelContext<Self>) {
1380 let range = self.range();
1381 let snapshot = self.snapshot.clone();
1382 let selected_text = snapshot
1383 .text_for_range(range.start..range.end)
1384 .collect::<Rope>();
1385
1386 let selection_start = range.start.to_point(&snapshot);
1387 let suggested_line_indent = snapshot
1388 .suggested_indents(selection_start.row..selection_start.row + 1, cx)
1389 .into_values()
1390 .next()
1391 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
1392
1393 let model_telemetry_id = prompt.model.telemetry_id();
1394 let response = CompletionProvider::global(cx).complete(prompt);
1395 let telemetry = self.telemetry.clone();
1396 self.edit_position = range.start;
1397 self.diff = Diff::default();
1398 self.status = CodegenStatus::Pending;
1399 self.generation = cx.spawn(|this, mut cx| {
1400 async move {
1401 let generate = async {
1402 let mut edit_start = range.start.to_offset(&snapshot);
1403
1404 let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1405 let diff: Task<anyhow::Result<()>> =
1406 cx.background_executor().spawn(async move {
1407 let mut response_latency = None;
1408 let request_start = Instant::now();
1409 let diff = async {
1410 let chunks = strip_invalid_spans_from_codeblock(response.await?);
1411 futures::pin_mut!(chunks);
1412 let mut diff = StreamingDiff::new(selected_text.to_string());
1413
1414 let mut new_text = String::new();
1415 let mut base_indent = None;
1416 let mut line_indent = None;
1417 let mut first_line = true;
1418
1419 while let Some(chunk) = chunks.next().await {
1420 if response_latency.is_none() {
1421 response_latency = Some(request_start.elapsed());
1422 }
1423 let chunk = chunk?;
1424
1425 let mut lines = chunk.split('\n').peekable();
1426 while let Some(line) = lines.next() {
1427 new_text.push_str(line);
1428 if line_indent.is_none() {
1429 if let Some(non_whitespace_ch_ix) =
1430 new_text.find(|ch: char| !ch.is_whitespace())
1431 {
1432 line_indent = Some(non_whitespace_ch_ix);
1433 base_indent = base_indent.or(line_indent);
1434
1435 let line_indent = line_indent.unwrap();
1436 let base_indent = base_indent.unwrap();
1437 let indent_delta =
1438 line_indent as i32 - base_indent as i32;
1439 let mut corrected_indent_len = cmp::max(
1440 0,
1441 suggested_line_indent.len as i32 + indent_delta,
1442 )
1443 as usize;
1444 if first_line {
1445 corrected_indent_len = corrected_indent_len
1446 .saturating_sub(
1447 selection_start.column as usize,
1448 );
1449 }
1450
1451 let indent_char = suggested_line_indent.char();
1452 let mut indent_buffer = [0; 4];
1453 let indent_str =
1454 indent_char.encode_utf8(&mut indent_buffer);
1455 new_text.replace_range(
1456 ..line_indent,
1457 &indent_str.repeat(corrected_indent_len),
1458 );
1459 }
1460 }
1461
1462 if line_indent.is_some() {
1463 hunks_tx.send(diff.push_new(&new_text)).await?;
1464 new_text.clear();
1465 }
1466
1467 if lines.peek().is_some() {
1468 hunks_tx.send(diff.push_new("\n")).await?;
1469 line_indent = None;
1470 first_line = false;
1471 }
1472 }
1473 }
1474 hunks_tx.send(diff.push_new(&new_text)).await?;
1475 hunks_tx.send(diff.finish()).await?;
1476
1477 anyhow::Ok(())
1478 };
1479
1480 let result = diff.await;
1481
1482 let error_message =
1483 result.as_ref().err().map(|error| error.to_string());
1484 if let Some(telemetry) = telemetry {
1485 telemetry.report_assistant_event(
1486 None,
1487 telemetry_events::AssistantKind::Inline,
1488 model_telemetry_id,
1489 response_latency,
1490 error_message,
1491 );
1492 }
1493
1494 result?;
1495 Ok(())
1496 });
1497
1498 while let Some(hunks) = hunks_rx.next().await {
1499 this.update(&mut cx, |this, cx| {
1500 this.last_equal_ranges.clear();
1501
1502 let transaction = this.buffer.update(cx, |buffer, cx| {
1503 // Avoid grouping assistant edits with user edits.
1504 buffer.finalize_last_transaction(cx);
1505
1506 buffer.start_transaction(cx);
1507 buffer.edit(
1508 hunks.into_iter().filter_map(|hunk| match hunk {
1509 Hunk::Insert { text } => {
1510 let edit_start = snapshot.anchor_after(edit_start);
1511 Some((edit_start..edit_start, text))
1512 }
1513 Hunk::Remove { len } => {
1514 let edit_end = edit_start + len;
1515 let edit_range = snapshot.anchor_after(edit_start)
1516 ..snapshot.anchor_before(edit_end);
1517 edit_start = edit_end;
1518 Some((edit_range, String::new()))
1519 }
1520 Hunk::Keep { len } => {
1521 let edit_end = edit_start + len;
1522 let edit_range = snapshot.anchor_after(edit_start)
1523 ..snapshot.anchor_before(edit_end);
1524 edit_start = edit_end;
1525 this.last_equal_ranges.push(edit_range);
1526 None
1527 }
1528 }),
1529 None,
1530 cx,
1531 );
1532 this.edit_position = snapshot.anchor_after(edit_start);
1533
1534 buffer.end_transaction(cx)
1535 });
1536
1537 if let Some(transaction) = transaction {
1538 if let Some(first_transaction) = this.transaction_id {
1539 // Group all assistant edits into the first transaction.
1540 this.buffer.update(cx, |buffer, cx| {
1541 buffer.merge_transactions(
1542 transaction,
1543 first_transaction,
1544 cx,
1545 )
1546 });
1547 } else {
1548 this.transaction_id = Some(transaction);
1549 this.buffer.update(cx, |buffer, cx| {
1550 buffer.finalize_last_transaction(cx)
1551 });
1552 }
1553 }
1554
1555 this.update_diff(cx);
1556 cx.notify();
1557 })?;
1558 }
1559
1560 diff.await?;
1561
1562 anyhow::Ok(())
1563 };
1564
1565 let result = generate.await;
1566 this.update(&mut cx, |this, cx| {
1567 this.last_equal_ranges.clear();
1568 if let Err(error) = result {
1569 this.status = CodegenStatus::Error(error);
1570 } else {
1571 this.status = CodegenStatus::Done;
1572 }
1573 cx.emit(CodegenEvent::Finished);
1574 cx.notify();
1575 })
1576 .ok();
1577 }
1578 });
1579 cx.notify();
1580 }
1581
1582 pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
1583 self.last_equal_ranges.clear();
1584 self.status = CodegenStatus::Done;
1585 self.generation = Task::ready(());
1586 cx.emit(CodegenEvent::Finished);
1587 cx.notify();
1588 }
1589
1590 pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
1591 if let Some(transaction_id) = self.transaction_id.take() {
1592 self.buffer
1593 .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
1594 }
1595 }
1596
1597 fn update_diff(&mut self, cx: &mut ModelContext<Self>) {
1598 if self.diff.task.is_some() {
1599 self.diff.should_update = true;
1600 } else {
1601 self.diff.should_update = false;
1602
1603 let old_snapshot = self.snapshot.clone();
1604 let old_range = self.range().to_point(&old_snapshot);
1605 let new_snapshot = self.buffer.read(cx).snapshot(cx);
1606 let new_range = self.range().to_point(&new_snapshot);
1607
1608 self.diff.task = Some(cx.spawn(|this, mut cx| async move {
1609 let (deleted_row_ranges, inserted_row_ranges) = cx
1610 .background_executor()
1611 .spawn(async move {
1612 let old_text = old_snapshot
1613 .text_for_range(
1614 Point::new(old_range.start.row, 0)
1615 ..Point::new(
1616 old_range.end.row,
1617 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
1618 ),
1619 )
1620 .collect::<String>();
1621 let new_text = new_snapshot
1622 .text_for_range(
1623 Point::new(new_range.start.row, 0)
1624 ..Point::new(
1625 new_range.end.row,
1626 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
1627 ),
1628 )
1629 .collect::<String>();
1630
1631 let mut old_row = old_range.start.row;
1632 let mut new_row = new_range.start.row;
1633 let diff = TextDiff::from_lines(old_text.as_str(), new_text.as_str());
1634
1635 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
1636 let mut inserted_row_ranges = Vec::new();
1637 for change in diff.iter_all_changes() {
1638 let line_count = change.value().lines().count() as u32;
1639 match change.tag() {
1640 similar::ChangeTag::Equal => {
1641 old_row += line_count;
1642 new_row += line_count;
1643 }
1644 similar::ChangeTag::Delete => {
1645 let old_end_row = old_row + line_count - 1;
1646 let new_row =
1647 new_snapshot.anchor_before(Point::new(new_row, 0));
1648
1649 if let Some((_, last_deleted_row_range)) =
1650 deleted_row_ranges.last_mut()
1651 {
1652 if *last_deleted_row_range.end() + 1 == old_row {
1653 *last_deleted_row_range =
1654 *last_deleted_row_range.start()..=old_end_row;
1655 } else {
1656 deleted_row_ranges
1657 .push((new_row, old_row..=old_end_row));
1658 }
1659 } else {
1660 deleted_row_ranges.push((new_row, old_row..=old_end_row));
1661 }
1662
1663 old_row += line_count;
1664 }
1665 similar::ChangeTag::Insert => {
1666 let new_end_row = new_row + line_count - 1;
1667 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
1668 let end = new_snapshot.anchor_before(Point::new(
1669 new_end_row,
1670 new_snapshot.line_len(MultiBufferRow(new_end_row)),
1671 ));
1672 inserted_row_ranges.push(start..=end);
1673 new_row += line_count;
1674 }
1675 }
1676 }
1677
1678 (deleted_row_ranges, inserted_row_ranges)
1679 })
1680 .await;
1681
1682 this.update(&mut cx, |this, cx| {
1683 this.diff.deleted_row_ranges = deleted_row_ranges;
1684 this.diff.inserted_row_ranges = inserted_row_ranges;
1685 this.diff.task = None;
1686 if this.diff.should_update {
1687 this.update_diff(cx);
1688 }
1689 cx.notify();
1690 })
1691 .ok();
1692 }));
1693 }
1694 }
1695}
1696
1697fn strip_invalid_spans_from_codeblock(
1698 stream: impl Stream<Item = Result<String>>,
1699) -> impl Stream<Item = Result<String>> {
1700 let mut first_line = true;
1701 let mut buffer = String::new();
1702 let mut starts_with_markdown_codeblock = false;
1703 let mut includes_start_or_end_span = false;
1704 stream.filter_map(move |chunk| {
1705 let chunk = match chunk {
1706 Ok(chunk) => chunk,
1707 Err(err) => return future::ready(Some(Err(err))),
1708 };
1709 buffer.push_str(&chunk);
1710
1711 if buffer.len() > "<|S|".len() && buffer.starts_with("<|S|") {
1712 includes_start_or_end_span = true;
1713
1714 buffer = buffer
1715 .strip_prefix("<|S|>")
1716 .or_else(|| buffer.strip_prefix("<|S|"))
1717 .unwrap_or(&buffer)
1718 .to_string();
1719 } else if buffer.ends_with("|E|>") {
1720 includes_start_or_end_span = true;
1721 } else if buffer.starts_with("<|")
1722 || buffer.starts_with("<|S")
1723 || buffer.starts_with("<|S|")
1724 || buffer.ends_with('|')
1725 || buffer.ends_with("|E")
1726 || buffer.ends_with("|E|")
1727 {
1728 return future::ready(None);
1729 }
1730
1731 if first_line {
1732 if buffer.is_empty() || buffer == "`" || buffer == "``" {
1733 return future::ready(None);
1734 } else if buffer.starts_with("```") {
1735 starts_with_markdown_codeblock = true;
1736 if let Some(newline_ix) = buffer.find('\n') {
1737 buffer.replace_range(..newline_ix + 1, "");
1738 first_line = false;
1739 } else {
1740 return future::ready(None);
1741 }
1742 }
1743 }
1744
1745 let mut text = buffer.to_string();
1746 if starts_with_markdown_codeblock {
1747 text = text
1748 .strip_suffix("\n```\n")
1749 .or_else(|| text.strip_suffix("\n```"))
1750 .or_else(|| text.strip_suffix("\n``"))
1751 .or_else(|| text.strip_suffix("\n`"))
1752 .or_else(|| text.strip_suffix('\n'))
1753 .unwrap_or(&text)
1754 .to_string();
1755 }
1756
1757 if includes_start_or_end_span {
1758 text = text
1759 .strip_suffix("|E|>")
1760 .or_else(|| text.strip_suffix("E|>"))
1761 .or_else(|| text.strip_prefix("|>"))
1762 .or_else(|| text.strip_prefix('>'))
1763 .unwrap_or(&text)
1764 .to_string();
1765 };
1766
1767 if text.contains('\n') {
1768 first_line = false;
1769 }
1770
1771 let remainder = buffer.split_off(text.len());
1772 let result = if buffer.is_empty() {
1773 None
1774 } else {
1775 Some(Ok(buffer.clone()))
1776 };
1777
1778 buffer = remainder;
1779 future::ready(result)
1780 })
1781}
1782
1783fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
1784 ranges.sort_unstable_by(|a, b| {
1785 a.start
1786 .cmp(&b.start, buffer)
1787 .then_with(|| b.end.cmp(&a.end, buffer))
1788 });
1789
1790 let mut ix = 0;
1791 while ix + 1 < ranges.len() {
1792 let b = ranges[ix + 1].clone();
1793 let a = &mut ranges[ix];
1794 if a.end.cmp(&b.start, buffer).is_gt() {
1795 if a.end.cmp(&b.end, buffer).is_lt() {
1796 a.end = b.end;
1797 }
1798 ranges.remove(ix + 1);
1799 } else {
1800 ix += 1;
1801 }
1802 }
1803}
1804
1805#[cfg(test)]
1806mod tests {
1807 use std::sync::Arc;
1808
1809 use crate::FakeCompletionProvider;
1810
1811 use super::*;
1812 use futures::stream::{self};
1813 use gpui::{Context, TestAppContext};
1814 use indoc::indoc;
1815 use language::{
1816 language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
1817 Point,
1818 };
1819 use rand::prelude::*;
1820 use serde::Serialize;
1821 use settings::SettingsStore;
1822
1823 #[derive(Serialize)]
1824 pub struct DummyCompletionRequest {
1825 pub name: String,
1826 }
1827
1828 #[gpui::test(iterations = 10)]
1829 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
1830 let provider = FakeCompletionProvider::default();
1831 cx.set_global(cx.update(SettingsStore::test));
1832 cx.set_global(CompletionProvider::Fake(provider.clone()));
1833 cx.update(language_settings::init);
1834
1835 let text = indoc! {"
1836 fn main() {
1837 let x = 0;
1838 for _ in 0..10 {
1839 x += 1;
1840 }
1841 }
1842 "};
1843 let buffer =
1844 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1845 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
1846 let range = buffer.read_with(cx, |buffer, cx| {
1847 let snapshot = buffer.snapshot(cx);
1848 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
1849 });
1850 let codegen = cx.new_model(|cx| {
1851 Codegen::new(buffer.clone(), CodegenKind::Transform { range }, None, cx)
1852 });
1853
1854 let request = LanguageModelRequest::default();
1855 codegen.update(cx, |codegen, cx| codegen.start(request, cx));
1856
1857 let mut new_text = concat!(
1858 " let mut x = 0;\n",
1859 " while x < 10 {\n",
1860 " x += 1;\n",
1861 " }",
1862 );
1863 while !new_text.is_empty() {
1864 let max_len = cmp::min(new_text.len(), 10);
1865 let len = rng.gen_range(1..=max_len);
1866 let (chunk, suffix) = new_text.split_at(len);
1867 provider.send_completion(chunk.into());
1868 new_text = suffix;
1869 cx.background_executor.run_until_parked();
1870 }
1871 provider.finish_completion();
1872 cx.background_executor.run_until_parked();
1873
1874 assert_eq!(
1875 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1876 indoc! {"
1877 fn main() {
1878 let mut x = 0;
1879 while x < 10 {
1880 x += 1;
1881 }
1882 }
1883 "}
1884 );
1885 }
1886
1887 #[gpui::test(iterations = 10)]
1888 async fn test_autoindent_when_generating_past_indentation(
1889 cx: &mut TestAppContext,
1890 mut rng: StdRng,
1891 ) {
1892 let provider = FakeCompletionProvider::default();
1893 cx.set_global(CompletionProvider::Fake(provider.clone()));
1894 cx.set_global(cx.update(SettingsStore::test));
1895 cx.update(language_settings::init);
1896
1897 let text = indoc! {"
1898 fn main() {
1899 le
1900 }
1901 "};
1902 let buffer =
1903 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1904 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
1905 let position = buffer.read_with(cx, |buffer, cx| {
1906 let snapshot = buffer.snapshot(cx);
1907 snapshot.anchor_before(Point::new(1, 6))
1908 });
1909 let codegen = cx.new_model(|cx| {
1910 Codegen::new(buffer.clone(), CodegenKind::Generate { position }, None, cx)
1911 });
1912
1913 let request = LanguageModelRequest::default();
1914 codegen.update(cx, |codegen, cx| codegen.start(request, cx));
1915
1916 let mut new_text = concat!(
1917 "t mut x = 0;\n",
1918 "while x < 10 {\n",
1919 " x += 1;\n",
1920 "}", //
1921 );
1922 while !new_text.is_empty() {
1923 let max_len = cmp::min(new_text.len(), 10);
1924 let len = rng.gen_range(1..=max_len);
1925 let (chunk, suffix) = new_text.split_at(len);
1926 provider.send_completion(chunk.into());
1927 new_text = suffix;
1928 cx.background_executor.run_until_parked();
1929 }
1930 provider.finish_completion();
1931 cx.background_executor.run_until_parked();
1932
1933 assert_eq!(
1934 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1935 indoc! {"
1936 fn main() {
1937 let mut x = 0;
1938 while x < 10 {
1939 x += 1;
1940 }
1941 }
1942 "}
1943 );
1944 }
1945
1946 #[gpui::test(iterations = 10)]
1947 async fn test_autoindent_when_generating_before_indentation(
1948 cx: &mut TestAppContext,
1949 mut rng: StdRng,
1950 ) {
1951 let provider = FakeCompletionProvider::default();
1952 cx.set_global(CompletionProvider::Fake(provider.clone()));
1953 cx.set_global(cx.update(SettingsStore::test));
1954 cx.update(language_settings::init);
1955
1956 let text = concat!(
1957 "fn main() {\n",
1958 " \n",
1959 "}\n" //
1960 );
1961 let buffer =
1962 cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1963 let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
1964 let position = buffer.read_with(cx, |buffer, cx| {
1965 let snapshot = buffer.snapshot(cx);
1966 snapshot.anchor_before(Point::new(1, 2))
1967 });
1968 let codegen = cx.new_model(|cx| {
1969 Codegen::new(buffer.clone(), CodegenKind::Generate { position }, None, cx)
1970 });
1971
1972 let request = LanguageModelRequest::default();
1973 codegen.update(cx, |codegen, cx| codegen.start(request, cx));
1974
1975 let mut new_text = concat!(
1976 "let mut x = 0;\n",
1977 "while x < 10 {\n",
1978 " x += 1;\n",
1979 "}", //
1980 );
1981 while !new_text.is_empty() {
1982 let max_len = cmp::min(new_text.len(), 10);
1983 let len = rng.gen_range(1..=max_len);
1984 let (chunk, suffix) = new_text.split_at(len);
1985 provider.send_completion(chunk.into());
1986 new_text = suffix;
1987 cx.background_executor.run_until_parked();
1988 }
1989 provider.finish_completion();
1990 cx.background_executor.run_until_parked();
1991
1992 assert_eq!(
1993 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1994 indoc! {"
1995 fn main() {
1996 let mut x = 0;
1997 while x < 10 {
1998 x += 1;
1999 }
2000 }
2001 "}
2002 );
2003 }
2004
2005 #[gpui::test]
2006 async fn test_strip_invalid_spans_from_codeblock() {
2007 assert_eq!(
2008 strip_invalid_spans_from_codeblock(chunks("Lorem ipsum dolor", 2))
2009 .map(|chunk| chunk.unwrap())
2010 .collect::<String>()
2011 .await,
2012 "Lorem ipsum dolor"
2013 );
2014 assert_eq!(
2015 strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor", 2))
2016 .map(|chunk| chunk.unwrap())
2017 .collect::<String>()
2018 .await,
2019 "Lorem ipsum dolor"
2020 );
2021 assert_eq!(
2022 strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor\n```", 2))
2023 .map(|chunk| chunk.unwrap())
2024 .collect::<String>()
2025 .await,
2026 "Lorem ipsum dolor"
2027 );
2028 assert_eq!(
2029 strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor\n```\n", 2))
2030 .map(|chunk| chunk.unwrap())
2031 .collect::<String>()
2032 .await,
2033 "Lorem ipsum dolor"
2034 );
2035 assert_eq!(
2036 strip_invalid_spans_from_codeblock(chunks(
2037 "```html\n```js\nLorem ipsum dolor\n```\n```",
2038 2
2039 ))
2040 .map(|chunk| chunk.unwrap())
2041 .collect::<String>()
2042 .await,
2043 "```js\nLorem ipsum dolor\n```"
2044 );
2045 assert_eq!(
2046 strip_invalid_spans_from_codeblock(chunks("``\nLorem ipsum dolor\n```", 2))
2047 .map(|chunk| chunk.unwrap())
2048 .collect::<String>()
2049 .await,
2050 "``\nLorem ipsum dolor\n```"
2051 );
2052 assert_eq!(
2053 strip_invalid_spans_from_codeblock(chunks("<|S|Lorem ipsum|E|>", 2))
2054 .map(|chunk| chunk.unwrap())
2055 .collect::<String>()
2056 .await,
2057 "Lorem ipsum"
2058 );
2059
2060 assert_eq!(
2061 strip_invalid_spans_from_codeblock(chunks("<|S|>Lorem ipsum", 2))
2062 .map(|chunk| chunk.unwrap())
2063 .collect::<String>()
2064 .await,
2065 "Lorem ipsum"
2066 );
2067
2068 assert_eq!(
2069 strip_invalid_spans_from_codeblock(chunks("```\n<|S|>Lorem ipsum\n```", 2))
2070 .map(|chunk| chunk.unwrap())
2071 .collect::<String>()
2072 .await,
2073 "Lorem ipsum"
2074 );
2075 assert_eq!(
2076 strip_invalid_spans_from_codeblock(chunks("```\n<|S|Lorem ipsum|E|>\n```", 2))
2077 .map(|chunk| chunk.unwrap())
2078 .collect::<String>()
2079 .await,
2080 "Lorem ipsum"
2081 );
2082 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
2083 stream::iter(
2084 text.chars()
2085 .collect::<Vec<_>>()
2086 .chunks(size)
2087 .map(|chunk| Ok(chunk.iter().collect::<String>()))
2088 .collect::<Vec<_>>(),
2089 )
2090 }
2091 }
2092
2093 fn rust_lang() -> Language {
2094 Language::new(
2095 LanguageConfig {
2096 name: "Rust".into(),
2097 matcher: LanguageMatcher {
2098 path_suffixes: vec!["rs".to_string()],
2099 ..Default::default()
2100 },
2101 ..Default::default()
2102 },
2103 Some(tree_sitter_rust::language()),
2104 )
2105 .with_indents_query(
2106 r#"
2107 (call_expression) @indent
2108 (field_expression) @indent
2109 (_ "(" ")" @end) @indent
2110 (_ "{" "}" @end) @indent
2111 "#,
2112 )
2113 .unwrap()
2114 }
2115}