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, Default, Debug, PartialEq, Eq, Hash)]
1284pub struct InlineAssistId(usize);
1285
1286impl InlineAssistId {
1287 fn post_inc(&mut self) -> InlineAssistId {
1288 let id = *self;
1289 self.0 += 1;
1290 id
1291 }
1292}
1293
1294#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1295struct InlineAssistGroupId(usize);
1296
1297impl InlineAssistGroupId {
1298 fn post_inc(&mut self) -> InlineAssistGroupId {
1299 let id = *self;
1300 self.0 += 1;
1301 id
1302 }
1303}
1304
1305enum PromptEditorEvent {
1306 StartRequested,
1307 StopRequested,
1308 ConfirmRequested,
1309 CancelRequested,
1310 DismissRequested,
1311}
1312
1313struct PromptEditor {
1314 id: InlineAssistId,
1315 fs: Arc<dyn Fs>,
1316 editor: View<Editor>,
1317 edited_since_done: bool,
1318 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1319 prompt_history: VecDeque<String>,
1320 prompt_history_ix: Option<usize>,
1321 pending_prompt: String,
1322 codegen: Model<Codegen>,
1323 _codegen_subscription: Subscription,
1324 editor_subscriptions: Vec<Subscription>,
1325 pending_token_count: Task<Result<()>>,
1326 token_counts: Option<TokenCounts>,
1327 _token_count_subscriptions: Vec<Subscription>,
1328 workspace: Option<WeakView<Workspace>>,
1329 show_rate_limit_notice: bool,
1330}
1331
1332#[derive(Copy, Clone)]
1333pub struct TokenCounts {
1334 total: usize,
1335 assistant_panel: usize,
1336}
1337
1338impl EventEmitter<PromptEditorEvent> for PromptEditor {}
1339
1340impl Render for PromptEditor {
1341 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1342 let gutter_dimensions = *self.gutter_dimensions.lock();
1343 let status = &self.codegen.read(cx).status;
1344 let buttons = match status {
1345 CodegenStatus::Idle => {
1346 vec![
1347 IconButton::new("cancel", IconName::Close)
1348 .icon_color(Color::Muted)
1349 .shape(IconButtonShape::Square)
1350 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1351 .on_click(
1352 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1353 ),
1354 IconButton::new("start", IconName::SparkleAlt)
1355 .icon_color(Color::Muted)
1356 .shape(IconButtonShape::Square)
1357 .tooltip(|cx| Tooltip::for_action("Transform", &menu::Confirm, cx))
1358 .on_click(
1359 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
1360 ),
1361 ]
1362 }
1363 CodegenStatus::Pending => {
1364 vec![
1365 IconButton::new("cancel", IconName::Close)
1366 .icon_color(Color::Muted)
1367 .shape(IconButtonShape::Square)
1368 .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
1369 .on_click(
1370 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1371 ),
1372 IconButton::new("stop", IconName::Stop)
1373 .icon_color(Color::Error)
1374 .shape(IconButtonShape::Square)
1375 .tooltip(|cx| {
1376 Tooltip::with_meta(
1377 "Interrupt Transformation",
1378 Some(&menu::Cancel),
1379 "Changes won't be discarded",
1380 cx,
1381 )
1382 })
1383 .on_click(
1384 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
1385 ),
1386 ]
1387 }
1388 CodegenStatus::Error(_) | CodegenStatus::Done => {
1389 vec![
1390 IconButton::new("cancel", IconName::Close)
1391 .icon_color(Color::Muted)
1392 .shape(IconButtonShape::Square)
1393 .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1394 .on_click(
1395 cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1396 ),
1397 if self.edited_since_done || matches!(status, CodegenStatus::Error(_)) {
1398 IconButton::new("restart", IconName::RotateCw)
1399 .icon_color(Color::Info)
1400 .shape(IconButtonShape::Square)
1401 .tooltip(|cx| {
1402 Tooltip::with_meta(
1403 "Restart Transformation",
1404 Some(&menu::Confirm),
1405 "Changes will be discarded",
1406 cx,
1407 )
1408 })
1409 .on_click(cx.listener(|_, _, cx| {
1410 cx.emit(PromptEditorEvent::StartRequested);
1411 }))
1412 } else {
1413 IconButton::new("confirm", IconName::Check)
1414 .icon_color(Color::Info)
1415 .shape(IconButtonShape::Square)
1416 .tooltip(|cx| Tooltip::for_action("Confirm Assist", &menu::Confirm, cx))
1417 .on_click(cx.listener(|_, _, cx| {
1418 cx.emit(PromptEditorEvent::ConfirmRequested);
1419 }))
1420 },
1421 ]
1422 }
1423 };
1424
1425 h_flex()
1426 .bg(cx.theme().colors().editor_background)
1427 .border_y_1()
1428 .border_color(cx.theme().status().info_border)
1429 .size_full()
1430 .py(cx.line_height() / 2.)
1431 .on_action(cx.listener(Self::confirm))
1432 .on_action(cx.listener(Self::cancel))
1433 .on_action(cx.listener(Self::move_up))
1434 .on_action(cx.listener(Self::move_down))
1435 .child(
1436 h_flex()
1437 .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
1438 .justify_center()
1439 .gap_2()
1440 .child(
1441 ModelSelector::new(
1442 self.fs.clone(),
1443 IconButton::new("context", IconName::SlidersAlt)
1444 .shape(IconButtonShape::Square)
1445 .icon_size(IconSize::Small)
1446 .icon_color(Color::Muted)
1447 .tooltip(move |cx| {
1448 Tooltip::with_meta(
1449 format!(
1450 "Using {}",
1451 LanguageModelRegistry::read_global(cx)
1452 .active_model()
1453 .map(|model| model.name().0)
1454 .unwrap_or_else(|| "No model selected".into()),
1455 ),
1456 None,
1457 "Change Model",
1458 cx,
1459 )
1460 }),
1461 )
1462 .with_info_text(
1463 "Inline edits use context\n\
1464 from the currently selected\n\
1465 assistant panel tab.",
1466 ),
1467 )
1468 .map(|el| {
1469 let CodegenStatus::Error(error) = &self.codegen.read(cx).status else {
1470 return el;
1471 };
1472
1473 let error_message = SharedString::from(error.to_string());
1474 if error.error_code() == proto::ErrorCode::RateLimitExceeded
1475 && cx.has_flag::<ZedPro>()
1476 {
1477 el.child(
1478 v_flex()
1479 .child(
1480 IconButton::new("rate-limit-error", IconName::XCircle)
1481 .selected(self.show_rate_limit_notice)
1482 .shape(IconButtonShape::Square)
1483 .icon_size(IconSize::Small)
1484 .on_click(cx.listener(Self::toggle_rate_limit_notice)),
1485 )
1486 .children(self.show_rate_limit_notice.then(|| {
1487 deferred(
1488 anchored()
1489 .position_mode(gpui::AnchoredPositionMode::Local)
1490 .position(point(px(0.), px(24.)))
1491 .anchor(gpui::AnchorCorner::TopLeft)
1492 .child(self.render_rate_limit_notice(cx)),
1493 )
1494 })),
1495 )
1496 } else {
1497 el.child(
1498 div()
1499 .id("error")
1500 .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
1501 .child(
1502 Icon::new(IconName::XCircle)
1503 .size(IconSize::Small)
1504 .color(Color::Error),
1505 ),
1506 )
1507 }
1508 }),
1509 )
1510 .child(div().flex_1().child(self.render_prompt_editor(cx)))
1511 .child(
1512 h_flex()
1513 .gap_2()
1514 .pr_6()
1515 .children(self.render_token_count(cx))
1516 .children(buttons),
1517 )
1518 }
1519}
1520
1521impl FocusableView for PromptEditor {
1522 fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
1523 self.editor.focus_handle(cx)
1524 }
1525}
1526
1527impl PromptEditor {
1528 const MAX_LINES: u8 = 8;
1529
1530 #[allow(clippy::too_many_arguments)]
1531 fn new(
1532 id: InlineAssistId,
1533 gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1534 prompt_history: VecDeque<String>,
1535 prompt_buffer: Model<MultiBuffer>,
1536 codegen: Model<Codegen>,
1537 parent_editor: &View<Editor>,
1538 assistant_panel: Option<&View<AssistantPanel>>,
1539 workspace: Option<WeakView<Workspace>>,
1540 fs: Arc<dyn Fs>,
1541 cx: &mut ViewContext<Self>,
1542 ) -> Self {
1543 let prompt_editor = cx.new_view(|cx| {
1544 let mut editor = Editor::new(
1545 EditorMode::AutoHeight {
1546 max_lines: Self::MAX_LINES as usize,
1547 },
1548 prompt_buffer,
1549 None,
1550 false,
1551 cx,
1552 );
1553 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1554 // Since the prompt editors for all inline assistants are linked,
1555 // always show the cursor (even when it isn't focused) because
1556 // typing in one will make what you typed appear in all of them.
1557 editor.set_show_cursor_when_unfocused(true, cx);
1558 editor.set_placeholder_text("Add a prompt…", cx);
1559 editor
1560 });
1561
1562 let mut token_count_subscriptions = Vec::new();
1563 token_count_subscriptions
1564 .push(cx.subscribe(parent_editor, Self::handle_parent_editor_event));
1565 if let Some(assistant_panel) = assistant_panel {
1566 token_count_subscriptions
1567 .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
1568 }
1569
1570 let mut this = Self {
1571 id,
1572 editor: prompt_editor,
1573 edited_since_done: false,
1574 gutter_dimensions,
1575 prompt_history,
1576 prompt_history_ix: None,
1577 pending_prompt: String::new(),
1578 _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
1579 editor_subscriptions: Vec::new(),
1580 codegen,
1581 fs,
1582 pending_token_count: Task::ready(Ok(())),
1583 token_counts: None,
1584 _token_count_subscriptions: token_count_subscriptions,
1585 workspace,
1586 show_rate_limit_notice: false,
1587 };
1588 this.count_tokens(cx);
1589 this.subscribe_to_editor(cx);
1590 this
1591 }
1592
1593 fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
1594 self.editor_subscriptions.clear();
1595 self.editor_subscriptions
1596 .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
1597 }
1598
1599 fn set_show_cursor_when_unfocused(
1600 &mut self,
1601 show_cursor_when_unfocused: bool,
1602 cx: &mut ViewContext<Self>,
1603 ) {
1604 self.editor.update(cx, |editor, cx| {
1605 editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
1606 });
1607 }
1608
1609 fn unlink(&mut self, cx: &mut ViewContext<Self>) {
1610 let prompt = self.prompt(cx);
1611 let focus = self.editor.focus_handle(cx).contains_focused(cx);
1612 self.editor = cx.new_view(|cx| {
1613 let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
1614 editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1615 editor.set_placeholder_text("Add a prompt…", cx);
1616 editor.set_text(prompt, cx);
1617 if focus {
1618 editor.focus(cx);
1619 }
1620 editor
1621 });
1622 self.subscribe_to_editor(cx);
1623 }
1624
1625 fn prompt(&self, cx: &AppContext) -> String {
1626 self.editor.read(cx).text(cx)
1627 }
1628
1629 fn toggle_rate_limit_notice(&mut self, _: &ClickEvent, cx: &mut ViewContext<Self>) {
1630 self.show_rate_limit_notice = !self.show_rate_limit_notice;
1631 if self.show_rate_limit_notice {
1632 cx.focus_view(&self.editor);
1633 }
1634 cx.notify();
1635 }
1636
1637 fn handle_parent_editor_event(
1638 &mut self,
1639 _: View<Editor>,
1640 event: &EditorEvent,
1641 cx: &mut ViewContext<Self>,
1642 ) {
1643 if let EditorEvent::BufferEdited { .. } = event {
1644 self.count_tokens(cx);
1645 }
1646 }
1647
1648 fn handle_assistant_panel_event(
1649 &mut self,
1650 _: View<AssistantPanel>,
1651 event: &AssistantPanelEvent,
1652 cx: &mut ViewContext<Self>,
1653 ) {
1654 let AssistantPanelEvent::ContextEdited { .. } = event;
1655 self.count_tokens(cx);
1656 }
1657
1658 fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
1659 let assist_id = self.id;
1660 self.pending_token_count = cx.spawn(|this, mut cx| async move {
1661 cx.background_executor().timer(Duration::from_secs(1)).await;
1662 let token_count = cx
1663 .update_global(|inline_assistant: &mut InlineAssistant, cx| {
1664 let assist = inline_assistant
1665 .assists
1666 .get(&assist_id)
1667 .context("assist not found")?;
1668 anyhow::Ok(assist.count_tokens(cx))
1669 })??
1670 .await?;
1671
1672 this.update(&mut cx, |this, cx| {
1673 this.token_counts = Some(token_count);
1674 cx.notify();
1675 })
1676 })
1677 }
1678
1679 fn handle_prompt_editor_events(
1680 &mut self,
1681 _: View<Editor>,
1682 event: &EditorEvent,
1683 cx: &mut ViewContext<Self>,
1684 ) {
1685 match event {
1686 EditorEvent::Edited { .. } => {
1687 let prompt = self.editor.read(cx).text(cx);
1688 if self
1689 .prompt_history_ix
1690 .map_or(true, |ix| self.prompt_history[ix] != prompt)
1691 {
1692 self.prompt_history_ix.take();
1693 self.pending_prompt = prompt;
1694 }
1695
1696 self.edited_since_done = true;
1697 cx.notify();
1698 }
1699 EditorEvent::BufferEdited => {
1700 self.count_tokens(cx);
1701 }
1702 EditorEvent::Blurred => {
1703 if self.show_rate_limit_notice {
1704 self.show_rate_limit_notice = false;
1705 cx.notify();
1706 }
1707 }
1708 _ => {}
1709 }
1710 }
1711
1712 fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
1713 match &self.codegen.read(cx).status {
1714 CodegenStatus::Idle => {
1715 self.editor
1716 .update(cx, |editor, _| editor.set_read_only(false));
1717 }
1718 CodegenStatus::Pending => {
1719 self.editor
1720 .update(cx, |editor, _| editor.set_read_only(true));
1721 }
1722 CodegenStatus::Done => {
1723 self.edited_since_done = false;
1724 self.editor
1725 .update(cx, |editor, _| editor.set_read_only(false));
1726 }
1727 CodegenStatus::Error(error) => {
1728 if cx.has_flag::<ZedPro>()
1729 && error.error_code() == proto::ErrorCode::RateLimitExceeded
1730 && !dismissed_rate_limit_notice()
1731 {
1732 self.show_rate_limit_notice = true;
1733 cx.notify();
1734 }
1735
1736 self.edited_since_done = false;
1737 self.editor
1738 .update(cx, |editor, _| editor.set_read_only(false));
1739 }
1740 }
1741 }
1742
1743 fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
1744 match &self.codegen.read(cx).status {
1745 CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1746 cx.emit(PromptEditorEvent::CancelRequested);
1747 }
1748 CodegenStatus::Pending => {
1749 cx.emit(PromptEditorEvent::StopRequested);
1750 }
1751 }
1752 }
1753
1754 fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
1755 match &self.codegen.read(cx).status {
1756 CodegenStatus::Idle => {
1757 cx.emit(PromptEditorEvent::StartRequested);
1758 }
1759 CodegenStatus::Pending => {
1760 cx.emit(PromptEditorEvent::DismissRequested);
1761 }
1762 CodegenStatus::Done | CodegenStatus::Error(_) => {
1763 if self.edited_since_done {
1764 cx.emit(PromptEditorEvent::StartRequested);
1765 } else {
1766 cx.emit(PromptEditorEvent::ConfirmRequested);
1767 }
1768 }
1769 }
1770 }
1771
1772 fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
1773 if let Some(ix) = self.prompt_history_ix {
1774 if ix > 0 {
1775 self.prompt_history_ix = Some(ix - 1);
1776 let prompt = self.prompt_history[ix - 1].as_str();
1777 self.editor.update(cx, |editor, cx| {
1778 editor.set_text(prompt, cx);
1779 editor.move_to_beginning(&Default::default(), cx);
1780 });
1781 }
1782 } else if !self.prompt_history.is_empty() {
1783 self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1784 let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1785 self.editor.update(cx, |editor, cx| {
1786 editor.set_text(prompt, cx);
1787 editor.move_to_beginning(&Default::default(), cx);
1788 });
1789 }
1790 }
1791
1792 fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
1793 if let Some(ix) = self.prompt_history_ix {
1794 if ix < self.prompt_history.len() - 1 {
1795 self.prompt_history_ix = Some(ix + 1);
1796 let prompt = self.prompt_history[ix + 1].as_str();
1797 self.editor.update(cx, |editor, cx| {
1798 editor.set_text(prompt, cx);
1799 editor.move_to_end(&Default::default(), cx)
1800 });
1801 } else {
1802 self.prompt_history_ix = None;
1803 let prompt = self.pending_prompt.as_str();
1804 self.editor.update(cx, |editor, cx| {
1805 editor.set_text(prompt, cx);
1806 editor.move_to_end(&Default::default(), cx)
1807 });
1808 }
1809 }
1810 }
1811
1812 fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
1813 let model = LanguageModelRegistry::read_global(cx).active_model()?;
1814 let token_counts = self.token_counts?;
1815 let max_token_count = model.max_token_count();
1816
1817 let remaining_tokens = max_token_count as isize - token_counts.total as isize;
1818 let token_count_color = if remaining_tokens <= 0 {
1819 Color::Error
1820 } else if token_counts.total as f32 / max_token_count as f32 >= 0.8 {
1821 Color::Warning
1822 } else {
1823 Color::Muted
1824 };
1825
1826 let mut token_count = h_flex()
1827 .id("token_count")
1828 .gap_0p5()
1829 .child(
1830 Label::new(humanize_token_count(token_counts.total))
1831 .size(LabelSize::Small)
1832 .color(token_count_color),
1833 )
1834 .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
1835 .child(
1836 Label::new(humanize_token_count(max_token_count))
1837 .size(LabelSize::Small)
1838 .color(Color::Muted),
1839 );
1840 if let Some(workspace) = self.workspace.clone() {
1841 token_count = token_count
1842 .tooltip(move |cx| {
1843 Tooltip::with_meta(
1844 format!(
1845 "Tokens Used ({} from the Assistant Panel)",
1846 humanize_token_count(token_counts.assistant_panel)
1847 ),
1848 None,
1849 "Click to open the Assistant Panel",
1850 cx,
1851 )
1852 })
1853 .cursor_pointer()
1854 .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
1855 .on_click(move |_, cx| {
1856 cx.stop_propagation();
1857 workspace
1858 .update(cx, |workspace, cx| {
1859 workspace.focus_panel::<AssistantPanel>(cx)
1860 })
1861 .ok();
1862 });
1863 } else {
1864 token_count = token_count
1865 .cursor_default()
1866 .tooltip(|cx| Tooltip::text("Tokens used", cx));
1867 }
1868
1869 Some(token_count)
1870 }
1871
1872 fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1873 let settings = ThemeSettings::get_global(cx);
1874 let text_style = TextStyle {
1875 color: if self.editor.read(cx).read_only(cx) {
1876 cx.theme().colors().text_disabled
1877 } else {
1878 cx.theme().colors().text
1879 },
1880 font_family: settings.ui_font.family.clone(),
1881 font_features: settings.ui_font.features.clone(),
1882 font_fallbacks: settings.ui_font.fallbacks.clone(),
1883 font_size: rems(0.875).into(),
1884 font_weight: settings.ui_font.weight,
1885 line_height: relative(1.3),
1886 ..Default::default()
1887 };
1888 EditorElement::new(
1889 &self.editor,
1890 EditorStyle {
1891 background: cx.theme().colors().editor_background,
1892 local_player: cx.theme().players().local(),
1893 text: text_style,
1894 ..Default::default()
1895 },
1896 )
1897 }
1898
1899 fn render_rate_limit_notice(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1900 Popover::new().child(
1901 v_flex()
1902 .occlude()
1903 .p_2()
1904 .child(
1905 Label::new("Out of Tokens")
1906 .size(LabelSize::Small)
1907 .weight(FontWeight::BOLD),
1908 )
1909 .child(Label::new(
1910 "Try Zed Pro for higher limits, a wider range of models, and more.",
1911 ))
1912 .child(
1913 h_flex()
1914 .justify_between()
1915 .child(CheckboxWithLabel::new(
1916 "dont-show-again",
1917 Label::new("Don't show again"),
1918 if dismissed_rate_limit_notice() {
1919 ui::Selection::Selected
1920 } else {
1921 ui::Selection::Unselected
1922 },
1923 |selection, cx| {
1924 let is_dismissed = match selection {
1925 ui::Selection::Unselected => false,
1926 ui::Selection::Indeterminate => return,
1927 ui::Selection::Selected => true,
1928 };
1929
1930 set_rate_limit_notice_dismissed(is_dismissed, cx)
1931 },
1932 ))
1933 .child(
1934 h_flex()
1935 .gap_2()
1936 .child(
1937 Button::new("dismiss", "Dismiss")
1938 .style(ButtonStyle::Transparent)
1939 .on_click(cx.listener(Self::toggle_rate_limit_notice)),
1940 )
1941 .child(Button::new("more-info", "More Info").on_click(
1942 |_event, cx| {
1943 cx.dispatch_action(Box::new(
1944 zed_actions::OpenAccountSettings,
1945 ))
1946 },
1947 )),
1948 ),
1949 ),
1950 )
1951 }
1952}
1953
1954const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
1955
1956fn dismissed_rate_limit_notice() -> bool {
1957 db::kvp::KEY_VALUE_STORE
1958 .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
1959 .log_err()
1960 .map_or(false, |s| s.is_some())
1961}
1962
1963fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut AppContext) {
1964 db::write_and_log(cx, move || async move {
1965 if is_dismissed {
1966 db::kvp::KEY_VALUE_STORE
1967 .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
1968 .await
1969 } else {
1970 db::kvp::KEY_VALUE_STORE
1971 .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
1972 .await
1973 }
1974 })
1975}
1976
1977struct InlineAssist {
1978 group_id: InlineAssistGroupId,
1979 range: Range<Anchor>,
1980 editor: WeakView<Editor>,
1981 decorations: Option<InlineAssistDecorations>,
1982 codegen: Model<Codegen>,
1983 _subscriptions: Vec<Subscription>,
1984 workspace: Option<WeakView<Workspace>>,
1985 include_context: bool,
1986}
1987
1988impl InlineAssist {
1989 #[allow(clippy::too_many_arguments)]
1990 fn new(
1991 assist_id: InlineAssistId,
1992 group_id: InlineAssistGroupId,
1993 include_context: bool,
1994 editor: &View<Editor>,
1995 prompt_editor: &View<PromptEditor>,
1996 prompt_block_id: CustomBlockId,
1997 end_block_id: CustomBlockId,
1998 range: Range<Anchor>,
1999 codegen: Model<Codegen>,
2000 workspace: Option<WeakView<Workspace>>,
2001 cx: &mut WindowContext,
2002 ) -> Self {
2003 let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
2004 InlineAssist {
2005 group_id,
2006 include_context,
2007 editor: editor.downgrade(),
2008 decorations: Some(InlineAssistDecorations {
2009 prompt_block_id,
2010 prompt_editor: prompt_editor.clone(),
2011 removed_line_block_ids: HashSet::default(),
2012 end_block_id,
2013 }),
2014 range,
2015 codegen: codegen.clone(),
2016 workspace: workspace.clone(),
2017 _subscriptions: vec![
2018 cx.on_focus_in(&prompt_editor_focus_handle, move |cx| {
2019 InlineAssistant::update_global(cx, |this, cx| {
2020 this.handle_prompt_editor_focus_in(assist_id, cx)
2021 })
2022 }),
2023 cx.on_focus_out(&prompt_editor_focus_handle, move |_, cx| {
2024 InlineAssistant::update_global(cx, |this, cx| {
2025 this.handle_prompt_editor_focus_out(assist_id, cx)
2026 })
2027 }),
2028 cx.subscribe(prompt_editor, |prompt_editor, event, cx| {
2029 InlineAssistant::update_global(cx, |this, cx| {
2030 this.handle_prompt_editor_event(prompt_editor, event, cx)
2031 })
2032 }),
2033 cx.observe(&codegen, {
2034 let editor = editor.downgrade();
2035 move |_, cx| {
2036 if let Some(editor) = editor.upgrade() {
2037 InlineAssistant::update_global(cx, |this, cx| {
2038 if let Some(editor_assists) =
2039 this.assists_by_editor.get(&editor.downgrade())
2040 {
2041 editor_assists.highlight_updates.send(()).ok();
2042 }
2043
2044 this.update_editor_blocks(&editor, assist_id, cx);
2045 })
2046 }
2047 }
2048 }),
2049 cx.subscribe(&codegen, move |codegen, event, cx| {
2050 InlineAssistant::update_global(cx, |this, cx| match event {
2051 CodegenEvent::Undone => this.finish_assist(assist_id, false, cx),
2052 CodegenEvent::Finished => {
2053 let assist = if let Some(assist) = this.assists.get(&assist_id) {
2054 assist
2055 } else {
2056 return;
2057 };
2058
2059 if let CodegenStatus::Error(error) = &codegen.read(cx).status {
2060 if assist.decorations.is_none() {
2061 if let Some(workspace) = assist
2062 .workspace
2063 .as_ref()
2064 .and_then(|workspace| workspace.upgrade())
2065 {
2066 let error = format!("Inline assistant error: {}", error);
2067 workspace.update(cx, |workspace, cx| {
2068 struct InlineAssistantError;
2069
2070 let id =
2071 NotificationId::identified::<InlineAssistantError>(
2072 assist_id.0,
2073 );
2074
2075 workspace.show_toast(Toast::new(id, error), cx);
2076 })
2077 }
2078 }
2079 }
2080
2081 if assist.decorations.is_none() {
2082 this.finish_assist(assist_id, false, cx);
2083 } else if let Some(tx) = this.assist_observations.get(&assist_id) {
2084 tx.0.send(()).ok();
2085 }
2086 }
2087 })
2088 }),
2089 ],
2090 }
2091 }
2092
2093 fn user_prompt(&self, cx: &AppContext) -> Option<String> {
2094 let decorations = self.decorations.as_ref()?;
2095 Some(decorations.prompt_editor.read(cx).prompt(cx))
2096 }
2097
2098 fn assistant_panel_context(&self, cx: &WindowContext) -> Option<LanguageModelRequest> {
2099 if self.include_context {
2100 let workspace = self.workspace.as_ref()?;
2101 let workspace = workspace.upgrade()?.read(cx);
2102 let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
2103 Some(
2104 assistant_panel
2105 .read(cx)
2106 .active_context(cx)?
2107 .read(cx)
2108 .to_completion_request(cx),
2109 )
2110 } else {
2111 None
2112 }
2113 }
2114
2115 pub fn count_tokens(&self, cx: &WindowContext) -> BoxFuture<'static, Result<TokenCounts>> {
2116 let Some(user_prompt) = self.user_prompt(cx) else {
2117 return future::ready(Err(anyhow!("no user prompt"))).boxed();
2118 };
2119 let assistant_panel_context = self.assistant_panel_context(cx);
2120 self.codegen
2121 .read(cx)
2122 .count_tokens(user_prompt, assistant_panel_context, cx)
2123 }
2124}
2125
2126struct InlineAssistDecorations {
2127 prompt_block_id: CustomBlockId,
2128 prompt_editor: View<PromptEditor>,
2129 removed_line_block_ids: HashSet<CustomBlockId>,
2130 end_block_id: CustomBlockId,
2131}
2132
2133#[derive(Debug)]
2134pub enum CodegenEvent {
2135 Finished,
2136 Undone,
2137}
2138
2139pub struct Codegen {
2140 buffer: Model<MultiBuffer>,
2141 old_buffer: Model<Buffer>,
2142 snapshot: MultiBufferSnapshot,
2143 transform_range: Range<Anchor>,
2144 selected_ranges: Vec<Range<Anchor>>,
2145 edit_position: Option<Anchor>,
2146 last_equal_ranges: Vec<Range<Anchor>>,
2147 initial_transaction_id: Option<TransactionId>,
2148 transformation_transaction_id: Option<TransactionId>,
2149 status: CodegenStatus,
2150 generation: Task<()>,
2151 diff: Diff,
2152 telemetry: Option<Arc<Telemetry>>,
2153 _subscription: gpui::Subscription,
2154 prompt_builder: Arc<PromptBuilder>,
2155}
2156
2157enum CodegenStatus {
2158 Idle,
2159 Pending,
2160 Done,
2161 Error(anyhow::Error),
2162}
2163
2164#[derive(Default)]
2165struct Diff {
2166 deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
2167 inserted_row_ranges: Vec<RangeInclusive<Anchor>>,
2168}
2169
2170impl Diff {
2171 fn is_empty(&self) -> bool {
2172 self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
2173 }
2174}
2175
2176impl EventEmitter<CodegenEvent> for Codegen {}
2177
2178impl Codegen {
2179 pub fn new(
2180 buffer: Model<MultiBuffer>,
2181 transform_range: Range<Anchor>,
2182 selected_ranges: Vec<Range<Anchor>>,
2183 initial_transaction_id: Option<TransactionId>,
2184 telemetry: Option<Arc<Telemetry>>,
2185 builder: Arc<PromptBuilder>,
2186 cx: &mut ModelContext<Self>,
2187 ) -> Self {
2188 let snapshot = buffer.read(cx).snapshot(cx);
2189
2190 let (old_buffer, _, _) = buffer
2191 .read(cx)
2192 .range_to_buffer_ranges(transform_range.clone(), cx)
2193 .pop()
2194 .unwrap();
2195 let old_buffer = cx.new_model(|cx| {
2196 let old_buffer = old_buffer.read(cx);
2197 let text = old_buffer.as_rope().clone();
2198 let line_ending = old_buffer.line_ending();
2199 let language = old_buffer.language().cloned();
2200 let language_registry = old_buffer.language_registry();
2201
2202 let mut buffer = Buffer::local_normalized(text, line_ending, cx);
2203 buffer.set_language(language, cx);
2204 if let Some(language_registry) = language_registry {
2205 buffer.set_language_registry(language_registry)
2206 }
2207 buffer
2208 });
2209
2210 Self {
2211 buffer: buffer.clone(),
2212 old_buffer,
2213 edit_position: None,
2214 snapshot,
2215 last_equal_ranges: Default::default(),
2216 transformation_transaction_id: None,
2217 status: CodegenStatus::Idle,
2218 generation: Task::ready(()),
2219 diff: Diff::default(),
2220 telemetry,
2221 _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
2222 initial_transaction_id,
2223 prompt_builder: builder,
2224 transform_range,
2225 selected_ranges,
2226 }
2227 }
2228
2229 fn handle_buffer_event(
2230 &mut self,
2231 _buffer: Model<MultiBuffer>,
2232 event: &multi_buffer::Event,
2233 cx: &mut ModelContext<Self>,
2234 ) {
2235 if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
2236 if self.transformation_transaction_id == Some(*transaction_id) {
2237 self.transformation_transaction_id = None;
2238 self.generation = Task::ready(());
2239 cx.emit(CodegenEvent::Undone);
2240 }
2241 }
2242 }
2243
2244 pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
2245 &self.last_equal_ranges
2246 }
2247
2248 pub fn count_tokens(
2249 &self,
2250 user_prompt: String,
2251 assistant_panel_context: Option<LanguageModelRequest>,
2252 cx: &AppContext,
2253 ) -> BoxFuture<'static, Result<TokenCounts>> {
2254 if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
2255 let request = self.build_request(user_prompt, assistant_panel_context.clone(), cx);
2256 match request {
2257 Ok(request) => {
2258 let total_count = model.count_tokens(request.clone(), cx);
2259 let assistant_panel_count = assistant_panel_context
2260 .map(|context| model.count_tokens(context, cx))
2261 .unwrap_or_else(|| future::ready(Ok(0)).boxed());
2262
2263 async move {
2264 Ok(TokenCounts {
2265 total: total_count.await?,
2266 assistant_panel: assistant_panel_count.await?,
2267 })
2268 }
2269 .boxed()
2270 }
2271 Err(error) => futures::future::ready(Err(error)).boxed(),
2272 }
2273 } else {
2274 future::ready(Err(anyhow!("no active model"))).boxed()
2275 }
2276 }
2277
2278 pub fn start(
2279 &mut self,
2280 user_prompt: String,
2281 assistant_panel_context: Option<LanguageModelRequest>,
2282 cx: &mut ModelContext<Self>,
2283 ) -> Result<()> {
2284 let model = LanguageModelRegistry::read_global(cx)
2285 .active_model()
2286 .context("no active model")?;
2287
2288 if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {
2289 self.buffer.update(cx, |buffer, cx| {
2290 buffer.undo_transaction(transformation_transaction_id, cx);
2291 });
2292 }
2293
2294 self.edit_position = Some(self.transform_range.start.bias_right(&self.snapshot));
2295
2296 let telemetry_id = model.telemetry_id();
2297 let chunks: LocalBoxFuture<Result<BoxStream<Result<String>>>> =
2298 if user_prompt.trim().to_lowercase() == "delete" {
2299 async { Ok(stream::empty().boxed()) }.boxed_local()
2300 } else {
2301 let request = self.build_request(user_prompt, assistant_panel_context, cx)?;
2302
2303 let chunks =
2304 cx.spawn(|_, cx| async move { model.stream_completion(request, &cx).await });
2305 async move { Ok(chunks.await?.boxed()) }.boxed_local()
2306 };
2307 self.handle_stream(telemetry_id, self.transform_range.clone(), chunks, cx);
2308 Ok(())
2309 }
2310
2311 fn build_request(
2312 &self,
2313 user_prompt: String,
2314 assistant_panel_context: Option<LanguageModelRequest>,
2315 cx: &AppContext,
2316 ) -> Result<LanguageModelRequest> {
2317 let buffer = self.buffer.read(cx).snapshot(cx);
2318 let language = buffer.language_at(self.transform_range.start);
2319 let language_name = if let Some(language) = language.as_ref() {
2320 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
2321 None
2322 } else {
2323 Some(language.name())
2324 }
2325 } else {
2326 None
2327 };
2328
2329 // Higher Temperature increases the randomness of model outputs.
2330 // If Markdown or No Language is Known, increase the randomness for more creative output
2331 // If Code, decrease temperature to get more deterministic outputs
2332 let temperature = if let Some(language) = language_name.clone() {
2333 if language.as_ref() == "Markdown" {
2334 1.0
2335 } else {
2336 0.5
2337 }
2338 } else {
2339 1.0
2340 };
2341
2342 let language_name = language_name.as_deref();
2343 let start = buffer.point_to_buffer_offset(self.transform_range.start);
2344 let end = buffer.point_to_buffer_offset(self.transform_range.end);
2345 let (transform_buffer, transform_range) = if let Some((start, end)) = start.zip(end) {
2346 let (start_buffer, start_buffer_offset) = start;
2347 let (end_buffer, end_buffer_offset) = end;
2348 if start_buffer.remote_id() == end_buffer.remote_id() {
2349 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
2350 } else {
2351 return Err(anyhow::anyhow!("invalid transformation range"));
2352 }
2353 } else {
2354 return Err(anyhow::anyhow!("invalid transformation range"));
2355 };
2356
2357 let selected_ranges = self
2358 .selected_ranges
2359 .iter()
2360 .filter_map(|selected_range| {
2361 let start = buffer
2362 .point_to_buffer_offset(selected_range.start)
2363 .map(|(_, offset)| offset)?;
2364 let end = buffer
2365 .point_to_buffer_offset(selected_range.end)
2366 .map(|(_, offset)| offset)?;
2367 Some(start..end)
2368 })
2369 .collect::<Vec<_>>();
2370
2371 let prompt = self
2372 .prompt_builder
2373 .generate_content_prompt(
2374 user_prompt,
2375 language_name,
2376 transform_buffer,
2377 transform_range,
2378 selected_ranges,
2379 )
2380 .map_err(|e| anyhow::anyhow!("Failed to generate content prompt: {}", e))?;
2381
2382 let mut messages = Vec::new();
2383 if let Some(context_request) = assistant_panel_context {
2384 messages = context_request.messages;
2385 }
2386
2387 messages.push(LanguageModelRequestMessage {
2388 role: Role::User,
2389 content: vec![prompt.into()],
2390 });
2391
2392 Ok(LanguageModelRequest {
2393 messages,
2394 stop: vec!["|END|>".to_string()],
2395 temperature,
2396 })
2397 }
2398
2399 pub fn handle_stream(
2400 &mut self,
2401 model_telemetry_id: String,
2402 edit_range: Range<Anchor>,
2403 stream: impl 'static + Future<Output = Result<BoxStream<'static, Result<String>>>>,
2404 cx: &mut ModelContext<Self>,
2405 ) {
2406 let snapshot = self.snapshot.clone();
2407 let selected_text = snapshot
2408 .text_for_range(edit_range.start..edit_range.end)
2409 .collect::<Rope>();
2410
2411 let selection_start = edit_range.start.to_point(&snapshot);
2412
2413 // Start with the indentation of the first line in the selection
2414 let mut suggested_line_indent = snapshot
2415 .suggested_indents(selection_start.row..=selection_start.row, cx)
2416 .into_values()
2417 .next()
2418 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
2419
2420 // If the first line in the selection does not have indentation, check the following lines
2421 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
2422 for row in selection_start.row..=edit_range.end.to_point(&snapshot).row {
2423 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
2424 // Prefer tabs if a line in the selection uses tabs as indentation
2425 if line_indent.kind == IndentKind::Tab {
2426 suggested_line_indent.kind = IndentKind::Tab;
2427 break;
2428 }
2429 }
2430 }
2431
2432 let telemetry = self.telemetry.clone();
2433 self.diff = Diff::default();
2434 self.status = CodegenStatus::Pending;
2435 let mut edit_start = edit_range.start.to_offset(&snapshot);
2436 self.generation = cx.spawn(|this, mut cx| {
2437 async move {
2438 let chunks = stream.await;
2439 let generate = async {
2440 let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
2441 let diff: Task<anyhow::Result<()>> =
2442 cx.background_executor().spawn(async move {
2443 let mut response_latency = None;
2444 let request_start = Instant::now();
2445 let diff = async {
2446 let chunks = StripInvalidSpans::new(chunks?);
2447 futures::pin_mut!(chunks);
2448 let mut diff = StreamingDiff::new(selected_text.to_string());
2449 let mut line_diff = LineDiff::default();
2450
2451 while let Some(chunk) = chunks.next().await {
2452 if response_latency.is_none() {
2453 response_latency = Some(request_start.elapsed());
2454 }
2455 let chunk = chunk?;
2456 let char_ops = diff.push_new(&chunk);
2457 line_diff.push_char_operations(&char_ops, &selected_text);
2458 diff_tx
2459 .send((char_ops, line_diff.line_operations()))
2460 .await?;
2461 }
2462
2463 let char_ops = diff.finish();
2464 line_diff.push_char_operations(&char_ops, &selected_text);
2465 line_diff.finish(&selected_text);
2466 diff_tx
2467 .send((char_ops, line_diff.line_operations()))
2468 .await?;
2469
2470 anyhow::Ok(())
2471 };
2472
2473 let result = diff.await;
2474
2475 let error_message =
2476 result.as_ref().err().map(|error| error.to_string());
2477 if let Some(telemetry) = telemetry {
2478 telemetry.report_assistant_event(
2479 None,
2480 telemetry_events::AssistantKind::Inline,
2481 model_telemetry_id,
2482 response_latency,
2483 error_message,
2484 );
2485 }
2486
2487 result?;
2488 Ok(())
2489 });
2490
2491 while let Some((char_ops, line_diff)) = diff_rx.next().await {
2492 this.update(&mut cx, |this, cx| {
2493 this.last_equal_ranges.clear();
2494
2495 let transaction = this.buffer.update(cx, |buffer, cx| {
2496 // Avoid grouping assistant edits with user edits.
2497 buffer.finalize_last_transaction(cx);
2498
2499 buffer.start_transaction(cx);
2500 buffer.edit(
2501 char_ops
2502 .into_iter()
2503 .filter_map(|operation| match operation {
2504 CharOperation::Insert { text } => {
2505 let edit_start = snapshot.anchor_after(edit_start);
2506 Some((edit_start..edit_start, text))
2507 }
2508 CharOperation::Delete { bytes } => {
2509 let edit_end = edit_start + bytes;
2510 let edit_range = snapshot.anchor_after(edit_start)
2511 ..snapshot.anchor_before(edit_end);
2512 edit_start = edit_end;
2513 Some((edit_range, String::new()))
2514 }
2515 CharOperation::Keep { bytes } => {
2516 let edit_end = edit_start + bytes;
2517 let edit_range = snapshot.anchor_after(edit_start)
2518 ..snapshot.anchor_before(edit_end);
2519 edit_start = edit_end;
2520 this.last_equal_ranges.push(edit_range);
2521 None
2522 }
2523 }),
2524 None,
2525 cx,
2526 );
2527 this.edit_position = Some(snapshot.anchor_after(edit_start));
2528
2529 buffer.end_transaction(cx)
2530 });
2531
2532 if let Some(transaction) = transaction {
2533 if let Some(first_transaction) = this.transformation_transaction_id
2534 {
2535 // Group all assistant edits into the first transaction.
2536 this.buffer.update(cx, |buffer, cx| {
2537 buffer.merge_transactions(
2538 transaction,
2539 first_transaction,
2540 cx,
2541 )
2542 });
2543 } else {
2544 this.transformation_transaction_id = Some(transaction);
2545 this.buffer.update(cx, |buffer, cx| {
2546 buffer.finalize_last_transaction(cx)
2547 });
2548 }
2549 }
2550
2551 this.update_diff(edit_range.clone(), line_diff, cx);
2552
2553 cx.notify();
2554 })?;
2555 }
2556
2557 diff.await?;
2558
2559 anyhow::Ok(())
2560 };
2561
2562 let result = generate.await;
2563 this.update(&mut cx, |this, cx| {
2564 this.last_equal_ranges.clear();
2565 if let Err(error) = result {
2566 this.status = CodegenStatus::Error(error);
2567 } else {
2568 this.status = CodegenStatus::Done;
2569 }
2570 cx.emit(CodegenEvent::Finished);
2571 cx.notify();
2572 })
2573 .ok();
2574 }
2575 });
2576 cx.notify();
2577 }
2578
2579 pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
2580 self.last_equal_ranges.clear();
2581 if self.diff.is_empty() {
2582 self.status = CodegenStatus::Idle;
2583 } else {
2584 self.status = CodegenStatus::Done;
2585 }
2586 self.generation = Task::ready(());
2587 cx.emit(CodegenEvent::Finished);
2588 cx.notify();
2589 }
2590
2591 pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
2592 self.buffer.update(cx, |buffer, cx| {
2593 if let Some(transaction_id) = self.transformation_transaction_id.take() {
2594 buffer.undo_transaction(transaction_id, cx);
2595 buffer.refresh_preview(cx);
2596 }
2597
2598 if let Some(transaction_id) = self.initial_transaction_id.take() {
2599 buffer.undo_transaction(transaction_id, cx);
2600 buffer.refresh_preview(cx);
2601 }
2602 });
2603 }
2604
2605 fn update_diff(
2606 &mut self,
2607 edit_range: Range<Anchor>,
2608 line_operations: Vec<LineOperation>,
2609 cx: &mut ModelContext<Self>,
2610 ) {
2611 let old_snapshot = self.snapshot.clone();
2612 let old_range = edit_range.to_point(&old_snapshot);
2613 let new_snapshot = self.buffer.read(cx).snapshot(cx);
2614 let new_range = edit_range.to_point(&new_snapshot);
2615
2616 let mut old_row = old_range.start.row;
2617 let mut new_row = new_range.start.row;
2618
2619 self.diff.deleted_row_ranges.clear();
2620 self.diff.inserted_row_ranges.clear();
2621 for operation in line_operations {
2622 match operation {
2623 LineOperation::Keep { lines } => {
2624 old_row += lines;
2625 new_row += lines;
2626 }
2627 LineOperation::Delete { lines } => {
2628 let old_end_row = old_row + lines - 1;
2629 let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
2630
2631 if let Some((_, last_deleted_row_range)) =
2632 self.diff.deleted_row_ranges.last_mut()
2633 {
2634 if *last_deleted_row_range.end() + 1 == old_row {
2635 *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
2636 } else {
2637 self.diff
2638 .deleted_row_ranges
2639 .push((new_row, old_row..=old_end_row));
2640 }
2641 } else {
2642 self.diff
2643 .deleted_row_ranges
2644 .push((new_row, old_row..=old_end_row));
2645 }
2646
2647 old_row += lines;
2648 }
2649 LineOperation::Insert { lines } => {
2650 let new_end_row = new_row + lines - 1;
2651 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
2652 let end = new_snapshot.anchor_before(Point::new(
2653 new_end_row,
2654 new_snapshot.line_len(MultiBufferRow(new_end_row)),
2655 ));
2656 self.diff.inserted_row_ranges.push(start..=end);
2657 new_row += lines;
2658 }
2659 }
2660
2661 cx.notify();
2662 }
2663 }
2664}
2665
2666struct StripInvalidSpans<T> {
2667 stream: T,
2668 stream_done: bool,
2669 buffer: String,
2670 first_line: bool,
2671 line_end: bool,
2672 starts_with_code_block: bool,
2673}
2674
2675impl<T> StripInvalidSpans<T>
2676where
2677 T: Stream<Item = Result<String>>,
2678{
2679 fn new(stream: T) -> Self {
2680 Self {
2681 stream,
2682 stream_done: false,
2683 buffer: String::new(),
2684 first_line: true,
2685 line_end: false,
2686 starts_with_code_block: false,
2687 }
2688 }
2689}
2690
2691impl<T> Stream for StripInvalidSpans<T>
2692where
2693 T: Stream<Item = Result<String>>,
2694{
2695 type Item = Result<String>;
2696
2697 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
2698 const CODE_BLOCK_DELIMITER: &str = "```";
2699 const CURSOR_SPAN: &str = "<|CURSOR|>";
2700
2701 let this = unsafe { self.get_unchecked_mut() };
2702 loop {
2703 if !this.stream_done {
2704 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
2705 match stream.as_mut().poll_next(cx) {
2706 Poll::Ready(Some(Ok(chunk))) => {
2707 this.buffer.push_str(&chunk);
2708 }
2709 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
2710 Poll::Ready(None) => {
2711 this.stream_done = true;
2712 }
2713 Poll::Pending => return Poll::Pending,
2714 }
2715 }
2716
2717 let mut chunk = String::new();
2718 let mut consumed = 0;
2719 if !this.buffer.is_empty() {
2720 let mut lines = this.buffer.split('\n').enumerate().peekable();
2721 while let Some((line_ix, line)) = lines.next() {
2722 if line_ix > 0 {
2723 this.first_line = false;
2724 }
2725
2726 if this.first_line {
2727 let trimmed_line = line.trim();
2728 if lines.peek().is_some() {
2729 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
2730 consumed += line.len() + 1;
2731 this.starts_with_code_block = true;
2732 continue;
2733 }
2734 } else if trimmed_line.is_empty()
2735 || prefixes(CODE_BLOCK_DELIMITER)
2736 .any(|prefix| trimmed_line.starts_with(prefix))
2737 {
2738 break;
2739 }
2740 }
2741
2742 let line_without_cursor = line.replace(CURSOR_SPAN, "");
2743 if lines.peek().is_some() {
2744 if this.line_end {
2745 chunk.push('\n');
2746 }
2747
2748 chunk.push_str(&line_without_cursor);
2749 this.line_end = true;
2750 consumed += line.len() + 1;
2751 } else if this.stream_done {
2752 if !this.starts_with_code_block
2753 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
2754 {
2755 if this.line_end {
2756 chunk.push('\n');
2757 }
2758
2759 chunk.push_str(&line);
2760 }
2761
2762 consumed += line.len();
2763 } else {
2764 let trimmed_line = line.trim();
2765 if trimmed_line.is_empty()
2766 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
2767 || prefixes(CODE_BLOCK_DELIMITER)
2768 .any(|prefix| trimmed_line.ends_with(prefix))
2769 {
2770 break;
2771 } else {
2772 if this.line_end {
2773 chunk.push('\n');
2774 this.line_end = false;
2775 }
2776
2777 chunk.push_str(&line_without_cursor);
2778 consumed += line.len();
2779 }
2780 }
2781 }
2782 }
2783
2784 this.buffer = this.buffer.split_off(consumed);
2785 if !chunk.is_empty() {
2786 return Poll::Ready(Some(Ok(chunk)));
2787 } else if this.stream_done {
2788 return Poll::Ready(None);
2789 }
2790 }
2791 }
2792}
2793
2794fn prefixes(text: &str) -> impl Iterator<Item = &str> {
2795 (0..text.len() - 1).map(|ix| &text[..ix + 1])
2796}
2797
2798fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
2799 ranges.sort_unstable_by(|a, b| {
2800 a.start
2801 .cmp(&b.start, buffer)
2802 .then_with(|| b.end.cmp(&a.end, buffer))
2803 });
2804
2805 let mut ix = 0;
2806 while ix + 1 < ranges.len() {
2807 let b = ranges[ix + 1].clone();
2808 let a = &mut ranges[ix];
2809 if a.end.cmp(&b.start, buffer).is_gt() {
2810 if a.end.cmp(&b.end, buffer).is_lt() {
2811 a.end = b.end;
2812 }
2813 ranges.remove(ix + 1);
2814 } else {
2815 ix += 1;
2816 }
2817 }
2818}
2819
2820#[cfg(test)]
2821mod tests {
2822 use super::*;
2823 use futures::stream::{self};
2824 use serde::Serialize;
2825
2826 #[derive(Serialize)]
2827 pub struct DummyCompletionRequest {
2828 pub name: String,
2829 }
2830
2831 #[gpui::test]
2832 async fn test_strip_invalid_spans_from_codeblock() {
2833 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
2834 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
2835 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
2836 assert_chunks(
2837 "```html\n```js\nLorem ipsum dolor\n```\n```",
2838 "```js\nLorem ipsum dolor\n```",
2839 )
2840 .await;
2841 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
2842 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
2843 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
2844 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
2845
2846 async fn assert_chunks(text: &str, expected_text: &str) {
2847 for chunk_size in 1..=text.len() {
2848 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
2849 .map(|chunk| chunk.unwrap())
2850 .collect::<String>()
2851 .await;
2852 assert_eq!(
2853 actual_text, expected_text,
2854 "failed to strip invalid spans, chunk size: {}",
2855 chunk_size
2856 );
2857 }
2858 }
2859
2860 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
2861 stream::iter(
2862 text.chars()
2863 .collect::<Vec<_>>()
2864 .chunks(size)
2865 .map(|chunk| Ok(chunk.iter().collect::<String>()))
2866 .collect::<Vec<_>>(),
2867 )
2868 }
2869 }
2870}