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