tool_loop_guard/

lib.rs

//! # tool-loop-guard
//!
//! Detect when an LLM agent gets stuck calling the same tool with the same
//! args.
//!
//! When an LLM agent gets confused, the most common failure mode is calling
//! the same tool with the same args over and over. A circuit breaker won't
//! catch this (no errors), a deadline won't catch it (single calls are fast),
//! but the wall clock keeps ticking and the bill keeps growing.
//!
//! [`LoopGuard`] watches a sliding window of recent tool calls and returns
//! [`LoopDetectedError`] the moment a `(tool_name, args)` pair repeats more
//! than `threshold` times within the last `window` calls.
//!
//! ```rust
//! use serde_json::json;
//! use tool_loop_guard::{LoopGuard, LoopDetectedError};
//!
//! let mut guard = LoopGuard::with_capacity(8, 3);
//!
//! let args = json!({"q": "anthropic prompt cache"});
//! guard.record("search", &args).unwrap();
//! guard.record("search", &args).unwrap();
//! match guard.record("search", &args) {
//!     Ok(()) => unreachable!(),
//!     Err(LoopDetectedError { tool_name, count, window, .. }) => {
//!         assert_eq!(tool_name, "search");
//!         assert_eq!(count, 3);
//!         assert_eq!(window, 8);
//!     }
//! }
//! ```
//!
//! The default key is `(tool_name, canonical_json(args))`. For more exotic
//! matching (for example, ignore a `request_id` field in args) pass a custom
//! `key_fn` via [`LoopGuard::with_key_fn`].
//!
//! ## Why not just `serde_json::to_string`?
//!
//! `serde_json::to_string` does not sort object keys, so
//! `{"a":1,"b":2}` and `{"b":2,"a":1}` would hash differently. This crate
//! re-walks the [`serde_json::Value`] and emits a stable canonical form
//! (sorted object keys, compact separators) so semantically identical
//! argument blobs collide as expected. That matches Python's
//! `json.dumps(args, sort_keys=True, separators=(",", ":"))`.

#![deny(missing_docs)]

use serde_json::Value;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt::{self, Write as _};

/// A canonical (tool_name, args) key used by the default key function.
///
/// First element is the tool name. Second element is the canonical JSON
/// representation of the args (sorted object keys, compact separators).
pub type Key = (String, String);

/// A user-supplied key function.
///
/// Given a tool name and its args, return a [`Key`] that two semantically
/// identical calls should produce. The default implementation is
/// [`default_key_fn`]. Override it to ignore noisy fields (`request_id`,
/// `timestamp`, etc.) before comparison.
pub type KeyFn = Box<dyn Fn(&str, &Value) -> Key + Send + Sync + 'static>;

/// Compute the canonical default key for a `(tool_name, args)` pair.
///
/// `args` is serialized to JSON with sorted object keys and compact
/// separators so that `{"a":1,"b":2}` and `{"b":2,"a":1}` collide. A
/// [`Value::Null`] is treated as an empty object (`"{}"`) to match the
/// Python library's behavior, where `None` args are also normalized to an
/// empty dict.
pub fn default_key_fn(tool_name: &str, args: &Value) -> Key {
    let canon = if args.is_null() {
        "{}".to_string()
    } else {
        canonical_json(args)
    };
    (tool_name.to_string(), canon)
}

/// Write `value` to `out` in canonical form (sorted keys, compact).
fn write_canonical(out: &mut String, value: &Value) {
    match value {
        Value::Null => out.push_str("null"),
        Value::Bool(b) => out.push_str(if *b { "true" } else { "false" }),
        Value::Number(n) => {
            // serde_json::Number's Display already produces a stable form.
            let _ = write!(out, "{n}");
        }
        Value::String(s) => {
            // Use serde_json to handle escapes correctly.
            let encoded = serde_json::to_string(s).unwrap_or_else(|_| String::from("\"\""));
            out.push_str(&encoded);
        }
        Value::Array(items) => {
            out.push('[');
            for (i, item) in items.iter().enumerate() {
                if i > 0 {
                    out.push(',');
                }
                write_canonical(out, item);
            }
            out.push(']');
        }
        Value::Object(map) => {
            // Sort keys for canonical form.
            let mut keys: Vec<&String> = map.keys().collect();
            keys.sort();
            out.push('{');
            for (i, k) in keys.iter().enumerate() {
                if i > 0 {
                    out.push(',');
                }
                let encoded_key =
                    serde_json::to_string(k.as_str()).unwrap_or_else(|_| String::from("\"\""));
                out.push_str(&encoded_key);
                out.push(':');
                write_canonical(out, &map[*k]);
            }
            out.push('}');
        }
    }
}

