Skip to repository content2018 lines · 78.0 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T01:48:41.839Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
buffer_codegen.rs
1use crate::{context::LoadedContext, inline_prompt_editor::CodegenStatus};
2use agent_settings::AgentSettings;
3use anyhow::{Context as _, Result};
4use collections::{HashMap, HashSet};
5use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};
6use futures::FutureExt;
7use futures::{
8 SinkExt, Stream, StreamExt, TryStreamExt as _,
9 channel::mpsc,
10 future::{LocalBoxFuture, Shared},
11 join,
12 stream::BoxStream,
13};
14use gpui::{App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Subscription, Task};
15use language::{
16 Buffer, BufferEditSource, IndentKind, LanguageName, Point, TransactionId, line_diff,
17};
18use language_model::{
19 CompletionIntent, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
20 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
21 LanguageModelRequestTool, LanguageModelTextStream, LanguageModelToolChoice,
22 LanguageModelToolUse, LanguageModelToolUseId, Role, TokenUsage,
23};
24use language_models::provider::anthropic::telemetry::{
25 AnthropicCompletionType, AnthropicEventData, AnthropicEventReporter, AnthropicEventType,
26};
27use multi_buffer::MultiBufferRow;
28use parking_lot::Mutex;
29use prompt_store::PromptBuilder;
30use rope::Rope;
31use schemars::JsonSchema;
32use serde::{Deserialize, Serialize};
33use settings::Settings as _;
34use std::{
35 cmp,
36 future::Future,
37 iter,
38 ops::{Range, RangeInclusive},
39 pin::Pin,
40 sync::Arc,
41 task::{self, Poll},
42 time::Instant,
43};
44use streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff};
45use uuid::Uuid;
46
47/// Use this tool when you cannot or should not make a rewrite. This includes:
48/// - The user's request is unclear, ambiguous, or nonsensical
49/// - The requested change cannot be made by only editing the <rewrite_this> section
50#[derive(Debug, Serialize, Deserialize, JsonSchema)]
51pub struct FailureMessageInput {
52 /// A brief message to the user explaining why you're unable to fulfill the request or to ask a question about the request.
53 #[serde(default)]
54 pub message: String,
55}
56
57/// Replaces text in <rewrite_this></rewrite_this> tags with your replacement_text.
58/// Only use this tool when you are confident you understand the user's request and can fulfill it
59/// by editing the marked section.
60#[derive(Debug, Serialize, Deserialize, JsonSchema)]
61pub struct RewriteSectionInput {
62 /// The text to replace the section with.
63 #[serde(default)]
64 pub replacement_text: String,
65}
66
67pub struct BufferCodegen {
68 alternatives: Vec<Entity<CodegenAlternative>>,
69 pub active_alternative: usize,
70 seen_alternatives: HashSet<usize>,
71 subscriptions: Vec<Subscription>,
72 buffer: Entity<MultiBuffer>,
73 range: Range<Anchor>,
74 initial_transaction_id: Option<TransactionId>,
75 builder: Arc<PromptBuilder>,
76 pub is_insertion: bool,
77 session_id: Uuid,
78}
79
80pub const REWRITE_SECTION_TOOL_NAME: &str = "rewrite_section";
81pub const FAILURE_MESSAGE_TOOL_NAME: &str = "failure_message";
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.start..range.end)
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 pub fn use_streaming_tools(model: &dyn LanguageModel, cx: &App) -> bool {
405 model.supports_streaming_tools()
406 && AgentSettings::get_global(cx).inline_assistant_use_streaming_tools
407 }
408
409 pub fn start(
410 &mut self,
411 user_prompt: String,
412 context_task: Shared<Task<Option<LoadedContext>>>,
413 model: Arc<dyn LanguageModel>,
414 cx: &mut Context<Self>,
415 ) -> Result<()> {
416 // Clear the model explanation since the user has started a new generation.
417 self.description = None;
418
419 if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {
420 self.buffer.update(cx, |buffer, cx| {
421 buffer.undo_transaction(transformation_transaction_id, cx);
422 });
423 }
424
425 self.edit_position = Some(self.range.start.bias_right(&self.snapshot));
426
427 if Self::use_streaming_tools(model.as_ref(), cx) {
428 let request = self.build_request(&model, user_prompt, context_task, cx)?;
429 let completion_events = cx.spawn({
430 let model = model.clone();
431 async move |_, cx| model.stream_completion(request.await, cx).await
432 });
433 self.generation = self.handle_completion(model, completion_events, cx);
434 } else {
435 let stream: LocalBoxFuture<Result<LanguageModelTextStream>> =
436 if user_prompt.trim().to_lowercase() == "delete" {
437 async { Ok(LanguageModelTextStream::default()) }.boxed_local()
438 } else {
439 let request = self.build_request(&model, user_prompt, context_task, cx)?;
440 cx.spawn({
441 let model = model.clone();
442 async move |_, cx| {
443 Ok(model.stream_completion_text(request.await, cx).await?)
444 }
445 })
446 .boxed_local()
447 };
448 self.generation =
449 self.handle_stream(model, /* strip_invalid_spans: */ true, stream, cx);
450 }
451
452 Ok(())
453 }
454
455 fn build_request_tools(
456 &self,
457 model: &Arc<dyn LanguageModel>,
458 user_prompt: String,
459 context_task: Shared<Task<Option<LoadedContext>>>,
460 cx: &mut App,
461 ) -> Result<Task<LanguageModelRequest>> {
462 let buffer = self.buffer.read(cx).snapshot(cx);
463 let language = buffer.language_at(self.range.start);
464 let language_name = if let Some(language) = language.as_ref() {
465 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
466 None
467 } else {
468 Some(language.name())
469 }
470 } else {
471 None
472 };
473
474 let language_name = language_name.as_ref();
475 let start = buffer.point_to_buffer_offset(self.range.start);
476 let end = buffer.point_to_buffer_offset(self.range.end);
477 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
478 let (start_buffer, start_buffer_offset) = start;
479 let (end_buffer, end_buffer_offset) = end;
480 if start_buffer.remote_id() == end_buffer.remote_id() {
481 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
482 } else {
483 anyhow::bail!("invalid transformation range");
484 }
485 } else {
486 anyhow::bail!("invalid transformation range");
487 };
488
489 let system_prompt = self
490 .builder
491 .generate_inline_transformation_prompt_tools(
492 language_name,
493 buffer,
494 range.start.0..range.end.0,
495 )
496 .context("generating content prompt")?;
497
498 let temperature = AgentSettings::temperature_for_model(model, cx);
499
500 let tool_input_format = model.tool_input_format();
501 let tool_choice = model
502 .supports_tool_choice(LanguageModelToolChoice::Any)
503 .then_some(LanguageModelToolChoice::Any);
504
505 Ok(cx.spawn(async move |_cx| {
506 let mut messages = vec![LanguageModelRequestMessage {
507 role: Role::System,
508 content: vec![system_prompt.into()],
509 cache: false,
510 reasoning_details: None,
511 }];
512
513 let mut user_message = LanguageModelRequestMessage {
514 role: Role::User,
515 content: Vec::new(),
516 cache: false,
517 reasoning_details: None,
518 };
519
520 if let Some(context) = context_task.await {
521 context.add_to_request_message(&mut user_message);
522 }
523
524 user_message.content.push(user_prompt.into());
525 messages.push(user_message);
526
527 let tools = vec![
528 LanguageModelRequestTool::function(
529 REWRITE_SECTION_TOOL_NAME.to_string(),
530 "Replaces text in <rewrite_this></rewrite_this> tags with your replacement_text.".to_string(),
531 language_model::tool_schema::root_schema_for::<RewriteSectionInput>(tool_input_format).to_value(),
532 false,
533 ),
534 LanguageModelRequestTool::function(
535 FAILURE_MESSAGE_TOOL_NAME.to_string(),
536 "Use this tool to provide a message to the user when you're unable to complete a task.".to_string(),
537 language_model::tool_schema::root_schema_for::<FailureMessageInput>(tool_input_format).to_value(),
538 false,
539 ),
540 ];
541
542 LanguageModelRequest {
543 thread_id: None,
544 prompt_id: None,
545 intent: Some(CompletionIntent::InlineAssist),
546 tools,
547 tool_choice,
548 stop: Vec::new(),
549 temperature,
550 messages,
551 thinking_allowed: false,
552 thinking_effort: None,
553 speed: None,
554 compact_at_tokens: None,
555 }
556 }))
557 }
558
559 fn build_request(
560 &self,
561 model: &Arc<dyn LanguageModel>,
562 user_prompt: String,
563 context_task: Shared<Task<Option<LoadedContext>>>,
564 cx: &mut App,
565 ) -> Result<Task<LanguageModelRequest>> {
566 if Self::use_streaming_tools(model.as_ref(), cx) {
567 return self.build_request_tools(model, user_prompt, context_task, cx);
568 }
569
570 let buffer = self.buffer.read(cx).snapshot(cx);
571 let language = buffer.language_at(self.range.start);
572 let language_name = if let Some(language) = language.as_ref() {
573 if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
574 None
575 } else {
576 Some(language.name())
577 }
578 } else {
579 None
580 };
581
582 let language_name = language_name.as_ref();
583 let start = buffer.point_to_buffer_offset(self.range.start);
584 let end = buffer.point_to_buffer_offset(self.range.end);
585 let (buffer, range) = if let Some((start, end)) = start.zip(end) {
586 let (start_buffer, start_buffer_offset) = start;
587 let (end_buffer, end_buffer_offset) = end;
588 if start_buffer.remote_id() == end_buffer.remote_id() {
589 (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
590 } else {
591 anyhow::bail!("invalid transformation range");
592 }
593 } else {
594 anyhow::bail!("invalid transformation range");
595 };
596
597 let prompt = self
598 .builder
599 .generate_inline_transformation_prompt(
600 user_prompt,
601 language_name,
602 buffer,
603 range.start.0..range.end.0,
604 )
605 .context("generating content prompt")?;
606
607 let temperature = AgentSettings::temperature_for_model(model, cx);
608
609 Ok(cx.spawn(async move |_cx| {
610 let mut request_message = LanguageModelRequestMessage {
611 role: Role::User,
612 content: Vec::new(),
613 cache: false,
614 reasoning_details: None,
615 };
616
617 if let Some(context) = context_task.await {
618 context.add_to_request_message(&mut request_message);
619 }
620
621 request_message.content.push(prompt.into());
622
623 LanguageModelRequest {
624 thread_id: None,
625 prompt_id: None,
626 intent: Some(CompletionIntent::InlineAssist),
627 tools: Vec::new(),
628 tool_choice: None,
629 stop: Vec::new(),
630 temperature,
631 messages: vec![request_message],
632 thinking_allowed: false,
633 thinking_effort: None,
634 speed: None,
635 compact_at_tokens: None,
636 }
637 }))
638 }
639
640 pub fn handle_stream(
641 &mut self,
642 model: Arc<dyn LanguageModel>,
643 strip_invalid_spans: bool,
644 stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,
645 cx: &mut Context<Self>,
646 ) -> Task<()> {
647 let anthropic_reporter = AnthropicEventReporter::new(&model, cx);
648 let session_id = self.session_id;
649 let model_telemetry_id = model.telemetry_id();
650 let model_provider_id = model.provider_id().to_string();
651 let start_time = Instant::now();
652
653 // Make a new snapshot and re-resolve anchor in case the document was modified.
654 // This can happen often if the editor loses focus and is saved + reformatted,
655 // as in https://github.com/zed-industries/zed/issues/39088
656 self.snapshot = self.buffer.read(cx).snapshot(cx);
657 self.range = self.snapshot.anchor_after(self.range.start)
658 ..self.snapshot.anchor_after(self.range.end);
659
660 let snapshot = self.snapshot.clone();
661 let selected_text = snapshot
662 .text_for_range(self.range.start..self.range.end)
663 .collect::<Rope>();
664
665 self.selected_text = Some(selected_text.to_string());
666
667 let selection_start = self.range.start.to_point(&snapshot);
668
669 // Start with the indentation of the first line in the selection
670 let mut suggested_line_indent = snapshot
671 .suggested_indents(selection_start.row..=selection_start.row, cx)
672 .into_values()
673 .next()
674 .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
675
676 // If the first line in the selection does not have indentation, check the following lines
677 if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
678 for row in selection_start.row..=self.range.end.to_point(&snapshot).row {
679 let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
680 // Prefer tabs if a line in the selection uses tabs as indentation
681 if line_indent.kind == IndentKind::Tab {
682 suggested_line_indent.kind = IndentKind::Tab;
683 break;
684 }
685 }
686 }
687
688 let language_name = {
689 let multibuffer = self.buffer.read(cx);
690 let snapshot = multibuffer.snapshot(cx);
691 let ranges = snapshot.range_to_buffer_ranges(self.range.start..self.range.end);
692 ranges
693 .first()
694 .and_then(|(buffer, _, _)| buffer.language())
695 .map(|language| language.name())
696 };
697
698 self.diff = Diff::default();
699 self.status = CodegenStatus::Pending;
700 let mut edit_start = self.range.start.to_offset(&snapshot);
701 let completion = Arc::new(Mutex::new(String::new()));
702 let completion_clone = completion.clone();
703
704 cx.notify();
705 cx.spawn(async move |codegen, cx| {
706 let stream = stream.await;
707
708 let token_usage = stream
709 .as_ref()
710 .ok()
711 .map(|stream| stream.last_token_usage.clone());
712 let message_id = stream
713 .as_ref()
714 .ok()
715 .and_then(|stream| stream.message_id.clone());
716 let generate = async {
717 let model_telemetry_id = model_telemetry_id.clone();
718 let model_provider_id = model_provider_id.clone();
719 let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
720 let message_id = message_id.clone();
721 let line_based_stream_diff: Task<anyhow::Result<()>> = cx.background_spawn({
722 let anthropic_reporter = anthropic_reporter.clone();
723 let language_name = language_name.clone();
724 async move {
725 let mut response_latency = None;
726 let request_start = Instant::now();
727 let diff = async {
728 let raw_stream = stream?.stream.map_err(|error| error.into());
729
730 let stripped;
731 let mut chunks: Pin<Box<dyn Stream<Item = Result<String>> + Send>> =
732 if strip_invalid_spans {
733 stripped = StripInvalidSpans::new(raw_stream);
734 Box::pin(stripped)
735 } else {
736 Box::pin(raw_stream)
737 };
738
739 let mut diff = StreamingDiff::new(selected_text.to_string());
740 let mut line_diff = LineDiff::default();
741
742 let mut new_text = String::new();
743 let mut base_indent = None;
744 let mut line_indent = None;
745 let mut first_line = true;
746
747 while let Some(chunk) = chunks.next().await {
748 if response_latency.is_none() {
749 response_latency = Some(request_start.elapsed());
750 }
751 let chunk = chunk?;
752 completion_clone.lock().push_str(&chunk);
753
754 let mut lines = chunk.split('\n').peekable();
755 while let Some(line) = lines.next() {
756 new_text.push_str(line);
757 if line_indent.is_none()
758 && let Some(non_whitespace_ch_ix) =
759 new_text.find(|ch: char| !ch.is_whitespace())
760 {
761 line_indent = Some(non_whitespace_ch_ix);
762 base_indent = base_indent.or(line_indent);
763
764 let line_indent = line_indent.unwrap();
765 let base_indent = base_indent.unwrap();
766 let indent_delta = line_indent as i32 - base_indent as i32;
767 let mut corrected_indent_len = cmp::max(
768 0,
769 suggested_line_indent.len as i32 + indent_delta,
770 )
771 as usize;
772 if first_line {
773 corrected_indent_len = corrected_indent_len
774 .saturating_sub(selection_start.column as usize);
775 }
776
777 let indent_char = suggested_line_indent.char();
778 let mut indent_buffer = [0; 4];
779 let indent_str =
780 indent_char.encode_utf8(&mut indent_buffer);
781 new_text.replace_range(
782 ..line_indent,
783 &indent_str.repeat(corrected_indent_len),
784 );
785 }
786
787 if line_indent.is_some() {
788 let char_ops = diff.push_new(&new_text);
789 line_diff.push_char_operations(&char_ops, &selected_text);
790 diff_tx
791 .send((char_ops, line_diff.line_operations()))
792 .await?;
793 new_text.clear();
794 }
795
796 if lines.peek().is_some() {
797 let char_ops = diff.push_new("\n");
798 line_diff.push_char_operations(&char_ops, &selected_text);
799 diff_tx
800 .send((char_ops, line_diff.line_operations()))
801 .await?;
802 if line_indent.is_none() {
803 // Don't write out the leading indentation in empty lines on the next line
804 // This is the case where the above if statement didn't clear the buffer
805 new_text.clear();
806 }
807 line_indent = None;
808 first_line = false;
809 }
810 }
811 }
812
813 let mut char_ops = diff.push_new(&new_text);
814 char_ops.extend(diff.finish());
815 line_diff.push_char_operations(&char_ops, &selected_text);
816 line_diff.finish(&selected_text);
817 diff_tx
818 .send((char_ops, line_diff.line_operations()))
819 .await?;
820
821 anyhow::Ok(())
822 };
823
824 let result = diff.await;
825
826 let error_message = result.as_ref().err().map(|error| error.to_string());
827 telemetry::event!(
828 "Assistant Responded",
829 kind = "inline",
830 phase = "response",
831 session_id = session_id.to_string(),
832 model = model_telemetry_id,
833 model_provider = model_provider_id,
834 language_name = language_name.as_ref().map(|n| n.to_string()),
835 message_id = message_id.as_deref(),
836 response_latency = response_latency,
837 error_message = error_message.as_deref(),
838 );
839
840 anthropic_reporter.report(AnthropicEventData {
841 completion_type: AnthropicCompletionType::Editor,
842 event: AnthropicEventType::Response,
843 language_name: language_name.map(|n| n.to_string()),
844 message_id,
845 });
846
847 result?;
848 Ok(())
849 }
850 });
851
852 while let Some((char_ops, line_ops)) = diff_rx.next().await {
853 codegen.update(cx, |codegen, cx| {
854 codegen.last_equal_ranges.clear();
855
856 let edits = char_ops
857 .into_iter()
858 .filter_map(|operation| match operation {
859 CharOperation::Insert { text } => {
860 let edit_start = snapshot.anchor_after(edit_start);
861 Some((edit_start..edit_start, text))
862 }
863 CharOperation::Delete { bytes } => {
864 let edit_end = edit_start + bytes;
865 let edit_range = snapshot.anchor_after(edit_start)
866 ..snapshot.anchor_before(edit_end);
867 edit_start = edit_end;
868 Some((edit_range, String::new()))
869 }
870 CharOperation::Keep { bytes } => {
871 let edit_end = edit_start + bytes;
872 let edit_range = snapshot.anchor_after(edit_start)
873 ..snapshot.anchor_before(edit_end);
874 edit_start = edit_end;
875 codegen.last_equal_ranges.push(edit_range);
876 None
877 }
878 })
879 .collect::<Vec<_>>();
880
881 if codegen.active {
882 codegen.apply_edits(edits.iter().cloned(), cx);
883 codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
884 }
885 codegen.edits.extend(edits);
886 codegen.line_operations = line_ops;
887 codegen.edit_position = Some(snapshot.anchor_after(edit_start));
888
889 cx.notify();
890 })?;
891 }
892
893 // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
894 // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
895 // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
896 let batch_diff_task =
897 codegen.update(cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
898 let (line_based_stream_diff, ()) = join!(line_based_stream_diff, batch_diff_task);
899 line_based_stream_diff?;
900
901 anyhow::Ok(())
902 };
903
904 let result = generate.await;
905 let elapsed_time = start_time.elapsed().as_secs_f64();
906
907 codegen
908 .update(cx, |this, cx| {
909 this.message_id = message_id;
910 this.last_equal_ranges.clear();
911 if let Err(error) = result {
912 this.status = CodegenStatus::Error(error);
913 } else {
914 this.status = CodegenStatus::Done;
915 }
916 this.elapsed_time = Some(elapsed_time);
917 this.completion = Some(completion.lock().clone());
918 if let Some(usage) = token_usage {
919 let usage = usage.lock();
920 telemetry::event!(
921 "Inline Assistant Completion",
922 model = model_telemetry_id,
923 model_provider = model_provider_id,
924 input_tokens = usage.input_tokens,
925 output_tokens = usage.output_tokens,
926 )
927 }
928
929 cx.emit(CodegenEvent::Finished);
930 cx.notify();
931 })
932 .ok();
933 })
934 }
935
936 pub fn current_completion(&self) -> Option<String> {
937 self.completion.clone()
938 }
939
940 #[cfg(any(test, feature = "test-support"))]
941 pub fn current_description(&self) -> Option<String> {
942 self.description.clone()
943 }
944
945 #[cfg(any(test, feature = "test-support"))]
946 pub fn current_failure(&self) -> Option<String> {
947 self.failure.clone()
948 }
949
950 pub fn selected_text(&self) -> Option<&str> {
951 self.selected_text.as_deref()
952 }
953
954 pub fn stop(&mut self, cx: &mut Context<Self>) {
955 self.last_equal_ranges.clear();
956 if self.diff.is_empty() {
957 self.status = CodegenStatus::Idle;
958 } else {
959 self.status = CodegenStatus::Done;
960 }
961 self.generation = Task::ready(());
962 cx.emit(CodegenEvent::Finished);
963 cx.notify();
964 }
965
966 pub fn undo(&mut self, cx: &mut Context<Self>) {
967 self.buffer.update(cx, |buffer, cx| {
968 if let Some(transaction_id) = self.transformation_transaction_id.take() {
969 buffer.undo_transaction(transaction_id, cx);
970 buffer.refresh_preview(cx);
971 }
972 });
973 }
974
975 fn apply_edits(
976 &mut self,
977 edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
978 cx: &mut Context<CodegenAlternative>,
979 ) {
980 let transaction = self.buffer.update(cx, |buffer, cx| {
981 // Avoid grouping agent edits with user edits.
982 buffer.finalize_last_transaction(cx);
983 buffer.start_transaction(cx);
984 buffer.edit(edits, None, cx);
985 buffer.end_transaction_with_source(BufferEditSource::Agent, cx)
986 });
987
988 if let Some(transaction) = transaction {
989 if let Some(first_transaction) = self.transformation_transaction_id {
990 // Group all agent edits into the first transaction.
991 self.buffer.update(cx, |buffer, cx| {
992 buffer.merge_transactions(transaction, first_transaction, cx)
993 });
994 } else {
995 self.transformation_transaction_id = Some(transaction);
996 self.buffer
997 .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
998 }
999 }
1000 }
1001
1002 fn reapply_line_based_diff(
1003 &mut self,
1004 line_operations: impl IntoIterator<Item = LineOperation>,
1005 cx: &mut Context<Self>,
1006 ) {
1007 let old_snapshot = self.snapshot.clone();
1008 let old_range = self.range.to_point(&old_snapshot);
1009 let new_snapshot = self.buffer.read(cx).snapshot(cx);
1010 let new_range = self.range.to_point(&new_snapshot);
1011
1012 let mut old_row = old_range.start.row;
1013 let mut new_row = new_range.start.row;
1014
1015 self.diff.deleted_row_ranges.clear();
1016 self.diff.inserted_row_ranges.clear();
1017 for operation in line_operations {
1018 match operation {
1019 LineOperation::Keep { lines } => {
1020 old_row += lines;
1021 new_row += lines;
1022 }
1023 LineOperation::Delete { lines } => {
1024 let old_end_row = old_row + lines - 1;
1025 let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
1026
1027 if let Some((_, last_deleted_row_range)) =
1028 self.diff.deleted_row_ranges.last_mut()
1029 {
1030 if *last_deleted_row_range.end() + 1 == old_row {
1031 *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
1032 } else {
1033 self.diff
1034 .deleted_row_ranges
1035 .push((new_row, old_row..=old_end_row));
1036 }
1037 } else {
1038 self.diff
1039 .deleted_row_ranges
1040 .push((new_row, old_row..=old_end_row));
1041 }
1042
1043 old_row += lines;
1044 }
1045 LineOperation::Insert { lines } => {
1046 let new_end_row = new_row + lines - 1;
1047 let start = new_snapshot.anchor_before(Point::new(new_row, 0));
1048 let end = new_snapshot.anchor_before(Point::new(
1049 new_end_row,
1050 new_snapshot.line_len(MultiBufferRow(new_end_row)),
1051 ));
1052 self.diff.inserted_row_ranges.push(start..end);
1053 new_row += lines;
1054 }
1055 }
1056
1057 cx.notify();
1058 }
1059 }
1060
1061 fn reapply_batch_diff(&mut self, cx: &mut Context<Self>) -> Task<()> {
1062 let old_snapshot = self.snapshot.clone();
1063 let old_range = self.range.to_point(&old_snapshot);
1064 let new_snapshot = self.buffer.read(cx).snapshot(cx);
1065 let new_range = self.range.to_point(&new_snapshot);
1066
1067 cx.spawn(async move |codegen, cx| {
1068 let (deleted_row_ranges, inserted_row_ranges) = cx
1069 .background_spawn(async move {
1070 let old_text = old_snapshot
1071 .text_for_range(
1072 Point::new(old_range.start.row, 0)
1073 ..Point::new(
1074 old_range.end.row,
1075 old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
1076 ),
1077 )
1078 .collect::<String>();
1079 let new_text = new_snapshot
1080 .text_for_range(
1081 Point::new(new_range.start.row, 0)
1082 ..Point::new(
1083 new_range.end.row,
1084 new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
1085 ),
1086 )
1087 .collect::<String>();
1088
1089 let old_start_row = old_range.start.row;
1090 let new_start_row = new_range.start.row;
1091 let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
1092 let mut inserted_row_ranges = Vec::new();
1093 for (old_rows, new_rows) in line_diff(&old_text, &new_text) {
1094 let old_rows = old_start_row + old_rows.start..old_start_row + old_rows.end;
1095 let new_rows = new_start_row + new_rows.start..new_start_row + new_rows.end;
1096 if !old_rows.is_empty() {
1097 deleted_row_ranges.push((
1098 new_snapshot.anchor_before(Point::new(new_rows.start, 0)),
1099 old_rows.start..=old_rows.end - 1,
1100 ));
1101 }
1102 if !new_rows.is_empty() {
1103 let start = new_snapshot.anchor_before(Point::new(new_rows.start, 0));
1104 let new_end_row = new_rows.end - 1;
1105 let end = new_snapshot.anchor_before(Point::new(
1106 new_end_row,
1107 new_snapshot.line_len(MultiBufferRow(new_end_row)),
1108 ));
1109 inserted_row_ranges.push(start..end);
1110 }
1111 }
1112 (deleted_row_ranges, inserted_row_ranges)
1113 })
1114 .await;
1115
1116 codegen
1117 .update(cx, |codegen, cx| {
1118 codegen.diff.deleted_row_ranges = deleted_row_ranges;
1119 codegen.diff.inserted_row_ranges = inserted_row_ranges;
1120 cx.notify();
1121 })
1122 .ok();
1123 })
1124 }
1125
1126 fn handle_completion(
1127 &mut self,
1128 model: Arc<dyn LanguageModel>,
1129 completion_stream: Task<
1130 Result<
1131 BoxStream<
1132 'static,
1133 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
1134 >,
1135 LanguageModelCompletionError,
1136 >,
1137 >,
1138 cx: &mut Context<Self>,
1139 ) -> Task<()> {
1140 self.diff = Diff::default();
1141 self.status = CodegenStatus::Pending;
1142
1143 cx.notify();
1144 // Leaving this in generation so that STOP equivalent events are respected even
1145 // while we're still pre-processing the completion event
1146 cx.spawn(async move |codegen, cx| {
1147 let finish_with_status = |status: CodegenStatus, cx: &mut AsyncApp| {
1148 let _ = codegen.update(cx, |this, cx| {
1149 this.status = status;
1150 cx.emit(CodegenEvent::Finished);
1151 cx.notify();
1152 });
1153 };
1154
1155 let mut completion_events = match completion_stream.await {
1156 Ok(events) => events,
1157 Err(err) => {
1158 finish_with_status(CodegenStatus::Error(err.into()), cx);
1159 return;
1160 }
1161 };
1162
1163 enum ToolUseOutput {
1164 Rewrite {
1165 text: String,
1166 description: Option<String>,
1167 },
1168 Failure(String),
1169 }
1170
1171 enum ModelUpdate {
1172 Description(String),
1173 Failure(String),
1174 }
1175
1176 let chars_read_by_tool_id: Arc<Mutex<HashMap<LanguageModelToolUseId, usize>>> =
1177 Arc::new(Mutex::new(HashMap::default()));
1178 let process_tool_use = move |tool_use: LanguageModelToolUse| -> Option<ToolUseOutput> {
1179 let mut chars_read_by_tool_id = chars_read_by_tool_id.lock();
1180 match tool_use.name.as_ref() {
1181 REWRITE_SECTION_TOOL_NAME => {
1182 let Ok(input) = tool_use.input.parse::<RewriteSectionInput>() else {
1183 return None;
1184 };
1185 let chars_read_so_far =
1186 chars_read_by_tool_id.entry(tool_use.id).or_insert(0);
1187 let Some(text_slice) = input.replacement_text.get(*chars_read_so_far..)
1188 else {
1189 return None;
1190 };
1191 let text = text_slice.to_string();
1192 *chars_read_so_far = input.replacement_text.len();
1193 Some(ToolUseOutput::Rewrite {
1194 text,
1195 description: None,
1196 })
1197 }
1198 FAILURE_MESSAGE_TOOL_NAME => {
1199 let Ok(mut input) = tool_use.input.parse::<FailureMessageInput>() else {
1200 return None;
1201 };
1202 Some(ToolUseOutput::Failure(std::mem::take(&mut input.message)))
1203 }
1204 _ => None,
1205 }
1206 };
1207
1208 let (message_tx, mut message_rx) = futures::channel::mpsc::unbounded::<ModelUpdate>();
1209
1210 cx.spawn({
1211 let codegen = codegen.clone();
1212 async move |cx| {
1213 while let Some(update) = message_rx.next().await {
1214 let _ = codegen.update(cx, |this, _cx| match update {
1215 ModelUpdate::Description(d) => this.description = Some(d),
1216 ModelUpdate::Failure(f) => this.failure = Some(f),
1217 });
1218 }
1219 }
1220 })
1221 .detach();
1222
1223 let mut message_id = None;
1224 let mut first_text = None;
1225 let last_token_usage = Arc::new(Mutex::new(TokenUsage::default()));
1226 let total_text = Arc::new(Mutex::new(String::new()));
1227
1228 loop {
1229 if let Some(first_event) = completion_events.next().await {
1230 match first_event {
1231 Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => {
1232 message_id = Some(id);
1233 }
1234 Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
1235 if let Some(output) = process_tool_use(tool_use) {
1236 let (text, update) = match output {
1237 ToolUseOutput::Rewrite { text, description } => {
1238 (Some(text), description.map(ModelUpdate::Description))
1239 }
1240 ToolUseOutput::Failure(message) => {
1241 (None, Some(ModelUpdate::Failure(message)))
1242 }
1243 };
1244 if let Some(update) = update {
1245 let _ = message_tx.unbounded_send(update);
1246 }
1247 first_text = text;
1248 if first_text.is_some() {
1249 break;
1250 }
1251 }
1252 }
1253 Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
1254 *last_token_usage.lock() = token_usage;
1255 }
1256 Ok(LanguageModelCompletionEvent::Text(text)) => {
1257 let mut lock = total_text.lock();
1258 lock.push_str(&text);
1259 }
1260 Ok(e) => {
1261 log::warn!("Unexpected event: {:?}", e);
1262 break;
1263 }
1264 Err(e) => {
1265 finish_with_status(CodegenStatus::Error(e.into()), cx);
1266 break;
1267 }
1268 }
1269 }
1270 }
1271
1272 let Some(first_text) = first_text else {
1273 finish_with_status(CodegenStatus::Done, cx);
1274 return;
1275 };
1276
1277 let move_last_token_usage = last_token_usage.clone();
1278
1279 let text_stream = Box::pin(futures::stream::once(async { Ok(first_text) }).chain(
1280 completion_events.filter_map(move |e| {
1281 let process_tool_use = process_tool_use.clone();
1282 let last_token_usage = move_last_token_usage.clone();
1283 let total_text = total_text.clone();
1284 let mut message_tx = message_tx.clone();
1285 async move {
1286 match e {
1287 Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
1288 let Some(output) = process_tool_use(tool_use) else {
1289 return None;
1290 };
1291 let (text, update) = match output {
1292 ToolUseOutput::Rewrite { text, description } => {
1293 (Some(text), description.map(ModelUpdate::Description))
1294 }
1295 ToolUseOutput::Failure(message) => {
1296 (None, Some(ModelUpdate::Failure(message)))
1297 }
1298 };
1299 if let Some(update) = update {
1300 let _ = message_tx.send(update).await;
1301 }
1302 text.map(Ok)
1303 }
1304 Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
1305 *last_token_usage.lock() = token_usage;
1306 None
1307 }
1308 Ok(LanguageModelCompletionEvent::Text(text)) => {
1309 let mut lock = total_text.lock();
1310 lock.push_str(&text);
1311 None
1312 }
1313 Ok(LanguageModelCompletionEvent::Stop(_reason)) => None,
1314 e => {
1315 log::error!("UNEXPECTED EVENT {:?}", e);
1316 None
1317 }
1318 }
1319 }
1320 }),
1321 ));
1322
1323 let language_model_text_stream = LanguageModelTextStream {
1324 message_id: message_id,
1325 stream: text_stream,
1326 last_token_usage,
1327 };
1328
1329 let Some(task) = codegen
1330 .update(cx, move |codegen, cx| {
1331 codegen.handle_stream(
1332 model,
1333 /* strip_invalid_spans: */ false,
1334 async { Ok(language_model_text_stream) },
1335 cx,
1336 )
1337 })
1338 .ok()
1339 else {
1340 return;
1341 };
1342
1343 task.await;
1344 })
1345 }
1346}
1347
1348#[derive(Copy, Clone, Debug)]
1349pub enum CodegenEvent {
1350 Finished,
1351 Undone,
1352}
1353
1354struct StripInvalidSpans<T> {
1355 stream: T,
1356 stream_done: bool,
1357 buffer: String,
1358 first_line: bool,
1359 line_end: bool,
1360 starts_with_code_block: bool,
1361}
1362
1363impl<T> StripInvalidSpans<T>
1364where
1365 T: Stream<Item = Result<String>>,
1366{
1367 fn new(stream: T) -> Self {
1368 Self {
1369 stream,
1370 stream_done: false,
1371 buffer: String::new(),
1372 first_line: true,
1373 line_end: false,
1374 starts_with_code_block: false,
1375 }
1376 }
1377}
1378
1379impl<T> Stream for StripInvalidSpans<T>
1380where
1381 T: Stream<Item = Result<String>>,
1382{
1383 type Item = Result<String>;
1384
1385 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
1386 const CODE_BLOCK_DELIMITER: &str = "```";
1387 const CURSOR_SPAN: &str = "<|CURSOR|>";
1388
1389 let this = unsafe { self.get_unchecked_mut() };
1390 loop {
1391 if !this.stream_done {
1392 let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
1393 match stream.as_mut().poll_next(cx) {
1394 Poll::Ready(Some(Ok(chunk))) => {
1395 this.buffer.push_str(&chunk);
1396 }
1397 Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
1398 Poll::Ready(None) => {
1399 this.stream_done = true;
1400 }
1401 Poll::Pending => return Poll::Pending,
1402 }
1403 }
1404
1405 let mut chunk = String::new();
1406 let mut consumed = 0;
1407 if !this.buffer.is_empty() {
1408 let mut lines = this.buffer.split('\n').enumerate().peekable();
1409 while let Some((line_ix, line)) = lines.next() {
1410 if line_ix > 0 {
1411 this.first_line = false;
1412 }
1413
1414 if this.first_line {
1415 let trimmed_line = line.trim();
1416 if lines.peek().is_some() {
1417 if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
1418 consumed += line.len() + 1;
1419 this.starts_with_code_block = true;
1420 continue;
1421 }
1422 } else if trimmed_line.is_empty()
1423 || prefixes(CODE_BLOCK_DELIMITER)
1424 .any(|prefix| trimmed_line.starts_with(prefix))
1425 {
1426 break;
1427 }
1428 }
1429
1430 let line_without_cursor = line.replace(CURSOR_SPAN, "");
1431 if lines.peek().is_some() {
1432 if this.line_end {
1433 chunk.push('\n');
1434 }
1435
1436 chunk.push_str(&line_without_cursor);
1437 this.line_end = true;
1438 consumed += line.len() + 1;
1439 } else if this.stream_done {
1440 if !this.starts_with_code_block
1441 || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
1442 {
1443 if this.line_end {
1444 chunk.push('\n');
1445 }
1446
1447 chunk.push_str(line);
1448 }
1449
1450 consumed += line.len();
1451 } else {
1452 let trimmed_line = line.trim();
1453 if trimmed_line.is_empty()
1454 || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
1455 || prefixes(CODE_BLOCK_DELIMITER)
1456 .any(|prefix| trimmed_line.ends_with(prefix))
1457 {
1458 break;
1459 } else {
1460 if this.line_end {
1461 chunk.push('\n');
1462 this.line_end = false;
1463 }
1464
1465 chunk.push_str(&line_without_cursor);
1466 consumed += line.len();
1467 }
1468 }
1469 }
1470 }
1471
1472 this.buffer = this.buffer.split_off(consumed);
1473 if !chunk.is_empty() {
1474 return Poll::Ready(Some(Ok(chunk)));
1475 } else if this.stream_done {
1476 return Poll::Ready(None);
1477 }
1478 }
1479 }
1480}
1481
1482fn prefixes(text: &str) -> impl Iterator<Item = &str> {
1483 (0..text.len() - 1).map(|ix| &text[..ix + 1])
1484}
1485
1486#[derive(Default)]
1487pub struct Diff {
1488 pub deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1489 pub inserted_row_ranges: Vec<Range<Anchor>>,
1490}
1491
1492impl Diff {
1493 fn is_empty(&self) -> bool {
1494 self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
1495 }
1496}
1497
1498#[cfg(test)]
1499mod tests {
1500 use super::*;
1501 use futures::{
1502 Stream,
1503 stream::{self},
1504 };
1505 use gpui::TestAppContext;
1506 use indoc::indoc;
1507 use language::{Buffer, Point};
1508 use language_model::fake_provider::FakeLanguageModel;
1509 use language_model::{
1510 LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelRegistry,
1511 LanguageModelToolUse, StopReason, TokenUsage,
1512 };
1513 use languages::rust_lang;
1514 use rand::prelude::*;
1515 use settings::SettingsStore;
1516 use std::{future, sync::Arc};
1517
1518 #[gpui::test(iterations = 10)]
1519 async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
1520 init_test(cx);
1521
1522 let text = indoc! {"
1523 fn main() {
1524 let x = 0;
1525 for _ in 0..10 {
1526 x += 1;
1527 }
1528 }
1529 "};
1530 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1531 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1532 let range = buffer.read_with(cx, |buffer, cx| {
1533 let snapshot = buffer.snapshot(cx);
1534 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
1535 });
1536 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1537 let codegen = cx.new(|cx| {
1538 CodegenAlternative::new(
1539 buffer.clone(),
1540 range.clone(),
1541 true,
1542 prompt_builder,
1543 Uuid::new_v4(),
1544 cx,
1545 )
1546 });
1547
1548 let chunks_tx = simulate_response_stream(&codegen, cx);
1549
1550 let mut new_text = concat!(
1551 " let mut x = 0;\n",
1552 " while x < 10 {\n",
1553 " x += 1;\n",
1554 " }",
1555 );
1556 while !new_text.is_empty() {
1557 let max_len = cmp::min(new_text.len(), 10);
1558 let len = rng.random_range(1..=max_len);
1559 let (chunk, suffix) = new_text.split_at(len);
1560 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1561 new_text = suffix;
1562 cx.background_executor.run_until_parked();
1563 }
1564 drop(chunks_tx);
1565 cx.background_executor.run_until_parked();
1566
1567 assert_eq!(
1568 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1569 indoc! {"
1570 fn main() {
1571 let mut x = 0;
1572 while x < 10 {
1573 x += 1;
1574 }
1575 }
1576 "}
1577 );
1578 }
1579
1580 #[gpui::test(iterations = 10)]
1581 async fn test_autoindent_when_generating_past_indentation(
1582 cx: &mut TestAppContext,
1583 mut rng: StdRng,
1584 ) {
1585 init_test(cx);
1586
1587 let text = indoc! {"
1588 fn main() {
1589 le
1590 }
1591 "};
1592 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1593 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1594 let range = buffer.read_with(cx, |buffer, cx| {
1595 let snapshot = buffer.snapshot(cx);
1596 snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
1597 });
1598 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1599 let codegen = cx.new(|cx| {
1600 CodegenAlternative::new(
1601 buffer.clone(),
1602 range.clone(),
1603 true,
1604 prompt_builder,
1605 Uuid::new_v4(),
1606 cx,
1607 )
1608 });
1609
1610 let chunks_tx = simulate_response_stream(&codegen, cx);
1611
1612 cx.background_executor.run_until_parked();
1613
1614 let mut new_text = concat!(
1615 "t mut x = 0;\n",
1616 "while x < 10 {\n",
1617 " x += 1;\n",
1618 "}", //
1619 );
1620 while !new_text.is_empty() {
1621 let max_len = cmp::min(new_text.len(), 10);
1622 let len = rng.random_range(1..=max_len);
1623 let (chunk, suffix) = new_text.split_at(len);
1624 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1625 new_text = suffix;
1626 cx.background_executor.run_until_parked();
1627 }
1628 drop(chunks_tx);
1629 cx.background_executor.run_until_parked();
1630
1631 assert_eq!(
1632 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1633 indoc! {"
1634 fn main() {
1635 let mut x = 0;
1636 while x < 10 {
1637 x += 1;
1638 }
1639 }
1640 "}
1641 );
1642 }
1643
1644 #[gpui::test(iterations = 10)]
1645 async fn test_autoindent_when_generating_before_indentation(
1646 cx: &mut TestAppContext,
1647 mut rng: StdRng,
1648 ) {
1649 init_test(cx);
1650
1651 let text = concat!(
1652 "fn main() {\n",
1653 " \n",
1654 "}\n" //
1655 );
1656 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1657 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1658 let range = buffer.read_with(cx, |buffer, cx| {
1659 let snapshot = buffer.snapshot(cx);
1660 snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
1661 });
1662 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1663 let codegen = cx.new(|cx| {
1664 CodegenAlternative::new(
1665 buffer.clone(),
1666 range.clone(),
1667 true,
1668 prompt_builder,
1669 Uuid::new_v4(),
1670 cx,
1671 )
1672 });
1673
1674 let chunks_tx = simulate_response_stream(&codegen, cx);
1675
1676 cx.background_executor.run_until_parked();
1677
1678 let mut new_text = concat!(
1679 "let mut x = 0;\n",
1680 "while x < 10 {\n",
1681 " x += 1;\n",
1682 "}", //
1683 );
1684 while !new_text.is_empty() {
1685 let max_len = cmp::min(new_text.len(), 10);
1686 let len = rng.random_range(1..=max_len);
1687 let (chunk, suffix) = new_text.split_at(len);
1688 chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1689 new_text = suffix;
1690 cx.background_executor.run_until_parked();
1691 }
1692 drop(chunks_tx);
1693 cx.background_executor.run_until_parked();
1694
1695 assert_eq!(
1696 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1697 indoc! {"
1698 fn main() {
1699 let mut x = 0;
1700 while x < 10 {
1701 x += 1;
1702 }
1703 }
1704 "}
1705 );
1706 }
1707
1708 #[gpui::test(iterations = 10)]
1709 async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
1710 init_test(cx);
1711
1712 let text = indoc! {"
1713 func main() {
1714 \tx := 0
1715 \tfor i := 0; i < 10; i++ {
1716 \t\tx++
1717 \t}
1718 }
1719 "};
1720 let buffer = cx.new(|cx| Buffer::local(text, cx));
1721 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1722 let range = buffer.read_with(cx, |buffer, cx| {
1723 let snapshot = buffer.snapshot(cx);
1724 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
1725 });
1726 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1727 let codegen = cx.new(|cx| {
1728 CodegenAlternative::new(
1729 buffer.clone(),
1730 range.clone(),
1731 true,
1732 prompt_builder,
1733 Uuid::new_v4(),
1734 cx,
1735 )
1736 });
1737
1738 let chunks_tx = simulate_response_stream(&codegen, cx);
1739 let new_text = concat!(
1740 "func main() {\n",
1741 "\tx := 0\n",
1742 "\tfor x < 10 {\n",
1743 "\t\tx++\n",
1744 "\t}", //
1745 );
1746 chunks_tx.unbounded_send(new_text.to_string()).unwrap();
1747 drop(chunks_tx);
1748 cx.background_executor.run_until_parked();
1749
1750 assert_eq!(
1751 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1752 indoc! {"
1753 func main() {
1754 \tx := 0
1755 \tfor x < 10 {
1756 \t\tx++
1757 \t}
1758 }
1759 "}
1760 );
1761 }
1762
1763 #[gpui::test]
1764 async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
1765 init_test(cx);
1766
1767 let text = indoc! {"
1768 fn main() {
1769 let x = 0;
1770 }
1771 "};
1772 let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
1773 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1774 let range = buffer.read_with(cx, |buffer, cx| {
1775 let snapshot = buffer.snapshot(cx);
1776 snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
1777 });
1778 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1779 let codegen = cx.new(|cx| {
1780 CodegenAlternative::new(
1781 buffer.clone(),
1782 range.clone(),
1783 false,
1784 prompt_builder,
1785 Uuid::new_v4(),
1786 cx,
1787 )
1788 });
1789
1790 let chunks_tx = simulate_response_stream(&codegen, cx);
1791 chunks_tx
1792 .unbounded_send("let mut x = 0;\nx += 1;".to_string())
1793 .unwrap();
1794 drop(chunks_tx);
1795 cx.run_until_parked();
1796
1797 // The codegen is inactive, so the buffer doesn't get modified.
1798 assert_eq!(
1799 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1800 text
1801 );
1802
1803 // Activating the codegen applies the changes.
1804 codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
1805 assert_eq!(
1806 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1807 indoc! {"
1808 fn main() {
1809 let mut x = 0;
1810 x += 1;
1811 }
1812 "}
1813 );
1814
1815 // Deactivating the codegen undoes the changes.
1816 codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
1817 cx.run_until_parked();
1818 assert_eq!(
1819 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1820 text
1821 );
1822 }
1823
1824 // When not streaming tool calls, we strip backticks as part of parsing the model's
1825 // plain text response. This is a regression test for a bug where we stripped
1826 // backticks incorrectly.
1827 #[gpui::test]
1828 async fn test_allows_model_to_output_backticks(cx: &mut TestAppContext) {
1829 init_test(cx);
1830 let text = "- Improved; `cmd+click` behavior. Now requires `cmd` to be pressed before the click starts or it doesn't run. ([#44579](https://github.com/zed-industries/zed/pull/44579); thanks [Zachiah](https://github.com/Zachiah))";
1831 let buffer = cx.new(|cx| Buffer::local("", cx));
1832 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1833 let range = buffer.read_with(cx, |buffer, cx| {
1834 let snapshot = buffer.snapshot(cx);
1835 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(0, 0))
1836 });
1837 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1838 let codegen = cx.new(|cx| {
1839 CodegenAlternative::new(
1840 buffer.clone(),
1841 range.clone(),
1842 true,
1843 prompt_builder,
1844 Uuid::new_v4(),
1845 cx,
1846 )
1847 });
1848
1849 let events_tx = simulate_tool_based_completion(&codegen, cx);
1850 let chunk_len = text.find('`').unwrap();
1851 events_tx
1852 .unbounded_send(rewrite_tool_use("tool_1", &text[..chunk_len], false))
1853 .unwrap();
1854 events_tx
1855 .unbounded_send(rewrite_tool_use("tool_1", &text, true))
1856 .unwrap();
1857 events_tx
1858 .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
1859 .unwrap();
1860 drop(events_tx);
1861 cx.run_until_parked();
1862
1863 assert_eq!(
1864 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1865 text
1866 );
1867 }
1868
1869 // Regression test: a second rewrite tool use with a *shorter* replacement_text
1870 // than the first would cause an index-out-of-bounds panic because the
1871 // chars_read_so_far counter was shared across all tool use IDs.
1872 #[gpui::test]
1873 async fn test_separate_tool_uses_have_independent_char_counters(cx: &mut TestAppContext) {
1874 init_test(cx);
1875 let buffer = cx.new(|cx| Buffer::local("", cx));
1876 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1877 let range = buffer.read_with(cx, |buffer, cx| {
1878 let snapshot = buffer.snapshot(cx);
1879 snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(0, 0))
1880 });
1881 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1882 let codegen = cx.new(|cx| {
1883 CodegenAlternative::new(
1884 buffer.clone(),
1885 range.clone(),
1886 true,
1887 prompt_builder,
1888 Uuid::new_v4(),
1889 cx,
1890 )
1891 });
1892
1893 let events_tx = simulate_tool_based_completion(&codegen, cx);
1894 // tool_1 has longer text; tool_2 has shorter text. With the old shared
1895 // counter, processing tool_2 would attempt replacement_text[N..] where
1896 // N > replacement_text.len(), panicking with index out of bounds.
1897 events_tx
1898 .unbounded_send(rewrite_tool_use("tool_1", "longer replacement text", true))
1899 .unwrap();
1900 events_tx
1901 .unbounded_send(rewrite_tool_use("tool_2", "short", true))
1902 .unwrap();
1903 events_tx
1904 .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
1905 .unwrap();
1906 drop(events_tx);
1907 cx.run_until_parked();
1908
1909 assert_eq!(
1910 buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1911 "longer replacement textshort"
1912 );
1913 }
1914
1915 #[gpui::test]
1916 async fn test_strip_invalid_spans_from_codeblock() {
1917 assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
1918 assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
1919 assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
1920 assert_chunks(
1921 "```html\n```js\nLorem ipsum dolor\n```\n```",
1922 "```js\nLorem ipsum dolor\n```",
1923 )
1924 .await;
1925 assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
1926 assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
1927 assert_chunks("Lorem ipsum", "Lorem ipsum").await;
1928 assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
1929
1930 async fn assert_chunks(text: &str, expected_text: &str) {
1931 for chunk_size in 1..=text.len() {
1932 let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
1933 .map(|chunk| chunk.unwrap())
1934 .collect::<String>()
1935 .await;
1936 assert_eq!(
1937 actual_text, expected_text,
1938 "failed to strip invalid spans, chunk size: {}",
1939 chunk_size
1940 );
1941 }
1942 }
1943
1944 fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
1945 stream::iter(
1946 text.chars()
1947 .collect::<Vec<_>>()
1948 .chunks(size)
1949 .map(|chunk| Ok(chunk.iter().collect::<String>()))
1950 .collect::<Vec<_>>(),
1951 )
1952 }
1953 }
1954
1955 fn init_test(cx: &mut TestAppContext) {
1956 cx.update(LanguageModelRegistry::test);
1957 cx.set_global(cx.update(SettingsStore::test));
1958 }
1959
1960 fn simulate_response_stream(
1961 codegen: &Entity<CodegenAlternative>,
1962 cx: &mut TestAppContext,
1963 ) -> mpsc::UnboundedSender<String> {
1964 let (chunks_tx, chunks_rx) = mpsc::unbounded();
1965 let model = Arc::new(FakeLanguageModel::default());
1966 codegen.update(cx, |codegen, cx| {
1967 codegen.generation = codegen.handle_stream(
1968 model,
1969 /* strip_invalid_spans: */ false,
1970 future::ready(Ok(LanguageModelTextStream {
1971 message_id: None,
1972 stream: chunks_rx.map(Ok).boxed(),
1973 last_token_usage: Arc::new(Mutex::new(TokenUsage::default())),
1974 })),
1975 cx,
1976 );
1977 });
1978 chunks_tx
1979 }
1980
1981 fn simulate_tool_based_completion(
1982 codegen: &Entity<CodegenAlternative>,
1983 cx: &mut TestAppContext,
1984 ) -> mpsc::UnboundedSender<LanguageModelCompletionEvent> {
1985 let (events_tx, events_rx) = mpsc::unbounded();
1986 let model = Arc::new(FakeLanguageModel::default());
1987 codegen.update(cx, |codegen, cx| {
1988 let completion_stream = Task::ready(Ok(events_rx.map(Ok).boxed()
1989 as BoxStream<
1990 'static,
1991 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
1992 >));
1993 codegen.generation = codegen.handle_completion(model, completion_stream, cx);
1994 });
1995 events_tx
1996 }
1997
1998 fn rewrite_tool_use(
1999 id: &str,
2000 replacement_text: &str,
2001 is_complete: bool,
2002 ) -> LanguageModelCompletionEvent {
2003 let input = RewriteSectionInput {
2004 replacement_text: replacement_text.into(),
2005 };
2006 LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse {
2007 id: id.into(),
2008 name: REWRITE_SECTION_TOOL_NAME.into(),
2009 raw_input: serde_json::to_string(&input).unwrap(),
2010 input: language_model::LanguageModelToolUseInput::Json(
2011 serde_json::to_value(&input).unwrap(),
2012 ),
2013 is_input_complete: is_complete,
2014 thought_signature: None,
2015 })
2016 }
2017}
2018