1mod create_file_parser;
2mod edit_parser;
3#[cfg(test)]
4mod evals;
5pub mod 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 thinking_allowed: bool,
85}
86
87impl EditAgent {
88 pub fn new(
89 model: Arc<dyn LanguageModel>,
90 project: Entity<Project>,
91 action_log: Entity<ActionLog>,
92 templates: Arc<Templates>,
93 edit_format: EditFormat,
94 allow_thinking: bool,
95 ) -> Self {
96 EditAgent {
97 model,
98 project,
99 action_log,
100 templates,
101 edit_format,
102 thinking_allowed: allow_thinking,
103 }
104 }
105
106 pub fn overwrite(
107 &self,
108 buffer: Entity<Buffer>,
109 edit_description: String,
110 conversation: &LanguageModelRequest,
111 cx: &mut AsyncApp,
112 ) -> (
113 Task<Result<EditAgentOutput>>,
114 mpsc::UnboundedReceiver<EditAgentOutputEvent>,
115 ) {
116 let this = self.clone();
117 let (events_tx, events_rx) = mpsc::unbounded();
118 let conversation = conversation.clone();
119 let output = cx.spawn(async move |cx| {
120 let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
121 let path = cx.update(|cx| snapshot.resolve_file_path(true, cx));
122 let prompt = CreateFilePromptTemplate {
123 path,
124 edit_description,
125 }
126 .render(&this.templates)?;
127 let new_chunks = this
128 .request(conversation, CompletionIntent::CreateFile, prompt, cx)
129 .await?;
130
131 let (output, mut inner_events) = this.overwrite_with_chunks(buffer, new_chunks, cx);
132 while let Some(event) = inner_events.next().await {
133 events_tx.unbounded_send(event).ok();
134 }
135 output.await
136 });
137 (output, events_rx)
138 }
139
140 fn overwrite_with_chunks(
141 &self,
142 buffer: Entity<Buffer>,
143 edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
144 cx: &mut AsyncApp,
145 ) -> (
146 Task<Result<EditAgentOutput>>,
147 mpsc::UnboundedReceiver<EditAgentOutputEvent>,
148 ) {
149 let (output_events_tx, output_events_rx) = mpsc::unbounded();
150 let (parse_task, parse_rx) = Self::parse_create_file_chunks(edit_chunks, cx);
151 let this = self.clone();
152 let task = cx.spawn(async move |cx| {
153 this.action_log
154 .update(cx, |log, cx| log.buffer_created(buffer.clone(), cx));
155 this.overwrite_with_chunks_internal(buffer, parse_rx, output_events_tx, cx)
156 .await?;
157 parse_task.await
158 });
159 (task, output_events_rx)
160 }
161
162 async fn overwrite_with_chunks_internal(
163 &self,
164 buffer: Entity<Buffer>,
165 mut parse_rx: UnboundedReceiver<Result<CreateFileParserEvent>>,
166 output_events_tx: mpsc::UnboundedSender<EditAgentOutputEvent>,
167 cx: &mut AsyncApp,
168 ) -> Result<()> {
169 cx.update(|cx| {
170 buffer.update(cx, |buffer, cx| buffer.set_text("", cx));
171 self.action_log.update(cx, |log, cx| {
172 log.buffer_edited(buffer.clone(), cx);
173 });
174 self.project.update(cx, |project, cx| {
175 project.set_agent_location(
176 Some(AgentLocation {
177 buffer: buffer.downgrade(),
178 position: language::Anchor::max_for_buffer(buffer.read(cx).remote_id()),
179 }),
180 cx,
181 )
182 });
183 output_events_tx
184 .unbounded_send(EditAgentOutputEvent::Edited(
185 Anchor::min_max_range_for_buffer(buffer.read(cx).remote_id()),
186 ))
187 .ok();
188 });
189
190 while let Some(event) = parse_rx.next().await {
191 match event? {
192 CreateFileParserEvent::NewTextChunk { chunk } => {
193 let buffer_id = cx.update(|cx| {
194 buffer.update(cx, |buffer, cx| buffer.append(chunk, cx));
195 self.action_log
196 .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
197 self.project.update(cx, |project, cx| {
198 project.set_agent_location(
199 Some(AgentLocation {
200 buffer: buffer.downgrade(),
201 position: language::Anchor::max_for_buffer(
202 buffer.read(cx).remote_id(),
203 ),
204 }),
205 cx,
206 )
207 });
208 buffer.read(cx).remote_id()
209 });
210 output_events_tx
211 .unbounded_send(EditAgentOutputEvent::Edited(
212 Anchor::min_max_range_for_buffer(buffer_id),
213 ))
214 .ok();
215 }
216 }
217 }
218
219 Ok(())
220 }
221
222 pub fn edit(
223 &self,
224 buffer: Entity<Buffer>,
225 edit_description: String,
226 conversation: &LanguageModelRequest,
227 cx: &mut AsyncApp,
228 ) -> (
229 Task<Result<EditAgentOutput>>,
230 mpsc::UnboundedReceiver<EditAgentOutputEvent>,
231 ) {
232 let this = self.clone();
233 let (events_tx, events_rx) = mpsc::unbounded();
234 let conversation = conversation.clone();
235 let edit_format = self.edit_format;
236 let output = cx.spawn(async move |cx| {
237 let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
238 let path = cx.update(|cx| snapshot.resolve_file_path(true, cx));
239 let prompt = match edit_format {
240 EditFormat::XmlTags => EditFileXmlPromptTemplate {
241 path,
242 edit_description,
243 }
244 .render(&this.templates)?,
245 EditFormat::DiffFenced => EditFileDiffFencedPromptTemplate {
246 path,
247 edit_description,
248 }
249 .render(&this.templates)?,
250 };
251
252 let edit_chunks = this
253 .request(conversation, CompletionIntent::EditFile, prompt, cx)
254 .await?;
255 this.apply_edit_chunks(buffer, edit_chunks, events_tx, cx)
256 .await
257 });
258 (output, events_rx)
259 }
260
261 async fn apply_edit_chunks(
262 &self,
263 buffer: Entity<Buffer>,
264 edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
265 output_events: mpsc::UnboundedSender<EditAgentOutputEvent>,
266 cx: &mut AsyncApp,
267 ) -> Result<EditAgentOutput> {
268 self.action_log
269 .update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
270
271 let (output, edit_events) = Self::parse_edit_chunks(edit_chunks, self.edit_format, cx);
272 let mut edit_events = edit_events.peekable();
273 while let Some(edit_event) = Pin::new(&mut edit_events).peek().await {
274 // Skip events until we're at the start of a new edit.
275 let Ok(EditParserEvent::OldTextChunk { .. }) = edit_event else {
276 edit_events.next().await.unwrap()?;
277 continue;
278 };
279
280 let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
281
282 // Resolve the old text in the background, updating the agent
283 // location as we keep refining which range it corresponds to.
284 let (resolve_old_text, mut old_range) =
285 Self::resolve_old_text(snapshot.text.clone(), edit_events, cx);
286 while let Ok(old_range) = old_range.recv().await {
287 if let Some(old_range) = old_range {
288 let old_range = snapshot.anchor_before(old_range.start)
289 ..snapshot.anchor_before(old_range.end);
290 self.project.update(cx, |project, cx| {
291 project.set_agent_location(
292 Some(AgentLocation {
293 buffer: buffer.downgrade(),
294 position: old_range.end,
295 }),
296 cx,
297 );
298 });
299 output_events
300 .unbounded_send(EditAgentOutputEvent::ResolvingEditRange(old_range))
301 .ok();
302 }
303 }
304
305 let (edit_events_, mut resolved_old_text) = resolve_old_text.await?;
306 edit_events = edit_events_;
307
308 // If we can't resolve the old text, restart the loop waiting for a
309 // new edit (or for the stream to end).
310 let resolved_old_text = match resolved_old_text.len() {
311 1 => resolved_old_text.pop().unwrap(),
312 0 => {
313 output_events
314 .unbounded_send(EditAgentOutputEvent::UnresolvedEditRange)
315 .ok();
316 continue;
317 }
318 _ => {
319 let ranges = resolved_old_text
320 .into_iter()
321 .map(|text| {
322 let start_line =
323 (snapshot.offset_to_point(text.range.start).row + 1) as usize;
324 let end_line =
325 (snapshot.offset_to_point(text.range.end).row + 1) as usize;
326 start_line..end_line
327 })
328 .collect();
329 output_events
330 .unbounded_send(EditAgentOutputEvent::AmbiguousEditRange(ranges))
331 .ok();
332 continue;
333 }
334 };
335
336 // Compute edits in the background and apply them as they become
337 // available.
338 let (compute_edits, edits) =
339 Self::compute_edits(snapshot, resolved_old_text, edit_events, cx);
340 let mut edits = edits.ready_chunks(32);
341 while let Some(edits) = edits.next().await {
342 if edits.is_empty() {
343 continue;
344 }
345
346 // Edit the buffer and report edits to the action log as part of the
347 // same effect cycle, otherwise the edit will be reported as if the
348 // user made it.
349 let (min_edit_start, max_edit_end) = cx.update(|cx| {
350 let (min_edit_start, max_edit_end) = buffer.update(cx, |buffer, cx| {
351 buffer.edit(edits.iter().cloned(), None, cx);
352 let max_edit_end = buffer
353 .summaries_for_anchors::<Point, _>(
354 edits.iter().map(|(range, _)| &range.end),
355 )
356 .max()
357 .unwrap();
358 let min_edit_start = buffer
359 .summaries_for_anchors::<Point, _>(
360 edits.iter().map(|(range, _)| &range.start),
361 )
362 .min()
363 .unwrap();
364 (
365 buffer.anchor_after(min_edit_start),
366 buffer.anchor_before(max_edit_end),
367 )
368 });
369 self.action_log
370 .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
371 self.project.update(cx, |project, cx| {
372 project.set_agent_location(
373 Some(AgentLocation {
374 buffer: buffer.downgrade(),
375 position: max_edit_end,
376 }),
377 cx,
378 );
379 });
380 (min_edit_start, max_edit_end)
381 });
382 output_events
383 .unbounded_send(EditAgentOutputEvent::Edited(min_edit_start..max_edit_end))
384 .ok();
385 }
386
387 edit_events = compute_edits.await?;
388 }
389
390 output.await
391 }
392
393 fn parse_edit_chunks(
394 chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
395 edit_format: EditFormat,
396 cx: &mut AsyncApp,
397 ) -> (
398 Task<Result<EditAgentOutput>>,
399 UnboundedReceiver<Result<EditParserEvent>>,
400 ) {
401 let (tx, rx) = mpsc::unbounded();
402 let output = cx.background_spawn(async move {
403 pin_mut!(chunks);
404
405 let mut parser = EditParser::new(edit_format);
406 let mut raw_edits = String::new();
407 while let Some(chunk) = chunks.next().await {
408 match chunk {
409 Ok(chunk) => {
410 raw_edits.push_str(&chunk);
411 for event in parser.push(&chunk) {
412 tx.unbounded_send(Ok(event))?;
413 }
414 }
415 Err(error) => {
416 tx.unbounded_send(Err(error.into()))?;
417 }
418 }
419 }
420 Ok(EditAgentOutput {
421 raw_edits,
422 parser_metrics: parser.finish(),
423 })
424 });
425 (output, rx)
426 }
427
428 fn parse_create_file_chunks(
429 chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
430 cx: &mut AsyncApp,
431 ) -> (
432 Task<Result<EditAgentOutput>>,
433 UnboundedReceiver<Result<CreateFileParserEvent>>,
434 ) {
435 let (tx, rx) = mpsc::unbounded();
436 let output = cx.background_spawn(async move {
437 pin_mut!(chunks);
438
439 let mut parser = CreateFileParser::new();
440 let mut raw_edits = String::new();
441 while let Some(chunk) = chunks.next().await {
442 match chunk {
443 Ok(chunk) => {
444 raw_edits.push_str(&chunk);
445 for event in parser.push(Some(&chunk)) {
446 tx.unbounded_send(Ok(event))?;
447 }
448 }
449 Err(error) => {
450 tx.unbounded_send(Err(error.into()))?;
451 }
452 }
453 }
454 // Send final events with None to indicate completion
455 for event in parser.push(None) {
456 tx.unbounded_send(Ok(event))?;
457 }
458 Ok(EditAgentOutput {
459 raw_edits,
460 parser_metrics: EditParserMetrics::default(),
461 })
462 });
463 (output, rx)
464 }
465
466 fn resolve_old_text<T>(
467 snapshot: TextBufferSnapshot,
468 mut edit_events: T,
469 cx: &mut AsyncApp,
470 ) -> (
471 Task<Result<(T, Vec<ResolvedOldText>)>>,
472 watch::Receiver<Option<Range<usize>>>,
473 )
474 where
475 T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,
476 {
477 let (mut old_range_tx, old_range_rx) = watch::channel(None);
478 let task = cx.background_spawn(async move {
479 let mut matcher = StreamingFuzzyMatcher::new(snapshot);
480 while let Some(edit_event) = edit_events.next().await {
481 let EditParserEvent::OldTextChunk {
482 chunk,
483 done,
484 line_hint,
485 } = edit_event?
486 else {
487 break;
488 };
489
490 old_range_tx.send(matcher.push(&chunk, line_hint))?;
491 if done {
492 break;
493 }
494 }
495
496 let matches = matcher.finish();
497 let best_match = matcher.select_best_match();
498
499 old_range_tx.send(best_match.clone())?;
500
501 let indent = LineIndent::from_iter(
502 matcher
503 .query_lines()
504 .first()
505 .unwrap_or(&String::new())
506 .chars(),
507 );
508
509 let resolved_old_texts = if let Some(best_match) = best_match {
510 vec![ResolvedOldText {
511 range: best_match,
512 indent,
513 }]
514 } else {
515 matches
516 .into_iter()
517 .map(|range| ResolvedOldText { range, indent })
518 .collect::<Vec<_>>()
519 };
520
521 Ok((edit_events, resolved_old_texts))
522 });
523
524 (task, old_range_rx)
525 }
526
527 fn compute_edits<T>(
528 snapshot: BufferSnapshot,
529 resolved_old_text: ResolvedOldText,
530 mut edit_events: T,
531 cx: &mut AsyncApp,
532 ) -> (
533 Task<Result<T>>,
534 UnboundedReceiver<(Range<Anchor>, Arc<str>)>,
535 )
536 where
537 T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,
538 {
539 let (edits_tx, edits_rx) = mpsc::unbounded();
540 let compute_edits = cx.background_spawn(async move {
541 let buffer_start_indent = snapshot
542 .line_indent_for_row(snapshot.offset_to_point(resolved_old_text.range.start).row);
543 let indent_delta = if buffer_start_indent.tabs > 0 {
544 IndentDelta::Tabs(
545 buffer_start_indent.tabs as isize - resolved_old_text.indent.tabs as isize,
546 )
547 } else {
548 IndentDelta::Spaces(
549 buffer_start_indent.spaces as isize - resolved_old_text.indent.spaces as isize,
550 )
551 };
552
553 let old_text = snapshot
554 .text_for_range(resolved_old_text.range.clone())
555 .collect::<String>();
556 let mut diff = StreamingDiff::new(old_text);
557 let mut edit_start = resolved_old_text.range.start;
558 let mut new_text_chunks =
559 Self::reindent_new_text_chunks(indent_delta, &mut edit_events);
560 let mut done = false;
561 while !done {
562 let char_operations = if let Some(new_text_chunk) = new_text_chunks.next().await {
563 diff.push_new(&new_text_chunk?)
564 } else {
565 done = true;
566 mem::take(&mut diff).finish()
567 };
568
569 for op in char_operations {
570 match op {
571 CharOperation::Insert { text } => {
572 let edit_start = snapshot.anchor_after(edit_start);
573 edits_tx.unbounded_send((edit_start..edit_start, Arc::from(text)))?;
574 }
575 CharOperation::Delete { bytes } => {
576 let edit_end = edit_start + bytes;
577 let edit_range =
578 snapshot.anchor_after(edit_start)..snapshot.anchor_before(edit_end);
579 edit_start = edit_end;
580 edits_tx.unbounded_send((edit_range, Arc::from("")))?;
581 }
582 CharOperation::Keep { bytes } => edit_start += bytes,
583 }
584 }
585 }
586
587 drop(new_text_chunks);
588 anyhow::Ok(edit_events)
589 });
590
591 (compute_edits, edits_rx)
592 }
593
594 fn reindent_new_text_chunks(
595 delta: IndentDelta,
596 mut stream: impl Unpin + Stream<Item = Result<EditParserEvent>>,
597 ) -> impl Stream<Item = Result<String>> {
598 let mut buffer = String::new();
599 let mut in_leading_whitespace = true;
600 let mut done = false;
601 futures::stream::poll_fn(move |cx| {
602 while !done {
603 let (chunk, is_last_chunk) = match stream.poll_next_unpin(cx) {
604 Poll::Ready(Some(Ok(EditParserEvent::NewTextChunk { chunk, done }))) => {
605 (chunk, done)
606 }
607 Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err))),
608 Poll::Pending => return Poll::Pending,
609 _ => return Poll::Ready(None),
610 };
611
612 buffer.push_str(&chunk);
613
614 let mut indented_new_text = String::new();
615 let mut start_ix = 0;
616 let mut newlines = buffer.match_indices('\n').peekable();
617 loop {
618 let (line_end, is_pending_line) = match newlines.next() {
619 Some((ix, _)) => (ix, false),
620 None => (buffer.len(), true),
621 };
622 let line = &buffer[start_ix..line_end];
623
624 if in_leading_whitespace {
625 if let Some(non_whitespace_ix) = line.find(|c| delta.character() != c) {
626 // We found a non-whitespace character, adjust
627 // indentation based on the delta.
628 let new_indent_len =
629 cmp::max(0, non_whitespace_ix as isize + delta.len()) as usize;
630 indented_new_text
631 .extend(iter::repeat(delta.character()).take(new_indent_len));
632 indented_new_text.push_str(&line[non_whitespace_ix..]);
633 in_leading_whitespace = false;
634 } else if is_pending_line {
635 // We're still in leading whitespace and this line is incomplete.
636 // Stop processing until we receive more input.
637 break;
638 } else {
639 // This line is entirely whitespace. Push it without indentation.
640 indented_new_text.push_str(line);
641 }
642 } else {
643 indented_new_text.push_str(line);
644 }
645
646 if is_pending_line {
647 start_ix = line_end;
648 break;
649 } else {
650 in_leading_whitespace = true;
651 indented_new_text.push('\n');
652 start_ix = line_end + 1;
653 }
654 }
655 buffer.replace_range(..start_ix, "");
656
657 // This was the last chunk, push all the buffered content as-is.
658 if is_last_chunk {
659 indented_new_text.push_str(&buffer);
660 buffer.clear();
661 done = true;
662 }
663
664 if !indented_new_text.is_empty() {
665 return Poll::Ready(Some(Ok(indented_new_text)));
666 }
667 }
668
669 Poll::Ready(None)
670 })
671 }
672
673 async fn request(
674 &self,
675 mut conversation: LanguageModelRequest,
676 intent: CompletionIntent,
677 prompt: String,
678 cx: &mut AsyncApp,
679 ) -> Result<BoxStream<'static, Result<String, LanguageModelCompletionError>>> {
680 let mut messages_iter = conversation.messages.iter_mut();
681 if let Some(last_message) = messages_iter.next_back()
682 && last_message.role == Role::Assistant
683 {
684 let old_content_len = last_message.content.len();
685 last_message
686 .content
687 .retain(|content| !matches!(content, MessageContent::ToolUse(_)));
688 let new_content_len = last_message.content.len();
689
690 // We just removed pending tool uses from the content of the
691 // last message, so it doesn't make sense to cache it anymore
692 // (e.g., the message will look very different on the next
693 // request). Thus, we move the flag to the message prior to it,
694 // as it will still be a valid prefix of the conversation.
695 if old_content_len != new_content_len
696 && last_message.cache
697 && let Some(prev_message) = messages_iter.next_back()
698 {
699 last_message.cache = false;
700 prev_message.cache = true;
701 }
702
703 if last_message.content.is_empty() {
704 conversation.messages.pop();
705 }
706 }
707
708 conversation.messages.push(LanguageModelRequestMessage {
709 role: Role::User,
710 content: vec![MessageContent::Text(prompt)],
711 cache: false,
712 reasoning_details: None,
713 });
714
715 // Include tools in the request so that we can take advantage of
716 // caching when ToolChoice::None is supported.
717 let mut tool_choice = None;
718 let mut tools = Vec::new();
719 if !conversation.tools.is_empty()
720 && self
721 .model
722 .supports_tool_choice(LanguageModelToolChoice::None)
723 {
724 tool_choice = Some(LanguageModelToolChoice::None);
725 tools = conversation.tools.clone();
726 }
727
728 let request = LanguageModelRequest {
729 thread_id: conversation.thread_id,
730 prompt_id: conversation.prompt_id,
731 intent: Some(intent),
732 messages: conversation.messages,
733 tool_choice,
734 tools,
735 stop: Vec::new(),
736 temperature: None,
737 thinking_allowed: self.thinking_allowed,
738 thinking_effort: None,
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 init_test_with_thinking(cx, true).await
1414 }
1415
1416 async fn init_test_with_thinking(cx: &mut TestAppContext, thinking_allowed: bool) -> EditAgent {
1417 cx.update(settings::init);
1418
1419 let project = Project::test(FakeFs::new(cx.executor()), [], cx).await;
1420 let model = Arc::new(FakeLanguageModel::default());
1421 let action_log = cx.new(|_| ActionLog::new(project.clone()));
1422 EditAgent::new(
1423 model,
1424 project,
1425 action_log,
1426 Templates::new(),
1427 EditFormat::XmlTags,
1428 thinking_allowed,
1429 )
1430 }
1431
1432 #[gpui::test(iterations = 10)]
1433 async fn test_non_unique_text_error(cx: &mut TestAppContext, mut rng: StdRng) {
1434 let agent = init_test(cx).await;
1435 let original_text = indoc! {"
1436 function foo() {
1437 return 42;
1438 }
1439
1440 function bar() {
1441 return 42;
1442 }
1443
1444 function baz() {
1445 return 42;
1446 }
1447 "};
1448 let buffer = cx.new(|cx| Buffer::local(original_text, cx));
1449 let (apply, mut events) = agent.edit(
1450 buffer.clone(),
1451 String::new(),
1452 &LanguageModelRequest::default(),
1453 &mut cx.to_async(),
1454 );
1455 cx.run_until_parked();
1456
1457 // When <old_text> matches text in more than one place
1458 simulate_llm_output(
1459 &agent,
1460 indoc! {"
1461 <old_text>
1462 return 42;
1463 }
1464 </old_text>
1465 <new_text>
1466 return 100;
1467 }
1468 </new_text>
1469 "},
1470 &mut rng,
1471 cx,
1472 );
1473 apply.await.unwrap();
1474
1475 // Then the text should remain unchanged
1476 let result_text = buffer.read_with(cx, |buffer, _| buffer.snapshot().text());
1477 assert_eq!(
1478 result_text,
1479 indoc! {"
1480 function foo() {
1481 return 42;
1482 }
1483
1484 function bar() {
1485 return 42;
1486 }
1487
1488 function baz() {
1489 return 42;
1490 }
1491 "},
1492 "Text should remain unchanged when there are multiple matches"
1493 );
1494
1495 // And AmbiguousEditRange even should be emitted
1496 let events = drain_events(&mut events);
1497 let ambiguous_ranges = vec![2..3, 6..7, 10..11];
1498 assert!(
1499 events.contains(&EditAgentOutputEvent::AmbiguousEditRange(ambiguous_ranges)),
1500 "Should emit AmbiguousEditRange for non-unique text"
1501 );
1502 }
1503
1504 #[gpui::test]
1505 async fn test_thinking_allowed_forwarded_to_request(cx: &mut TestAppContext) {
1506 let agent = init_test_with_thinking(cx, false).await;
1507 let buffer = cx.new(|cx| Buffer::local("hello\n", cx));
1508 let (_apply, _events) = agent.edit(
1509 buffer.clone(),
1510 String::new(),
1511 &LanguageModelRequest::default(),
1512 &mut cx.to_async(),
1513 );
1514 cx.run_until_parked();
1515
1516 let pending = agent.model.as_fake().pending_completions();
1517 assert_eq!(pending.len(), 1);
1518 assert!(
1519 !pending[0].thinking_allowed,
1520 "Expected thinking_allowed to be false when EditAgent is constructed with allow_thinking=false"
1521 );
1522 agent.model.as_fake().end_last_completion_stream();
1523
1524 let agent = init_test_with_thinking(cx, true).await;
1525 let buffer = cx.new(|cx| Buffer::local("hello\n", cx));
1526 let (_apply, _events) = agent.edit(
1527 buffer,
1528 String::new(),
1529 &LanguageModelRequest::default(),
1530 &mut cx.to_async(),
1531 );
1532 cx.run_until_parked();
1533
1534 let pending = agent.model.as_fake().pending_completions();
1535 assert_eq!(pending.len(), 1);
1536 assert!(
1537 pending[0].thinking_allowed,
1538 "Expected thinking_allowed to be true when EditAgent is constructed with allow_thinking=true"
1539 );
1540 agent.model.as_fake().end_last_completion_stream();
1541 }
1542
1543 fn drain_events(
1544 stream: &mut UnboundedReceiver<EditAgentOutputEvent>,
1545 ) -> Vec<EditAgentOutputEvent> {
1546 let mut events = Vec::new();
1547 while let Ok(Some(event)) = stream.try_next() {
1548 events.push(event);
1549 }
1550 events
1551 }
1552}