use std::sync::mpsc;
use std::time::Duration;
use crate::tools::CrpMode;
use lean_ctx::core::tokens::count_tokens;
const DEFAULT_SEMANTIC_BUDGET_MS: u64 = 4000;
const ONE_SHOT_SEMANTIC_BUDGET_MS: u64 = 120_000;
static RANKING_WORKERS: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(4);
fn semantic_budget() -> Duration {
let ms = std::env::var("LEAN_CTX_COMPOSE_BUDGET_MS")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.filter(|&v| v > 0)
.unwrap_or(if lean_ctx::core::conversation::is_one_shot_process() {
ONE_SHOT_SEMANTIC_BUDGET_MS
} else {
DEFAULT_SEMANTIC_BUDGET_MS
});
Duration::from_millis(ms)
}
const DEFAULT_SYMBOL_BUDGET_TOKENS: usize = 600;
fn symbol_budget_tokens() -> usize {
std::env::var("LEAN_CTX_COMPOSE_SYMBOL_TOKENS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|&v| v > 0)
.unwrap_or(DEFAULT_SYMBOL_BUDGET_TOKENS)
}
pub(crate) fn kernel_supplement_budget(project_root: &str) -> usize {
let config = lean_ctx::core::context_kernel::activation::load_config(project_root);
(symbol_budget_tokens() / 5)
.min(lean_ctx::core::context_kernel::activation::supplement_budget(&config))
}
const STOPWORDS: &[&str] = &[
"the",
"and",
"for",
"with",
"that",
"this",
"from",
"into",
"how",
"where",
"what",
"does",
"are",
"was",
"use",
"used",
"uses",
"add",
"all",
"any",
"can",
"get",
"set",
"via",
"out",
"its",
"his",
"her",
"you",
"your",
"our",
"find",
"show",
"list",
"make",
"when",
"then",
"has",
"have",
"had",
"not",
"but",
"see",
"function",
"method",
"class",
"code",
"file",
"files",
"implement",
"implementation",
];
fn extract_keywords(task: &str, max: usize) -> Vec<String> {
let mut seen = std::collections::HashSet::new();
let mut out = Vec::new();
for raw in task.split(|c: char| !(c.is_alphanumeric() || c == '_')) {
if raw.len() < 3 {
continue;
}
if STOPWORDS.contains(&raw.to_ascii_lowercase().as_str()) {
continue;
}
if seen.insert(raw.to_string()) {
out.push(raw.to_string());
if out.len() >= max {
break;
}
}
}
out
}
fn order_by_specificity(keywords: &[String], project_root: &str) -> Vec<String> {
let Some(index) = resident_index(project_root) else {
return keywords.to_vec();
};
rank_by_doc_freq(keywords, &index.doc_freqs)
}
fn rank_by_doc_freq(
keywords: &[String],
doc_freqs: &std::collections::HashMap<String, usize>,
) -> Vec<String> {
let df = |kw: &String| match doc_freqs.get(&kw.to_ascii_lowercase()) {
Some(&n) if n > 0 => n,
_ => usize::MAX,
};
let rank_key = |kw: &String| (u8::from(!is_code_identifier(kw)), df(kw));
let mut ranked = keywords.to_vec();
ranked.sort_by_key(rank_key);
ranked
}
fn is_code_identifier(kw: &str) -> bool {
if kw.contains('_') {
return true;
}
let has_lower = kw.chars().any(|c| c.is_ascii_lowercase());
let internal_upper = kw.chars().skip(1).any(|c| c.is_ascii_uppercase());
has_lower && internal_upper
}
fn resident_index(
project_root: &str,
) -> Option<std::sync::Arc<lean_ctx::core::bm25_index::BM25Index>> {
let cache = crate::tools::ctx_semantic_search::get_thread_cache()?;
lean_ctx::core::bm25_cache::get_or_background(&cache, std::path::Path::new(project_root))
}
fn ranked_files_budgeted(task: &str, project_root: &str, crp_mode: CrpMode) -> String {
ranked_files_with_slots(task, project_root, crp_mode, &RANKING_WORKERS)
}
fn ranked_files_with_slots(
task: &str,
project_root: &str,
crp_mode: CrpMode,
slots: &'static tokio::sync::Semaphore,
) -> String {
let role = lean_ctx::core::roles::active_role();
if role
.tools
.denied
.iter()
.any(|denied| denied == "ctx_search")
{
return "ERR: source search is not authorized".into();
}
let Ok(slot) = slots.try_acquire() else {
return "(local BM25 ranking busy: retry after active requests finish. Other sections are independently admitted.)".into();
};
let (tx, rx) = mpsc::channel::<String>();
let task_owned = task.to_string();
let root_owned = project_root.to_string();
let resident_cache = crate::tools::ctx_semantic_search::get_thread_cache();
lean_ctx::core::task_spine::TaskSpine::spawn_thread(move || {
if let Some(cache) = resident_cache {
crate::tools::ctx_semantic_search::set_thread_cache(cache);
}
let _slot = slot;
let ranked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
crate::tools::ctx_semantic_search::handle_for_tool(
"ctx_compose",
&task_owned,
&root_owned,
8,
crp_mode,
None,
None,
Some("bm25"),
Some(false),
Some(false),
)
}))
.unwrap_or_else(|_| {
tracing::warn!("[ctx_compose: semantic ranking panicked; omitting section]");
String::new()
});
let _ = tx.send(ranked);
});
match rx.recv_timeout(semantic_budget()) {
Ok(ranked) => ranked.trim().to_string(),
Err(_) => deferred_ranking_note().to_string(),
}
}
fn deferred_ranking_note() -> &'static str {
"(local BM25 ranking deferred: source admission did not finish within this call's budget; \
late ranking output is discarded. Other sections are independently admitted. \
LEAN_CTX_COMPOSE_BUDGET_MS raises the budget.)"
}
const PROSE_GLOBS: &str = "*.{md,mdx,markdown,rst,txt,adoc}";
fn exact_matches_section(
primary: &str,
project_root: &str,
crp_mode: CrpMode,
allow_secret: bool,
) -> String {
let search = |pattern: &str, max: usize, exclude: Option<&str>| {
crate::tools::ctx_search::handle_filtered(
pattern,
project_root,
None,
max,
crp_mode,
true,
allow_secret,
false,
exclude,
None,
)
.text
};
let found = |text: &str| {
let t = text.trim_start();
!(t.is_empty() || t.starts_with("0 matches") || t.starts_with("ERR"))
};
let mut out = String::new();
let definition = format!(
r"\b(?:fn|struct|enum|trait|type|const|static|mod|class|interface|def|func|function)\s+{}\b",
regex::escape(primary)
);
let defs = search(&definition, 5, None);
if found(&defs) {
out.push_str("Definitions:\n");
out.push_str(defs.trim());
out.push('\n');
}
let code = search(primary, 10, Some(PROSE_GLOBS));
if found(&code) {
out.push_str("Uses in code:\n");
out.push_str(code.trim());
out.push('\n');
} else {
out.push_str(search(primary, 10, None).trim());
out.push('\n');
}
out
}
const SMALL_FILE_MAX_LINES: usize = 150;
const SMALL_FILE_MAX_TOKENS: usize = 800;
const SMALL_FILES_MAX: usize = 4;
const DEFAULT_SMALL_FILES_BUDGET_TOKENS: usize = 2_000;
fn small_files_budget_tokens() -> usize {
std::env::var("LEAN_CTX_COMPOSE_INLINE_TOKENS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(DEFAULT_SMALL_FILES_BUDGET_TOKENS)
}
fn named_files(composed: &str) -> Vec<String> {
let mut seen = std::collections::HashSet::new();
let mut files = Vec::new();
for line in composed.lines() {
let candidate = if let Some((path, _)) = line.split_once(" :: ") {
path.trim()
} else if let Some(rest) = line.strip_prefix("File: ") {
rest.split(" (").next().unwrap_or("").trim()
} else {
let Some((path, rest)) = line.split_once(':') else {
continue;
};
if !rest.starts_with(|c: char| c.is_ascii_digit()) {
continue;
}
path.trim()
};
if candidate.is_empty()
|| candidate.contains(char::is_whitespace)
|| !candidate.contains('.')
|| candidate.starts_with('[')
{
continue;
}
if seen.insert(candidate.to_string()) {
files.push(candidate.to_string());
}
}
files
}
fn small_files_section(composed: &str, project_root: &str, allow_secret: bool) -> String {
let budget = small_files_budget_tokens();
if budget == 0 {
return String::new();
}
let mut used = 0usize;
let mut inlined = 0usize;
let mut section = String::new();
let mut remaining_bytes = lean_ctx::core::limits::max_read_bytes();
for rel in named_files(composed) {
if inlined >= SMALL_FILES_MAX {
break;
}
let path = std::path::Path::new(project_root).join(&rel);
if !allow_secret && lean_ctx::core::io_boundary::is_secret_like(&path).is_some() {
continue;
}
let Ok(meta) = std::fs::metadata(&path) else {
continue;
};
if !meta.is_file() || meta.len() > 16 * 1024 {
continue;
}
let Ok(content) = crate::tools::ctx_read::read_file_for_tool_rooted_budgeted(
&path.to_string_lossy(),
project_root,
"ctx_read",
&mut remaining_bytes,
) else {
continue;
};
let lines = content.lines().count();
let tokens = count_tokens(&content);
if lines == 0
|| lines > SMALL_FILE_MAX_LINES
|| tokens > SMALL_FILE_MAX_TOKENS
|| used + tokens > budget
{
continue;
}
let longest_run = content.split(|c| c != '`').map(str::len).max().unwrap_or(0);
let fence = "`".repeat(longest_run.max(2) + 1);
if section.is_empty() {
section.push_str("\n## Small files (full, verbatim)\n");
}
section.push_str(&format!("File: {rel} ({lines} lines)\n{fence}\n{content}"));
if !content.ends_with('\n') {
section.push('\n');
}
section.push_str(&fence);
section.push('\n');
used += tokens;
inlined += 1;
}
section
}
pub fn handle(task: &str, project_root: &str, crp_mode: CrpMode) -> (String, usize) {
lean_ctx::core::policy::runtime::with_project_source_view(project_root, || {
handle_in_view(task, project_root, crp_mode)
})
.unwrap_or_else(|reason| {
(
format!(
"ERROR: context withheld: source authority changed or could not be verified \
({reason})"
),
0,
)
})
}
fn handle_in_view(task: &str, project_root: &str, crp_mode: CrpMode) -> (String, usize) {
let task = task.trim();
if task.is_empty() {
return ("ERROR: task is required".to_string(), 0);
}
let protected = lean_ctx::core::policy::runtime::is_active();
let keywords = extract_keywords(task, 6);
let allow_secret = lean_ctx::core::roles::active_role().io.allow_secret_paths;
let mut out = String::new();
out.push_str(&format!("TASK: {task}\n"));
if keywords.is_empty() {
out.push_str("KEYWORDS: (none extracted — using full task for ranking)\n");
} else {
out.push_str(&format!("KEYWORDS: {}\n", keywords.join(", ")));
}
out.push_str("\n## Ranked files (local BM25)\n");
{
let _phase = lean_ctx::core::phase_timing::enter("ctx_compose: ranking");
out.push_str(&ranked_files_budgeted(task, project_root, crp_mode));
}
if protected {
out.push_str("\nStored graph and memory enrichment withheld: source authorization is unavailable. Fresh exact matches and symbols follow.");
}
out.push('\n');
let exact_phase = lean_ctx::core::phase_timing::enter("ctx_compose: exact matches");
let ranked_keywords = if protected {
keywords.clone()
} else {
order_by_specificity(&keywords, project_root)
};
if let Some(primary) = ranked_keywords
.iter()
.find(|keyword| is_code_identifier(keyword))
{
out.push_str(&format!("\n## Exact matches: '{primary}'\n"));
out.push_str(&exact_matches_section(
primary,
project_root,
crp_mode,
allow_secret,
));
}
drop(exact_phase);
let symbols_phase = lean_ctx::core::phase_timing::enter("ctx_compose: symbol bodies");
use lean_ctx::core::context_packing::CoverageItem;
let mut snippets: Vec<String> = Vec::new();
let mut items: Vec<CoverageItem> = Vec::new();
for kw in &keywords {
if let Some((rendered, toks)) =
crate::tools::ctx_symbol::best_symbol_snippet_for_task(kw, task, project_root)
{
let mut terms: std::collections::HashSet<String> =
std::collections::HashSet::from([kw.clone()]);
for other in &keywords {
if other != kw && rendered.contains(other.as_str()) {
terms.insert(other.clone());
}
}
if let Some(index) = snippets.iter().position(|snippet| snippet == &rendered) {
items[index].terms.extend(terms);
continue;
}
items.push(CoverageItem {
terms,
cost: toks.max(1),
});
snippets.push(rendered);
}
}
if !items.is_empty() {
let selected = lean_ctx::core::context_packing::greedy_max_coverage(
&items,
symbol_budget_tokens(),
|_| 1.0,
);
let mut seen = std::collections::HashSet::new();
let mut header_written = false;
for idx in selected {
let rendered = snippets[idx].trim();
if rendered.is_empty() || !seen.insert(rendered.to_string()) {
continue;
}
if !header_written {
out.push_str("\n## Top symbols (bodies)\n");
header_written = true;
}
out.push_str(rendered);
out.push('\n');
}
}
drop(symbols_phase);
{
let _phase = lean_ctx::core::phase_timing::enter("ctx_compose: small files");
out.push_str(&small_files_section(&out, project_root, allow_secret));
}
if !protected {
use lean_ctx::core::context_kernel::context_dedup::dedup_kernel_blocks;
let _phase = lean_ctx::core::phase_timing::enter("ctx_compose: kernel enrichment");
let budget = kernel_supplement_budget(project_root);
if let Some(enrichment) =
lean_ctx::core::context_kernel::bridge::kernel_enrich(task, project_root, budget)
.filter(|enrichment| !enrichment.blocks.is_empty())
{
let blocks =
dedup_kernel_blocks(&enrichment.blocks, &mut std::collections::HashSet::new());
if !blocks.is_empty() {
out.push_str("\n## Context Kernel\n");
out.push_str(&blocks);
out.push('\n');
}
}
}
let sent = count_tokens(&out);
(out, sent)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ranking_worker_keeps_original_source_policy_and_full_role() {
let _env = lean_ctx::core::data_dir::test_env_lock();
let policy = lean_ctx::core::policy::load(
"name = \"compose-worker\"\nversion = \"1.0.0\"\ndescription = \"test\"\n\
[filters]\nclassification = \"block\"\nblocked_labels = [\"CONFIDENTIAL\"]\n\
[redaction]\ncustomer = \"ACC-1234\"\n",
)
.unwrap();
let _policy = lean_ctx::core::policy::runtime::TestPolicyOverride::set(Some(policy));
let dir = tempfile::tempdir().unwrap();
std::fs::write(
dir.path().join("public.rs"),
"fn authenticate() { let customer = \"ACC-1234\"; }\n",
)
.unwrap();
std::fs::write(
dir.path().join("private.rs"),
"// CONFIDENTIAL\nfn authenticate_private() { let hidden = \"CLASSIFIED_CANARY\"; }\n",
)
.unwrap();
let root = dir.path().to_str().unwrap();
let rank = || {
lean_ctx::core::policy::runtime::with_project_source_view(root, || {
ranked_files_budgeted("authenticate", root, CrpMode::Off)
})
.unwrap()
};
let role = lean_ctx::core::roles::load_role("coder").unwrap();
lean_ctx::core::roles::with_test_active_role(role.clone(), || {
let ranked = rank();
assert!(ranked.contains("public.rs"), "{ranked}");
assert!(ranked.contains("REDACTED"), "{ranked}");
assert!(!ranked.contains("ACC-1234"));
assert!(!ranked.contains("private.rs"));
assert!(!ranked.contains("CLASSIFIED_CANARY"));
static ONE_SLOT: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(1);
let occupied = ONE_SLOT.try_acquire().unwrap();
let rank_with_limit = || {
lean_ctx::core::policy::runtime::with_project_source_view(root, || {
ranked_files_with_slots("authenticate", root, CrpMode::Off, &ONE_SLOT)
})
.unwrap()
};
let busy = rank_with_limit();
assert!(busy.contains("ranking busy"));
assert!(!busy.contains("public.rs"));
drop(occupied);
let recovered = rank_with_limit();
assert!(recovered.contains("public.rs"));
assert!(recovered.contains("REDACTED"));
assert!(!recovered.contains("CLASSIFIED_CANARY"));
});
let mut denied = role;
denied.tools.denied.push("ctx_search".into());
lean_ctx::core::roles::with_test_active_role(denied, || {
let ranked = rank();
assert!(ranked.starts_with("ERR:"), "{ranked}");
assert!(!ranked.contains("public.rs"));
});
}
#[test]
fn rank_by_doc_freq_puts_rare_identifier_first() {
let keywords = vec![
"OCPP".to_string(),
"GetMaxCurrent".to_string(),
"Current".to_string(),
];
let doc_freqs = std::collections::HashMap::from([
("ocpp".to_string(), 120),
("current".to_string(), 400),
("getmaxcurrent".to_string(), 3),
]);
let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
assert_eq!(ranked.first().unwrap(), "GetMaxCurrent");
assert_eq!(ranked.last().unwrap(), "Current");
}
#[test]
fn named_small_files_are_inlined_verbatim_within_budget() {
let _env = lean_ctx::core::data_dir::test_env_lock();
let _policy = lean_ctx::core::policy::runtime::TestPolicyOverride::set(None);
let dir = tempfile::tempdir().unwrap();
std::fs::write(
dir.path().join("pyproject.toml"),
"[project]\nname = \"x\"\n",
)
.unwrap();
std::fs::create_dir(dir.path().join("pkg")).unwrap();
let cli = "import argparse\n\n\ndef main():\n return \"```\"\n";
std::fs::write(dir.path().join("pkg/cli.py"), cli).unwrap();
let big: String = (0..400).fold(String::new(), |mut acc, i| {
use std::fmt::Write as _;
let _ = writeln!(acc, "value_{i} = {i}");
acc
});
std::fs::write(dir.path().join("pkg/big.py"), &big).unwrap();
std::fs::write(dir.path().join(".env"), "TOKEN=abc\n").unwrap();
let composed = "## Ranked files (local BM25)\n\
pkg/big.py :: value_1 [Variable] (L2-2)\n\
pkg/cli.py :: main [Function] (L4-5)\n\
## Exact matches: 'main'\n\
.env:1 TOKEN=abc\n\
pkg/cli.py:4 def main():\n";
let root = dir.path().to_str().unwrap();
let section = lean_ctx::core::policy::runtime::with_project_source_view(root, || {
small_files_section(composed, root, false)
})
.unwrap();
assert!(
section.contains(&format!("File: pkg/cli.py (5 lines)\n````\n{cli}````\n")),
"{section}"
);
assert!(!section.contains("value_399"), "{section}");
assert!(!section.contains("TOKEN=abc"), "{section}");
assert_eq!(section.matches("File: ").count(), 1, "{section}");
}
#[test]
fn rank_by_doc_freq_sinks_absent_tokens_and_is_stable() {
let keywords = vec![
"absent".to_string(),
"alpha".to_string(),
"beta".to_string(),
];
let doc_freqs =
std::collections::HashMap::from([("alpha".to_string(), 5), ("beta".to_string(), 5)]);
let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
assert_eq!(ranked, vec!["alpha", "beta", "absent"]);
}
#[test]
fn rank_prefers_code_identifier_over_rarer_prose_word() {
let keywords = vec!["measurand".to_string(), "GetMaxCurrent".to_string()];
let doc_freqs = std::collections::HashMap::from([
("measurand".to_string(), 4),
("getmaxcurrent".to_string(), 30),
]);
let ranked = rank_by_doc_freq(&keywords, &doc_freqs);
assert_eq!(ranked.first().unwrap(), "GetMaxCurrent");
}
#[test]
fn is_code_identifier_classifies_camel_snake_vs_prose_and_acronym() {
assert!(is_code_identifier("GetMaxCurrent"));
assert!(is_code_identifier("CurrentGetter"));
assert!(is_code_identifier("get_max_current"));
assert!(!is_code_identifier("Current"));
assert!(!is_code_identifier("OCPP"));
assert!(!is_code_identifier("charger"));
}
#[test]
fn extract_keywords_drops_stopwords_and_short_tokens() {
let kw = extract_keywords("How does the BM25Index cache work for ctx_search?", 6);
assert!(kw.contains(&"BM25Index".to_string()));
assert!(kw.contains(&"cache".to_string()));
assert!(kw.contains(&"ctx_search".to_string()));
assert!(!kw.iter().any(|k| k == "the" || k == "How" || k == "for"));
}
#[test]
fn extract_keywords_dedups_and_caps() {
let kw = extract_keywords("alpha alpha beta gamma delta epsilon zeta eta", 3);
assert_eq!(kw.len(), 3);
assert_eq!(kw[0], "alpha");
}
#[test]
fn exact_matches_choose_specific_identifier_not_first_broad_keyword() {
let keywords = extract_keywords(
"OCPP charger GetMaxCurrent Current.Offered measurand CurrentGetter",
6,
);
assert!(keywords.iter().any(|keyword| keyword == "GetMaxCurrent"));
assert!(keywords.iter().any(|keyword| is_code_identifier(keyword)));
assert!(!is_code_identifier("OCPP"));
let prose = extract_keywords("Fix semantic ranking exact matches", 6);
assert!(prose.iter().all(|keyword| !is_code_identifier(keyword)));
}
#[test]
fn exact_matches_rank_definition_and_code_above_prose() {
let _env = lean_ctx::core::data_dir::test_env_lock();
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::write(root.join("AGENTS.md"), "use widget_factory\n".repeat(30)).unwrap();
std::fs::create_dir_all(root.join("src")).unwrap();
std::fs::write(
root.join("src/lib.rs"),
"pub fn widget_factory() -> u8 { 1 }\n",
)
.unwrap();
std::fs::write(
root.join("src/main.rs"),
"fn main() { lib::widget_factory(); }\n",
)
.unwrap();
let root_str = root.to_str().unwrap();
let out = lean_ctx::core::policy::runtime::with_project_source_view(root_str, || {
exact_matches_section("widget_factory", root_str, CrpMode::Off, false)
})
.expect("source view");
let defs = out.find("Definitions:").expect("definition section");
let uses = out.find("Uses in code:").expect("code section");
assert!(defs < uses, "{out}");
assert!(out.contains("lib.rs:1"), "definition found: {out}");
assert!(out.contains("main.rs:1"), "use found: {out}");
assert!(
!out.contains("AGENTS.md"),
"prose only when code has none: {out}"
);
let prose_only =
lean_ctx::core::policy::runtime::with_project_source_view(root_str, || {
exact_matches_section("use", root_str, CrpMode::Off, false)
})
.expect("source view");
assert!(prose_only.contains("AGENTS.md"), "{prose_only}");
}
#[test]
fn empty_task_is_rejected() {
let (out, tok) = handle(" ", "/tmp", CrpMode::Off);
assert!(out.starts_with("ERROR"));
assert_eq!(tok, 0);
}
#[test]
fn handle_withholds_context_when_project_authority_cannot_be_verified() {
let root = tempfile::tempdir().unwrap();
let missing = root.path().join("missing-project");
let (output, tokens) = handle(
"find authentication bugs",
missing.to_str().unwrap(),
CrpMode::Tdd,
);
assert_eq!(tokens, 0);
assert!(output.starts_with("ERROR"));
assert!(output.contains("source authority"));
assert!(!output.contains("TASK:"));
}
#[test]
fn deferred_ranking_note_is_deterministic_and_has_no_timing() {
let a = deferred_ranking_note();
let b = deferred_ranking_note();
assert_eq!(a, b, "deferred note must be byte-stable across calls");
assert!(
!a.contains("elapsed"),
"deferred note must not embed timing data: {a}"
);
}
#[test]
fn deferred_note_does_not_promise_cached_or_late_delivery() {
let note = deferred_ranking_note();
assert!(note.contains("late ranking output is discarded"));
assert!(note.contains("independently admitted"));
assert!(!note.contains("warming"));
assert!(!note.contains("next call"));
}
}