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