buffer_codegen.rs

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