fn canonical_json(value: &Value) -> String {
    let mut out = String::new();
    write_canonical(&mut out, value);
    out
}

/// Raised when the same `(tool_name, args)` pair fires too often.
///
/// Carries the full context needed to surface a useful diagnostic to the
/// caller: which tool was looping, what args it was called with on the
/// offending call, how many times the key has been seen inside the window,
/// and the guard's configured `window` / `threshold`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LoopDetectedError {
    /// Name of the offending tool.
    pub tool_name: String,
    /// The args object that was matched (last recorded form).
    pub tool_args: Value,
    /// How many times the key was seen inside the window.
    pub count: usize,
    /// Configured window size.
    pub window: usize,
    /// Configured threshold.
    pub threshold: usize,
}

impl fmt::Display for LoopDetectedError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "tool '{}' called {} times in the last {} calls (threshold={})",
            self.tool_name, self.count, self.window, self.threshold
        )
    }
}

impl Error for LoopDetectedError {}

#[derive(Clone, Debug)]
struct Recent {
    key: Key,
    // Kept for symmetry with the Python implementation; not currently read
    // back, but cheap and aids future Debug/inspection use.
    #[allow(dead_code)]
    args: Value,
}

/// Sliding-window detector for repeated tool calls.
///
/// The guard keeps the last `window` `(tool_name, args)` records in a
/// FIFO ring. On every [`LoopGuard::record`], if the new key already appears
/// `threshold` times (including the call being recorded) in the buffer,
/// [`LoopDetectedError`] is returned. The new call is *still* recorded so
/// repeated calls to `record` after a detection stay deterministic.
///
/// Use [`LoopGuard::would_raise`] to peek at the next decision without
/// mutating the buffer.
pub struct LoopGuard {
    window: usize,
    threshold: usize,
    key_fn: KeyFn,
    buf: VecDeque<Recent>,
}

impl LoopGuard {
    /// Construct a guard.
    ///
    /// # Panics
    /// Panics on invalid construction (matches the spirit of Python's
    /// `ValueError`, but in Rust idiom we treat invalid construction as a
    /// programming bug rather than a recoverable error):
    /// * `window < 2`
    /// * `threshold < 2`
    /// * `threshold > window`
    ///
    /// If you want a non-panicking constructor, see [`LoopGuard::try_new`].
    pub fn with_capacity(window: usize, threshold: usize) -> Self {
        Self::try_new(window, threshold).expect("invalid LoopGuard configuration")
    }

    /// Fallible constructor. Returns `Err(ConfigError)` instead of panicking.
    pub fn try_new(window: usize, threshold: usize) -> Result<Self, ConfigError> {
        if window < 2 {
            return Err(ConfigError::WindowTooSmall { window });
        }
        if threshold < 2 {
            return Err(ConfigError::ThresholdTooSmall { threshold });
        }
        if threshold > window {
            return Err(ConfigError::ThresholdExceedsWindow { window, threshold });
        }
        Ok(Self {
            window,
            threshold,
            key_fn: Box::new(default_key_fn),
            buf: VecDeque::with_capacity(window),
        })
    }

    /// Construct a guard with a custom key function.
    ///
    /// The key function decides which calls should be considered "the
    /// same". Two calls that produce equal [`Key`]s count toward the
    /// threshold. Use this to ignore noisy fields like `request_id` or
    /// `timestamp` that change every call but don't represent real intent.
    pub fn with_key_fn<F>(window: usize, threshold: usize, key_fn: F) -> Self
    where
        F: Fn(&str, &Value) -> Key + Send + Sync + 'static,
    {
        let mut g = Self::with_capacity(window, threshold);
        g.key_fn = Box::new(key_fn);
        g
    }

