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

lean_ctx/tools/
ctx_prefetch.rs
