Skip to main content

lean_ctx/core/
loop_detection.rs

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/// Classification of why an agent re-requested data it already had.
16#[derive(Debug, Clone, PartialEq)]
17pub enum CorrectionKind {
18    FreshReRead,
19    ShellReRun,
20    ModeBounce,
21}
22
23/// Tracks repeated tool calls within a time window to detect and throttle agent loops.
24#[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-loop tracking (Fix A)
38    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/// Severity of throttling applied to a repeated call: normal, reduced, or blocked.
45#[derive(Debug, Clone, PartialEq)]
46pub enum ThrottleLevel {
47    Normal,
48    Reduced,
49    Blocked,
50}
51
52/// Outcome of a loop detection check: throttle level, count, and optional warning.
53#[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    /// Creates a loop detector with default thresholds.
78    pub fn new() -> Self {
79        Self::with_config(&LoopDetectionConfig::default())
80    }
81
82    /// Creates a loop detector with custom thresholds from config.
83    /// Set blocked_threshold to 0 to disable blocking entirely (LeanCTX philosophy).
84    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    /// Records a tool call and returns the throttle result based on repetition count.
109    pub fn record_call(&mut self, tool: &str, args_fingerprint: &str) -> ThrottleResult {
110        let now = Instant::now();
111        self.prune_window(now);
112
113        // Per-tool total count (regardless of args)
114        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    /// Undo the pre-dispatch count for a call that resulted in an error.
190    /// Prevents failed retries from triggering throttling prematurely.
191    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    /// Record a search-category call and check the cross-tool search group limit.
201    /// `search_pattern` is the extracted query/regex the agent is looking for (if available).
202    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        // blocked_threshold == 0 means blocking is disabled (LeanCTX default)
227        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        // search_group_limit == u32::MAX when blocking is disabled
236        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    /// Returns `true` if the tool name is a known search tool (ctx_search, etc.).
289    pub fn is_search_tool(tool: &str) -> bool {
290        SEARCH_TOOLS.contains(&tool)
291    }
292
293    /// Returns `true` if the shell command starts with a search tool (grep, rg, find, etc.).
294    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    /// Computes a deterministic hash fingerprint of JSON tool arguments.
300    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    /// Returns duplicate call entries sorted by count (descending), filtered to count > 1.
311    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    /// Records a ctx_read call and detects correction signals:
323    /// - `fresh=true` re-read of a previously cached file
324    /// - Mode bounce: map/signatures followed by full within 30s
325    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    /// Records a ctx_shell command and detects re-runs of the same command within 60s.
360    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    /// Returns the number of correction signals in the sliding window.
380    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    /// Returns the correction rate: signals per minute within the window.
389    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    /// Prunes expired correction signals and stale read/command entries.
399    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    /// Clears all tracking state (call history, search patterns, counters).
410    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        // Drop keys whose window emptied, plus their orphaned duplicate counts, so the
426        // per-fingerprint maps don't grow unbounded over a long session. Behavior-neutral:
427        // empty Vecs contribute 0 to record_call's count, and duplicate_counts is only
428        // read by stats() (already filtered to count > 1). tool_total_counts is left
429        // intact — it is a cumulative per-tool-name guard (bounded by the tool set).
430        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        // Blocking must be explicitly enabled (blocked_threshold > 0)
585        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        // Default config has blocked_threshold = 0, so blocking never happens
601        let cfg = LoopDetectionConfig::default();
602        assert_eq!(cfg.blocked_threshold, 0);
603        let mut detector = LoopDetector::with_config(&cfg);
604        // Even 100 calls should not block when blocking is disabled
605        for _ in 0..100 {
606            detector.record_call("ctx_shell", "same_fp");
607        }
608        let result = detector.record_call("ctx_shell", "same_fp");
609        // Should be Reduced (warning) but never Blocked
610        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        // Blocking must be explicitly enabled for search group limits to block
686        let cfg = LoopDetectionConfig {
687            search_group_limit: 5,
688            blocked_threshold: 6, // Enable blocking
689            ..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        // Blocking must be explicitly enabled
705        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        // First read (cold start period, skipped)
751        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        // 4th call: past cold start
755        detector.record_read_for_correction("src/main.rs", "full", false);
756        assert_eq!(detector.correction_count(), 0);
757        // fresh=true re-read of previously read file = correction signal
758        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        // Cold start
766        for i in 0..COLD_START_CALLS {
767            detector.record_read_for_correction(&format!("f{i}.rs"), "full", false);
768        }
769        // Read with map mode
770        detector.record_read_for_correction("src/cache.rs", "map", false);
771        assert_eq!(detector.correction_count(), 0);
772        // Immediately bounce to full mode = correction
773        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        // Cold start
781        for i in 0..COLD_START_CALLS {
782            detector.record_shell_for_correction(&format!("echo {i}"));
783        }
784        // First run
785        detector.record_shell_for_correction("cargo test --lib");
786        assert_eq!(detector.correction_count(), 0);
787        // Same command again within 60s = correction
788        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        // 2 corrections (first run doesn't count)
802        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        // During cold start, same-command re-runs are not counted
810        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        // Blocking must be explicitly enabled to get block messages
819        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        // 3 pre-dispatch, 2 error undos -> effective count = 1
879        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        // After prune, count is still > 0 because signal is within window
942        // but the mechanism works: once window expires, count drops
943        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}