1use std::collections::HashMap;
2use std::time::{Duration, Instant};
3
4use super::config::LoopDetectionConfig;
5
6const SEARCH_TOOLS: &[&str] = &["ctx_search", "ctx_semantic_search"];
7
8const SEARCH_SHELL_PREFIXES: &[&str] = &["grep ", "rg ", "find ", "fd ", "ag ", "ack "];
9
10const CORRECTION_WINDOW: Duration = Duration::from_mins(2);
11const MODE_BOUNCE_WINDOW: Duration = Duration::from_secs(30);
12const SHELL_RERUN_WINDOW: Duration = Duration::from_mins(1);
13const COLD_START_CALLS: u32 = 3;
14
15#[derive(Debug, Clone, PartialEq)]
17pub enum CorrectionKind {
18 FreshReRead,
19 ShellReRun,
20 ModeBounce,
21}
22
23#[derive(Debug, Clone)]
25pub struct LoopDetector {
26 call_history: HashMap<String, Vec<Instant>>,
27 duplicate_counts: HashMap<String, u32>,
28 tool_total_counts: HashMap<String, u32>,
29 tool_total_limits: HashMap<String, u32>,
30 search_group_history: Vec<Instant>,
31 recent_search_patterns: Vec<String>,
32 normal_threshold: u32,
33 reduced_threshold: u32,
34 blocked_threshold: u32,
35 window: Duration,
36 search_group_limit: u32,
37 correction_signals: Vec<(Instant, CorrectionKind)>,
39 recent_reads: HashMap<String, (Instant, String)>,
40 recent_commands: HashMap<String, Instant>,
41 total_calls: u32,
42}
43
44#[derive(Debug, Clone, PartialEq)]
46pub enum ThrottleLevel {
47 Normal,
48 Reduced,
49 Blocked,
50}
51
52#[derive(Debug, Clone)]
54pub struct ThrottleResult {
55 pub level: ThrottleLevel,
56 pub call_count: u32,
57 pub message: Option<String>,
58}
59
60impl Default for ThrottleResult {
61 fn default() -> Self {
62 Self {
63 level: ThrottleLevel::Normal,
64 call_count: 0,
65 message: None,
66 }
67 }
68}
69
70impl Default for LoopDetector {
71 fn default() -> Self {
72 Self::new()
73 }
74}
75
76impl LoopDetector {
77 pub fn new() -> Self {
79 Self::with_config(&LoopDetectionConfig::default())
80 }
81
82 pub fn with_config(cfg: &LoopDetectionConfig) -> Self {
85 Self {
86 call_history: HashMap::new(),
87 duplicate_counts: HashMap::new(),
88 tool_total_counts: HashMap::new(),
89 tool_total_limits: cfg.tool_total_limits.clone(),
90 search_group_history: Vec::new(),
91 recent_search_patterns: Vec::new(),
92 normal_threshold: cfg.normal_threshold.max(1),
93 reduced_threshold: cfg.reduced_threshold.max(2),
94 blocked_threshold: cfg.blocked_threshold,
95 window: Duration::from_secs(cfg.window_secs),
96 search_group_limit: if cfg.blocked_threshold == 0 {
97 u32::MAX
98 } else {
99 cfg.search_group_limit.max(3)
100 },
101 correction_signals: Vec::new(),
102 recent_reads: HashMap::new(),
103 recent_commands: HashMap::new(),
104 total_calls: 0,
105 }
106 }
107
108 pub fn record_call(&mut self, tool: &str, args_fingerprint: &str) -> ThrottleResult {
110 let now = Instant::now();
111 self.prune_window(now);
112
113 let total = self.tool_total_counts.entry(tool.to_string()).or_insert(0);
115 *total += 1;
116 let total_count = *total;
117
118 if let Some(&limit) = self.tool_total_limits.get(tool) {
119 if total_count > limit {
120 let msg = if crate::core::protocol::meta_visible() {
121 Some(format!(
122 "Warning: {tool} called {total_count}x total (limit: {limit}). \
123 Consider ctx_compress or narrowing scope."
124 ))
125 } else {
126 None
127 };
128 return ThrottleResult {
129 level: ThrottleLevel::Reduced,
130 call_count: total_count,
131 message: msg,
132 };
133 }
134 }
135
136 let key = format!("{tool}:{args_fingerprint}");
137 let entries = self.call_history.entry(key.clone()).or_default();
138 entries.push(now);
139 let count = entries.len() as u32;
140 *self.duplicate_counts.entry(key).or_default() = count;
141
142 if self.blocked_threshold > 0 && count > self.blocked_threshold {
143 return ThrottleResult {
144 level: ThrottleLevel::Blocked,
145 call_count: count,
146 message: Some(self.block_message(tool, count)),
147 };
148 }
149 if count > self.reduced_threshold {
150 if !crate::core::protocol::meta_visible() {
151 return ThrottleResult {
152 level: ThrottleLevel::Reduced,
153 call_count: count,
154 message: None,
155 };
156 }
157 return ThrottleResult {
158 level: ThrottleLevel::Reduced,
159 call_count: count,
160 message: Some(format!(
161 "Warning: {tool} called {count}x with same args. \
162 Results reduced. Try a different approach or narrow your scope."
163 )),
164 };
165 }
166 if count > self.normal_threshold {
167 if !crate::core::protocol::meta_visible() {
168 return ThrottleResult {
169 level: ThrottleLevel::Reduced,
170 call_count: count,
171 message: None,
172 };
173 }
174 return ThrottleResult {
175 level: ThrottleLevel::Reduced,
176 call_count: count,
177 message: Some(format!(
178 "Note: {tool} called {count}x with similar args. Consider narrowing scope."
179 )),
180 };
181 }
182 ThrottleResult {
183 level: ThrottleLevel::Normal,
184 call_count: count,
185 message: None,
186 }
187 }
188
189 pub fn record_error_outcome(&mut self, tool: &str, args_fingerprint: &str) {
192 let key = format!("{tool}:{args_fingerprint}");
193 if let Some(entries) = self.call_history.get_mut(&key) {
194 entries.pop();
195 let count = entries.len() as u32;
196 self.duplicate_counts.insert(key, count);
197 }
198 }
199
200 pub fn record_search(
203 &mut self,
204 tool: &str,
205 args_fingerprint: &str,
206 search_pattern: Option<&str>,
207 ) -> ThrottleResult {
208 let now = Instant::now();
209
210 self.search_group_history.push(now);
211 let search_count = self.search_group_history.len() as u32;
212
213 let similar_count = if let Some(pat) = search_pattern {
214 let sc = self.count_similar_patterns(pat);
215 if !pat.is_empty() {
216 self.recent_search_patterns.push(pat.to_string());
217 if self.recent_search_patterns.len() > 15 {
218 self.recent_search_patterns.remove(0);
219 }
220 }
221 sc
222 } else {
223 0
224 };
225
226 if self.blocked_threshold > 0 && similar_count >= self.blocked_threshold {
228 return ThrottleResult {
229 level: ThrottleLevel::Blocked,
230 call_count: similar_count,
231 message: Some(self.search_block_message(similar_count)),
232 };
233 }
234
235 if self.blocked_threshold > 0 && search_count > self.search_group_limit {
237 return ThrottleResult {
238 level: ThrottleLevel::Blocked,
239 call_count: search_count,
240 message: Some(self.search_group_block_message(search_count)),
241 };
242 }
243
244 if similar_count >= self.reduced_threshold {
245 if !crate::core::protocol::meta_visible() {
246 return ThrottleResult {
247 level: ThrottleLevel::Reduced,
248 call_count: similar_count,
249 message: None,
250 };
251 }
252 return ThrottleResult {
253 level: ThrottleLevel::Reduced,
254 call_count: similar_count,
255 message: Some(format!(
256 "Warning: You've searched for similar patterns {similar_count}x. \
257 Narrow your search with the 'path' parameter or try ctx_tree first."
258 )),
259 };
260 }
261
262 if search_count > self.search_group_limit.saturating_sub(3) {
263 let per_fp = self.record_call(tool, args_fingerprint);
264 if per_fp.level != ThrottleLevel::Normal {
265 return per_fp;
266 }
267 if !crate::core::protocol::meta_visible() {
268 return ThrottleResult {
269 level: ThrottleLevel::Reduced,
270 call_count: search_count,
271 message: None,
272 };
273 }
274 return ThrottleResult {
275 level: ThrottleLevel::Reduced,
276 call_count: search_count,
277 message: Some(format!(
278 "Note: {search_count} search calls in the last {}s. \
279 Use ctx_tree to orient first, then scope searches with 'path'.",
280 self.window.as_secs()
281 )),
282 };
283 }
284
285 self.record_call(tool, args_fingerprint)
286 }
287
288 pub fn is_search_tool(tool: &str) -> bool {
290 SEARCH_TOOLS.contains(&tool)
291 }
292
293 pub fn is_search_shell_command(command: &str) -> bool {
295 let cmd = command.trim_start();
296 SEARCH_SHELL_PREFIXES.iter().any(|p| cmd.starts_with(p))
297 }
298
299 pub fn fingerprint(args: &serde_json::Value) -> String {
301 use std::collections::hash_map::DefaultHasher;
302 use std::hash::{Hash, Hasher};
303
304 let canonical = canonical_json(args);
305 let mut hasher = DefaultHasher::new();
306 canonical.hash(&mut hasher);
307 format!("{:016x}", hasher.finish())
308 }
309
310 pub fn stats(&self) -> Vec<(String, u32)> {
312 let mut entries: Vec<(String, u32)> = self
313 .duplicate_counts
314 .iter()
315 .filter(|(_, &count)| count > 1)
316 .map(|(k, &v)| (k.clone(), v))
317 .collect();
318 entries.sort_by_key(|x| std::cmp::Reverse(x.1));
319 entries
320 }
321
322 pub fn record_read_for_correction(&mut self, path: &str, mode: &str, fresh: bool) {
326 self.total_calls += 1;
327 let now = Instant::now();
328
329 if self.total_calls <= COLD_START_CALLS {
330 self.recent_reads
331 .insert(path.to_string(), (now, mode.to_string()));
332 return;
333 }
334
335 if fresh {
336 if let Some((prev_time, _)) = self.recent_reads.get(path) {
337 if now.duration_since(*prev_time) < CORRECTION_WINDOW {
338 self.correction_signals
339 .push((now, CorrectionKind::FreshReRead));
340 }
341 }
342 }
343
344 if mode == "full" {
345 if let Some((prev_time, prev_mode)) = self.recent_reads.get(path) {
346 let is_bounce = (prev_mode == "map" || prev_mode == "signatures")
347 && now.duration_since(*prev_time) < MODE_BOUNCE_WINDOW;
348 if is_bounce {
349 self.correction_signals
350 .push((now, CorrectionKind::ModeBounce));
351 }
352 }
353 }
354
355 self.recent_reads
356 .insert(path.to_string(), (now, mode.to_string()));
357 }
358
359 pub fn record_shell_for_correction(&mut self, command: &str) {
361 self.total_calls += 1;
362 let now = Instant::now();
363
364 if self.total_calls <= COLD_START_CALLS {
365 self.recent_commands.insert(command.to_string(), now);
366 return;
367 }
368
369 let key = normalize_shell_command(command);
370 if let Some(prev_time) = self.recent_commands.get(&key) {
371 if now.duration_since(*prev_time) < SHELL_RERUN_WINDOW {
372 self.correction_signals
373 .push((now, CorrectionKind::ShellReRun));
374 }
375 }
376 self.recent_commands.insert(key, now);
377 }
378
379 pub fn correction_count(&self) -> u32 {
381 let now = Instant::now();
382 self.correction_signals
383 .iter()
384 .filter(|(t, _)| now.duration_since(*t) < CORRECTION_WINDOW)
385 .count() as u32
386 }
387
388 pub fn correction_rate(&self) -> f64 {
390 let count = self.correction_count();
391 if count == 0 {
392 return 0.0;
393 }
394 let window_mins = CORRECTION_WINDOW.as_secs_f64() / 60.0;
395 f64::from(count) / window_mins
396 }
397
398 pub fn prune_corrections(&mut self) {
400 let now = Instant::now();
401 self.correction_signals
402 .retain(|(t, _)| now.duration_since(*t) < CORRECTION_WINDOW);
403 self.recent_reads
404 .retain(|_, (t, _)| now.duration_since(*t) < CORRECTION_WINDOW);
405 self.recent_commands
406 .retain(|_, t| now.duration_since(*t) < CORRECTION_WINDOW);
407 }
408
409 pub fn reset(&mut self) {
411 self.call_history.clear();
412 self.duplicate_counts.clear();
413 self.search_group_history.clear();
414 self.recent_search_patterns.clear();
415 self.correction_signals.clear();
416 self.recent_reads.clear();
417 self.recent_commands.clear();
418 self.total_calls = 0;
419 }
420
421 fn prune_window(&mut self, now: Instant) {
422 for entries in self.call_history.values_mut() {
423 entries.retain(|t| now.duration_since(*t) < self.window);
424 }
425 self.call_history.retain(|_, v| !v.is_empty());
431 let live = &self.call_history;
432 self.duplicate_counts.retain(|k, _| live.contains_key(k));
433 self.search_group_history
434 .retain(|t| now.duration_since(*t) < self.window);
435 }
436
437 fn count_similar_patterns(&self, new_pattern: &str) -> u32 {
438 let new_lower = new_pattern.to_lowercase();
439 let new_root = extract_alpha_root(&new_lower);
440
441 let mut count = 0u32;
442 for existing in &self.recent_search_patterns {
443 let existing_lower = existing.to_lowercase();
444 if patterns_are_similar(&new_lower, &existing_lower) {
445 count += 1;
446 } else if new_root.len() >= 4 {
447 let existing_root = extract_alpha_root(&existing_lower);
448 if existing_root.len() >= 4
449 && (new_root.starts_with(&existing_root)
450 || existing_root.starts_with(&new_root))
451 {
452 count += 1;
453 }
454 }
455 }
456 count
457 }
458
459 fn block_message(&self, tool: &str, count: u32) -> String {
460 if Self::is_search_tool(tool) {
461 self.search_block_message(count)
462 } else {
463 format!(
464 "LOOP DETECTED: {tool} called {count}x with same/similar args. \
465 Call blocked. Change your approach — the current strategy is not working."
466 )
467 }
468 }
469
470 #[allow(clippy::unused_self)]
471 fn search_block_message(&self, count: u32) -> String {
472 format!(
473 "LOOP DETECTED: You've searched {count}x with similar patterns. STOP searching and change strategy. \
474 1) Use ctx_tree to understand the project structure first. \
475 2) Narrow your search with the 'path' parameter to a specific directory. \
476 3) Use ctx_read with mode='map' to understand a file before searching more."
477 )
478 }
479
480 fn search_group_block_message(&self, count: u32) -> String {
481 format!(
482 "LOOP DETECTED: {count} search calls in {}s — too many. STOP and rethink. \
483 1) Use ctx_tree to map the project structure. \
484 2) Pick ONE specific directory and search there with the 'path' parameter. \
485 3) Read files with ctx_read mode='map' instead of searching blindly.",
486 self.window.as_secs()
487 )
488 }
489}
490
491fn normalize_shell_command(cmd: &str) -> String {
492 cmd.split_whitespace()
493 .take(5)
494 .collect::<Vec<_>>()
495 .join(" ")
496 .to_lowercase()
497}
498
499fn extract_alpha_root(pattern: &str) -> String {
500 pattern
501 .chars()
502 .take_while(|c| c.is_alphanumeric())
503 .collect()
504}
505
506fn patterns_are_similar(a: &str, b: &str) -> bool {
507 if a == b {
508 return true;
509 }
510 if a.contains(b) || b.contains(a) {
511 return true;
512 }
513 let a_alpha: String = a.chars().filter(|c| c.is_alphanumeric()).collect();
514 let b_alpha: String = b.chars().filter(|c| c.is_alphanumeric()).collect();
515 if a_alpha.len() >= 3
516 && b_alpha.len() >= 3
517 && (a_alpha.contains(&b_alpha) || b_alpha.contains(&a_alpha))
518 {
519 return true;
520 }
521 false
522}
523
524fn canonical_json(value: &serde_json::Value) -> String {
525 match value {
526 serde_json::Value::Object(map) => {
527 let mut keys: Vec<&String> = map.keys().collect();
528 keys.sort();
529 let entries: Vec<String> = keys
530 .iter()
531 .map(|k| format!("{}:{}", k, canonical_json(&map[*k])))
532 .collect();
533 format!("{{{}}}", entries.join(","))
534 }
535 serde_json::Value::Array(arr) => {
536 let entries: Vec<String> = arr.iter().map(canonical_json).collect();
537 format!("[{}]", entries.join(","))
538 }
539 _ => value.to_string(),
540 }
541}
542
543#[cfg(test)]
544mod tests {
545 use super::*;
546
547 fn test_config(normal: u32, reduced: u32, blocked: u32) -> LoopDetectionConfig {
548 LoopDetectionConfig {
549 normal_threshold: normal,
550 reduced_threshold: reduced,
551 blocked_threshold: blocked,
552 window_secs: 300,
553 search_group_limit: 10,
554 tool_total_limits: std::collections::HashMap::new(),
555 }
556 }
557
558 #[test]
559 fn normal_calls_pass_through() {
560 let mut detector = LoopDetector::new();
561 let r1 = detector.record_call("ctx_read", "abc123");
562 assert_eq!(r1.level, ThrottleLevel::Normal);
563 assert_eq!(r1.call_count, 1);
564 assert!(r1.message.is_none());
565 }
566
567 #[test]
568 fn repeated_calls_trigger_reduced() {
569 let _lock = crate::core::data_dir::test_env_lock();
570 std::env::set_var("LEAN_CTX_META", "1");
571 let cfg = LoopDetectionConfig::default();
572 let mut detector = LoopDetector::with_config(&cfg);
573 for _ in 0..cfg.normal_threshold {
574 detector.record_call("ctx_read", "same_fp");
575 }
576 let result = detector.record_call("ctx_read", "same_fp");
577 assert_eq!(result.level, ThrottleLevel::Reduced);
578 assert!(result.message.is_some());
579 std::env::remove_var("LEAN_CTX_META");
580 }
581
582 #[test]
583 fn excessive_calls_get_blocked_when_enabled() {
584 let cfg = LoopDetectionConfig {
586 blocked_threshold: 6,
587 ..Default::default()
588 };
589 let mut detector = LoopDetector::with_config(&cfg);
590 for _ in 0..cfg.blocked_threshold {
591 detector.record_call("ctx_shell", "same_fp");
592 }
593 let result = detector.record_call("ctx_shell", "same_fp");
594 assert_eq!(result.level, ThrottleLevel::Blocked);
595 assert!(result.message.unwrap().contains("LOOP DETECTED"));
596 }
597
598 #[test]
599 fn blocking_disabled_by_default() {
600 let cfg = LoopDetectionConfig::default();
602 assert_eq!(cfg.blocked_threshold, 0);
603 let mut detector = LoopDetector::with_config(&cfg);
604 for _ in 0..100 {
606 detector.record_call("ctx_shell", "same_fp");
607 }
608 let result = detector.record_call("ctx_shell", "same_fp");
609 assert_ne!(result.level, ThrottleLevel::Blocked);
611 }
612
613 #[test]
614 fn different_args_tracked_separately() {
615 let mut detector = LoopDetector::new();
616 for _ in 0..10 {
617 detector.record_call("ctx_read", "fp_a");
618 }
619 let result = detector.record_call("ctx_read", "fp_b");
620 assert_eq!(result.level, ThrottleLevel::Normal);
621 assert_eq!(result.call_count, 1);
622 }
623
624 #[test]
625 fn fingerprint_deterministic() {
626 let args = serde_json::json!({"path": "test.rs", "mode": "full"});
627 let fp1 = LoopDetector::fingerprint(&args);
628 let fp2 = LoopDetector::fingerprint(&args);
629 assert_eq!(fp1, fp2);
630 }
631
632 #[test]
633 fn fingerprint_order_independent() {
634 let a = serde_json::json!({"mode": "full", "path": "test.rs"});
635 let b = serde_json::json!({"path": "test.rs", "mode": "full"});
636 assert_eq!(LoopDetector::fingerprint(&a), LoopDetector::fingerprint(&b));
637 }
638
639 #[test]
640 fn stats_shows_duplicates() {
641 let mut detector = LoopDetector::new();
642 for _ in 0..5 {
643 detector.record_call("ctx_read", "fp_a");
644 }
645 detector.record_call("ctx_shell", "fp_b");
646 let stats = detector.stats();
647 assert_eq!(stats.len(), 1);
648 assert_eq!(stats[0].1, 5);
649 }
650
651 #[test]
652 fn reset_clears_state() {
653 let mut detector = LoopDetector::new();
654 for _ in 0..5 {
655 detector.record_call("ctx_read", "fp_a");
656 }
657 detector.reset();
658 let result = detector.record_call("ctx_read", "fp_a");
659 assert_eq!(result.call_count, 1);
660 }
661
662 #[test]
663 fn custom_thresholds_from_config() {
664 let cfg = test_config(1, 2, 3);
665 let mut detector = LoopDetector::with_config(&cfg);
666 detector.record_call("ctx_read", "fp");
667 let r = detector.record_call("ctx_read", "fp");
668 assert_eq!(r.level, ThrottleLevel::Reduced);
669 detector.record_call("ctx_read", "fp");
670 let r = detector.record_call("ctx_read", "fp");
671 assert_eq!(r.level, ThrottleLevel::Blocked);
672 }
673
674 #[test]
675 fn similar_patterns_detected() {
676 assert!(patterns_are_similar("compress", "compress"));
677 assert!(patterns_are_similar("compress", "compression"));
678 assert!(patterns_are_similar("compress.*data", "compress"));
679 assert!(!patterns_are_similar("foo", "bar"));
680 assert!(!patterns_are_similar("ab", "cd"));
681 }
682
683 #[test]
684 fn search_group_tracking_when_blocking_enabled() {
685 let cfg = LoopDetectionConfig {
687 search_group_limit: 5,
688 blocked_threshold: 6, ..Default::default()
690 };
691 let mut detector = LoopDetector::with_config(&cfg);
692 for i in 0..5 {
693 let fp = format!("fp_{i}");
694 let r = detector.record_search("ctx_search", &fp, Some(&format!("pattern_{i}")));
695 assert_ne!(r.level, ThrottleLevel::Blocked, "call {i} should not block");
696 }
697 let r = detector.record_search("ctx_search", "fp_5", Some("pattern_5"));
698 assert_eq!(r.level, ThrottleLevel::Blocked);
699 assert!(r.message.unwrap().contains("search calls"));
700 }
701
702 #[test]
703 fn similar_search_patterns_trigger_block_when_enabled() {
704 let cfg = LoopDetectionConfig {
706 blocked_threshold: 6,
707 ..Default::default()
708 };
709 let mut detector = LoopDetector::with_config(&cfg);
710 let variants = [
711 "compress",
712 "compression",
713 "compress.*data",
714 "compress_output",
715 "compressor",
716 "compress_result",
717 "compress_file",
718 ];
719 for (i, pat) in variants
720 .iter()
721 .enumerate()
722 .take(cfg.blocked_threshold as usize)
723 {
724 detector.record_search("ctx_search", &format!("fp_{i}"), Some(pat));
725 }
726 let r = detector.record_search("ctx_search", "fp_new", Some("compress_all"));
727 assert_eq!(r.level, ThrottleLevel::Blocked);
728 }
729
730 #[test]
731 fn is_search_tool_detection() {
732 assert!(LoopDetector::is_search_tool("ctx_search"));
733 assert!(LoopDetector::is_search_tool("ctx_semantic_search"));
734 assert!(!LoopDetector::is_search_tool("ctx_read"));
735 assert!(!LoopDetector::is_search_tool("ctx_shell"));
736 }
737
738 #[test]
739 fn is_search_shell_command_detection() {
740 assert!(LoopDetector::is_search_shell_command("grep -r foo ."));
741 assert!(LoopDetector::is_search_shell_command("rg pattern src/"));
742 assert!(LoopDetector::is_search_shell_command("find . -name '*.rs'"));
743 assert!(!LoopDetector::is_search_shell_command("cargo build"));
744 assert!(!LoopDetector::is_search_shell_command("git status"));
745 }
746
747 #[test]
748 fn correction_fresh_reread_detected() {
749 let mut detector = LoopDetector::new();
750 detector.record_read_for_correction("src/main.rs", "full", false);
752 detector.record_read_for_correction("src/lib.rs", "full", false);
753 detector.record_read_for_correction("src/util.rs", "full", false);
754 detector.record_read_for_correction("src/main.rs", "full", false);
756 assert_eq!(detector.correction_count(), 0);
757 detector.record_read_for_correction("src/main.rs", "full", true);
759 assert_eq!(detector.correction_count(), 1);
760 }
761
762 #[test]
763 fn correction_mode_bounce_detected() {
764 let mut detector = LoopDetector::new();
765 for i in 0..COLD_START_CALLS {
767 detector.record_read_for_correction(&format!("f{i}.rs"), "full", false);
768 }
769 detector.record_read_for_correction("src/cache.rs", "map", false);
771 assert_eq!(detector.correction_count(), 0);
772 detector.record_read_for_correction("src/cache.rs", "full", false);
774 assert_eq!(detector.correction_count(), 1);
775 }
776
777 #[test]
778 fn correction_shell_rerun_detected() {
779 let mut detector = LoopDetector::new();
780 for i in 0..COLD_START_CALLS {
782 detector.record_shell_for_correction(&format!("echo {i}"));
783 }
784 detector.record_shell_for_correction("cargo test --lib");
786 assert_eq!(detector.correction_count(), 0);
787 detector.record_shell_for_correction("cargo test --lib");
789 assert_eq!(detector.correction_count(), 1);
790 }
791
792 #[test]
793 fn correction_rate_calculation() {
794 let mut detector = LoopDetector::new();
795 for i in 0..COLD_START_CALLS {
796 detector.record_shell_for_correction(&format!("init{i}"));
797 }
798 detector.record_shell_for_correction("cargo check");
799 detector.record_shell_for_correction("cargo check");
800 detector.record_shell_for_correction("cargo check");
801 assert_eq!(detector.correction_count(), 2);
803 assert!(detector.correction_rate() > 0.0);
804 }
805
806 #[test]
807 fn correction_cold_start_ignored() {
808 let mut detector = LoopDetector::new();
809 detector.record_shell_for_correction("cargo check");
811 detector.record_shell_for_correction("cargo check");
812 detector.record_shell_for_correction("cargo check");
813 assert_eq!(detector.correction_count(), 0);
814 }
815
816 #[test]
817 fn search_block_message_has_guidance_when_blocking_enabled() {
818 let cfg = LoopDetectionConfig {
820 blocked_threshold: 6,
821 search_group_limit: 8,
822 ..Default::default()
823 };
824 let mut detector = LoopDetector::with_config(&cfg);
825 for i in 0..10 {
826 detector.record_search("ctx_search", &format!("fp_{i}"), Some("compress"));
827 }
828 let r = detector.record_search("ctx_search", "fp_new", Some("compress"));
829 assert_eq!(r.level, ThrottleLevel::Blocked);
830 let msg = r.message.unwrap();
831 assert!(msg.contains("ctx_tree"));
832 assert!(msg.contains("path"));
833 assert!(msg.contains("ctx_read"));
834 }
835
836 #[test]
837 fn error_outcome_undoes_pre_dispatch_count() {
838 let cfg = test_config(2, 4, 0);
839 let mut detector = LoopDetector::with_config(&cfg);
840
841 detector.record_call("ctx_read", "fp1");
842 detector.record_call("ctx_read", "fp1");
843 detector.record_error_outcome("ctx_read", "fp1");
844
845 let r = detector.record_call("ctx_read", "fp1");
846 assert_eq!(r.call_count, 2, "error should have undone one count");
847 assert_eq!(r.level, ThrottleLevel::Normal);
848 }
849
850 #[test]
851 fn repeated_errors_dont_trigger_reduced() {
852 let cfg = test_config(2, 4, 0);
853 let mut detector = LoopDetector::with_config(&cfg);
854
855 for _ in 0..5 {
856 detector.record_call("ctx_read", "fp1");
857 detector.record_error_outcome("ctx_read", "fp1");
858 }
859
860 let r = detector.record_call("ctx_read", "fp1");
861 assert_eq!(
862 r.level,
863 ThrottleLevel::Normal,
864 "5 failed retries should not throttle"
865 );
866 }
867
868 #[test]
869 fn mixed_success_and_error_correct_count() {
870 let cfg = test_config(2, 4, 0);
871 let mut detector = LoopDetector::with_config(&cfg);
872
873 detector.record_call("ctx_read", "fp1");
874 detector.record_error_outcome("ctx_read", "fp1");
875 detector.record_call("ctx_read", "fp1");
876 detector.record_error_outcome("ctx_read", "fp1");
877 detector.record_call("ctx_read", "fp1");
878 assert_eq!(detector.record_call("ctx_read", "fp1").call_count, 2);
880 }
881
882 #[test]
883 fn error_outcome_on_nonexistent_key_is_noop() {
884 let mut detector = LoopDetector::new();
885 detector.record_error_outcome("ctx_read", "never_called");
886 let r = detector.record_call("ctx_read", "never_called");
887 assert_eq!(r.call_count, 1);
888 }
889
890 #[test]
891 fn error_outcome_doesnt_go_negative() {
892 let mut detector = LoopDetector::new();
893 detector.record_call("ctx_read", "fp1");
894 detector.record_error_outcome("ctx_read", "fp1");
895 detector.record_error_outcome("ctx_read", "fp1");
896 let r = detector.record_call("ctx_read", "fp1");
897 assert_eq!(r.call_count, 1, "count should never go below 0");
898 }
899
900 #[test]
901 fn error_in_tool_a_doesnt_affect_tool_b() {
902 let cfg = test_config(2, 4, 0);
903 let mut detector = LoopDetector::with_config(&cfg);
904
905 for _ in 0..5 {
906 detector.record_call("ctx_read", "fp1");
907 detector.record_error_outcome("ctx_read", "fp1");
908 }
909
910 let r = detector.record_call("ctx_shell", "fp_shell");
911 assert_eq!(r.call_count, 1);
912 assert_eq!(r.level, ThrottleLevel::Normal);
913 }
914
915 #[test]
916 fn different_fingerprints_independent_after_errors() {
917 let cfg = test_config(2, 4, 0);
918 let mut detector = LoopDetector::with_config(&cfg);
919
920 detector.record_call("ctx_read", "fp_a");
921 detector.record_error_outcome("ctx_read", "fp_a");
922
923 detector.record_call("ctx_read", "fp_b");
924 let r = detector.record_call("ctx_read", "fp_b");
925 assert_eq!(r.call_count, 2);
926
927 let r_a = detector.record_call("ctx_read", "fp_a");
928 assert_eq!(r_a.call_count, 1, "fp_a count should be reset to 0 then +1");
929 }
930
931 #[test]
932 fn correction_degrade_recovery_after_prune() {
933 let mut detector = LoopDetector::new();
934 for i in 0..4u32 {
935 detector.record_read_for_correction(&format!("warmup{i}.rs"), "full", false);
936 }
937 detector.record_read_for_correction("target.rs", "full", false);
938 detector.record_read_for_correction("target.rs", "full", true);
939 assert!(detector.correction_count() > 0);
940 detector.prune_corrections();
941 assert!(detector.correction_count() >= 1);
944 }
945
946 #[test]
947 fn success_after_errors_resets_to_normal() {
948 let cfg = test_config(2, 4, 0);
949 let mut detector = LoopDetector::with_config(&cfg);
950
951 for _ in 0..3 {
952 detector.record_call("ctx_read", "fp1");
953 detector.record_error_outcome("ctx_read", "fp1");
954 }
955
956 let r = detector.record_call("ctx_read", "fp1");
957 assert_eq!(r.level, ThrottleLevel::Normal);
958 assert_eq!(r.call_count, 1);
959 }
960
961 #[test]
962 fn throttle_result_default_is_normal() {
963 let r = ThrottleResult::default();
964 assert_eq!(r.level, ThrottleLevel::Normal);
965 assert_eq!(r.call_count, 0);
966 assert!(r.message.is_none());
967 }
968}