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