1use std::collections::HashMap;
2use std::sync::Mutex;
3
4use crate::core::cache::SessionCache;
5use crate::core::context_ledger::PressureAction;
6use crate::core::mode_predictor::{FileSignature, ModePredictor};
7
8static SOURCE_COUNTS: Mutex<Option<HashMap<&'static str, u64>>> = Mutex::new(None);
11
12fn count_source(source: &'static str) {
13 if let Ok(mut guard) = SOURCE_COUNTS.lock() {
14 *guard
15 .get_or_insert_with(HashMap::new)
16 .entry(source)
17 .or_insert(0) += 1;
18 }
19}
20
21pub fn source_counts() -> Vec<(&'static str, u64)> {
23 let Ok(guard) = SOURCE_COUNTS.lock() else {
24 return Vec::new();
25 };
26 let mut items: Vec<(&'static str, u64)> = guard
27 .as_ref()
28 .map(|m| m.iter().map(|(k, v)| (*k, *v)).collect())
29 .unwrap_or_default();
30 items.sort_by_key(|(_, n)| std::cmp::Reverse(*n));
31 items
32}
33
34fn sources_path() -> Option<std::path::PathBuf> {
35 crate::core::data_dir::lean_ctx_data_dir()
36 .ok()
37 .map(|d| d.join("auto_mode_sources.json"))
38}
39
40pub fn flush_sources() {
45 let drained: Vec<(String, u64)> = {
46 let Ok(mut guard) = SOURCE_COUNTS.lock() else {
47 return;
48 };
49 match guard.take() {
50 Some(m) if !m.is_empty() => m.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
51 _ => return,
52 }
53 };
54 let Some(path) = sources_path() else {
55 return;
56 };
57 let mut on_disk: HashMap<String, u64> = std::fs::read_to_string(&path)
58 .ok()
59 .and_then(|s| serde_json::from_str(&s).ok())
60 .unwrap_or_default();
61 for (k, v) in drained {
62 *on_disk.entry(k).or_insert(0) += v;
63 }
64 let Ok(json) = serde_json::to_string_pretty(&on_disk) else {
65 return;
66 };
67 let tmp = path.with_extension("json.tmp");
68 if std::fs::write(&tmp, json).is_ok() {
69 let _ = std::fs::rename(&tmp, &path);
70 }
71}
72
73pub fn persisted_source_counts() -> Vec<(String, u64)> {
76 let Some(path) = sources_path() else {
77 return Vec::new();
78 };
79 let map: HashMap<String, u64> = std::fs::read_to_string(&path)
80 .ok()
81 .and_then(|s| serde_json::from_str(&s).ok())
82 .unwrap_or_default();
83 let mut items: Vec<(String, u64)> = map.into_iter().collect();
84 items.sort_by_key(|(_, n)| std::cmp::Reverse(*n));
85 items
86}
87
88pub struct AutoModeContext<'a> {
89 pub path: &'a str,
90 pub token_count: usize,
91 pub task: Option<&'a str>,
92 pub cache: Option<&'a SessionCache>,
93}
94
95pub struct ResolvedMode {
96 pub mode: String,
97 pub source: &'static str,
98}
99
100pub fn resolve(ctx: &AutoModeContext) -> ResolvedMode {
103 if crate::core::edit_quality::take_pending_escalation(ctx.path) {
106 return resolved("full", "edit_fail_escalation");
107 }
108
109 let r = resolve_inner(ctx);
110
111 if r.mode != "full" && crate::core::edit_quality::is_risky_mode(ctx.path, &r.mode) {
114 return resolved("full", "edit_quality_penalty");
115 }
116 r
117}
118
119fn resolve_inner(ctx: &AutoModeContext) -> ResolvedMode {
120 if crate::tools::ctx_read::is_instruction_file(ctx.path) {
121 return resolved("full", "instruction_file");
122 }
123
124 if crate::core::binary_detect::is_binary_file(ctx.path) {
125 return resolved("full", "binary");
126 }
127
128 if let Some(cache) = ctx.cache
129 && let Some(cached) = cache.get(ctx.path)
130 {
131 if !file_unchanged(ctx.path, cached) {
132 return resolved("diff", "cache_changed");
133 }
134 if cache.is_full_delivered(ctx.path) {
146 return resolved("full", "cache_hit");
147 }
148 }
149
150 if ctx.token_count <= 200 {
151 return resolved("full", "small_file");
152 }
153
154 let ext = std::path::Path::new(ctx.path)
155 .extension()
156 .and_then(|e| e.to_str())
157 .unwrap_or("");
158
159 if is_config_or_data(ext, ctx.path) {
160 return resolved("full", "config_data");
161 }
162
163 if let Ok(bt) = crate::core::bounce_tracker::global().lock()
164 && bt.should_force_full(ctx.path)
165 {
166 return resolved("full", "bounce_tracker");
167 }
168
169 if crate::core::path_mode_memory::should_force_full(ctx.path) {
173 return resolved("full", "path_bounce_memory");
174 }
175
176 if crate::core::diagnostics_store::has_error(ctx.path) {
179 return resolved("full", "active_diagnostic");
180 }
181
182 if task_names_file(ctx.task, ctx.path) {
187 return resolved("full", "task_suspect_file");
188 }
189
190 if let Some(mode) = intent_recommended_mode(ctx.task) {
191 return resolved(&mode, "intent");
192 }
193
194 let sig = FileSignature::from_path(ctx.path, ctx.token_count);
195 let predictor = ModePredictor::new();
196 let mut predicted = predictor
197 .predict_best_mode(&sig)
198 .unwrap_or_else(|| "full".to_string());
199 if predicted == "auto" {
200 predicted = "full".to_string();
201 }
202
203 if predicted != "full"
204 && let Some(bandit_override) = bandit_explore(ctx.path, ctx.token_count)
205 {
206 predicted = bandit_override;
207 }
208
209 if predicted != "full"
213 && let Some((access_count, avg_ratio)) = crate::core::heatmap::entry_stats(ctx.path)
214 && access_count >= 5
215 && avg_ratio < 0.30
216 {
217 let conservative = match predicted.as_str() {
218 "signatures" | "aggressive" | "entropy" => "map".to_string(),
219 "map" if ctx.token_count <= 6000 => "full".to_string(),
220 other => other.to_string(),
221 };
222 if conservative != predicted {
223 return resolved(&conservative, "heatmap_conservative");
224 }
225 }
226
227 let policy = crate::core::adaptive_mode_policy::AdaptiveModePolicyStore::load();
228 let chosen = policy.choose_auto_mode(ctx.task, &predicted);
229
230 if ctx.token_count > 2000 {
231 if (predicted == "map" || predicted == "signatures")
232 && chosen != "map"
233 && chosen != "signatures"
234 {
235 return resolved(&predicted, "predictor_guard");
236 }
237 if chosen == "full" && predicted != "full" {
238 return resolved(&predicted, "predictor_override");
239 }
240 }
241
242 if chosen != predicted {
243 return resolved(&chosen, "adaptive_policy");
244 }
245
246 if predicted != "full" {
247 return resolved(&predicted, "predictor");
248 }
249
250 let structure_first = crate::core::config::Config::load().structure_first_effective();
257 let heuristic = heuristic_mode(ext, ctx.token_count, structure_first);
258 let source = if structure_first && heuristic == "map" && ctx.token_count <= 6000 {
259 "structure_first"
260 } else {
261 "heuristic"
262 };
263 resolved(&heuristic, source)
264}
265
266pub fn pressure_downgrade(requested_mode: &str, action: &PressureAction) -> Option<String> {
269 match action {
270 PressureAction::SuggestCompression => match requested_mode {
271 "auto" | "full" => Some("map".to_string()),
272 _ => None,
273 },
274 PressureAction::ForceCompression => match requested_mode {
275 "full" => Some("map".to_string()),
276 "auto" | "map" => Some("signatures".to_string()),
277 _ => None,
278 },
279 PressureAction::EvictLeastRelevant => match requested_mode {
280 "full" => Some("map".to_string()),
281 "auto" | "map" => Some("signatures".to_string()),
282 "signatures" => Some("reference".to_string()),
283 _ => None,
284 },
285 PressureAction::NoAction => None,
286 }
287}
288
289fn task_names_file(task: Option<&str>, path: &str) -> bool {
296 let Some(task) = task else {
297 return false;
298 };
299 let basename = std::path::Path::new(path)
300 .file_name()
301 .and_then(|n| n.to_str())
302 .unwrap_or("");
303 if basename.len() < 4 || !basename.contains('.') {
304 return false;
305 }
306 task.to_ascii_lowercase()
307 .contains(&basename.to_ascii_lowercase())
308}
309
310fn intent_recommended_mode(task: Option<&str>) -> Option<String> {
311 let task_desc = task?;
312 let classification = crate::core::intent_engine::classify(task_desc);
313 if classification.confidence < 0.4 {
314 return None;
315 }
316 let route = crate::core::intent_engine::route_intent(task_desc, &classification);
317 let mode =
318 crate::core::intent_router::read_mode_for_tier(route.model_tier, classification.task_type);
319 if mode == "auto" {
320 return None;
321 }
322 Some(mode)
323}
324
325fn bandit_explore(file_path: &str, token_count: usize) -> Option<String> {
326 let project_root =
327 crate::core::session::SessionState::load_latest().and_then(|s| s.project_root)?;
328 let ext = std::path::Path::new(file_path)
329 .extension()
330 .and_then(|e| e.to_str())
331 .unwrap_or("");
332 let bucket = match token_count {
333 0..=2000 => "sm",
334 2001..=10000 => "md",
335 10001..=50000 => "lg",
336 _ => "xl",
337 };
338 let bandit_key = format!("{ext}_{bucket}");
339 let mut store = crate::core::bandit::BanditStore::load(&project_root);
340 let bandit = store.get_or_create(&bandit_key);
341 let arm = bandit.select_arm();
342 if arm.budget_ratio < 0.25 && token_count > 2000 {
343 Some("aggressive".to_string())
344 } else {
345 None
346 }
347}
348
349fn heuristic_mode(ext: &str, token_count: usize, structure_first: bool) -> String {
350 if token_count > 8000 {
351 if is_code(ext) {
352 return "map".to_string();
353 }
354 return "aggressive".to_string();
355 }
356 if token_count > 6000 && is_code(ext) {
361 return "map".to_string();
362 }
363 if structure_first && token_count > 500 && is_code(ext) {
370 return "map".to_string();
371 }
372 "full".to_string()
373}
374
375fn file_unchanged(path: &str, cached: &crate::core::cache::CacheEntry) -> bool {
384 let Some(stored_mtime) = cached.stored_mtime else {
385 return false;
386 };
387 let Ok(meta) = std::fs::metadata(path) else {
388 return false;
389 };
390 let Ok(current_mtime) = meta.modified() else {
391 return false;
392 };
393 current_mtime == stored_mtime
394}
395
396fn is_code(ext: &str) -> bool {
397 matches!(
398 ext,
399 "rs" | "ts"
400 | "tsx"
401 | "js"
402 | "jsx"
403 | "py"
404 | "go"
405 | "java"
406 | "c"
407 | "cpp"
408 | "cc"
409 | "h"
410 | "hpp"
411 | "rb"
412 | "cs"
413 | "kt"
414 | "swift"
415 | "php"
416 | "zig"
417 | "ex"
418 | "exs"
419 | "scala"
420 | "sc"
421 | "dart"
422 | "sh"
423 | "bash"
424 | "svelte"
425 | "vue"
426 )
427}
428
429fn is_config_or_data(ext: &str, path: &str) -> bool {
430 if matches!(ext, "xml" | "ini" | "cfg" | "env") {
431 return true;
432 }
433 let name = std::path::Path::new(path)
434 .file_name()
435 .and_then(|n| n.to_str())
436 .unwrap_or("");
437 matches!(
438 name,
439 "Cargo.toml"
440 | "package.json"
441 | "tsconfig.json"
442 | "Makefile"
443 | "Dockerfile"
444 | "docker-compose.yml"
445 | ".gitignore"
446 | ".env"
447 | "pyproject.toml"
448 | "go.mod"
449 | "build.gradle"
450 | "pom.xml"
451 )
452}
453
454fn resolved(mode: &str, source: &'static str) -> ResolvedMode {
455 count_source(source);
456 ResolvedMode {
457 mode: mode.to_string(),
458 source,
459 }
460}
461
462#[cfg(test)]
463mod tests {
464 use super::*;
465
466 #[test]
467 fn pressure_suggest_full_to_map() {
468 assert_eq!(
469 pressure_downgrade("full", &PressureAction::SuggestCompression),
470 Some("map".to_string())
471 );
472 }
473
474 #[test]
475 fn pressure_suggest_auto_to_map() {
476 assert_eq!(
477 pressure_downgrade("auto", &PressureAction::SuggestCompression),
478 Some("map".to_string())
479 );
480 }
481
482 #[test]
483 fn pressure_suggest_does_not_touch_signatures() {
484 assert!(pressure_downgrade("signatures", &PressureAction::SuggestCompression).is_none());
485 }
486
487 #[test]
488 fn pressure_force_full_to_map() {
489 assert_eq!(
490 pressure_downgrade("full", &PressureAction::ForceCompression),
491 Some("map".to_string())
492 );
493 }
494
495 #[test]
496 fn pressure_force_map_to_signatures() {
497 assert_eq!(
498 pressure_downgrade("map", &PressureAction::ForceCompression),
499 Some("signatures".to_string())
500 );
501 }
502
503 #[test]
504 fn pressure_evict_signatures_to_reference() {
505 assert_eq!(
506 pressure_downgrade("signatures", &PressureAction::EvictLeastRelevant),
507 Some("reference".to_string())
508 );
509 }
510
511 #[test]
512 fn pressure_noaction_returns_none() {
513 assert!(pressure_downgrade("full", &PressureAction::NoAction).is_none());
514 }
515
516 #[test]
517 fn flush_sources_merges_additively_into_disk_file() {
518 let _lock = crate::core::data_dir::test_env_lock();
519 let dir = std::env::temp_dir().join(format!("lctx-amr-flush-{}", std::process::id()));
520 let _ = std::fs::create_dir_all(&dir);
521 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
522 let _ = std::fs::remove_file(dir.join("auto_mode_sources.json"));
523
524 count_source("test_flush_alpha");
527 count_source("test_flush_alpha");
528 count_source("test_flush_beta");
529 flush_sources();
530
531 count_source("test_flush_alpha");
532 flush_sources();
533
534 let persisted = persisted_source_counts();
535 let get = |k: &str| {
536 persisted
537 .iter()
538 .find(|(s, _)| s == k)
539 .map_or(0, |(_, n)| *n)
540 };
541 assert_eq!(
542 get("test_flush_alpha"),
543 3,
544 "two flushes must merge additively"
545 );
546 assert_eq!(get("test_flush_beta"), 1);
547
548 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
549 let _ = std::fs::remove_dir_all(&dir);
550 }
551
552 #[test]
553 fn small_file_always_full() {
554 let ctx = AutoModeContext {
555 path: "test.rs",
556 token_count: 100,
557 task: None,
558 cache: None,
559 };
560 let result = resolve(&ctx);
561 assert_eq!(result.mode, "full");
562 assert_eq!(result.source, "small_file");
563 }
564
565 #[test]
566 fn config_file_returns_full() {
567 let ctx = AutoModeContext {
568 path: "config.ini",
569 token_count: 500,
570 task: None,
571 cache: None,
572 };
573 let result = resolve(&ctx);
574 assert_eq!(result.mode, "full");
575 assert_eq!(result.source, "config_data");
576 }
577
578 #[test]
579 fn cached_compressed_only_file_does_not_escalate_to_full() {
580 let dir = tempfile::tempdir().unwrap();
586 let file = dir.path().join("medium.rs");
587 let body = "fn placeholder() { let _ = 1; }\n".repeat(400);
588 std::fs::write(&file, &body).unwrap();
589 let path = file.to_str().unwrap();
590
591 let mut cache = SessionCache::new();
592 cache.store(path, &body);
593 let ctx = AutoModeContext {
596 path,
597 token_count: 3000,
598 task: None,
599 cache: Some(&cache),
600 };
601 let result = resolve(&ctx);
602 assert_ne!(
603 result.mode, "full",
604 "compressed-only cached file must not escalate to full on re-read"
605 );
606 assert_ne!(result.source, "cache_hit");
607 }
608
609 #[test]
610 fn cached_full_delivered_file_short_circuits_to_stub() {
611 let dir = tempfile::tempdir().unwrap();
614 let file = dir.path().join("medium.rs");
615 let body = "fn placeholder() { let _ = 1; }\n".repeat(400);
616 std::fs::write(&file, &body).unwrap();
617 let path = file.to_str().unwrap();
618
619 let mut cache = SessionCache::new();
620 cache.store(path, &body);
621 cache.mark_full_delivered(path);
622
623 let ctx = AutoModeContext {
624 path,
625 token_count: 3000,
626 task: None,
627 cache: Some(&cache),
628 };
629 let result = resolve(&ctx);
630 assert_eq!(result.mode, "full");
631 assert_eq!(result.source, "cache_hit");
632 }
633
634 #[test]
635 fn intent_explore_returns_map() {
636 let ctx = AutoModeContext {
637 path: "large.rs",
638 token_count: 5000,
639 task: Some("how does the cache work?"),
640 cache: None,
641 };
642 let result = resolve(&ctx);
643 assert_eq!(result.mode, "map");
644 assert_eq!(result.source, "intent");
645 }
646
647 #[test]
648 fn task_names_file_matches_explicit_filename() {
649 assert!(task_names_file(
650 Some("fix the version sort in versioncmp.c"),
651 "src/versioncmp.c"
652 ));
653 assert!(task_names_file(
654 Some("why does graph.ts loop?"),
655 "web/src/graph.ts"
656 ));
657 }
658
659 #[test]
660 fn task_names_file_ignores_bare_stems_and_trivia() {
661 assert!(!task_names_file(
663 Some("improve the parser"),
664 "src/parser.rs"
665 ));
666 assert!(!task_names_file(None, "src/parser.rs"));
667 assert!(!task_names_file(Some("touch a.c"), "a.c"));
669 assert!(!task_names_file(Some("look at Makefile"), "Makefile"));
670 }
671
672 #[test]
673 fn task_suspect_file_overrides_intent() {
674 let _lock = crate::core::data_dir::test_env_lock();
675 let dir = std::env::temp_dir().join(format!("lctx-amr-suspect-{}", std::process::id()));
676 let _ = std::fs::create_dir_all(&dir);
677 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
678
679 let ctx = AutoModeContext {
683 path: "large.rs",
684 token_count: 5000,
685 task: Some("how does large.rs build the cache?"),
686 cache: None,
687 };
688 let result = resolve(&ctx);
689 assert_eq!(result.mode, "full");
690 assert_eq!(result.source, "task_suspect_file");
691
692 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
693 let _ = std::fs::remove_dir_all(&dir);
694 }
695
696 #[test]
697 fn heuristic_full_for_medium_code_by_default() {
698 assert_eq!(heuristic_mode("rs", 1500, false), "full");
701 assert_eq!(heuristic_mode("ts", 1000, false), "full");
702 }
703
704 #[test]
705 fn heuristic_structure_first_maps_medium_code() {
706 assert_eq!(heuristic_mode("rs", 1500, true), "map");
708 assert_eq!(heuristic_mode("c", 800, true), "map");
709 }
710
711 #[test]
712 fn heuristic_structure_first_keeps_tiny_and_prose_full() {
713 assert_eq!(heuristic_mode("rs", 400, true), "full");
715 assert_eq!(heuristic_mode("md", 4000, true), "full");
717 assert_eq!(heuristic_mode("txt", 1000, true), "full");
718 }
719
720 #[test]
721 fn heuristic_large_code_maps_regardless() {
722 assert_eq!(heuristic_mode("rs", 9000, false), "map");
723 assert_eq!(heuristic_mode("rs", 9000, true), "map");
724 }
725
726 #[test]
732 fn structure_first_resolve_bugfix_cold_read() {
733 let _lock = crate::core::data_dir::test_env_lock();
734 let dir = std::env::temp_dir().join(format!("lctx-amr-sf-{}", std::process::id()));
735 let _ = std::fs::create_dir_all(&dir);
736 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
737 crate::test_env::set_var("LEAN_CTX_STRUCTURE_FIRST", "1");
738
739 let suspect = AutoModeContext {
740 path: "src/versioncmp.c",
741 token_count: 1500,
742 task: None,
743 cache: None,
744 };
745 let result = resolve(&suspect);
746 assert_eq!(result.mode, "map");
747 assert_eq!(result.source, "structure_first");
748
749 let tiny = AutoModeContext {
750 path: "src/util.c",
751 token_count: 120,
752 task: None,
753 cache: None,
754 };
755 assert_eq!(resolve(&tiny).mode, "full");
756
757 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
758 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
759 let _ = std::fs::remove_dir_all(&dir);
760 }
761
762 #[test]
763 fn structure_first_off_keeps_medium_code_full() {
764 let _lock = crate::core::data_dir::test_env_lock();
765 let dir = std::env::temp_dir().join(format!("lctx-amr-sfoff-{}", std::process::id()));
766 let _ = std::fs::create_dir_all(&dir);
767 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
768 crate::test_env::set_var("LEAN_CTX_STRUCTURE_FIRST", "0");
769
770 let ctx = AutoModeContext {
771 path: "src/versioncmp.c",
772 token_count: 1500,
773 task: None,
774 cache: None,
775 };
776 let result = resolve(&ctx);
777 assert_eq!(result.mode, "full");
778 assert_eq!(result.source, "heuristic");
779
780 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
781 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
782 let _ = std::fs::remove_dir_all(&dir);
783 }
784}