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