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