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