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