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