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