Skip to main content

vtcode_core/tools/registry/
resiliency.rs

1//! Tool failure tracking and circuit breaker logic.
2//!
3//! This module provides resiliency patterns for tool execution,
4//! including consecutive failure tracking and backoff mechanisms.
5
6use hashbrown::HashMap;
7use std::time::Duration;
8
9use super::timeout::{AdaptiveTimeoutTuning, ToolLatencyStats, ToolTimeoutCategory};
10
11/// Tracks consecutive failures for a tool to enable circuit breaking.
12#[derive(Debug, Clone, Default)]
13pub struct ToolFailureTracker {
14    pub(super) consecutive_failures: u32,
15}
16
17impl ToolFailureTracker {
18    /// Record a failure.
19    pub fn record_failure(&mut self) {
20        self.consecutive_failures = self.consecutive_failures.saturating_add(1);
21    }
22
23    /// Reset the failure counter (on success).
24    pub fn reset(&mut self) {
25        self.consecutive_failures = 0;
26    }
27
28    /// Check if the circuit breaker should trip.
29    pub fn should_circuit_break(&self) -> bool {
30        self.consecutive_failures >= 5
31    }
32
33    /// Calculate backoff duration based on failure count.
34    pub fn backoff_duration(&self) -> Duration {
35        let base_ms = 500;
36        let max_ms = 5_000;
37        let failures = self.consecutive_failures.saturating_sub(5);
38        let backoff_ms = base_ms * 2_u64.pow(failures.min(8));
39        Duration::from_millis(backoff_ms.min(max_ms))
40    }
41}
42
43/// Internal state for resiliency tracking across tool categories.
44#[derive(Clone, Debug, Default)]
45pub struct ResiliencyContext {
46    pub(super) adaptive_timeout_ceiling: HashMap<ToolTimeoutCategory, Duration>,
47    pub(super) failure_trackers: HashMap<ToolTimeoutCategory, ToolFailureTracker>,
48    pub(super) success_trackers: HashMap<ToolTimeoutCategory, u32>,
49    pub(super) latency_stats: HashMap<ToolTimeoutCategory, ToolLatencyStats>,
50    pub(super) adaptive_tuning: AdaptiveTimeoutTuning,
51}
52
53#[cfg(test)]
54mod tests {
55    use super::ToolFailureTracker;
56    use std::time::Duration;
57
58    #[test]
59    fn circuit_breaker_threshold_and_backoff_ramp() {
60        let mut tracker = ToolFailureTracker::default();
61
62        for _ in 0..5 {
63            tracker.record_failure();
64        }
65        assert!(tracker.should_circuit_break());
66        assert_eq!(tracker.backoff_duration(), Duration::from_millis(500));
67
68        tracker.record_failure();
69        assert_eq!(tracker.backoff_duration(), Duration::from_millis(1000));
70
71        for _ in 0..2 {
72            tracker.record_failure();
73        }
74        assert_eq!(tracker.backoff_duration(), Duration::from_millis(4000));
75
76        tracker.record_failure();
77        assert_eq!(tracker.backoff_duration(), Duration::from_millis(5000));
78    }
79}