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