1use crate::{
2 context::load_context, context_store::ContextStore, inline_prompt_editor::CodegenStatus,
3};
4use agent_settings::AgentSettings;
5use anyhow::{Context as _, Result};
6use client::telemetry::Telemetry;
7use cloud_llm_client::CompletionIntent;
8use collections::HashSet;
9use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};
10use futures::{
11 SinkExt, Stream, StreamExt, TryStreamExt as _, channel::mpsc, future::LocalBoxFuture, join,
12};
13use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Subscription, Task, WeakEntity};
14use language::{Buffer, IndentKind, Point, TransactionId, line_diff};
15use language_model::{
16 LanguageModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
17 LanguageModelTextStream, Role, report_assistant_event,
18};
19use multi_buffer::MultiBufferRow;
20use parking_lot::Mutex;
21use project::Project;
22use prompt_store::{PromptBuilder, PromptStore};
23use rope::Rope;
24use smol::future::FutureExt;
25use std::{
26 cmp,
27 future::Future,
28 iter,
29 ops::{Range, RangeInclusive},
30 pin::Pin,
31 sync::Arc,
32 task::{self, Poll},
33 time::Instant,
34};
35use streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff};
36use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
37
38pub struct BufferCodegen {
39 alternatives: Vec<Entity<CodegenAlternative>>,
40 pub active_alternative: usize,
41 seen_alternatives: HashSet<usize>,
42 subscriptions: Vec<Subscription>,
43 buffer: Entity<MultiBuffer>,
44 range: Range<Anchor>,
45 initial_transaction_id: Option<TransactionId>,
46 context_store: Entity<ContextStore>,
47 project: WeakEntity<Project>,
48 prompt_store: Option<Entity<PromptStore>>,
49 telemetry: Arc<Telemetry>,
50 builder: Arc<PromptBuilder>,
51 pub is_insertion: bool,
52}
53
54impl BufferCodegen {
55 pub fn new(
56 buffer: Entity<MultiBuffer>,
57 range: Range<Anchor>,
58 initial_transaction_id: Option<TransactionId>,
59 context_store: Entity<ContextStore>,
60 project: WeakEntity<Project>,
61 prompt_store: Option<Entity<PromptStore>>,
62 telemetry: Arc<Telemetry>,
63 builder: Arc<PromptBuilder>,
64 cx: &mut Context<Self>,
65 ) -> Self {
66 let codegen = cx.new(|cx| {
67 CodegenAlternative::new(
68 buffer.clone(),
69 range.clone(),
70 false,
71 Some(context_store.clone()),
72 project.clone(),
73 prompt_store.clone(),
74 Some(telemetry.clone()),
75 builder.clone(),
76 cx,
77 )
78 });
79 let mut this = Self {
80 is_insertion: range.to_offset(&buffer.read(cx).snapshot(cx)).is_empty(),
81 alternatives: vec![codegen],
82 active_alternative: 0,
83 seen_alternatives: HashSet::default(),
84 subscriptions: Vec::new(),
85 buffer,
86 range,
87 initial_transaction_id,
88 context_store,
89 project,
90 prompt_store,
91 telemetry,
92 builder,
93 };
94 this.activate(0, cx);
95 this
96 }
97
98 fn subscribe_to_alternative(&mut self, cx: &mut Context<Self>) {
99 let codegen = self.active_alternative().clone();
100 self.subscriptions.clear();
101 self.subscriptions
102 .push(cx.observe(&codegen, |_, _, cx| cx.notify()));
103 self.subscriptions
104 .push(cx.subscribe(&codegen, |_, _, event, cx| cx.emit(*event)));
105 }
106
107 pub fn active_alternative(&self) -> &Entity<CodegenAlternative> {
108 &self.alternatives[self.active_alternative]
109 }
110
111 pub fn status<'a>(&self, cx: &'a App) -> &'a CodegenStatus {
112 &self.active_alternative().read(cx).status
113 }
114
115 pub fn alternative_count(&self, cx: &App) -> usize {
116 LanguageModelRegistry::read_global(cx)
117 .inline_alternative_models()
118 .len()
119 + 1
120 }
121
122 pub fn cycle_prev(&mut self, cx: &mut Context<Self>) {
123 let next_active_ix = if self.active_alternative == 0 {
124 self.alternatives.len() - 1
125 } else {
126 self.active_alternative - 1
127 };
128 self.activate(next_active_ix, cx);
129 }
130
131 pub fn cycle_next(&mut self, cx: &mut Context<Self>) {
132 let next_active_ix = (self.active_alternative + 1) % self.alternatives.len();
133 self.activate(next_active_ix, cx);
134 }
135
136 fn activate(&mut self, index: usize, cx: &mut Context<Self>) {
137 self.active_alternative()
138 .update(cx, |codegen, cx| codegen.set_active(false, cx));
139 self.seen_alternatives.insert(index);
140 self.active_alternative = index;
141 self.active_alternative()
142 .update(cx, |codegen, cx| codegen.set_active(true, cx));
143 self.subscribe_to_alternative(cx);
144 cx.notify();
145 }
146
147 pub fn start(
148 &mut self,
149 primary_model: Arc<dyn LanguageModel>,
150 user_prompt: String,
151 cx: &mut Context<Self>,
152 ) -> Result<()> {
153 let alternative_models = LanguageModelRegistry::read_global(cx)
154 .inline_alternative_models()
155 .to_vec();
156
157 self.active_alternative()
158 .update(cx, |alternative, cx| alternative.undo(cx));
159 self.activate(0, cx);
160 self.alternatives.truncate(1);
161
162 for _ in 0..alternative_models.len() {
163 self.alternatives.push(cx.new(|cx| {
164 CodegenAlternative::new(
165 self.buffer.clone(),
166 self.range.clone(),
167 false,
168 Some(self.context_store.clone()),
169 self.project.clone(),
170 self.prompt_store.clone(),
171 Some(self.telemetry.clone()),
172 self.builder.clone(),
173 cx,
174 )
175 }));
176 }
177
178 for (model, alternative) in iter::once(primary_model)
179 .chain(alternative_models)
180 .zip(&self.alternatives)
181 {
182 alternative.update(cx, |alternative, cx| {
183 alternative.start(user_prompt.clone(), model.clone(), cx)
184 })?;
185 }
186
187 Ok(())
188 }
189
190 pub fn stop(&mut self, cx: &mut Context<Self>) {
191 for codegen in &self.alternatives {
192 codegen.update(cx, |codegen, cx| codegen.stop(cx));
193 }
194 }
195
196 pub fn undo(&mut self, cx: &mut Context<Self>) {
197 self.active_alternative()
198 .update(cx, |codegen, cx| codegen.undo(cx));
199
200 self.buffer.update(cx, |buffer, cx| {
201 if let Some(transaction_id) = self.initial_transaction_id.take() {
202 buffer.undo_transaction(transaction_id, cx);
203 buffer.refresh_preview(cx);
204 }
205 });
206 }
207
208 pub fn buffer(&self, cx: &App) -> Entity<MultiBuffer> {
209 self.active_alternative().read(cx).buffer.clone()
210 }
211
212 pub fn old_buffer(&self, cx: &App) -> Entity<Buffer> {
213 self.active_alternative().read(cx).old_buffer.clone()
214 }
215
216 pub fn snapshot(&self, cx: &App) -> MultiBufferSnapshot {
217 self.active_alternative().read(cx).snapshot.clone()
218 }
219
220 pub fn edit_position(&self, cx: &App) -> Option<Anchor> {
221 self.active_alternative().read(cx).edit_position
222 }
223
224 pub fn diff<'a>(&self, cx: &'a App) -> &'a Diff {
225 &self.active_alternative().read(cx).diff
226 }
227
228 pub fn last_equal_ranges<'a>(&self, cx: &'a App) -> &'a [Range<Anchor>] {
229 self.active_alternative().read(cx).last_equal_ranges()
230 }
231}
232
233impl EventEmitter<CodegenEvent> for BufferCodegen {}
234
235pub struct CodegenAlternative {
236 buffer: Entity<MultiBuffer>,
237 old_buffer: Entity<Buffer>,
238 snapshot: MultiBufferSnapshot,
239 edit_position: Option<Anchor>,
240 range: Range<Anchor>,
241 last_equal_ranges: Vec<Range<Anchor>>,
242 transformation_transaction_id: Option<TransactionId>,
243 status: CodegenStatus,
244 generation: Task<()>,
245 diff: Diff,
246 context_store: Option<Entity<ContextStore>>,
247 project: WeakEntity<Project>,
248 prompt_store: Option<Entity<PromptStore>>,
249 telemetry: Option<Arc<Telemetry>>,
250 _subscription: gpui::Subscription,
251 builder: Arc<PromptBuilder>,
252 active: bool,
253 edits: Vec<(Range<Anchor>, String)>,
254 line_operations: Vec<LineOperation>,
255 elapsed_time: Option<f64>,
256 completion: Option<String>,
257 pub message_id: Option<String>,
258}
259
260impl EventEmitter<CodegenEvent> for CodegenAlternative {}
261
262impl CodegenAlternative {
263 pub fn new(
264 buffer: Entity<MultiBuffer>,
265 range: Range<Anchor>,
266 active: bool,
267 context_store: Option<Entity<ContextStore>>,
268 project: WeakEntity<Project>,
269 prompt_store: Option<Entity<PromptStore>>,
270 telemetry: Option<Arc<Telemetry>>,
271 builder: Arc<PromptBuilder>,
272 cx: &mut Context<Self>,
273 ) -> Self {
274 let snapshot = buffer.read(cx).snapshot(cx);
275
276 let (old_buffer, _, _) = snapshot
277 .range_to_buffer_ranges(range.clone())
278 .pop()
279 .unwrap();
280 let old_buffer = cx.new(|cx| {
281 let text = old_buffer.as_rope().clone();
282 let line_ending = old_buffer.line_ending();
283 let language = old_buffer.language().cloned();
284 let language_registry = buffer
285 .read(cx)
286 .buffer(old_buffer.remote_id())
287 .unwrap()
288 .read(cx)
289 .language_registry();
290
291 let mut buffer = Buffer::local_normalized(text, line_ending, cx);
292 buffer.set_language(language, cx);
293 if let Some(language_registry) = language_registry {
294 buffer.set_language_registry(language_registry)
295 }
296 buffer
297 });
298
299 Self {
300 buffer: buffer.clone(),
301 old_buffer,
302 edit_position: None,
303 message_id: None,
304 snapshot,
305 last_equal_ranges: Default::default(),
306 transformation_transaction_id: None,
307 status: CodegenStatus::Idle,
308 generation: Task::ready(()),
309 diff: Diff::default(),
310 context_store,
311 project,
312 prompt_store,
313 telemetry,
314 _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
315 builder,
316 active,
317 edits: Vec::new(),
318 line_operations: Vec::new(),
319 range,
320 elapsed_time: None,
321 completion: None,
322 }
323 }
324
325 pub fn set_active(&mut self, active: bool, cx: &mut Context<Self>) {
326 if active != self.active {
327 self.active = active;
328
329 if self.active {
330 let edits = self.edits.clone();
331 self.apply_edits(edits, cx);
332 if matches!(self.status, CodegenStatus::Pending) {
333 let line_operations = self.line_operations.clone();
334 self.reapply_line_based_diff(line_operations, cx);
335 } else {
336 self.reapply_batch_diff(cx).detach();
337 }
338 } else if let Some(transaction_id) = self.transformation_transaction_id.take() {
339 self.buffer.update(cx, |buffer, cx| {
340 buffer.undo_transaction(transaction_id, cx);
341 buffer.forget_transaction(transaction_id, cx);
342 });
343 }
344 }
345 }
346
347 fn handle_buffer_event(
348 &mut self,
349 _buffer: Entity<MultiBuffer>,
350 event: &multi_buffer::Event,
351 cx: &mut Context<Self>,
352 ) {
353 if let multi_buffer::Event::TransactionUndone { transaction_id } = event
354 && self.transformation_transaction_id == Some(*transaction_id)
355 {
356 self.transformation_transaction_id = None;
357 self.generation = Task::ready(());
358 cx.emit(CodegenEvent::Undone);
359 }
360 }
361
362 pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
363 &self.last_equal_ranges
364 }
365
366 pub fn start(
367 &mut self,
368 user_prompt: String,
369 model: Arc<dyn LanguageModel>,
370 cx: &mut Context<Self>,
371 ) -> Result<()> {
372 if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {
373 self.buffer.update(cx, |buffer, cx| {
374 buffer.undo_transaction(transformation_transaction_id, cx);
375 });
376 }
377
378 self.edit_position = Some(self.range.start.bias_right(&self.snapshot));
379
380 let api_key = model.api_key(cx);
381 let telemetry_id = model.telemetry_id();
382 let provider_id = model.provider_id();
383 let stream: LocalBoxFuture<Result<LanguageModelTextStream>> =
384 if user_prompt.trim().to_lowercase() == "delete" {
385 async { Ok(LanguageModelTextStream::default()) }.boxed_local()
386 } else {
387 let request = self.build_request(&model, user_prompt, cx)?;
388 cx.spawn(async move |_, cx| {
389 Ok(model.stream_completion_text(request.await, cx).await?)
390 })
391 .boxed_local()
392 };
393 self.handle_stream(telemetry_id, provider_id.to_string(), api_key, stream, cx);
394 Ok(())
395 }
396
397 fn build_request(
398 &self,
399 model: &Arc<dyn LanguageModel>,
400 user_prompt: String,
401 cx: &mut App,
402 ) -> Result<Task<LanguageModelRequest>> {
403 let buffer = self.buffer.read(cx).snapshot(cx);
404 let language = buffer.language_at(self.range.start);
405 let language_name = if let Some(language) = language.as_ref() {
406 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
407 None
408 } else {
409 Some(language.name())
410 }
411 } else {
412 None
413 };
414
415 let language_name = language_name.as_ref();
416 let start = buffer.point_to_buffer_offset(self.range.start);
417 let end = buffer.point_to_buffer_offset(self.range.end);
418 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
419 let (start_buffer, start_buffer_offset) = start;
420 let (end_buffer, end_buffer_offset) = end;
421 if start_buffer.remote_id() == end_buffer.remote_id() {
422 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
423 } else {
424 anyhow::bail!("invalid transformation range");
425 }
426 } else {
427 anyhow::bail!("invalid transformation range");
428 };
429
430 let prompt = self
431 .builder
432 .generate_inline_transformation_prompt(
433 user_prompt,
434 language_name,
435 buffer,
436 range.start.0..range.end.0,
437 )
438 .context("generating content prompt")?;
439
440 let context_task = self.context_store.as_ref().and_then(|context_store| {
441 if let Some(project) = self.project.upgrade() {
442 let context = context_store
443 .read(cx)
444 .context()
445 .cloned()
446 .collect::<Vec<_>>();
447 Some(load_context(context, &project, &self.prompt_store, cx))
448 } else {
449 None
450 }
451 });
452
453 let temperature = AgentSettings::temperature_for_model(model, cx);
454
455 Ok(cx.spawn(async move |_cx| {
456 let mut request_message = LanguageModelRequestMessage {
457 role: Role::User,
458 content: Vec::new(),
459 cache: false,
460 };
461
462 if let Some(context_task) = context_task {
463 context_task
464 .await
465 .add_to_request_message(&mut request_message);
466 }
467
468 request_message.content.push(prompt.into());
469
470 LanguageModelRequest {
471 thread_id: None,
472 prompt_id: None,
473 intent: Some(CompletionIntent::InlineAssist),
474 mode: None,
475 tools: Vec::new(),
476 tool_choice: None,
477 stop: Vec::new(),
478 temperature,
479 messages: vec![request_message],
480 thinking_allowed: false,
481 }
482 }))
483 }
484
485 pub fn handle_stream(
486 &mut self,
487 model_telemetry_id: String,
488 model_provider_id: String,
489 model_api_key: Option<String>,
490 stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,
491 cx: &mut Context<Self>,
492 ) {
493 let start_time = Instant::now();
494
495 // Make a new snapshot and re-resolve anchor in case the document was modified.
496 // This can happen often if the editor loses focus and is saved + reformatted,
497 // as in https://github.com/zed-industries/zed/issues/39088
498 self.snapshot = self.buffer.read(cx).snapshot(cx);
499 self.range = self.snapshot.anchor_after(self.range.start)
500 ..self.snapshot.anchor_after(self.range.end);
501
502 let snapshot = self.snapshot.clone();
503 let selected_text = snapshot
504 .text_for_range(self.range.start..self.range.end)
505 .collect::<Rope>();
506
507 let selection_start = self.range.start.to_point(&snapshot);
508
509 // Start with the indentation of the first line in the selection
510 let mut suggested_line_indent = snapshot
511 .suggested_indents(selection_start.row..=selection_start.row, cx)
512 .into_values()
513 .next()
514 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
515
516 // If the first line in the selection does not have indentation, check the following lines
517 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
518 for row in selection_start.row..=self.range.end.to_point(&snapshot).row {
519 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
520 // Prefer tabs if a line in the selection uses tabs as indentation
521 if line_indent.kind == IndentKind::Tab {
522 suggested_line_indent.kind = IndentKind::Tab;
523 break;
524 }
525 }
526 }
527
528 let http_client = cx.http_client();
529 let telemetry = self.telemetry.clone();
530 let language_name = {
531 let multibuffer = self.buffer.read(cx);
532 let snapshot = multibuffer.snapshot(cx);
533 let ranges = snapshot.range_to_buffer_ranges(self.range.clone());
534 ranges
535 .first()
536 .and_then(|(buffer, _, _)| buffer.language())
537 .map(|language| language.name())
538 };
539
540 self.diff = Diff::default();
541 self.status = CodegenStatus::Pending;
542 let mut edit_start = self.range.start.to_offset(&snapshot);
543 let completion = Arc::new(Mutex::new(String::new()));
544 let completion_clone = completion.clone();
545
546 self.generation = cx.spawn(async move |codegen, cx| {
547 let stream = stream.await;
548 let token_usage = stream
549 .as_ref()
550 .ok()
551 .map(|stream| stream.last_token_usage.clone());
552 let message_id = stream
553 .as_ref()
554 .ok()
555 .and_then(|stream| stream.message_id.clone());
556 let generate = async {
557 let model_telemetry_id = model_telemetry_id.clone();
558 let model_provider_id = model_provider_id.clone();
559 let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
560 let executor = cx.background_executor().clone();
561 let message_id = message_id.clone();
562 let line_based_stream_diff: Task<anyhow::Result<()>> =
563 cx.background_spawn(async move {
564 let mut response_latency = None;
565 let request_start = Instant::now();
566 let diff = async {
567 let chunks = StripInvalidSpans::new(
568 stream?.stream.map_err(|error| error.into()),
569 );
570 futures::pin_mut!(chunks);
571 let mut diff = StreamingDiff::new(selected_text.to_string());
572 let mut line_diff = LineDiff::default();
573
574 let mut new_text = String::new();
575 let mut base_indent = None;
576 let mut line_indent = None;
577 let mut first_line = true;
578
579 while let Some(chunk) = chunks.next().await {
580 if response_latency.is_none() {
581 response_latency = Some(request_start.elapsed());
582 }
583 let chunk = chunk?;
584 completion_clone.lock().push_str(&chunk);
585
586 let mut lines = chunk.split('\n').peekable();
587 while let Some(line) = lines.next() {
588 new_text.push_str(line);
589 if line_indent.is_none()
590 && let Some(non_whitespace_ch_ix) =
591 new_text.find(|ch: char| !ch.is_whitespace())
592 {
593 line_indent = Some(non_whitespace_ch_ix);
594 base_indent = base_indent.or(line_indent);
595
596 let line_indent = line_indent.unwrap();
597 let base_indent = base_indent.unwrap();
598 let indent_delta = line_indent as i32 - base_indent as i32;
599 let mut corrected_indent_len = cmp::max(
600 0,
601 suggested_line_indent.len as i32 + indent_delta,
602 )
603 as usize;
604 if first_line {
605 corrected_indent_len = corrected_indent_len
606 .saturating_sub(selection_start.column as usize);
607 }
608
609 let indent_char = suggested_line_indent.char();
610 let mut indent_buffer = [0; 4];
611 let indent_str =
612 indent_char.encode_utf8(&mut indent_buffer);
613 new_text.replace_range(
614 ..line_indent,
615 &indent_str.repeat(corrected_indent_len),
616 );
617 }
618
619 if line_indent.is_some() {
620 let char_ops = diff.push_new(&new_text);
621 line_diff.push_char_operations(&char_ops, &selected_text);
622 diff_tx
623 .send((char_ops, line_diff.line_operations()))
624 .await?;
625 new_text.clear();
626 }
627
628 if lines.peek().is_some() {
629 let char_ops = diff.push_new("\n");
630 line_diff.push_char_operations(&char_ops, &selected_text);
631 diff_tx
632 .send((char_ops, line_diff.line_operations()))
633 .await?;
634 if line_indent.is_none() {
635 // Don't write out the leading indentation in empty lines on the next line
636 // This is the case where the above if statement didn't clear the buffer
637 new_text.clear();
638 }
639 line_indent = None;
640 first_line = false;
641 }
642 }
643 }
644
645 let mut char_ops = diff.push_new(&new_text);
646 char_ops.extend(diff.finish());
647 line_diff.push_char_operations(&char_ops, &selected_text);
648 line_diff.finish(&selected_text);
649 diff_tx
650 .send((char_ops, line_diff.line_operations()))
651 .await?;
652
653 anyhow::Ok(())
654 };
655
656 let result = diff.await;
657
658 let error_message = result.as_ref().err().map(|error| error.to_string());
659 report_assistant_event(
660 AssistantEventData {
661 conversation_id: None,
662 message_id,
663 kind: AssistantKind::Inline,
664 phase: AssistantPhase::Response,
665 model: model_telemetry_id,
666 model_provider: model_provider_id,
667 response_latency,
668 error_message,
669 language_name: language_name.map(|name| name.to_proto()),
670 },
671 telemetry,
672 http_client,
673 model_api_key,
674 &executor,
675 );
676
677 result?;
678 Ok(())
679 });
680
681 while let Some((char_ops, line_ops)) = diff_rx.next().await {
682 codegen.update(cx, |codegen, cx| {
683 codegen.last_equal_ranges.clear();
684
685 let edits = char_ops
686 .into_iter()
687 .filter_map(|operation| match operation {
688 CharOperation::Insert { text } => {
689 let edit_start = snapshot.anchor_after(edit_start);
690 Some((edit_start..edit_start, text))
691 }
692 CharOperation::Delete { bytes } => {
693 let edit_end = edit_start + bytes;
694 let edit_range = snapshot.anchor_after(edit_start)
695 ..snapshot.anchor_before(edit_end);
696 edit_start = edit_end;
697 Some((edit_range, String::new()))
698 }
699 CharOperation::Keep { bytes } => {
700 let edit_end = edit_start + bytes;
701 let edit_range = snapshot.anchor_after(edit_start)
702 ..snapshot.anchor_before(edit_end);
703 edit_start = edit_end;
704 codegen.last_equal_ranges.push(edit_range);
705 None
706 }
707 })
708 .collect::<Vec<_>>();
709
710 if codegen.active {
711 codegen.apply_edits(edits.iter().cloned(), cx);
712 codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
713 }
714 codegen.edits.extend(edits);
715 codegen.line_operations = line_ops;
716 codegen.edit_position = Some(snapshot.anchor_after(edit_start));
717
718 cx.notify();
719 })?;
720 }
721
722 // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
723 // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
724 // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
725 let batch_diff_task =
726 codegen.update(cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
727 let (line_based_stream_diff, ()) = join!(line_based_stream_diff, batch_diff_task);
728 line_based_stream_diff?;
729
730 anyhow::Ok(())
731 };
732
733 let result = generate.await;
734 let elapsed_time = start_time.elapsed().as_secs_f64();
735
736 codegen
737 .update(cx, |this, cx| {
738 this.message_id = message_id;
739 this.last_equal_ranges.clear();
740 if let Err(error) = result {
741 this.status = CodegenStatus::Error(error);
742 } else {
743 this.status = CodegenStatus::Done;
744 }
745 this.elapsed_time = Some(elapsed_time);
746 this.completion = Some(completion.lock().clone());
747 if let Some(usage) = token_usage {
748 let usage = usage.lock();
749 telemetry::event!(
750 "Inline Assistant Completion",
751 model = model_telemetry_id,
752 model_provider = model_provider_id,
753 input_tokens = usage.input_tokens,
754 output_tokens = usage.output_tokens,
755 )
756 }
757 cx.emit(CodegenEvent::Finished);
758 cx.notify();
759 })
760 .ok();
761 });
762 cx.notify();
763 }
764
765 pub fn stop(&mut self, cx: &mut Context<Self>) {
766 self.last_equal_ranges.clear();
767 if self.diff.is_empty() {
768 self.status = CodegenStatus::Idle;
769 } else {
770 self.status = CodegenStatus::Done;
771 }
772 self.generation = Task::ready(());
773 cx.emit(CodegenEvent::Finished);
774 cx.notify();
775 }
776
777 pub fn undo(&mut self, cx: &mut Context<Self>) {
778 self.buffer.update(cx, |buffer, cx| {
779 if let Some(transaction_id) = self.transformation_transaction_id.take() {
780 buffer.undo_transaction(transaction_id, cx);
781 buffer.refresh_preview(cx);
782 }
783 });
784 }
785
786 fn apply_edits(
787 &mut self,
788 edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
789 cx: &mut Context<CodegenAlternative>,
790 ) {
791 let transaction = self.buffer.update(cx, |buffer, cx| {
792 // Avoid grouping agent edits with user edits.
793 buffer.finalize_last_transaction(cx);
794 buffer.start_transaction(cx);
795 buffer.edit(edits, None, cx);
796 buffer.end_transaction(cx)
797 });
798
799 if let Some(transaction) = transaction {
800 if let Some(first_transaction) = self.transformation_transaction_id {
801 // Group all agent edits into the first transaction.
802 self.buffer.update(cx, |buffer, cx| {
803 buffer.merge_transactions(transaction, first_transaction, cx)
804 });
805 } else {
806 self.transformation_transaction_id = Some(transaction);
807 self.buffer
808 .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
809 }
810 }
811 }
812
813 fn reapply_line_based_diff(
814 &mut self,
815 line_operations: impl IntoIterator<Item = LineOperation>,
816 cx: &mut Context<Self>,
817 ) {
818 let old_snapshot = self.snapshot.clone();
819 let old_range = self.range.to_point(&old_snapshot);
820 let new_snapshot = self.buffer.read(cx).snapshot(cx);
821 let new_range = self.range.to_point(&new_snapshot);
822
823 let mut old_row = old_range.start.row;
824 let mut new_row = new_range.start.row;
825
826 self.diff.deleted_row_ranges.clear();
827 self.diff.inserted_row_ranges.clear();
828 for operation in line_operations {
829 match operation {
830 LineOperation::Keep { lines } => {
831 old_row += lines;
832 new_row += lines;
833 }
834 LineOperation::Delete { lines } => {
835 let old_end_row = old_row + lines - 1;
836 let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
837
838 if let Some((_, last_deleted_row_range)) =
839 self.diff.deleted_row_ranges.last_mut()
840 {
841 if *last_deleted_row_range.end() + 1 == old_row {
842 *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
843 } else {
844 self.diff
845 .deleted_row_ranges
846 .push((new_row, old_row..=old_end_row));
847 }
848 } else {
849 self.diff
850 .deleted_row_ranges
851 .push((new_row, old_row..=old_end_row));
852 }
853
854 old_row += lines;
855 }
856 LineOperation::Insert { lines } => {
857 let new_end_row = new_row + lines - 1;
858 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
859 let end = new_snapshot.anchor_before(Point::new(
860 new_end_row,
861 new_snapshot.line_len(MultiBufferRow(new_end_row)),
862 ));
863 self.diff.inserted_row_ranges.push(start..end);
864 new_row += lines;
865 }
866 }
867
868 cx.notify();
869 }
870 }
871
872 fn reapply_batch_diff(&mut self, cx: &mut Context<Self>) -> Task<()> {
873 let old_snapshot = self.snapshot.clone();
874 let old_range = self.range.to_point(&old_snapshot);
875 let new_snapshot = self.buffer.read(cx).snapshot(cx);
876 let new_range = self.range.to_point(&new_snapshot);
877
878 cx.spawn(async move |codegen, cx| {
879 let (deleted_row_ranges, inserted_row_ranges) = cx
880 .background_spawn(async move {
881 let old_text = old_snapshot
882 .text_for_range(
883 Point::new(old_range.start.row, 0)
884 ..Point::new(
885 old_range.end.row,
886 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
887 ),
888 )
889 .collect::<String>();
890 let new_text = new_snapshot
891 .text_for_range(
892 Point::new(new_range.start.row, 0)
893 ..Point::new(
894 new_range.end.row,
895 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
896 ),
897 )
898 .collect::<String>();
899
900 let old_start_row = old_range.start.row;
901 let new_start_row = new_range.start.row;
902 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
903 let mut inserted_row_ranges = Vec::new();
904 for (old_rows, new_rows) in line_diff(&old_text, &new_text) {
905 let old_rows = old_start_row + old_rows.start..old_start_row + old_rows.end;
906 let new_rows = new_start_row + new_rows.start..new_start_row + new_rows.end;
907 if !old_rows.is_empty() {
908 deleted_row_ranges.push((
909 new_snapshot.anchor_before(Point::new(new_rows.start, 0)),
910 old_rows.start..=old_rows.end - 1,
911 ));
912 }
913 if !new_rows.is_empty() {
914 let start = new_snapshot.anchor_before(Point::new(new_rows.start, 0));
915 let new_end_row = new_rows.end - 1;
916 let end = new_snapshot.anchor_before(Point::new(
917 new_end_row,
918 new_snapshot.line_len(MultiBufferRow(new_end_row)),
919 ));
920 inserted_row_ranges.push(start..end);
921 }
922 }
923 (deleted_row_ranges, inserted_row_ranges)
924 })
925 .await;
926
927 codegen
928 .update(cx, |codegen, cx| {
929 codegen.diff.deleted_row_ranges = deleted_row_ranges;
930 codegen.diff.inserted_row_ranges = inserted_row_ranges;
931 cx.notify();
932 })
933 .ok();
934 })
935 }
936}
937
938#[derive(Copy, Clone, Debug)]
939pub enum CodegenEvent {
940 Finished,
941 Undone,
942}
943
944struct StripInvalidSpans<T> {
945 stream: T,
946 stream_done: bool,
947 buffer: String,
948 first_line: bool,
949 line_end: bool,
950 starts_with_code_block: bool,
951}
952
953impl<T> StripInvalidSpans<T>
954where
955 T: Stream<Item = Result<String>>,
956{
957 fn new(stream: T) -> Self {
958 Self {
959 stream,
960 stream_done: false,
961 buffer: String::new(),
962 first_line: true,
963 line_end: false,
964 starts_with_code_block: false,
965 }
966 }
967}
968
969impl<T> Stream for StripInvalidSpans<T>
970where
971 T: Stream<Item = Result<String>>,
972{
973 type Item = Result<String>;
974
975 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
976 const CODE_BLOCK_DELIMITER: &str = "```";
977 const CURSOR_SPAN: &str = "<|CURSOR|>";
978
979 let this = unsafe { self.get_unchecked_mut() };
980 loop {
981 if !this.stream_done {
982 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
983 match stream.as_mut().poll_next(cx) {
984 Poll::Ready(Some(Ok(chunk))) => {
985 this.buffer.push_str(&chunk);
986 }
987 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
988 Poll::Ready(None) => {
989 this.stream_done = true;
990 }
991 Poll::Pending => return Poll::Pending,
992 }
993 }
994
995 let mut chunk = String::new();
996 let mut consumed = 0;
997 if !this.buffer.is_empty() {
998 let mut lines = this.buffer.split('\n').enumerate().peekable();
999 while let Some((line_ix, line)) = lines.next() {
1000 if line_ix > 0 {
1001 this.first_line = false;
1002 }
1003
1004 if this.first_line {
1005 let trimmed_line = line.trim();
1006 if lines.peek().is_some() {
1007 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
1008 consumed += line.len() + 1;
1009 this.starts_with_code_block = true;
1010 continue;
1011 }
1012 } else if trimmed_line.is_empty()
1013 || prefixes(CODE_BLOCK_DELIMITER)
1014 .any(|prefix| trimmed_line.starts_with(prefix))
1015 {
1016 break;
1017 }
1018 }
1019
1020 let line_without_cursor = line.replace(CURSOR_SPAN, "");
1021 if lines.peek().is_some() {
1022 if this.line_end {
1023 chunk.push('\n');
1024 }
1025
1026 chunk.push_str(&line_without_cursor);
1027 this.line_end = true;
1028 consumed += line.len() + 1;
1029 } else if this.stream_done {
1030 if !this.starts_with_code_block
1031 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
1032 {
1033 if this.line_end {
1034 chunk.push('\n');
1035 }
1036
1037 chunk.push_str(line);
1038 }
1039
1040 consumed += line.len();
1041 } else {
1042 let trimmed_line = line.trim();
1043 if trimmed_line.is_empty()
1044 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
1045 || prefixes(CODE_BLOCK_DELIMITER)
1046 .any(|prefix| trimmed_line.ends_with(prefix))
1047 {
1048 break;
1049 } else {
1050 if this.line_end {
1051 chunk.push('\n');
1052 this.line_end = false;
1053 }
1054
1055 chunk.push_str(&line_without_cursor);
1056 consumed += line.len();
1057 }
1058 }
1059 }
1060 }
1061
1062 this.buffer = this.buffer.split_off(consumed);
1063 if !chunk.is_empty() {
1064 return Poll::Ready(Some(Ok(chunk)));
1065 } else if this.stream_done {
1066 return Poll::Ready(None);
1067 }
1068 }
1069 }
1070}
1071
1072fn prefixes(text: &str) -> impl Iterator<Item = &str> {
1073 (0..text.len() - 1).map(|ix| &text[..ix + 1])
1074}
1075
1076#[derive(Default)]
1077pub struct Diff {
1078 pub deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1079 pub inserted_row_ranges: Vec<Range<Anchor>>,
1080}
1081
1082impl Diff {
1083 fn is_empty(&self) -> bool {
1084 self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
1085 }
1086}
1087
1088#[cfg(test)]
1089mod tests {
1090 use super::*;
1091 use fs::FakeFs;
1092 use futures::{
1093 Stream,
1094 stream::{self},
1095 };
1096 use gpui::TestAppContext;
1097 use indoc::indoc;
1098 use language::{Buffer, Language, LanguageConfig, LanguageMatcher, Point, tree_sitter_rust};
1099 use language_model::{LanguageModelRegistry, TokenUsage};
1100 use rand::prelude::*;
1101 use settings::SettingsStore;
1102 use std::{future, sync::Arc};
1103
1104 #[gpui::test(iterations = 10)]
1105 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
1106 init_test(cx);
1107
1108 let text = indoc! {"
1109 fn main() {
1110 let x = 0;
1111 for _ in 0..10 {
1112 x += 1;
1113 }
1114 }
1115 "};
1116 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1117 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1118 let range = buffer.read_with(cx, |buffer, cx| {
1119 let snapshot = buffer.snapshot(cx);
1120 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
1121 });
1122 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1123 let fs = FakeFs::new(cx.executor());
1124 let project = Project::test(fs, vec![], cx).await;
1125 let codegen = cx.new(|cx| {
1126 CodegenAlternative::new(
1127 buffer.clone(),
1128 range.clone(),
1129 true,
1130 None,
1131 project.downgrade(),
1132 None,
1133 None,
1134 prompt_builder,
1135 cx,
1136 )
1137 });
1138
1139 let chunks_tx = simulate_response_stream(&codegen, cx);
1140
1141 let mut new_text = concat!(
1142 " let mut x = 0;\n",
1143 " while x < 10 {\n",
1144 " x += 1;\n",
1145 " }",
1146 );
1147 while !new_text.is_empty() {
1148 let max_len = cmp::min(new_text.len(), 10);
1149 let len = rng.random_range(1..=max_len);
1150 let (chunk, suffix) = new_text.split_at(len);
1151 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1152 new_text = suffix;
1153 cx.background_executor.run_until_parked();
1154 }
1155 drop(chunks_tx);
1156 cx.background_executor.run_until_parked();
1157
1158 assert_eq!(
1159 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1160 indoc! {"
1161 fn main() {
1162 let mut x = 0;
1163 while x < 10 {
1164 x += 1;
1165 }
1166 }
1167 "}
1168 );
1169 }
1170
1171 #[gpui::test(iterations = 10)]
1172 async fn test_autoindent_when_generating_past_indentation(
1173 cx: &mut TestAppContext,
1174 mut rng: StdRng,
1175 ) {
1176 init_test(cx);
1177
1178 let text = indoc! {"
1179 fn main() {
1180 le
1181 }
1182 "};
1183 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1184 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1185 let range = buffer.read_with(cx, |buffer, cx| {
1186 let snapshot = buffer.snapshot(cx);
1187 snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
1188 });
1189 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1190 let fs = FakeFs::new(cx.executor());
1191 let project = Project::test(fs, vec![], cx).await;
1192 let codegen = cx.new(|cx| {
1193 CodegenAlternative::new(
1194 buffer.clone(),
1195 range.clone(),
1196 true,
1197 None,
1198 project.downgrade(),
1199 None,
1200 None,
1201 prompt_builder,
1202 cx,
1203 )
1204 });
1205
1206 let chunks_tx = simulate_response_stream(&codegen, cx);
1207
1208 cx.background_executor.run_until_parked();
1209
1210 let mut new_text = concat!(
1211 "t mut x = 0;\n",
1212 "while x < 10 {\n",
1213 " x += 1;\n",
1214 "}", //
1215 );
1216 while !new_text.is_empty() {
1217 let max_len = cmp::min(new_text.len(), 10);
1218 let len = rng.random_range(1..=max_len);
1219 let (chunk, suffix) = new_text.split_at(len);
1220 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1221 new_text = suffix;
1222 cx.background_executor.run_until_parked();
1223 }
1224 drop(chunks_tx);
1225 cx.background_executor.run_until_parked();
1226
1227 assert_eq!(
1228 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1229 indoc! {"
1230 fn main() {
1231 let mut x = 0;
1232 while x < 10 {
1233 x += 1;
1234 }
1235 }
1236 "}
1237 );
1238 }
1239
1240 #[gpui::test(iterations = 10)]
1241 async fn test_autoindent_when_generating_before_indentation(
1242 cx: &mut TestAppContext,
1243 mut rng: StdRng,
1244 ) {
1245 init_test(cx);
1246
1247 let text = concat!(
1248 "fn main() {\n",
1249 " \n",
1250 "}\n" //
1251 );
1252 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1253 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1254 let range = buffer.read_with(cx, |buffer, cx| {
1255 let snapshot = buffer.snapshot(cx);
1256 snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
1257 });
1258 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1259 let fs = FakeFs::new(cx.executor());
1260 let project = Project::test(fs, vec![], cx).await;
1261 let codegen = cx.new(|cx| {
1262 CodegenAlternative::new(
1263 buffer.clone(),
1264 range.clone(),
1265 true,
1266 None,
1267 project.downgrade(),
1268 None,
1269 None,
1270 prompt_builder,
1271 cx,
1272 )
1273 });
1274
1275 let chunks_tx = simulate_response_stream(&codegen, cx);
1276
1277 cx.background_executor.run_until_parked();
1278
1279 let mut new_text = concat!(
1280 "let mut x = 0;\n",
1281 "while x < 10 {\n",
1282 " x += 1;\n",
1283 "}", //
1284 );
1285 while !new_text.is_empty() {
1286 let max_len = cmp::min(new_text.len(), 10);
1287 let len = rng.random_range(1..=max_len);
1288 let (chunk, suffix) = new_text.split_at(len);
1289 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1290 new_text = suffix;
1291 cx.background_executor.run_until_parked();
1292 }
1293 drop(chunks_tx);
1294 cx.background_executor.run_until_parked();
1295
1296 assert_eq!(
1297 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1298 indoc! {"
1299 fn main() {
1300 let mut x = 0;
1301 while x < 10 {
1302 x += 1;
1303 }
1304 }
1305 "}
1306 );
1307 }
1308
1309 #[gpui::test(iterations = 10)]
1310 async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
1311 init_test(cx);
1312
1313 let text = indoc! {"
1314 func main() {
1315 \tx := 0
1316 \tfor i := 0; i < 10; i++ {
1317 \t\tx++
1318 \t}
1319 }
1320 "};
1321 let buffer = cx.new(|cx| Buffer::local(text, cx));
1322 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1323 let range = buffer.read_with(cx, |buffer, cx| {
1324 let snapshot = buffer.snapshot(cx);
1325 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
1326 });
1327 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1328 let fs = FakeFs::new(cx.executor());
1329 let project = Project::test(fs, vec![], cx).await;
1330 let codegen = cx.new(|cx| {
1331 CodegenAlternative::new(
1332 buffer.clone(),
1333 range.clone(),
1334 true,
1335 None,
1336 project.downgrade(),
1337 None,
1338 None,
1339 prompt_builder,
1340 cx,
1341 )
1342 });
1343
1344 let chunks_tx = simulate_response_stream(&codegen, cx);
1345 let new_text = concat!(
1346 "func main() {\n",
1347 "\tx := 0\n",
1348 "\tfor x < 10 {\n",
1349 "\t\tx++\n",
1350 "\t}", //
1351 );
1352 chunks_tx.unbounded_send(new_text.to_string()).unwrap();
1353 drop(chunks_tx);
1354 cx.background_executor.run_until_parked();
1355
1356 assert_eq!(
1357 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1358 indoc! {"
1359 func main() {
1360 \tx := 0
1361 \tfor x < 10 {
1362 \t\tx++
1363 \t}
1364 }
1365 "}
1366 );
1367 }
1368
1369 #[gpui::test]
1370 async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
1371 init_test(cx);
1372
1373 let text = indoc! {"
1374 fn main() {
1375 let x = 0;
1376 }
1377 "};
1378 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1379 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1380 let range = buffer.read_with(cx, |buffer, cx| {
1381 let snapshot = buffer.snapshot(cx);
1382 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
1383 });
1384 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1385 let fs = FakeFs::new(cx.executor());
1386 let project = Project::test(fs, vec![], cx).await;
1387 let codegen = cx.new(|cx| {
1388 CodegenAlternative::new(
1389 buffer.clone(),
1390 range.clone(),
1391 false,
1392 None,
1393 project.downgrade(),
1394 None,
1395 None,
1396 prompt_builder,
1397 cx,
1398 )
1399 });
1400
1401 let chunks_tx = simulate_response_stream(&codegen, cx);
1402 chunks_tx
1403 .unbounded_send("let mut x = 0;\nx += 1;".to_string())
1404 .unwrap();
1405 drop(chunks_tx);
1406 cx.run_until_parked();
1407
1408 // The codegen is inactive, so the buffer doesn't get modified.
1409 assert_eq!(
1410 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1411 text
1412 );
1413
1414 // Activating the codegen applies the changes.
1415 codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
1416 assert_eq!(
1417 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1418 indoc! {"
1419 fn main() {
1420 let mut x = 0;
1421 x += 1;
1422 }
1423 "}
1424 );
1425
1426 // Deactivating the codegen undoes the changes.
1427 codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
1428 cx.run_until_parked();
1429 assert_eq!(
1430 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1431 text
1432 );
1433 }
1434
1435 #[gpui::test]
1436 async fn test_strip_invalid_spans_from_codeblock() {
1437 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
1438 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
1439 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
1440 assert_chunks(
1441 "```html\n```js\nLorem ipsum dolor\n```\n```",
1442 "```js\nLorem ipsum dolor\n```",
1443 )
1444 .await;
1445 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
1446 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
1447 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
1448 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
1449
1450 async fn assert_chunks(text: &str, expected_text: &str) {
1451 for chunk_size in 1..=text.len() {
1452 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
1453 .map(|chunk| chunk.unwrap())
1454 .collect::<String>()
1455 .await;
1456 assert_eq!(
1457 actual_text, expected_text,
1458 "failed to strip invalid spans, chunk size: {}",
1459 chunk_size
1460 );
1461 }
1462 }
1463
1464 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
1465 stream::iter(
1466 text.chars()
1467 .collect::<Vec<_>>()
1468 .chunks(size)
1469 .map(|chunk| Ok(chunk.iter().collect::<String>()))
1470 .collect::<Vec<_>>(),
1471 )
1472 }
1473 }
1474
1475 fn init_test(cx: &mut TestAppContext) {
1476 cx.update(LanguageModelRegistry::test);
1477 cx.set_global(cx.update(SettingsStore::test));
1478 }
1479
1480 fn simulate_response_stream(
1481 codegen: &Entity<CodegenAlternative>,
1482 cx: &mut TestAppContext,
1483 ) -> mpsc::UnboundedSender<String> {
1484 let (chunks_tx, chunks_rx) = mpsc::unbounded();
1485 codegen.update(cx, |codegen, cx| {
1486 codegen.handle_stream(
1487 String::new(),
1488 String::new(),
1489 None,
1490 future::ready(Ok(LanguageModelTextStream {
1491 message_id: None,
1492 stream: chunks_rx.map(Ok).boxed(),
1493 last_token_usage: Arc::new(Mutex::new(TokenUsage::default())),
1494 })),
1495 cx,
1496 );
1497 });
1498 chunks_tx
1499 }
1500
1501 fn rust_lang() -> Language {
1502 Language::new(
1503 LanguageConfig {
1504 name: "Rust".into(),
1505 matcher: LanguageMatcher {
1506 path_suffixes: vec!["rs".to_string()],
1507 ..Default::default()
1508 },
1509 ..Default::default()
1510 },
1511 Some(tree_sitter_rust::LANGUAGE.into()),
1512 )
1513 .with_indents_query(
1514 r#"
1515 (call_expression) @indent
1516 (field_expression) @indent
1517 (_ "(" ")" @end) @indent
1518 (_ "{" "}" @end) @indent
1519 "#,
1520 )
1521 .unwrap()
1522 }
1523}