diff --git a/CHANGELOG.md b/CHANGELOG.md index 4859016..a83fc2d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,10 @@ from [Conventional Commits](https://www.conventionalcommits.org/). ## [Unreleased] +### Added + +- `StuckLoopHook`: a new `ToolDispatchHook` to prevent infinite tool-dispatch loops by terminating dispatch when the model requests the exact same tool invocation repeatedly. + ## [0.5.0](https://github.com/ForeverAngry/rig-compose/compare/v0.4.3...v0.5.0) - 2026-05-28 ### Fixed diff --git a/ROADMAP.md b/ROADMAP.md index e27908c..32f2db6 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -35,7 +35,7 @@ It must stay independent from `rig-core`, Memvid, MCP, concrete resource stores, memory-like, resource-like, and tool-result-like examples, but downstream crates still need to project concrete memory, resource, graph, and lineage data into it. -- Dispatch and agent lifecycle policy have hooks and budget accounting, but richer approval, retry, result-size, stuck-loop, and trace policies still live downstream or remain unbuilt. +- Dispatch and agent lifecycle policy have hooks and budget accounting, but richer approval and trace policies still live downstream or remain unbuilt. ## Next Work diff --git a/src/lib.rs b/src/lib.rs index bf4b494..e59ece7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -97,8 +97,8 @@ pub use normalizer::{ }; pub use registry::{KernelError, SkillDescriptor, SkillRegistry, ToolRegistry}; pub use reliability::{ - DefaultRetryClassifier, HistoryEntry, RetryClass, RetryClassifier, ToolCallFingerprint, - repair_history, + DefaultRetryClassifier, HistoryEntry, RetryClass, RetryClassifier, StuckLoopHook, + ToolCallFingerprint, repair_history, }; pub use skill::{Skill, SkillId, SkillOutcome}; pub use tool::{ diff --git a/src/reliability.rs b/src/reliability.rs index be7ff2c..e426539 100644 --- a/src/reliability.rs +++ b/src/reliability.rs @@ -17,6 +17,9 @@ //! model should see. Multiple retries of the same fingerprint collapse to //! a single canonical entry; the host stays in control of how many //! physical retries actually happened. +//! 4. [`StuckLoopHook`] — a [`ToolDispatchHook`] that tracks invocation +//! fingerprints and terminates dispatch if the model requests identical +//! calls too many times, preventing runaway costs. //! //! These primitives are intentionally synchronous and infallible: they //! operate on already-materialized invocations and outcomes, never on live @@ -49,12 +52,17 @@ //! assert!(matches!(repaired[0], HistoryEntry::Completed { .. })); //! ``` +use std::collections::HashMap; use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; +use std::sync::Mutex; +use async_trait::async_trait; use serde_json::Value; -use crate::normalizer::{ToolInvocation, ToolInvocationResult}; +use crate::normalizer::{ + ToolDispatchAction, ToolDispatchHook, ToolInvocation, ToolInvocationResult, +}; use crate::registry::KernelError; // ── Retry classification ───────────────────────────────────────────────────── @@ -166,6 +174,59 @@ fn canonicalize_value(value: &Value) -> Value { } } +// ── Stuck loop detection ─────────────────────────────────────────────────── + +/// A dispatch hook that cuts off repeated identical tool calls. +/// +/// Models sometimes enter a stuck loop where they repeatedly dispatch the +/// exact same tool invocation (same tool name, same arguments) without +/// making progress. This hook tracks invocation fingerprints across +/// dispatches and terminates the loop if the same fingerprint is invoked +/// more than `max_identical_calls` times. +pub struct StuckLoopHook { + max_identical_calls: usize, + counts: Mutex>, +} + +impl StuckLoopHook { + /// Create a new hook that terminates dispatch if any fingerprint + /// is requested more than `max_identical_calls` times. + pub fn new(max_identical_calls: usize) -> Self { + Self { + max_identical_calls, + counts: Mutex::new(HashMap::new()), + } + } +} + +#[async_trait] +impl ToolDispatchHook for StuckLoopHook { + async fn before_invocation( + &self, + invocation: &ToolInvocation, + ) -> Result { + let fingerprint = invocation.fingerprint(); + let mut counts = self + .counts + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + + let count = counts.entry(fingerprint).or_insert(0); + *count = count.saturating_add(1); + + if *count > self.max_identical_calls { + Ok(ToolDispatchAction::Terminate { + reason: format!( + "stuck loop detected: `{}` has been called {} times", + invocation.name, count + ), + }) + } else { + Ok(ToolDispatchAction::Continue) + } + } +} + // ── History repair ─────────────────────────────────────────────────────────── /// One entry in a tool-call history slice fed back to the model. @@ -484,4 +545,58 @@ mod tests { other => panic!("expected Failed, got {other:?}"), } } + + // ── Stuck loop detection tests ─────────────────────────────────────────── + + #[tokio::test] + async fn stuck_loop_terminates_after_threshold() { + let hook = StuckLoopHook::new(2); + let i = inv("search", json!({"q": "rig"})); + + // First call: allowed + let action = hook.before_invocation(&i).await.unwrap(); + assert!(matches!(action, ToolDispatchAction::Continue)); + + // Second call: allowed + let action = hook.before_invocation(&i).await.unwrap(); + assert!(matches!(action, ToolDispatchAction::Continue)); + + // Third call: terminated + let action = hook.before_invocation(&i).await.unwrap(); + match action { + ToolDispatchAction::Terminate { reason } => { + assert!(reason.contains("stuck loop detected")); + assert!(reason.contains("3 times")); + } + other => panic!("expected Terminate, got {other:?}"), + } + } + + #[tokio::test] + async fn stuck_loop_tracks_multiple_distinct_tools() { + let hook = StuckLoopHook::new(1); + let i1 = inv("search", json!({"q": "1"})); + let i2 = inv("search", json!({"q": "2"})); + + assert!(matches!( + hook.before_invocation(&i1).await.unwrap(), + ToolDispatchAction::Continue + )); + assert!(matches!( + hook.before_invocation(&i2).await.unwrap(), + ToolDispatchAction::Continue + )); + + // Calling i1 again triggers termination + assert!(matches!( + hook.before_invocation(&i1).await.unwrap(), + ToolDispatchAction::Terminate { .. } + )); + + // Calling i2 again triggers termination + assert!(matches!( + hook.before_invocation(&i2).await.unwrap(), + ToolDispatchAction::Terminate { .. } + )); + } }