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