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