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