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