1use std::collections::HashMap;
15use std::sync::mpsc;
16use std::time::Duration;
17
18use crate::core::graph_provider;
19use crate::core::tokens::count_tokens;
20use crate::tools::CrpMode;
21
22const DEFAULT_SEMANTIC_BUDGET_MS: u64 = 4000;
28
29fn semantic_budget() -> Duration {
30 let ms = std::env::var("LEAN_CTX_COMPOSE_BUDGET_MS")
31 .ok()
32 .and_then(|v| v.parse::<u64>().ok())
33 .filter(|&v| v > 0)
34 .unwrap_or(DEFAULT_SEMANTIC_BUDGET_MS);
35 Duration::from_millis(ms)
36}
37
38const DEFAULT_SYMBOL_BUDGET_TOKENS: usize = 600;
42
43fn symbol_budget_tokens() -> usize {
44 std::env::var("LEAN_CTX_COMPOSE_SYMBOL_TOKENS")
45 .ok()
46 .and_then(|v| v.parse::<usize>().ok())
47 .filter(|&v| v > 0)
48 .unwrap_or(DEFAULT_SYMBOL_BUDGET_TOKENS)
49}
50
51const DEFAULT_GRAPH_BUDGET_MS: u64 = 1500;
56
57fn graph_budget() -> Duration {
58 let ms = std::env::var("LEAN_CTX_COMPOSE_GRAPH_BUDGET_MS")
59 .ok()
60 .and_then(|v| v.parse::<u64>().ok())
61 .filter(|&v| v > 0)
62 .unwrap_or(DEFAULT_GRAPH_BUDGET_MS);
63 Duration::from_millis(ms)
64}
65
66const SPREAD_DECAY: f64 = 0.6;
70const SPREAD_HOPS: usize = 3;
71const SPREAD_TOP_K: usize = 8;
73
74fn build_associative_block(project_root: &str, keywords: &[String]) -> String {
80 let Some(open) = graph_provider::open_or_build(project_root) else {
81 return String::new();
82 };
83 let gp = &open.provider;
84
85 let mut seed_files: Vec<String> = Vec::new();
87 for kw in keywords {
88 for sym in gp.find_symbols(kw, None, None) {
89 if !seed_files.contains(&sym.file) {
90 seed_files.push(sym.file);
91 }
92 }
93 }
94 if seed_files.is_empty() {
95 return String::new();
96 }
97
98 crate::core::cooccurrence::record_access(project_root, &seed_files);
101
102 let mut adjacency: HashMap<String, Vec<(String, f64)>> = HashMap::new();
105 let mut add_edge = |a: &str, b: &str, w: f64| {
106 adjacency
107 .entry(a.to_string())
108 .or_default()
109 .push((b.to_string(), w));
110 adjacency
111 .entry(b.to_string())
112 .or_default()
113 .push((a.to_string(), w));
114 };
115 for e in gp.edges() {
116 add_edge(&e.from, &e.to, if e.weight > 0.0 { e.weight } else { 1.0 });
117 }
118 let coaccess = crate::core::cooccurrence::load(project_root);
119 for sf in &seed_files {
120 for (nbr, w) in coaccess.related(sf, 16) {
121 add_edge(sf, &nbr, w);
122 }
123 }
124
125 let seeds: HashMap<String, f64> = seed_files.iter().map(|f| (f.clone(), 1.0)).collect();
126 let ranked = crate::core::spreading_activation::related_ranked(
127 &seeds,
128 &adjacency,
129 SPREAD_DECAY,
130 SPREAD_HOPS,
131 SPREAD_TOP_K,
132 );
133 if ranked.is_empty() {
134 return String::new();
135 }
136
137 let mut s = String::from("\n## Related (associative: import/call graph + learned co-access)\n");
138 for (file, activation) in ranked {
139 let file = crate::core::protocol::display_path(&file);
142 s.push_str(&format!("- {file} (activation {activation:.2})\n"));
143 }
144 s
145}
146
147fn associative_block_budgeted(project_root: &str, keywords: &[String]) -> String {
151 if keywords.is_empty() {
152 return String::new();
153 }
154 let (tx, rx) = mpsc::channel::<String>();
155 let root = project_root.to_string();
156 let kws = keywords.to_vec();
157 std::thread::spawn(move || {
158 let block = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
159 build_associative_block(&root, &kws)
160 }))
161 .unwrap_or_else(|_| {
162 tracing::warn!("[ctx_compose: associative block panicked; omitting section]");
163 String::new()
164 });
165 let _ = tx.send(block);
166 });
167 rx.recv_timeout(graph_budget()).unwrap_or_default()
168}
169
170const STOPWORDS: &[&str] = &[
172 "the",
173 "and",
174 "for",
175 "with",
176 "that",
177 "this",
178 "from",
179 "into",
180 "how",
181 "where",
182 "what",
183 "does",
184 "are",
185 "was",
186 "use",
187 "used",
188 "uses",
189 "add",
190 "all",
191 "any",
192 "can",
193 "get",
194 "set",
195 "via",
196 "out",
197 "its",
198 "his",
199 "her",
200 "you",
201 "your",
202 "our",
203 "find",
204 "show",
205 "list",
206 "make",
207 "when",
208 "then",
209 "has",
210 "have",
211 "had",
212 "not",
213 "but",
214 "see",
215 "function",
216 "method",
217 "class",
218 "code",
219 "file",
220 "files",
221 "implement",
222 "implementation",
223];
224
225fn extract_keywords(task: &str, max: usize) -> Vec<String> {
228 let mut seen = std::collections::HashSet::new();
229 let mut out = Vec::new();
230 for raw in task.split(|c: char| !(c.is_alphanumeric() || c == '_')) {
231 if raw.len() < 3 {
232 continue;
233 }
234 if STOPWORDS.contains(&raw.to_ascii_lowercase().as_str()) {
235 continue;
236 }
237 if seen.insert(raw.to_string()) {
238 out.push(raw.to_string());
239 if out.len() >= max {
240 break;
241 }
242 }
243 }
244 out
245}
246
247fn order_by_specificity(keywords: &[String], project_root: &str) -> Vec<String> {
256 let Some(index) = resident_index(project_root) else {
257 return keywords.to_vec();
258 };
259 rank_by_doc_freq(keywords, &index.doc_freqs)
260}
261
262fn rank_by_doc_freq(
274 keywords: &[String],
275 doc_freqs: &std::collections::HashMap<String, usize>,
276) -> Vec<String> {
277 let df = |kw: &String| match doc_freqs.get(&kw.to_ascii_lowercase()) {
278 Some(&n) if n > 0 => n,
279 _ => usize::MAX,
280 };
281 let rank_key = |kw: &String| (u8::from(!is_code_identifier(kw)), df(kw));
283 let mut ranked = keywords.to_vec();
284 ranked.sort_by_key(rank_key);
285 ranked
286}
287
288fn is_code_identifier(kw: &str) -> bool {
293 if kw.contains('_') {
294 return true;
295 }
296 let has_lower = kw.chars().any(|c| c.is_ascii_lowercase());
297 let internal_upper = kw.chars().skip(1).any(|c| c.is_ascii_uppercase());
298 has_lower && internal_upper
299}
300
301fn resident_index(
304 project_root: &str,
305) -> Option<std::sync::Arc<crate::core::bm25_index::BM25Index>> {
306 let cache = crate::tools::ctx_semantic_search::get_thread_cache()?;
307 crate::core::bm25_cache::get_or_background(&cache, std::path::Path::new(project_root))
308}
309
310fn ranked_files_budgeted(task: &str, project_root: &str, crp_mode: CrpMode) -> String {
314 let shared_cache = crate::tools::ctx_semantic_search::get_thread_cache();
315 let (tx, rx) = mpsc::channel::<String>();
316 let task_owned = task.to_string();
317 let root_owned = project_root.to_string();
318
319 std::thread::spawn(move || {
320 if let Some(cache) = shared_cache {
321 crate::tools::ctx_semantic_search::set_thread_cache(cache);
322 }
323 let ranked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
324 crate::tools::ctx_semantic_search::handle(
325 &task_owned,
326 &root_owned,
327 8,
328 crp_mode,
329 None,
330 None,
331 None,
332 Some(false),
333 Some(false),
334 )
335 }))
336 .unwrap_or_else(|_| {
337 tracing::warn!("[ctx_compose: semantic ranking panicked; omitting section]");
338 String::new()
339 });
340 let _ = tx.send(ranked);
343 });
344
345 match rx.recv_timeout(semantic_budget()) {
346 Ok(ranked) => ranked.trim().to_string(),
347 Err(_) => deferred_ranking_note(project_root),
348 }
349}
350
351fn deferred_ranking_note(project_root: &str) -> String {
360 let exact = "the exact matches below are authoritative for this call";
361 let s = crate::core::index_orchestrator::bm25_summary(project_root);
362 match s.state {
363 "failed" => {
364 let why = s
365 .last_error
366 .or(s.note)
367 .unwrap_or_else(|| "unknown error".to_string());
368 format!(
369 "(semantic ranking unavailable — index build FAILED: {why}. {exact}. \
370 Inspect with `ctx_index status` / `lean-ctx doctor`, then `lean-ctx reindex`)"
371 )
372 }
373 "building" => format!(
374 "(deferred — semantic index is building; {exact}, \
375 and ranking becomes available once the build finishes)"
376 ),
377 _ => match s.note {
381 Some(note) if note.contains("NOT persisted") => {
382 format!("(semantic ranking deferred — {note} {exact}.)")
383 }
384 _ => format!(
385 "(deferred — semantic index is warming; {exact}, \
386 and ranking will be fast on the next call once the index is cached)"
387 ),
388 },
389 }
390}
391
392pub fn handle(task: &str, project_root: &str, crp_mode: CrpMode) -> (String, usize) {
394 let task = task.trim();
395 if task.is_empty() {
396 return ("ERROR: task is required".to_string(), 0);
397 }
398
399 let keywords = extract_keywords(task, 6);
400 let allow_secret = crate::core::roles::active_role().io.allow_secret_paths;
401
402 let mut out = String::new();
403 out.push_str(&format!("TASK: {task}\n"));
404 if keywords.is_empty() {
405 out.push_str("KEYWORDS: (none extracted — using full task for ranking)\n");
406 } else {
407 out.push_str(&format!("KEYWORDS: {}\n", keywords.join(", ")));
408 }
409
410 out.push_str("\n## Ranked files (semantic)\n");
415 out.push_str(&ranked_files_budgeted(task, project_root, crp_mode));
416 out.push('\n');
417
418 let ranked_keywords = order_by_specificity(&keywords, project_root);
422 if let Some(primary) = ranked_keywords
423 .iter()
424 .find(|keyword| is_code_identifier(keyword))
425 {
426 let grep = crate::tools::ctx_search::handle(
427 primary,
428 project_root,
429 None,
430 10,
431 crp_mode,
432 true,
433 allow_secret,
434 false,
435 )
436 .text;
437 out.push_str(&format!("\n## Exact matches: '{primary}'\n"));
438 out.push_str(grep.trim());
439 out.push('\n');
440 }
441
442 use crate::core::context_packing::{CoverageItem, greedy_max_coverage};
448 let mut snippets: Vec<String> = Vec::new();
449 let mut items: Vec<CoverageItem> = Vec::new();
450 for kw in &keywords {
451 if let Some((rendered, toks)) =
452 crate::tools::ctx_symbol::best_symbol_snippet_for_task(kw, task, project_root)
453 {
454 let mut terms: std::collections::HashSet<String> =
457 std::collections::HashSet::from([kw.clone()]);
458 for other in &keywords {
459 if other != kw && rendered.contains(other.as_str()) {
460 terms.insert(other.clone());
461 }
462 }
463 items.push(CoverageItem {
464 terms,
465 cost: toks.max(1),
466 });
467 snippets.push(rendered);
468 }
469 }
470 if !items.is_empty() {
471 let chosen = greedy_max_coverage(&items, symbol_budget_tokens(), |_| 1.0);
472 let mut seen = std::collections::HashSet::new();
473 let mut header_written = false;
474 for idx in chosen {
475 let rendered = snippets[idx].trim();
476 if rendered.is_empty() || !seen.insert(rendered.to_string()) {
477 continue;
478 }
479 if !header_written {
480 out.push_str("\n## Top symbols (bodies)\n");
481 header_written = true;
482 }
483 out.push_str(rendered);
484 out.push('\n');
485 }
486 }
487
488 out.push_str(&associative_block_budgeted(project_root, &keywords));
492
493 {
497 use crate::core::context_kernel::activation::{load_config, supplement_budget};
498 use crate::core::context_kernel::context_dedup::dedup_kernel_blocks;
499
500 let config = load_config(project_root);
501 let budget = symbol_budget_tokens() / 5;
502 let budget = budget
503 .min(config.max_supplement_tokens)
504 .min(supplement_budget(&config));
505 if let Some(enrichment) =
506 crate::core::context_kernel::bridge::kernel_enrich(task, project_root, budget)
507 .filter(|enrichment| !enrichment.blocks.is_empty())
508 {
509 let blocks =
510 dedup_kernel_blocks(&enrichment.blocks, &mut std::collections::HashSet::new());
511 if !blocks.is_empty() {
512 out.push_str("\n## Context Kernel\n");
513 out.push_str(&blocks);
514 out.push('\n');
515 }
516 }
517 }
518
519 let sent = count_tokens(&out);
520 (out, sent)
521}
522
523#[cfg(test)]
524mod tests {
525 use super::*;
526
527 #[test]
528 fn rank_by_doc_freq_puts_rare_identifier_first() {
529 let keywords = vec![
533 "OCPP".to_string(),
534 "GetMaxCurrent".to_string(),
535 "Current".to_string(),
536 ];
537 let doc_freqs = std::collections::HashMap::from([
538 ("ocpp".to_string(), 120),
539 ("current".to_string(), 400),
540 ("getmaxcurrent".to_string(), 3),
541 ]);
542 let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
543 assert_eq!(ranked.first().unwrap(), "GetMaxCurrent");
544 assert_eq!(ranked.last().unwrap(), "Current");
545 }
546
547 #[test]
548 fn rank_by_doc_freq_sinks_absent_tokens_and_is_stable() {
549 let keywords = vec![
552 "absent".to_string(),
553 "alpha".to_string(),
554 "beta".to_string(),
555 ];
556 let doc_freqs =
557 std::collections::HashMap::from([("alpha".to_string(), 5), ("beta".to_string(), 5)]);
558 let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
559 assert_eq!(ranked, vec!["alpha", "beta", "absent"]);
560 }
561
562 #[test]
563 fn rank_prefers_code_identifier_over_rarer_prose_word() {
564 let keywords = vec!["measurand".to_string(), "GetMaxCurrent".to_string()];
569 let doc_freqs = std::collections::HashMap::from([
570 ("measurand".to_string(), 4),
571 ("getmaxcurrent".to_string(), 30),
572 ]);
573 let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
574 assert_eq!(ranked.first().unwrap(), "GetMaxCurrent");
575 }
576
577 #[test]
578 fn is_code_identifier_classifies_camel_snake_vs_prose_and_acronym() {
579 assert!(is_code_identifier("GetMaxCurrent"));
580 assert!(is_code_identifier("CurrentGetter"));
581 assert!(is_code_identifier("get_max_current"));
582 assert!(!is_code_identifier("Current"));
584 assert!(!is_code_identifier("OCPP"));
585 assert!(!is_code_identifier("charger"));
586 }
587
588 #[test]
589 fn extract_keywords_drops_stopwords_and_short_tokens() {
590 let kw = extract_keywords("How does the BM25Index cache work for ctx_search?", 6);
591 assert!(kw.contains(&"BM25Index".to_string()));
592 assert!(kw.contains(&"cache".to_string()));
593 assert!(kw.contains(&"ctx_search".to_string()));
594 assert!(!kw.iter().any(|k| k == "the" || k == "How" || k == "for"));
595 }
596
597 #[test]
598 fn extract_keywords_dedups_and_caps() {
599 let kw = extract_keywords("alpha alpha beta gamma delta epsilon zeta eta", 3);
600 assert_eq!(kw.len(), 3);
601 assert_eq!(kw[0], "alpha");
602 }
603
604 #[test]
605 fn exact_matches_choose_specific_identifier_not_first_broad_keyword() {
606 let keywords = extract_keywords(
607 "OCPP charger GetMaxCurrent Current.Offered measurand CurrentGetter",
608 6,
609 );
610 assert!(keywords.iter().any(|keyword| keyword == "GetMaxCurrent"));
611 assert!(keywords.iter().any(|keyword| is_code_identifier(keyword)));
612 assert!(!is_code_identifier("OCPP"));
613
614 let prose = extract_keywords("Fix semantic ranking exact matches", 6);
615 assert!(prose.iter().all(|keyword| !is_code_identifier(keyword)));
616 }
617
618 #[test]
619 fn empty_task_is_rejected() {
620 let (out, tok) = handle(" ", "/tmp", CrpMode::Off);
621 assert!(out.starts_with("ERROR"));
622 assert_eq!(tok, 0);
623 }
624
625 #[test]
626 fn handle_includes_context_kernel_section_when_available() {
627 let (output, tokens) = handle("find authentication bugs", "/tmp/nonexistent", CrpMode::Tdd);
628 assert!(tokens > 0);
631 assert!(output.contains("TASK:"));
632 }
633
634 #[test]
635 fn deferred_ranking_note_is_deterministic_and_has_no_timing() {
636 let tmp = tempfile::tempdir().unwrap();
639 let root = tmp.path().to_string_lossy();
640 let a = deferred_ranking_note(root.as_ref());
641 let b = deferred_ranking_note(root.as_ref());
642 assert_eq!(a, b, "deferred note must be byte-stable across calls");
643 assert!(
644 !a.contains("elapsed"),
645 "deferred note must not embed timing data: {a}"
646 );
647 }
648
649 #[test]
650 fn deferred_note_for_idle_index_is_optimistic_but_honest() {
651 let tmp = tempfile::tempdir().unwrap();
655 let note = deferred_ranking_note(tmp.path().to_string_lossy().as_ref());
656 assert!(
657 note.contains("warming") || note.contains("building"),
658 "note: {note}"
659 );
660 assert!(
661 note.contains("authoritative"),
662 "note must reassure that exact matches are authoritative: {note}"
663 );
664 assert!(
665 !note.contains("instant on the next call"),
666 "must not repeat the dishonest 'instant next call' promise: {note}"
667 );
668 }
669}