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