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 && total_count > limit
120 {
121 let msg = if crate::core::protocol::meta_visible() {
122 Some(format!(
123 "Warning: {tool} called {total_count}x total (limit: {limit}). \
124 Consider ctx_compress or narrowing scope."
125 ))
126 } else {
127 None
128 };
129 return ThrottleResult {
130 level: ThrottleLevel::Reduced,
131 call_count: total_count,
132 message: msg,
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 && let Some((prev_time, _)) = self.recent_reads.get(path)
337 && now.duration_since(*prev_time) < CORRECTION_WINDOW
338 {
339 self.correction_signals
340 .push((now, CorrectionKind::FreshReRead));
341 }
342
343 if mode == "full"
344 && let Some((prev_time, prev_mode)) = self.recent_reads.get(path)
345 {
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 self.recent_reads
355 .insert(path.to_string(), (now, mode.to_string()));
356 }
357
358 pub fn record_shell_for_correction(&mut self, command: &str) {
360 self.total_calls += 1;
361 let now = Instant::now();
362
363 if self.total_calls <= COLD_START_CALLS {
364 self.recent_commands.insert(command.to_string(), now);
365 return;
366 }
367
368 let key = normalize_shell_command(command);
369 if let Some(prev_time) = self.recent_commands.get(&key)
370 && now.duration_since(*prev_time) < SHELL_RERUN_WINDOW
371 {
372 self.correction_signals
373 .push((now, CorrectionKind::ShellReRun));
374 }
375 self.recent_commands.insert(key, now);
376 }
377
378 pub fn correction_count(&self) -> u32 {
380 let now = Instant::now();
381 self.correction_signals
382 .iter()
383 .filter(|(t, _)| now.duration_since(*t) < CORRECTION_WINDOW)
384 .count() as u32
385 }
386
387 pub fn correction_rate(&self) -> f64 {
389 let count = self.correction_count();
390 if count == 0 {
391 return 0.0;
392 }
393 let window_mins = CORRECTION_WINDOW.as_secs_f64() / 60.0;
394 f64::from(count) / window_mins
395 }
396
397 pub fn prune_corrections(&mut self) {
399 let now = Instant::now();
400 self.correction_signals
401 .retain(|(t, _)| now.duration_since(*t) < CORRECTION_WINDOW);
402 self.recent_reads
403 .retain(|_, (t, _)| now.duration_since(*t) < CORRECTION_WINDOW);
404 self.recent_commands
405 .retain(|_, t| now.duration_since(*t) < CORRECTION_WINDOW);
406 }
407
408 pub fn reset(&mut self) {
410 self.call_history.clear();
411 self.duplicate_counts.clear();
412 self.search_group_history.clear();
413 self.recent_search_patterns.clear();
414 self.correction_signals.clear();
415 self.recent_reads.clear();
416 self.recent_commands.clear();
417 self.total_calls = 0;
418 }
419
420 fn prune_window(&mut self, now: Instant) {
421 for entries in self.call_history.values_mut() {
422 entries.retain(|t| now.duration_since(*t) < self.window);
423 }
424 self.call_history.retain(|_, v| !v.is_empty());
430 let live = &self.call_history;
431 self.duplicate_counts.retain(|k, _| live.contains_key(k));
432 self.search_group_history
433 .retain(|t| now.duration_since(*t) < self.window);
434 }
435
436 fn count_similar_patterns(&self, new_pattern: &str) -> u32 {
437 let new_lower = new_pattern.to_lowercase();
438 let new_root = extract_alpha_root(&new_lower);
439
440 let mut count = 0u32;
441 for existing in &self.recent_search_patterns {
442 let existing_lower = existing.to_lowercase();
443 if patterns_are_similar(&new_lower, &existing_lower) {
444 count += 1;
445 } else if new_root.len() >= 4 {
446 let existing_root = extract_alpha_root(&existing_lower);
447 if existing_root.len() >= 4
448 && (new_root.starts_with(&existing_root)
449 || existing_root.starts_with(&new_root))
450 {
451 count += 1;
452 }
453 }
454 }
455 count
456 }
457
458 fn block_message(&self, tool: &str, count: u32) -> String {
459 if Self::is_search_tool(tool) {
460 self.search_block_message(count)
461 } else {
462 format!(
463 "LOOP DETECTED: {tool} called {count}x with same/similar args. \
464 Call blocked. Change your approach — the current strategy is not working."
465 )
466 }
467 }
468
469 #[allow(clippy::unused_self)]
470 fn search_block_message(&self, count: u32) -> String {
471 format!(
472 "LOOP DETECTED: You've searched {count}x with similar patterns. STOP searching and change strategy. \
473 1) Use ctx_tree to understand the project structure first. \
474 2) Narrow your search with the 'path' parameter to a specific directory. \
475 3) Use ctx_read with mode='map' to understand a file before searching more."
476 )
477 }
478
479 fn search_group_block_message(&self, count: u32) -> String {
480 format!(
481 "LOOP DETECTED: {count} search calls in {}s — too many. STOP and rethink. \
482 1) Use ctx_tree to map the project structure. \
483 2) Pick ONE specific directory and search there with the 'path' parameter. \
484 3) Read files with ctx_read mode='map' instead of searching blindly.",
485 self.window.as_secs()
486 )
487 }
488}
489
490fn normalize_shell_command(cmd: &str) -> String {
491 cmd.split_whitespace()
492 .take(5)
493 .collect::<Vec<_>>()
494 .join(" ")
495 .to_lowercase()
496}
497
498fn extract_alpha_root(pattern: &str) -> String {
499 pattern
500 .chars()
501 .take_while(|c| c.is_alphanumeric())
502 .collect()
503}
504
505fn patterns_are_similar(a: &str, b: &str) -> bool {
506 if a == b {
507 return true;
508 }
509 if a.contains(b) || b.contains(a) {
510 return true;
511 }
512 let a_alpha: String = a.chars().filter(|c| c.is_alphanumeric()).collect();
513 let b_alpha: String = b.chars().filter(|c| c.is_alphanumeric()).collect();
514 if a_alpha.len() >= 3
515 && b_alpha.len() >= 3
516 && (a_alpha.contains(&b_alpha) || b_alpha.contains(&a_alpha))
517 {
518 return true;
519 }
520 false
521}
522
523fn canonical_json(value: &serde_json::Value) -> String {
524 match value {
525 serde_json::Value::Object(map) => {
526 let mut keys: Vec<&String> = map.keys().collect();
527 keys.sort();
528 let entries: Vec<String> = keys
529 .iter()
530 .map(|k| format!("{}:{}", k, canonical_json(&map[*k])))
531 .collect();
532 format!("{{{}}}", entries.join(","))
533 }
534 serde_json::Value::Array(arr) => {
535 let entries: Vec<String> = arr.iter().map(canonical_json).collect();
536 format!("[{}]", entries.join(","))
537 }
538 _ => value.to_string(),
539 }
540}
541
542#[cfg(test)]
543mod tests {
544 use super::*;
545
546 fn test_config(normal: u32, reduced: u32, blocked: u32) -> LoopDetectionConfig {
547 LoopDetectionConfig {
548 normal_threshold: normal,
549 reduced_threshold: reduced,
550 blocked_threshold: blocked,
551 window_secs: 300,
552 search_group_limit: 10,
553 tool_total_limits: std::collections::HashMap::new(),
554 }
555 }
556
557 #[test]
558 fn normal_calls_pass_through() {
559 let mut detector = LoopDetector::new();
560 let r1 = detector.record_call("ctx_read", "abc123");
561 assert_eq!(r1.level, ThrottleLevel::Normal);
562 assert_eq!(r1.call_count, 1);
563 assert!(r1.message.is_none());
564 }
565
566 #[test]
567 fn repeated_calls_trigger_reduced() {
568 let _lock = crate::core::data_dir::test_env_lock();
569 crate::test_env::set_var("LEAN_CTX_META", "1");
570 let cfg = LoopDetectionConfig::default();
571 let mut detector = LoopDetector::with_config(&cfg);
572 for _ in 0..cfg.normal_threshold {
573 detector.record_call("ctx_read", "same_fp");
574 }
575 let result = detector.record_call("ctx_read", "same_fp");
576 assert_eq!(result.level, ThrottleLevel::Reduced);
577 assert!(result.message.is_some());
578 crate::test_env::remove_var("LEAN_CTX_META");
579 }
580
581 #[test]
582 fn excessive_calls_get_blocked_when_enabled() {
583 let cfg = LoopDetectionConfig {
585 blocked_threshold: 6,
586 ..Default::default()
587 };
588 let mut detector = LoopDetector::with_config(&cfg);
589 for _ in 0..cfg.blocked_threshold {
590 detector.record_call("ctx_shell", "same_fp");
591 }
592 let result = detector.record_call("ctx_shell", "same_fp");
593 assert_eq!(result.level, ThrottleLevel::Blocked);
594 assert!(result.message.unwrap().contains("LOOP DETECTED"));
595 }
596
597 #[test]
598 fn blocking_disabled_by_default() {
599 let cfg = LoopDetectionConfig::default();
601 assert_eq!(cfg.blocked_threshold, 0);
602 let mut detector = LoopDetector::with_config(&cfg);
603 for _ in 0..100 {
605 detector.record_call("ctx_shell", "same_fp");
606 }
607 let result = detector.record_call("ctx_shell", "same_fp");
608 assert_ne!(result.level, ThrottleLevel::Blocked);
610 }
611
612 #[test]
613 fn different_args_tracked_separately() {
614 let mut detector = LoopDetector::new();
615 for _ in 0..10 {
616 detector.record_call("ctx_read", "fp_a");
617 }
618 let result = detector.record_call("ctx_read", "fp_b");
619 assert_eq!(result.level, ThrottleLevel::Normal);
620 assert_eq!(result.call_count, 1);
621 }
622
623 #[test]
624 fn fingerprint_deterministic() {
625 let args = serde_json::json!({"path": "test.rs", "mode": "full"});
626 let fp1 = LoopDetector::fingerprint(&args);
627 let fp2 = LoopDetector::fingerprint(&args);
628 assert_eq!(fp1, fp2);
629 }
630
631 #[test]
632 fn fingerprint_order_independent() {
633 let a = serde_json::json!({"mode": "full", "path": "test.rs"});
634 let b = serde_json::json!({"path": "test.rs", "mode": "full"});
635 assert_eq!(LoopDetector::fingerprint(&a), LoopDetector::fingerprint(&b));
636 }
637
638 #[test]
639 fn stats_shows_duplicates() {
640 let mut detector = LoopDetector::new();
641 for _ in 0..5 {
642 detector.record_call("ctx_read", "fp_a");
643 }
644 detector.record_call("ctx_shell", "fp_b");
645 let stats = detector.stats();
646 assert_eq!(stats.len(), 1);
647 assert_eq!(stats[0].1, 5);
648 }
649
650 #[test]
651 fn reset_clears_state() {
652 let mut detector = LoopDetector::new();
653 for _ in 0..5 {
654 detector.record_call("ctx_read", "fp_a");
655 }
656 detector.reset();
657 let result = detector.record_call("ctx_read", "fp_a");
658 assert_eq!(result.call_count, 1);
659 }
660
661 #[test]
662 fn custom_thresholds_from_config() {
663 let cfg = test_config(1, 2, 3);
664 let mut detector = LoopDetector::with_config(&cfg);
665 detector.record_call("ctx_read", "fp");
666 let r = detector.record_call("ctx_read", "fp");
667 assert_eq!(r.level, ThrottleLevel::Reduced);
668 detector.record_call("ctx_read", "fp");
669 let r = detector.record_call("ctx_read", "fp");
670 assert_eq!(r.level, ThrottleLevel::Blocked);
671 }
672
673 #[test]
674 fn similar_patterns_detected() {
675 assert!(patterns_are_similar("compress", "compress"));
676 assert!(patterns_are_similar("compress", "compression"));
677 assert!(patterns_are_similar("compress.*data", "compress"));
678 assert!(!patterns_are_similar("foo", "bar"));
679 assert!(!patterns_are_similar("ab", "cd"));
680 }
681
682 #[test]
683 fn search_group_tracking_when_blocking_enabled() {
684 let cfg = LoopDetectionConfig {
686 search_group_limit: 5,
687 blocked_threshold: 6, ..Default::default()
689 };
690 let mut detector = LoopDetector::with_config(&cfg);
691 for i in 0..5 {
692 let fp = format!("fp_{i}");
693 let r = detector.record_search("ctx_search", &fp, Some(&format!("pattern_{i}")));
694 assert_ne!(r.level, ThrottleLevel::Blocked, "call {i} should not block");
695 }
696 let r = detector.record_search("ctx_search", "fp_5", Some("pattern_5"));
697 assert_eq!(r.level, ThrottleLevel::Blocked);
698 assert!(r.message.unwrap().contains("search calls"));
699 }
700
701 #[test]
702 fn similar_search_patterns_trigger_block_when_enabled() {
703 let cfg = LoopDetectionConfig {
705 blocked_threshold: 6,
706 ..Default::default()
707 };
708 let mut detector = LoopDetector::with_config(&cfg);
709 let variants = [
710 "compress",
711 "compression",
712 "compress.*data",
713 "compress_output",
714 "compressor",
715 "compress_result",
716 "compress_file",
717 ];
718 for (i, pat) in variants
719 .iter()
720 .enumerate()
721 .take(cfg.blocked_threshold as usize)
722 {
723 detector.record_search("ctx_search", &format!("fp_{i}"), Some(pat));
724 }
725 let r = detector.record_search("ctx_search", "fp_new", Some("compress_all"));
726 assert_eq!(r.level, ThrottleLevel::Blocked);
727 }
728
729 #[test]
730 fn is_search_tool_detection() {
731 assert!(LoopDetector::is_search_tool("ctx_search"));
732 assert!(LoopDetector::is_search_tool("ctx_semantic_search"));
733 assert!(!LoopDetector::is_search_tool("ctx_read"));
734 assert!(!LoopDetector::is_search_tool("ctx_shell"));
735 }
736
737 #[test]
738 fn is_search_shell_command_detection() {
739 assert!(LoopDetector::is_search_shell_command("grep -r foo ."));
740 assert!(LoopDetector::is_search_shell_command("rg pattern src/"));
741 assert!(LoopDetector::is_search_shell_command("find . -name '*.rs'"));
742 assert!(!LoopDetector::is_search_shell_command("cargo build"));
743 assert!(!LoopDetector::is_search_shell_command("git status"));
744 }
745
746 #[test]
747 fn correction_fresh_reread_detected() {
748 let mut detector = LoopDetector::new();
749 detector.record_read_for_correction("src/main.rs", "full", false);
751 detector.record_read_for_correction("src/lib.rs", "full", false);
752 detector.record_read_for_correction("src/util.rs", "full", false);
753 detector.record_read_for_correction("src/main.rs", "full", false);
755 assert_eq!(detector.correction_count(), 0);
756 detector.record_read_for_correction("src/main.rs", "full", true);
758 assert_eq!(detector.correction_count(), 1);
759 }
760
761 #[test]
762 fn correction_mode_bounce_detected() {
763 let mut detector = LoopDetector::new();
764 for i in 0..COLD_START_CALLS {
766 detector.record_read_for_correction(&format!("f{i}.rs"), "full", false);
767 }
768 detector.record_read_for_correction("src/cache.rs", "map", false);
770 assert_eq!(detector.correction_count(), 0);
771 detector.record_read_for_correction("src/cache.rs", "full", false);
773 assert_eq!(detector.correction_count(), 1);
774 }
775
776 #[test]
777 fn correction_shell_rerun_detected() {
778 let mut detector = LoopDetector::new();
779 for i in 0..COLD_START_CALLS {
781 detector.record_shell_for_correction(&format!("echo {i}"));
782 }
783 detector.record_shell_for_correction("cargo test --lib");
785 assert_eq!(detector.correction_count(), 0);
786 detector.record_shell_for_correction("cargo test --lib");
788 assert_eq!(detector.correction_count(), 1);
789 }
790
791 #[test]
792 fn correction_rate_calculation() {
793 let mut detector = LoopDetector::new();
794 for i in 0..COLD_START_CALLS {
795 detector.record_shell_for_correction(&format!("init{i}"));
796 }
797 detector.record_shell_for_correction("cargo check");
798 detector.record_shell_for_correction("cargo check");
799 detector.record_shell_for_correction("cargo check");
800 assert_eq!(detector.correction_count(), 2);
802 assert!(detector.correction_rate() > 0.0);
803 }
804
805 #[test]
806 fn correction_cold_start_ignored() {
807 let mut detector = LoopDetector::new();
808 detector.record_shell_for_correction("cargo check");
810 detector.record_shell_for_correction("cargo check");
811 detector.record_shell_for_correction("cargo check");
812 assert_eq!(detector.correction_count(), 0);
813 }
814
815 #[test]
816 fn search_block_message_has_guidance_when_blocking_enabled() {
817 let cfg = LoopDetectionConfig {
819 blocked_threshold: 6,
820 search_group_limit: 8,
821 ..Default::default()
822 };
823 let mut detector = LoopDetector::with_config(&cfg);
824 for i in 0..10 {
825 detector.record_search("ctx_search", &format!("fp_{i}"), Some("compress"));
826 }
827 let r = detector.record_search("ctx_search", "fp_new", Some("compress"));
828 assert_eq!(r.level, ThrottleLevel::Blocked);
829 let msg = r.message.unwrap();
830 assert!(msg.contains("ctx_tree"));
831 assert!(msg.contains("path"));
832 assert!(msg.contains("ctx_read"));
833 }
834
835 #[test]
836 fn error_outcome_undoes_pre_dispatch_count() {
837 let cfg = test_config(2, 4, 0);
838 let mut detector = LoopDetector::with_config(&cfg);
839
840 detector.record_call("ctx_read", "fp1");
841 detector.record_call("ctx_read", "fp1");
842 detector.record_error_outcome("ctx_read", "fp1");
843
844 let r = detector.record_call("ctx_read", "fp1");
845 assert_eq!(r.call_count, 2, "error should have undone one count");
846 assert_eq!(r.level, ThrottleLevel::Normal);
847 }
848
849 #[test]
850 fn repeated_errors_dont_trigger_reduced() {
851 let cfg = test_config(2, 4, 0);
852 let mut detector = LoopDetector::with_config(&cfg);
853
854 for _ in 0..5 {
855 detector.record_call("ctx_read", "fp1");
856 detector.record_error_outcome("ctx_read", "fp1");
857 }
858
859 let r = detector.record_call("ctx_read", "fp1");
860 assert_eq!(
861 r.level,
862 ThrottleLevel::Normal,
863 "5 failed retries should not throttle"
864 );
865 }
866
867 #[test]
868 fn mixed_success_and_error_correct_count() {
869 let cfg = test_config(2, 4, 0);
870 let mut detector = LoopDetector::with_config(&cfg);
871
872 detector.record_call("ctx_read", "fp1");
873 detector.record_error_outcome("ctx_read", "fp1");
874 detector.record_call("ctx_read", "fp1");
875 detector.record_error_outcome("ctx_read", "fp1");
876 detector.record_call("ctx_read", "fp1");
877 assert_eq!(detector.record_call("ctx_read", "fp1").call_count, 2);
879 }
880
881 #[test]
882 fn error_outcome_on_nonexistent_key_is_noop() {
883 let mut detector = LoopDetector::new();
884 detector.record_error_outcome("ctx_read", "never_called");
885 let r = detector.record_call("ctx_read", "never_called");
886 assert_eq!(r.call_count, 1);
887 }
888
889 #[test]
890 fn error_outcome_doesnt_go_negative() {
891 let mut detector = LoopDetector::new();
892 detector.record_call("ctx_read", "fp1");
893 detector.record_error_outcome("ctx_read", "fp1");
894 detector.record_error_outcome("ctx_read", "fp1");
895 let r = detector.record_call("ctx_read", "fp1");
896 assert_eq!(r.call_count, 1, "count should never go below 0");
897 }
898
899 #[test]
900 fn error_in_tool_a_doesnt_affect_tool_b() {
901 let cfg = test_config(2, 4, 0);
902 let mut detector = LoopDetector::with_config(&cfg);
903
904 for _ in 0..5 {
905 detector.record_call("ctx_read", "fp1");
906 detector.record_error_outcome("ctx_read", "fp1");
907 }
908
909 let r = detector.record_call("ctx_shell", "fp_shell");
910 assert_eq!(r.call_count, 1);
911 assert_eq!(r.level, ThrottleLevel::Normal);
912 }
913
914 #[test]
915 fn different_fingerprints_independent_after_errors() {
916 let cfg = test_config(2, 4, 0);
917 let mut detector = LoopDetector::with_config(&cfg);
918
919 detector.record_call("ctx_read", "fp_a");
920 detector.record_error_outcome("ctx_read", "fp_a");
921
922 detector.record_call("ctx_read", "fp_b");
923 let r = detector.record_call("ctx_read", "fp_b");
924 assert_eq!(r.call_count, 2);
925
926 let r_a = detector.record_call("ctx_read", "fp_a");
927 assert_eq!(r_a.call_count, 1, "fp_a count should be reset to 0 then +1");
928 }
929
930 #[test]
931 fn correction_degrade_recovery_after_prune() {
932 let mut detector = LoopDetector::new();
933 for i in 0..4u32 {
934 detector.record_read_for_correction(&format!("warmup{i}.rs"), "full", false);
935 }
936 detector.record_read_for_correction("target.rs", "full", false);
937 detector.record_read_for_correction("target.rs", "full", true);
938 assert!(detector.correction_count() > 0);
939 detector.prune_corrections();
940 assert!(detector.correction_count() >= 1);
943 }
944
945 #[test]
946 fn success_after_errors_resets_to_normal() {
947 let cfg = test_config(2, 4, 0);
948 let mut detector = LoopDetector::with_config(&cfg);
949
950 for _ in 0..3 {
951 detector.record_call("ctx_read", "fp1");
952 detector.record_error_outcome("ctx_read", "fp1");
953 }
954
955 let r = detector.record_call("ctx_read", "fp1");
956 assert_eq!(r.level, ThrottleLevel::Normal);
957 assert_eq!(r.call_count, 1);
958 }
959
960 #[test]
961 fn throttle_result_default_is_normal() {
962 let r = ThrottleResult::default();
963 assert_eq!(r.level, ThrottleLevel::Normal);
964 assert_eq!(r.call_count, 0);
965 assert!(r.message.is_none());
966 }
967}