lean_ctx/tools/
ctx_fill.rs1use std::path::Path;
2
3use crate::core::cache::SessionCache;
4use crate::core::signatures;
5use crate::core::tokens::count_tokens;
6use crate::tools::CrpMode;
7
8struct FileCandidate {
9 path: String,
10 score: f64,
11 tokens_full: usize,
12 tokens_map: usize,
13 tokens_sig: usize,
14}
15
16pub fn handle(
17 cache: &mut SessionCache,
18 paths: &[String],
19 budget: usize,
20 crp_mode: CrpMode,
21 task: Option<&str>,
22) -> String {
23 if paths.is_empty() {
24 return "No files specified.".to_string();
25 }
26
27 let pagerank_scores = load_pagerank_scores(paths);
28 let mut candidates: Vec<FileCandidate> = Vec::new();
29
30 for path in paths {
31 let Ok(content) = std::fs::read_to_string(path) else {
32 continue;
33 };
34
35 let ext = Path::new(path)
36 .extension()
37 .and_then(|e| e.to_str())
38 .unwrap_or("");
39 let tokens_full = count_tokens(&content);
40 let sigs = signatures::extract_signatures(&content, ext);
41 let sig_text: String = sigs
42 .iter()
43 .map(super::super::core::signatures::Signature::to_compact)
44 .collect::<Vec<_>>()
45 .join("\n");
46 let tokens_sig = count_tokens(&sig_text);
47
48 let map_text = format_map(&content, ext, &sigs);
49 let tokens_map = count_tokens(&map_text);
50
51 let mut score = compute_relevance_score(path, &content);
52 if let Some(pr_boost) = pagerank_scores.get(path) {
53 score *= 1.0 + pr_boost * 5.0;
54 }
55
56 candidates.push(FileCandidate {
57 path: path.clone(),
58 score,
59 tokens_full,
60 tokens_map,
61 tokens_sig,
62 });
63 }
64
65 candidates.sort_by(|a, b| {
66 b.score
67 .partial_cmp(&a.score)
68 .unwrap_or(std::cmp::Ordering::Equal)
69 });
70
71 let mut pop_lines: Vec<String> = Vec::new();
72 if let Some(t) = task
73 && let Some(root) = paths
74 .first()
75 .and_then(|p| crate::core::protocol::detect_project_root(p))
76 {
77 let rs: Vec<crate::core::task_relevance::RelevanceScore> = candidates
78 .iter()
79 .map(|c| crate::core::task_relevance::RelevanceScore {
80 path: c.path.clone(),
81 score: c.score,
82 recommended_mode: "signatures",
83 })
84 .collect();
85 let refs: Vec<&crate::core::task_relevance::RelevanceScore> = rs.iter().collect();
86 let pop = crate::core::pop_pruning::decide_for_candidates(t, &root, &refs);
87 if !pop.excluded_modules.is_empty() {
88 let excluded: std::collections::BTreeSet<&str> = pop
89 .excluded_modules
90 .iter()
91 .map(|e| e.module.as_str())
92 .collect();
93 candidates.retain(|c| {
94 let m = crate::core::pop_pruning::module_for_path(&c.path, &root);
95 !excluded.contains(m.as_str())
96 });
97 pop_lines.push("POP:".to_string());
98 for ex in &pop.excluded_modules {
99 pop_lines.push(format!(
100 " - exclude {}/ ({} candidates) — {}",
101 ex.module, ex.candidate_files, ex.reason
102 ));
103 }
104 }
105 }
106
107 let mut used_tokens = 0usize;
108 let mut selections: Vec<(String, String)> = Vec::new();
109
110 for candidate in &candidates {
111 if used_tokens >= budget {
112 break;
113 }
114
115 if crate::tools::ctx_read::is_instruction_file(&candidate.path) {
116 selections.push((candidate.path.clone(), "full".to_string()));
117 used_tokens += candidate.tokens_full;
118 continue;
119 }
120
121 let remaining = budget - used_tokens;
122 let (mode, cost) = select_best_fit(candidate, remaining);
123
124 if cost > remaining {
125 let sig_cost = candidate.tokens_sig;
126 if sig_cost <= remaining {
127 selections.push((candidate.path.clone(), "signatures".to_string()));
128 used_tokens += sig_cost;
129 }
130 continue;
131 }
132
133 selections.push((candidate.path.clone(), mode));
134 used_tokens += cost;
135 }
136
137 let mut output_parts = Vec::new();
138 output_parts.push(format!(
139 "ctx_fill: {budget} token budget, {} files analyzed, {} selected",
140 candidates.len(),
141 selections.len()
142 ));
143 if !pop_lines.is_empty() {
144 output_parts.push(pop_lines.join("\n"));
145 }
146 output_parts.push(String::new());
147
148 for (path, mode) in &selections {
149 let result = crate::tools::ctx_read::handle(cache, path, mode, crp_mode);
150 output_parts.push(result);
151 output_parts.push("---".to_string());
152 }
153
154 let skipped = candidates.len() - selections.len();
155 if skipped > 0 {
156 output_parts.push(format!("{skipped} files skipped (budget exhausted)"));
157 }
158 output_parts.push(format!("\nUsed: {used_tokens}/{budget} tokens"));
159
160 output_parts.join("\n")
161}
162
163fn select_best_fit(candidate: &FileCandidate, remaining: usize) -> (String, usize) {
164 if candidate.tokens_full <= remaining {
165 return ("full".to_string(), candidate.tokens_full);
166 }
167 if candidate.tokens_map <= remaining {
168 return ("map".to_string(), candidate.tokens_map);
169 }
170 if candidate.tokens_sig <= remaining {
171 return ("signatures".to_string(), candidate.tokens_sig);
172 }
173 ("signatures".to_string(), candidate.tokens_sig)
174}
175
176fn compute_relevance_score(path: &str, content: &str) -> f64 {
177 let mut score = 1.0;
178
179 let name = Path::new(path)
180 .file_name()
181 .and_then(|n| n.to_str())
182 .unwrap_or("");
183 if name.contains("test") || name.contains("spec") {
184 score *= 0.5;
185 }
186 if name.contains("config") || name.contains("types") || name.contains("schema") {
187 score *= 1.3;
188 }
189 if name == "mod.rs" || name == "index.ts" || name == "index.js" || name == "__init__.py" {
190 score *= 1.5;
191 }
192
193 let ext = Path::new(path)
194 .extension()
195 .and_then(|e| e.to_str())
196 .unwrap_or("");
197 if matches!(ext, "rs" | "ts" | "py" | "go" | "java") {
198 score *= 1.2;
199 }
200
201 let lines = content.lines().count();
202 if lines > 500 {
203 score *= 0.8;
204 }
205 if lines < 50 {
206 score *= 1.1;
207 }
208
209 let export_count = content
210 .lines()
211 .filter(|l| l.contains("pub ") || l.contains("export ") || l.contains("def "))
212 .count();
213 score *= 1.0 + (export_count as f64 * 0.02).min(0.5);
214
215 score
216}
217
218fn load_pagerank_scores(paths: &[String]) -> std::collections::HashMap<String, f64> {
219 let root = paths
220 .first()
221 .and_then(|p| crate::core::protocol::detect_project_root(p));
222
223 let Some(root) = root else {
224 return std::collections::HashMap::new();
225 };
226
227 let Ok(graph) = crate::core::property_graph::CodeGraph::open(&root) else {
228 return std::collections::HashMap::new();
229 };
230
231 if graph.node_count().unwrap_or(0) == 0 {
232 return std::collections::HashMap::new();
233 }
234
235 let top = crate::core::pagerank::top_files(graph.connection(), 200);
236 top.into_iter().collect()
237}
238
239fn format_map(content: &str, ext: &str, sigs: &[crate::core::signatures::Signature]) -> String {
240 let deps = crate::core::deps::extract_deps(content, ext);
241 let mut parts = Vec::new();
242 if !deps.imports.is_empty() {
243 parts.push(format!("deps: {}", deps.imports.join(", ")));
244 }
245 if !deps.exports.is_empty() {
246 parts.push(format!("exports: {}", deps.exports.join(", ")));
247 }
248 let key_sigs: Vec<_> = sigs
249 .iter()
250 .filter(|s| s.is_exported || s.indent == 0)
251 .collect();
252 for sig in &key_sigs {
253 parts.push(sig.to_compact());
254 }
255 parts.join("\n")
256}