    /// Configured window size.
    pub fn window(&self) -> usize {
        self.window
    }

    /// Configured threshold.
    pub fn threshold(&self) -> usize {
        self.threshold
    }

    /// Number of calls currently in the buffer.
    #[allow(clippy::len_without_is_empty)]
    pub fn len(&self) -> usize {
        self.buf.len()
    }

    /// True if no calls have been recorded since construction or last reset.
    pub fn is_empty(&self) -> bool {
        self.buf.is_empty()
    }

    /// Recent keys, oldest first.
    ///
    /// Returned as borrowed references so callers do not pay an allocation
    /// cost just to inspect the ring. Clone individual entries with
    /// `key.clone()` if you need to keep them past the next `record` call.
    pub fn recent_keys(&self) -> Vec<&Key> {
        self.buf.iter().map(|r| &r.key).collect()
    }

    /// Record a call. Returns `Err(LoopDetectedError)` if the threshold is met.
    ///
    /// The call is recorded into the ring before the count is checked, so
    /// subsequent calls observe the same offending entry until evicted.
    pub fn record(&mut self, tool_name: &str, args: &Value) -> Result<(), LoopDetectedError> {
        let key = (self.key_fn)(tool_name, args);
        if self.buf.len() == self.window {
            self.buf.pop_front();
        }
        self.buf.push_back(Recent {
            key: key.clone(),
            args: args.clone(),
        });
        let count = self.buf.iter().filter(|r| r.key == key).count();
        if count >= self.threshold {
            return Err(LoopDetectedError {
                tool_name: tool_name.to_string(),
                tool_args: args.clone(),
                count,
                window: self.window,
                threshold: self.threshold,
            });
        }
        Ok(())
    }

    /// Return `true` if the *next* `record` for this key would raise.
    ///
    /// Lets callers preview the outcome without mutating the buffer. The
    /// simulation accounts for buffer eviction: if the buffer is already
    /// full and the oldest entry has the same key, the projected count
    /// loses one before gaining one.
    pub fn would_raise(&self, tool_name: &str, args: &Value) -> bool {
        let key = (self.key_fn)(tool_name, args);
        // Simulate appending without overflow side effects.
        let skip_oldest = self.buf.len() == self.window;
        let projected = self
            .buf
            .iter()
            .enumerate()
            .filter(|(i, _)| !(skip_oldest && *i == 0))
            .filter(|(_, r)| r.key == key)
            .count()
            + 1;
        projected >= self.threshold
    }

    /// Forget all recorded calls.
    pub fn reset(&mut self) {
        self.buf.clear();
    }
}

impl fmt::Debug for LoopGuard {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LoopGuard")
            .field("window", &self.window)
            .field("threshold", &self.threshold)
            .field("len", &self.buf.len())
            .finish()
    }
}

/// Errors from invalid [`LoopGuard`] construction.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConfigError {
    /// `window` was less than 2.
    WindowTooSmall {
        /// The offending value.
        window: usize,
    },
    /// `threshold` was less than 2.
    ThresholdTooSmall {
        /// The offending value.
        threshold: usize,
    },
    /// `threshold` was greater than `window`.
    ThresholdExceedsWindow {
        /// Configured window size.
        window: usize,
        /// Configured threshold.
        threshold: usize,
    },
}

impl fmt::Display for ConfigError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ConfigError::WindowTooSmall { window } => {
                write!(f, "window must be >= 2 (got {window})")
            }
            ConfigError::ThresholdTooSmall { threshold } => {
                write!(f, "threshold must be >= 2 (got {threshold})")
            }
            ConfigError::ThresholdExceedsWindow { window, threshold } => write!(
                f,
                "threshold must be <= window (got threshold={threshold}, window={window})"
            ),
        }
    }
}

impl Error for ConfigError {}