Skip to main content

tool_loop_break/
lib.rs

1//! # tool-loop-break
2//!
3//! Detect repeated agent tool invocations and signal "loop" before the
4//! agent burns the budget.
5//!
6//! Approach: maintain a sliding window of the last `window` recorded
7//! `(tool_name, args_fingerprint)` tuples. When any tuple has appeared
8//! `>= threshold` times inside the window, [`record`] returns `true`.
9//!
10//! `args_fingerprint` is whatever you decide — typically a short hash
11//! of the JSON arguments — so identical arg sets collide.
12//!
13//! ## Example
14//!
15//! ```
16//! use tool_loop_break::LoopDetector;
17//! let mut d = LoopDetector::new(10, 3); // window=10, threshold=3
18//! assert!(!d.record("read_file", "abc"));
19//! assert!(!d.record("read_file", "abc"));
20//! assert!(d.record("read_file", "abc")); // third within window -> loop
21//! ```
22
23#![deny(missing_docs)]
24
25use std::collections::{HashMap, VecDeque};
26
27/// Sliding-window loop detector.
28#[derive(Debug)]
29pub struct LoopDetector {
30    window: usize,
31    threshold: u32,
32    buf: VecDeque<(String, String)>,
33    counts: HashMap<(String, String), u32>,
34}
35
36impl LoopDetector {
37    /// Build a detector with the given sliding window size and trip
38    /// threshold.
39    pub fn new(window: usize, threshold: u32) -> Self {
40        Self {
41            window,
42            threshold,
43            buf: VecDeque::with_capacity(window),
44            counts: HashMap::new(),
45        }
46    }
47
48    /// Record one invocation. Returns true when this invocation pushed
49    /// `(tool, args_fingerprint)` to or past the trip threshold.
50    pub fn record(&mut self, tool: &str, args_fingerprint: &str) -> bool {
51        let key = (tool.to_string(), args_fingerprint.to_string());
52        // Evict the oldest if window full.
53        if self.buf.len() == self.window {
54            if let Some(old) = self.buf.pop_front() {
55                if let Some(c) = self.counts.get_mut(&old) {
56                    *c -= 1;
57                    if *c == 0 {
58                        self.counts.remove(&old);
59                    }
60                }
61            }
62        }
63        self.buf.push_back(key.clone());
64        let entry = self.counts.entry(key).or_insert(0);
65        *entry += 1;
66        *entry >= self.threshold
67    }
68
69    /// True when any tuple is currently at or past the threshold.
70    pub fn is_looping(&self) -> bool {
71        self.counts.values().any(|&v| v >= self.threshold)
72    }
73
74    /// Reset the detector.
75    pub fn reset(&mut self) {
76        self.buf.clear();
77        self.counts.clear();
78    }
79}