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(|_, cx| async move { 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(|codegen, mut cx| {
 484            async move {
 485                let stream = stream.await;
 486                let message_id = stream
 487                    .as_ref()
 488                    .ok()
 489                    .and_then(|stream| stream.message_id.clone());
 490                let generate = async {
 491                    let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
 492                    let executor = cx.background_executor().clone();
 493                    let message_id = message_id.clone();
 494                    let line_based_stream_diff: Task<anyhow::Result<()>> =
 495                        cx.background_spawn(async move {
 496                            let mut response_latency = None;
 497                            let request_start = Instant::now();
 498                            let diff = async {
 499                                let chunks = StripInvalidSpans::new(stream?.stream);
 500                                futures::pin_mut!(chunks);
 501                                let mut diff = StreamingDiff::new(selected_text.to_string());
 502                                let mut line_diff = LineDiff::default();
 503
 504                                let mut new_text = String::new();
 505                                let mut base_indent = None;
 506                                let mut line_indent = None;
 507                                let mut first_line = true;
 508
 509                                while let Some(chunk) = chunks.next().await {
 510                                    if response_latency.is_none() {
 511                                        response_latency = Some(request_start.elapsed());
 512                                    }
 513                                    let chunk = chunk?;
 514                                    completion_clone.lock().push_str(&chunk);
 515
 516                                    let mut lines = chunk.split('\n').peekable();
 517                                    while let Some(line) = lines.next() {
 518                                        new_text.push_str(line);
 519                                        if line_indent.is_none() {
 520                                            if let Some(non_whitespace_ch_ix) =
 521                                                new_text.find(|ch: char| !ch.is_whitespace())
 522                                            {
 523                                                line_indent = Some(non_whitespace_ch_ix);
 524                                                base_indent = base_indent.or(line_indent);
 525
 526                                                let line_indent = line_indent.unwrap();
 527                                                let base_indent = base_indent.unwrap();
 528                                                let indent_delta =
 529                                                    line_indent as i32 - base_indent as i32;
 530                                                let mut corrected_indent_len = cmp::max(
 531                                                    0,
 532                                                    suggested_line_indent.len as i32 + indent_delta,
 533                                                )
 534                                                    as usize;
 535                                                if first_line {
 536                                                    corrected_indent_len = corrected_indent_len
 537                                                        .saturating_sub(
 538                                                            selection_start.column as usize,
 539                                                        );
 540                                                }
 541
 542                                                let indent_char = suggested_line_indent.char();
 543                                                let mut indent_buffer = [0; 4];
 544                                                let indent_str =
 545                                                    indent_char.encode_utf8(&mut indent_buffer);
 546                                                new_text.replace_range(
 547                                                    ..line_indent,
 548                                                    &indent_str.repeat(corrected_indent_len),
 549                                                );
 550                                            }
 551                                        }
 552
 553                                        if line_indent.is_some() {
 554                                            let char_ops = diff.push_new(&new_text);
 555                                            line_diff
 556                                                .push_char_operations(&char_ops, &selected_text);
 557                                            diff_tx
 558                                                .send((char_ops, line_diff.line_operations()))
 559                                                .await?;
 560                                            new_text.clear();
 561                                        }
 562
 563                                        if lines.peek().is_some() {
 564                                            let char_ops = diff.push_new("\n");
 565                                            line_diff
 566                                                .push_char_operations(&char_ops, &selected_text);
 567                                            diff_tx
 568                                                .send((char_ops, line_diff.line_operations()))
 569                                                .await?;
 570                                            if line_indent.is_none() {
 571                                                // Don't write out the leading indentation in empty lines on the next line
 572                                                // This is the case where the above if statement didn't clear the buffer
 573                                                new_text.clear();
 574                                            }
 575                                            line_indent = None;
 576                                            first_line = false;
 577                                        }
 578                                    }
 579                                }
 580
 581                                let mut char_ops = diff.push_new(&new_text);
 582                                char_ops.extend(diff.finish());
 583                                line_diff.push_char_operations(&char_ops, &selected_text);
 584                                line_diff.finish(&selected_text);
 585                                diff_tx
 586                                    .send((char_ops, line_diff.line_operations()))
 587                                    .await?;
 588
 589                                anyhow::Ok(())
 590                            };
 591
 592                            let result = diff.await;
 593
 594                            let error_message =
 595                                result.as_ref().err().map(|error| error.to_string());
 596                            report_assistant_event(
 597                                AssistantEvent {
 598                                    conversation_id: None,
 599                                    message_id,
 600                                    kind: AssistantKind::Inline,
 601                                    phase: AssistantPhase::Response,
 602                                    model: model_telemetry_id,
 603                                    model_provider: model_provider_id.to_string(),
 604                                    response_latency,
 605                                    error_message,
 606                                    language_name: language_name.map(|name| name.to_proto()),
 607                                },
 608                                telemetry,
 609                                http_client,
 610                                model_api_key,
 611                                &executor,
 612                            );
 613
 614                            result?;
 615                            Ok(())
 616                        });
 617
 618                    while let Some((char_ops, line_ops)) = diff_rx.next().await {
 619                        codegen.update(&mut cx, |codegen, cx| {
 620                            codegen.last_equal_ranges.clear();
 621
 622                            let edits = char_ops
 623                                .into_iter()
 624                                .filter_map(|operation| match operation {
 625                                    CharOperation::Insert { text } => {
 626                                        let edit_start = snapshot.anchor_after(edit_start);
 627                                        Some((edit_start..edit_start, text))
 628                                    }
 629                                    CharOperation::Delete { bytes } => {
 630                                        let edit_end = edit_start + bytes;
 631                                        let edit_range = snapshot.anchor_after(edit_start)
 632                                            ..snapshot.anchor_before(edit_end);
 633                                        edit_start = edit_end;
 634                                        Some((edit_range, String::new()))
 635                                    }
 636                                    CharOperation::Keep { bytes } => {
 637                                        let edit_end = edit_start + bytes;
 638                                        let edit_range = snapshot.anchor_after(edit_start)
 639                                            ..snapshot.anchor_before(edit_end);
 640                                        edit_start = edit_end;
 641                                        codegen.last_equal_ranges.push(edit_range);
 642                                        None
 643                                    }
 644                                })
 645                                .collect::<Vec<_>>();
 646
 647                            if codegen.active {
 648                                codegen.apply_edits(edits.iter().cloned(), cx);
 649                                codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
 650                            }
 651                            codegen.edits.extend(edits);
 652                            codegen.line_operations = line_ops;
 653                            codegen.edit_position = Some(snapshot.anchor_after(edit_start));
 654
 655                            cx.notify();
 656                        })?;
 657                    }
 658
 659                    // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
 660                    // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
 661                    // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
 662                    let batch_diff_task =
 663                        codegen.update(&mut cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
 664                    let (line_based_stream_diff, ()) =
 665                        join!(line_based_stream_diff, batch_diff_task);
 666                    line_based_stream_diff?;
 667
 668                    anyhow::Ok(())
 669                };
 670
 671                let result = generate.await;
 672                let elapsed_time = start_time.elapsed().as_secs_f64();
 673
 674                codegen
 675                    .update(&mut cx, |this, cx| {
 676                        this.message_id = message_id;
 677                        this.last_equal_ranges.clear();
 678                        if let Err(error) = result {
 679                            this.status = CodegenStatus::Error(error);
 680                        } else {
 681                            this.status = CodegenStatus::Done;
 682                        }
 683                        this.elapsed_time = Some(elapsed_time);
 684                        this.completion = Some(completion.lock().clone());
 685                        cx.emit(CodegenEvent::Finished);
 686                        cx.notify();
 687                    })
 688                    .ok();
 689            }
 690        });
 691        cx.notify();
 692    }
 693
 694    pub fn stop(&mut self, cx: &mut Context<Self>) {
 695        self.last_equal_ranges.clear();
 696        if self.diff.is_empty() {
 697            self.status = CodegenStatus::Idle;
 698        } else {
 699            self.status = CodegenStatus::Done;
 700        }
 701        self.generation = Task::ready(());
 702        cx.emit(CodegenEvent::Finished);
 703        cx.notify();
 704    }
 705
 706    pub fn undo(&mut self, cx: &mut Context<Self>) {
 707        self.buffer.update(cx, |buffer, cx| {
 708            if let Some(transaction_id) = self.transformation_transaction_id.take() {
 709                buffer.undo_transaction(transaction_id, cx);
 710                buffer.refresh_preview(cx);
 711            }
 712        });
 713    }
 714
 715    fn apply_edits(
 716        &mut self,
 717        edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
 718        cx: &mut Context<CodegenAlternative>,
 719    ) {
 720        let transaction = self.buffer.update(cx, |buffer, cx| {
 721            // Avoid grouping assistant edits with user edits.
 722            buffer.finalize_last_transaction(cx);
 723            buffer.start_transaction(cx);
 724            buffer.edit(edits, None, cx);
 725            buffer.end_transaction(cx)
 726        });
 727
 728        if let Some(transaction) = transaction {
 729            if let Some(first_transaction) = self.transformation_transaction_id {
 730                // Group all assistant edits into the first transaction.
 731                self.buffer.update(cx, |buffer, cx| {
 732                    buffer.merge_transactions(transaction, first_transaction, cx)
 733                });
 734            } else {
 735                self.transformation_transaction_id = Some(transaction);
 736                self.buffer
 737                    .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
 738            }
 739        }
 740    }
 741
 742    fn reapply_line_based_diff(
 743        &mut self,
 744        line_operations: impl IntoIterator<Item = LineOperation>,
 745        cx: &mut Context<Self>,
 746    ) {
 747        let old_snapshot = self.snapshot.clone();
 748        let old_range = self.range.to_point(&old_snapshot);
 749        let new_snapshot = self.buffer.read(cx).snapshot(cx);
 750        let new_range = self.range.to_point(&new_snapshot);
 751
 752        let mut old_row = old_range.start.row;
 753        let mut new_row = new_range.start.row;
 754
 755        self.diff.deleted_row_ranges.clear();
 756        self.diff.inserted_row_ranges.clear();
 757        for operation in line_operations {
 758            match operation {
 759                LineOperation::Keep { lines } => {
 760                    old_row += lines;
 761                    new_row += lines;
 762                }
 763                LineOperation::Delete { lines } => {
 764                    let old_end_row = old_row + lines - 1;
 765                    let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
 766
 767                    if let Some((_, last_deleted_row_range)) =
 768                        self.diff.deleted_row_ranges.last_mut()
 769                    {
 770                        if *last_deleted_row_range.end() + 1 == old_row {
 771                            *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
 772                        } else {
 773                            self.diff
 774                                .deleted_row_ranges
 775                                .push((new_row, old_row..=old_end_row));
 776                        }
 777                    } else {
 778                        self.diff
 779                            .deleted_row_ranges
 780                            .push((new_row, old_row..=old_end_row));
 781                    }
 782
 783                    old_row += lines;
 784                }
 785                LineOperation::Insert { lines } => {
 786                    let new_end_row = new_row + lines - 1;
 787                    let start = new_snapshot.anchor_before(Point::new(new_row, 0));
 788                    let end = new_snapshot.anchor_before(Point::new(
 789                        new_end_row,
 790                        new_snapshot.line_len(MultiBufferRow(new_end_row)),
 791                    ));
 792                    self.diff.inserted_row_ranges.push(start..end);
 793                    new_row += lines;
 794                }
 795            }
 796
 797            cx.notify();
 798        }
 799    }
 800
 801    fn reapply_batch_diff(&mut self, cx: &mut Context<Self>) -> Task<()> {
 802        let old_snapshot = self.snapshot.clone();
 803        let old_range = self.range.to_point(&old_snapshot);
 804        let new_snapshot = self.buffer.read(cx).snapshot(cx);
 805        let new_range = self.range.to_point(&new_snapshot);
 806
 807        cx.spawn(|codegen, mut cx| async move {
 808            let (deleted_row_ranges, inserted_row_ranges) = cx
 809                .background_spawn(async move {
 810                    let old_text = old_snapshot
 811                        .text_for_range(
 812                            Point::new(old_range.start.row, 0)
 813                                ..Point::new(
 814                                    old_range.end.row,
 815                                    old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
 816                                ),
 817                        )
 818                        .collect::<String>();
 819                    let new_text = new_snapshot
 820                        .text_for_range(
 821                            Point::new(new_range.start.row, 0)
 822                                ..Point::new(
 823                                    new_range.end.row,
 824                                    new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
 825                                ),
 826                        )
 827                        .collect::<String>();
 828
 829                    let old_start_row = old_range.start.row;
 830                    let new_start_row = new_range.start.row;
 831                    let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
 832                    let mut inserted_row_ranges = Vec::new();
 833                    for (old_rows, new_rows) in line_diff(&old_text, &new_text) {
 834                        let old_rows = old_start_row + old_rows.start..old_start_row + old_rows.end;
 835                        let new_rows = new_start_row + new_rows.start..new_start_row + new_rows.end;
 836                        if !old_rows.is_empty() {
 837                            deleted_row_ranges.push((
 838                                new_snapshot.anchor_before(Point::new(new_rows.start, 0)),
 839                                old_rows.start..=old_rows.end - 1,
 840                            ));
 841                        }
 842                        if !new_rows.is_empty() {
 843                            let start = new_snapshot.anchor_before(Point::new(new_rows.start, 0));
 844                            let new_end_row = new_rows.end - 1;
 845                            let end = new_snapshot.anchor_before(Point::new(
 846                                new_end_row,
 847                                new_snapshot.line_len(MultiBufferRow(new_end_row)),
 848                            ));
 849                            inserted_row_ranges.push(start..end);
 850                        }
 851                    }
 852                    (deleted_row_ranges, inserted_row_ranges)
 853                })
 854                .await;
 855
 856            codegen
 857                .update(&mut cx, |codegen, cx| {
 858                    codegen.diff.deleted_row_ranges = deleted_row_ranges;
 859                    codegen.diff.inserted_row_ranges = inserted_row_ranges;
 860                    cx.notify();
 861                })
 862                .ok();
 863        })
 864    }
 865}
 866
 867#[derive(Copy, Clone, Debug)]
 868pub enum CodegenEvent {
 869    Finished,
 870    Undone,
 871}
 872
 873struct StripInvalidSpans<T> {
 874    stream: T,
 875    stream_done: bool,
 876    buffer: String,
 877    first_line: bool,
 878    line_end: bool,
 879    starts_with_code_block: bool,
 880}
 881
 882impl<T> StripInvalidSpans<T>
 883where
 884    T: Stream<Item = Result<String>>,
 885{
 886    fn new(stream: T) -> Self {
 887        Self {
 888            stream,
 889            stream_done: false,
 890            buffer: String::new(),
 891            first_line: true,
 892            line_end: false,
 893            starts_with_code_block: false,
 894        }
 895    }
 896}
 897
 898impl<T> Stream for StripInvalidSpans<T>
 899where
 900    T: Stream<Item = Result<String>>,
 901{
 902    type Item = Result<String>;
 903
 904    fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
 905        const CODE_BLOCK_DELIMITER: &str = "```";
 906        const CURSOR_SPAN: &str = "<|CURSOR|>";
 907
 908        let this = unsafe { self.get_unchecked_mut() };
 909        loop {
 910            if !this.stream_done {
 911                let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
 912                match stream.as_mut().poll_next(cx) {
 913                    Poll::Ready(Some(Ok(chunk))) => {
 914                        this.buffer.push_str(&chunk);
 915                    }
 916                    Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
 917                    Poll::Ready(None) => {
 918                        this.stream_done = true;
 919                    }
 920                    Poll::Pending => return Poll::Pending,
 921                }
 922            }
 923
 924            let mut chunk = String::new();
 925            let mut consumed = 0;
 926            if !this.buffer.is_empty() {
 927                let mut lines = this.buffer.split('\n').enumerate().peekable();
 928                while let Some((line_ix, line)) = lines.next() {
 929                    if line_ix > 0 {
 930                        this.first_line = false;
 931                    }
 932
 933                    if this.first_line {
 934                        let trimmed_line = line.trim();
 935                        if lines.peek().is_some() {
 936                            if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
 937                                consumed += line.len() + 1;
 938                                this.starts_with_code_block = true;
 939                                continue;
 940                            }
 941                        } else if trimmed_line.is_empty()
 942                            || prefixes(CODE_BLOCK_DELIMITER)
 943                                .any(|prefix| trimmed_line.starts_with(prefix))
 944                        {
 945                            break;
 946                        }
 947                    }
 948
 949                    let line_without_cursor = line.replace(CURSOR_SPAN, "");
 950                    if lines.peek().is_some() {
 951                        if this.line_end {
 952                            chunk.push('\n');
 953                        }
 954
 955                        chunk.push_str(&line_without_cursor);
 956                        this.line_end = true;
 957                        consumed += line.len() + 1;
 958                    } else if this.stream_done {
 959                        if !this.starts_with_code_block
 960                            || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
 961                        {
 962                            if this.line_end {
 963                                chunk.push('\n');
 964                            }
 965
 966                            chunk.push_str(&line);
 967                        }
 968
 969                        consumed += line.len();
 970                    } else {
 971                        let trimmed_line = line.trim();
 972                        if trimmed_line.is_empty()
 973                            || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
 974                            || prefixes(CODE_BLOCK_DELIMITER)
 975                                .any(|prefix| trimmed_line.ends_with(prefix))
 976                        {
 977                            break;
 978                        } else {
 979                            if this.line_end {
 980                                chunk.push('\n');
 981                                this.line_end = false;
 982                            }
 983
 984                            chunk.push_str(&line_without_cursor);
 985                            consumed += line.len();
 986                        }
 987                    }
 988                }
 989            }
 990
 991            this.buffer = this.buffer.split_off(consumed);
 992            if !chunk.is_empty() {
 993                return Poll::Ready(Some(Ok(chunk)));
 994            } else if this.stream_done {
 995                return Poll::Ready(None);
 996            }
 997        }
 998    }
 999}
