strange_loop/nano_agent/
budget.rs

1//! Budget enforcement for nano-agents
2
3use super::rdtsc;
4use std::time::Instant;
5
6/// Budget configuration for an agent
7#[derive(Debug, Clone, Copy)]
8pub struct Budget {
9    pub per_tick_ns: u128,
10    pub max_jitter_ns: u128,
11    pub max_violations: u32,
12}
13
14impl Budget {
15    pub fn new(per_tick_ns: u128) -> Self {
16        Self {
17            per_tick_ns,
18            max_jitter_ns: per_tick_ns / 10, // 10% jitter tolerance
19            max_violations: 3,
20        }
21    }
22
23    pub fn with_jitter(mut self, max_jitter_ns: u128) -> Self {
24        self.max_jitter_ns = max_jitter_ns;
25        self
26    }
27
28    pub fn with_max_violations(mut self, max: u32) -> Self {
29        self.max_violations = max;
30        self
31    }
32}
33
34/// Guard for enforcing time budgets
35pub struct BudgetGuard {
36    start_tsc: u64,
37    start_ns: u128,
38    budget_ns: u128,
39    jitter_threshold: u128,
40    violations: u32,
41    max_violations: u32,
42}
43
44impl BudgetGuard {
45    pub fn new(budget: Budget) -> Self {
46        let start_instant = Instant::now();
47        Self {
48            start_tsc: rdtsc(),
49            start_ns: 0, // Will be set on first check
50            budget_ns: budget.per_tick_ns,
51            jitter_threshold: budget.max_jitter_ns,
52            violations: 0,
53            max_violations: budget.max_violations,
54        }
55    }
56
57    /// Check if budget is exhausted
58    #[inline(always)]
59    pub fn is_exhausted(&self, now_ns: u128) -> bool {
60        now_ns - self.start_ns >= self.budget_ns
61    }
62
63    /// Check for budget violation
64    #[inline(always)]
65    pub fn check_violation(&mut self, elapsed_ns: u128) -> bool {
66        if elapsed_ns > self.budget_ns + self.jitter_threshold {
67            self.violations += 1;
68            return self.violations > self.max_violations;
69        }
70        false
71    }
72
73    /// Reset for next tick
74    #[inline(always)]
75    pub fn reset(&mut self, now_ns: u128) {
76        self.start_tsc = rdtsc();
77        self.start_ns = now_ns;
78    }
79
80    /// Get elapsed cycles since start
81    #[inline(always)]
82    pub fn elapsed_cycles(&self) -> u64 {
83        rdtsc() - self.start_tsc
84    }
85}
86
87/// Kill switch for runaway agents
88pub struct KillSwitch {
89    enabled: bool,
90    max_runtime_ns: u128,
91    max_ticks: u64,
92    max_messages: u64,
93    start_time: Instant,
94    tick_count: u64,
95    message_count: u64,
96}
97
98impl KillSwitch {
99    pub fn new(max_runtime_ns: u128) -> Self {
100        Self {
101            enabled: true,
102            max_runtime_ns,
103            max_ticks: 1_000_000_000, // 1 billion ticks
104            max_messages: 10_000_000,  // 10 million messages
105            start_time: Instant::now(),
106            tick_count: 0,
107            message_count: 0,
108        }
109    }
110
111    /// Check if kill condition is met
112    pub fn should_kill(&self) -> bool {
113        if !self.enabled {
114            return false;
115        }
116
117        // Check runtime
118        if self.start_time.elapsed().as_nanos() > self.max_runtime_ns {
119            return true;
120        }
121
122        // Check tick count
123        if self.tick_count > self.max_ticks {
124            return true;
125        }
126
127        // Check message count
128        if self.message_count > self.max_messages {
129            return true;
130        }
131
132        false
133    }
134
135    /// Record a tick
136    pub fn record_tick(&mut self) {
137        self.tick_count += 1;
138    }
139
140    /// Record messages
141    pub fn record_messages(&mut self, count: u64) {
142        self.message_count += count;
143    }
144
145    /// Disable the kill switch (dangerous!)
146    pub fn disable(&mut self) {
147        self.enabled = false;
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use std::thread::sleep;
155    use std::time::Duration;
156
157    #[test]
158    fn test_budget() {
159        let budget = Budget::new(1000)
160            .with_jitter(100)
161            .with_max_violations(5);
162
163        assert_eq!(budget.per_tick_ns, 1000);
164        assert_eq!(budget.max_jitter_ns, 100);
165        assert_eq!(budget.max_violations, 5);
166    }
167
168    #[test]
169    fn test_budget_guard() {
170        let budget = Budget::new(1_000_000); // 1ms
171        let mut guard = BudgetGuard::new(budget);
172
173        guard.reset(0);
174        assert!(!guard.is_exhausted(500_000)); // 0.5ms
175        assert!(guard.is_exhausted(1_500_000)); // 1.5ms
176    }
177
178    #[test]
179    fn test_kill_switch() {
180        let mut kill = KillSwitch::new(100_000_000); // 100ms
181
182        assert!(!kill.should_kill());
183
184        // Simulate many ticks
185        for _ in 0..1000 {
186            kill.record_tick();
187        }
188        assert!(!kill.should_kill());
189
190        // Exceed tick limit
191        kill.tick_count = 2_000_000_000;
192        assert!(kill.should_kill());
193    }
194}
strange_loop/nano_agent/budget.rs

strange_loop/nano_agent/
budget.rs
