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