lean_ctx/tools/
ctx_prefetch.rs

1use std::collections::{BTreeMap, BTreeSet, VecDeque};
2use std::path::Path;
3
4use crate::core::cache::SessionCache;
5use crate::core::graph_index::ProjectIndex;
6use crate::core::protocol;
7use crate::core::task_relevance::{compute_relevance, parse_task_hints};
8use crate::tools::CrpMode;
9
10const DEFAULT_MAX_FILES: usize = 10;
11
12pub fn handle(
13    cache: &mut SessionCache,
14    root: &str,
15    task: Option<&str>,
16    changed_files: Option<&[String]>,
17    budget_tokens: usize,
18    max_files: Option<usize>,
19    crp_mode: CrpMode,
20) -> String {
21    let project_root = if root.trim().is_empty() { "." } else { root };
22    let index = crate::core::graph_index::load_or_build(project_root);
23
24    let mut candidates: BTreeMap<String, f64> = BTreeMap::new(); // path -> score
25
26    if let Some(t) = task {
27        let (task_files, task_keywords) = parse_task_hints(t);
28        let relevance = compute_relevance(&index, &task_files, &task_keywords);
29        for r in relevance.iter().take(50) {
30            if r.score < 0.1 {
31                break;
32            }
33            candidates.insert(r.path.clone(), r.score);
34        }
35    }
36
37    if let Some(changed) = changed_files {
38        for p in changed {
39            let rel = normalize_rel_path(p, project_root);
40            for (path, dist) in blast_radius(&index, &rel, 2) {
41                let boost = 1.0 / (dist.max(1) as f64);
42                candidates
43                    .entry(path)
44                    .and_modify(|s| *s = (*s + boost).min(1.0))
45                    .or_insert(boost.min(1.0));
46            }
47        }
48    }
49
50    if candidates.is_empty() {
51        return "ctx_prefetch: no candidates (provide task or changed_files)".to_string();
52    }
53
54    let mut scored: Vec<(String, f64)> = candidates.into_iter().collect();
55    scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
56
57    let max_files = max_files.unwrap_or(DEFAULT_MAX_FILES).max(1);
58    let mut picked: Vec<String> = Vec::new();
59    for (p, _s) in scored {
60        picked.push(p);
61        if picked.len() >= max_files {
62            break;
63        }
64    }
65
66    let mut total = 0usize;
67    let mut prefetched: Vec<(String, String)> = Vec::new(); // (path, mode)
68    let jail_root = Path::new(project_root);
69    for p in &picked {
70        let full = to_fs_path(project_root, p);
71        let Ok((jailed, warning)) = crate::core::io_boundary::jail_and_check_path(
72            "ctx_prefetch",
73            Path::new(&full),
74            jail_root,
75        ) else {
76            continue;
77        };
78        if warning.is_some() {
79            continue;
80        }
81        let jailed_s = jailed.to_string_lossy().to_string();
82
83        let Ok(content) = std::fs::read_to_string(&jailed) else {
84            continue;
85        };
86        let tokens = crate::core::tokens::count_tokens(&content);
87        total = total.saturating_add(tokens);
88
89        let mode = if budget_tokens > 0 {
90            let ratio = budget_tokens as f64 / total.max(1) as f64;
91            if ratio >= 0.8 {
92                "full"
93            } else if ratio >= 0.4 {
94                "map"
95            } else {
96                "signatures"
97            }
98        } else {
99            "signatures"
100        };
101
102        let _ = crate::tools::ctx_read::handle_with_task_resolved(
103            cache, &jailed_s, mode, crp_mode, task,
104        );
105        prefetched.push((jailed_s, mode.to_string()));
106    }
107
108    let mut lines = vec![
109        format!(
110            "ctx_prefetch: prefetched {} file(s) (max_files={})",
111            prefetched.len(),
112            max_files
113        ),
114        format!("  root: {}", project_root),
115    ];
116    for (p, mode) in prefetched.iter().take(20) {
117        let r = cache.get_file_ref(p);
118        let short = protocol::shorten_path(p);
119        lines.push(format!("  - [{r}] {short} mode={mode}"));
120    }
121    lines.join("\n")
122}
123
124fn blast_radius(index: &ProjectIndex, start_rel: &str, max_depth: usize) -> Vec<(String, usize)> {
125    let mut adj: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
126    for e in &index.edges {
127        adj.entry(e.from.as_str()).or_default().push(e.to.as_str());
128        adj.entry(e.to.as_str()).or_default().push(e.from.as_str());
129    }
130
131    let mut out = Vec::new();
132    let mut q: VecDeque<(String, usize)> = VecDeque::new();
133    let mut seen: BTreeSet<String> = BTreeSet::new();
134
135    q.push_back((start_rel.to_string(), 0));
136    seen.insert(start_rel.to_string());
137
138    while let Some((node, depth)) = q.pop_front() {
139        out.push((node.clone(), depth));
140        if depth >= max_depth {
141            continue;
142        }
143        if let Some(nbrs) = adj.get(node.as_str()) {
144            for &n in nbrs {
145                let ns = n.to_string();
146                if seen.insert(ns.clone()) {
147                    q.push_back((ns, depth + 1));
148                }
149            }
150        }
151    }
152    out
153}
154
155fn normalize_rel_path(path: &str, project_root: &str) -> String {
156    let p = Path::new(path);
157    if p.is_absolute() {
158        if let Ok(stripped) = p.strip_prefix(project_root) {
159            return stripped
160                .to_string_lossy()
161                .trim_start_matches('/')
162                .to_string();
163        }
164    }
165    path.trim_start_matches('/').to_string()
166}
167
168fn to_fs_path(project_root: &str, rel_or_abs: &str) -> String {
169    let p = Path::new(rel_or_abs);
170    if p.is_absolute() {
171        return rel_or_abs.to_string();
172    }
173    Path::new(project_root)
174        .join(rel_or_abs)
175        .to_string_lossy()
176        .to_string()
177}
lean_ctx/tools/ctx_prefetch.rs

lean_ctx/tools/
ctx_prefetch.rs
