edit_agent.rs

   1mod create_file_parser;
   2mod edit_parser;
   3#[cfg(test)]
   4mod evals;
   5mod streaming_fuzzy_matcher;
   6
   7use crate::{Template, Templates};
   8use anyhow::Result;
   9use assistant_tool::ActionLog;
  10use cloud_llm_client::CompletionIntent;
  11use create_file_parser::{CreateFileParser, CreateFileParserEvent};
  12pub use edit_parser::EditFormat;
  13use edit_parser::{EditParser, EditParserEvent, EditParserMetrics};
  14use futures::{
  15    Stream, StreamExt,
  16    channel::mpsc::{self, UnboundedReceiver},
  17    pin_mut,
  18    stream::BoxStream,
  19};
  20use gpui::{AppContext, AsyncApp, Entity, Task};
  21use language::{Anchor, Buffer, BufferSnapshot, LineIndent, Point, TextBufferSnapshot};
  22use language_model::{
  23    LanguageModel, LanguageModelCompletionError, LanguageModelRequest, LanguageModelRequestMessage,
  24    LanguageModelToolChoice, MessageContent, Role,
  25};
  26use project::{AgentLocation, Project};
  27use schemars::JsonSchema;
  28use serde::{Deserialize, Serialize};
  29use std::{cmp, iter, mem, ops::Range, path::PathBuf, pin::Pin, sync::Arc, task::Poll};
  30use streaming_diff::{CharOperation, StreamingDiff};
  31use streaming_fuzzy_matcher::StreamingFuzzyMatcher;
  32use util::debug_panic;
  33
  34#[derive(Serialize)]
  35struct CreateFilePromptTemplate {
  36    path: Option<PathBuf>,
  37    edit_description: String,
  38}
  39
  40impl Template for CreateFilePromptTemplate {
  41    const TEMPLATE_NAME: &'static str = "create_file_prompt.hbs";
  42}
  43
  44#[derive(Serialize)]
  45struct EditFileXmlPromptTemplate {
  46    path: Option<PathBuf>,
  47    edit_description: String,
  48}
  49
  50impl Template for EditFileXmlPromptTemplate {
  51    const TEMPLATE_NAME: &'static str = "edit_file_prompt_xml.hbs";
  52}
  53
  54#[derive(Serialize)]
  55struct EditFileDiffFencedPromptTemplate {
  56    path: Option<PathBuf>,
  57    edit_description: String,
  58}
  59
  60impl Template for EditFileDiffFencedPromptTemplate {
  61    const TEMPLATE_NAME: &'static str = "edit_file_prompt_diff_fenced.hbs";
  62}
  63
  64#[derive(Clone, Debug, PartialEq, Eq)]
  65pub enum EditAgentOutputEvent {
  66    ResolvingEditRange(Range<Anchor>),
  67    UnresolvedEditRange,
  68    AmbiguousEditRange(Vec<Range<usize>>),
  69    Edited,
  70}
  71
  72#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
  73pub struct EditAgentOutput {
  74    pub raw_edits: String,
  75    pub parser_metrics: EditParserMetrics,
  76}
  77
  78#[derive(Clone)]
  79pub struct EditAgent {
  80    model: Arc<dyn LanguageModel>,
  81    action_log: Entity<ActionLog>,
  82    project: Entity<Project>,
  83    templates: Arc<Templates>,
  84    edit_format: EditFormat,
  85}
  86
  87impl EditAgent {
  88    pub fn new(
  89        model: Arc<dyn LanguageModel>,
  90        project: Entity<Project>,
  91        action_log: Entity<ActionLog>,
  92        templates: Arc<Templates>,
  93        edit_format: EditFormat,
  94    ) -> Self {
  95        EditAgent {
  96            model,
  97            project,
  98            action_log,
  99            templates,
 100            edit_format,
 101        }
 102    }
 103
 104    pub fn overwrite(
 105        &self,
 106        buffer: Entity<Buffer>,
 107        edit_description: String,
 108        conversation: &LanguageModelRequest,
 109        cx: &mut AsyncApp,
 110    ) -> (
 111        Task<Result<EditAgentOutput>>,
 112        mpsc::UnboundedReceiver<EditAgentOutputEvent>,
 113    ) {
 114        let this = self.clone();
 115        let (events_tx, events_rx) = mpsc::unbounded();
 116        let conversation = conversation.clone();
 117        let output = cx.spawn(async move |cx| {
 118            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
 119            let path = cx.update(|cx| snapshot.resolve_file_path(cx, true))?;
 120            let prompt = CreateFilePromptTemplate {
 121                path,
 122                edit_description,
 123            }
 124            .render(&this.templates)?;
 125            let new_chunks = this
 126                .request(conversation, CompletionIntent::CreateFile, prompt, cx)
 127                .await?;
 128
 129            let (output, mut inner_events) = this.overwrite_with_chunks(buffer, new_chunks, cx);
 130            while let Some(event) = inner_events.next().await {
 131                events_tx.unbounded_send(event).ok();
 132            }
 133            output.await
 134        });
 135        (output, events_rx)
 136    }
 137
 138    fn overwrite_with_chunks(
 139        &self,
 140        buffer: Entity<Buffer>,
 141        edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
 142        cx: &mut AsyncApp,
 143    ) -> (
 144        Task<Result<EditAgentOutput>>,
 145        mpsc::UnboundedReceiver<EditAgentOutputEvent>,
 146    ) {
 147        let (output_events_tx, output_events_rx) = mpsc::unbounded();
 148        let (parse_task, parse_rx) = Self::parse_create_file_chunks(edit_chunks, cx);
 149        let this = self.clone();
 150        let task = cx.spawn(async move |cx| {
 151            this.action_log
 152                .update(cx, |log, cx| log.buffer_created(buffer.clone(), cx))?;
 153            this.overwrite_with_chunks_internal(buffer, parse_rx, output_events_tx, cx)
 154                .await?;
 155            parse_task.await
 156        });
 157        (task, output_events_rx)
 158    }
 159
 160    async fn overwrite_with_chunks_internal(
 161        &self,
 162        buffer: Entity<Buffer>,
 163        mut parse_rx: UnboundedReceiver<Result<CreateFileParserEvent>>,
 164        output_events_tx: mpsc::UnboundedSender<EditAgentOutputEvent>,
 165        cx: &mut AsyncApp,
 166    ) -> Result<()> {
 167        cx.update(|cx| {
 168            buffer.update(cx, |buffer, cx| buffer.set_text("", cx));
 169            self.action_log.update(cx, |log, cx| {
 170                log.buffer_edited(buffer.clone(), cx);
 171            });
 172            self.project.update(cx, |project, cx| {
 173                project.set_agent_location(
 174                    Some(AgentLocation {
 175                        buffer: buffer.downgrade(),
 176                        position: language::Anchor::MAX,
 177                    }),
 178                    cx,
 179                )
 180            });
 181            output_events_tx
 182                .unbounded_send(EditAgentOutputEvent::Edited)
 183                .ok();
 184        })?;
 185
 186        while let Some(event) = parse_rx.next().await {
 187            match event? {
 188                CreateFileParserEvent::NewTextChunk { chunk } => {
 189                    cx.update(|cx| {
 190                        buffer.update(cx, |buffer, cx| buffer.append(chunk, cx));
 191                        self.action_log
 192                            .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
 193                        self.project.update(cx, |project, cx| {
 194                            project.set_agent_location(
 195                                Some(AgentLocation {
 196                                    buffer: buffer.downgrade(),
 197                                    position: language::Anchor::MAX,
 198                                }),
 199                                cx,
 200                            )
 201                        });
 202                    })?;
 203                    output_events_tx
 204                        .unbounded_send(EditAgentOutputEvent::Edited)
 205                        .ok();
 206                }
 207            }
 208        }
 209
 210        Ok(())
 211    }
 212
 213    pub fn edit(
 214        &self,
 215        buffer: Entity<Buffer>,
 216        edit_description: String,
 217        conversation: &LanguageModelRequest,
 218        cx: &mut AsyncApp,
 219    ) -> (
 220        Task<Result<EditAgentOutput>>,
 221        mpsc::UnboundedReceiver<EditAgentOutputEvent>,
 222    ) {
 223        let this = self.clone();
 224        let (events_tx, events_rx) = mpsc::unbounded();
 225        let conversation = conversation.clone();
 226        let edit_format = self.edit_format;
 227        let output = cx.spawn(async move |cx| {
 228            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
 229            let path = cx.update(|cx| snapshot.resolve_file_path(cx, true))?;
 230            let prompt = match edit_format {
 231                EditFormat::XmlTags => EditFileXmlPromptTemplate {
 232                    path,
 233                    edit_description,
 234                }
 235                .render(&this.templates)?,
 236                EditFormat::DiffFenced => EditFileDiffFencedPromptTemplate {
 237                    path,
 238                    edit_description,
 239                }
 240                .render(&this.templates)?,
 241            };
 242
 243            let edit_chunks = this
 244                .request(conversation, CompletionIntent::EditFile, prompt, cx)
 245                .await?;
 246            this.apply_edit_chunks(buffer, edit_chunks, events_tx, cx)
 247                .await
 248        });
 249        (output, events_rx)
 250    }
 251
 252    async fn apply_edit_chunks(
 253        &self,
 254        buffer: Entity<Buffer>,
 255        edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
 256        output_events: mpsc::UnboundedSender<EditAgentOutputEvent>,
 257        cx: &mut AsyncApp,
 258    ) -> Result<EditAgentOutput> {
 259        self.action_log
 260            .update(cx, |log, cx| log.buffer_read(buffer.clone(), cx))?;
 261
 262        let (output, edit_events) = Self::parse_edit_chunks(edit_chunks, self.edit_format, cx);
 263        let mut edit_events = edit_events.peekable();
 264        while let Some(edit_event) = Pin::new(&mut edit_events).peek().await {
 265            // Skip events until we're at the start of a new edit.
 266            let Ok(EditParserEvent::OldTextChunk { .. }) = edit_event else {
 267                edit_events.next().await.unwrap()?;
 268                continue;
 269            };
 270
 271            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
 272
 273            // Resolve the old text in the background, updating the agent
 274            // location as we keep refining which range it corresponds to.
 275            let (resolve_old_text, mut old_range) =
 276                Self::resolve_old_text(snapshot.text.clone(), edit_events, cx);
 277            while let Ok(old_range) = old_range.recv().await {
 278                if let Some(old_range) = old_range {
 279                    let old_range = snapshot.anchor_before(old_range.start)
 280                        ..snapshot.anchor_before(old_range.end);
 281                    self.project.update(cx, |project, cx| {
 282                        project.set_agent_location(
 283                            Some(AgentLocation {
 284                                buffer: buffer.downgrade(),
 285                                position: old_range.end,
 286                            }),
 287                            cx,
 288                        );
 289                    })?;
 290                    output_events
 291                        .unbounded_send(EditAgentOutputEvent::ResolvingEditRange(old_range))
 292                        .ok();
 293                }
 294            }
 295
 296            let (edit_events_, mut resolved_old_text) = resolve_old_text.await?;
 297            edit_events = edit_events_;
 298
 299            // If we can't resolve the old text, restart the loop waiting for a
 300            // new edit (or for the stream to end).
 301            let resolved_old_text = match resolved_old_text.len() {
 302                1 => resolved_old_text.pop().unwrap(),
 303                0 => {
 304                    output_events
 305                        .unbounded_send(EditAgentOutputEvent::UnresolvedEditRange)
 306                        .ok();
 307                    continue;
 308                }
 309                _ => {
 310                    let ranges = resolved_old_text
 311                        .into_iter()
 312                        .map(|text| {
 313                            let start_line =
 314                                (snapshot.offset_to_point(text.range.start).row + 1) as usize;
 315                            let end_line =
 316                                (snapshot.offset_to_point(text.range.end).row + 1) as usize;
 317                            start_line..end_line
 318                        })
 319                        .collect();
 320                    output_events
 321                        .unbounded_send(EditAgentOutputEvent::AmbiguousEditRange(ranges))
 322                        .ok();
 323                    continue;
 324                }
 325            };
 326
 327            // Compute edits in the background and apply them as they become
 328            // available.
 329            let (compute_edits, edits) =
 330                Self::compute_edits(snapshot, resolved_old_text, edit_events, cx);
 331            let mut edits = edits.ready_chunks(32);
 332            while let Some(edits) = edits.next().await {
 333                if edits.is_empty() {
 334                    continue;
 335                }
 336
 337                // Edit the buffer and report edits to the action log as part of the
 338                // same effect cycle, otherwise the edit will be reported as if the
 339                // user made it.
 340                cx.update(|cx| {
 341                    let max_edit_end = buffer.update(cx, |buffer, cx| {
 342                        buffer.edit(edits.iter().cloned(), None, cx);
 343                        let max_edit_end = buffer
 344                            .summaries_for_anchors::<Point, _>(
 345                                edits.iter().map(|(range, _)| &range.end),
 346                            )
 347                            .max()
 348                            .unwrap();
 349                        buffer.anchor_before(max_edit_end)
 350                    });
 351                    self.action_log
 352                        .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
 353                    self.project.update(cx, |project, cx| {
 354                        project.set_agent_location(
 355                            Some(AgentLocation {
 356                                buffer: buffer.downgrade(),
 357                                position: max_edit_end,
 358                            }),
 359                            cx,
 360                        );
 361                    });
 362                })?;
 363                output_events
 364                    .unbounded_send(EditAgentOutputEvent::Edited)
 365                    .ok();
 366            }
 367
 368            edit_events = compute_edits.await?;
 369        }
 370
 371        output.await
 372    }
 373
 374    fn parse_edit_chunks(
 375        chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
 376        edit_format: EditFormat,
 377        cx: &mut AsyncApp,
 378    ) -> (
 379        Task<Result<EditAgentOutput>>,
 380        UnboundedReceiver<Result<EditParserEvent>>,
 381    ) {
 382        let (tx, rx) = mpsc::unbounded();
 383        let output = cx.background_spawn(async move {
 384            pin_mut!(chunks);
 385
 386            let mut parser = EditParser::new(edit_format);
 387            let mut raw_edits = String::new();
 388            while let Some(chunk) = chunks.next().await {
 389                match chunk {
 390                    Ok(chunk) => {
 391                        raw_edits.push_str(&chunk);
 392                        for event in parser.push(&chunk) {
 393                            tx.unbounded_send(Ok(event))?;
 394                        }
 395                    }
 396                    Err(error) => {
 397                        tx.unbounded_send(Err(error.into()))?;
 398                    }
 399                }
 400            }
 401            Ok(EditAgentOutput {
 402                raw_edits,
 403                parser_metrics: parser.finish(),
 404            })
 405        });
 406        (output, rx)
 407    }
 408
 409    fn parse_create_file_chunks(
 410        chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
 411        cx: &mut AsyncApp,
 412    ) -> (
 413        Task<Result<EditAgentOutput>>,
 414        UnboundedReceiver<Result<CreateFileParserEvent>>,
 415    ) {
 416        let (tx, rx) = mpsc::unbounded();
 417        let output = cx.background_spawn(async move {
 418            pin_mut!(chunks);
 419
 420            let mut parser = CreateFileParser::new();
 421            let mut raw_edits = String::new();
 422            while let Some(chunk) = chunks.next().await {
 423                match chunk {
 424                    Ok(chunk) => {
 425                        raw_edits.push_str(&chunk);
 426                        for event in parser.push(Some(&chunk)) {
 427                            tx.unbounded_send(Ok(event))?;
 428                        }
 429                    }
 430                    Err(error) => {
 431                        tx.unbounded_send(Err(error.into()))?;
 432                    }
 433                }
 434            }
 435            // Send final events with None to indicate completion
 436            for event in parser.push(None) {
 437                tx.unbounded_send(Ok(event))?;
 438            }
 439            Ok(EditAgentOutput {
 440                raw_edits,
 441                parser_metrics: EditParserMetrics::default(),
 442            })
 443        });
 444        (output, rx)
 445    }
 446
 447    fn resolve_old_text<T>(
 448        snapshot: TextBufferSnapshot,
 449        mut edit_events: T,
 450        cx: &mut AsyncApp,
 451    ) -> (
 452        Task<Result<(T, Vec<ResolvedOldText>)>>,
 453        watch::Receiver<Option<Range<usize>>>,
 454    )
 455    where
 456        T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,
 457    {
 458        let (mut old_range_tx, old_range_rx) = watch::channel(None);
 459        let task = cx.background_spawn(async move {
 460            let mut matcher = StreamingFuzzyMatcher::new(snapshot);
 461            while let Some(edit_event) = edit_events.next().await {
 462                let EditParserEvent::OldTextChunk {
 463                    chunk,
 464                    done,
 465                    line_hint,
 466                } = edit_event?
 467                else {
 468                    break;
 469                };
 470
 471                old_range_tx.send(matcher.push(&chunk, line_hint))?;
 472                if done {
 473                    break;
 474                }
 475            }
 476
 477            let matches = matcher.finish();
 478            let best_match = matcher.select_best_match();
 479
 480            old_range_tx.send(best_match.clone())?;
 481
 482            let indent = LineIndent::from_iter(
 483                matcher
 484                    .query_lines()
 485                    .first()
 486                    .unwrap_or(&String::new())
 487                    .chars(),
 488            );
 489
 490            let resolved_old_texts = if let Some(best_match) = best_match {
 491                vec![ResolvedOldText {
 492                    range: best_match,
 493                    indent,
 494                }]
 495            } else {
 496                matches
 497                    .into_iter()
 498                    .map(|range| ResolvedOldText { range, indent })
 499                    .collect::<Vec<_>>()
 500            };
 501
 502            Ok((edit_events, resolved_old_texts))
 503        });
 504
 505        (task, old_range_rx)
 506    }
 507
 508    fn compute_edits<T>(
 509        snapshot: BufferSnapshot,
 510        resolved_old_text: ResolvedOldText,
 511        mut edit_events: T,
 512        cx: &mut AsyncApp,
 513    ) -> (
 514        Task<Result<T>>,
 515        UnboundedReceiver<(Range<Anchor>, Arc<str>)>,
 516    )
 517    where
 518        T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,
 519    {
 520        let (edits_tx, edits_rx) = mpsc::unbounded();
 521        let compute_edits = cx.background_spawn(async move {
 522            let buffer_start_indent = snapshot
 523                .line_indent_for_row(snapshot.offset_to_point(resolved_old_text.range.start).row);
 524            let indent_delta = if buffer_start_indent.tabs > 0 {
 525                IndentDelta::Tabs(
 526                    buffer_start_indent.tabs as isize - resolved_old_text.indent.tabs as isize,
 527                )
 528            } else {
 529                IndentDelta::Spaces(
 530                    buffer_start_indent.spaces as isize - resolved_old_text.indent.spaces as isize,
 531                )
 532            };
 533
 534            let old_text = snapshot
 535                .text_for_range(resolved_old_text.range.clone())
 536                .collect::<String>();
 537            let mut diff = StreamingDiff::new(old_text);
 538            let mut edit_start = resolved_old_text.range.start;
 539            let mut new_text_chunks =
 540                Self::reindent_new_text_chunks(indent_delta, &mut edit_events);
 541            let mut done = false;
 542            while !done {
 543                let char_operations = if let Some(new_text_chunk) = new_text_chunks.next().await {
 544                    diff.push_new(&new_text_chunk?)
 545                } else {
 546                    done = true;
 547                    mem::take(&mut diff).finish()
 548                };
 549
 550                for op in char_operations {
 551                    match op {
 552                        CharOperation::Insert { text } => {
 553                            let edit_start = snapshot.anchor_after(edit_start);
 554                            edits_tx.unbounded_send((edit_start..edit_start, Arc::from(text)))?;
 555                        }
 556                        CharOperation::Delete { bytes } => {
 557                            let edit_end = edit_start + bytes;
 558                            let edit_range =
 559                                snapshot.anchor_after(edit_start)..snapshot.anchor_before(edit_end);
 560                            edit_start = edit_end;
 561                            edits_tx.unbounded_send((edit_range, Arc::from("")))?;
 562                        }
 563                        CharOperation::Keep { bytes } => edit_start += bytes,
 564                    }
 565                }
 566            }
 567
 568            drop(new_text_chunks);
 569            anyhow::Ok(edit_events)
 570        });
 571
 572        (compute_edits, edits_rx)
 573    }
 574
 575    fn reindent_new_text_chunks(
 576        delta: IndentDelta,
 577        mut stream: impl Unpin + Stream<Item = Result<EditParserEvent>>,
 578    ) -> impl Stream<Item = Result<String>> {
 579        let mut buffer = String::new();
 580        let mut in_leading_whitespace = true;
 581        let mut done = false;
 582        futures::stream::poll_fn(move |cx| {
 583            while !done {
 584                let (chunk, is_last_chunk) = match stream.poll_next_unpin(cx) {
 585                    Poll::Ready(Some(Ok(EditParserEvent::NewTextChunk { chunk, done }))) => {
 586                        (chunk, done)
 587                    }
 588                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err))),
 589                    Poll::Pending => return Poll::Pending,
 590                    _ => return Poll::Ready(None),
 591                };
 592
 593                buffer.push_str(&chunk);
 594
 595                let mut indented_new_text = String::new();
 596                let mut start_ix = 0;
 597                let mut newlines = buffer.match_indices('\n').peekable();
 598                loop {
 599                    let (line_end, is_pending_line) = match newlines.next() {
 600                        Some((ix, _)) => (ix, false),
 601                        None => (buffer.len(), true),
 602                    };
 603                    let line = &buffer[start_ix..line_end];
 604
 605                    if in_leading_whitespace {
 606                        if let Some(non_whitespace_ix) = line.find(|c| delta.character() != c) {
 607                            // We found a non-whitespace character, adjust
 608                            // indentation based on the delta.
 609                            let new_indent_len =
 610                                cmp::max(0, non_whitespace_ix as isize + delta.len()) as usize;
 611                            indented_new_text
 612                                .extend(iter::repeat(delta.character()).take(new_indent_len));
 613                            indented_new_text.push_str(&line[non_whitespace_ix..]);
 614                            in_leading_whitespace = false;
 615                        } else if is_pending_line {
 616                            // We're still in leading whitespace and this line is incomplete.
 617                            // Stop processing until we receive more input.
 618                            break;
 619                        } else {
 620                            // This line is entirely whitespace. Push it without indentation.
 621                            indented_new_text.push_str(line);
 622                        }
 623                    } else {
 624                        indented_new_text.push_str(line);
 625                    }
 626
 627                    if is_pending_line {
 628                        start_ix = line_end;
 629                        break;
 630                    } else {
 631                        in_leading_whitespace = true;
 632                        indented_new_text.push('\n');
 633                        start_ix = line_end + 1;
 634                    }
 635                }
 636                buffer.replace_range(..start_ix, "");
 637
 638                // This was the last chunk, push all the buffered content as-is.
 639                if is_last_chunk {
 640                    indented_new_text.push_str(&buffer);
 641                    buffer.clear();
 642                    done = true;
 643                }
 644
 645                if !indented_new_text.is_empty() {
 646                    return Poll::Ready(Some(Ok(indented_new_text)));
 647                }
 648            }
 649
 650            Poll::Ready(None)
 651        })
 652    }
 653
 654    async fn request(
 655        &self,
 656        mut conversation: LanguageModelRequest,
 657        intent: CompletionIntent,
 658        prompt: String,
 659        cx: &mut AsyncApp,
 660    ) -> Result<BoxStream<'static, Result<String, LanguageModelCompletionError>>> {
 661        let mut messages_iter = conversation.messages.iter_mut();
 662        if let Some(last_message) = messages_iter.next_back() {
 663            if last_message.role == Role::Assistant {
 664                let old_content_len = last_message.content.len();
 665                last_message
 666                    .content
 667                    .retain(|content| !matches!(content, MessageContent::ToolUse(_)));
 668                let new_content_len = last_message.content.len();
 669
 670                // We just removed pending tool uses from the content of the
 671                // last message, so it doesn't make sense to cache it anymore
 672                // (e.g., the message will look very different on the next
 673                // request). Thus, we move the flag to the message prior to it,
 674                // as it will still be a valid prefix of the conversation.
 675                if old_content_len != new_content_len && last_message.cache {
 676                    if let Some(prev_message) = messages_iter.next_back() {
 677                        last_message.cache = false;
 678                        prev_message.cache = true;
 679                    }
 680                }
 681
 682                if last_message.content.is_empty() {
 683                    conversation.messages.pop();
 684                }
 685            } else {
 686                debug_panic!(
 687                    "Last message must be an Assistant tool calling! Got {:?}",
 688                    last_message.content
 689                );
 690            }
 691        }
 692
 693        conversation.messages.push(LanguageModelRequestMessage {
 694            role: Role::User,
 695            content: vec![MessageContent::Text(prompt)],
 696            cache: false,
 697        });
 698
 699        // Include tools in the request so that we can take advantage of
 700        // caching when ToolChoice::None is supported.
 701        let mut tool_choice = None;
 702        let mut tools = Vec::new();
 703        if !conversation.tools.is_empty()
 704            && self
 705                .model
 706                .supports_tool_choice(LanguageModelToolChoice::None)
 707        {
 708            tool_choice = Some(LanguageModelToolChoice::None);
 709            tools = conversation.tools.clone();
 710        }
 711
 712        let request = LanguageModelRequest {
 713            thread_id: conversation.thread_id,
 714            prompt_id: conversation.prompt_id,
 715            intent: Some(intent),
 716            mode: conversation.mode,
 717            messages: conversation.messages,
 718            tool_choice,
 719            tools,
 720            stop: Vec::new(),
 721            temperature: None,
 722            thinking_allowed: true,
 723        };
 724
 725        Ok(self.model.stream_completion_text(request, cx).await?.stream)
 726    }
 727}
 728
 729struct ResolvedOldText {
 730    range: Range<usize>,
 731    indent: LineIndent,
 732}
 733
 734#[derive(Copy, Clone, Debug)]
 735enum IndentDelta {
 736    Spaces(isize),
 737    Tabs(isize),
 738}
 739
 740impl IndentDelta {
 741    fn character(&self) -> char {
 742        match self {
 743            IndentDelta::Spaces(_) => ' ',
 744            IndentDelta::Tabs(_) => '\t',
 745        }
 746    }
 747
 748    fn len(&self) -> isize {
 749        match self {
 750            IndentDelta::Spaces(n) => *n,
 751            IndentDelta::Tabs(n) => *n,
 752        }
 753    }
 754}
 755
 756#[cfg(test)]
 757mod tests {
 758    use super::*;
 759    use fs::FakeFs;
 760    use futures::stream;
 761    use gpui::{AppContext, TestAppContext};
 762    use indoc::indoc;
 763    use language_model::fake_provider::FakeLanguageModel;
 764    use project::{AgentLocation, Project};
 765    use rand::prelude::*;
 766    use rand::rngs::StdRng;
 767    use std::cmp;
 768
 769    #[gpui::test(iterations = 100)]
 770    async fn test_empty_old_text(cx: &mut TestAppContext, mut rng: StdRng) {
 771        let agent = init_test(cx).await;
 772        let buffer = cx.new(|cx| {
 773            Buffer::local(
 774                indoc! {"
 775                    abc
 776                    def
 777                    ghi
 778                "},
 779                cx,
 780            )
 781        });
 782        let (apply, _events) = agent.edit(
 783            buffer.clone(),
 784            String::new(),
 785            &LanguageModelRequest::default(),
 786            &mut cx.to_async(),
 787        );
 788        cx.run_until_parked();
 789
 790        simulate_llm_output(
 791            &agent,
 792            indoc! {"
 793                <old_text></old_text>
 794                <new_text>jkl</new_text>
 795                <old_text>def</old_text>
 796                <new_text>DEF</new_text>
 797            "},
 798            &mut rng,
 799            cx,
 800        );
 801        apply.await.unwrap();
 802
 803        pretty_assertions::assert_eq!(
 804            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 805            indoc! {"
 806                abc
 807                DEF
 808                ghi
 809            "}
 810        );
 811    }
 812
 813    #[gpui::test(iterations = 100)]
 814    async fn test_indentation(cx: &mut TestAppContext, mut rng: StdRng) {
 815        let agent = init_test(cx).await;
 816        let buffer = cx.new(|cx| {
 817            Buffer::local(
 818                indoc! {"
 819                    lorem
 820                            ipsum
 821                            dolor
 822                            sit
 823                "},
 824                cx,
 825            )
 826        });
 827        let (apply, _events) = agent.edit(
 828            buffer.clone(),
 829            String::new(),
 830            &LanguageModelRequest::default(),
 831            &mut cx.to_async(),
 832        );
 833        cx.run_until_parked();
 834
 835        simulate_llm_output(
 836            &agent,
 837            indoc! {"
 838                <old_text>
 839                    ipsum
 840                    dolor
 841                    sit
 842                </old_text>
 843                <new_text>
 844                    ipsum
 845                    dolor
 846                    sit
 847                amet
 848                </new_text>
 849            "},
 850            &mut rng,
 851            cx,
 852        );
 853        apply.await.unwrap();
 854
 855        pretty_assertions::assert_eq!(
 856            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 857            indoc! {"
 858                lorem
 859                        ipsum
 860                        dolor
 861                        sit
 862                    amet
 863            "}
 864        );
 865    }
 866
 867    #[gpui::test(iterations = 100)]
 868    async fn test_dependent_edits(cx: &mut TestAppContext, mut rng: StdRng) {
 869        let agent = init_test(cx).await;
 870        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
 871        let (apply, _events) = agent.edit(
 872            buffer.clone(),
 873            String::new(),
 874            &LanguageModelRequest::default(),
 875            &mut cx.to_async(),
 876        );
 877        cx.run_until_parked();
 878
 879        simulate_llm_output(
 880            &agent,
 881            indoc! {"
 882                <old_text>
 883                def
 884                </old_text>
 885                <new_text>
 886                DEF
 887                </new_text>
 888
 889                <old_text>
 890                DEF
 891                </old_text>
 892                <new_text>
 893                DeF
 894                </new_text>
 895            "},
 896            &mut rng,
 897            cx,
 898        );
 899        apply.await.unwrap();
 900
 901        assert_eq!(
 902            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 903            "abc\nDeF\nghi"
 904        );
 905    }
 906
 907    #[gpui::test(iterations = 100)]
 908    async fn test_old_text_hallucination(cx: &mut TestAppContext, mut rng: StdRng) {
 909        let agent = init_test(cx).await;
 910        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
 911        let (apply, _events) = agent.edit(
 912            buffer.clone(),
 913            String::new(),
 914            &LanguageModelRequest::default(),
 915            &mut cx.to_async(),
 916        );
 917        cx.run_until_parked();
 918
 919        simulate_llm_output(
 920            &agent,
 921            indoc! {"
 922                <old_text>
 923                jkl
 924                </old_text>
 925                <new_text>
 926                mno
 927                </new_text>
 928
 929                <old_text>
 930                abc
 931                </old_text>
 932                <new_text>
 933                ABC
 934                </new_text>
 935            "},
 936            &mut rng,
 937            cx,
 938        );
 939        apply.await.unwrap();
 940
 941        assert_eq!(
 942            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 943            "ABC\ndef\nghi"
 944        );
 945    }
 946
 947    #[gpui::test]
 948    async fn test_edit_events(cx: &mut TestAppContext) {
 949        let agent = init_test(cx).await;
 950        let model = agent.model.as_fake();
 951        let project = agent
 952            .action_log
 953            .read_with(cx, |log, _| log.project().clone());
 954        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl", cx));
 955
 956        let mut async_cx = cx.to_async();
 957        let (apply, mut events) = agent.edit(
 958            buffer.clone(),
 959            String::new(),
 960            &LanguageModelRequest::default(),
 961            &mut async_cx,
 962        );
 963        cx.run_until_parked();
 964
 965        model.send_last_completion_stream_text_chunk("<old_text>a");
 966        cx.run_until_parked();
 967        assert_eq!(drain_events(&mut events), vec![]);
 968        assert_eq!(
 969            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 970            "abc\ndef\nghi\njkl"
 971        );
 972        assert_eq!(
 973            project.read_with(cx, |project, _| project.agent_location()),
 974            None
 975        );
 976
 977        model.send_last_completion_stream_text_chunk("bc</old_text>");
 978        cx.run_until_parked();
 979        assert_eq!(
 980            drain_events(&mut events),
 981            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
 982                cx,
 983                |buffer, _| buffer.anchor_before(Point::new(0, 0))
 984                    ..buffer.anchor_before(Point::new(0, 3))
 985            ))]
 986        );
 987        assert_eq!(
 988            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 989            "abc\ndef\nghi\njkl"
 990        );
 991        assert_eq!(
 992            project.read_with(cx, |project, _| project.agent_location()),
 993            Some(AgentLocation {
 994                buffer: buffer.downgrade(),
 995                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
 996            })
 997        );
 998
 999        model.send_last_completion_stream_text_chunk("<new_text>abX");
1000        cx.run_until_parked();
1001        assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
1002        assert_eq!(
1003            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1004            "abXc\ndef\nghi\njkl"
1005        );
1006        assert_eq!(
1007            project.read_with(cx, |project, _| project.agent_location()),
1008            Some(AgentLocation {
1009                buffer: buffer.downgrade(),
1010                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
1011            })
1012        );
1013
1014        model.send_last_completion_stream_text_chunk("cY");
1015        cx.run_until_parked();
1016        assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
1017        assert_eq!(
1018            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1019            "abXcY\ndef\nghi\njkl"
1020        );
1021        assert_eq!(
1022            project.read_with(cx, |project, _| project.agent_location()),
1023            Some(AgentLocation {
1024                buffer: buffer.downgrade(),
1025                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1026            })
1027        );
1028
1029        model.send_last_completion_stream_text_chunk("</new_text>");
1030        model.send_last_completion_stream_text_chunk("<old_text>hall");
1031        cx.run_until_parked();
1032        assert_eq!(drain_events(&mut events), vec![]);
1033        assert_eq!(
1034            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1035            "abXcY\ndef\nghi\njkl"
1036        );
1037        assert_eq!(
1038            project.read_with(cx, |project, _| project.agent_location()),
1039            Some(AgentLocation {
1040                buffer: buffer.downgrade(),
1041                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1042            })
1043        );
1044
1045        model.send_last_completion_stream_text_chunk("ucinated old</old_text>");
1046        model.send_last_completion_stream_text_chunk("<new_text>");
1047        cx.run_until_parked();
1048        assert_eq!(
1049            drain_events(&mut events),
1050            vec![EditAgentOutputEvent::UnresolvedEditRange]
1051        );
1052        assert_eq!(
1053            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1054            "abXcY\ndef\nghi\njkl"
1055        );
1056        assert_eq!(
1057            project.read_with(cx, |project, _| project.agent_location()),
1058            Some(AgentLocation {
1059                buffer: buffer.downgrade(),
1060                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1061            })
1062        );
1063
1064        model.send_last_completion_stream_text_chunk("hallucinated new</new_");
1065        model.send_last_completion_stream_text_chunk("text>");
1066        cx.run_until_parked();
1067        assert_eq!(drain_events(&mut events), vec![]);
1068        assert_eq!(
1069            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1070            "abXcY\ndef\nghi\njkl"
1071        );
1072        assert_eq!(
1073            project.read_with(cx, |project, _| project.agent_location()),
1074            Some(AgentLocation {
1075                buffer: buffer.downgrade(),
1076                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1077            })
1078        );
1079
1080        model.send_last_completion_stream_text_chunk("<old_text>\nghi\nj");
1081        cx.run_until_parked();
1082        assert_eq!(
1083            drain_events(&mut events),
1084            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
1085                cx,
1086                |buffer, _| buffer.anchor_before(Point::new(2, 0))
1087                    ..buffer.anchor_before(Point::new(2, 3))
1088            ))]
1089        );
1090        assert_eq!(
1091            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1092            "abXcY\ndef\nghi\njkl"
1093        );
1094        assert_eq!(
1095            project.read_with(cx, |project, _| project.agent_location()),
1096            Some(AgentLocation {
1097                buffer: buffer.downgrade(),
1098                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
1099            })
1100        );
1101
1102        model.send_last_completion_stream_text_chunk("kl</old_text>");
1103        model.send_last_completion_stream_text_chunk("<new_text>");
1104        cx.run_until_parked();
1105        assert_eq!(
1106            drain_events(&mut events),
1107            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
1108                cx,
1109                |buffer, _| buffer.anchor_before(Point::new(2, 0))
1110                    ..buffer.anchor_before(Point::new(3, 3))
1111            ))]
1112        );
1113        assert_eq!(
1114            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1115            "abXcY\ndef\nghi\njkl"
1116        );
1117        assert_eq!(
1118            project.read_with(cx, |project, _| project.agent_location()),
1119            Some(AgentLocation {
1120                buffer: buffer.downgrade(),
1121                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(3, 3)))
1122            })
1123        );
1124
1125        model.send_last_completion_stream_text_chunk("GHI</new_text>");
1126        cx.run_until_parked();
1127        assert_eq!(
1128            drain_events(&mut events),
1129            vec![EditAgentOutputEvent::Edited]
1130        );
1131        assert_eq!(
1132            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1133            "abXcY\ndef\nGHI"
1134        );
1135        assert_eq!(
1136            project.read_with(cx, |project, _| project.agent_location()),
1137            Some(AgentLocation {
1138                buffer: buffer.downgrade(),
1139                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
1140            })
1141        );
1142
1143        model.end_last_completion_stream();
1144        apply.await.unwrap();
1145        assert_eq!(
1146            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1147            "abXcY\ndef\nGHI"
1148        );
1149        assert_eq!(drain_events(&mut events), vec![]);
1150        assert_eq!(
1151            project.read_with(cx, |project, _| project.agent_location()),
1152            Some(AgentLocation {
1153                buffer: buffer.downgrade(),
1154                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
1155            })
1156        );
1157    }
1158
1159    #[gpui::test]
1160    async fn test_overwrite_events(cx: &mut TestAppContext) {
1161        let agent = init_test(cx).await;
1162        let project = agent
1163            .action_log
1164            .read_with(cx, |log, _| log.project().clone());
1165        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
1166        let (chunks_tx, chunks_rx) = mpsc::unbounded();
1167        let (apply, mut events) = agent.overwrite_with_chunks(
1168            buffer.clone(),
1169            chunks_rx.map(|chunk: &str| Ok(chunk.to_string())),
1170            &mut cx.to_async(),
1171        );
1172
1173        cx.run_until_parked();
1174        assert_eq!(
1175            drain_events(&mut events),
1176            vec![EditAgentOutputEvent::Edited]
1177        );
1178        assert_eq!(
1179            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1180            ""
1181        );
1182        assert_eq!(
1183            project.read_with(cx, |project, _| project.agent_location()),
1184            Some(AgentLocation {
1185                buffer: buffer.downgrade(),
1186                position: language::Anchor::MAX
1187            })
1188        );
1189
1190        chunks_tx.unbounded_send("```\njkl\n").unwrap();
1191        cx.run_until_parked();
1192        assert_eq!(
1193            drain_events(&mut events),
1194            vec![EditAgentOutputEvent::Edited]
1195        );
1196        assert_eq!(
1197            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1198            "jkl"
1199        );
1200        assert_eq!(
1201            project.read_with(cx, |project, _| project.agent_location()),
1202            Some(AgentLocation {
1203                buffer: buffer.downgrade(),
1204                position: language::Anchor::MAX
1205            })
1206        );
1207
1208        chunks_tx.unbounded_send("mno\n").unwrap();
1209        cx.run_until_parked();
1210        assert_eq!(
1211            drain_events(&mut events),
1212            vec![EditAgentOutputEvent::Edited]
1213        );
1214        assert_eq!(
1215            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1216            "jkl\nmno"
1217        );
1218        assert_eq!(
1219            project.read_with(cx, |project, _| project.agent_location()),
1220            Some(AgentLocation {
1221                buffer: buffer.downgrade(),
1222                position: language::Anchor::MAX
1223            })
1224        );
1225
1226        chunks_tx.unbounded_send("pqr\n```").unwrap();
1227        cx.run_until_parked();
1228        assert_eq!(
1229            drain_events(&mut events),
1230            vec![EditAgentOutputEvent::Edited]
1231        );
1232        assert_eq!(
1233            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1234            "jkl\nmno\npqr"
1235        );
1236        assert_eq!(
1237            project.read_with(cx, |project, _| project.agent_location()),
1238            Some(AgentLocation {
1239                buffer: buffer.downgrade(),
1240                position: language::Anchor::MAX
1241            })
1242        );
1243
1244        drop(chunks_tx);
1245        apply.await.unwrap();
1246        assert_eq!(
1247            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1248            "jkl\nmno\npqr"
1249        );
1250        assert_eq!(drain_events(&mut events), vec![]);
1251        assert_eq!(
1252            project.read_with(cx, |project, _| project.agent_location()),
1253            Some(AgentLocation {
1254                buffer: buffer.downgrade(),
1255                position: language::Anchor::MAX
1256            })
1257        );
1258    }
1259
1260    #[gpui::test(iterations = 100)]
1261    async fn test_indent_new_text_chunks(mut rng: StdRng) {
1262        let chunks = to_random_chunks(&mut rng, "    abc\n  def\n      ghi");
1263        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
1264            Ok(EditParserEvent::NewTextChunk {
1265                chunk: chunk.clone(),
1266                done: index == chunks.len() - 1,
1267            })
1268        }));
1269        let indented_chunks =
1270            EditAgent::reindent_new_text_chunks(IndentDelta::Spaces(2), new_text_chunks)
1271                .collect::<Vec<_>>()
1272                .await;
1273        let new_text = indented_chunks
1274            .into_iter()
1275            .collect::<Result<String>>()
1276            .unwrap();
1277        assert_eq!(new_text, "      abc\n    def\n        ghi");
1278    }
1279
1280    #[gpui::test(iterations = 100)]
1281    async fn test_outdent_new_text_chunks(mut rng: StdRng) {
1282        let chunks = to_random_chunks(&mut rng, "\t\t\t\tabc\n\t\tdef\n\t\t\t\t\t\tghi");
1283        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
1284            Ok(EditParserEvent::NewTextChunk {
1285                chunk: chunk.clone(),
1286                done: index == chunks.len() - 1,
1287            })
1288        }));
1289        let indented_chunks =
1290            EditAgent::reindent_new_text_chunks(IndentDelta::Tabs(-2), new_text_chunks)
1291                .collect::<Vec<_>>()
1292                .await;
1293        let new_text = indented_chunks
1294            .into_iter()
1295            .collect::<Result<String>>()
1296            .unwrap();
1297        assert_eq!(new_text, "\t\tabc\ndef\n\t\t\t\tghi");
1298    }
1299
1300    #[gpui::test(iterations = 100)]
1301    async fn test_random_indents(mut rng: StdRng) {
1302        let len = rng.gen_range(1..=100);
1303        let new_text = util::RandomCharIter::new(&mut rng)
1304            .with_simple_text()
1305            .take(len)
1306            .collect::<String>();
1307        let new_text = new_text
1308            .split('\n')
1309            .map(|line| format!("{}{}", " ".repeat(rng.gen_range(0..=8)), line))
1310            .collect::<Vec<_>>()
1311            .join("\n");
1312        let delta = IndentDelta::Spaces(rng.gen_range(-4..=4));
1313
1314        let chunks = to_random_chunks(&mut rng, &new_text);
1315        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
1316            Ok(EditParserEvent::NewTextChunk {
1317                chunk: chunk.clone(),
1318                done: index == chunks.len() - 1,
1319            })
1320        }));
1321        let reindented_chunks = EditAgent::reindent_new_text_chunks(delta, new_text_chunks)
1322            .collect::<Vec<_>>()
1323            .await;
1324        let actual_reindented_text = reindented_chunks
1325            .into_iter()
1326            .collect::<Result<String>>()
1327            .unwrap();
1328        let expected_reindented_text = new_text
1329            .split('\n')
1330            .map(|line| {
1331                if let Some(ix) = line.find(|c| c != ' ') {
1332                    let new_indent = cmp::max(0, ix as isize + delta.len()) as usize;
1333                    format!("{}{}", " ".repeat(new_indent), &line[ix..])
1334                } else {
1335                    line.to_string()
1336                }
1337            })
1338            .collect::<Vec<_>>()
1339            .join("\n");
1340        assert_eq!(actual_reindented_text, expected_reindented_text);
1341    }
1342
1343    fn to_random_chunks(rng: &mut StdRng, input: &str) -> Vec<String> {
1344        let chunk_count = rng.gen_range(1..=cmp::min(input.len(), 50));
1345        let mut chunk_indices = (0..input.len()).choose_multiple(rng, chunk_count);
1346        chunk_indices.sort();
1347        chunk_indices.push(input.len());
1348
1349        let mut chunks = Vec::new();
1350        let mut last_ix = 0;
1351        for chunk_ix in chunk_indices {
1352            chunks.push(input[last_ix..chunk_ix].to_string());
1353            last_ix = chunk_ix;
1354        }
1355        chunks
1356    }
1357
1358    fn simulate_llm_output(
1359        agent: &EditAgent,
1360        output: &str,
1361        rng: &mut StdRng,
1362        cx: &mut TestAppContext,
1363    ) {
1364        let executor = cx.executor();
1365        let chunks = to_random_chunks(rng, output);
1366        let model = agent.model.clone();
1367        cx.background_spawn(async move {
1368            for chunk in chunks {
1369                executor.simulate_random_delay().await;
1370                model
1371                    .as_fake()
1372                    .send_last_completion_stream_text_chunk(chunk);
1373            }
1374            model.as_fake().end_last_completion_stream();
1375        })
1376        .detach();
1377    }
1378
1379    async fn init_test(cx: &mut TestAppContext) -> EditAgent {
1380        cx.update(settings::init);
1381        cx.update(Project::init_settings);
1382        let project = Project::test(FakeFs::new(cx.executor()), [], cx).await;
1383        let model = Arc::new(FakeLanguageModel::default());
1384        let action_log = cx.new(|_| ActionLog::new(project.clone()));
1385        EditAgent::new(
1386            model,
1387            project,
1388            action_log,
1389            Templates::new(),
1390            EditFormat::XmlTags,
1391        )
1392    }
1393
1394    #[gpui::test(iterations = 10)]
1395    async fn test_non_unique_text_error(cx: &mut TestAppContext, mut rng: StdRng) {
1396        let agent = init_test(cx).await;
1397        let original_text = indoc! {"
1398                function foo() {
1399                    return 42;
1400                }
1401
1402                function bar() {
1403                    return 42;
1404                }
1405
1406                function baz() {
1407                    return 42;
1408                }
1409            "};
1410        let buffer = cx.new(|cx| Buffer::local(original_text, cx));
1411        let (apply, mut events) = agent.edit(
1412            buffer.clone(),
1413            String::new(),
1414            &LanguageModelRequest::default(),
1415            &mut cx.to_async(),
1416        );
1417        cx.run_until_parked();
1418
1419        // When <old_text> matches text in more than one place
1420        simulate_llm_output(
1421            &agent,
1422            indoc! {"
1423                <old_text>
1424                    return 42;
1425                }
1426                </old_text>
1427                <new_text>
1428                    return 100;
1429                }
1430                </new_text>
1431            "},
1432            &mut rng,
1433            cx,
1434        );
1435        apply.await.unwrap();
1436
1437        // Then the text should remain unchanged
1438        let result_text = buffer.read_with(cx, |buffer, _| buffer.snapshot().text());
1439        assert_eq!(
1440            result_text,
1441            indoc! {"
1442                function foo() {
1443                    return 42;
1444                }
1445
1446                function bar() {
1447                    return 42;
1448                }
1449
1450                function baz() {
1451                    return 42;
1452                }
1453            "},
1454            "Text should remain unchanged when there are multiple matches"
1455        );
1456
1457        // And AmbiguousEditRange even should be emitted
1458        let events = drain_events(&mut events);
1459        let ambiguous_ranges = vec![2..3, 6..7, 10..11];
1460        assert!(
1461            events.contains(&EditAgentOutputEvent::AmbiguousEditRange(ambiguous_ranges)),
1462            "Should emit AmbiguousEditRange for non-unique text"
1463        );
1464    }
1465
1466    fn drain_events(
1467        stream: &mut UnboundedReceiver<EditAgentOutputEvent>,
1468    ) -> Vec<EditAgentOutputEvent> {
1469        let mut events = Vec::new();
1470        while let Ok(Some(event)) = stream.try_next() {
1471            events.push(event);
1472        }
1473        events
1474    }
1475}