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