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