1use crate::context::ContextLoadResult;
2use crate::inline_prompt_editor::CodegenStatus;
3use crate::{context::load_context, context_store::ContextStore};
4use anyhow::Result;
5use assistant_settings::AssistantSettings;
6use client::telemetry::Telemetry;
7use collections::HashSet;
8use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};
9use futures::{
10 SinkExt, Stream, StreamExt, TryStreamExt as _, channel::mpsc, future::LocalBoxFuture, join,
11};
12use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Subscription, Task, WeakEntity};
13use language::{Buffer, IndentKind, Point, TransactionId, line_diff};
14use language_model::{
15 LanguageModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
16 LanguageModelTextStream, Role, report_assistant_event,
17};
18use multi_buffer::MultiBufferRow;
19use parking_lot::Mutex;
20use project::Project;
21use prompt_store::PromptBuilder;
22use prompt_store::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 if self.transformation_transaction_id == Some(*transaction_id) {
355 self.transformation_transaction_id = None;
356 self.generation = Task::ready(());
357 cx.emit(CodegenEvent::Undone);
358 }
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| model.stream_completion_text(request.await, &cx).await)
389 .boxed_local()
390 };
391 self.handle_stream(telemetry_id, provider_id.to_string(), api_key, stream, cx);
392 Ok(())
393 }
394
395 fn build_request(
396 &self,
397 model: &Arc<dyn LanguageModel>,
398 user_prompt: String,
399 cx: &mut App,
400 ) -> Result<Task<LanguageModelRequest>> {
401 let buffer = self.buffer.read(cx).snapshot(cx);
402 let language = buffer.language_at(self.range.start);
403 let language_name = if let Some(language) = language.as_ref() {
404 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
405 None
406 } else {
407 Some(language.name())
408 }
409 } else {
410 None
411 };
412
413 let language_name = language_name.as_ref();
414 let start = buffer.point_to_buffer_offset(self.range.start);
415 let end = buffer.point_to_buffer_offset(self.range.end);
416 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
417 let (start_buffer, start_buffer_offset) = start;
418 let (end_buffer, end_buffer_offset) = end;
419 if start_buffer.remote_id() == end_buffer.remote_id() {
420 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
421 } else {
422 return Err(anyhow::anyhow!("invalid transformation range"));
423 }
424 } else {
425 return Err(anyhow::anyhow!("invalid transformation range"));
426 };
427
428 let prompt = self
429 .builder
430 .generate_inline_transformation_prompt(user_prompt, language_name, buffer, range)
431 .map_err(|e| anyhow::anyhow!("Failed to generate content prompt: {}", e))?;
432
433 let context_task = self.context_store.as_ref().map(|context_store| {
434 if let Some(project) = self.project.upgrade() {
435 let context = context_store
436 .read(cx)
437 .context()
438 .cloned()
439 .collect::<Vec<_>>();
440 load_context(context, &project, &self.prompt_store, cx)
441 } else {
442 Task::ready(ContextLoadResult::default())
443 }
444 });
445
446 let temperature = AssistantSettings::temperature_for_model(&model, cx);
447
448 Ok(cx.spawn(async move |_cx| {
449 let mut request_message = LanguageModelRequestMessage {
450 role: Role::User,
451 content: Vec::new(),
452 cache: false,
453 };
454
455 if let Some(context_task) = context_task {
456 context_task
457 .await
458 .loaded_context
459 .add_to_request_message(&mut request_message);
460 }
461
462 request_message.content.push(prompt.into());
463
464 LanguageModelRequest {
465 thread_id: None,
466 prompt_id: None,
467 mode: None,
468 tools: Vec::new(),
469 stop: Vec::new(),
470 temperature,
471 messages: vec![request_message],
472 }
473 }))
474 }
475
476 pub fn handle_stream(
477 &mut self,
478 model_telemetry_id: String,
479 model_provider_id: String,
480 model_api_key: Option<String>,
481 stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,
482 cx: &mut Context<Self>,
483 ) {
484 let start_time = Instant::now();
485 let snapshot = self.snapshot.clone();
486 let selected_text = snapshot
487 .text_for_range(self.range.start..self.range.end)
488 .collect::<Rope>();
489
490 let selection_start = self.range.start.to_point(&snapshot);
491
492 // Start with the indentation of the first line in the selection
493 let mut suggested_line_indent = snapshot
494 .suggested_indents(selection_start.row..=selection_start.row, cx)
495 .into_values()
496 .next()
497 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
498
499 // If the first line in the selection does not have indentation, check the following lines
500 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
501 for row in selection_start.row..=self.range.end.to_point(&snapshot).row {
502 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
503 // Prefer tabs if a line in the selection uses tabs as indentation
504 if line_indent.kind == IndentKind::Tab {
505 suggested_line_indent.kind = IndentKind::Tab;
506 break;
507 }
508 }
509 }
510
511 let http_client = cx.http_client();
512 let telemetry = self.telemetry.clone();
513 let language_name = {
514 let multibuffer = self.buffer.read(cx);
515 let snapshot = multibuffer.snapshot(cx);
516 let ranges = snapshot.range_to_buffer_ranges(self.range.clone());
517 ranges
518 .first()
519 .and_then(|(buffer, _, _)| buffer.language())
520 .map(|language| language.name())
521 };
522
523 self.diff = Diff::default();
524 self.status = CodegenStatus::Pending;
525 let mut edit_start = self.range.start.to_offset(&snapshot);
526 let completion = Arc::new(Mutex::new(String::new()));
527 let completion_clone = completion.clone();
528
529 self.generation = cx.spawn(async move |codegen, cx| {
530 let stream = stream.await;
531 let token_usage = stream
532 .as_ref()
533 .ok()
534 .map(|stream| stream.last_token_usage.clone());
535 let message_id = stream
536 .as_ref()
537 .ok()
538 .and_then(|stream| stream.message_id.clone());
539 let generate = async {
540 let model_telemetry_id = model_telemetry_id.clone();
541 let model_provider_id = model_provider_id.clone();
542 let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
543 let executor = cx.background_executor().clone();
544 let message_id = message_id.clone();
545 let line_based_stream_diff: Task<anyhow::Result<()>> =
546 cx.background_spawn(async move {
547 let mut response_latency = None;
548 let request_start = Instant::now();
549 let diff = async {
550 let chunks = StripInvalidSpans::new(
551 stream?.stream.map_err(|error| error.into()),
552 );
553 futures::pin_mut!(chunks);
554 let mut diff = StreamingDiff::new(selected_text.to_string());
555 let mut line_diff = LineDiff::default();
556
557 let mut new_text = String::new();
558 let mut base_indent = None;
559 let mut line_indent = None;
560 let mut first_line = true;
561
562 while let Some(chunk) = chunks.next().await {
563 if response_latency.is_none() {
564 response_latency = Some(request_start.elapsed());
565 }
566 let chunk = chunk?;
567 completion_clone.lock().push_str(&chunk);
568
569 let mut lines = chunk.split('\n').peekable();
570 while let Some(line) = lines.next() {
571 new_text.push_str(line);
572 if line_indent.is_none() {
573 if let Some(non_whitespace_ch_ix) =
574 new_text.find(|ch: char| !ch.is_whitespace())
575 {
576 line_indent = Some(non_whitespace_ch_ix);
577 base_indent = base_indent.or(line_indent);
578
579 let line_indent = line_indent.unwrap();
580 let base_indent = base_indent.unwrap();
581 let indent_delta =
582 line_indent as i32 - base_indent as i32;
583 let mut corrected_indent_len = cmp::max(
584 0,
585 suggested_line_indent.len as i32 + indent_delta,
586 )
587 as usize;
588 if first_line {
589 corrected_indent_len = corrected_indent_len
590 .saturating_sub(
591 selection_start.column as usize,
592 );
593 }
594
595 let indent_char = suggested_line_indent.char();
596 let mut indent_buffer = [0; 4];
597 let indent_str =
598 indent_char.encode_utf8(&mut indent_buffer);
599 new_text.replace_range(
600 ..line_indent,
601 &indent_str.repeat(corrected_indent_len),
602 );
603 }
604 }
605
606 if line_indent.is_some() {
607 let char_ops = diff.push_new(&new_text);
608 line_diff.push_char_operations(&char_ops, &selected_text);
609 diff_tx
610 .send((char_ops, line_diff.line_operations()))
611 .await?;
612 new_text.clear();
613 }
614
615 if lines.peek().is_some() {
616 let char_ops = diff.push_new("\n");
617 line_diff.push_char_operations(&char_ops, &selected_text);
618 diff_tx
619 .send((char_ops, line_diff.line_operations()))
620 .await?;
621 if line_indent.is_none() {
622 // Don't write out the leading indentation in empty lines on the next line
623 // This is the case where the above if statement didn't clear the buffer
624 new_text.clear();
625 }
626 line_indent = None;
627 first_line = false;
628 }
629 }
630 }
631
632 let mut char_ops = diff.push_new(&new_text);
633 char_ops.extend(diff.finish());
634 line_diff.push_char_operations(&char_ops, &selected_text);
635 line_diff.finish(&selected_text);
636 diff_tx
637 .send((char_ops, line_diff.line_operations()))
638 .await?;
639
640 anyhow::Ok(())
641 };
642
643 let result = diff.await;
644
645 let error_message = result.as_ref().err().map(|error| error.to_string());
646 report_assistant_event(
647 AssistantEventData {
648 conversation_id: None,
649 message_id,
650 kind: AssistantKind::Inline,
651 phase: AssistantPhase::Response,
652 model: model_telemetry_id,
653 model_provider: model_provider_id,
654 response_latency,
655 error_message,
656 language_name: language_name.map(|name| name.to_proto()),
657 },
658 telemetry,
659 http_client,
660 model_api_key,
661 &executor,
662 );
663
664 result?;
665 Ok(())
666 });
667
668 while let Some((char_ops, line_ops)) = diff_rx.next().await {
669 codegen.update(cx, |codegen, cx| {
670 codegen.last_equal_ranges.clear();
671
672 let edits = char_ops
673 .into_iter()
674 .filter_map(|operation| match operation {
675 CharOperation::Insert { text } => {
676 let edit_start = snapshot.anchor_after(edit_start);
677 Some((edit_start..edit_start, text))
678 }
679 CharOperation::Delete { bytes } => {
680 let edit_end = edit_start + bytes;
681 let edit_range = snapshot.anchor_after(edit_start)
682 ..snapshot.anchor_before(edit_end);
683 edit_start = edit_end;
684 Some((edit_range, String::new()))
685 }
686 CharOperation::Keep { bytes } => {
687 let edit_end = edit_start + bytes;
688 let edit_range = snapshot.anchor_after(edit_start)
689 ..snapshot.anchor_before(edit_end);
690 edit_start = edit_end;
691 codegen.last_equal_ranges.push(edit_range);
692 None
693 }
694 })
695 .collect::<Vec<_>>();
696
697 if codegen.active {
698 codegen.apply_edits(edits.iter().cloned(), cx);
699 codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
700 }
701 codegen.edits.extend(edits);
702 codegen.line_operations = line_ops;
703 codegen.edit_position = Some(snapshot.anchor_after(edit_start));
704
705 cx.notify();
706 })?;
707 }
708
709 // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
710 // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
711 // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
712 let batch_diff_task =
713 codegen.update(cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
714 let (line_based_stream_diff, ()) = join!(line_based_stream_diff, batch_diff_task);
715 line_based_stream_diff?;
716
717 anyhow::Ok(())
718 };
719
720 let result = generate.await;
721 let elapsed_time = start_time.elapsed().as_secs_f64();
722
723 codegen
724 .update(cx, |this, cx| {
725 this.message_id = message_id;
726 this.last_equal_ranges.clear();
727 if let Err(error) = result {
728 this.status = CodegenStatus::Error(error);
729 } else {
730 this.status = CodegenStatus::Done;
731 }
732 this.elapsed_time = Some(elapsed_time);
733 this.completion = Some(completion.lock().clone());
734 if let Some(usage) = token_usage {
735 let usage = usage.lock();
736 telemetry::event!(
737 "Inline Assistant Completion",
738 model = model_telemetry_id,
739 model_provider = model_provider_id,
740 input_tokens = usage.input_tokens,
741 output_tokens = usage.output_tokens,
742 )
743 }
744 cx.emit(CodegenEvent::Finished);
745 cx.notify();
746 })
747 .ok();
748 });
749 cx.notify();
750 }
751
752 pub fn stop(&mut self, cx: &mut Context<Self>) {
753 self.last_equal_ranges.clear();
754 if self.diff.is_empty() {
755 self.status = CodegenStatus::Idle;
756 } else {
757 self.status = CodegenStatus::Done;
758 }
759 self.generation = Task::ready(());
760 cx.emit(CodegenEvent::Finished);
761 cx.notify();
762 }
763
764 pub fn undo(&mut self, cx: &mut Context<Self>) {
765 self.buffer.update(cx, |buffer, cx| {
766 if let Some(transaction_id) = self.transformation_transaction_id.take() {
767 buffer.undo_transaction(transaction_id, cx);
768 buffer.refresh_preview(cx);
769 }
770 });
771 }
772
773 fn apply_edits(
774 &mut self,
775 edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
776 cx: &mut Context<CodegenAlternative>,
777 ) {
778 let transaction = self.buffer.update(cx, |buffer, cx| {
779 // Avoid grouping agent edits with user edits.
780 buffer.finalize_last_transaction(cx);
781 buffer.start_transaction(cx);
782 buffer.edit(edits, None, cx);
783 buffer.end_transaction(cx)
784 });
785
786 if let Some(transaction) = transaction {
787 if let Some(first_transaction) = self.transformation_transaction_id {
788 // Group all agent edits into the first transaction.
789 self.buffer.update(cx, |buffer, cx| {
790 buffer.merge_transactions(transaction, first_transaction, cx)
791 });
792 } else {
793 self.transformation_transaction_id = Some(transaction);
794 self.buffer
795 .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
796 }
797 }
798 }
799
800 fn reapply_line_based_diff(
801 &mut self,
802 line_operations: impl IntoIterator<Item = LineOperation>,
803 cx: &mut Context<Self>,
804 ) {
805 let old_snapshot = self.snapshot.clone();
806 let old_range = self.range.to_point(&old_snapshot);
807 let new_snapshot = self.buffer.read(cx).snapshot(cx);
808 let new_range = self.range.to_point(&new_snapshot);
809
810 let mut old_row = old_range.start.row;
811 let mut new_row = new_range.start.row;
812
813 self.diff.deleted_row_ranges.clear();
814 self.diff.inserted_row_ranges.clear();
815 for operation in line_operations {
816 match operation {
817 LineOperation::Keep { lines } => {
818 old_row += lines;
819 new_row += lines;
820 }
821 LineOperation::Delete { lines } => {
822 let old_end_row = old_row + lines - 1;
823 let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
824
825 if let Some((_, last_deleted_row_range)) =
826 self.diff.deleted_row_ranges.last_mut()
827 {
828 if *last_deleted_row_range.end() + 1 == old_row {
829 *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
830 } else {
831 self.diff
832 .deleted_row_ranges
833 .push((new_row, old_row..=old_end_row));
834 }
835 } else {
836 self.diff
837 .deleted_row_ranges
838 .push((new_row, old_row..=old_end_row));
839 }
840
841 old_row += lines;
842 }
843 LineOperation::Insert { lines } => {
844 let new_end_row = new_row + lines - 1;
845 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
846 let end = new_snapshot.anchor_before(Point::new(
847 new_end_row,
848 new_snapshot.line_len(MultiBufferRow(new_end_row)),
849 ));
850 self.diff.inserted_row_ranges.push(start..end);
851 new_row += lines;
852 }
853 }
854
855 cx.notify();
856 }
857 }
858
859 fn reapply_batch_diff(&mut self, cx: &mut Context<Self>) -> Task<()> {
860 let old_snapshot = self.snapshot.clone();
861 let old_range = self.range.to_point(&old_snapshot);
862 let new_snapshot = self.buffer.read(cx).snapshot(cx);
863 let new_range = self.range.to_point(&new_snapshot);
864
865 cx.spawn(async move |codegen, cx| {
866 let (deleted_row_ranges, inserted_row_ranges) = cx
867 .background_spawn(async move {
868 let old_text = old_snapshot
869 .text_for_range(
870 Point::new(old_range.start.row, 0)
871 ..Point::new(
872 old_range.end.row,
873 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
874 ),
875 )
876 .collect::<String>();
877 let new_text = new_snapshot
878 .text_for_range(
879 Point::new(new_range.start.row, 0)
880 ..Point::new(
881 new_range.end.row,
882 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
883 ),
884 )
885 .collect::<String>();
886
887 let old_start_row = old_range.start.row;
888 let new_start_row = new_range.start.row;
889 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
890 let mut inserted_row_ranges = Vec::new();
891 for (old_rows, new_rows) in line_diff(&old_text, &new_text) {
892 let old_rows = old_start_row + old_rows.start..old_start_row + old_rows.end;
893 let new_rows = new_start_row + new_rows.start..new_start_row + new_rows.end;
894 if !old_rows.is_empty() {
895 deleted_row_ranges.push((
896 new_snapshot.anchor_before(Point::new(new_rows.start, 0)),
897 old_rows.start..=old_rows.end - 1,
898 ));
899 }
900 if !new_rows.is_empty() {
901 let start = new_snapshot.anchor_before(Point::new(new_rows.start, 0));
902 let new_end_row = new_rows.end - 1;
903 let end = new_snapshot.anchor_before(Point::new(
904 new_end_row,
905 new_snapshot.line_len(MultiBufferRow(new_end_row)),
906 ));
907 inserted_row_ranges.push(start..end);
908 }
909 }
910 (deleted_row_ranges, inserted_row_ranges)
911 })
912 .await;
913
914 codegen
915 .update(cx, |codegen, cx| {
916 codegen.diff.deleted_row_ranges = deleted_row_ranges;
917 codegen.diff.inserted_row_ranges = inserted_row_ranges;
918 cx.notify();
919 })
920 .ok();
921 })
922 }
923}
924
925#[derive(Copy, Clone, Debug)]
926pub enum CodegenEvent {
927 Finished,
928 Undone,
929}
930
931struct StripInvalidSpans<T> {
932 stream: T,
933 stream_done: bool,
934 buffer: String,
935 first_line: bool,
936 line_end: bool,
937 starts_with_code_block: bool,
938}
939
940impl<T> StripInvalidSpans<T>
941where
942 T: Stream<Item = Result<String>>,
943{
944 fn new(stream: T) -> Self {
945 Self {
946 stream,
947 stream_done: false,
948 buffer: String::new(),
949 first_line: true,
950 line_end: false,
951 starts_with_code_block: false,
952 }
953 }
954}
955
956impl<T> Stream for StripInvalidSpans<T>
957where
958 T: Stream<Item = Result<String>>,
959{
960 type Item = Result<String>;
961
962 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
963 const CODE_BLOCK_DELIMITER: &str = "```";
964 const CURSOR_SPAN: &str = "<|CURSOR|>";
965
966 let this = unsafe { self.get_unchecked_mut() };
967 loop {
968 if !this.stream_done {
969 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
970 match stream.as_mut().poll_next(cx) {
971 Poll::Ready(Some(Ok(chunk))) => {
972 this.buffer.push_str(&chunk);
973 }
974 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
975 Poll::Ready(None) => {
976 this.stream_done = true;
977 }
978 Poll::Pending => return Poll::Pending,
979 }
980 }
981
982 let mut chunk = String::new();
983 let mut consumed = 0;
984 if !this.buffer.is_empty() {
985 let mut lines = this.buffer.split('\n').enumerate().peekable();
986 while let Some((line_ix, line)) = lines.next() {
987 if line_ix > 0 {
988 this.first_line = false;
989 }
990
991 if this.first_line {
992 let trimmed_line = line.trim();
993 if lines.peek().is_some() {
994 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
995 consumed += line.len() + 1;
996 this.starts_with_code_block = true;
997 continue;
998 }
999 } else if trimmed_line.is_empty()
1000 || prefixes(CODE_BLOCK_DELIMITER)
1001 .any(|prefix| trimmed_line.starts_with(prefix))
1002 {
1003 break;
1004 }
1005 }
1006
1007 let line_without_cursor = line.replace(CURSOR_SPAN, "");
1008 if lines.peek().is_some() {
1009 if this.line_end {
1010 chunk.push('\n');
1011 }
1012
1013 chunk.push_str(&line_without_cursor);
1014 this.line_end = true;
1015 consumed += line.len() + 1;
1016 } else if this.stream_done {
1017 if !this.starts_with_code_block
1018 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
1019 {
1020 if this.line_end {
1021 chunk.push('\n');
1022 }
1023
1024 chunk.push_str(&line);
1025 }
1026
1027 consumed += line.len();
1028 } else {
1029 let trimmed_line = line.trim();
1030 if trimmed_line.is_empty()
1031 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
1032 || prefixes(CODE_BLOCK_DELIMITER)
1033 .any(|prefix| trimmed_line.ends_with(prefix))
1034 {
1035 break;
1036 } else {
1037 if this.line_end {
1038 chunk.push('\n');
1039 this.line_end = false;
1040 }
1041
1042 chunk.push_str(&line_without_cursor);
1043 consumed += line.len();
1044 }
1045 }
1046 }
1047 }
1048
1049 this.buffer = this.buffer.split_off(consumed);
1050 if !chunk.is_empty() {
1051 return Poll::Ready(Some(Ok(chunk)));
1052 } else if this.stream_done {
1053 return Poll::Ready(None);
1054 }
1055 }
1056 }
1057}
1058
1059fn prefixes(text: &str) -> impl Iterator<Item = &str> {
1060 (0..text.len() - 1).map(|ix| &text[..ix + 1])
1061}
1062
1063#[derive(Default)]
1064pub struct Diff {
1065 pub deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1066 pub inserted_row_ranges: Vec<Range<Anchor>>,
1067}
1068
1069impl Diff {
1070 fn is_empty(&self) -> bool {
1071 self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
1072 }
1073}
1074
1075#[cfg(test)]
1076mod tests {
1077 use super::*;
1078 use fs::FakeFs;
1079 use futures::{
1080 Stream,
1081 stream::{self},
1082 };
1083 use gpui::TestAppContext;
1084 use indoc::indoc;
1085 use language::{
1086 Buffer, Language, LanguageConfig, LanguageMatcher, Point, language_settings,
1087 tree_sitter_rust,
1088 };
1089 use language_model::{LanguageModelRegistry, TokenUsage};
1090 use rand::prelude::*;
1091 use serde::Serialize;
1092 use settings::SettingsStore;
1093 use std::{future, sync::Arc};
1094
1095 #[derive(Serialize)]
1096 pub struct DummyCompletionRequest {
1097 pub name: String,
1098 }
1099
1100 #[gpui::test(iterations = 10)]
1101 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
1102 init_test(cx);
1103
1104 let text = indoc! {"
1105 fn main() {
1106 let x = 0;
1107 for _ in 0..10 {
1108 x += 1;
1109 }
1110 }
1111 "};
1112 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1113 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1114 let range = buffer.read_with(cx, |buffer, cx| {
1115 let snapshot = buffer.snapshot(cx);
1116 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
1117 });
1118 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1119 let fs = FakeFs::new(cx.executor());
1120 let project = Project::test(fs, vec![], cx).await;
1121 let codegen = cx.new(|cx| {
1122 CodegenAlternative::new(
1123 buffer.clone(),
1124 range.clone(),
1125 true,
1126 None,
1127 project.downgrade(),
1128 None,
1129 None,
1130 prompt_builder,
1131 cx,
1132 )
1133 });
1134
1135 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1136
1137 let mut new_text = concat!(
1138 " let mut x = 0;\n",
1139 " while x < 10 {\n",
1140 " x += 1;\n",
1141 " }",
1142 );
1143 while !new_text.is_empty() {
1144 let max_len = cmp::min(new_text.len(), 10);
1145 let len = rng.gen_range(1..=max_len);
1146 let (chunk, suffix) = new_text.split_at(len);
1147 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1148 new_text = suffix;
1149 cx.background_executor.run_until_parked();
1150 }
1151 drop(chunks_tx);
1152 cx.background_executor.run_until_parked();
1153
1154 assert_eq!(
1155 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1156 indoc! {"
1157 fn main() {
1158 let mut x = 0;
1159 while x < 10 {
1160 x += 1;
1161 }
1162 }
1163 "}
1164 );
1165 }
1166
1167 #[gpui::test(iterations = 10)]
1168 async fn test_autoindent_when_generating_past_indentation(
1169 cx: &mut TestAppContext,
1170 mut rng: StdRng,
1171 ) {
1172 init_test(cx);
1173
1174 let text = indoc! {"
1175 fn main() {
1176 le
1177 }
1178 "};
1179 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1180 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1181 let range = buffer.read_with(cx, |buffer, cx| {
1182 let snapshot = buffer.snapshot(cx);
1183 snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
1184 });
1185 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1186 let fs = FakeFs::new(cx.executor());
1187 let project = Project::test(fs, vec![], cx).await;
1188 let codegen = cx.new(|cx| {
1189 CodegenAlternative::new(
1190 buffer.clone(),
1191 range.clone(),
1192 true,
1193 None,
1194 project.downgrade(),
1195 None,
1196 None,
1197 prompt_builder,
1198 cx,
1199 )
1200 });
1201
1202 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1203
1204 cx.background_executor.run_until_parked();
1205
1206 let mut new_text = concat!(
1207 "t mut x = 0;\n",
1208 "while x < 10 {\n",
1209 " x += 1;\n",
1210 "}", //
1211 );
1212 while !new_text.is_empty() {
1213 let max_len = cmp::min(new_text.len(), 10);
1214 let len = rng.gen_range(1..=max_len);
1215 let (chunk, suffix) = new_text.split_at(len);
1216 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1217 new_text = suffix;
1218 cx.background_executor.run_until_parked();
1219 }
1220 drop(chunks_tx);
1221 cx.background_executor.run_until_parked();
1222
1223 assert_eq!(
1224 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1225 indoc! {"
1226 fn main() {
1227 let mut x = 0;
1228 while x < 10 {
1229 x += 1;
1230 }
1231 }
1232 "}
1233 );
1234 }
1235
1236 #[gpui::test(iterations = 10)]
1237 async fn test_autoindent_when_generating_before_indentation(
1238 cx: &mut TestAppContext,
1239 mut rng: StdRng,
1240 ) {
1241 init_test(cx);
1242
1243 let text = concat!(
1244 "fn main() {\n",
1245 " \n",
1246 "}\n" //
1247 );
1248 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1249 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1250 let range = buffer.read_with(cx, |buffer, cx| {
1251 let snapshot = buffer.snapshot(cx);
1252 snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
1253 });
1254 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1255 let fs = FakeFs::new(cx.executor());
1256 let project = Project::test(fs, vec![], cx).await;
1257 let codegen = cx.new(|cx| {
1258 CodegenAlternative::new(
1259 buffer.clone(),
1260 range.clone(),
1261 true,
1262 None,
1263 project.downgrade(),
1264 None,
1265 None,
1266 prompt_builder,
1267 cx,
1268 )
1269 });
1270
1271 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1272
1273 cx.background_executor.run_until_parked();
1274
1275 let mut new_text = concat!(
1276 "let mut x = 0;\n",
1277 "while x < 10 {\n",
1278 " x += 1;\n",
1279 "}", //
1280 );
1281 while !new_text.is_empty() {
1282 let max_len = cmp::min(new_text.len(), 10);
1283 let len = rng.gen_range(1..=max_len);
1284 let (chunk, suffix) = new_text.split_at(len);
1285 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1286 new_text = suffix;
1287 cx.background_executor.run_until_parked();
1288 }
1289 drop(chunks_tx);
1290 cx.background_executor.run_until_parked();
1291
1292 assert_eq!(
1293 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1294 indoc! {"
1295 fn main() {
1296 let mut x = 0;
1297 while x < 10 {
1298 x += 1;
1299 }
1300 }
1301 "}
1302 );
1303 }
1304
1305 #[gpui::test(iterations = 10)]
1306 async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
1307 init_test(cx);
1308
1309 let text = indoc! {"
1310 func main() {
1311 \tx := 0
1312 \tfor i := 0; i < 10; i++ {
1313 \t\tx++
1314 \t}
1315 }
1316 "};
1317 let buffer = cx.new(|cx| Buffer::local(text, cx));
1318 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1319 let range = buffer.read_with(cx, |buffer, cx| {
1320 let snapshot = buffer.snapshot(cx);
1321 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
1322 });
1323 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1324 let fs = FakeFs::new(cx.executor());
1325 let project = Project::test(fs, vec![], cx).await;
1326 let codegen = cx.new(|cx| {
1327 CodegenAlternative::new(
1328 buffer.clone(),
1329 range.clone(),
1330 true,
1331 None,
1332 project.downgrade(),
1333 None,
1334 None,
1335 prompt_builder,
1336 cx,
1337 )
1338 });
1339
1340 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1341 let new_text = concat!(
1342 "func main() {\n",
1343 "\tx := 0\n",
1344 "\tfor x < 10 {\n",
1345 "\t\tx++\n",
1346 "\t}", //
1347 );
1348 chunks_tx.unbounded_send(new_text.to_string()).unwrap();
1349 drop(chunks_tx);
1350 cx.background_executor.run_until_parked();
1351
1352 assert_eq!(
1353 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1354 indoc! {"
1355 func main() {
1356 \tx := 0
1357 \tfor x < 10 {
1358 \t\tx++
1359 \t}
1360 }
1361 "}
1362 );
1363 }
1364
1365 #[gpui::test]
1366 async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
1367 init_test(cx);
1368
1369 let text = indoc! {"
1370 fn main() {
1371 let x = 0;
1372 }
1373 "};
1374 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1375 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1376 let range = buffer.read_with(cx, |buffer, cx| {
1377 let snapshot = buffer.snapshot(cx);
1378 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
1379 });
1380 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1381 let fs = FakeFs::new(cx.executor());
1382 let project = Project::test(fs, vec![], cx).await;
1383 let codegen = cx.new(|cx| {
1384 CodegenAlternative::new(
1385 buffer.clone(),
1386 range.clone(),
1387 false,
1388 None,
1389 project.downgrade(),
1390 None,
1391 None,
1392 prompt_builder,
1393 cx,
1394 )
1395 });
1396
1397 let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1398 chunks_tx
1399 .unbounded_send("let mut x = 0;\nx += 1;".to_string())
1400 .unwrap();
1401 drop(chunks_tx);
1402 cx.run_until_parked();
1403
1404 // The codegen is inactive, so the buffer doesn't get modified.
1405 assert_eq!(
1406 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1407 text
1408 );
1409
1410 // Activating the codegen applies the changes.
1411 codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
1412 assert_eq!(
1413 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1414 indoc! {"
1415 fn main() {
1416 let mut x = 0;
1417 x += 1;
1418 }
1419 "}
1420 );
1421
1422 // Deactivating the codegen undoes the changes.
1423 codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
1424 cx.run_until_parked();
1425 assert_eq!(
1426 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1427 text
1428 );
1429 }
1430
1431 #[gpui::test]
1432 async fn test_strip_invalid_spans_from_codeblock() {
1433 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
1434 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
1435 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
1436 assert_chunks(
1437 "```html\n```js\nLorem ipsum dolor\n```\n```",
1438 "```js\nLorem ipsum dolor\n```",
1439 )
1440 .await;
1441 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
1442 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
1443 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
1444 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
1445
1446 async fn assert_chunks(text: &str, expected_text: &str) {
1447 for chunk_size in 1..=text.len() {
1448 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
1449 .map(|chunk| chunk.unwrap())
1450 .collect::<String>()
1451 .await;
1452 assert_eq!(
1453 actual_text, expected_text,
1454 "failed to strip invalid spans, chunk size: {}",
1455 chunk_size
1456 );
1457 }
1458 }
1459
1460 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
1461 stream::iter(
1462 text.chars()
1463 .collect::<Vec<_>>()
1464 .chunks(size)
1465 .map(|chunk| Ok(chunk.iter().collect::<String>()))
1466 .collect::<Vec<_>>(),
1467 )
1468 }
1469 }
1470
1471 fn init_test(cx: &mut TestAppContext) {
1472 cx.update(LanguageModelRegistry::test);
1473 cx.set_global(cx.update(SettingsStore::test));
1474 cx.update(Project::init_settings);
1475 cx.update(language_settings::init);
1476 }
1477
1478 fn simulate_response_stream(
1479 codegen: Entity<CodegenAlternative>,
1480 cx: &mut TestAppContext,
1481 ) -> mpsc::UnboundedSender<String> {
1482 let (chunks_tx, chunks_rx) = mpsc::unbounded();
1483 codegen.update(cx, |codegen, cx| {
1484 codegen.handle_stream(
1485 String::new(),
1486 String::new(),
1487 None,
1488 future::ready(Ok(LanguageModelTextStream {
1489 message_id: None,
1490 stream: chunks_rx.map(Ok).boxed(),
1491 last_token_usage: Arc::new(Mutex::new(TokenUsage::default())),
1492 })),
1493 cx,
1494 );
1495 });
1496 chunks_tx
1497 }
1498
1499 fn rust_lang() -> Language {
1500 Language::new(
1501 LanguageConfig {
1502 name: "Rust".into(),
1503 matcher: LanguageMatcher {
1504 path_suffixes: vec!["rs".to_string()],
1505 ..Default::default()
1506 },
1507 ..Default::default()
1508 },
1509 Some(tree_sitter_rust::LANGUAGE.into()),
1510 )
1511 .with_indents_query(
1512 r#"
1513 (call_expression) @indent
1514 (field_expression) @indent
1515 (_ "(" ")" @end) @indent
1516 (_ "{" "}" @end) @indent
1517 "#,
1518 )
1519 .unwrap()
1520 }
1521}