tenant.openagents/omega

1use crate::{
2    example::Example,
3    format_prompt::line_start_offset,
4    headless::EpAppState,
5    load_project::run_load_project,
6    progress::{ExampleProgress, InfoStyle, Step, StepProgress},
7};
8use anyhow::Context as _;
9use clap::ValueEnum;
10use collections::{HashMap, HashSet};
11use edit_prediction::{DebugEvent, EditPredictionStore, udiff::refresh_worktree_entries};
12use edit_prediction_context::OracleTarget;
13use futures::{FutureExt as _, StreamExt as _, channel::mpsc};
14use gpui::{AsyncApp, Entity};
15use language::Buffer;
16use project::Project;
17use std::{
18    path::{Path, PathBuf},
19    sync::Arc,
20    time::Duration,
21};
22use zeta_prompt::{ContextSource, udiff::DiffLine};
23
24#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, ValueEnum)]
25pub enum ContextRetrievalType {
26    Lsp,
27    Editable,
28    CurrentFile,
29    EditHistory,
30    EditHistoryFile,
31    GitLog,
32    Bm25,
33    OracleFile,
34    OracleSnippet,
35    #[default]
36    All,
37    None,
38}
39
40impl std::fmt::Display for ContextRetrievalType {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            ContextRetrievalType::Lsp => write!(f, "lsp"),
44            ContextRetrievalType::Editable => write!(f, "editable"),
45            ContextRetrievalType::CurrentFile => write!(f, "current-file"),
46            ContextRetrievalType::EditHistory => write!(f, "edit-history"),
47            ContextRetrievalType::EditHistoryFile => write!(f, "edit-history-file"),
48            ContextRetrievalType::GitLog => write!(f, "git-log"),
49            ContextRetrievalType::Bm25 => write!(f, "bm25"),
50            ContextRetrievalType::OracleFile => write!(f, "oracle-file"),
51            ContextRetrievalType::OracleSnippet => write!(f, "oracle-snippet"),
52            ContextRetrievalType::All => write!(f, "all"),
53            ContextRetrievalType::None => write!(f, "none"),
54        }
55    }
56}
57
58impl ContextRetrievalType {
59    pub fn context_sources(self) -> Vec<ContextSource> {
60        match self {
61            ContextRetrievalType::Lsp => vec![ContextSource::Lsp],
62            ContextRetrievalType::Editable => editable_context_sources(),
63            ContextRetrievalType::CurrentFile => vec![ContextSource::CurrentFile],
64            ContextRetrievalType::EditHistory => vec![ContextSource::EditHistory],
65            ContextRetrievalType::EditHistoryFile => vec![ContextSource::EditHistoryFile],
66            ContextRetrievalType::GitLog => vec![ContextSource::GitLog],
67            ContextRetrievalType::Bm25 => vec![ContextSource::Bm25],
68            ContextRetrievalType::OracleFile => vec![ContextSource::OracleFile],
69            ContextRetrievalType::OracleSnippet => vec![ContextSource::OracleSnippet],
70            ContextRetrievalType::All => {
71                let mut sources = vec![ContextSource::Lsp];
72                sources.extend(editable_context_sources());
73                sources
74            }
75            ContextRetrievalType::None => Vec::new(),
76        }
77    }
78}
79
80pub fn context_sources_for_types(context_types: &[ContextRetrievalType]) -> Vec<ContextSource> {
81    let mut context_sources = Vec::new();
82    for context_type in context_types {
83        for context_source in context_type.context_sources() {
84            if !context_sources.contains(&context_source) {
85                context_sources.push(context_source);
86            }
87        }
88    }
89    context_sources
90}
91
92fn editable_context_sources() -> Vec<ContextSource> {
93    vec![
94        ContextSource::CursorExcerpt,
95        ContextSource::CurrentFile,
96        ContextSource::EditHistory,
97        ContextSource::EditHistoryFile,
98        ContextSource::GitLog,
99        ContextSource::Bm25,
100    ]
101}
102
103pub async fn run_context_retrieval(
104    example: &mut Example,
105    app_state: Arc<EpAppState>,
106    example_progress: &ExampleProgress,
107    context_types: Vec<ContextRetrievalType>,
108    force: bool,
109    mut cx: AsyncApp,
110) -> anyhow::Result<()> {
111    if (!force
112        && example
113            .prompt_inputs
114            .as_ref()
115            .is_some_and(|inputs| inputs.related_files.is_some()))
116        || example.spec.repository_url.is_empty()
117    {
118        return Ok(());
119    }
120
121    run_load_project(example, app_state.clone(), example_progress, cx.clone()).await?;
122
123    let step_progress: Arc<StepProgress> = example_progress.start(Step::Context).into();
124
125    let state = example.state.as_ref().unwrap();
126    let project = state.project.clone();
127
128    let ep_store = cx
129        .update(|cx| EditPredictionStore::try_global(cx))
130        .context("EditPredictionStore not initialized")?;
131
132    let mut context_files = Vec::new();
133    let context_sources = context_sources_for_types(&context_types);
134
135    if context_sources.contains(&ContextSource::Lsp) {
136        let _lsp_handle = project.update(&mut cx, |project, cx| {
137            project.register_buffer_with_language_servers(&state.buffer, cx)
138        });
139        wait_for_language_servers_to_start(&project, &state.buffer, &step_progress, &mut cx)
140            .await?;
141
142        let mut events = ep_store.update(&mut cx, |store, cx| {
143            store.register_buffer(&state.buffer, &project, cx);
144            store.refresh_context(&project, &state.buffer, state.cursor_position, cx);
145            store.debug_info(&project, cx)
146        });
147
148        while let Some(event) = events.next().await {
149            match event {
150                DebugEvent::ContextRetrievalFinished(_) => {
151                    break;
152                }
153                _ => {}
154            }
155        }
156
157        context_files
158            .extend(ep_store.update(&mut cx, |store, cx| store.context_for_project(&project, cx)));
159    }
160
161    let editable_context_sources = context_sources
162        .into_iter()
163        .filter(|context_source| *context_source != ContextSource::Lsp)
164        .collect::<Vec<_>>();
165    if !editable_context_sources.is_empty() {
166        let oracle_targets = if editable_context_sources.contains(&ContextSource::OracleFile)
167            || editable_context_sources.contains(&ContextSource::OracleSnippet)
168        {
169            let oracle_targets = oracle_targets_from_expected_patches(example);
170            let oracle_paths = oracle_targets
171                .iter()
172                .map(|target| target.path.clone())
173                .collect::<Vec<_>>();
174            refresh_paths(&project, &oracle_paths, &mut cx).await?;
175            oracle_targets
176        } else {
177            Vec::new()
178        };
179
180        let editable_context = ep_store
181            .update(&mut cx, |store, cx| {
182                store.collect_editable_context(
183                    project.clone(),
184                    state.buffer.clone(),
185                    state.cursor_position,
186                    oracle_targets,
187                    editable_context_sources,
188                    cx,
189                )
190            })
191            .await?;
192        merge_context_files(&mut context_files, editable_context);
193    }
194
195    let excerpt_count: usize = context_files.iter().map(|f| f.excerpts.len()).sum();
196    step_progress.set_info(format!("{} excerpts", excerpt_count), InfoStyle::Normal);
197
198    if let Some(prompt_inputs) = example.prompt_inputs.as_mut() {
199        prompt_inputs.related_files = Some(context_files);
200    }
201    Ok(())
202}
203
204fn merge_context_files(
205    context_files: &mut Vec<zeta_prompt::RelatedFile>,
206    new_files: Vec<zeta_prompt::RelatedFile>,
207) {
208    for mut new_file in new_files {
209        if let Some(existing_file) = context_files
210            .iter_mut()
211            .find(|existing_file| existing_file.path == new_file.path)
212        {
213            existing_file.max_row = existing_file.max_row.max(new_file.max_row);
214            existing_file.excerpts.append(&mut new_file.excerpts);
215            existing_file.in_open_source_repo =
216                existing_file.in_open_source_repo && new_file.in_open_source_repo;
217        } else {
218            context_files.push(new_file);
219        }
220    }
221    for file in context_files.iter_mut() {
222        coalesce_touching_excerpts(&mut file.excerpts);
223    }
224}
225
226/// Sort a file's excerpts by position and merge those whose row ranges touch
227/// or overlap. Touching excerpts render seamlessly in prompts (no `...`
228/// separator), so keeping them separate splits edit addressing across a
229/// boundary the model can't see and wastes adjacent marker tags (see
230/// `TeacherJumpsPrompt::parse`).
231///
232/// A merged excerpt keeps the minimum `order` (and that excerpt's context
233/// source), so the highest-priority content still survives budget-based
234/// selection downstream; the tradeoff is that lower-priority touching
235/// content is now selected together with it.
236fn coalesce_touching_excerpts(excerpts: &mut Vec<zeta_prompt::RelatedExcerpt>) {
237    excerpts.sort_by_key(|excerpt| (excerpt.row_range.start, excerpt.row_range.end));
238    let mut coalesced: Vec<zeta_prompt::RelatedExcerpt> = Vec::with_capacity(excerpts.len());
239    for excerpt in excerpts.drain(..) {
240        let Some(last) = coalesced.last_mut() else {
241            coalesced.push(excerpt);
242            continue;
243        };
244        if excerpt.row_range.start > last.row_range.end {
245            coalesced.push(excerpt);
246            continue;
247        }
248        if excerpt.row_range.end > last.row_range.end {
249            // Touching or overlapping. Shared rows come from the same buffer
250            // snapshot, so drop the duplicated prefix of the new excerpt and
251            // append the rest.
252            let mut overlap_rows = (last.row_range.end - excerpt.row_range.start) as usize;
253            if !last.text.ends_with('\n') && !last.text.is_empty() {
254                // An unterminated final line means `last` includes the
255                // content of its `row_range.end` row itself.
256                overlap_rows += 1;
257            }
258            let Some(tail_start) = line_start_offset(&excerpt.text, overlap_rows) else {
259                // The excerpt has fewer lines than its row range claims;
260                // leave it unmerged rather than corrupt the text.
261                coalesced.push(excerpt);
262                continue;
263            };
264            let mut text =
265                String::with_capacity(last.text.len() + excerpt.text.len() - tail_start + 1);
266            text.push_str(&last.text);
267            if !text.is_empty() && !text.ends_with('\n') {
268                text.push('\n');
269            }
270            text.push_str(&excerpt.text[tail_start..]);
271            last.text = text.into();
272            last.row_range.end = excerpt.row_range.end;
273        }
274        if excerpt.order < last.order {
275            last.order = excerpt.order;
276            last.context_source = excerpt.context_source;
277        }
278    }
279    *excerpts = coalesced;
280}
281
282fn oracle_targets_from_expected_patches(example: &Example) -> Vec<OracleTarget> {
283    let mut target_indices: HashMap<PathBuf, usize> = HashMap::default();
284    let mut targets: Vec<(PathBuf, Vec<std::ops::Range<u32>>)> = Vec::new();
285
286    let mut target_index = |path: &str, targets: &mut Vec<(PathBuf, Vec<std::ops::Range<u32>>)>| {
287        let path = Path::new(path).to_path_buf();
288        *target_indices.entry(path.clone()).or_insert_with(|| {
289            targets.push((path, Vec::new()));
290            targets.len() - 1
291        })
292    };
293
294    for patch in &example.spec.expected_patches {
295        // Index of the target whose old-side rows the current hunk headers
296        // refer to. Hunk old rows refer to the file's current state only for
297        // the first expected patch; later patches drift, but the snippet
298        // padding absorbs small offsets.
299        let mut current_target: Option<usize> = None;
300        for line in patch.lines() {
301            match DiffLine::parse(line) {
302                DiffLine::OldPath { path } => {
303                    current_target = (path.as_ref() != "/dev/null")
304                        .then(|| target_index(path.as_ref(), &mut targets));
305                }
306                DiffLine::NewPath { path } => {
307                    if path.as_ref() != "/dev/null" {
308                        let index = target_index(path.as_ref(), &mut targets);
309                        if current_target.is_none() {
310                            current_target = Some(index);
311                        }
312                    }
313                }
314                DiffLine::HunkHeader(Some(location)) => {
315                    if let Some(index) = current_target {
316                        let start = location.start_line_old;
317                        targets[index].1.push(start..start + location.count_old);
318                    }
319                }
320                _ => {}
321            }
322        }
323    }
324
325    targets
326        .into_iter()
327        .map(|(path, row_ranges)| OracleTarget {
328            path: path.into(),
329            row_ranges,
330        })
331        .collect()
332}
333
334async fn refresh_paths(
335    project: &Entity<Project>,
336    paths: &[Arc<Path>],
337    cx: &mut AsyncApp,
338) -> anyhow::Result<()> {
339    if paths.is_empty() {
340        return Ok(());
341    }
342
343    let Some(worktree) = project.read_with(cx, |project, cx| project.visible_worktrees(cx).next())
344    else {
345        return Ok(());
346    };
347
348    refresh_worktree_entries(&worktree, paths.iter().map(|path| path.as_ref()), cx).await
349}
350
351async fn wait_for_language_servers_to_start(
352    project: &Entity<Project>,
353    buffer: &Entity<Buffer>,
354    step_progress: &Arc<StepProgress>,
355    cx: &mut AsyncApp,
356) -> anyhow::Result<()> {
357    let lsp_store = project.read_with(cx, |project, _| project.lsp_store());
358
359    // Determine which servers exist for this buffer, and which are still starting.
360    let mut servers_pending_start = HashSet::default();
361    let mut servers_pending_diagnostics = HashSet::default();
362    buffer.update(cx, |buffer, cx| {
363        lsp_store.update(cx, |lsp_store, cx| {
364            let ids = lsp_store.language_servers_for_local_buffer(buffer, cx);
365            for &id in &ids {
366                match lsp_store.language_server_statuses.get(&id) {
367                    None => {
368                        servers_pending_start.insert(id);
369                        servers_pending_diagnostics.insert(id);
370                    }
371                    Some(status) if status.has_pending_diagnostic_updates => {
372                        servers_pending_diagnostics.insert(id);
373                    }
374                    Some(_) => {}
375                }
376            }
377        });
378    });
379
380    step_progress.set_substatus(format!(
381        "waiting for {} LSPs",
382        servers_pending_diagnostics.len()
383    ));
384
385    let timeout_duration = if servers_pending_start.is_empty() {
386        Duration::from_secs(30)
387    } else {
388        Duration::from_secs(60 * 5)
389    };
390    let timeout = cx.background_executor().timer(timeout_duration).shared();
391
392    let (mut started_tx, mut started_rx) = mpsc::channel(servers_pending_start.len().max(1));
393    let (mut diag_tx, mut diag_rx) = mpsc::channel(servers_pending_diagnostics.len().max(1));
394    let subscriptions = [cx.subscribe(&lsp_store, {
395        let step_progress = step_progress.clone();
396        move |lsp_store, event, cx| match event {
397            project::LspStoreEvent::LanguageServerAdded(id, name, _) => {
398                step_progress.set_substatus(format!("LSP started: {}", name));
399                started_tx.try_send(*id).ok();
400            }
401            project::LspStoreEvent::DiskBasedDiagnosticsFinished { language_server_id } => {
402                let name = lsp_store
403                    .read(cx)
404                    .language_server_adapter_for_id(*language_server_id)
405                    .unwrap()
406                    .name();
407                step_progress.set_substatus(format!("LSP idle: {}", name));
408                diag_tx.try_send(*language_server_id).ok();
409            }
410            project::LspStoreEvent::LanguageServerUpdate {
411                message:
412                    client::proto::update_language_server::Variant::WorkProgress(
413                        client::proto::LspWorkProgress {
414                            message: Some(message),
415                            ..
416                        },
417                    ),
418                ..
419            } => {
420                step_progress.set_substatus(message.clone());
421            }
422            _ => {}
423        }
424    })];
425
426    // Phase 1: wait for all servers to start.
427    while !servers_pending_start.is_empty() {
428        futures::select! {
429            id = started_rx.next() => {
430                if let Some(id) = id {
431                    servers_pending_start.remove(&id);
432                }
433            },
434            _ = timeout.clone().fuse() => {
435                return Err(anyhow::anyhow!("LSP wait timed out after {} minutes", timeout_duration.as_secs() / 60));
436            }
437        }
438    }
439
440    // Save the buffer so the server sees the current content and kicks off diagnostics.
441    project
442        .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))
443        .await?;
444
445    // Phase 2: wait for all servers to finish their diagnostic pass.
446    while !servers_pending_diagnostics.is_empty() {
447        futures::select! {
448            id = diag_rx.next() => {
449                if let Some(id) = id {
450                    servers_pending_diagnostics.remove(&id);
451                }
452            },
453            _ = timeout.clone().fuse() => {
454                return Err(anyhow::anyhow!("LSP wait timed out after {} minutes", timeout_duration.as_secs() / 60));
455            }
456        }
457    }
458
459    drop(subscriptions);
460    step_progress.clear_substatus();
461    Ok(())
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use zeta_prompt::{ContextSource, RelatedExcerpt, RelatedFile};
468
469    fn excerpt(
470        row_range: std::ops::Range<u32>,
471        text: &str,
472        order: usize,
473        context_source: ContextSource,
474    ) -> RelatedExcerpt {
475        RelatedExcerpt {
476            row_range,
477            text: Arc::from(text),
478            order,
479            context_source,
480        }
481    }
482
483    #[test]
484    fn test_coalesce_touching_excerpts() {
485        let mut excerpts = vec![
486            excerpt(2..4, "line2\nline3\n", 1, ContextSource::EditHistory),
487            excerpt(0..2, "line0\nline1\n", 3, ContextSource::Bm25),
488            excerpt(4..6, "line4\nline5\n", 2, ContextSource::OracleSnippet),
489            excerpt(10..12, "line10\nline11\n", 0, ContextSource::Bm25),
490        ];
491        coalesce_touching_excerpts(&mut excerpts);
492
493        assert_eq!(excerpts.len(), 2);
494        assert_eq!(excerpts[0].row_range, 0..6);
495        assert_eq!(
496            excerpts[0].text.as_ref(),
497            "line0\nline1\nline2\nline3\nline4\nline5\n"
498        );
499        // The merged excerpt keeps the highest priority (minimum order) and
500        // its context source.
501        assert_eq!(excerpts[0].order, 1);
502        assert_eq!(excerpts[0].context_source, ContextSource::EditHistory);
503        // The gapped excerpt stays separate.
504        assert_eq!(excerpts[1].row_range, 10..12);
505    }
506
507    #[test]
508    fn test_coalesce_overlapping_excerpts_drops_duplicated_rows() {
509        let mut excerpts = vec![
510            excerpt(0..3, "line0\nline1\nline2\n", 0, ContextSource::Bm25),
511            excerpt(
512                2..5,
513                "line2\nline3\nline4\n",
514                1,
515                ContextSource::OracleSnippet,
516            ),
517        ];
518        coalesce_touching_excerpts(&mut excerpts);
519
520        assert_eq!(excerpts.len(), 1);
521        assert_eq!(excerpts[0].row_range, 0..5);
522        assert_eq!(
523            excerpts[0].text.as_ref(),
524            "line0\nline1\nline2\nline3\nline4\n"
525        );
526        assert_eq!(excerpts[0].order, 0);
527        assert_eq!(excerpts[0].context_source, ContextSource::Bm25);
528    }
529
530    #[test]
531    fn test_coalesce_contained_excerpt_upgrades_order() {
532        let mut excerpts = vec![
533            excerpt(0..4, "line0\nline1\nline2\nline3\n", 5, ContextSource::Bm25),
534            excerpt(1..3, "line1\nline2\n", 2, ContextSource::OracleSnippet),
535        ];
536        coalesce_touching_excerpts(&mut excerpts);
537
538        assert_eq!(excerpts.len(), 1);
539        assert_eq!(excerpts[0].row_range, 0..4);
540        assert_eq!(excerpts[0].text.as_ref(), "line0\nline1\nline2\nline3\n");
541        assert_eq!(excerpts[0].order, 2);
542        assert_eq!(excerpts[0].context_source, ContextSource::OracleSnippet);
543    }
544
545    #[test]
546    fn test_coalesce_handles_unterminated_final_line() {
547        // An excerpt ending without a newline includes the content of its
548        // `row_range.end` row (e.g. git-log excerpts ending at EOF), so a
549        // touching excerpt's first row duplicates it.
550        let mut excerpts = vec![
551            excerpt(0..2, "line0\nline1\nline2", 0, ContextSource::GitLog),
552            excerpt(2..4, "line2\nline3\n", 1, ContextSource::Bm25),
553        ];
554        coalesce_touching_excerpts(&mut excerpts);
555
556        assert_eq!(excerpts.len(), 1);
557        assert_eq!(excerpts[0].row_range, 0..4);
558        assert_eq!(excerpts[0].text.as_ref(), "line0\nline1\nline2\nline3\n");
559    }
560
561    #[test]
562    fn test_merge_context_files_coalesces_across_sources() {
563        let path: Arc<Path> = Path::new("root/src/lib.rs").into();
564        let mut context_files = vec![RelatedFile {
565            path: path.clone(),
566            max_row: 100,
567            excerpts: vec![excerpt(
568                0..2,
569                "line0\nline1\n",
570                0,
571                ContextSource::CurrentFile,
572            )],
573            in_open_source_repo: false,
574        }];
575        let new_files = vec![RelatedFile {
576            path,
577            max_row: 100,
578            excerpts: vec![excerpt(2..4, "line2\nline3\n", 1, ContextSource::Bm25)],
579            in_open_source_repo: false,
580        }];
581        merge_context_files(&mut context_files, new_files);
582
583        assert_eq!(context_files.len(), 1);
584        assert_eq!(context_files[0].excerpts.len(), 1);
585        assert_eq!(context_files[0].excerpts[0].row_range, 0..4);
586        assert_eq!(
587            context_files[0].excerpts[0].text.as_ref(),
588            "line0\nline1\nline2\nline3\n"
589        );
590    }
591}
592