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