1000
1001fn prefixes(text: &str) -> impl Iterator<Item = &str> {
1002    (0..text.len() - 1).map(|ix| &text[..ix + 1])
1003}
1004
1005#[derive(Default)]
1006pub struct Diff {
1007    pub deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1008    pub inserted_row_ranges: Vec<Range<Anchor>>,
1009}
1010
1011impl Diff {
1012    fn is_empty(&self) -> bool {
1013        self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
1014    }
1015}
1016
1017#[cfg(test)]
1018mod tests {
1019    use super::*;
1020    use futures::{
1021        stream::{self},
1022        Stream,
1023    };
1024    use gpui::TestAppContext;
1025    use indoc::indoc;
1026    use language::{
1027        language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
1028        Point,
1029    };
1030    use language_model::LanguageModelRegistry;
1031    use rand::prelude::*;
1032    use serde::Serialize;
1033    use settings::SettingsStore;
1034    use std::{future, sync::Arc};
1035
1036    #[derive(Serialize)]
1037    pub struct DummyCompletionRequest {
1038        pub name: String,
1039    }
1040
1041    #[gpui::test(iterations = 10)]
1042    async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
1043        cx.set_global(cx.update(SettingsStore::test));
1044        cx.update(language_model::LanguageModelRegistry::test);
1045        cx.update(language_settings::init);
1046
1047        let text = indoc! {"
1048            fn main() {
1049                let x = 0;
1050                for _ in 0..10 {
1051                    x += 1;
1052                }
1053            }
1054        "};
1055        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1056        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1057        let range = buffer.read_with(cx, |buffer, cx| {
1058            let snapshot = buffer.snapshot(cx);
1059            snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
1060        });
1061        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1062        let codegen = cx.new(|cx| {
1063            CodegenAlternative::new(
1064                buffer.clone(),
1065                range.clone(),
1066                true,
1067                None,
1068                None,
1069                prompt_builder,
1070                cx,
1071            )
1072        });
1073
1074        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1075
1076        let mut new_text = concat!(
1077            "       let mut x = 0;\n",
1078            "       while x < 10 {\n",
1079            "           x += 1;\n",
1080            "       }",
1081        );
1082        while !new_text.is_empty() {
1083            let max_len = cmp::min(new_text.len(), 10);
1084            let len = rng.gen_range(1..=max_len);
1085            let (chunk, suffix) = new_text.split_at(len);
1086            chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1087            new_text = suffix;
1088            cx.background_executor.run_until_parked();
1089        }
1090        drop(chunks_tx);
1091        cx.background_executor.run_until_parked();
1092
1093        assert_eq!(
1094            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1095            indoc! {"
1096                fn main() {
1097                    let mut x = 0;
1098                    while x < 10 {
1099                        x += 1;
1100                    }
1101                }
1102            "}
1103        );
1104    }
1105
1106    #[gpui::test(iterations = 10)]
1107    async fn test_autoindent_when_generating_past_indentation(
1108        cx: &mut TestAppContext,
1109        mut rng: StdRng,
1110    ) {
1111        cx.set_global(cx.update(SettingsStore::test));
1112        cx.update(language_settings::init);
1113
1114        let text = indoc! {"
1115            fn main() {
1116                le
1117            }
1118        "};
1119        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1120        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1121        let range = buffer.read_with(cx, |buffer, cx| {
1122            let snapshot = buffer.snapshot(cx);
1123            snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
1124        });
1125        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1126        let codegen = cx.new(|cx| {
1127            CodegenAlternative::new(
1128                buffer.clone(),
1129                range.clone(),
1130                true,
1131                None,
1132                None,
1133                prompt_builder,
1134                cx,
1135            )
1136        });
1137
1138        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1139
1140        cx.background_executor.run_until_parked();
1141
1142        let mut new_text = concat!(
1143            "t mut x = 0;\n",
1144            "while x < 10 {\n",
1145            "    x += 1;\n",
1146            "}", //
1147        );
1148        while !new_text.is_empty() {
1149            let max_len = cmp::min(new_text.len(), 10);
1150            let len = rng.gen_range(1..=max_len);
1151            let (chunk, suffix) = new_text.split_at(len);
1152            chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1153            new_text = suffix;
1154            cx.background_executor.run_until_parked();
1155        }
1156        drop(chunks_tx);
1157        cx.background_executor.run_until_parked();
1158
1159        assert_eq!(
1160            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1161            indoc! {"
1162                fn main() {
1163                    let mut x = 0;
1164                    while x < 10 {
1165                        x += 1;
1166                    }
1167                }
1168            "}
1169        );
1170    }
1171
1172    #[gpui::test(iterations = 10)]
1173    async fn test_autoindent_when_generating_before_indentation(
1174        cx: &mut TestAppContext,
1175        mut rng: StdRng,
1176    ) {
1177        cx.update(LanguageModelRegistry::test);
1178        cx.set_global(cx.update(SettingsStore::test));
1179        cx.update(language_settings::init);
1180
1181        let text = concat!(
1182            "fn main() {\n",
1183            "  \n",
1184            "}\n" //
1185        );
1186        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1187        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1188        let range = buffer.read_with(cx, |buffer, cx| {
1189            let snapshot = buffer.snapshot(cx);
1190            snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
1191        });
1192        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1193        let codegen = cx.new(|cx| {
1194            CodegenAlternative::new(
1195                buffer.clone(),
1196                range.clone(),
1197                true,
1198                None,
1199                None,
1200                prompt_builder,
1201                cx,
1202            )
1203        });
1204
1205        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1206
1207        cx.background_executor.run_until_parked();
1208
1209        let mut new_text = concat!(
1210            "let mut x = 0;\n",
1211            "while x < 10 {\n",
1212            "    x += 1;\n",
1213            "}", //
1214        );
1215        while !new_text.is_empty() {
1216            let max_len = cmp::min(new_text.len(), 10);
1217            let len = rng.gen_range(1..=max_len);
1218            let (chunk, suffix) = new_text.split_at(len);
1219            chunks_tx.unbounded_send(chunk.to_string()).unwrap();
1220            new_text = suffix;
1221            cx.background_executor.run_until_parked();
1222        }
1223        drop(chunks_tx);
1224        cx.background_executor.run_until_parked();
1225
1226        assert_eq!(
1227            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1228            indoc! {"
1229                fn main() {
1230                    let mut x = 0;
1231                    while x < 10 {
1232                        x += 1;
1233                    }
1234                }
1235            "}
1236        );
1237    }
1238
1239    #[gpui::test(iterations = 10)]
1240    async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
1241        cx.update(LanguageModelRegistry::test);
1242        cx.set_global(cx.update(SettingsStore::test));
1243        cx.update(language_settings::init);
1244
1245        let text = indoc! {"
1246            func main() {
1247            \tx := 0
1248            \tfor i := 0; i < 10; i++ {
1249            \t\tx++
1250            \t}
1251            }
1252        "};
1253        let buffer = cx.new(|cx| Buffer::local(text, cx));
1254        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1255        let range = buffer.read_with(cx, |buffer, cx| {
1256            let snapshot = buffer.snapshot(cx);
1257            snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
1258        });
1259        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1260        let codegen = cx.new(|cx| {
1261            CodegenAlternative::new(
1262                buffer.clone(),
1263                range.clone(),
1264                true,
1265                None,
1266                None,
1267                prompt_builder,
1268                cx,
1269            )
1270        });
1271
1272        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1273        let new_text = concat!(
1274            "func main() {\n",
1275            "\tx := 0\n",
1276            "\tfor x < 10 {\n",
1277            "\t\tx++\n",
1278            "\t}", //
1279        );
1280        chunks_tx.unbounded_send(new_text.to_string()).unwrap();
1281        drop(chunks_tx);
1282        cx.background_executor.run_until_parked();
1283
1284        assert_eq!(
1285            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1286            indoc! {"
1287                func main() {
1288                \tx := 0
1289                \tfor x < 10 {
1290                \t\tx++
1291                \t}
1292                }
1293            "}
1294        );
1295    }
1296
1297    #[gpui::test]
1298    async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
1299        cx.update(LanguageModelRegistry::test);
1300        cx.set_global(cx.update(SettingsStore::test));
1301        cx.update(language_settings::init);
1302
1303        let text = indoc! {"
1304            fn main() {
1305                let x = 0;
1306            }
1307        "};
1308        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1309        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
1310        let range = buffer.read_with(cx, |buffer, cx| {
1311            let snapshot = buffer.snapshot(cx);
1312            snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
1313        });
1314        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
1315        let codegen = cx.new(|cx| {
1316            CodegenAlternative::new(
1317                buffer.clone(),
1318                range.clone(),
1319                false,
1320                None,
1321                None,
1322                prompt_builder,
1323                cx,
1324            )
1325        });
1326
1327        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
1328        chunks_tx
1329            .unbounded_send("let mut x = 0;\nx += 1;".to_string())
1330            .unwrap();
1331        drop(chunks_tx);
1332        cx.run_until_parked();
1333
1334        // The codegen is inactive, so the buffer doesn't get modified.
1335        assert_eq!(
1336            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1337            text
1338        );
1339
1340        // Activating the codegen applies the changes.
1341        codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
1342        assert_eq!(
1343            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1344            indoc! {"
1345                fn main() {
1346                    let mut x = 0;
1347                    x += 1;
1348                }
1349            "}
1350        );
1351
1352        // Deactivating the codegen undoes the changes.
1353        codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
1354        cx.run_until_parked();
1355        assert_eq!(
1356            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1357            text
1358        );
1359    }
1360
1361    #[gpui::test]
1362    async fn test_strip_invalid_spans_from_codeblock() {
1363        assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
1364        assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
1365        assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
1366        assert_chunks(
1367            "```html\n```js\nLorem ipsum dolor\n```\n```",
1368            "```js\nLorem ipsum dolor\n```",
1369        )
1370        .await;
1371        assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
1372        assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
1373        assert_chunks("Lorem ipsum", "Lorem ipsum").await;
1374        assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
1375
1376        async fn assert_chunks(text: &str, expected_text: &str) {
1377            for chunk_size in 1..=text.len() {
1378                let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
1379                    .map(|chunk| chunk.unwrap())
1380                    .collect::<String>()
1381                    .await;
1382                assert_eq!(
1383                    actual_text, expected_text,
1384                    "failed to strip invalid spans, chunk size: {}",
1385                    chunk_size
1386                );
1387            }
1388        }
1389
1390        fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
1391            stream::iter(
1392                text.chars()
1393                    .collect::<Vec<_>>()
1394                    .chunks(size)
1395                    .map(|chunk| Ok(chunk.iter().collect::<String>()))
1396                    .collect::<Vec<_>>(),
1397            )
1398        }
1399    }
1400
1401    fn simulate_response_stream(
1402        codegen: Entity<CodegenAlternative>,
1403        cx: &mut TestAppContext,
1404    ) -> mpsc::UnboundedSender<String> {
1405        let (chunks_tx, chunks_rx) = mpsc::unbounded();
1406        codegen.update(cx, |codegen, cx| {
1407            codegen.handle_stream(
1408                String::new(),
1409                String::new(),
1410                None,
1411                future::ready(Ok(LanguageModelTextStream {
1412                    message_id: None,
1413                    stream: chunks_rx.map(Ok).boxed(),
1414                })),
1415                cx,
1416            );
1417        });
1418        chunks_tx
1419    }
1420
1421    fn rust_lang() -> Language {
1422        Language::new(
1423            LanguageConfig {
1424                name: "Rust".into(),
1425                matcher: LanguageMatcher {
1426                    path_suffixes: vec!["rs".to_string()],
1427                    ..Default::default()
1428                },
1429                ..Default::default()
1430            },
1431            Some(tree_sitter_rust::LANGUAGE.into()),
1432        )
1433        .with_indents_query(
1434            r#"
1435            (call_expression) @indent
1436            (field_expression) @indent
1437            (_ "(" ")" @end) @indent
1438            (_ "{" "}" @end) @indent
1439            "#,
1440        )
1441        .unwrap()
1442    }
1443}