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