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.choose_arm();
362 if arm.budget_ratio < 0.25 && token_count > 2000 {
363 Some("aggressive".to_string())
364 } else {
365 None
366 }
367}
368
369fn heuristic_mode(ext: &str, token_count: usize, structure_first: bool) -> String {
370 if token_count > 8000 {
371 if is_code(ext) {
372 return "map".to_string();
373 }
374 return "aggressive".to_string();
375 }
376 if token_count > 6000 && is_code(ext) {
381 return "map".to_string();
382 }
383 if structure_first && token_count > 500 && is_code(ext) {
390 return "map".to_string();
391 }
392 "full".to_string()
393}
394
395fn file_unchanged(path: &str, cached: &crate::core::cache::CacheEntry) -> bool {
404 let Some(stored_mtime) = cached.stored_mtime else {
405 return false;
406 };
407 let Ok(meta) = std::fs::metadata(path) else {
408 return false;
409 };
410 let Ok(current_mtime) = meta.modified() else {
411 return false;
412 };
413 current_mtime == stored_mtime
414}
415
416fn is_code(ext: &str) -> bool {
417 matches!(
418 ext,
419 "rs" | "ts"
420 | "tsx"
421 | "js"
422 | "jsx"
423 | "py"
424 | "go"
425 | "java"
426 | "c"
427 | "cpp"
428 | "cc"
429 | "h"
430 | "hpp"
431 | "rb"
432 | "cs"
433 | "kt"
434 | "swift"
435 | "php"
436 | "zig"
437 | "ex"
438 | "exs"
439 | "scala"
440 | "sc"
441 | "dart"
442 | "sh"
443 | "bash"
444 | "svelte"
445 | "vue"
446 )
447}
448
449fn is_config_or_data(ext: &str, path: &str) -> bool {
450 if matches!(ext, "xml" | "ini" | "cfg" | "env") {
451 return true;
452 }
453 let name = std::path::Path::new(path)
454 .file_name()
455 .and_then(|n| n.to_str())
456 .unwrap_or("");
457 matches!(
458 name,
459 "Cargo.toml"
460 | "package.json"
461 | "tsconfig.json"
462 | "Makefile"
463 | "Dockerfile"
464 | "docker-compose.yml"
465 | ".gitignore"
466 | ".env"
467 | "pyproject.toml"
468 | "go.mod"
469 | "build.gradle"
470 | "pom.xml"
471 )
472}
473
474fn resolved(mode: &str, source: &'static str) -> ResolvedMode {
475 count_source(source);
476 ResolvedMode {
477 mode: mode.to_string(),
478 source,
479 }
480}
481
482#[cfg(test)]
483mod tests {
484 use super::*;
485
486 #[test]
487 fn pressure_suggest_full_to_map() {
488 assert_eq!(
489 pressure_downgrade("full", &PressureAction::SuggestCompression),
490 Some("map".to_string())
491 );
492 }
493
494 #[test]
495 fn pressure_suggest_auto_to_map() {
496 assert_eq!(
497 pressure_downgrade("auto", &PressureAction::SuggestCompression),
498 Some("map".to_string())
499 );
500 }
501
502 #[test]
503 fn pressure_suggest_does_not_touch_signatures() {
504 assert!(pressure_downgrade("signatures", &PressureAction::SuggestCompression).is_none());
505 }
506
507 #[test]
508 fn pressure_force_full_to_map() {
509 assert_eq!(
510 pressure_downgrade("full", &PressureAction::ForceCompression),
511 Some("map".to_string())
512 );
513 }
514
515 #[test]
516 fn pressure_force_map_to_signatures() {
517 assert_eq!(
518 pressure_downgrade("map", &PressureAction::ForceCompression),
519 Some("signatures".to_string())
520 );
521 }
522
523 #[test]
524 fn pressure_evict_signatures_to_reference() {
525 assert_eq!(
526 pressure_downgrade("signatures", &PressureAction::EvictLeastRelevant),
527 Some("reference".to_string())
528 );
529 }
530
531 #[test]
532 fn pressure_noaction_returns_none() {
533 assert!(pressure_downgrade("full", &PressureAction::NoAction).is_none());
534 }
535
536 #[test]
537 fn flush_sources_merges_additively_into_disk_file() {
538 let _lock = crate::core::data_dir::test_env_lock();
539 let dir = std::env::temp_dir().join(format!("lctx-amr-flush-{}", std::process::id()));
540 let _ = std::fs::create_dir_all(&dir);
541 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
542 let _ = std::fs::remove_file(dir.join("auto_mode_sources.json"));
543
544 count_source("test_flush_alpha");
547 count_source("test_flush_alpha");
548 count_source("test_flush_beta");
549 flush_sources();
550
551 count_source("test_flush_alpha");
552 flush_sources();
553
554 let persisted = persisted_source_counts();
555 let get = |k: &str| {
556 persisted
557 .iter()
558 .find(|(s, _)| s == k)
559 .map_or(0, |(_, n)| *n)
560 };
561 assert_eq!(
562 get("test_flush_alpha"),
563 3,
564 "two flushes must merge additively"
565 );
566 assert_eq!(get("test_flush_beta"), 1);
567
568 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
569 let _ = std::fs::remove_dir_all(&dir);
570 }
571
572 #[test]
573 fn small_file_always_full() {
574 let ctx = AutoModeContext {
575 path: "test.rs",
576 token_count: 100,
577 task: None,
578 cache: None,
579 };
580 let result = resolve(&ctx);
581 assert_eq!(result.mode, "full");
582 assert_eq!(result.source, "small_file");
583 }
584
585 #[test]
586 fn config_file_returns_full() {
587 let ctx = AutoModeContext {
588 path: "config.ini",
589 token_count: 500,
590 task: None,
591 cache: None,
592 };
593 let result = resolve(&ctx);
594 assert_eq!(result.mode, "full");
595 assert_eq!(result.source, "config_data");
596 }
597
598 #[test]
599 fn cached_compressed_only_file_does_not_escalate_to_full() {
600 let dir = tempfile::tempdir().unwrap();
612 let file = dir.path().join("large.rs");
613 let body = "fn placeholder() { let _ = 1; }\n".repeat(900);
614 std::fs::write(&file, &body).unwrap();
615 let path = file.to_str().unwrap();
616
617 let mut cache = SessionCache::new();
618 cache.store(path, &body);
619 let ctx = AutoModeContext {
622 path,
623 token_count: 7000,
624 task: None,
625 cache: Some(&cache),
626 };
627 let result = resolve(&ctx);
628 assert_ne!(
629 result.mode, "full",
630 "compressed-only cached file must not escalate to full on re-read"
631 );
632 assert_ne!(result.source, "cache_hit");
633 }
634
635 #[test]
636 fn cached_full_delivered_file_short_circuits_to_stub() {
637 let dir = tempfile::tempdir().unwrap();
640 let file = dir.path().join("medium.rs");
641 let body = "fn placeholder() { let _ = 1; }\n".repeat(400);
642 std::fs::write(&file, &body).unwrap();
643 let path = file.to_str().unwrap();
644
645 let mut cache = SessionCache::new();
646 cache.store(path, &body);
647 cache.mark_full_delivered(path);
648
649 let ctx = AutoModeContext {
650 path,
651 token_count: 3000,
652 task: None,
653 cache: Some(&cache),
654 };
655 let result = resolve(&ctx);
656 assert_eq!(result.mode, "full");
657 assert_eq!(result.source, "cache_hit");
658 }
659
660 #[test]
661 fn intent_explore_returns_map() {
662 let ctx = AutoModeContext {
663 path: "large.rs",
664 token_count: 5000,
665 task: Some("how does the cache work?"),
666 cache: None,
667 };
668 let result = resolve(&ctx);
669 assert_eq!(result.mode, "map");
670 assert_eq!(result.source, "intent");
671 }
672
673 #[test]
674 fn task_names_file_matches_explicit_filename() {
675 assert!(task_names_file(
676 Some("fix the version sort in versioncmp.c"),
677 "src/versioncmp.c"
678 ));
679 assert!(task_names_file(
680 Some("why does graph.ts loop?"),
681 "web/src/graph.ts"
682 ));
683 }
684
685 #[test]
686 fn task_names_file_ignores_bare_stems_and_trivia() {
687 assert!(!task_names_file(
689 Some("improve the parser"),
690 "src/parser.rs"
691 ));
692 assert!(!task_names_file(None, "src/parser.rs"));
693 assert!(!task_names_file(Some("touch a.c"), "a.c"));
695 assert!(!task_names_file(Some("look at Makefile"), "Makefile"));
696 }
697
698 #[test]
699 fn task_suspect_file_overrides_intent() {
700 let _lock = crate::core::data_dir::test_env_lock();
701 let dir = std::env::temp_dir().join(format!("lctx-amr-suspect-{}", std::process::id()));
702 let _ = std::fs::create_dir_all(&dir);
703 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
704
705 let ctx = AutoModeContext {
709 path: "large.rs",
710 token_count: 5000,
711 task: Some("how does large.rs build the cache?"),
712 cache: None,
713 };
714 let result = resolve(&ctx);
715 assert_eq!(result.mode, "full");
716 assert_eq!(result.source, "task_suspect_file");
717
718 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
719 let _ = std::fs::remove_dir_all(&dir);
720 }
721
722 #[test]
723 fn heuristic_full_for_medium_code_by_default() {
724 assert_eq!(heuristic_mode("rs", 1500, false), "full");
727 assert_eq!(heuristic_mode("ts", 1000, false), "full");
728 }
729
730 #[test]
731 fn heuristic_structure_first_maps_medium_code() {
732 assert_eq!(heuristic_mode("rs", 1500, true), "map");
734 assert_eq!(heuristic_mode("c", 800, true), "map");
735 }
736
737 #[test]
738 fn heuristic_structure_first_keeps_tiny_and_prose_full() {
739 assert_eq!(heuristic_mode("rs", 400, true), "full");
741 assert_eq!(heuristic_mode("md", 4000, true), "full");
743 assert_eq!(heuristic_mode("txt", 1000, true), "full");
744 }
745
746 #[test]
747 fn heuristic_large_code_maps_regardless() {
748 assert_eq!(heuristic_mode("rs", 9000, false), "map");
749 assert_eq!(heuristic_mode("rs", 9000, true), "map");
750 }
751
752 #[test]
758 fn structure_first_resolve_bugfix_cold_read() {
759 let _lock = crate::core::data_dir::test_env_lock();
760 let dir = std::env::temp_dir().join(format!("lctx-amr-sf-{}", std::process::id()));
761 let _ = std::fs::create_dir_all(&dir);
762 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
763 crate::test_env::set_var("LEAN_CTX_STRUCTURE_FIRST", "1");
764
765 let suspect = AutoModeContext {
766 path: "src/versioncmp.c",
767 token_count: 1500,
768 task: None,
769 cache: None,
770 };
771 let result = resolve(&suspect);
772 assert_eq!(result.mode, "map");
773 assert_eq!(result.source, "structure_first");
774
775 let tiny = AutoModeContext {
776 path: "src/util.c",
777 token_count: 120,
778 task: None,
779 cache: None,
780 };
781 assert_eq!(resolve(&tiny).mode, "full");
782
783 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
784 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
785 let _ = std::fs::remove_dir_all(&dir);
786 }
787
788 #[test]
789 fn structure_first_off_keeps_medium_code_full() {
790 let _lock = crate::core::data_dir::test_env_lock();
791 let dir = std::env::temp_dir().join(format!("lctx-amr-sfoff-{}", std::process::id()));
792 let _ = std::fs::create_dir_all(&dir);
793 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
794 crate::test_env::set_var("LEAN_CTX_STRUCTURE_FIRST", "0");
795
796 let ctx = AutoModeContext {
797 path: "src/versioncmp.c",
798 token_count: 1500,
799 task: None,
800 cache: None,
801 };
802 let result = resolve(&ctx);
803 assert_eq!(result.mode, "full");
804 assert_eq!(result.source, "heuristic");
805
806 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
807 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
808 let _ = std::fs::remove_dir_all(&dir);
809 }
810
811 #[test]
815 fn learning_off_by_default_is_deterministic() {
816 let _lock = crate::core::data_dir::test_env_lock();
817 let dir = std::env::temp_dir().join(format!("lctx-amr-det-{}", std::process::id()));
818 let _ = std::fs::create_dir_all(&dir);
819 crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
820 crate::test_env::remove_var("LEAN_CTX_AUTO_MODE_LEARNING");
821 crate::test_env::remove_var("LEAN_CTX_STRUCTURE_FIRST");
822
823 let ctx = AutoModeContext {
824 path: "src/widget.rs",
825 token_count: 1500,
826 task: None,
827 cache: None,
828 };
829 let a = resolve(&ctx);
830 let b = resolve(&ctx);
831 assert_eq!(a.mode, "full");
832 assert_eq!(a.source, "heuristic");
833 assert_eq!((a.mode, a.source), (b.mode, b.source));
834
835 crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
836 let _ = std::fs::remove_dir_all(&dir);
837 }
838
839 #[test]
842 fn auto_mode_learning_env_opt_in_is_honored() {
843 let _lock = crate::core::data_dir::test_env_lock();
844 crate::test_env::set_var("LEAN_CTX_AUTO_MODE_LEARNING", "1");
845 assert!(crate::core::config::Config::default().auto_mode_learning_effective());
846 crate::test_env::set_var("LEAN_CTX_AUTO_MODE_LEARNING", "0");
847 assert!(!crate::core::config::Config::default().auto_mode_learning_effective());
848 crate::test_env::remove_var("LEAN_CTX_AUTO_MODE_LEARNING");
849 }
850}