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    for p in &picked {
69        let full = to_fs_path(project_root, p);
70        let Ok(content) = std::fs::read_to_string(&full) else {
71            continue;
72        };
73        let tokens = crate::core::tokens::count_tokens(&content);
74        total = total.saturating_add(tokens);
75
76        let mode = if budget_tokens > 0 {
77            let ratio = budget_tokens as f64 / total.max(1) as f64;
78            if ratio >= 0.8 {
79                "full"
80            } else if ratio >= 0.4 {
81                "map"
82            } else {
83                "signatures"
84            }
85        } else {
86            "signatures"
87        };
88
89        let _ =
90            crate::tools::ctx_read::handle_with_task_resolved(cache, &full, mode, crp_mode, task);
91        prefetched.push((full.clone(), mode.to_string()));
92    }
93
94    let mut lines = vec![
95        format!(
96            "ctx_prefetch: prefetched {} file(s) (max_files={})",
97            prefetched.len(),
98            max_files
99        ),
100        format!("  root: {}", project_root),
101    ];
102    for (p, mode) in prefetched.iter().take(20) {
103        let r = cache.get_file_ref(p);
104        let short = protocol::shorten_path(p);
105        lines.push(format!("  - [{r}] {short} mode={mode}"));
106    }
107    lines.join("\n")
108}
109
110fn blast_radius(index: &ProjectIndex, start_rel: &str, max_depth: usize) -> Vec<(String, usize)> {
111    let mut adj: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
112    for e in &index.edges {
113        adj.entry(e.from.as_str()).or_default().push(e.to.as_str());
114        adj.entry(e.to.as_str()).or_default().push(e.from.as_str());
115    }
116
117    let mut out = Vec::new();
118    let mut q: VecDeque<(String, usize)> = VecDeque::new();
119    let mut seen: BTreeSet<String> = BTreeSet::new();
120
121    q.push_back((start_rel.to_string(), 0));
122    seen.insert(start_rel.to_string());
123
124    while let Some((node, depth)) = q.pop_front() {
125        out.push((node.clone(), depth));
126        if depth >= max_depth {
127            continue;
128        }
129        if let Some(nbrs) = adj.get(node.as_str()) {
130            for &n in nbrs {
131                let ns = n.to_string();
132                if seen.insert(ns.clone()) {
133                    q.push_back((ns, depth + 1));
134                }
135            }
136        }
137    }
138    out
139}
140
141fn normalize_rel_path(path: &str, project_root: &str) -> String {
142    let p = Path::new(path);
143    if p.is_absolute() {
144        if let Ok(stripped) = p.strip_prefix(project_root) {
145            return stripped
146                .to_string_lossy()
147                .trim_start_matches('/')
148                .to_string();
149        }
150    }
151    path.trim_start_matches('/').to_string()
152}
153
154fn to_fs_path(project_root: &str, rel_or_abs: &str) -> String {
155    let p = Path::new(rel_or_abs);
156    if p.is_absolute() {
157        return rel_or_abs.to_string();
158    }
159    Path::new(project_root)
160        .join(rel_or_abs)
161        .to_string_lossy()
162        .to_string()
163}
lean_ctx/tools/ctx_prefetch.rs

lean_ctx/tools/
ctx_prefetch.rs
