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