tenant.openagents/omega

1use collections::{HashMap, HashSet};
2use gpui::{App, AppContext as _, AsyncApp, Entity, EntityId};
3use language::{Buffer, BufferSnapshot, Point, ToPoint as _};
4use project::{Project, ProjectPath};
5use std::{
6    ops::Range,
7    path::{Path, PathBuf},
8    sync::Arc,
9};
10use text::Anchor;
11use util::{paths::PathStyle, rel_path::RelPath};
12use zeta_prompt::{ContextSource, RelatedExcerpt, RelatedFile, multi_region::is_good_block_start};
13
14use crate::{
15    bm25_context::{Bm25ContextCandidate, collect_bm25_context},
16    git_log_context::build_git_log_index,
17};
18
19/// This module contains collectors for editable context:
20/// excerpts or full files that are likely to be edited.
21const CURSOR_CONTEXT_LINE_COUNT: u32 = 20;
22const EDIT_HISTORY_CONTEXT_LINE_COUNT: u32 = 20;
23const GIT_LOG_CONTEXT_LINE_COUNT: u32 = 10000;
24const GIT_LOG_CONTEXT_FILE_COUNT: usize = 10;
25const ORACLE_SNIPPET_MIN_CONTEXT_LINE_COUNT: u32 = 10;
26const ORACLE_SNIPPET_MAX_CONTEXT_LINE_COUNT: u32 = 40;
27/// How far excerpt boundaries may be nudged to land on a natural block
28/// boundary, mirroring `zeta_prompt::multi_region`'s marker placement.
29const BOUNDARY_SNAP_LINE_COUNT: u32 = 5;
30/// Maximum number of rows between two excerpts of the same buffer that get
31/// bridged into one contiguous excerpt instead of rendering an elision
32/// marker between them.
33const BRIDGED_GAP_LINE_COUNT: u32 = 3;
34
35type RangesByBuffer = HashMap<EntityId, (Entity<Buffer>, Vec<EditableContextRange>)>;
36
37#[derive(Clone)]
38pub struct EditHistoryContextEntry {
39    pub buffer: Entity<Buffer>,
40    pub edited_range: Range<Anchor>,
41}
42
43/// A file known (from expected patches) to be edited next, used by the
44/// oracle context sources when generating training data.
45#[derive(Clone, Debug)]
46pub struct OracleTarget {
47    pub path: Arc<Path>,
48    /// 0-based, end-exclusive row ranges of the expected edit hunks.
49    /// Only used by `ContextSource::OracleSnippet`.
50    pub row_ranges: Vec<Range<u32>>,
51}
52
53struct EditableContextRange {
54    range: Range<Anchor>,
55    order: usize,
56    context_source: ContextSource,
57}
58
59struct ResolvedEditableContextRange {
60    range: Range<Point>,
61    order: usize,
62    context_source: ContextSource,
63}
64
65pub async fn collect_editable_context(
66    project: Entity<Project>,
67    active_buffer: Entity<Buffer>,
68    cursor_position: Anchor,
69    edit_history: Vec<EditHistoryContextEntry>,
70    oracle_targets: Vec<OracleTarget>,
71    context_sources: Vec<ContextSource>,
72    cx: &mut AsyncApp,
73) -> anyhow::Result<Vec<RelatedFile>> {
74    let mut ranges_by_buffer = RangesByBuffer::default();
75
76    if context_sources.contains(&ContextSource::CursorExcerpt) {
77        collect_cursor_excerpt_context(
78            &mut ranges_by_buffer,
79            active_buffer.clone(),
80            cursor_position,
81            cx,
82        );
83    }
84    if context_sources.contains(&ContextSource::CurrentFile) {
85        collect_current_file_context(&mut ranges_by_buffer, active_buffer.clone(), cx);
86    }
87    if context_sources.contains(&ContextSource::EditHistory) {
88        collect_edit_history_context(&mut ranges_by_buffer, &edit_history, cx);
89    }
90    if context_sources.contains(&ContextSource::EditHistoryFile) {
91        collect_edit_history_file_context(&mut ranges_by_buffer, &edit_history, cx);
92    }
93    if context_sources.contains(&ContextSource::GitLog) {
94        collect_git_log_context(
95            &mut ranges_by_buffer,
96            project.clone(),
97            active_buffer.clone(),
98            cx,
99        )
100        .await;
101    }
102
103    // Collected before bm25 so that, under a related-files byte budget, the
104    // small snippets containing the expected edits are never trimmed away in
105    // favor of bm25 excerpts.
106    if context_sources.contains(&ContextSource::OracleSnippet) {
107        collect_oracle_snippet_context(&mut ranges_by_buffer, project.clone(), &oracle_targets, cx)
108            .await;
109    }
110
111    if context_sources.contains(&ContextSource::Bm25) {
112        collect_bm25_context_ranges(
113            &mut ranges_by_buffer,
114            project.clone(),
115            active_buffer,
116            cursor_position,
117            &edit_history,
118            cx,
119        )
120        .await;
121    }
122
123    if context_sources.contains(&ContextSource::OracleFile) {
124        collect_oracle_file_context(&mut ranges_by_buffer, project.clone(), &oracle_targets, cx)
125            .await;
126    }
127
128    Ok(cx.update(|cx| {
129        let project = project.read(cx);
130        let mut related_files = ranges_by_buffer
131            .into_values()
132            .filter_map(|(buffer, ranges)| related_file_for_ranges(&project, &buffer, ranges, cx))
133            .collect::<Vec<_>>();
134        related_files.sort_by_key(|file| {
135            file.excerpts
136                .iter()
137                .map(|excerpt| excerpt.order)
138                .min()
139                .unwrap_or(usize::MAX)
140        });
141        related_files
142    }))
143}
144
145pub fn limit_retrieved_context_to_bytes(
146    related_files: &[RelatedFile],
147    max_bytes: usize,
148) -> Vec<RelatedFile> {
149    struct ExcerptCandidate {
150        file_index: usize,
151        excerpt_index: usize,
152        order: usize,
153    }
154
155    let mut candidates = related_files
156        .iter()
157        .enumerate()
158        .flat_map(|(file_index, file)| {
159            file.excerpts
160                .iter()
161                .enumerate()
162                .map(move |(excerpt_index, excerpt)| ExcerptCandidate {
163                    file_index,
164                    excerpt_index,
165                    order: excerpt.order,
166                })
167        })
168        .collect::<Vec<_>>();
169    candidates.sort_by_key(|candidate| {
170        (
171            candidate.order,
172            candidate.file_index,
173            candidate.excerpt_index,
174        )
175    });
176
177    let mut selected_excerpts = related_files
178        .iter()
179        .map(|file| vec![false; file.excerpts.len()])
180        .collect::<Vec<_>>();
181    let mut covered_ranges_by_file = vec![Vec::<Range<u32>>::new(); related_files.len()];
182    let mut selected_bytes: usize = 0;
183
184    for candidate in candidates {
185        let file = &related_files[candidate.file_index];
186        let excerpt = &file.excerpts[candidate.excerpt_index];
187        let added_bytes =
188            uncovered_excerpt_bytes(excerpt, &covered_ranges_by_file[candidate.file_index]);
189        if added_bytes == 0 || selected_bytes.saturating_add(added_bytes) > max_bytes {
190            continue;
191        }
192
193        selected_bytes += added_bytes;
194        selected_excerpts[candidate.file_index][candidate.excerpt_index] = true;
195        push_covered_range(
196            &mut covered_ranges_by_file[candidate.file_index],
197            excerpt.row_range.clone(),
198        );
199    }
200
201    related_files
202        .iter()
203        .enumerate()
204        .filter_map(|(file_index, file)| {
205            let excerpts = file
206                .excerpts
207                .iter()
208                .enumerate()
209                .filter_map(|(excerpt_index, excerpt)| {
210                    selected_excerpts[file_index][excerpt_index].then(|| excerpt.clone())
211                })
212                .collect::<Vec<_>>();
213            if excerpts.is_empty() {
214                return None;
215            }
216
217            Some(RelatedFile {
218                path: file.path.clone(),
219                max_row: file.max_row,
220                excerpts,
221                in_open_source_repo: file.in_open_source_repo,
222            })
223        })
224        .collect()
225}
226
227fn uncovered_excerpt_bytes(excerpt: &RelatedExcerpt, covered_ranges: &[Range<u32>]) -> usize {
228    let mut bytes = 0;
229
230    for (row, line) in (excerpt.row_range.start..).zip(excerpt.text.split_inclusive('\n')) {
231        if row >= excerpt.row_range.end {
232            break;
233        }
234        if !covered_ranges
235            .iter()
236            .any(|covered_range| covered_range.contains(&row))
237        {
238            bytes += line.len();
239        }
240    }
241
242    bytes
243}
244
245fn push_covered_range(covered_ranges: &mut Vec<Range<u32>>, range: Range<u32>) {
246    covered_ranges.push(range);
247    covered_ranges.sort_by_key(|range| (range.start, range.end));
248
249    let mut merged_ranges: Vec<Range<u32>> = Vec::new();
250    for range in covered_ranges.drain(..) {
251        if let Some(last_range) = merged_ranges.last_mut()
252            && range.start <= last_range.end
253        {
254            last_range.end = last_range.end.max(range.end);
255            continue;
256        }
257
258        merged_ranges.push(range);
259    }
260
261    *covered_ranges = merged_ranges;
262}
263
264fn collect_cursor_excerpt_context(
265    ranges_by_buffer: &mut RangesByBuffer,
266    active_buffer: Entity<Buffer>,
267    cursor_position: Anchor,
268    cx: &mut AsyncApp,
269) {
270    let cursor_range = active_buffer.read_with(cx, |buffer, _cx| {
271        let snapshot = buffer.snapshot();
272        expanded_anchor_range(
273            &snapshot,
274            cursor_position..cursor_position,
275            CURSOR_CONTEXT_LINE_COUNT,
276        )
277    });
278
279    push_context_range(
280        ranges_by_buffer,
281        active_buffer,
282        cursor_range,
283        0,
284        ContextSource::CursorExcerpt,
285    );
286}
287
288fn collect_current_file_context(
289    ranges_by_buffer: &mut RangesByBuffer,
290    active_buffer: Entity<Buffer>,
291    cx: &mut AsyncApp,
292) {
293    collect_full_buffer_context(
294        ranges_by_buffer,
295        active_buffer,
296        0,
297        ContextSource::CurrentFile,
298        cx,
299    );
300}
301
302fn collect_edit_history_context(
303    ranges_by_buffer: &mut RangesByBuffer,
304    edit_history: &[EditHistoryContextEntry],
305    cx: &mut AsyncApp,
306) {
307    for (index, entry) in edit_history.iter().enumerate() {
308        let edit_history_range = entry.buffer.read_with(cx, |buffer, _cx| {
309            expanded_anchor_range(
310                &buffer.snapshot(),
311                entry.edited_range.clone(),
312                EDIT_HISTORY_CONTEXT_LINE_COUNT,
313            )
314        });
315
316        push_context_range(
317            ranges_by_buffer,
318            entry.buffer.clone(),
319            edit_history_range,
320            index + 1,
321            ContextSource::EditHistory,
322        );
323    }
324}
325
326fn collect_edit_history_file_context(
327    ranges_by_buffer: &mut RangesByBuffer,
328    edit_history: &[EditHistoryContextEntry],
329    cx: &mut AsyncApp,
330) {
331    let next_order = next_context_order(ranges_by_buffer);
332    let mut seen_buffers = HashSet::default();
333    let mut index = 0;
334
335    for entry in edit_history {
336        if !seen_buffers.insert(entry.buffer.entity_id()) {
337            continue;
338        }
339
340        collect_full_buffer_context(
341            ranges_by_buffer,
342            entry.buffer.clone(),
343            next_order + index,
344            ContextSource::EditHistoryFile,
345            cx,
346        );
347        index += 1;
348    }
349}
350
351async fn collect_bm25_context_ranges(
352    ranges_by_buffer: &mut RangesByBuffer,
353    project: Entity<Project>,
354    active_buffer: Entity<Buffer>,
355    cursor_position: Anchor,
356    edit_history: &[EditHistoryContextEntry],
357    cx: &mut AsyncApp,
358) {
359    let next_order = next_context_order(ranges_by_buffer);
360    let candidates = collect_bm25_context(
361        project.clone(),
362        active_buffer,
363        cursor_position,
364        edit_history,
365        next_order,
366        cx,
367    )
368    .await;
369
370    for candidate in candidates {
371        collect_bm25_candidate_context(ranges_by_buffer, &project, candidate, cx).await;
372    }
373}
374
375async fn collect_bm25_candidate_context(
376    ranges_by_buffer: &mut RangesByBuffer,
377    project: &Entity<Project>,
378    candidate: Bm25ContextCandidate,
379    cx: &mut AsyncApp,
380) {
381    let buffer = match open_buffer_for_path(project, &candidate.path, cx).await {
382        Ok(Some(buffer)) => buffer,
383        Ok(None) => {
384            log::debug!(
385                "failed to find BM25 context path: {}",
386                candidate.path.display()
387            );
388            return;
389        }
390        Err(error) => {
391            log::debug!(
392                "failed to open BM25 context path {}: {error:#}",
393                candidate.path.display()
394            );
395            return;
396        }
397    };
398
399    let Some(range) = buffer.read_with(cx, |buffer, _cx| {
400        anchor_range_for_row_range(&buffer.snapshot(), candidate.row_range.clone())
401    }) else {
402        return;
403    };
404
405    push_context_range(
406        ranges_by_buffer,
407        buffer,
408        range,
409        candidate.order,
410        ContextSource::Bm25,
411    );
412}
413
414fn anchor_range_for_row_range(
415    snapshot: &BufferSnapshot,
416    row_range: Range<u32>,
417) -> Option<Range<Anchor>> {
418    if row_range.start >= row_range.end || row_range.start > snapshot.max_point().row {
419        return None;
420    }
421
422    let max_point = snapshot.max_point();
423    let start = snapshot.anchor_before(Point::new(row_range.start, 0));
424    let end_point = if row_range.end > max_point.row {
425        max_point
426    } else {
427        Point::new(row_range.end, 0)
428    };
429    let end = snapshot.anchor_after(end_point);
430    Some(start..end)
431}
432
433async fn collect_oracle_file_context(
434    ranges_by_buffer: &mut RangesByBuffer,
435    project: Entity<Project>,
436    oracle_targets: &[OracleTarget],
437    cx: &mut AsyncApp,
438) {
439    let next_order = next_context_order(ranges_by_buffer);
440    let mut seen_buffers = HashSet::default();
441    let mut index = 0;
442
443    for target in oracle_targets {
444        let buffer = match open_buffer_for_path(&project, &target.path, cx).await {
445            Ok(Some(buffer)) => buffer,
446            Ok(None) => {
447                log::debug!("failed to find oracle file path: {}", target.path.display());
448                continue;
449            }
450            Err(error) => {
451                log::debug!(
452                    "failed to open oracle file path {}: {error:#}",
453                    target.path.display()
454                );
455                continue;
456            }
457        };
458
459        if !seen_buffers.insert(buffer.entity_id()) {
460            continue;
461        }
462
463        collect_full_buffer_context(
464            ranges_by_buffer,
465            buffer,
466            next_order + index,
467            ContextSource::OracleFile,
468            cx,
469        );
470        index += 1;
471    }
472}
473
474async fn collect_oracle_snippet_context(
475    ranges_by_buffer: &mut RangesByBuffer,
476    project: Entity<Project>,
477    oracle_targets: &[OracleTarget],
478    cx: &mut AsyncApp,
479) {
480    let next_order = next_context_order(ranges_by_buffer);
481    let mut index = 0;
482
483    for target in oracle_targets {
484        if target.row_ranges.is_empty() {
485            continue;
486        }
487
488        let buffer = match open_buffer_for_path(&project, &target.path, cx).await {
489            Ok(Some(buffer)) => buffer,
490            Ok(None) => {
491                log::debug!(
492                    "failed to find oracle snippet path: {}",
493                    target.path.display()
494                );
495                continue;
496            }
497            Err(error) => {
498                log::debug!(
499                    "failed to open oracle snippet path {}: {error:#}",
500                    target.path.display()
501                );
502                continue;
503            }
504        };
505
506        for row_range in &target.row_ranges {
507            let Some(range) = buffer.read_with(cx, |buffer, _cx| {
508                let snapshot = buffer.snapshot();
509                let padding_above = oracle_snippet_padding(&target.path, row_range.start, 0);
510                let padding_below = oracle_snippet_padding(&target.path, row_range.end, 1);
511                let start_row = row_range.start.saturating_sub(padding_above);
512                // Empty hunk row ranges (pure insertions) still cover one row.
513                let core_end_row = row_range.end.max(row_range.start + 1);
514                let end_row = core_end_row.saturating_add(padding_below);
515                let start_row =
516                    snap_start_row_to_block_boundary(&snapshot, start_row, row_range.start);
517                // `end_row` is exclusive while snapping operates on the last
518                // included row.
519                let end_row = snap_end_row_to_block_boundary(
520                    &snapshot,
521                    end_row.saturating_sub(1),
522                    core_end_row.saturating_sub(1),
523                ) + 1;
524                anchor_range_for_row_range(&snapshot, start_row..end_row)
525            }) else {
526                continue;
527            };
528
529            push_context_range(
530                ranges_by_buffer,
531                buffer.clone(),
532                range,
533                next_order + index,
534                ContextSource::OracleSnippet,
535            );
536            index += 1;
537        }
538    }
539}
540
541/// Deterministic pseudo-random padding, so that the expected edit is not
542/// always centered in the snippet (a student model could otherwise learn the
543/// excerpt center as a position prior), while keeping context retrieval
544/// reproducible across runs.
545fn oracle_snippet_padding(path: &Path, row: u32, salt: u32) -> u32 {
546    use std::hash::{Hash as _, Hasher as _};
547
548    let mut hasher = collections::FxHasher::default();
549    path.hash(&mut hasher);
550    row.hash(&mut hasher);
551    salt.hash(&mut hasher);
552    let span =
553        u64::from(ORACLE_SNIPPET_MAX_CONTEXT_LINE_COUNT - ORACLE_SNIPPET_MIN_CONTEXT_LINE_COUNT);
554    ORACLE_SNIPPET_MIN_CONTEXT_LINE_COUNT + (hasher.finish() % (span + 1)) as u32
555}
556
557async fn open_buffer_for_path(
558    project: &Entity<Project>,
559    path: &Path,
560    cx: &mut AsyncApp,
561) -> anyhow::Result<Option<Entity<Buffer>>> {
562    let path = path.to_path_buf();
563    let path_without_prefix: PathBuf = path.components().skip(1).collect();
564    let project_path = project.update(cx, |project, cx| {
565        project.find_project_path(&path, cx).or_else(|| {
566            if path_without_prefix.as_os_str().is_empty() {
567                None
568            } else {
569                project.find_project_path(&path_without_prefix, cx)
570            }
571        })
572    });
573
574    let Some(project_path) = project_path else {
575        return Ok(None);
576    };
577
578    project
579        .update(cx, |project, cx| project.open_buffer(project_path, cx))
580        .await
581        .map(Some)
582}
583
584fn collect_full_buffer_context(
585    ranges_by_buffer: &mut RangesByBuffer,
586    buffer: Entity<Buffer>,
587    order: usize,
588    context_source: ContextSource,
589    cx: &mut AsyncApp,
590) {
591    let range = buffer.read_with(cx, |buffer, _cx| full_file_anchor_range(&buffer.snapshot()));
592    push_context_range(ranges_by_buffer, buffer, range, order, context_source);
593}
594
595fn full_file_anchor_range(snapshot: &BufferSnapshot) -> Range<Anchor> {
596    let start = snapshot.anchor_before(Point::new(0, 0));
597    let max_point = snapshot.max_point();
598    let end = snapshot.anchor_after(max_point);
599    start..end
600}
601
602fn next_context_order(ranges_by_buffer: &RangesByBuffer) -> usize {
603    ranges_by_buffer
604        .values()
605        .flat_map(|(_, ranges)| ranges.iter().map(|range| range.order))
606        .max()
607        .map_or(0, |order| order + 1)
608}
609
610async fn collect_git_log_context(
611    ranges_by_buffer: &mut RangesByBuffer,
612    project: Entity<Project>,
613    active_buffer: Entity<Buffer>,
614    cx: &mut AsyncApp,
615) {
616    let Some((worktree_id, active_path, worktree_abs_path)) = cx.update(|cx| {
617        let buffer = active_buffer.read(cx);
618        let file = buffer.file()?;
619        let project = project.read(cx);
620        if !project.is_local() {
621            return None;
622        }
623        let worktree = project.worktree_for_id(file.worktree_id(cx), cx)?;
624        let worktree = worktree.read(cx);
625        if !worktree.is_local() {
626            return None;
627        }
628        Some((
629            file.worktree_id(cx),
630            file.path().clone(),
631            worktree.abs_path(),
632        ))
633    }) else {
634        return;
635    };
636
637    let index_result = cx
638        .background_spawn(async move { build_git_log_index(&worktree_abs_path).await })
639        .await;
640    let index = match index_result {
641        Ok(index) => index,
642        Err(error) => {
643            log::debug!("failed to build git log context index: {error:#}");
644            return;
645        }
646    };
647
648    let next_order = next_context_order(ranges_by_buffer);
649
650    for (index, related_path) in index
651        .get_related(active_path.as_std_path(), GIT_LOG_CONTEXT_FILE_COUNT)
652        .into_iter()
653        .enumerate()
654    {
655        let Ok(related_path) = RelPath::new(&related_path, PathStyle::Unix) else {
656            continue;
657        };
658        let project_path = ProjectPath {
659            worktree_id,
660            path: related_path.into_owned().into(),
661        };
662        let buffer = match project
663            .update(cx, |project, cx| project.open_buffer(project_path, cx))
664            .await
665        {
666            Ok(buffer) => buffer,
667            Err(error) => {
668                log::debug!("failed to open git log related buffer: {error:#}");
669                continue;
670            }
671        };
672
673        let range = buffer.read_with(cx, |buffer, _cx| {
674            let snapshot = buffer.snapshot();
675            let max_row = GIT_LOG_CONTEXT_LINE_COUNT.min(snapshot.max_point().row);
676            let end = snapshot.anchor_after(Point::new(max_row, snapshot.line_len(max_row)));
677            snapshot.anchor_before(Point::new(0, 0))..end
678        });
679
680        push_context_range(
681            ranges_by_buffer,
682            buffer,
683            range,
684            next_order + index,
685            ContextSource::GitLog,
686        );
687    }
688}
689
690fn expanded_anchor_range(
691    snapshot: &BufferSnapshot,
692    range: Range<Anchor>,
693    context_line_count: u32,
694) -> Range<Anchor> {
695    let start = range.start.to_point(snapshot);
696    let end = range.end.to_point(snapshot);
697    let start_row = start.row.saturating_sub(context_line_count);
698    let end_row = end
699        .row
700        .saturating_add(context_line_count)
701        .min(snapshot.max_point().row);
702    let start_row = snap_start_row_to_block_boundary(snapshot, start_row, start.row);
703    let end_row = snap_end_row_to_block_boundary(snapshot, end_row, end.row);
704    let start = snapshot.anchor_before(Point::new(start_row, 0));
705    let end = snapshot.anchor_after(Point::new(end_row, snapshot.line_len(end_row)));
706    start..end
707}
708
709/// Nudge an excerpt's first row forward (up to `BOUNDARY_SNAP_LINE_COUNT`
710/// lines, never past `core_row`) so the excerpt starts at a natural block
711/// boundary: preferably a good block start right after blank line(s), or
712/// failing that any good block start. Mirrors the marker placement
713/// heuristics of `zeta_prompt::multi_region`.
714fn snap_start_row_to_block_boundary(
715    snapshot: &text::BufferSnapshot,
716    row: u32,
717    core_row: u32,
718) -> u32 {
719    let limit = core_row
720        .min(row.saturating_add(BOUNDARY_SNAP_LINE_COUNT))
721        .min(snapshot.max_point().row);
722    let mut first_good_start = None;
723    for candidate in row..=limit {
724        if snapshot.is_line_blank(candidate) {
725            continue;
726        }
727        if !is_good_block_start(line_text(snapshot, candidate).trim()) {
728            continue;
729        }
730        if candidate > 0 && snapshot.is_line_blank(candidate - 1) {
731            return candidate;
732        }
733        if first_good_start.is_none() {
734            first_good_start = Some(candidate);
735        }
736    }
737    first_good_start.unwrap_or(row)
738}
739
740/// Nudge an excerpt's last row backward (up to `BOUNDARY_SNAP_LINE_COUNT`
741/// lines, never before `core_row`) so the excerpt ends at the last non-blank
742/// line before a blank line or at the end of the file.
743fn snap_end_row_to_block_boundary(snapshot: &text::BufferSnapshot, row: u32, core_row: u32) -> u32 {
744    let max_row = snapshot.max_point().row;
745    let row = row.min(max_row);
746    if row == max_row {
747        return row;
748    }
749    let limit = core_row.max(row.saturating_sub(BOUNDARY_SNAP_LINE_COUNT));
750    for candidate in (limit..=row).rev() {
751        if snapshot.is_line_blank(candidate) {
752            continue;
753        }
754        if snapshot.is_line_blank(candidate + 1) {
755            return candidate;
756        }
757    }
758    row
759}
760
761fn line_text(snapshot: &text::BufferSnapshot, row: u32) -> String {
762    snapshot
763        .text_for_range(Point::new(row, 0)..Point::new(row, snapshot.line_len(row)))
764        .collect()
765}
766
767fn push_context_range(
768    ranges_by_buffer: &mut RangesByBuffer,
769    buffer: Entity<Buffer>,
770    range: Range<Anchor>,
771    order: usize,
772    context_source: ContextSource,
773) {
774    ranges_by_buffer
775        .entry(buffer.entity_id())
776        .or_insert_with(|| (buffer.clone(), Vec::new()))
777        .1
778        .push(EditableContextRange {
779            range,
780            order,
781            context_source,
782        });
783}
784
785fn related_file_for_ranges(
786    project: &Project,
787    buffer: &Entity<Buffer>,
788    ranges: Vec<EditableContextRange>,
789    cx: &App,
790) -> Option<RelatedFile> {
791    let buffer = buffer.read(cx);
792    let snapshot = buffer.snapshot();
793    let file = snapshot.file()?;
794    let worktree = project.worktree_for_id(file.worktree_id(cx), cx)?;
795    let path: Arc<Path> = Path::new(&format!(
796        "{}/{}",
797        worktree.read(cx).root_name().as_unix_str(),
798        file.path().as_unix_str()
799    ))
800    .into();
801
802    let mut ranges = resolved_context_ranges(ranges, &snapshot);
803    split_overlapping_ranges(&mut ranges);
804
805    let excerpts = ranges
806        .into_iter()
807        .map(|range| RelatedExcerpt {
808            row_range: range.range.start.row..range.range.end.row,
809            text: snapshot
810                .text_for_range(range.range)
811                .collect::<String>()
812                .into(),
813            order: range.order,
814            context_source: range.context_source,
815        })
816        .collect::<Vec<_>>();
817
818    Some(RelatedFile {
819        path,
820        max_row: snapshot.max_point().row,
821        excerpts,
822        in_open_source_repo: false,
823    })
824}
825
826fn resolved_context_ranges(
827    ranges: Vec<EditableContextRange>,
828    snapshot: &BufferSnapshot,
829) -> Vec<ResolvedEditableContextRange> {
830    ranges
831        .into_iter()
832        .filter_map(|range| {
833            let start = range.range.start.to_point(snapshot);
834            let end = range.range.end.to_point(snapshot);
835            if start >= end {
836                return None;
837            }
838
839            Some(ResolvedEditableContextRange {
840                range: start..end,
841                order: range.order,
842                context_source: range.context_source,
843            })
844        })
845        .collect()
846}
847
848/// Split overlapping ranges into disjoint segments instead of merging them
849/// into one range. Each segment keeps the minimum order and highest-priority
850/// source among the ranges covering it, and adjacent segments with equal
851/// order are coalesced. This preserves priority granularity: a small
852/// high-priority snippet inside a large low-priority range remains its own
853/// excerpt, so byte-budget selection can retain it even when the surrounding
854/// range doesn't fit. The resulting segments are disjoint and sorted by
855/// position.
856///
857/// Ranges separated by at most `BRIDGED_GAP_LINE_COUNT` rows are bridged:
858/// the small gap is attached to the preceding segment so the excerpts render
859/// as one contiguous block instead of being separated by an elision marker.
860fn split_overlapping_ranges(ranges: &mut Vec<ResolvedEditableContextRange>) {
861    ranges.sort_by_key(|range| (range.range.start, range.range.end));
862    let mut output: Vec<ResolvedEditableContextRange> = Vec::new();
863    let mut cluster: Vec<ResolvedEditableContextRange> = Vec::new();
864    let mut cluster_end = Point::zero();
865
866    for range in ranges.drain(..) {
867        let bridge_limit_row = row_aligned_end(cluster_end)
868            .row
869            .saturating_add(BRIDGED_GAP_LINE_COUNT);
870        if cluster.is_empty() || range.range.start.row <= bridge_limit_row {
871            cluster_end = cluster_end.max(range.range.end);
872            cluster.push(range);
873        } else {
874            split_cluster(std::mem::take(&mut cluster), cluster_end, &mut output);
875            cluster_end = range.range.end;
876            cluster.push(range);
877        }
878    }
879    if !cluster.is_empty() {
880        split_cluster(cluster, cluster_end, &mut output);
881    }
882
883    *ranges = output;
884}
885
886fn split_cluster(
887    cluster: Vec<ResolvedEditableContextRange>,
888    cluster_end: Point,
889    output: &mut Vec<ResolvedEditableContextRange>,
890) {
891    if cluster.len() == 1 {
892        output.extend(cluster);
893        return;
894    }
895
896    let cluster_start = cluster[0].range.start;
897    let mut boundaries = Vec::with_capacity(cluster.len() * 2 + 1);
898    boundaries.push(cluster_start);
899    for range in &cluster {
900        boundaries.push(range.range.start);
901        boundaries.push(row_aligned_end(range.range.end));
902    }
903    boundaries.retain(|boundary| *boundary >= cluster_start && *boundary < cluster_end);
904    boundaries.sort();
905    boundaries.dedup();
906    boundaries.push(cluster_end);
907
908    for window in boundaries.windows(2) {
909        let segment = window[0]..window[1];
910        // The segment is attributed to the lowest-order covering range
911        // (breaking ties by source priority), so that its source label is
912        // consistent with the order that drives budget selection.
913        let mut order_and_source: Option<(usize, ContextSource)> = None;
914        for range in &cluster {
915            if range.range.start <= segment.start && row_aligned_end(range.range.end) >= segment.end
916            {
917                let candidate = (range.order, range.context_source);
918                if order_and_source.is_none_or(|(order, context_source)| {
919                    (candidate.0, context_source_order(candidate.1))
920                        < (order, context_source_order(context_source))
921                }) {
922                    order_and_source = Some(candidate);
923                }
924            }
925        }
926        let Some((order, context_source)) = order_and_source else {
927            // A bridged gap between two nearby ranges: no range covers it, so
928            // attach it to the preceding segment to form contiguous output.
929            if let Some(last) = output.last_mut()
930                && last.range.end == segment.start
931            {
932                last.range.end = segment.end;
933            }
934            continue;
935        };
936
937        if let Some(last) = output.last_mut()
938            && last.range.end == segment.start
939            && last.order == order
940        {
941            last.range.end = segment.end;
942            if context_source_order(context_source) < context_source_order(last.context_source) {
943                last.context_source = context_source;
944            }
945            continue;
946        }
947
948        output.push(ResolvedEditableContextRange {
949            range: segment,
950            order,
951            context_source,
952        });
953    }
954}
955
956/// Context ranges always cover whole lines: their ends sit either at a line
957/// start (column 0) or at the end of a line's content. Cutting a segment at
958/// the end of a line's content would attribute the trailing newline to the
959/// next segment, so nudge such ends forward to the next line start.
960fn row_aligned_end(point: Point) -> Point {
961    if point.column > 0 {
962        Point::new(point.row + 1, 0)
963    } else {
964        point
965    }
966}
967
968#[allow(dead_code)]
969fn push_context_source(context_sources: &mut Vec<ContextSource>, context_source: ContextSource) {
970    if !context_sources.contains(&context_source) {
971        context_sources.push(context_source);
972        context_sources.sort_by_key(|context_source| context_source_order(*context_source));
973    }
974}
975
976fn context_source_order(context_source: ContextSource) -> usize {
977    match context_source {
978        ContextSource::Lsp => 0,
979        ContextSource::CursorExcerpt => 1,
980        ContextSource::CurrentFile => 2,
981        ContextSource::EditHistory => 3,
982        ContextSource::EditHistoryFile => 4,
983        ContextSource::GitLog => 5,
984        ContextSource::Bm25 => 6,
985        ContextSource::OracleSnippet => 7,
986        ContextSource::OracleFile => 8,
987    }
988}
989
990#[cfg(test)]
991mod tests {
992    use super::*;
993
994    fn text_snapshot(text: &str) -> text::BufferSnapshot {
995        text::Buffer::new(
996            text::ReplicaId::LOCAL,
997            text::BufferId::new(1).unwrap(),
998            text,
999        )
1000        .snapshot()
1001        .clone()
1002    }
1003
1004    fn resolved_range(
1005        start_row: u32,
1006        end_row: u32,
1007        order: usize,
1008        context_source: ContextSource,
1009    ) -> ResolvedEditableContextRange {
1010        ResolvedEditableContextRange {
1011            range: Point::new(start_row, 0)..Point::new(end_row, 0),
1012            order,
1013            context_source,
1014        }
1015    }
1016
1017    #[test]
1018    fn test_split_overlapping_ranges_coalesces_equal_orders() {
1019        // A full-file current-file range plus edit-history windows inside it:
1020        // every segment is covered by the order-0 full-file range, so they
1021        // coalesce back into a single range with the highest-priority source.
1022        let mut ranges = vec![
1023            resolved_range(6, 46, 1, ContextSource::EditHistory),
1024            resolved_range(0, 281, 0, ContextSource::CurrentFile),
1025            resolved_range(17, 79, 2, ContextSource::EditHistory),
1026            resolved_range(85, 125, 3, ContextSource::EditHistory),
1027        ];
1028        split_overlapping_ranges(&mut ranges);
1029        assert_eq!(ranges.len(), 1);
1030        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(281, 0));
1031        assert_eq!(ranges[0].order, 0);
1032        assert_eq!(ranges[0].context_source, ContextSource::CurrentFile);
1033    }
1034
1035    #[test]
1036    fn test_split_overlapping_ranges_keeps_disjoint_ranges() {
1037        let mut ranges = vec![
1038            resolved_range(50, 60, 1, ContextSource::EditHistory),
1039            resolved_range(0, 10, 0, ContextSource::CursorExcerpt),
1040            resolved_range(5, 12, 2, ContextSource::EditHistory),
1041        ];
1042        split_overlapping_ranges(&mut ranges);
1043        assert_eq!(ranges.len(), 3);
1044        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(10, 0));
1045        assert_eq!(ranges[0].order, 0);
1046        assert_eq!(ranges[0].context_source, ContextSource::CursorExcerpt);
1047        assert_eq!(ranges[1].range, Point::new(10, 0)..Point::new(12, 0));
1048        assert_eq!(ranges[1].order, 2);
1049        assert_eq!(ranges[1].context_source, ContextSource::EditHistory);
1050        assert_eq!(ranges[2].range, Point::new(50, 0)..Point::new(60, 0));
1051        assert_eq!(ranges[2].order, 1);
1052    }
1053
1054    #[test]
1055    fn test_split_overlapping_ranges_keeps_adjacent_ranges_with_distinct_orders() {
1056        let mut ranges = vec![
1057            resolved_range(0, 10, 0, ContextSource::EditHistory),
1058            resolved_range(10, 20, 1, ContextSource::EditHistory),
1059        ];
1060        split_overlapping_ranges(&mut ranges);
1061        assert_eq!(ranges.len(), 2);
1062        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(10, 0));
1063        assert_eq!(ranges[0].order, 0);
1064        assert_eq!(ranges[1].range, Point::new(10, 0)..Point::new(20, 0));
1065        assert_eq!(ranges[1].order, 1);
1066    }
1067
1068    #[test]
1069    fn test_split_overlapping_ranges_preserves_high_priority_snippet_inside_large_range() {
1070        // A small high-priority snippet inside a large low-priority range
1071        // stays its own segment, so byte-budget selection can keep it even
1072        // when the surrounding range doesn't fit.
1073        let mut ranges = vec![
1074            resolved_range(0, 200, 8, ContextSource::GitLog),
1075            resolved_range(50, 60, 2, ContextSource::OracleSnippet),
1076        ];
1077        split_overlapping_ranges(&mut ranges);
1078        assert_eq!(ranges.len(), 3);
1079        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(50, 0));
1080        assert_eq!(ranges[0].order, 8);
1081        assert_eq!(ranges[0].context_source, ContextSource::GitLog);
1082        assert_eq!(ranges[1].range, Point::new(50, 0)..Point::new(60, 0));
1083        assert_eq!(ranges[1].order, 2);
1084        assert_eq!(ranges[1].context_source, ContextSource::OracleSnippet);
1085        assert_eq!(ranges[2].range, Point::new(60, 0)..Point::new(200, 0));
1086        assert_eq!(ranges[2].order, 8);
1087        assert_eq!(ranges[2].context_source, ContextSource::GitLog);
1088    }
1089
1090    #[test]
1091    fn test_split_overlapping_ranges_aligns_mid_line_ends_to_row_starts() {
1092        // Ranges ending at a line's content end (column > 0) are treated as
1093        // covering through that whole line, so equal-order overlapping ranges
1094        // still coalesce into one segment.
1095        let mut ranges = vec![
1096            ResolvedEditableContextRange {
1097                range: Point::new(0, 0)..Point::new(10, 5),
1098                order: 1,
1099                context_source: ContextSource::EditHistory,
1100            },
1101            resolved_range(5, 20, 1, ContextSource::EditHistory),
1102        ];
1103        split_overlapping_ranges(&mut ranges);
1104        assert_eq!(ranges.len(), 1);
1105        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(20, 0));
1106        assert_eq!(ranges[0].order, 1);
1107    }
1108
1109    #[test]
1110    fn test_split_overlapping_ranges_bridges_small_gaps() {
1111        // Ranges separated by a few rows are bridged: the gap is attached to
1112        // the preceding segment, producing contiguous excerpts.
1113        let mut ranges = vec![
1114            resolved_range(0, 10, 0, ContextSource::EditHistory),
1115            resolved_range(13, 20, 1, ContextSource::Bm25),
1116        ];
1117        split_overlapping_ranges(&mut ranges);
1118        assert_eq!(ranges.len(), 2);
1119        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(13, 0));
1120        assert_eq!(ranges[0].order, 0);
1121        assert_eq!(ranges[0].context_source, ContextSource::EditHistory);
1122        assert_eq!(ranges[1].range, Point::new(13, 0)..Point::new(20, 0));
1123        assert_eq!(ranges[1].order, 1);
1124        assert_eq!(ranges[1].context_source, ContextSource::Bm25);
1125    }
1126
1127    #[test]
1128    fn test_split_overlapping_ranges_bridged_equal_orders_coalesce() {
1129        let mut ranges = vec![
1130            resolved_range(0, 10, 1, ContextSource::EditHistory),
1131            resolved_range(12, 20, 1, ContextSource::EditHistory),
1132        ];
1133        split_overlapping_ranges(&mut ranges);
1134        assert_eq!(ranges.len(), 1);
1135        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(20, 0));
1136        assert_eq!(ranges[0].order, 1);
1137    }
1138
1139    #[test]
1140    fn test_split_overlapping_ranges_does_not_bridge_large_gaps() {
1141        let mut ranges = vec![
1142            resolved_range(0, 10, 0, ContextSource::EditHistory),
1143            resolved_range(14, 20, 1, ContextSource::Bm25),
1144        ];
1145        split_overlapping_ranges(&mut ranges);
1146        assert_eq!(ranges.len(), 2);
1147        assert_eq!(ranges[0].range, Point::new(0, 0)..Point::new(10, 0));
1148        assert_eq!(ranges[1].range, Point::new(14, 0)..Point::new(20, 0));
1149    }
1150
1151    #[test]
1152    fn test_snap_start_row_prefers_block_start_after_blank_line() {
1153        let snapshot = text_snapshot(concat!(
1154            "    body\n",   // 0
1155            "}\n",          // 1
1156            "\n",           // 2
1157            "fn bar() {\n", // 3
1158            "    body\n",   // 4
1159            "    core\n",   // 5
1160        ));
1161        // Row 1 is a structural tail and row 2 is blank; row 3 follows a
1162        // blank line and is a good block start.
1163        assert_eq!(snap_start_row_to_block_boundary(&snapshot, 1, 5), 3);
1164        // Snapping never moves past the core row.
1165        assert_eq!(snap_start_row_to_block_boundary(&snapshot, 1, 2), 1);
1166        // A row that is already a good start after a blank line stays put.
1167        assert_eq!(snap_start_row_to_block_boundary(&snapshot, 3, 5), 3);
1168    }
1169
1170    #[test]
1171    fn test_snap_start_row_falls_back_to_first_good_start() {
1172        let snapshot = text_snapshot(concat!(
1173            "}\n",          // 0
1174            "let x = 1;\n", // 1
1175            "let y = 2;\n", // 2
1176            "core\n",       // 3
1177        ));
1178        // No after-blank boundary in the window; the first good start wins.
1179        assert_eq!(snap_start_row_to_block_boundary(&snapshot, 0, 3), 1);
1180    }
1181
1182    #[test]
1183    fn test_snap_end_row_ends_before_blank_line() {
1184        let snapshot = text_snapshot(concat!(
1185            "core\n", // 0
1186            "a\n",    // 1
1187            "b\n",    // 2
1188            "\n",     // 3
1189            "c\n",    // 4
1190            "d\n",    // 5
1191        ));
1192        // Row 4 is followed by non-blank row 5, so scan back to row 2 which
1193        // precedes the blank row 3.
1194        assert_eq!(snap_end_row_to_block_boundary(&snapshot, 4, 0), 2);
1195        // Snapping never moves before the core row.
1196        assert_eq!(snap_end_row_to_block_boundary(&snapshot, 4, 4), 4);
1197        // The last row of the file stays put.
1198        let max_row = snapshot.max_point().row;
1199        assert_eq!(
1200            snap_end_row_to_block_boundary(&snapshot, max_row, 0),
1201            max_row
1202        );
1203    }
1204
1205    #[test]
1206    fn test_limit_retrieved_context_keeps_high_priority_snippet_under_tight_budget() {
1207        let line = "0123456789\n";
1208        let excerpt = |row_range: Range<u32>, order: usize, context_source: ContextSource| {
1209            let text = line.repeat((row_range.end - row_range.start) as usize);
1210            RelatedExcerpt {
1211                row_range,
1212                text: text.into(),
1213                order,
1214                context_source,
1215            }
1216        };
1217        let related_files = vec![RelatedFile {
1218            path: Path::new("root/file.rs").into(),
1219            max_row: 200,
1220            excerpts: vec![
1221                excerpt(0..50, 8, ContextSource::GitLog),
1222                excerpt(50..60, 2, ContextSource::OracleSnippet),
1223                excerpt(60..200, 8, ContextSource::GitLog),
1224            ],
1225            in_open_source_repo: false,
1226        }];
1227
1228        // Budget fits the snippet but not the surrounding segments.
1229        let limited = limit_retrieved_context_to_bytes(&related_files, 20 * line.len());
1230        assert_eq!(limited.len(), 1);
1231        assert_eq!(limited[0].excerpts.len(), 1);
1232        assert_eq!(limited[0].excerpts[0].row_range, 50..60);
1233        assert_eq!(
1234            limited[0].excerpts[0].context_source,
1235            ContextSource::OracleSnippet
1236        );
1237    }
1238}
1239