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