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