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            // Bypass the rate limiter for nested requests (edit agent requests spawned
 736            // from within a tool call) to avoid deadlocks when multiple subagents try
 737            // to use edit_file simultaneously.
 738            bypass_rate_limit: true,
 739        };
 740
 741        Ok(self.model.stream_completion_text(request, cx).await?.stream)
 742    }
 743}
 744
 745struct ResolvedOldText {
 746    range: Range<usize>,
 747    indent: LineIndent,
 748}
 749
 750#[derive(Copy, Clone, Debug)]
 751enum IndentDelta {
 752    Spaces(isize),
 753    Tabs(isize),
 754}
 755
 756impl IndentDelta {
 757    fn character(&self) -> char {
 758        match self {
 759            IndentDelta::Spaces(_) => ' ',
 760            IndentDelta::Tabs(_) => '\t',
 761        }
 762    }
 763
 764    fn len(&self) -> isize {
 765        match self {
 766            IndentDelta::Spaces(n) => *n,
 767            IndentDelta::Tabs(n) => *n,
 768        }
 769    }
 770}
 771
 772#[cfg(test)]
 773mod tests {
 774    use super::*;
 775    use fs::FakeFs;
 776    use futures::stream;
 777    use gpui::{AppContext, TestAppContext};
 778    use indoc::indoc;
 779    use language_model::fake_provider::FakeLanguageModel;
 780    use pretty_assertions::assert_matches;
 781    use project::{AgentLocation, Project};
 782    use rand::prelude::*;
 783    use rand::rngs::StdRng;
 784    use std::cmp;
 785
 786    #[gpui::test(iterations = 100)]
 787    async fn test_empty_old_text(cx: &mut TestAppContext, mut rng: StdRng) {
 788        let agent = init_test(cx).await;
 789        let buffer = cx.new(|cx| {
 790            Buffer::local(
 791                indoc! {"
 792                    abc
 793                    def
 794                    ghi
 795                "},
 796                cx,
 797            )
 798        });
 799        let (apply, _events) = agent.edit(
 800            buffer.clone(),
 801            String::new(),
 802            &LanguageModelRequest::default(),
 803            &mut cx.to_async(),
 804        );
 805        cx.run_until_parked();
 806
 807        simulate_llm_output(
 808            &agent,
 809            indoc! {"
 810                <old_text></old_text>
 811                <new_text>jkl</new_text>
 812                <old_text>def</old_text>
 813                <new_text>DEF</new_text>
 814            "},
 815            &mut rng,
 816            cx,
 817        );
 818        apply.await.unwrap();
 819
 820        pretty_assertions::assert_eq!(
 821            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 822            indoc! {"
 823                abc
 824                DEF
 825                ghi
 826            "}
 827        );
 828    }
 829
 830    #[gpui::test(iterations = 100)]
 831    async fn test_indentation(cx: &mut TestAppContext, mut rng: StdRng) {
 832        let agent = init_test(cx).await;
 833        let buffer = cx.new(|cx| {
 834            Buffer::local(
 835                indoc! {"
 836                    lorem
 837                            ipsum
 838                            dolor
 839                            sit
 840                "},
 841                cx,
 842            )
 843        });
 844        let (apply, _events) = agent.edit(
 845            buffer.clone(),
 846            String::new(),
 847            &LanguageModelRequest::default(),
 848            &mut cx.to_async(),
 849        );
 850        cx.run_until_parked();
 851
 852        simulate_llm_output(
 853            &agent,
 854            indoc! {"
 855                <old_text>
 856                    ipsum
 857                    dolor
 858                    sit
 859                </old_text>
 860                <new_text>
 861                    ipsum
 862                    dolor
 863                    sit
 864                amet
 865                </new_text>
 866            "},
 867            &mut rng,
 868            cx,
 869        );
 870        apply.await.unwrap();
 871
 872        pretty_assertions::assert_eq!(
 873            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 874            indoc! {"
 875                lorem
 876                        ipsum
 877                        dolor
 878                        sit
 879                    amet
 880            "}
 881        );
 882    }
 883
 884    #[gpui::test(iterations = 100)]
 885    async fn test_dependent_edits(cx: &mut TestAppContext, mut rng: StdRng) {
 886        let agent = init_test(cx).await;
 887        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
 888        let (apply, _events) = agent.edit(
 889            buffer.clone(),
 890            String::new(),
 891            &LanguageModelRequest::default(),
 892            &mut cx.to_async(),
 893        );
 894        cx.run_until_parked();
 895
 896        simulate_llm_output(
 897            &agent,
 898            indoc! {"
 899                <old_text>
 900                def
 901                </old_text>
 902                <new_text>
 903                DEF
 904                </new_text>
 905
 906                <old_text>
 907                DEF
 908                </old_text>
 909                <new_text>
 910                DeF
 911                </new_text>
 912            "},
 913            &mut rng,
 914            cx,
 915        );
 916        apply.await.unwrap();
 917
 918        assert_eq!(
 919            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 920            "abc\nDeF\nghi"
 921        );
 922    }
 923
 924    #[gpui::test(iterations = 100)]
 925    async fn test_old_text_hallucination(cx: &mut TestAppContext, mut rng: StdRng) {
 926        let agent = init_test(cx).await;
 927        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
 928        let (apply, _events) = agent.edit(
 929            buffer.clone(),
 930            String::new(),
 931            &LanguageModelRequest::default(),
 932            &mut cx.to_async(),
 933        );
 934        cx.run_until_parked();
 935
 936        simulate_llm_output(
 937            &agent,
 938            indoc! {"
 939                <old_text>
 940                jkl
 941                </old_text>
 942                <new_text>
 943                mno
 944                </new_text>
 945
 946                <old_text>
 947                abc
 948                </old_text>
 949                <new_text>
 950                ABC
 951                </new_text>
 952            "},
 953            &mut rng,
 954            cx,
 955        );
 956        apply.await.unwrap();
 957
 958        assert_eq!(
 959            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 960            "ABC\ndef\nghi"
 961        );
 962    }
 963
 964    #[gpui::test]
 965    async fn test_edit_events(cx: &mut TestAppContext) {
 966        let agent = init_test(cx).await;
 967        let model = agent.model.as_fake();
 968        let project = agent
 969            .action_log
 970            .read_with(cx, |log, _| log.project().clone());
 971        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl", cx));
 972
 973        let mut async_cx = cx.to_async();
 974        let (apply, mut events) = agent.edit(
 975            buffer.clone(),
 976            String::new(),
 977            &LanguageModelRequest::default(),
 978            &mut async_cx,
 979        );
 980        cx.run_until_parked();
 981
 982        model.send_last_completion_stream_text_chunk("<old_text>a");
 983        cx.run_until_parked();
 984        assert_eq!(drain_events(&mut events), vec![]);
 985        assert_eq!(
 986            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
 987            "abc\ndef\nghi\njkl"
 988        );
 989        assert_eq!(
 990            project.read_with(cx, |project, _| project.agent_location()),
 991            None
 992        );
 993
 994        model.send_last_completion_stream_text_chunk("bc</old_text>");
 995        cx.run_until_parked();
 996        assert_eq!(
 997            drain_events(&mut events),
 998            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
 999                cx,
1000                |buffer, _| buffer.anchor_before(Point::new(0, 0))
1001                    ..buffer.anchor_before(Point::new(0, 3))
1002            ))]
1003        );
1004        assert_eq!(
1005            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1006            "abc\ndef\nghi\njkl"
1007        );
1008        assert_eq!(
1009            project.read_with(cx, |project, _| project.agent_location()),
1010            Some(AgentLocation {
1011                buffer: buffer.downgrade(),
1012                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
1013            })
1014        );
1015
1016        model.send_last_completion_stream_text_chunk("<new_text>abX");
1017        cx.run_until_parked();
1018        assert_matches!(
1019            drain_events(&mut events).as_slice(),
1020            [EditAgentOutputEvent::Edited(_)]
1021        );
1022        assert_eq!(
1023            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1024            "abXc\ndef\nghi\njkl"
1025        );
1026        assert_eq!(
1027            project.read_with(cx, |project, _| project.agent_location()),
1028            Some(AgentLocation {
1029                buffer: buffer.downgrade(),
1030                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
1031            })
1032        );
1033
1034        model.send_last_completion_stream_text_chunk("cY");
1035        cx.run_until_parked();
1036        assert_matches!(
1037            drain_events(&mut events).as_slice(),
1038            [EditAgentOutputEvent::Edited { .. }]
1039        );
1040        assert_eq!(
1041            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1042            "abXcY\ndef\nghi\njkl"
1043        );
1044        assert_eq!(
1045            project.read_with(cx, |project, _| project.agent_location()),
1046            Some(AgentLocation {
1047                buffer: buffer.downgrade(),
1048                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1049            })
1050        );
1051
1052        model.send_last_completion_stream_text_chunk("</new_text>");
1053        model.send_last_completion_stream_text_chunk("<old_text>hall");
1054        cx.run_until_parked();
1055        assert_eq!(drain_events(&mut events), vec![]);
1056        assert_eq!(
1057            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1058            "abXcY\ndef\nghi\njkl"
1059        );
1060        assert_eq!(
1061            project.read_with(cx, |project, _| project.agent_location()),
1062            Some(AgentLocation {
1063                buffer: buffer.downgrade(),
1064                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1065            })
1066        );
1067
1068        model.send_last_completion_stream_text_chunk("ucinated old</old_text>");
1069        model.send_last_completion_stream_text_chunk("<new_text>");
1070        cx.run_until_parked();
1071        assert_eq!(
1072            drain_events(&mut events),
1073            vec![EditAgentOutputEvent::UnresolvedEditRange]
1074        );
1075        assert_eq!(
1076            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1077            "abXcY\ndef\nghi\njkl"
1078        );
1079        assert_eq!(
1080            project.read_with(cx, |project, _| project.agent_location()),
1081            Some(AgentLocation {
1082                buffer: buffer.downgrade(),
1083                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1084            })
1085        );
1086
1087        model.send_last_completion_stream_text_chunk("hallucinated new</new_");
1088        model.send_last_completion_stream_text_chunk("text>");
1089        cx.run_until_parked();
1090        assert_eq!(drain_events(&mut events), vec![]);
1091        assert_eq!(
1092            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1093            "abXcY\ndef\nghi\njkl"
1094        );
1095        assert_eq!(
1096            project.read_with(cx, |project, _| project.agent_location()),
1097            Some(AgentLocation {
1098                buffer: buffer.downgrade(),
1099                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
1100            })
1101        );
1102
1103        model.send_last_completion_stream_text_chunk("<old_text>\nghi\nj");
1104        cx.run_until_parked();
1105        assert_eq!(
1106            drain_events(&mut events),
1107            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
1108                cx,
1109                |buffer, _| buffer.anchor_before(Point::new(2, 0))
1110                    ..buffer.anchor_before(Point::new(2, 3))
1111            ))]
1112        );
1113        assert_eq!(
1114            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1115            "abXcY\ndef\nghi\njkl"
1116        );
1117        assert_eq!(
1118            project.read_with(cx, |project, _| project.agent_location()),
1119            Some(AgentLocation {
1120                buffer: buffer.downgrade(),
1121                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
1122            })
1123        );
1124
1125        model.send_last_completion_stream_text_chunk("kl</old_text>");
1126        model.send_last_completion_stream_text_chunk("<new_text>");
1127        cx.run_until_parked();
1128        assert_eq!(
1129            drain_events(&mut events),
1130            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
1131                cx,
1132                |buffer, _| buffer.anchor_before(Point::new(2, 0))
1133                    ..buffer.anchor_before(Point::new(3, 3))
1134            ))]
1135        );
1136        assert_eq!(
1137            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1138            "abXcY\ndef\nghi\njkl"
1139        );
1140        assert_eq!(
1141            project.read_with(cx, |project, _| project.agent_location()),
1142            Some(AgentLocation {
1143                buffer: buffer.downgrade(),
1144                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(3, 3)))
1145            })
1146        );
1147
1148        model.send_last_completion_stream_text_chunk("GHI</new_text>");
1149        cx.run_until_parked();
1150        assert_matches!(
1151            drain_events(&mut events).as_slice(),
1152            [EditAgentOutputEvent::Edited { .. }]
1153        );
1154        assert_eq!(
1155            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1156            "abXcY\ndef\nGHI"
1157        );
1158        assert_eq!(
1159            project.read_with(cx, |project, _| project.agent_location()),
1160            Some(AgentLocation {
1161                buffer: buffer.downgrade(),
1162                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
1163            })
1164        );
1165
1166        model.end_last_completion_stream();
1167        apply.await.unwrap();
1168        assert_eq!(
1169            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1170            "abXcY\ndef\nGHI"
1171        );
1172        assert_eq!(drain_events(&mut events), vec![]);
1173        assert_eq!(
1174            project.read_with(cx, |project, _| project.agent_location()),
1175            Some(AgentLocation {
1176                buffer: buffer.downgrade(),
1177                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
1178            })
1179        );
1180    }
1181
1182    #[gpui::test]
1183    async fn test_overwrite_events(cx: &mut TestAppContext) {
1184        let agent = init_test(cx).await;
1185        let project = agent
1186            .action_log
1187            .read_with(cx, |log, _| log.project().clone());
1188        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
1189        let (chunks_tx, chunks_rx) = mpsc::unbounded();
1190        let (apply, mut events) = agent.overwrite_with_chunks(
1191            buffer.clone(),
1192            chunks_rx.map(|chunk: &str| Ok(chunk.to_string())),
1193            &mut cx.to_async(),
1194        );
1195
1196        cx.run_until_parked();
1197        assert_matches!(
1198            drain_events(&mut events).as_slice(),
1199            [EditAgentOutputEvent::Edited(_)]
1200        );
1201        assert_eq!(
1202            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1203            ""
1204        );
1205        assert_eq!(
1206            project.read_with(cx, |project, _| project.agent_location()),
1207            Some(AgentLocation {
1208                buffer: buffer.downgrade(),
1209                position: language::Anchor::max_for_buffer(
1210                    cx.update(|cx| buffer.read(cx).remote_id())
1211                ),
1212            })
1213        );
1214
1215        chunks_tx.unbounded_send("```\njkl\n").unwrap();
1216        cx.run_until_parked();
1217        assert_matches!(
1218            drain_events(&mut events).as_slice(),
1219            [EditAgentOutputEvent::Edited { .. }]
1220        );
1221        assert_eq!(
1222            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1223            "jkl"
1224        );
1225        assert_eq!(
1226            project.read_with(cx, |project, _| project.agent_location()),
1227            Some(AgentLocation {
1228                buffer: buffer.downgrade(),
1229                position: language::Anchor::max_for_buffer(
1230                    cx.update(|cx| buffer.read(cx).remote_id())
1231                ),
1232            })
1233        );
1234
1235        chunks_tx.unbounded_send("mno\n").unwrap();
1236        cx.run_until_parked();
1237        assert_matches!(
1238            drain_events(&mut events).as_slice(),
1239            [EditAgentOutputEvent::Edited { .. }]
1240        );
1241        assert_eq!(
1242            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1243            "jkl\nmno"
1244        );
1245        assert_eq!(
1246            project.read_with(cx, |project, _| project.agent_location()),
1247            Some(AgentLocation {
1248                buffer: buffer.downgrade(),
1249                position: language::Anchor::max_for_buffer(
1250                    cx.update(|cx| buffer.read(cx).remote_id())
1251                ),
1252            })
1253        );
1254
1255        chunks_tx.unbounded_send("pqr\n```").unwrap();
1256        cx.run_until_parked();
1257        assert_matches!(
1258            drain_events(&mut events).as_slice(),
1259            [EditAgentOutputEvent::Edited(_)],
1260        );
1261        assert_eq!(
1262            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1263            "jkl\nmno\npqr"
1264        );
1265        assert_eq!(
1266            project.read_with(cx, |project, _| project.agent_location()),
1267            Some(AgentLocation {
1268                buffer: buffer.downgrade(),
1269                position: language::Anchor::max_for_buffer(
1270                    cx.update(|cx| buffer.read(cx).remote_id())
1271                ),
1272            })
1273        );
1274
1275        drop(chunks_tx);
1276        apply.await.unwrap();
1277        assert_eq!(
1278            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
1279            "jkl\nmno\npqr"
1280        );
1281        assert_eq!(drain_events(&mut events), vec![]);
1282        assert_eq!(
1283            project.read_with(cx, |project, _| project.agent_location()),
1284            Some(AgentLocation {
1285                buffer: buffer.downgrade(),
1286                position: language::Anchor::max_for_buffer(
1287                    cx.update(|cx| buffer.read(cx).remote_id())
1288                ),
1289            })
1290        );
1291    }
1292
1293    #[gpui::test(iterations = 100)]
1294    async fn test_indent_new_text_chunks(mut rng: StdRng) {
1295        let chunks = to_random_chunks(&mut rng, "    abc\n  def\n      ghi");
1296        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
1297            Ok(EditParserEvent::NewTextChunk {
1298                chunk: chunk.clone(),
1299                done: index == chunks.len() - 1,
1300            })
1301        }));
1302        let indented_chunks =
1303            EditAgent::reindent_new_text_chunks(IndentDelta::Spaces(2), new_text_chunks)
1304                .collect::<Vec<_>>()
1305                .await;
1306        let new_text = indented_chunks
1307            .into_iter()
1308            .collect::<Result<String>>()
1309            .unwrap();
1310        assert_eq!(new_text, "      abc\n    def\n        ghi");
1311    }
1312
1313    #[gpui::test(iterations = 100)]
1314    async fn test_outdent_new_text_chunks(mut rng: StdRng) {
1315        let chunks = to_random_chunks(&mut rng, "\t\t\t\tabc\n\t\tdef\n\t\t\t\t\t\tghi");
1316        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
1317            Ok(EditParserEvent::NewTextChunk {
1318                chunk: chunk.clone(),
1319                done: index == chunks.len() - 1,
1320            })
1321        }));
1322        let indented_chunks =
1323            EditAgent::reindent_new_text_chunks(IndentDelta::Tabs(-2), new_text_chunks)
1324                .collect::<Vec<_>>()
1325                .await;
1326        let new_text = indented_chunks
1327            .into_iter()
1328            .collect::<Result<String>>()
1329            .unwrap();
1330        assert_eq!(new_text, "\t\tabc\ndef\n\t\t\t\tghi");
1331    }
1332
1333    #[gpui::test(iterations = 100)]
1334    async fn test_random_indents(mut rng: StdRng) {
1335        let len = rng.random_range(1..=100);
1336        let new_text = util::RandomCharIter::new(&mut rng)
1337            .with_simple_text()
1338            .take(len)
1339            .collect::<String>();
1340        let new_text = new_text
1341            .split('\n')
1342            .map(|line| format!("{}{}", " ".repeat(rng.random_range(0..=8)), line))
1343            .collect::<Vec<_>>()
1344            .join("\n");
1345        let delta = IndentDelta::Spaces(rng.random_range(-4i8..=4i8) as isize);
1346
1347        let chunks = to_random_chunks(&mut rng, &new_text);
1348        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
1349            Ok(EditParserEvent::NewTextChunk {
1350                chunk: chunk.clone(),
1351                done: index == chunks.len() - 1,
1352            })
1353        }));
1354        let reindented_chunks = EditAgent::reindent_new_text_chunks(delta, new_text_chunks)
1355            .collect::<Vec<_>>()
1356            .await;
1357        let actual_reindented_text = reindented_chunks
1358            .into_iter()
1359            .collect::<Result<String>>()
1360            .unwrap();
1361        let expected_reindented_text = new_text
1362            .split('\n')
1363            .map(|line| {
1364                if let Some(ix) = line.find(|c| c != ' ') {
1365                    let new_indent = cmp::max(0, ix as isize + delta.len()) as usize;
1366                    format!("{}{}", " ".repeat(new_indent), &line[ix..])
1367                } else {
1368                    line.to_string()
1369                }
1370            })
1371            .collect::<Vec<_>>()
1372            .join("\n");
1373        assert_eq!(actual_reindented_text, expected_reindented_text);
1374    }
1375
1376    fn to_random_chunks(rng: &mut StdRng, input: &str) -> Vec<String> {
1377        let chunk_count = rng.random_range(1..=cmp::min(input.len(), 50));
1378        let mut chunk_indices = (0..input.len()).choose_multiple(rng, chunk_count);
1379        chunk_indices.sort();
1380        chunk_indices.push(input.len());
1381
1382        let mut chunks = Vec::new();
1383        let mut last_ix = 0;
1384        for chunk_ix in chunk_indices {
1385            chunks.push(input[last_ix..chunk_ix].to_string());
1386            last_ix = chunk_ix;
1387        }
1388        chunks
1389    }
1390
1391    fn simulate_llm_output(
1392        agent: &EditAgent,
1393        output: &str,
1394        rng: &mut StdRng,
1395        cx: &mut TestAppContext,
1396    ) {
1397        let executor = cx.executor();
1398        let chunks = to_random_chunks(rng, output);
1399        let model = agent.model.clone();
1400        cx.background_spawn(async move {
1401            for chunk in chunks {
1402                executor.simulate_random_delay().await;
1403                model
1404                    .as_fake()
1405                    .send_last_completion_stream_text_chunk(chunk);
1406            }
1407            model.as_fake().end_last_completion_stream();
1408        })
1409        .detach();
1410    }
1411
1412    async fn init_test(cx: &mut TestAppContext) -> EditAgent {
1413        cx.update(settings::init);
1414
1415        let project = Project::test(FakeFs::new(cx.executor()), [], cx).await;
1416        let model = Arc::new(FakeLanguageModel::default());
1417        let action_log = cx.new(|_| ActionLog::new(project.clone()));
1418        EditAgent::new(
1419            model,
1420            project,
1421            action_log,
1422            Templates::new(),
1423            EditFormat::XmlTags,
1424        )
1425    }
1426
1427    #[gpui::test(iterations = 10)]
1428    async fn test_non_unique_text_error(cx: &mut TestAppContext, mut rng: StdRng) {
1429        let agent = init_test(cx).await;
1430        let original_text = indoc! {"
1431                function foo() {
1432                    return 42;
1433                }
1434
1435                function bar() {
1436                    return 42;
1437                }
1438
1439                function baz() {
1440                    return 42;
1441                }
1442            "};
1443        let buffer = cx.new(|cx| Buffer::local(original_text, cx));
1444        let (apply, mut events) = agent.edit(
1445            buffer.clone(),
1446            String::new(),
1447            &LanguageModelRequest::default(),
1448            &mut cx.to_async(),
1449        );
1450        cx.run_until_parked();
1451
1452        // When <old_text> matches text in more than one place
1453        simulate_llm_output(
1454            &agent,
1455            indoc! {"
1456                <old_text>
1457                    return 42;
1458                }
1459                </old_text>
1460                <new_text>
1461                    return 100;
1462                }
1463                </new_text>
1464            "},
1465            &mut rng,
1466            cx,
1467        );
1468        apply.await.unwrap();
1469
1470        // Then the text should remain unchanged
1471        let result_text = buffer.read_with(cx, |buffer, _| buffer.snapshot().text());
1472        assert_eq!(
1473            result_text,
1474            indoc! {"
1475                function foo() {
1476                    return 42;
1477                }
1478
1479                function bar() {
1480                    return 42;
1481                }
1482
1483                function baz() {
1484                    return 42;
1485                }
1486            "},
1487            "Text should remain unchanged when there are multiple matches"
1488        );
1489
1490        // And AmbiguousEditRange even should be emitted
1491        let events = drain_events(&mut events);
1492        let ambiguous_ranges = vec![2..3, 6..7, 10..11];
1493        assert!(
1494            events.contains(&EditAgentOutputEvent::AmbiguousEditRange(ambiguous_ranges)),
1495            "Should emit AmbiguousEditRange for non-unique text"
1496        );
1497    }
1498
1499    fn drain_events(
1500        stream: &mut UnboundedReceiver<EditAgentOutputEvent>,
1501    ) -> Vec<EditAgentOutputEvent> {
1502        let mut events = Vec::new();
1503        while let Ok(Some(event)) = stream.try_next() {
1504            events.push(event);
1505        }
1506        events
1507    }
1508}