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