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 crate::core::diagnostics_store::has_error(ctx.path) {
166 return resolved("full", "active_diagnostic");
167 }
168
169 if task_names_file(ctx.task, ctx.path) {
174 return resolved("full", "task_suspect_file");
175 }
176
177 if let Some(mode) = intent_recommended_mode(ctx.task) {
178 return resolved(&mode, "intent");
179 }
180
181 if crate::core::config::Config::load().auto_mode_learning_effective()
187 && let Some(r) = resolve_adaptive(ctx)
188 {
189 return r;
190 }
191
192 let structure_first = crate::core::config::Config::load().structure_first_effective();
198 let heuristic = heuristic_mode(ext, ctx.token_count, structure_first);
199 let source = if structure_first && heuristic == "map" && ctx.token_count <= 6000 {
200 "structure_first"
201 } else {
202 "heuristic"
203 };
204 resolved(&heuristic, source)
205}
206
207fn resolve_adaptive(ctx: &AutoModeContext) -> Option<ResolvedMode> {
213 if let Ok(bt) = crate::core::bounce_tracker::global().lock()
214 && bt.should_force_full(ctx.path)
215 {
216 return Some(resolved("full", "bounce_tracker"));
217 }
218
219 if crate::core::path_mode_memory::should_force_full(ctx.path) {
223 return Some(resolved("full", "path_bounce_memory"));
224 }
225
226 let sig = FileSignature::from_path(ctx.path, ctx.token_count);
227 let predictor = ModePredictor::new();
228 let mut predicted = predictor
229 .predict_best_mode(&sig)
230 .unwrap_or_else(|| "full".to_string());
231 if predicted == "auto" {
232 predicted = "full".to_string();
233 }
234
235 if predicted != "full"
236 && let Some(bandit_override) = bandit_explore(ctx.path, ctx.token_count)
237 {
238 predicted = bandit_override;
239 }
240
241 if predicted != "full"
245 && let Some((access_count, avg_ratio)) = crate::core::heatmap::entry_stats(ctx.path)
246 && access_count >= 5
247 && avg_ratio < 0.30
248 {
249 let conservative = match predicted.as_str() {
250 "signatures" | "aggressive" | "entropy" => "map".to_string(),
251 "map" if ctx.token_count <= 6000 => "full".to_string(),
252 other => other.to_string(),
253 };
254 if conservative != predicted {
255 return Some(resolved(&conservative, "heatmap_conservative"));
256 }
257 }
258
259 let policy = crate::core::adaptive_mode_policy::AdaptiveModePolicyStore::load();
260 let chosen = policy.choose_auto_mode(ctx.task, &predicted);
261
262 if ctx.token_count > 2000 {
263 if (predicted == "map" || predicted == "signatures")
264 && chosen != "map"
265 && chosen != "signatures"
266 {
267 return Some(resolved(&predicted, "predictor_guard"));
268 }
269 if chosen == "full" && predicted != "full" {
270 return Some(resolved(&predicted, "predictor_override"));
271 }
272 }
273
274 if chosen != predicted {
275 return Some(resolved(&chosen, "adaptive_policy"));
276 }
277
278 if predicted != "full" {
279 return Some(resolved(&predicted, "predictor"));
280 }
281
282 None
283}
284
285pub fn pressure_downgrade(requested_mode: &str, action: &PressureAction) -> Option<String> {
288 match action {
289 PressureAction::SuggestCompression => match requested_mode {
290 "auto" | "full" => Some("map".to_string()),
291 _ => None,
292 },
293 PressureAction::ForceCompression => match requested_mode {
294 "full" => Some("map".to_string()),
295 "auto" | "map" => Some("signatures".to_string()),
296 _ => None,
297 },
298 PressureAction::EvictLeastRelevant => match requested_mode {
299 "full" => Some("map".to_string()),
300 "auto" | "map" => Some("signatures".to_string()),
301 "signatures" => Some("reference".to_string()),
302 _ => None,
303 },
304 PressureAction::NoAction => None,
305 }
306}
307
308fn task_names_file(task: Option<&str>, path: &str) -> bool {
315 let Some(task) = task else {
316 return false;
317 };
318 let basename = std::path::Path::new(path)
319 .file_name()
320 .and_then(|n| n.to_str())
321 .unwrap_or("");
322 if basename.len() < 4 || !basename.contains('.') {
323 return false;
324 }
325 task.to_ascii_lowercase()
326 .contains(&basename.to_ascii_lowercase())
327}
328
329fn intent_recommended_mode(task: Option<&str>) -> Option<String> {
330 let task_desc = task?;
331 let classification = crate::core::intent_engine::classify(task_desc);
332 if classification.confidence < 0.4 {
333 return None;
334 }
335 let route = crate::core::intent_engine::route_intent(task_desc, &classification);
336 let mode =
337 crate::core::intent_router::read_mode_for_tier(route.model_tier, classification.task_type);
338 if mode == "auto" {
339 return None;
340 }
341 Some(mode)
342}
343
344fn bandit_explore(file_path: &str, token_count: usize) -> Option<String> {
345 let project_root =
346 crate::core::session::SessionState::load_latest().and_then(|s| s.project_root)?;
347 let ext = std::path::Path::new(file_path)
348 .extension()
349 .and_then(|e| e.to_str())
350 .unwrap_or("");
351 let bucket = match token_count {
352 0..=2000 => "sm",
353 2001..=10000 => "md",
354 10001..=50000 => "lg",
355 _ => "xl",
356 };
357 let bandit_key = format!("{ext}_{bucket}");
358 let mut store = crate::core::bandit::BanditStore::load(&project_root);
359 let bandit = store.get_or_create(&bandit_key);
360 let arm = bandit.select_arm();
361 if arm.budget_ratio < 0.25 && token_count > 2000 {
362 Some("aggressive".to_string())
363 } else {
364 None
365 }
366}
367
368fn heuristic_mode(ext: &str, token_count: usize, structure_first: bool) -> String {
369 if token_count > 8000 {
370 if is_code(ext) {
371 return "map".to_string();
372 }
373 return "aggressive".to_string();
374 }
375 if token_count > 6000 && is_code(ext) {
380 return "map".to_string();
381 }
382 if structure_first && token_count > 500 && is_code(ext) {
389 return "map".to_string();
390 }
391 "full".to_string()
392}
393
394fn file_unchanged(path: &str, cached: &crate::core::cache::CacheEntry) -> bool {
403 let Some(stored_mtime) = cached.stored_mtime else {
404 return false;
405 };
406 let Ok(meta) = std::fs::metadata(path) else {
407 return false;
408 };
409 let Ok(current_mtime) = meta.modified() else {
410 return false;
411 };
412 current_mtime == stored_mtime
413}
414
415fn is_code(ext: &str) -> bool {
416 matches!(
417 ext,
418 "rs" | "ts"
419 | "tsx"
420 | "js"
421 | "jsx"
422 | "py"
423 | "go"
424 | "java"
425 | "c"
426 | "cpp"
427 | "cc"
428 | "h"
429 | "hpp"
430 | "rb"
431 | "cs"
432 | "kt"
433 | "swift"
434 | "php"
435 | "zig"
436 | "ex"
437 | "exs"
438 | "scala"
439 | "sc"
440 | "dart"
441 | "sh"
442 | "bash"
443 | "svelte"
444 | "vue"
445 )
446}
447
448fn is_config_or_data(ext: &str, path: &str) -> bool {
449 if matches!(ext, "xml" | "ini" | "cfg" | "env") {
450 return true;
451 }
452 let name = std::path::Path::new(path)
453 .file_name()
454 .and_then(|n| n.to_str())
455 .unwrap_or("");
456 matches!(
457 name,
458 "Cargo.toml"
459 | "package.json"
460 | "tsconfig.json"
461 | "Makefile"
462 | "Dockerfile"
463 | "docker-compose.yml"
464 | ".gitignore"
465 | ".env"
466 | "pyproject.toml"
467 | "go.mod"
468 | "build.gradle"
469 | "pom.xml"
470 )
471}
472
473fn resolved(mode: &str, source: &'static str) -> ResolvedMode {
474 count_source(source);
475 ResolvedMode {
476 mode: mode.to_string(),
477 source,
478 }
479}
480
481#[cfg(test)]
482mod tests {
483 use super::*;
484
485 #[test]
486 fn pressure_suggest_full_to_map() {
487 assert_eq!(
488 pressure_downgrade("full", &PressureAction::SuggestCompression),
489 Some("map".to_string())
490 );
491 }
492
493 #[test]
494 fn pressure_suggest_auto_to_map() {
495 assert_eq!(
496 pressure_downgrade("auto", &PressureAction::SuggestCompression),
497 Some("map".to_string())
498 );
499 }
500
501 #[test]
502 fn pressure_suggest_does_not_touch_signatures() {
503 assert!(pressure_downgrade("signatures", &PressureAction::SuggestCompression).is_none());
504 }
505
506 #[test]
507 fn pressure_force_full_to_map() {
508 assert_eq!(
509 pressure_downgrade("full", &PressureAction::ForceCompression),
510 Some("map".to_string())
511 );
512 }
513
514 #[test]
515 fn pressure_force_map_to_signatures() {
516 assert_eq!(
517 pressure_downgrade("map", &PressureAction::ForceCompression),
518 Some("signatures".to_string())
519 );
520 }
521
522 #[test]
523 fn pressure_evict_signatures_to_reference() {
524 assert_eq!(
525 pressure_downgrade("signatures", &PressureAction::EvictLeastRelevant),
526 Some("reference".to_string())
527 );
528 }
529
530 #[test]
531 fn pressure_noaction_returns_none() {
532 assert!(pressure_downgrade("full", &PressureAction::NoAction).is_none());
533 }
534
535 #[test]
536 fn flush_sources_merges_additively_into_disk_file() {
537 let _lock = crate::core::data_dir::test_env_lock();
538 let dir = std::env::temp_dir().join(format!("lctx-amr-flush-{}", std::process::id()));
539 let _ = std::fs::create_dir_all(&dir);
540 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
541 let _ = std::fs::remove_file(dir.join("auto_mode_sources.json"));
542
543 count_source("test_flush_alpha");
546 count_source("test_flush_alpha");
547 count_source("test_flush_beta");
548 flush_sources();
549
550 count_source("test_flush_alpha");
551 flush_sources();
552
553 let persisted = persisted_source_counts();
554 let get = |k: &str| {
555 persisted
556 .iter()
557 .find(|(s, _)| s == k)
558 .map_or(0, |(_, n)| *n)
559 };
560 assert_eq!(
561 get("test_flush_alpha"),
562 3,
563 "two flushes must merge additively"
564 );
565 assert_eq!(get("test_flush_beta"), 1);
566
567 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
568 let _ = std::fs::remove_dir_all(&dir);
569 }
570
571 #[test]
572 fn small_file_always_full() {
573 let ctx = AutoModeContext {
574 path: "test.rs",
575 token_count: 100,
576 task: None,
577 cache: None,
578 };
579 let result = resolve(&ctx);
580 assert_eq!(result.mode, "full");
581 assert_eq!(result.source, "small_file");
582 }
583
584 #[test]
585 fn config_file_returns_full() {
586 let ctx = AutoModeContext {
587 path: "config.ini",
588 token_count: 500,
589 task: None,
590 cache: None,
591 };
592 let result = resolve(&ctx);
593 assert_eq!(result.mode, "full");
594 assert_eq!(result.source, "config_data");
595 }
596
597 #[test]
598 fn cached_compressed_only_file_does_not_escalate_to_full() {
599 let dir = tempfile::tempdir().unwrap();
611 let file = dir.path().join("large.rs");
612 let body = "fn placeholder() { let _ = 1; }\n".repeat(900);
613 std::fs::write(&file, &body).unwrap();
614 let path = file.to_str().unwrap();
615
616 let mut cache = SessionCache::new();
617 cache.store(path, &body);
618 let ctx = AutoModeContext {
621 path,
622 token_count: 7000,
623 task: None,
624 cache: Some(&cache),
625 };
626 let result = resolve(&ctx);
627 assert_ne!(
628 result.mode, "full",
629 "compressed-only cached file must not escalate to full on re-read"
630 );
631 assert_ne!(result.source, "cache_hit");
632 }
633
634 #[test]
635 fn cached_full_delivered_file_short_circuits_to_stub() {
636 let dir = tempfile::tempdir().unwrap();
639 let file = dir.path().join("medium.rs");
640 let body = "fn placeholder() { let _ = 1; }\n".repeat(400);
641 std::fs::write(&file, &body).unwrap();
642 let path = file.to_str().unwrap();
643
644 let mut cache = SessionCache::new();
645 cache.store(path, &body);
646 cache.mark_full_delivered(path);
647
648 let ctx = AutoModeContext {
649 path,
650 token_count: 3000,
651 task: None,
652 cache: Some(&cache),
653 };
654 let result = resolve(&ctx);
655 assert_eq!(result.mode, "full");
656 assert_eq!(result.source, "cache_hit");
657 }
658
659 #[test]
660 fn intent_explore_returns_map() {
661 let ctx = AutoModeContext {
662 path: "large.rs",
663 token_count: 5000,
664 task: Some("how does the cache work?"),
665 cache: None,
666 };
667 let result = resolve(&ctx);
668 assert_eq!(result.mode, "map");
669 assert_eq!(result.source, "intent");
670 }
671
672 #[test]
673 fn task_names_file_matches_explicit_filename() {
674 assert!(task_names_file(
675 Some("fix the version sort in versioncmp.c"),
676 "src/versioncmp.c"
677 ));
678 assert!(task_names_file(
679 Some("why does graph.ts loop?"),
680 "web/src/graph.ts"
681 ));
682 }
683
684 #[test]
685 fn task_names_file_ignores_bare_stems_and_trivia() {
686 assert!(!task_names_file(
688 Some("improve the parser"),
689 "src/parser.rs"
690 ));
691 assert!(!task_names_file(None, "src/parser.rs"));
692 assert!(!task_names_file(Some("touch a.c"), "a.c"));
694 assert!(!task_names_file(Some("look at Makefile"), "Makefile"));
695 }
696
697 #[test]
698 fn task_suspect_file_overrides_intent() {
699 let _lock = crate::core::data_dir::test_env_lock();
700 let dir = std::env::temp_dir().join(format!("lctx-amr-suspect-{}", std::process::id()));
701 let _ = std::fs::create_dir_all(&dir);
702 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
703
704 let ctx = AutoModeContext {
708 path: "large.rs",
709 token_count: 5000,
710 task: Some("how does large.rs build the cache?"),
711 cache: None,
712 };
713 let result = resolve(&ctx);
714 assert_eq!(result.mode, "full");
715 assert_eq!(result.source, "task_suspect_file");
716
717 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
718 let _ = std::fs::remove_dir_all(&dir);
719 }
720
721 #[test]
722 fn heuristic_full_for_medium_code_by_default() {
723 assert_eq!(heuristic_mode("rs", 1500, false), "full");
726 assert_eq!(heuristic_mode("ts", 1000, false), "full");
727 }
728
729 #[test]
730 fn heuristic_structure_first_maps_medium_code() {
731 assert_eq!(heuristic_mode("rs", 1500, true), "map");
733 assert_eq!(heuristic_mode("c", 800, true), "map");
734 }
735
736 #[test]
737 fn heuristic_structure_first_keeps_tiny_and_prose_full() {
738 assert_eq!(heuristic_mode("rs", 400, true), "full");
740 assert_eq!(heuristic_mode("md", 4000, true), "full");
742 assert_eq!(heuristic_mode("txt", 1000, true), "full");
743 }
744
745 #[test]
746 fn heuristic_large_code_maps_regardless() {
747 assert_eq!(heuristic_mode("rs", 9000, false), "map");
748 assert_eq!(heuristic_mode("rs", 9000, true), "map");
749 }
750
751 #[test]
757 fn structure_first_resolve_bugfix_cold_read() {
758 let _lock = crate::core::data_dir::test_env_lock();
759 let dir = std::env::temp_dir().join(format!("lctx-amr-sf-{}", std::process::id()));
760 let _ = std::fs::create_dir_all(&dir);
761 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
762 crate::test_env::set_var("LEAN_CTX_STRUCTURE_FIRST", "1");
763
764 let suspect = AutoModeContext {
765 path: "src/versioncmp.c",
766 token_count: 1500,
767 task: None,
768 cache: None,
769 };
770 let result = resolve(&suspect);
771 assert_eq!(result.mode, "map");
772 assert_eq!(result.source, "structure_first");
773
774 let tiny = AutoModeContext {
775 path: "src/util.c",
776 token_count: 120,
777 task: None,
778 cache: None,
779 };
780 assert_eq!(resolve(&tiny).mode, "full");
781
782 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
783 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
784 let _ = std::fs::remove_dir_all(&dir);
785 }
786
787 #[test]
788 fn structure_first_off_keeps_medium_code_full() {
789 let _lock = crate::core::data_dir::test_env_lock();
790 let dir = std::env::temp_dir().join(format!("lctx-amr-sfoff-{}", std::process::id()));
791 let _ = std::fs::create_dir_all(&dir);
792 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
793 crate::test_env::set_var("LEAN_CTX_STRUCTURE_FIRST", "0");
794
795 let ctx = AutoModeContext {
796 path: "src/versioncmp.c",
797 token_count: 1500,
798 task: None,
799 cache: None,
800 };
801 let result = resolve(&ctx);
802 assert_eq!(result.mode, "full");
803 assert_eq!(result.source, "heuristic");
804
805 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
806 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
807 let _ = std::fs::remove_dir_all(&dir);
808 }
809
810 #[test]
814 fn learning_off_by_default_is_deterministic() {
815 let _lock = crate::core::data_dir::test_env_lock();
816 let dir = std::env::temp_dir().join(format!("lctx-amr-det-{}", std::process::id()));
817 let _ = std::fs::create_dir_all(&dir);
818 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
819 crate::test_env::remove_var("LEAN_CTX_AUTO_MODE_LEARNING");
820 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
821
822 let ctx = AutoModeContext {
823 path: "src/widget.rs",
824 token_count: 1500,
825 task: None,
826 cache: None,
827 };
828 let a = resolve(&ctx);
829 let b = resolve(&ctx);
830 assert_eq!(a.mode, "full");
831 assert_eq!(a.source, "heuristic");
832 assert_eq!((a.mode, a.source), (b.mode, b.source));
833
834 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
835 let _ = std::fs::remove_dir_all(&dir);
836 }
837
838 #[test]
841 fn auto_mode_learning_env_opt_in_is_honored() {
842 let _lock = crate::core::data_dir::test_env_lock();
843 crate::test_env::set_var("LEAN_CTX_AUTO_MODE_LEARNING", "1");
844 assert!(crate::core::config::Config::default().auto_mode_learning_effective());
845 crate::test_env::set_var("LEAN_CTX_AUTO_MODE_LEARNING", "0");
846 assert!(!crate::core::config::Config::default().auto_mode_learning_effective());
847 crate::test_env::remove_var("LEAN_CTX_AUTO_MODE_LEARNING");
848 }
849}