use dashmap::DashMap;
use std::time::{Duration, Instant};
const DEFAULT_WINDOW_SECS: u64 = 30;
const DEFAULT_MAX_REPEATS: usize = 5;
const MAX_SAMPLES_PER_SESSION: usize = 256;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LoopDecision {
Allow,
Break { repeats: usize },
}
pub struct LoopBreaker {
sessions: DashMap<String, Vec<(u64, Instant)>>,
window: Duration,
max_repeats: usize,
}
impl Default for LoopBreaker {
fn default() -> Self {
Self::new()
}
}
impl LoopBreaker {
pub fn new() -> Self {
let window_secs = std::env::var("LOOP_WINDOW_SECS")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(DEFAULT_WINDOW_SECS);
let max_repeats = std::env::var("LOOP_MAX_REPEATS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(DEFAULT_MAX_REPEATS);
Self::with_config(window_secs, max_repeats)
}
pub fn with_config(window_secs: u64, max_repeats: usize) -> Self {
Self {
sessions: DashMap::new(),
window: Duration::from_secs(window_secs),
max_repeats,
}
}
pub fn enabled(&self) -> bool {
self.max_repeats > 0
}
pub fn check(&self, session_id: &str, intent_fingerprint: u64) -> LoopDecision {
if !self.enabled() {
return LoopDecision::Allow;
}
let now = Instant::now();
let mut samples = self.sessions.entry(session_id.to_string()).or_default();
samples.retain(|(_, t)| now.duration_since(*t) < self.window);
let prior_repeats = samples
.iter()
.filter(|(h, _)| *h == intent_fingerprint)
.count();
samples.push((intent_fingerprint, now));
if samples.len() > MAX_SAMPLES_PER_SESSION {
let drain_to = samples.len() - MAX_SAMPLES_PER_SESSION;
samples.drain(0..drain_to);
}
if prior_repeats >= self.max_repeats {
LoopDecision::Break {
repeats: prior_repeats + 1,
}
} else {
LoopDecision::Allow
}
}
pub fn tracked_sessions(&self) -> usize {
self.sessions.len()
}
}
pub fn intent_fingerprint(path: &str, body: &[u8]) -> u64 {
let intent = crate::engine::parser::extract_intent(body)
.and_then(|v| serde_json::to_string(&v).ok())
.unwrap_or_else(|| String::from_utf8_lossy(body).to_string());
fingerprint_from_basis(path, &intent)
}
pub fn intent_fingerprint_value(path: &str, json: &serde_json::Value) -> u64 {
let intent = crate::engine::parser::extract_intent_value(json)
.and_then(|v| serde_json::to_string(&v).ok())
.unwrap_or_else(|| json.to_string());
fingerprint_from_basis(path, &intent)
}
fn fingerprint_from_basis(path: &str, intent: &str) -> u64 {
let basis = format!("{}\n{}", path, intent);
let hash = blake3::hash(basis.as_bytes());
let bytes = hash.as_bytes();
u64::from_le_bytes(bytes[0..8].try_into().unwrap())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_breaks_after_threshold() {
let lb = LoopBreaker::with_config(60, 3);
let fp = 42;
assert_eq!(lb.check("s1", fp), LoopDecision::Allow); assert_eq!(lb.check("s1", fp), LoopDecision::Allow); assert_eq!(lb.check("s1", fp), LoopDecision::Allow); assert_eq!(lb.check("s1", fp), LoopDecision::Break { repeats: 4 }); }
#[test]
fn test_distinct_intents_do_not_trip() {
let lb = LoopBreaker::with_config(60, 3);
for i in 0..10 {
assert_eq!(lb.check("s1", i as u64), LoopDecision::Allow);
}
}
#[test]
fn test_sessions_isolated() {
let lb = LoopBreaker::with_config(60, 1);
assert_eq!(lb.check("a", 7), LoopDecision::Allow);
assert_eq!(lb.check("b", 7), LoopDecision::Allow); assert_eq!(lb.check("a", 7), LoopDecision::Break { repeats: 2 });
}
#[test]
fn test_disabled_when_max_repeats_zero() {
let lb = LoopBreaker::with_config(60, 0);
assert!(!lb.enabled());
for _ in 0..100 {
assert_eq!(lb.check("s1", 1), LoopDecision::Allow);
}
}
#[test]
fn test_fingerprint_stable_and_intent_aware() {
let a =
br#"{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"temperature":0.1}"#;
let b =
br#"{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"temperature":0.9}"#;
assert_eq!(
intent_fingerprint("/v1/chat", a),
intent_fingerprint("/v1/chat", b)
);
let c = br#"{"model":"gpt-4o","messages":[{"role":"user","content":"bye"}]}"#;
assert_ne!(
intent_fingerprint("/v1/chat", a),
intent_fingerprint("/v1/chat", c)
);
assert_ne!(
intent_fingerprint("/v1/chat", a),
intent_fingerprint("/v2/chat", a)
);
}
}