1use crate::{context::LoadedContext, inline_prompt_editor::CodegenStatus};
2use agent_settings::AgentSettings;
3use anyhow::{Context as _, Result};
4use uuid::Uuid;
5
6use cloud_llm_client::CompletionIntent;
7use collections::HashSet;
8use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};
9use feature_flags::{FeatureFlagAppExt as _, InlineAssistantUseToolFeatureFlag};
10use futures::{
11 SinkExt, Stream, StreamExt, TryStreamExt as _,
12 channel::mpsc,
13 future::{LocalBoxFuture, Shared},
14 join,
15 stream::BoxStream,
16};
17use gpui::{App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Subscription, Task};
18use language::{Buffer, IndentKind, LanguageName, Point, TransactionId, line_diff};
19use language_model::{
20 LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
21 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
22 LanguageModelRequestTool, LanguageModelTextStream, LanguageModelToolChoice,
23 LanguageModelToolUse, Role, TokenUsage,
24};
25use multi_buffer::MultiBufferRow;
26use parking_lot::Mutex;
27use prompt_store::PromptBuilder;
28use rope::Rope;
29use schemars::JsonSchema;
30use serde::{Deserialize, Serialize};
31use settings::Settings as _;
32use smol::future::FutureExt;
33use std::{
34 cmp,
35 future::Future,
36 iter,
37 ops::{Range, RangeInclusive},
38 pin::Pin,
39 sync::Arc,
40 task::{self, Poll},
41 time::Instant,
42};
43use streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff};
44
45/// Use this tool when you cannot or should not make a rewrite. This includes:
46/// - The user's request is unclear, ambiguous, or nonsensical
47/// - The requested change cannot be made by only editing the <rewrite_this> section
48#[derive(Debug, Serialize, Deserialize, JsonSchema)]
49pub struct FailureMessageInput {
50 /// A brief message to the user explaining why you're unable to fulfill the request or to ask a question about the request.
51 #[serde(default)]
52 pub message: String,
53}
54
55/// Replaces text in <rewrite_this></rewrite_this> tags with your replacement_text.
56/// Only use this tool when you are confident you understand the user's request and can fulfill it
57/// by editing the marked section.
58#[derive(Debug, Serialize, Deserialize, JsonSchema)]
59pub struct RewriteSectionInput {
60 /// The text to replace the section with.
61 #[serde(default)]
62 pub replacement_text: String,
63}
64
65pub struct BufferCodegen {
66 alternatives: Vec<Entity<CodegenAlternative>>,
67 pub active_alternative: usize,
68 seen_alternatives: HashSet<usize>,
69 subscriptions: Vec<Subscription>,
70 buffer: Entity<MultiBuffer>,
71 range: Range<Anchor>,
72 initial_transaction_id: Option<TransactionId>,
73 builder: Arc<PromptBuilder>,
74 pub is_insertion: bool,
75 session_id: Uuid,
76}
77
78impl BufferCodegen {
79 pub fn new(
80 buffer: Entity<MultiBuffer>,
81 range: Range<Anchor>,
82 initial_transaction_id: Option<TransactionId>,
83 session_id: Uuid,
84 builder: Arc<PromptBuilder>,
85 cx: &mut Context<Self>,
86 ) -> Self {
87 let codegen = cx.new(|cx| {
88 CodegenAlternative::new(
89 buffer.clone(),
90 range.clone(),
91 false,
92 builder.clone(),
93 session_id,
94 cx,
95 )
96 });
97 let mut this = Self {
98 is_insertion: range.to_offset(&buffer.read(cx).snapshot(cx)).is_empty(),
99 alternatives: vec![codegen],
100 active_alternative: 0,
101 seen_alternatives: HashSet::default(),
102 subscriptions: Vec::new(),
103 buffer,
104 range,
105 initial_transaction_id,
106 builder,
107 session_id,
108 };
109 this.activate(0, cx);
110 this
111 }
112
113 fn subscribe_to_alternative(&mut self, cx: &mut Context<Self>) {
114 let codegen = self.active_alternative().clone();
115 self.subscriptions.clear();
116 self.subscriptions
117 .push(cx.observe(&codegen, |_, _, cx| cx.notify()));
118 self.subscriptions
119 .push(cx.subscribe(&codegen, |_, _, event, cx| cx.emit(*event)));
120 }
121
122 pub fn active_completion(&self, cx: &App) -> Option<String> {
123 self.active_alternative().read(cx).current_completion()
124 }
125
126 pub fn active_alternative(&self) -> &Entity<CodegenAlternative> {
127 &self.alternatives[self.active_alternative]
128 }
129
130 pub fn language_name(&self, cx: &App) -> Option<LanguageName> {
131 self.active_alternative().read(cx).language_name(cx)
132 }
133
134 pub fn status<'a>(&self, cx: &'a App) -> &'a CodegenStatus {
135 &self.active_alternative().read(cx).status
136 }
137
138 pub fn alternative_count(&self, cx: &App) -> usize {
139 LanguageModelRegistry::read_global(cx)
140 .inline_alternative_models()
141 .len()
142 + 1
143 }
144
145 pub fn cycle_prev(&mut self, cx: &mut Context<Self>) {
146 let next_active_ix = if self.active_alternative == 0 {
147 self.alternatives.len() - 1
148 } else {
149 self.active_alternative - 1
150 };
151 self.activate(next_active_ix, cx);
152 }
153
154 pub fn cycle_next(&mut self, cx: &mut Context<Self>) {
155 let next_active_ix = (self.active_alternative + 1) % self.alternatives.len();
156 self.activate(next_active_ix, cx);
157 }
158
159 fn activate(&mut self, index: usize, cx: &mut Context<Self>) {
160 self.active_alternative()
161 .update(cx, |codegen, cx| codegen.set_active(false, cx));
162 self.seen_alternatives.insert(index);
163 self.active_alternative = index;
164 self.active_alternative()
165 .update(cx, |codegen, cx| codegen.set_active(true, cx));
166 self.subscribe_to_alternative(cx);
167 cx.notify();
168 }
169
170 pub fn start(
171 &mut self,
172 primary_model: Arc<dyn LanguageModel>,
173 user_prompt: String,
174 context_task: Shared<Task<Option<LoadedContext>>>,
175 cx: &mut Context<Self>,
176 ) -> Result<()> {
177 let alternative_models = LanguageModelRegistry::read_global(cx)
178 .inline_alternative_models()
179 .to_vec();
180
181 self.active_alternative()
182 .update(cx, |alternative, cx| alternative.undo(cx));
183 self.activate(0, cx);
184 self.alternatives.truncate(1);
185
186 for _ in 0..alternative_models.len() {
187 self.alternatives.push(cx.new(|cx| {
188 CodegenAlternative::new(
189 self.buffer.clone(),
190 self.range.clone(),
191 false,
192 self.builder.clone(),
193 self.session_id,
194 cx,
195 )
196 }));
197 }
198
199 for (model, alternative) in iter::once(primary_model)
200 .chain(alternative_models)
201 .zip(&self.alternatives)
202 {
203 alternative.update(cx, |alternative, cx| {
204 alternative.start(user_prompt.clone(), context_task.clone(), model.clone(), cx)
205 })?;
206 }
207
208 Ok(())
209 }
210
211 pub fn stop(&mut self, cx: &mut Context<Self>) {
212 for codegen in &self.alternatives {
213 codegen.update(cx, |codegen, cx| codegen.stop(cx));
214 }
215 }
216
217 pub fn undo(&mut self, cx: &mut Context<Self>) {
218 self.active_alternative()
219 .update(cx, |codegen, cx| codegen.undo(cx));
220
221 self.buffer.update(cx, |buffer, cx| {
222 if let Some(transaction_id) = self.initial_transaction_id.take() {
223 buffer.undo_transaction(transaction_id, cx);
224 buffer.refresh_preview(cx);
225 }
226 });
227 }
228
229 pub fn buffer(&self, cx: &App) -> Entity<MultiBuffer> {
230 self.active_alternative().read(cx).buffer.clone()
231 }
232
233 pub fn old_buffer(&self, cx: &App) -> Entity<Buffer> {
234 self.active_alternative().read(cx).old_buffer.clone()
235 }
236
237 pub fn snapshot(&self, cx: &App) -> MultiBufferSnapshot {
238 self.active_alternative().read(cx).snapshot.clone()
239 }
240
241 pub fn edit_position(&self, cx: &App) -> Option<Anchor> {
242 self.active_alternative().read(cx).edit_position
243 }
244
245 pub fn diff<'a>(&self, cx: &'a App) -> &'a Diff {
246 &self.active_alternative().read(cx).diff
247 }
248
249 pub fn last_equal_ranges<'a>(&self, cx: &'a App) -> &'a [Range<Anchor>] {
250 self.active_alternative().read(cx).last_equal_ranges()
251 }
252
253 pub fn selected_text<'a>(&self, cx: &'a App) -> Option<&'a str> {
254 self.active_alternative().read(cx).selected_text()
255 }
256
257 pub fn session_id(&self) -> Uuid {
258 self.session_id
259 }
260}
261
262impl EventEmitter<CodegenEvent> for BufferCodegen {}
263
264pub struct CodegenAlternative {
265 buffer: Entity<MultiBuffer>,
266 old_buffer: Entity<Buffer>,
267 snapshot: MultiBufferSnapshot,
268 edit_position: Option<Anchor>,
269 range: Range<Anchor>,
270 last_equal_ranges: Vec<Range<Anchor>>,
271 transformation_transaction_id: Option<TransactionId>,
272 status: CodegenStatus,
273 generation: Task<()>,
274 diff: Diff,
275 _subscription: gpui::Subscription,
276 builder: Arc<PromptBuilder>,
277 active: bool,
278 edits: Vec<(Range<Anchor>, String)>,
279 line_operations: Vec<LineOperation>,
280 elapsed_time: Option<f64>,
281 completion: Option<String>,
282 selected_text: Option<String>,
283 pub message_id: Option<String>,
284 session_id: Uuid,
285 pub description: Option<String>,
286 pub failure: Option<String>,
287}
288
289impl EventEmitter<CodegenEvent> for CodegenAlternative {}
290
291impl CodegenAlternative {
292 pub fn new(
293 buffer: Entity<MultiBuffer>,
294 range: Range<Anchor>,
295 active: bool,
296 builder: Arc<PromptBuilder>,
297 session_id: Uuid,
298 cx: &mut Context<Self>,
299 ) -> Self {
300 let snapshot = buffer.read(cx).snapshot(cx);
301
302 let (old_buffer, _, _) = snapshot
303 .range_to_buffer_ranges(range.clone())
304 .pop()
305 .unwrap();
306 let old_buffer = cx.new(|cx| {
307 let text = old_buffer.as_rope().clone();
308 let line_ending = old_buffer.line_ending();
309 let language = old_buffer.language().cloned();
310 let language_registry = buffer
311 .read(cx)
312 .buffer(old_buffer.remote_id())
313 .unwrap()
314 .read(cx)
315 .language_registry();
316
317 let mut buffer = Buffer::local_normalized(text, line_ending, cx);
318 buffer.set_language(language, cx);
319 if let Some(language_registry) = language_registry {
320 buffer.set_language_registry(language_registry);
321 }
322 buffer
323 });
324
325 Self {
326 buffer: buffer.clone(),
327 old_buffer,
328 edit_position: None,
329 message_id: None,
330 snapshot,
331 last_equal_ranges: Default::default(),
332 transformation_transaction_id: None,
333 status: CodegenStatus::Idle,
334 generation: Task::ready(()),
335 diff: Diff::default(),
336 builder,
337 active: active,
338 edits: Vec::new(),
339 line_operations: Vec::new(),
340 range,
341 elapsed_time: None,
342 completion: None,
343 selected_text: None,
344 session_id,
345 description: None,
346 failure: None,
347 _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
348 }
349 }
350
351 pub fn language_name(&self, cx: &App) -> Option<LanguageName> {
352 self.old_buffer
353 .read(cx)
354 .language()
355 .map(|language| language.name())
356 }
357
358 pub fn set_active(&mut self, active: bool, cx: &mut Context<Self>) {
359 if active != self.active {
360 self.active = active;
361
362 if self.active {
363 let edits = self.edits.clone();
364 self.apply_edits(edits, cx);
365 if matches!(self.status, CodegenStatus::Pending) {
366 let line_operations = self.line_operations.clone();
367 self.reapply_line_based_diff(line_operations, cx);
368 } else {
369 self.reapply_batch_diff(cx).detach();
370 }
371 } else if let Some(transaction_id) = self.transformation_transaction_id.take() {
372 self.buffer.update(cx, |buffer, cx| {
373 buffer.undo_transaction(transaction_id, cx);
374 buffer.forget_transaction(transaction_id, cx);
375 });
376 }
377 }
378 }
379
380 fn handle_buffer_event(
381 &mut self,
382 _buffer: Entity<MultiBuffer>,
383 event: &multi_buffer::Event,
384 cx: &mut Context<Self>,
385 ) {
386 if let multi_buffer::Event::TransactionUndone { transaction_id } = event
387 && self.transformation_transaction_id == Some(*transaction_id)
388 {
389 self.transformation_transaction_id = None;
390 self.generation = Task::ready(());
391 cx.emit(CodegenEvent::Undone);
392 }
393 }
394
395 pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
396 &self.last_equal_ranges
397 }
398
399 pub fn use_streaming_tools(model: &dyn LanguageModel, cx: &App) -> bool {
400 model.supports_streaming_tools()
401 && cx.has_flag::<InlineAssistantUseToolFeatureFlag>()
402 && AgentSettings::get_global(cx).inline_assistant_use_streaming_tools
403 }
404
405 pub fn start(
406 &mut self,
407 user_prompt: String,
408 context_task: Shared<Task<Option<LoadedContext>>>,
409 model: Arc<dyn LanguageModel>,
410 cx: &mut Context<Self>,
411 ) -> Result<()> {
412 if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {
413 self.buffer.update(cx, |buffer, cx| {
414 buffer.undo_transaction(transformation_transaction_id, cx);
415 });
416 }
417
418 self.edit_position = Some(self.range.start.bias_right(&self.snapshot));
419
420 if Self::use_streaming_tools(model.as_ref(), cx) {
421 let request = self.build_request(&model, user_prompt, context_task, cx)?;
422 let completion_events = cx.spawn({
423 let model = model.clone();
424 async move |_, cx| model.stream_completion(request.await, cx).await
425 });
426 self.generation = self.handle_completion(model, completion_events, cx);
427 } else {
428 let stream: LocalBoxFuture<Result<LanguageModelTextStream>> =
429 if user_prompt.trim().to_lowercase() == "delete" {
430 async { Ok(LanguageModelTextStream::default()) }.boxed_local()
431 } else {
432 let request = self.build_request(&model, user_prompt, context_task, cx)?;
433 cx.spawn({
434 let model = model.clone();
435 async move |_, cx| {
436 Ok(model.stream_completion_text(request.await, cx).await?)
437 }
438 })
439 .boxed_local()
440 };
441 self.generation = self.handle_stream(model, stream, cx);
442 }
443
444 Ok(())
445 }
446
447 fn build_request_tools(
448 &self,
449 model: &Arc<dyn LanguageModel>,
450 user_prompt: String,
451 context_task: Shared<Task<Option<LoadedContext>>>,
452 cx: &mut App,
453 ) -> Result<Task<LanguageModelRequest>> {
454 let buffer = self.buffer.read(cx).snapshot(cx);
455 let language = buffer.language_at(self.range.start);
456 let language_name = if let Some(language) = language.as_ref() {
457 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
458 None
459 } else {
460 Some(language.name())
461 }
462 } else {
463 None
464 };
465
466 let language_name = language_name.as_ref();
467 let start = buffer.point_to_buffer_offset(self.range.start);
468 let end = buffer.point_to_buffer_offset(self.range.end);
469 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
470 let (start_buffer, start_buffer_offset) = start;
471 let (end_buffer, end_buffer_offset) = end;
472 if start_buffer.remote_id() == end_buffer.remote_id() {
473 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
474 } else {
475 anyhow::bail!("invalid transformation range");
476 }
477 } else {
478 anyhow::bail!("invalid transformation range");
479 };
480
481 let system_prompt = self
482 .builder
483 .generate_inline_transformation_prompt_tools(
484 language_name,
485 buffer,
486 range.start.0..range.end.0,
487 )
488 .context("generating content prompt")?;
489
490 let temperature = AgentSettings::temperature_for_model(model, cx);
491
492 let tool_input_format = model.tool_input_format();
493 let tool_choice = model
494 .supports_tool_choice(LanguageModelToolChoice::Any)
495 .then_some(LanguageModelToolChoice::Any);
496
497 Ok(cx.spawn(async move |_cx| {
498 let mut messages = vec![LanguageModelRequestMessage {
499 role: Role::System,
500 content: vec![system_prompt.into()],
501 cache: false,
502 reasoning_details: None,
503 }];
504
505 let mut user_message = LanguageModelRequestMessage {
506 role: Role::User,
507 content: Vec::new(),
508 cache: false,
509 reasoning_details: None,
510 };
511
512 if let Some(context) = context_task.await {
513 context.add_to_request_message(&mut user_message);
514 }
515
516 user_message.content.push(user_prompt.into());
517 messages.push(user_message);
518
519 let tools = vec![
520 LanguageModelRequestTool {
521 name: "rewrite_section".to_string(),
522 description: "Replaces text in <rewrite_this></rewrite_this> tags with your replacement_text.".to_string(),
523 input_schema: language_model::tool_schema::root_schema_for::<RewriteSectionInput>(tool_input_format).to_value(),
524 },
525 LanguageModelRequestTool {
526 name: "failure_message".to_string(),
527 description: "Use this tool to provide a message to the user when you're unable to complete a task.".to_string(),
528 input_schema: language_model::tool_schema::root_schema_for::<FailureMessageInput>(tool_input_format).to_value(),
529 },
530 ];
531
532 LanguageModelRequest {
533 thread_id: None,
534 prompt_id: None,
535 intent: Some(CompletionIntent::InlineAssist),
536 mode: None,
537 tools,
538 tool_choice,
539 stop: Vec::new(),
540 temperature,
541 messages,
542 thinking_allowed: false,
543 }
544 }))
545 }
546
547 fn build_request(
548 &self,
549 model: &Arc<dyn LanguageModel>,
550 user_prompt: String,
551 context_task: Shared<Task<Option<LoadedContext>>>,
552 cx: &mut App,
553 ) -> Result<Task<LanguageModelRequest>> {
554 if Self::use_streaming_tools(model.as_ref(), cx) {
555 return self.build_request_tools(model, user_prompt, context_task, cx);
556 }
557
558 let buffer = self.buffer.read(cx).snapshot(cx);
559 let language = buffer.language_at(self.range.start);
560 let language_name = if let Some(language) = language.as_ref() {
561 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
562 None
563 } else {
564 Some(language.name())
565 }
566 } else {
567 None
568 };
569
570 let language_name = language_name.as_ref();
571 let start = buffer.point_to_buffer_offset(self.range.start);
572 let end = buffer.point_to_buffer_offset(self.range.end);
573 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
574 let (start_buffer, start_buffer_offset) = start;
575 let (end_buffer, end_buffer_offset) = end;
576 if start_buffer.remote_id() == end_buffer.remote_id() {
577 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
578 } else {
579 anyhow::bail!("invalid transformation range");
580 }
581 } else {
582 anyhow::bail!("invalid transformation range");
583 };
584
585 let prompt = self
586 .builder
587 .generate_inline_transformation_prompt(
588 user_prompt,
589 language_name,
590 buffer,
591 range.start.0..range.end.0,
592 )
593 .context("generating content prompt")?;
594
595 let temperature = AgentSettings::temperature_for_model(model, cx);
596
597 Ok(cx.spawn(async move |_cx| {
598 let mut request_message = LanguageModelRequestMessage {
599 role: Role::User,
600 content: Vec::new(),
601 cache: false,
602 reasoning_details: None,
603 };
604
605 if let Some(context) = context_task.await {
606 context.add_to_request_message(&mut request_message);
607 }
608
609 request_message.content.push(prompt.into());
610
611 LanguageModelRequest {
612 thread_id: None,
613 prompt_id: None,
614 intent: Some(CompletionIntent::InlineAssist),
615 mode: None,
616 tools: Vec::new(),
617 tool_choice: None,
618 stop: Vec::new(),
619 temperature,
620 messages: vec![request_message],
621 thinking_allowed: false,
622 }
623 }))
624 }
625
626 pub fn handle_stream(
627 &mut self,
628 model: Arc<dyn LanguageModel>,
629 stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,
630 cx: &mut Context<Self>,
631 ) -> Task<()> {
632 let anthropic_reporter = language_model::AnthropicEventReporter::new(&model, cx);
633 let session_id = self.session_id;
634 let model_telemetry_id = model.telemetry_id();
635 let model_provider_id = model.provider_id().to_string();
636 let start_time = Instant::now();
637
638 // Make a new snapshot and re-resolve anchor in case the document was modified.
639 // This can happen often if the editor loses focus and is saved + reformatted,
640 // as in https://github.com/zed-industries/zed/issues/39088
641 self.snapshot = self.buffer.read(cx).snapshot(cx);
642 self.range = self.snapshot.anchor_after(self.range.start)
643 ..self.snapshot.anchor_after(self.range.end);
644
645 let snapshot = self.snapshot.clone();
646 let selected_text = snapshot
647 .text_for_range(self.range.start..self.range.end)
648 .collect::<Rope>();
649
650 self.selected_text = Some(selected_text.to_string());
651
652 let selection_start = self.range.start.to_point(&snapshot);
653
654 // Start with the indentation of the first line in the selection
655 let mut suggested_line_indent = snapshot
656 .suggested_indents(selection_start.row..=selection_start.row, cx)
657 .into_values()
658 .next()
659 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
660
661 // If the first line in the selection does not have indentation, check the following lines
662 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
663 for row in selection_start.row..=self.range.end.to_point(&snapshot).row {
664 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
665 // Prefer tabs if a line in the selection uses tabs as indentation
666 if line_indent.kind == IndentKind::Tab {
667 suggested_line_indent.kind = IndentKind::Tab;
668 break;
669 }
670 }
671 }
672
673 let language_name = {
674 let multibuffer = self.buffer.read(cx);
675 let snapshot = multibuffer.snapshot(cx);
676 let ranges = snapshot.range_to_buffer_ranges(self.range.clone());
677 ranges
678 .first()
679 .and_then(|(buffer, _, _)| buffer.language())
680 .map(|language| language.name())
681 };
682
683 self.diff = Diff::default();
684 self.status = CodegenStatus::Pending;
685 let mut edit_start = self.range.start.to_offset(&snapshot);
686 let completion = Arc::new(Mutex::new(String::new()));
687 let completion_clone = completion.clone();
688
689 cx.notify();
690 cx.spawn(async move |codegen, cx| {
691 let stream = stream.await;
692
693 let token_usage = stream
694 .as_ref()
695 .ok()
696 .map(|stream| stream.last_token_usage.clone());
697 let message_id = stream
698 .as_ref()
699 .ok()
700 .and_then(|stream| stream.message_id.clone());
701 let generate = async {
702 let model_telemetry_id = model_telemetry_id.clone();
703 let model_provider_id = model_provider_id.clone();
704 let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
705 let message_id = message_id.clone();
706 let line_based_stream_diff: Task<anyhow::Result<()>> = cx.background_spawn({
707 let anthropic_reporter = anthropic_reporter.clone();
708 let language_name = language_name.clone();
709 async move {
710 let mut response_latency = None;
711 let request_start = Instant::now();
712 let diff = async {
713 let chunks = StripInvalidSpans::new(
714 stream?.stream.map_err(|error| error.into()),
715 );
716 futures::pin_mut!(chunks);
717
718 let mut diff = StreamingDiff::new(selected_text.to_string());
719 let mut line_diff = LineDiff::default();
720
721 let mut new_text = String::new();
722 let mut base_indent = None;
723 let mut line_indent = None;
724 let mut first_line = true;
725
726 while let Some(chunk) = chunks.next().await {
727 if response_latency.is_none() {
728 response_latency = Some(request_start.elapsed());
729 }
730 let chunk = chunk?;
731 completion_clone.lock().push_str(&chunk);
732
733 let mut lines = chunk.split('\n').peekable();
734 while let Some(line) = lines.next() {
735 new_text.push_str(line);
736 if line_indent.is_none()
737 && let Some(non_whitespace_ch_ix) =
738 new_text.find(|ch: char| !ch.is_whitespace())
739 {
740 line_indent = Some(non_whitespace_ch_ix);
741 base_indent = base_indent.or(line_indent);
742
743 let line_indent = line_indent.unwrap();
744 let base_indent = base_indent.unwrap();
745 let indent_delta = line_indent as i32 - base_indent as i32;
746 let mut corrected_indent_len = cmp::max(
747 0,
748 suggested_line_indent.len as i32 + indent_delta,
749 )
750 as usize;
751 if first_line {
752 corrected_indent_len = corrected_indent_len
753 .saturating_sub(selection_start.column as usize);
754 }
755
756 let indent_char = suggested_line_indent.char();
757 let mut indent_buffer = [0; 4];
758 let indent_str =
759 indent_char.encode_utf8(&mut indent_buffer);
760 new_text.replace_range(
761 ..line_indent,
762 &indent_str.repeat(corrected_indent_len),
763 );
764 }
765
766 if line_indent.is_some() {
767 let char_ops = diff.push_new(&new_text);
768 line_diff.push_char_operations(&char_ops, &selected_text);
769 diff_tx
770 .send((char_ops, line_diff.line_operations()))
771 .await?;
772 new_text.clear();
773 }
774
775 if lines.peek().is_some() {
776 let char_ops = diff.push_new("\n");
777 line_diff.push_char_operations(&char_ops, &selected_text);
778 diff_tx
779 .send((char_ops, line_diff.line_operations()))
780 .await?;
781 if line_indent.is_none() {
782 // Don't write out the leading indentation in empty lines on the next line
783 // This is the case where the above if statement didn't clear the buffer
784 new_text.clear();
785 }
786 line_indent = None;
787 first_line = false;
788 }
789 }
790 }
791
792 let mut char_ops = diff.push_new(&new_text);
793 char_ops.extend(diff.finish());
794 line_diff.push_char_operations(&char_ops, &selected_text);
795 line_diff.finish(&selected_text);
796 diff_tx
797 .send((char_ops, line_diff.line_operations()))
798 .await?;
799
800 anyhow::Ok(())
801 };
802
803 let result = diff.await;
804
805 let error_message = result.as_ref().err().map(|error| error.to_string());
806 telemetry::event!(
807 "Assistant Responded",
808 kind = "inline",
809 phase = "response",
810 session_id = session_id.to_string(),
811 model = model_telemetry_id,
812 model_provider = model_provider_id,
813 language_name = language_name.as_ref().map(|n| n.to_string()),
814 message_id = message_id.as_deref(),
815 response_latency = response_latency,
816 error_message = error_message.as_deref(),
817 );
818
819 anthropic_reporter.report(language_model::AnthropicEventData {
820 completion_type: language_model::AnthropicCompletionType::Editor,
821 event: language_model::AnthropicEventType::Response,
822 language_name: language_name.map(|n| n.to_string()),
823 message_id,
824 });
825
826 result?;
827 Ok(())
828 }
829 });
830
831 while let Some((char_ops, line_ops)) = diff_rx.next().await {
832 codegen.update(cx, |codegen, cx| {
833 codegen.last_equal_ranges.clear();
834
835 let edits = char_ops
836 .into_iter()
837 .filter_map(|operation| match operation {
838 CharOperation::Insert { text } => {
839 let edit_start = snapshot.anchor_after(edit_start);
840 Some((edit_start..edit_start, text))
841 }
842 CharOperation::Delete { bytes } => {
843 let edit_end = edit_start + bytes;
844 let edit_range = snapshot.anchor_after(edit_start)
845 ..snapshot.anchor_before(edit_end);
846 edit_start = edit_end;
847 Some((edit_range, String::new()))
848 }
849 CharOperation::Keep { bytes } => {
850 let edit_end = edit_start + bytes;
851 let edit_range = snapshot.anchor_after(edit_start)
852 ..snapshot.anchor_before(edit_end);
853 edit_start = edit_end;
854 codegen.last_equal_ranges.push(edit_range);
855 None
856 }
857 })
858 .collect::<Vec<_>>();
859
860 if codegen.active {
861 codegen.apply_edits(edits.iter().cloned(), cx);
862 codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
863 }
864 codegen.edits.extend(edits);
865 codegen.line_operations = line_ops;
866 codegen.edit_position = Some(snapshot.anchor_after(edit_start));
867
868 cx.notify();
869 })?;
870 }
871
872 // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
873 // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
874 // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
875 let batch_diff_task =
876 codegen.update(cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
877 let (line_based_stream_diff, ()) = join!(line_based_stream_diff, batch_diff_task);
878 line_based_stream_diff?;
879
880 anyhow::Ok(())
881 };
882
883 let result = generate.await;
884 let elapsed_time = start_time.elapsed().as_secs_f64();
885
886 codegen
887 .update(cx, |this, cx| {
888 this.message_id = message_id;
889 this.last_equal_ranges.clear();
890 if let Err(error) = result {
891 this.status = CodegenStatus::Error(error);
892 } else {
893 this.status = CodegenStatus::Done;
894 }
895 this.elapsed_time = Some(elapsed_time);
896 this.completion = Some(completion.lock().clone());
897 if let Some(usage) = token_usage {
898 let usage = usage.lock();
899 telemetry::event!(
900 "Inline Assistant Completion",
901 model = model_telemetry_id,
902 model_provider = model_provider_id,
903 input_tokens = usage.input_tokens,
904 output_tokens = usage.output_tokens,
905 )
906 }
907
908 cx.emit(CodegenEvent::Finished);
909 cx.notify();
910 })
911 .ok();
912 })
913 }
914
915 pub fn current_completion(&self) -> Option<String> {
916 self.completion.clone()
917 }
918
919 #[cfg(any(test, feature = "test-support"))]
920 pub fn current_description(&self) -> Option<String> {
921 self.description.clone()
922 }
923
924 #[cfg(any(test, feature = "test-support"))]
925 pub fn current_failure(&self) -> Option<String> {
926 self.failure.clone()
927 }
928
929 pub fn selected_text(&self) -> Option<&str> {
930 self.selected_text.as_deref()
931 }
932
933 pub fn stop(&mut self, cx: &mut Context<Self>) {
934 self.last_equal_ranges.clear();
935 if self.diff.is_empty() {
936 self.status = CodegenStatus::Idle;
937 } else {
938 self.status = CodegenStatus::Done;
939 }
940 self.generation = Task::ready(());
941 cx.emit(CodegenEvent::Finished);
942 cx.notify();
943 }
944
945 pub fn undo(&mut self, cx: &mut Context<Self>) {
946 self.buffer.update(cx, |buffer, cx| {
947 if let Some(transaction_id) = self.transformation_transaction_id.take() {
948 buffer.undo_transaction(transaction_id, cx);
949 buffer.refresh_preview(cx);
950 }
951 });
952 }
953
954 fn apply_edits(
955 &mut self,
956 edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
957 cx: &mut Context<CodegenAlternative>,
958 ) {
959 let transaction = self.buffer.update(cx, |buffer, cx| {
960 // Avoid grouping agent edits with user edits.
961 buffer.finalize_last_transaction(cx);
962 buffer.start_transaction(cx);
963 buffer.edit(edits, None, cx);
964 buffer.end_transaction(cx)
965 });
966
967 if let Some(transaction) = transaction {
968 if let Some(first_transaction) = self.transformation_transaction_id {
969 // Group all agent edits into the first transaction.
970 self.buffer.update(cx, |buffer, cx| {
971 buffer.merge_transactions(transaction, first_transaction, cx)
972 });
973 } else {
974 self.transformation_transaction_id = Some(transaction);
975 self.buffer
976 .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
977 }
978 }
979 }
980
981 fn reapply_line_based_diff(
982 &mut self,
983 line_operations: impl IntoIterator<Item = LineOperation>,
984 cx: &mut Context<Self>,
985 ) {
986 let old_snapshot = self.snapshot.clone();
987 let old_range = self.range.to_point(&old_snapshot);
988 let new_snapshot = self.buffer.read(cx).snapshot(cx);
989 let new_range = self.range.to_point(&new_snapshot);
990
991 let mut old_row = old_range.start.row;
992 let mut new_row = new_range.start.row;
993
994 self.diff.deleted_row_ranges.clear();
995 self.diff.inserted_row_ranges.clear();
996 for operation in line_operations {
997 match operation {
998 LineOperation::Keep { lines } => {
999 old_row += lines;
1000 new_row += lines;
1001 }
1002 LineOperation::Delete { lines } => {
1003 let old_end_row = old_row + lines - 1;
1004 let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
1005
1006 if let Some((_, last_deleted_row_range)) =
1007 self.diff.deleted_row_ranges.last_mut()
1008 {
1009 if *last_deleted_row_range.end() + 1 == old_row {
1010 *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
1011 } else {
1012 self.diff
1013 .deleted_row_ranges
1014 .push((new_row, old_row..=old_end_row));
1015 }
1016 } else {
1017 self.diff
1018 .deleted_row_ranges
1019 .push((new_row, old_row..=old_end_row));
1020 }
1021
1022 old_row += lines;
1023 }
1024 LineOperation::Insert { lines } => {
1025 let new_end_row = new_row + lines - 1;
1026 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
1027 let end = new_snapshot.anchor_before(Point::new(
1028 new_end_row,
1029 new_snapshot.line_len(MultiBufferRow(new_end_row)),
1030 ));
1031 self.diff.inserted_row_ranges.push(start..end);
1032 new_row += lines;
1033 }
1034 }
1035
1036 cx.notify();
1037 }
1038 }
1039
1040 fn reapply_batch_diff(&mut self, cx: &mut Context<Self>) -> Task<()> {
1041 let old_snapshot = self.snapshot.clone();
1042 let old_range = self.range.to_point(&old_snapshot);
1043 let new_snapshot = self.buffer.read(cx).snapshot(cx);
1044 let new_range = self.range.to_point(&new_snapshot);
1045
1046 cx.spawn(async move |codegen, cx| {
1047 let (deleted_row_ranges, inserted_row_ranges) = cx
1048 .background_spawn(async move {
1049 let old_text = old_snapshot
1050 .text_for_range(
1051 Point::new(old_range.start.row, 0)
1052 ..Point::new(
1053 old_range.end.row,
1054 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
1055 ),
1056 )
1057 .collect::<String>();
1058 let new_text = new_snapshot
1059 .text_for_range(
1060 Point::new(new_range.start.row, 0)
1061 ..Point::new(
1062 new_range.end.row,
1063 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
1064 ),
1065 )
1066 .collect::<String>();
1067
1068 let old_start_row = old_range.start.row;
1069 let new_start_row = new_range.start.row;
1070 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
1071 let mut inserted_row_ranges = Vec::new();
1072 for (old_rows, new_rows) in line_diff(&old_text, &new_text) {
1073 let old_rows = old_start_row + old_rows.start..old_start_row + old_rows.end;
1074 let new_rows = new_start_row + new_rows.start..new_start_row + new_rows.end;
1075 if !old_rows.is_empty() {
1076 deleted_row_ranges.push((
1077 new_snapshot.anchor_before(Point::new(new_rows.start, 0)),
1078 old_rows.start..=old_rows.end - 1,
1079 ));
1080 }
1081 if !new_rows.is_empty() {
1082 let start = new_snapshot.anchor_before(Point::new(new_rows.start, 0));
1083 let new_end_row = new_rows.end - 1;
1084 let end = new_snapshot.anchor_before(Point::new(
1085 new_end_row,
1086 new_snapshot.line_len(MultiBufferRow(new_end_row)),
1087 ));
1088 inserted_row_ranges.push(start..end);
1089 }
1090 }
1091 (deleted_row_ranges, inserted_row_ranges)
1092 })
1093 .await;
1094
1095 codegen
1096 .update(cx, |codegen, cx| {
1097 codegen.diff.deleted_row_ranges = deleted_row_ranges;
1098 codegen.diff.inserted_row_ranges = inserted_row_ranges;
1099 cx.notify();
1100 })
1101 .ok();
1102 })
1103 }
1104
1105 fn handle_completion(
1106 &mut self,
1107 model: Arc<dyn LanguageModel>,
1108 completion_stream: Task<
1109 Result<
1110 BoxStream<
1111 'static,
1112 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
1113 >,
1114 LanguageModelCompletionError,
1115 >,
1116 >,
1117 cx: &mut Context<Self>,
1118 ) -> Task<()> {
1119 self.diff = Diff::default();
1120 self.status = CodegenStatus::Pending;
1121
1122 cx.notify();
1123 // Leaving this in generation so that STOP equivalent events are respected even
1124 // while we're still pre-processing the completion event
1125 cx.spawn(async move |codegen, cx| {
1126 let finish_with_status = |status: CodegenStatus, cx: &mut AsyncApp| {
1127 let _ = codegen.update(cx, |this, cx| {
1128 this.status = status;
1129 cx.emit(CodegenEvent::Finished);
1130 cx.notify();
1131 });
1132 };
1133
1134 let mut completion_events = match completion_stream.await {
1135 Ok(events) => events,
1136 Err(err) => {
1137 finish_with_status(CodegenStatus::Error(err.into()), cx);
1138 return;
1139 }
1140 };
1141
1142 enum ToolUseOutput {
1143 Rewrite {
1144 text: String,
1145 description: Option<String>,
1146 },
1147 Failure(String),
1148 }
1149
1150 enum ModelUpdate {
1151 Description(String),
1152 Failure(String),
1153 }
1154
1155 let chars_read_so_far = Arc::new(Mutex::new(0usize));
1156 let process_tool_use = move |tool_use: LanguageModelToolUse| -> Option<ToolUseOutput> {
1157 let mut chars_read_so_far = chars_read_so_far.lock();
1158 match tool_use.name.as_ref() {
1159 "rewrite_section" => {
1160 let Ok(input) =
1161 serde_json::from_value::<RewriteSectionInput>(tool_use.input)
1162 else {
1163 return None;
1164 };
1165 let text = input.replacement_text[*chars_read_so_far..].to_string();
1166 *chars_read_so_far = input.replacement_text.len();
1167 Some(ToolUseOutput::Rewrite {
1168 text,
1169 description: None,
1170 })
1171 }
1172 "failure_message" => {
1173 let Ok(mut input) =
1174 serde_json::from_value::<FailureMessageInput>(tool_use.input)
1175 else {
1176 return None;
1177 };
1178 Some(ToolUseOutput::Failure(std::mem::take(&mut input.message)))
1179 }
1180 _ => None,
1181 }
1182 };
1183
1184 let (message_tx, mut message_rx) = futures::channel::mpsc::unbounded::<ModelUpdate>();
1185
1186 cx.spawn({
1187 let codegen = codegen.clone();
1188 async move |cx| {
1189 while let Some(update) = message_rx.next().await {
1190 let _ = codegen.update(cx, |this, _cx| match update {
1191 ModelUpdate::Description(d) => this.description = Some(d),
1192 ModelUpdate::Failure(f) => this.failure = Some(f),
1193 });
1194 }
1195 }
1196 })
1197 .detach();
1198
1199 let mut message_id = None;
1200 let mut first_text = None;
1201 let last_token_usage = Arc::new(Mutex::new(TokenUsage::default()));
1202 let total_text = Arc::new(Mutex::new(String::new()));
1203
1204 loop {
1205 if let Some(first_event) = completion_events.next().await {
1206 match first_event {
1207 Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => {
1208 message_id = Some(id);
1209 }
1210 Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
1211 if let Some(output) = process_tool_use(tool_use) {
1212 let (text, update) = match output {
1213 ToolUseOutput::Rewrite { text, description } => {
1214 (Some(text), description.map(ModelUpdate::Description))
1215 }
1216 ToolUseOutput::Failure(message) => {
1217 (None, Some(ModelUpdate::Failure(message)))
1218 }
1219 };
1220 if let Some(update) = update {
1221 let _ = message_tx.unbounded_send(update);
1222 }
1223 first_text = text;
1224 if first_text.is_some() {
1225 break;
1226 }
1227 }
1228 }
1229 Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
1230 *last_token_usage.lock() = token_usage;
1231 }
1232 Ok(LanguageModelCompletionEvent::Text(text)) => {
1233 let mut lock = total_text.lock();
1234 lock.push_str(&text);
1235 }
1236 Ok(e) => {
1237 log::warn!("Unexpected event: {:?}", e);
1238 break;
1239 }
1240 Err(e) => {
1241 finish_with_status(CodegenStatus::Error(e.into()), cx);
1242 break;
1243 }
1244 }
1245 }
1246 }
1247
1248 let Some(first_text) = first_text else {
1249 finish_with_status(CodegenStatus::Done, cx);
1250 return;
1251 };
1252
1253 let move_last_token_usage = last_token_usage.clone();
1254
1255 let text_stream = Box::pin(futures::stream::once(async { Ok(first_text) }).chain(
1256 completion_events.filter_map(move |e| {
1257 let process_tool_use = process_tool_use.clone();
1258 let last_token_usage = move_last_token_usage.clone();
1259 let total_text = total_text.clone();
1260 let mut message_tx = message_tx.clone();
1261 async move {
1262 match e {
1263 Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
1264 let Some(output) = process_tool_use(tool_use) else {
1265 return None;
1266 };
1267 let (text, update) = match output {
1268 ToolUseOutput::Rewrite { text, description } => {
1269 (Some(text), description.map(ModelUpdate::Description))
1270 }
1271 ToolUseOutput::Failure(message) => {
1272 (None, Some(ModelUpdate::Failure(message)))
1273 }
1274 };
1275 if let Some(update) = update {
1276 let _ = message_tx.send(update).await;
1277 }
1278 text.map(Ok)
1279 }
1280 Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
1281 *last_token_usage.lock() = token_usage;
1282 None
1283 }
1284 Ok(LanguageModelCompletionEvent::Text(text)) => {
1285 let mut lock = total_text.lock();
1286 lock.push_str(&text);
1287 None
1288 }
1289 Ok(LanguageModelCompletionEvent::Stop(_reason)) => None,
1290 e => {
1291 log::error!("UNEXPECTED EVENT {:?}", e);
1292 None
1293 }
1294 }
1295 }
1296 }),
1297 ));
1298
1299 let language_model_text_stream = LanguageModelTextStream {
1300 message_id: message_id,
1301 stream: text_stream,
1302 last_token_usage,
1303 };
1304
1305 let Some(task) = codegen
1306 .update(cx, move |codegen, cx| {
1307 codegen.handle_stream(model, async { Ok(language_model_text_stream) }, cx)
1308 })
1309 .ok()
1310 else {
1311 return;
1312 };
1313
1314 task.await;
1315 })
1316 }
1317}
1318
1319#[derive(Copy, Clone, Debug)]
1320pub enum CodegenEvent {
1321 Finished,
1322 Undone,
1323}
1324
1325struct StripInvalidSpans<T> {
1326 stream: T,
1327 stream_done: bool,
1328 buffer: String,
1329 first_line: bool,
1330 line_end: bool,
1331 starts_with_code_block: bool,
1332}
1333
1334impl<T> StripInvalidSpans<T>
1335where
1336 T: Stream<Item = Result<String>>,
1337{
1338 fn new(stream: T) -> Self {
1339 Self {
1340 stream,
1341 stream_done: false,
1342 buffer: String::new(),
1343 first_line: true,
1344 line_end: false,
1345 starts_with_code_block: false,
1346 }
1347 }
1348}
1349
1350impl<T> Stream for StripInvalidSpans<T>
1351where
1352 T: Stream<Item = Result<String>>,
1353{
1354 type Item = Result<String>;
1355
1356 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
1357 const CODE_BLOCK_DELIMITER: &str = "```";
1358 const CURSOR_SPAN: &str = "<|CURSOR|>";
1359
1360 let this = unsafe { self.get_unchecked_mut() };
1361 loop {
1362 if !this.stream_done {
1363 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
1364 match stream.as_mut().poll_next(cx) {
1365 Poll::Ready(Some(Ok(chunk))) => {
1366 this.buffer.push_str(&chunk);
1367 }
1368 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
1369 Poll::Ready(None) => {
1370 this.stream_done = true;
1371 }
1372 Poll::Pending => return Poll::Pending,
1373 }
1374 }
1375
1376 let mut chunk = String::new();
1377 let mut consumed = 0;
1378 if !this.buffer.is_empty() {
1379 let mut lines = this.buffer.split('\n').enumerate().peekable();
1380 while let Some((line_ix, line)) = lines.next() {
1381 if line_ix > 0 {
1382 this.first_line = false;
1383 }
1384
1385 if this.first_line {
1386 let trimmed_line = line.trim();
1387 if lines.peek().is_some() {
1388 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
1389 consumed += line.len() + 1;
1390 this.starts_with_code_block = true;
1391 continue;
1392 }
1393 } else if trimmed_line.is_empty()
1394 || prefixes(CODE_BLOCK_DELIMITER)
1395 .any(|prefix| trimmed_line.starts_with(prefix))
1396 {
1397 break;
1398 }
1399 }
1400
1401 let line_without_cursor = line.replace(CURSOR_SPAN, "");
1402 if lines.peek().is_some() {
1403 if this.line_end {
1404 chunk.push('\n');
1405 }
1406
1407 chunk.push_str(&line_without_cursor);
1408 this.line_end = true;
1409 consumed += line.len() + 1;
1410 } else if this.stream_done {
1411 if !this.starts_with_code_block
1412 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
1413 {
1414 if this.line_end {
1415 chunk.push('\n');
1416 }
1417
1418 chunk.push_str(line);
1419 }
1420
1421 consumed += line.len();
1422 } else {
1423 let trimmed_line = line.trim();
1424 if trimmed_line.is_empty()
1425 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
1426 || prefixes(CODE_BLOCK_DELIMITER)
1427 .any(|prefix| trimmed_line.ends_with(prefix))
1428 {
1429 break;
1430 } else {
1431 if this.line_end {
1432 chunk.push('\n');
1433 this.line_end = false;
1434 }
1435
1436 chunk.push_str(&line_without_cursor);
1437 consumed += line.len();
1438 }
1439 }
1440 }
1441 }
1442
1443 this.buffer = this.buffer.split_off(consumed);
1444 if !chunk.is_empty() {
1445 return Poll::Ready(Some(Ok(chunk)));
1446 } else if this.stream_done {
1447 return Poll::Ready(None);
1448 }
1449 }
1450 }
1451}
1452
1453fn prefixes(text: &str) -> impl Iterator<Item = &str> {
1454 (0..text.len() - 1).map(|ix| &text[..ix + 1])
1455}
1456
1457#[derive(Default)]
1458pub struct Diff {
1459 pub deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1460 pub inserted_row_ranges: Vec<Range<Anchor>>,
1461}
1462
1463impl Diff {
1464 fn is_empty(&self) -> bool {
1465 self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
1466 }
1467}
1468
1469#[cfg(test)]
1470mod tests {
1471 use super::*;
1472 use futures::{
1473 Stream,
1474 stream::{self},
1475 };
1476 use gpui::TestAppContext;
1477 use indoc::indoc;
1478 use language::{Buffer, Point};
1479 use language_model::fake_provider::FakeLanguageModel;
1480 use language_model::{LanguageModelRegistry, TokenUsage};
1481 use languages::rust_lang;
1482 use rand::prelude::*;
1483 use settings::SettingsStore;
1484 use std::{future, sync::Arc};
1485
1486 #[gpui::test(iterations = 10)]
1487 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
1488 init_test(cx);
1489
1490 let text = indoc! {"
1491 fn main() {
1492 let x = 0;
1493 for _ in 0..10 {
1494 x += 1;
1495 }
1496 }
1497 "};
1498 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1499 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1500 let range = buffer.read_with(cx, |buffer, cx| {
1501 let snapshot = buffer.snapshot(cx);
1502 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
1503 });
1504 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1505 let codegen = cx.new(|cx| {
1506 CodegenAlternative::new(
1507 buffer.clone(),
1508 range.clone(),
1509 true,
1510 prompt_builder,
1511 Uuid::new_v4(),
1512 cx,
1513 )
1514 });
1515
1516 let chunks_tx = simulate_response_stream(&codegen, cx);
1517
1518 let mut new_text = concat!(
1519 " let mut x = 0;\n",
1520 " while x < 10 {\n",
1521 " x += 1;\n",
1522 " }",
1523 );
1524 while !new_text.is_empty() {
1525 let max_len = cmp::min(new_text.len(), 10);
1526 let len = rng.random_range(1..=max_len);
1527 let (chunk, suffix) = new_text.split_at(len);
1528 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1529 new_text = suffix;
1530 cx.background_executor.run_until_parked();
1531 }
1532 drop(chunks_tx);
1533 cx.background_executor.run_until_parked();
1534
1535 assert_eq!(
1536 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1537 indoc! {"
1538 fn main() {
1539 let mut x = 0;
1540 while x < 10 {
1541 x += 1;
1542 }
1543 }
1544 "}
1545 );
1546 }
1547
1548 #[gpui::test(iterations = 10)]
1549 async fn test_autoindent_when_generating_past_indentation(
1550 cx: &mut TestAppContext,
1551 mut rng: StdRng,
1552 ) {
1553 init_test(cx);
1554
1555 let text = indoc! {"
1556 fn main() {
1557 le
1558 }
1559 "};
1560 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1561 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1562 let range = buffer.read_with(cx, |buffer, cx| {
1563 let snapshot = buffer.snapshot(cx);
1564 snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
1565 });
1566 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1567 let codegen = cx.new(|cx| {
1568 CodegenAlternative::new(
1569 buffer.clone(),
1570 range.clone(),
1571 true,
1572 prompt_builder,
1573 Uuid::new_v4(),
1574 cx,
1575 )
1576 });
1577
1578 let chunks_tx = simulate_response_stream(&codegen, cx);
1579
1580 cx.background_executor.run_until_parked();
1581
1582 let mut new_text = concat!(
1583 "t mut x = 0;\n",
1584 "while x < 10 {\n",
1585 " x += 1;\n",
1586 "}", //
1587 );
1588 while !new_text.is_empty() {
1589 let max_len = cmp::min(new_text.len(), 10);
1590 let len = rng.random_range(1..=max_len);
1591 let (chunk, suffix) = new_text.split_at(len);
1592 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1593 new_text = suffix;
1594 cx.background_executor.run_until_parked();
1595 }
1596 drop(chunks_tx);
1597 cx.background_executor.run_until_parked();
1598
1599 assert_eq!(
1600 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1601 indoc! {"
1602 fn main() {
1603 let mut x = 0;
1604 while x < 10 {
1605 x += 1;
1606 }
1607 }
1608 "}
1609 );
1610 }
1611
1612 #[gpui::test(iterations = 10)]
1613 async fn test_autoindent_when_generating_before_indentation(
1614 cx: &mut TestAppContext,
1615 mut rng: StdRng,
1616 ) {
1617 init_test(cx);
1618
1619 let text = concat!(
1620 "fn main() {\n",
1621 " \n",
1622 "}\n" //
1623 );
1624 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1625 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1626 let range = buffer.read_with(cx, |buffer, cx| {
1627 let snapshot = buffer.snapshot(cx);
1628 snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
1629 });
1630 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1631 let codegen = cx.new(|cx| {
1632 CodegenAlternative::new(
1633 buffer.clone(),
1634 range.clone(),
1635 true,
1636 prompt_builder,
1637 Uuid::new_v4(),
1638 cx,
1639 )
1640 });
1641
1642 let chunks_tx = simulate_response_stream(&codegen, cx);
1643
1644 cx.background_executor.run_until_parked();
1645
1646 let mut new_text = concat!(
1647 "let mut x = 0;\n",
1648 "while x < 10 {\n",
1649 " x += 1;\n",
1650 "}", //
1651 );
1652 while !new_text.is_empty() {
1653 let max_len = cmp::min(new_text.len(), 10);
1654 let len = rng.random_range(1..=max_len);
1655 let (chunk, suffix) = new_text.split_at(len);
1656 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1657 new_text = suffix;
1658 cx.background_executor.run_until_parked();
1659 }
1660 drop(chunks_tx);
1661 cx.background_executor.run_until_parked();
1662
1663 assert_eq!(
1664 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1665 indoc! {"
1666 fn main() {
1667 let mut x = 0;
1668 while x < 10 {
1669 x += 1;
1670 }
1671 }
1672 "}
1673 );
1674 }
1675
1676 #[gpui::test(iterations = 10)]
1677 async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
1678 init_test(cx);
1679
1680 let text = indoc! {"
1681 func main() {
1682 \tx := 0
1683 \tfor i := 0; i < 10; i++ {
1684 \t\tx++
1685 \t}
1686 }
1687 "};
1688 let buffer = cx.new(|cx| Buffer::local(text, cx));
1689 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1690 let range = buffer.read_with(cx, |buffer, cx| {
1691 let snapshot = buffer.snapshot(cx);
1692 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
1693 });
1694 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1695 let codegen = cx.new(|cx| {
1696 CodegenAlternative::new(
1697 buffer.clone(),
1698 range.clone(),
1699 true,
1700 prompt_builder,
1701 Uuid::new_v4(),
1702 cx,
1703 )
1704 });
1705
1706 let chunks_tx = simulate_response_stream(&codegen, cx);
1707 let new_text = concat!(
1708 "func main() {\n",
1709 "\tx := 0\n",
1710 "\tfor x < 10 {\n",
1711 "\t\tx++\n",
1712 "\t}", //
1713 );
1714 chunks_tx.unbounded_send(new_text.to_string()).unwrap();
1715 drop(chunks_tx);
1716 cx.background_executor.run_until_parked();
1717
1718 assert_eq!(
1719 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1720 indoc! {"
1721 func main() {
1722 \tx := 0
1723 \tfor x < 10 {
1724 \t\tx++
1725 \t}
1726 }
1727 "}
1728 );
1729 }
1730
1731 #[gpui::test]
1732 async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
1733 init_test(cx);
1734
1735 let text = indoc! {"
1736 fn main() {
1737 let x = 0;
1738 }
1739 "};
1740 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1741 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1742 let range = buffer.read_with(cx, |buffer, cx| {
1743 let snapshot = buffer.snapshot(cx);
1744 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
1745 });
1746 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1747 let codegen = cx.new(|cx| {
1748 CodegenAlternative::new(
1749 buffer.clone(),
1750 range.clone(),
1751 false,
1752 prompt_builder,
1753 Uuid::new_v4(),
1754 cx,
1755 )
1756 });
1757
1758 let chunks_tx = simulate_response_stream(&codegen, cx);
1759 chunks_tx
1760 .unbounded_send("let mut x = 0;\nx += 1;".to_string())
1761 .unwrap();
1762 drop(chunks_tx);
1763 cx.run_until_parked();
1764
1765 // The codegen is inactive, so the buffer doesn't get modified.
1766 assert_eq!(
1767 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1768 text
1769 );
1770
1771 // Activating the codegen applies the changes.
1772 codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
1773 assert_eq!(
1774 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1775 indoc! {"
1776 fn main() {
1777 let mut x = 0;
1778 x += 1;
1779 }
1780 "}
1781 );
1782
1783 // Deactivating the codegen undoes the changes.
1784 codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
1785 cx.run_until_parked();
1786 assert_eq!(
1787 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1788 text
1789 );
1790 }
1791
1792 #[gpui::test]
1793 async fn test_strip_invalid_spans_from_codeblock() {
1794 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
1795 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
1796 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
1797 assert_chunks(
1798 "```html\n```js\nLorem ipsum dolor\n```\n```",
1799 "```js\nLorem ipsum dolor\n```",
1800 )
1801 .await;
1802 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
1803 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
1804 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
1805 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
1806
1807 async fn assert_chunks(text: &str, expected_text: &str) {
1808 for chunk_size in 1..=text.len() {
1809 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
1810 .map(|chunk| chunk.unwrap())
1811 .collect::<String>()
1812 .await;
1813 assert_eq!(
1814 actual_text, expected_text,
1815 "failed to strip invalid spans, chunk size: {}",
1816 chunk_size
1817 );
1818 }
1819 }
1820
1821 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
1822 stream::iter(
1823 text.chars()
1824 .collect::<Vec<_>>()
1825 .chunks(size)
1826 .map(|chunk| Ok(chunk.iter().collect::<String>()))
1827 .collect::<Vec<_>>(),
1828 )
1829 }
1830 }
1831
1832 fn init_test(cx: &mut TestAppContext) {
1833 cx.update(LanguageModelRegistry::test);
1834 cx.set_global(cx.update(SettingsStore::test));
1835 }
1836
1837 fn simulate_response_stream(
1838 codegen: &Entity<CodegenAlternative>,
1839 cx: &mut TestAppContext,
1840 ) -> mpsc::UnboundedSender<String> {
1841 let (chunks_tx, chunks_rx) = mpsc::unbounded();
1842 let model = Arc::new(FakeLanguageModel::default());
1843 codegen.update(cx, |codegen, cx| {
1844 codegen.generation = codegen.handle_stream(
1845 model,
1846 future::ready(Ok(LanguageModelTextStream {
1847 message_id: None,
1848 stream: chunks_rx.map(Ok).boxed(),
1849 last_token_usage: Arc::new(Mutex::new(TokenUsage::default())),
1850 })),
1851 cx,
1852 );
1853 });
1854 chunks_tx
1855 }
1856}