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