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