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
392fn enrich_keywords_with_ib_intent(task: &str, keywords: Vec<String>) -> Vec<String> {
397 if !crate::core::cognitive_gate::basic_science_enabled() {
398 return keywords;
399 }
400
401 let fallback = keywords.clone();
402 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
403 use crate::core::ib::{classify_intent, intent_query_terms};
404 use crate::core::session::{SessionState, TaskInfo};
405
406 let mut session = SessionState::new();
407 session.task = Some(TaskInfo {
408 description: task.to_owned(),
409 intent: None,
410 progress_pct: None,
411 });
412 let intent = classify_intent(&session);
413 let mut enriched = keywords;
414 for term in intent_query_terms(&intent) {
415 let term = term.to_string();
416 if !enriched
417 .iter()
418 .any(|keyword| keyword.eq_ignore_ascii_case(&term))
419 {
420 enriched.push(term);
421 }
422 }
423 enriched
424 }))
425 .unwrap_or_else(|_| {
426 tracing::warn!("[ctx_compose: IB intent enrichment failed; using original keywords]");
427 fallback
428 })
429}
430
431pub fn handle(task: &str, project_root: &str, crp_mode: CrpMode) -> (String, usize) {
433 let task = task.trim();
434 if task.is_empty() {
435 return ("ERROR: task is required".to_string(), 0);
436 }
437
438 let keywords = enrich_keywords_with_ib_intent(task, extract_keywords(task, 6));
439 let allow_secret = crate::core::roles::active_role().io.allow_secret_paths;
440
441 let mut out = String::new();
442 out.push_str(&format!("TASK: {task}\n"));
443 if keywords.is_empty() {
444 out.push_str("KEYWORDS: (none extracted — using full task for ranking)\n");
445 } else {
446 out.push_str(&format!("KEYWORDS: {}\n", keywords.join(", ")));
447 }
448
449 out.push_str("\n## Ranked files (semantic)\n");
454 out.push_str(&ranked_files_budgeted(task, project_root, crp_mode));
455 out.push('\n');
456
457 let ranked_keywords = order_by_specificity(&keywords, project_root);
461 if let Some(primary) = ranked_keywords
462 .iter()
463 .find(|keyword| is_code_identifier(keyword))
464 {
465 let grep = crate::tools::ctx_search::handle(
466 primary,
467 project_root,
468 None,
469 10,
470 crp_mode,
471 true,
472 allow_secret,
473 false,
474 )
475 .text;
476 out.push_str(&format!("\n## Exact matches: '{primary}'\n"));
477 out.push_str(grep.trim());
478 out.push('\n');
479 }
480
481 use crate::core::context_packing::{CoverageItem, greedy_max_coverage};
487 let mut snippets: Vec<String> = Vec::new();
488 let mut items: Vec<CoverageItem> = Vec::new();
489 for kw in &keywords {
490 if let Some((rendered, toks)) =
491 crate::tools::ctx_symbol::best_symbol_snippet_for_task(kw, task, project_root)
492 {
493 let mut terms: std::collections::HashSet<String> =
496 std::collections::HashSet::from([kw.clone()]);
497 for other in &keywords {
498 if other != kw && rendered.contains(other.as_str()) {
499 terms.insert(other.clone());
500 }
501 }
502 items.push(CoverageItem {
503 terms,
504 cost: toks.max(1),
505 });
506 snippets.push(rendered);
507 }
508 }
509 if !items.is_empty() {
510 let chosen = greedy_max_coverage(&items, symbol_budget_tokens(), |_| 1.0);
511 let mut seen = std::collections::HashSet::new();
512 let mut header_written = false;
513 for idx in chosen {
514 let rendered = snippets[idx].trim();
515 if rendered.is_empty() || !seen.insert(rendered.to_string()) {
516 continue;
517 }
518 if !header_written {
519 out.push_str("\n## Top symbols (bodies)\n");
520 header_written = true;
521 }
522 out.push_str(rendered);
523 out.push('\n');
524 }
525 }
526
527 out.push_str(&associative_block_budgeted(project_root, &keywords));
531
532 {
536 use crate::core::context_kernel::activation::{load_config, supplement_budget};
537 use crate::core::context_kernel::context_dedup::dedup_kernel_blocks;
538
539 let config = load_config(project_root);
540 let budget = symbol_budget_tokens() / 5;
541 let budget = budget
542 .min(config.max_supplement_tokens)
543 .min(supplement_budget(&config));
544 if let Some(enrichment) =
545 crate::core::context_kernel::bridge::kernel_enrich(task, project_root, budget)
546 .filter(|enrichment| !enrichment.blocks.is_empty())
547 {
548 let blocks =
549 dedup_kernel_blocks(&enrichment.blocks, &mut std::collections::HashSet::new());
550 if !blocks.is_empty() {
551 out.push_str("\n## Context Kernel\n");
552 out.push_str(&blocks);
553 out.push('\n');
554 }
555 }
556 }
557
558 let sent = count_tokens(&out);
559 (out, sent)
560}
561
562#[cfg(test)]
563mod tests {
564 use super::*;
565
566 #[test]
567 fn rank_by_doc_freq_puts_rare_identifier_first() {
568 let keywords = vec![
572 "OCPP".to_string(),
573 "GetMaxCurrent".to_string(),
574 "Current".to_string(),
575 ];
576 let doc_freqs = std::collections::HashMap::from([
577 ("ocpp".to_string(), 120),
578 ("current".to_string(), 400),
579 ("getmaxcurrent".to_string(), 3),
580 ]);
581 let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
582 assert_eq!(ranked.first().unwrap(), "GetMaxCurrent");
583 assert_eq!(ranked.last().unwrap(), "Current");
584 }
585
586 #[test]
587 fn rank_by_doc_freq_sinks_absent_tokens_and_is_stable() {
588 let keywords = vec![
591 "absent".to_string(),
592 "alpha".to_string(),
593 "beta".to_string(),
594 ];
595 let doc_freqs =
596 std::collections::HashMap::from([("alpha".to_string(), 5), ("beta".to_string(), 5)]);
597 let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
598 assert_eq!(ranked, vec!["alpha", "beta", "absent"]);
599 }
600
601 #[test]
602 fn rank_prefers_code_identifier_over_rarer_prose_word() {
603 let keywords = vec!["measurand".to_string(), "GetMaxCurrent".to_string()];
608 let doc_freqs = std::collections::HashMap::from([
609 ("measurand".to_string(), 4),
610 ("getmaxcurrent".to_string(), 30),
611 ]);
612 let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
613 assert_eq!(ranked.first().unwrap(), "GetMaxCurrent");
614 }
615
616 #[test]
617 fn is_code_identifier_classifies_camel_snake_vs_prose_and_acronym() {
618 assert!(is_code_identifier("GetMaxCurrent"));
619 assert!(is_code_identifier("CurrentGetter"));
620 assert!(is_code_identifier("get_max_current"));
621 assert!(!is_code_identifier("Current"));
623 assert!(!is_code_identifier("OCPP"));
624 assert!(!is_code_identifier("charger"));
625 }
626
627 #[test]
628 fn extract_keywords_drops_stopwords_and_short_tokens() {
629 let kw = extract_keywords("How does the BM25Index cache work for ctx_search?", 6);
630 assert!(kw.contains(&"BM25Index".to_string()));
631 assert!(kw.contains(&"cache".to_string()));
632 assert!(kw.contains(&"ctx_search".to_string()));
633 assert!(!kw.iter().any(|k| k == "the" || k == "How" || k == "for"));
634 }
635
636 #[test]
637 fn extract_keywords_dedups_and_caps() {
638 let kw = extract_keywords("alpha alpha beta gamma delta epsilon zeta eta", 3);
639 assert_eq!(kw.len(), 3);
640 assert_eq!(kw[0], "alpha");
641 }
642
643 #[test]
644 fn exact_matches_choose_specific_identifier_not_first_broad_keyword() {
645 let keywords = extract_keywords(
646 "OCPP charger GetMaxCurrent Current.Offered measurand CurrentGetter",
647 6,
648 );
649 assert!(keywords.iter().any(|keyword| keyword == "GetMaxCurrent"));
650 assert!(keywords.iter().any(|keyword| is_code_identifier(keyword)));
651 assert!(!is_code_identifier("OCPP"));
652
653 let prose = extract_keywords("Fix semantic ranking exact matches", 6);
654 assert!(prose.iter().all(|keyword| !is_code_identifier(keyword)));
655 }
656
657 #[test]
658 fn empty_task_is_rejected() {
659 let (out, tok) = handle(" ", "/tmp", CrpMode::Off);
660 assert!(out.starts_with("ERROR"));
661 assert_eq!(tok, 0);
662 }
663
664 #[test]
665 fn handle_includes_context_kernel_section_when_available() {
666 let (output, tokens) = handle("find authentication bugs", "/tmp/nonexistent", CrpMode::Tdd);
667 assert!(tokens > 0);
670 assert!(output.contains("TASK:"));
671 }
672
673 #[test]
674 fn deferred_ranking_note_is_deterministic_and_has_no_timing() {
675 let tmp = tempfile::tempdir().unwrap();
678 let root = tmp.path().to_string_lossy();
679 let a = deferred_ranking_note(root.as_ref());
680 let b = deferred_ranking_note(root.as_ref());
681 assert_eq!(a, b, "deferred note must be byte-stable across calls");
682 assert!(
683 !a.contains("elapsed"),
684 "deferred note must not embed timing data: {a}"
685 );
686 }
687
688 #[test]
689 fn deferred_note_for_idle_index_is_optimistic_but_honest() {
690 let tmp = tempfile::tempdir().unwrap();
694 let note = deferred_ranking_note(tmp.path().to_string_lossy().as_ref());
695 assert!(
696 note.contains("warming") || note.contains("building"),
697 "note: {note}"
698 );
699 assert!(
700 note.contains("authoritative"),
701 "note must reassure that exact matches are authoritative: {note}"
702 );
703 assert!(
704 !note.contains("instant on the next call"),
705 "must not repeat the dishonest 'instant next call' promise: {note}"
706 );
707 }
708}