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