tenant.openagents/omega

1use language::{Point, TextBufferSnapshot};
2use std::{cmp, ops::Range};
3
4const REPLACEMENT_COST: u32 = 1;
5const INSERTION_COST: u32 = 3;
6const DELETION_COST: u32 = 10;
7// Two matches are enough to prove ambiguity, so stop scanning early; exact
8// matches are byte-identical and can't be ranked.
9const MAX_EXACT_MATCHES: usize = 2;
10
11/// A streaming fuzzy matcher that can process text chunks incrementally
12/// and return the best match found so far at each step.
13pub struct StreamingFuzzyMatcher {
14    snapshot: TextBufferSnapshot,
15    query_lines: Vec<String>,
16    line_hint: Option<u32>,
17    incomplete_line: String,
18    raw_query: String,
19    matches: Vec<SearchMatch>,
20    matrix: SearchMatrix,
21}
22
23/// A match candidate: the matched byte range plus the 0-based
24/// `(query_row, buffer_row)` line pairs the search aligned to produce it.
25#[derive(Clone, Debug)]
26pub(super) struct SearchMatch {
27    pub(super) range: Range<usize>,
28    pub(super) line_pairs: Vec<(u32, u32)>,
29}
30
31#[derive(Debug)]
32pub(super) enum SearchMatches {
33    Exact(Vec<SearchMatch>),
34    Fuzzy(Vec<SearchMatch>),
35}
36
37impl StreamingFuzzyMatcher {
38    pub fn new(snapshot: TextBufferSnapshot) -> Self {
39        let buffer_line_count = snapshot.max_point().row as usize + 1;
40        Self {
41            snapshot,
42            query_lines: Vec::new(),
43            line_hint: None,
44            incomplete_line: String::new(),
45            raw_query: String::new(),
46            matches: Vec::new(),
47            matrix: SearchMatrix::new(buffer_line_count + 1),
48        }
49    }
50
51    /// Returns the query lines.
52    pub fn query_lines(&self) -> &[String] {
53        &self.query_lines
54    }
55
56    /// Push a new chunk of text and get the best match found so far.
57    ///
58    /// This method accumulates text chunks and processes complete lines.
59    /// Partial lines are buffered internally until a newline is received.
60    ///
61    /// # Returns
62    ///
63    /// Returns `Some(range)` if a match has been found with the accumulated
64    /// query so far, or `None` if no suitable match exists yet.
65    pub fn push(&mut self, chunk: &str, line_hint: Option<u32>) -> Option<Range<usize>> {
66        // Add the chunk to our incomplete line buffer
67        self.raw_query.push_str(chunk);
68        self.incomplete_line.push_str(chunk);
69        self.line_hint = line_hint;
70
71        if let Some((last_pos, _)) = self.incomplete_line.match_indices('\n').next_back() {
72            let complete_part = &self.incomplete_line[..=last_pos];
73
74            // Split into lines and add to query_lines
75            for line in complete_part.lines() {
76                self.query_lines.push(line.to_string());
77            }
78
79            self.incomplete_line.replace_range(..last_pos + 1, "");
80
81            self.matches = self.resolve_location_fuzzy();
82        }
83
84        let best_match = self.select_best_match();
85        best_match.or_else(|| {
86            self.matches
87                .first()
88                .map(|search_match| search_match.range.clone())
89        })
90    }
91
92    /// Finish processing and return the final best match(es).
93    ///
94    /// This processes any remaining incomplete line before returning the final
95    /// match result.
96    pub fn finish(&mut self) -> SearchMatches {
97        let exact_ranges = find_exact_matches(&self.snapshot, &self.raw_query);
98        if !exact_ranges.is_empty() {
99            if !self.incomplete_line.is_empty() {
100                self.query_lines
101                    .push(std::mem::take(&mut self.incomplete_line));
102            }
103            let matches = exact_ranges
104                .into_iter()
105                .map(|range| {
106                    let start_row = self.snapshot.offset_to_point(range.start).row;
107                    let line_pairs = (0..self.query_lines.len())
108                        .map(|query_row| (query_row as u32, start_row + query_row as u32))
109                        .collect();
110                    SearchMatch { range, line_pairs }
111                })
112                .collect();
113            return SearchMatches::Exact(matches);
114        }
115
116        // Process any remaining incomplete line
117        if !self.incomplete_line.is_empty() {
118            if let [only_match] = self.matches.as_mut_slice() {
119                let range = &mut only_match.range;
120                if range.end < self.snapshot.len()
121                    && self
122                        .snapshot
123                        .contains_str_at(range.end + 1, &self.incomplete_line)
124                {
125                    range.end += 1 + self.incomplete_line.len();
126                    // Record the line and its alignment so that `query_lines`
127                    // and `line_pairs` stay in sync with the lines covered by
128                    // the returned range.
129                    let extended_row = self.snapshot.offset_to_point(range.end).row;
130                    self.query_lines
131                        .push(std::mem::take(&mut self.incomplete_line));
132                    only_match
133                        .line_pairs
134                        .push(((self.query_lines.len() - 1) as u32, extended_row));
135                    return SearchMatches::Fuzzy(self.matches.clone());
136                }
137            }
138
139            self.query_lines
140                .push(std::mem::take(&mut self.incomplete_line));
141            self.matches = self.resolve_location_fuzzy();
142        }
143        SearchMatches::Fuzzy(self.matches.clone())
144    }
145
146    fn resolve_location_fuzzy(&mut self) -> Vec<SearchMatch> {
147        let new_query_line_count = self.query_lines.len();
148        let old_query_line_count = self.matrix.rows.saturating_sub(1);
149        if new_query_line_count == old_query_line_count {
150            return Vec::new();
151        }
152
153        self.matrix.resize_rows(new_query_line_count + 1);
154
155        // Process only the new query lines
156        for row in old_query_line_count..new_query_line_count {
157            let query_line = self.query_lines[row].trim();
158            let leading_deletion_cost = (row + 1) as u32 * DELETION_COST;
159
160            self.matrix.set(
161                row + 1,
162                0,
163                SearchState::new(leading_deletion_cost, SearchDirection::Up),
164            );
165
166            let mut buffer_lines = self.snapshot.as_rope().chunks().lines();
167            let mut col = 0;
168            while let Some(buffer_line) = buffer_lines.next() {
169                let buffer_line = buffer_line.trim();
170                let up = SearchState::new(
171                    self.matrix
172                        .get(row, col + 1)
173                        .cost
174                        .saturating_add(DELETION_COST),
175                    SearchDirection::Up,
176                );
177                let left = SearchState::new(
178                    self.matrix
179                        .get(row + 1, col)
180                        .cost
181                        .saturating_add(INSERTION_COST),
182                    SearchDirection::Left,
183                );
184                let diagonal = SearchState::new(
185                    if query_line == buffer_line {
186                        self.matrix.get(row, col).cost
187                    } else if fuzzy_eq(query_line, buffer_line) {
188                        self.matrix.get(row, col).cost + REPLACEMENT_COST
189                    } else {
190                        self.matrix
191                            .get(row, col)
192                            .cost
193                            .saturating_add(DELETION_COST + INSERTION_COST)
194                    },
195                    SearchDirection::Diagonal,
196                );
197                self.matrix
198                    .set(row + 1, col + 1, up.min(left).min(diagonal));
199                col += 1;
200            }
201        }
202
203        // Find all matches with the best cost
204        let buffer_line_count = self.snapshot.max_point().row as usize + 1;
205        let mut best_cost = u32::MAX;
206        let mut matches_with_best_cost = Vec::new();
207
208        for col in 1..=buffer_line_count {
209            let cost = self.matrix.get(new_query_line_count, col).cost;
210            if cost < best_cost {
211                best_cost = cost;
212                matches_with_best_cost.clear();
213                matches_with_best_cost.push(col as u32);
214            } else if cost == best_cost {
215                matches_with_best_cost.push(col as u32);
216            }
217        }
218
219        // Find ranges for the matches
220        let mut valid_matches = Vec::new();
221        for &buffer_row_end in &matches_with_best_cost {
222            let mut line_pairs = Vec::new();
223            let mut query_row = new_query_line_count;
224            let mut buffer_row_start = buffer_row_end;
225            while query_row > 0 && buffer_row_start > 0 {
226                let current = self.matrix.get(query_row, buffer_row_start as usize);
227                match current.direction {
228                    SearchDirection::Diagonal => {
229                        query_row -= 1;
230                        buffer_row_start -= 1;
231                        line_pairs.push((query_row as u32, buffer_row_start));
232                    }
233                    SearchDirection::Up => {
234                        query_row -= 1;
235                    }
236                    SearchDirection::Left => {
237                        buffer_row_start -= 1;
238                    }
239                }
240            }
241            line_pairs.reverse();
242
243            let matched_buffer_row_count = buffer_row_end - buffer_row_start;
244            let matched_ratio = line_pairs.len() as f32
245                / (matched_buffer_row_count as f32).max(new_query_line_count as f32);
246            if matched_ratio >= 0.8 {
247                let buffer_start_ix = self
248                    .snapshot
249                    .point_to_offset(Point::new(buffer_row_start, 0));
250                let buffer_end_ix = self.snapshot.point_to_offset(Point::new(
251                    buffer_row_end - 1,
252                    self.snapshot.line_len(buffer_row_end - 1),
253                ));
254                valid_matches.push(SearchMatch {
255                    range: buffer_start_ix..buffer_end_ix,
256                    line_pairs,
257                });
258            }
259        }
260
261        valid_matches
262    }
263
264    /// Return the best match with starting position close enough to line_hint.
265    pub fn select_best_match(&self) -> Option<Range<usize>> {
266        // Allow line hint to be off by that many lines.
267        // Higher values increase probability of applying edits to a wrong place,
268        // Lower values increase edits failures and overall conversation length.
269        const LINE_HINT_TOLERANCE: u32 = 200;
270
271        if self.matches.is_empty() {
272            return None;
273        }
274
275        if let [only_match] = self.matches.as_slice() {
276            return Some(only_match.range.clone());
277        }
278
279        let Some(line_hint) = self.line_hint else {
280            // Multiple ambiguous matches
281            return None;
282        };
283
284        let mut best_match = None;
285        let mut best_distance = u32::MAX;
286
287        for search_match in &self.matches {
288            let start_point = self.snapshot.offset_to_point(search_match.range.start);
289            let start_line = start_point.row;
290            let distance = start_line.abs_diff(line_hint);
291
292            if distance <= LINE_HINT_TOLERANCE && distance < best_distance {
293                best_distance = distance;
294                best_match = Some(search_match.range.clone());
295            }
296        }
297
298        best_match
299    }
300}
301
302// Aho-Corasick's overlapping search requires a contiguous haystack, while its
303// streaming search skips overlapping matches. KMP detects overlapping ambiguity
304// while scanning rope chunks without copying the buffer.
305fn find_exact_matches(snapshot: &TextBufferSnapshot, query: &str) -> Vec<Range<usize>> {
306    if query.is_empty() {
307        return Vec::new();
308    }
309
310    let query = query.as_bytes();
311    let trailing_line_ending_len = if query.ends_with(b"\r\n") {
312        2
313    } else if query.ends_with(b"\n") {
314        1
315    } else {
316        0
317    };
318    let mut prefix_lengths = vec![0; query.len()];
319    let mut prefix_length = 0;
320    for query_index in 1..query.len() {
321        while prefix_length > 0 && query[query_index] != query[prefix_length] {
322            prefix_length = prefix_lengths[prefix_length - 1];
323        }
324        if query[query_index] == query[prefix_length] {
325            prefix_length += 1;
326            prefix_lengths[query_index] = prefix_length;
327        }
328    }
329
330    let mut matches = Vec::new();
331    let mut matched_length = 0;
332    for (offset, byte) in snapshot
333        .bytes_in_range(0..snapshot.len())
334        .flatten()
335        .copied()
336        .enumerate()
337    {
338        while matched_length > 0 && byte != query[matched_length] {
339            matched_length = prefix_lengths[matched_length - 1];
340        }
341        if byte == query[matched_length] {
342            matched_length += 1;
343        }
344        if matched_length == query.len() {
345            let raw_end = offset + 1;
346            let start = raw_end - query.len();
347            let end = raw_end - trailing_line_ending_len;
348            matches.push(start..end);
349            if matches.len() == MAX_EXACT_MATCHES {
350                break;
351            }
352            matched_length = prefix_lengths[matched_length - 1];
353        }
354    }
355    matches
356}
357
358fn fuzzy_eq(left: &str, right: &str) -> bool {
359    const THRESHOLD: f64 = 0.8;
360
361    let min_levenshtein = left.len().abs_diff(right.len());
362    let min_normalized_levenshtein =
363        1. - (min_levenshtein as f64 / cmp::max(left.len(), right.len()) as f64);
364    if min_normalized_levenshtein < THRESHOLD {
365        return false;
366    }
367
368    strsim::normalized_levenshtein(left, right) >= THRESHOLD
369}
370
371#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
372enum SearchDirection {
373    Up,
374    Left,
375    Diagonal,
376}
377
378#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
379struct SearchState {
380    cost: u32,
381    direction: SearchDirection,
382}
383
384impl SearchState {
385    fn new(cost: u32, direction: SearchDirection) -> Self {
386        Self { cost, direction }
387    }
388}
389
390struct SearchMatrix {
391    cols: usize,
392    rows: usize,
393    data: Vec<SearchState>,
394}
395
396impl SearchMatrix {
397    fn new(cols: usize) -> Self {
398        SearchMatrix {
399            cols,
400            rows: 0,
401            data: Vec::new(),
402        }
403    }
404
405    fn resize_rows(&mut self, needed_rows: usize) {
406        debug_assert!(needed_rows > self.rows);
407        self.rows = needed_rows;
408        self.data.resize(
409            self.rows * self.cols,
410            SearchState::new(0, SearchDirection::Diagonal),
411        );
412    }
413
414    fn get(&self, row: usize, col: usize) -> SearchState {
415        debug_assert!(row < self.rows && col < self.cols);
416        self.data[row * self.cols + col]
417    }
418
419    fn set(&mut self, row: usize, col: usize, state: SearchState) {
420        debug_assert!(row < self.rows && col < self.cols);
421        self.data[row * self.cols + col] = state;
422    }
423}
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428    use indoc::indoc;
429    use language::{BufferId, TextBuffer};
430    use rand::prelude::*;
431    use text::ReplicaId;
432    use util::test::{generate_marked_text, marked_text_ranges};
433
434    #[test]
435    fn test_empty_query() {
436        let buffer = TextBuffer::new(
437            ReplicaId::LOCAL,
438            BufferId::new(1).unwrap(),
439            "Hello world\nThis is a test\nFoo bar baz",
440        );
441        let snapshot = buffer.snapshot();
442
443        let mut finder = StreamingFuzzyMatcher::new(snapshot.clone());
444        assert_eq!(push(&mut finder, ""), None);
445        assert_eq!(finish(finder), None);
446    }
447
448    #[test]
449    fn test_streaming_exact_match() {
450        let buffer = TextBuffer::new(
451            ReplicaId::LOCAL,
452            BufferId::new(1).unwrap(),
453            "Hello world\nThis is a test\nFoo bar baz",
454        );
455        let snapshot = buffer.snapshot();
456
457        let mut finder = StreamingFuzzyMatcher::new(snapshot.clone());
458
459        // Push partial query
460        assert_eq!(push(&mut finder, "This"), None);
461
462        // Complete the line
463        assert_eq!(
464            push(&mut finder, " is a test\n"),
465            Some("This is a test".to_string())
466        );
467
468        // Finish should return the same result
469        assert_eq!(finish(finder), Some("This is a test".to_string()));
470    }
471
472    #[test]
473    fn test_streaming_fuzzy_match() {
474        let buffer = TextBuffer::new(
475            ReplicaId::LOCAL,
476            BufferId::new(1).unwrap(),
477            indoc! {"
478                function foo(a, b) {
479                    return a + b;
480                }
481
482                function bar(x, y) {
483                    return x * y;
484                }
485            "},
486        );
487        let snapshot = buffer.snapshot();
488
489        let mut finder = StreamingFuzzyMatcher::new(snapshot.clone());
490
491        // Push a fuzzy query that should match the first function
492        assert_eq!(
493            push(&mut finder, "function foo(a, c) {\n").as_deref(),
494            Some("function foo(a, b) {")
495        );
496        assert_eq!(
497            push(&mut finder, "    return a + c;\n}\n").as_deref(),
498            Some(concat!(
499                "function foo(a, b) {\n",
500                "    return a + b;\n",
501                "}"
502            ))
503        );
504    }
505
506    #[test]
507    fn test_incremental_improvement() {
508        let buffer = TextBuffer::new(
509            ReplicaId::LOCAL,
510            BufferId::new(1).unwrap(),
511            "Line 1\nLine 2\nLine 3\nLine 4\nLine 5",
512        );
513        let snapshot = buffer.snapshot();
514
515        let mut finder = StreamingFuzzyMatcher::new(snapshot.clone());
516
517        // No match initially
518        assert_eq!(push(&mut finder, "Lin"), None);
519
520        // Get a match when we complete a line
521        assert_eq!(push(&mut finder, "e 3\n"), Some("Line 3".to_string()));
522
523        // The match might change if we add more specific content
524        assert_eq!(
525            push(&mut finder, "Line 4\n"),
526            Some("Line 3\nLine 4".to_string())
527        );
528        assert_eq!(finish(finder), Some("Line 3\nLine 4".to_string()));
529    }
530
531    #[test]
532    fn test_incomplete_lines_buffering() {
533        let buffer = TextBuffer::new(
534            ReplicaId::LOCAL,
535            BufferId::new(1).unwrap(),
536            indoc! {"
537                The quick brown fox
538                jumps over the lazy dog
539                Pack my box with five dozen liquor jugs
540            "},
541        );
542        let snapshot = buffer.snapshot();
543
544        let mut finder = StreamingFuzzyMatcher::new(snapshot.clone());
545
546        // Push text in small chunks across line boundaries
547        assert_eq!(push(&mut finder, "jumps "), None); // No newline yet
548        assert_eq!(push(&mut finder, "over the"), None); // Still no newline
549        assert_eq!(push(&mut finder, " lazy"), None); // Still incomplete
550
551        // Complete the line
552        assert_eq!(
553            push(&mut finder, " dog\n"),
554            Some("jumps over the lazy dog".to_string())
555        );
556    }
557
558    #[test]
559    fn test_multiline_fuzzy_match() {
560        let buffer = TextBuffer::new(
561            ReplicaId::LOCAL,
562            BufferId::new(1).unwrap(),
563            indoc! {r#"
564                impl Display for User {
565                    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
566                        write!(f, "User: {} ({})", self.name, self.email)
567                    }
568                }
569
570                impl Debug for User {
571                    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
572                        f.debug_struct("User")
573                            .field("name", &self.name)
574                            .field("email", &self.email)
575                            .finish()
576                    }
577                }
578            "#},
579        );
580        let snapshot = buffer.snapshot();
581
582        let mut finder = StreamingFuzzyMatcher::new(snapshot.clone());
583
584        assert_eq!(
585            push(&mut finder, "impl Debug for User {\n"),
586            Some("impl Debug for User {".to_string())
587        );
588        assert_eq!(
589            push(
590                &mut finder,
591                "    fn fmt(&self, f: &mut Formatter) -> Result {\n"
592            )
593            .as_deref(),
594            Some(concat!(
595                "impl Debug for User {\n",
596                "    fn fmt(&self, f: &mut Formatter) -> fmt::Result {"
597            ))
598        );
599        assert_eq!(
600            push(&mut finder, "        f.debug_struct(\"User\")\n").as_deref(),
601            Some(concat!(
602                "impl Debug for User {\n",
603                "    fn fmt(&self, f: &mut Formatter) -> fmt::Result {\n",
604                "        f.debug_struct(\"User\")"
605            ))
606        );
607        assert_eq!(
608            push(
609                &mut finder,
610                "            .field(\"name\", &self.username)\n"
611            )
612            .as_deref(),
613            Some(concat!(
614                "impl Debug for User {\n",
615                "    fn fmt(&self, f: &mut Formatter) -> fmt::Result {\n",
616                "        f.debug_struct(\"User\")\n",
617                "            .field(\"name\", &self.name)"
618            ))
619        );
620        assert_eq!(
621            finish(finder).as_deref(),
622            Some(concat!(
623                "impl Debug for User {\n",
624                "    fn fmt(&self, f: &mut Formatter) -> fmt::Result {\n",
625                "        f.debug_struct(\"User\")\n",
626                "            .field(\"name\", &self.name)"
627            ))
628        );
629    }
630
631    #[gpui::test(iterations = 100)]
632    fn test_resolve_location_single_line(mut rng: StdRng) {
633        assert_location_resolution(
634            concat!(
635                "    Lorem\n",
636                "    «ipsum»\n",
637                "    dolor sit amet\n",
638                "    consecteur",
639            ),
640            "ipsum",
641            &mut rng,
642        );
643    }
644
645    #[gpui::test(iterations = 100)]
646    fn test_resolve_location_multiline(mut rng: StdRng) {
647        assert_location_resolution(
648            concat!(
649                "    Lorem\n",
650                "«    ipsum\n",
651                "    dolor sit amet»\n",
652                "    consecteur",
653            ),
654            "ipsum\ndolor sit amet",
655            &mut rng,
656        );
657    }
658
659    #[gpui::test(iterations = 100)]
660    fn test_resolve_location_function_with_typo(mut rng: StdRng) {
661        assert_location_resolution(
662            indoc! {"
663                «fn foo1(a: usize) -> usize {
664                    40
665                }»
666
667                fn foo2(b: usize) -> usize {
668                    42
669                }
670            "},
671            "fn foo1(a: usize) -> u32 {\n40\n}",
672            &mut rng,
673        );
674    }
675
676    #[gpui::test(iterations = 100)]
677    fn test_resolve_location_class_methods(mut rng: StdRng) {
678        assert_location_resolution(
679            indoc! {"
680                class Something {
681                    one() { return 1; }
682                «    two() { return 2222; }
683                    three() { return 333; }
684                    four() { return 4444; }
685                    five() { return 5555; }
686                    six() { return 6666; }»
687                    seven() { return 7; }
688                    eight() { return 8; }
689                }
690            "},
691            indoc! {"
692                two() { return 2222; }
693                four() { return 4444; }
694                five() { return 5555; }
695                six() { return 6666; }
696            "},
697            &mut rng,
698        );
699    }
700
701    #[gpui::test(iterations = 100)]
702    fn test_resolve_location_imports_no_match(mut rng: StdRng) {
703        assert_location_resolution(
704            indoc! {"
705                use std::ops::Range;
706                use std::sync::Mutex;
707                use std::{
708                    collections::HashMap,
709                    env,
710                    ffi::{OsStr, OsString},
711                    fs,
712                    io::{BufRead, BufReader},
713                    mem,
714                    path::{Path, PathBuf},
715                    process::Command,
716                    sync::LazyLock,
717                    time::SystemTime,
718                };
719            "},
720            indoc! {"
721                use std::collections::{HashMap, HashSet};
722                use std::ffi::{OsStr, OsString};
723                use std::fmt::Write as _;
724                use std::fs;
725                use std::io::{BufReader, Read, Write};
726                use std::mem;
727                use std::path::{Path, PathBuf};
728                use std::process::Command;
729                use std::sync::Arc;
730            "},
731            &mut rng,
732        );
733    }
734
735    #[gpui::test(iterations = 100)]
736    fn test_resolve_location_nested_closure(mut rng: StdRng) {
737        assert_location_resolution(
738            indoc! {"
739                impl Foo {
740                    fn new() -> Self {
741                        Self {
742                            subscriptions: vec![
743                                cx.observe_window_activation(window, |editor, window, cx| {
744                                    let active = window.is_window_active();
745                                    editor.blink_manager.update(cx, |blink_manager, cx| {
746                                        if active {
747                                            blink_manager.enable(cx);
748                                        } else {
749                                            blink_manager.disable(cx);
750                                        }
751                                    });
752                                }),
753                            ];
754                        }
755                    }
756                }
757            "},
758            concat!(
759                "                    editor.blink_manager.update(cx, |blink_manager, cx| {\n",
760                "                        blink_manager.enable(cx);\n",
761                "                    });",
762            ),
763            &mut rng,
764        );
765    }
766
767    #[gpui::test(iterations = 100)]
768    fn test_resolve_location_tool_invocation(mut rng: StdRng) {
769        assert_location_resolution(
770            indoc! {r#"
771                let tool = cx
772                    .update(|cx| working_set.tool(&tool_name, cx))
773                    .map_err(|err| {
774                        anyhow!("Failed to look up tool '{}': {}", tool_name, err)
775                    })?;
776
777                let Some(tool) = tool else {
778                    return Err(anyhow!("Tool '{}' not found", tool_name));
779                };
780
781                let project = project.clone();
782                let action_log = action_log.clone();
783                let messages = messages.clone();
784                let tool_result = cx
785                    .update(|cx| tool.run(invocation.input, &messages, project, action_log, cx))
786                    .map_err(|err| anyhow!("Failed to start tool '{}': {}", tool_name, err))?;
787
788                tasks.push(tool_result.output);
789            "#},
790            concat!(
791                "let tool_result = cx\n",
792                "    .update(|cx| tool.run(invocation.input, &messages, project, action_log, cx))\n",
793                "    .output;",
794            ),
795            &mut rng,
796        );
797    }
798
799    #[gpui::test]
800    fn test_line_hint_selection() {
801        let text = indoc! {r#"
802            fn first_function() {
803                return 42;
804            }
805
806            fn second_function() {
807                return 42;
808            }
809
810            fn third_function() {
811                return 42;
812            }
813        "#};
814
815        let buffer = TextBuffer::new(
816            ReplicaId::LOCAL,
817            BufferId::new(1).unwrap(),
818            text.to_string(),
819        );
820        let snapshot = buffer.snapshot();
821        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
822
823        // Given a query that matches all three functions
824        let query = "return 42;\n";
825
826        // Test with line hint pointing to second function (around line 5)
827        let best_match = matcher.push(query, Some(5)).expect("Failed to match query");
828
829        let matched_text = snapshot
830            .text_for_range(best_match.clone())
831            .collect::<String>();
832        assert!(matched_text.contains("return 42;"));
833        assert_eq!(
834            best_match,
835            63..77,
836            "Expected to match `second_function` based on the line hint"
837        );
838
839        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
840        matcher.push(query, None);
841        matcher.finish();
842        let best_match = matcher.select_best_match();
843        assert!(
844            best_match.is_none(),
845            "Best match should be None when query cannot be uniquely resolved"
846        );
847    }
848
849    #[gpui::test]
850    fn test_exact_match_takes_precedence_over_fuzzy_match() {
851        let buffer = TextBuffer::new(
852            ReplicaId::LOCAL,
853            BufferId::new(1).unwrap(),
854            concat!(
855                "prefix keyboard WASD, voxel-based suffix\n",
856                "keyboard WASD, voxel-baseX\n",
857            ),
858        );
859        let snapshot = buffer.snapshot();
860        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
861
862        assert_eq!(matcher.push("keyboard WASD, voxel-based", None), None);
863        let SearchMatches::Exact(matches) = matcher.finish() else {
864            panic!("expected an exact match");
865        };
866        let [search_match] = matches.as_slice() else {
867            panic!("expected one match, got {}", matches.len());
868        };
869        assert_eq!(
870            snapshot
871                .text_for_range(search_match.range.clone())
872                .collect::<String>(),
873            "keyboard WASD, voxel-based"
874        );
875    }
876
877    #[gpui::test]
878    fn test_exact_match_uses_trailing_newline_to_disambiguate() {
879        let buffer = TextBuffer::new(
880            ReplicaId::LOCAL,
881            BufferId::new(1).unwrap(),
882            "foo suffix\nfoo\n",
883        );
884        let snapshot = buffer.snapshot();
885        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
886
887        matcher.push("foo\n", None);
888        let SearchMatches::Exact(matches) = matcher.finish() else {
889            panic!("expected an exact match");
890        };
891        let [search_match] = matches.as_slice() else {
892            panic!("expected one match, got {}", matches.len());
893        };
894        assert_eq!(
895            snapshot
896                .text_for_range(search_match.range.clone())
897                .collect::<String>(),
898            "foo"
899        );
900    }
901
902    #[gpui::test]
903    fn test_exact_newline_only_match_excludes_line_ending() {
904        let buffer = TextBuffer::new(ReplicaId::LOCAL, BufferId::new(1).unwrap(), "a\nb");
905        let mut matcher = StreamingFuzzyMatcher::new(buffer.snapshot().clone());
906
907        matcher.push("\n", None);
908        let SearchMatches::Exact(matches) = matcher.finish() else {
909            panic!("expected an exact match");
910        };
911        let [search_match] = matches.as_slice() else {
912            panic!("expected one match, got {}", matches.len());
913        };
914        assert_eq!(search_match.range, 1..1);
915    }
916
917    #[gpui::test]
918    fn test_exact_overlapping_matches_are_ambiguous() {
919        let buffer = TextBuffer::new(ReplicaId::LOCAL, BufferId::new(1).unwrap(), "aaaaa");
920        let mut matcher = StreamingFuzzyMatcher::new(buffer.snapshot().clone());
921
922        matcher.push("aaaa", None);
923        let matches = matcher.finish();
924
925        assert!(matches!(matches, SearchMatches::Exact(_)));
926        assert_eq!(match_ranges(&matches), vec![0..4, 1..5]);
927    }
928
929    #[gpui::test]
930    fn test_exact_multiline_match_does_not_extend_incomplete_line() {
931        let buffer = TextBuffer::new(
932            ReplicaId::LOCAL,
933            BufferId::new(1).unwrap(),
934            "prefix fragment\nnext\nnext\n",
935        );
936        let snapshot = buffer.snapshot();
937        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
938
939        assert_eq!(matcher.push("fragment\nnext", None), None);
940        let SearchMatches::Exact(matches) = matcher.finish() else {
941            panic!("expected an exact match");
942        };
943        let [search_match] = matches.as_slice() else {
944            panic!("expected one match, got {}", matches.len());
945        };
946        assert_eq!(
947            snapshot
948                .text_for_range(search_match.range.clone())
949                .collect::<String>(),
950            "fragment\nnext"
951        );
952    }
953
954    #[gpui::test]
955    fn test_prefix_of_last_line_resolves_to_correct_range() {
956        let text = indoc! {r#"
957            fn on_query_change(&mut self, cx: &mut Context<Self>) {
958                self.filter(cx);
959            }
960
961
962
963            fn render_search(&self, cx: &mut Context<Self>) -> Div {
964                div()
965            }
966        "#};
967
968        let buffer = TextBuffer::new(
969            ReplicaId::LOCAL,
970            BufferId::new(1).unwrap(),
971            text.to_string(),
972        );
973        let snapshot = buffer.snapshot();
974
975        // Query with a partial last line. This is a verbatim substring of the
976        // buffer, so it resolves through the exact-match path.
977        let query = "}\n\n\n\nfn render_search";
978
979        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
980        matcher.push(query, None);
981        let matches = matcher.finish();
982        let matched = search_matches(&matches);
983
984        // The match should include the line containing "fn render_search".
985        let matched_text = matched.first().map(|search_match| {
986            snapshot
987                .text_for_range(search_match.range.clone())
988                .collect::<String>()
989        });
990
991        assert!(
992            matched.len() == 1,
993            "Expected exactly one match, got {}: {:?}",
994            matched.len(),
995            matched_text,
996        );
997
998        let Some(matched_text) = matched_text else {
999            panic!("expected a match");
1000        };
1001        pretty_assertions::assert_eq!(
1002            matched_text,
1003            "}\n\n\n\nfn render_search",
1004            "Match should include the render_search line",
1005        );
1006    }
1007
1008    #[track_caller]
1009    fn assert_location_resolution(text_with_expected_range: &str, query: &str, rng: &mut StdRng) {
1010        let (text, expected_ranges) = marked_text_ranges(text_with_expected_range, false);
1011        let buffer = TextBuffer::new(ReplicaId::LOCAL, BufferId::new(1).unwrap(), text.clone());
1012        let snapshot = buffer.snapshot();
1013
1014        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
1015
1016        // Split query into random chunks
1017        let chunks = to_random_chunks(rng, query);
1018
1019        // Push chunks incrementally
1020        for chunk in &chunks {
1021            matcher.push(chunk, None);
1022        }
1023
1024        let actual_matches = matcher.finish();
1025        let actual_ranges = match_ranges(&actual_matches);
1026
1027        // If no expected ranges, we expect no match
1028        if expected_ranges.is_empty() {
1029            assert!(
1030                actual_ranges.is_empty(),
1031                "Expected no match for query: {:?}, but found: {:?}",
1032                query,
1033                actual_ranges
1034            );
1035        } else {
1036            let text_with_actual_range = generate_marked_text(&text, &actual_ranges, false);
1037            pretty_assertions::assert_eq!(
1038                text_with_actual_range,
1039                text_with_expected_range,
1040                indoc! {"
1041                    Query: {:?}
1042                    Chunks: {:?}
1043                    Expected marked text: {}
1044                    Actual marked text: {}
1045                    Expected ranges: {:?}
1046                    Actual ranges: {:?}"
1047                },
1048                query,
1049                chunks,
1050                text_with_expected_range,
1051                text_with_actual_range,
1052                expected_ranges,
1053                actual_ranges
1054            );
1055        }
1056    }
1057
1058    #[test]
1059    fn test_line_pairs_skip_unmatched_buffer_line() {
1060        let text = indoc! {r#"
1061            class Outer:
1062                def method(self):
1063                    self.kept = "unchanged"
1064                    self.target_a = "before"
1065                    self.extra = "row"
1066                    self.target_b = "before"
1067                    self.target_c = "before"
1068                    self.target_d = "before"
1069                    self.kept_2 = "unchanged"
1070        "#};
1071        let buffer = TextBuffer::new(
1072            ReplicaId::LOCAL,
1073            BufferId::new(1).unwrap(),
1074            text.to_string(),
1075        );
1076        let mut matcher = StreamingFuzzyMatcher::new(buffer.snapshot().clone());
1077
1078        // The query omits the `self.extra` row that sits between the matched
1079        // buffer lines.
1080        matcher.push(
1081            concat!(
1082                "        self.target_a = \"before\"\n",
1083                "        self.target_b = \"before\"\n",
1084                "        self.target_c = \"before\"\n",
1085                "        self.target_d = \"before\"\n",
1086            ),
1087            None,
1088        );
1089        let SearchMatches::Fuzzy(matches) = matcher.finish() else {
1090            panic!("expected a fuzzy match");
1091        };
1092        let [search_match] = matches.as_slice() else {
1093            panic!("expected one match, got {}", matches.len());
1094        };
1095        assert_eq!(search_match.line_pairs, [(0, 3), (1, 5), (2, 6), (3, 7)]);
1096    }
1097
1098    #[test]
1099    fn test_line_pairs_include_extended_incomplete_line() {
1100        let text = indoc! {r#"
1101            fn on_query_change(&mut self, cx: &mut Context<Self>) {
1102                self.filter(cx);
1103            }
1104
1105
1106
1107            fn render_search(&self, cx: &mut Context<Self>) -> Div {
1108                div()
1109            }
1110        "#};
1111        let buffer = TextBuffer::new(
1112            ReplicaId::LOCAL,
1113            BufferId::new(1).unwrap(),
1114            text.to_string(),
1115        );
1116        let mut matcher = StreamingFuzzyMatcher::new(buffer.snapshot().clone());
1117
1118        // The last query line is incomplete and gets appended to the match by
1119        // `finish` via verbatim comparison rather than the fuzzy search. The
1120        // trailing space after `}` keeps the query from being an exact
1121        // substring of the buffer, so the fuzzy extension path is exercised.
1122        matcher.push("} \n\n\n\nfn render_search", None);
1123        let SearchMatches::Fuzzy(matches) = matcher.finish() else {
1124            panic!("expected a fuzzy match");
1125        };
1126        let [search_match] = matches.as_slice() else {
1127            panic!("expected one match, got {}", matches.len());
1128        };
1129        assert_eq!(
1130            search_match.line_pairs,
1131            [(0, 2), (1, 3), (2, 4), (3, 5), (4, 6)]
1132        );
1133        assert_eq!(matcher.query_lines().len(), 5);
1134    }
1135
1136    fn to_random_chunks(rng: &mut StdRng, input: &str) -> Vec<String> {
1137        let chunk_count = rng.random_range(1..=cmp::min(input.len(), 50));
1138        let mut chunk_indices = (0..input.len()).choose_multiple(rng, chunk_count);
1139        chunk_indices.sort();
1140        chunk_indices.push(input.len());
1141
1142        let mut chunks = Vec::new();
1143        let mut last_ix = 0;
1144        for chunk_ix in chunk_indices {
1145            chunks.push(input[last_ix..chunk_ix].to_string());
1146            last_ix = chunk_ix;
1147        }
1148        chunks
1149    }
1150
1151    fn push(finder: &mut StreamingFuzzyMatcher, chunk: &str) -> Option<String> {
1152        finder
1153            .push(chunk, None)
1154            .map(|range| finder.snapshot.text_for_range(range).collect::<String>())
1155    }
1156
1157    fn search_matches(matches: &SearchMatches) -> &[SearchMatch] {
1158        match matches {
1159            SearchMatches::Exact(matches) | SearchMatches::Fuzzy(matches) => matches,
1160        }
1161    }
1162
1163    fn match_ranges(matches: &SearchMatches) -> Vec<Range<usize>> {
1164        search_matches(matches)
1165            .iter()
1166            .map(|search_match| search_match.range.clone())
1167            .collect()
1168    }
1169
1170    fn finish(mut finder: StreamingFuzzyMatcher) -> Option<String> {
1171        let snapshot = finder.snapshot.clone();
1172        let matches = finder.finish();
1173        search_matches(&matches).first().map(|search_match| {
1174            snapshot
1175                .text_for_range(search_match.range.clone())
1176                .collect::<String>()
1177        })
1178    }
1179}
1180