deepstrike_core/governance/
rate_limit.rs

1use std::collections::{HashMap, VecDeque};
2
3use compact_str::CompactString;
4
5use crate::types::message::ToolCall;
6use crate::types::policy::GovernanceVerdict;
7
8/// Rate limit configuration for a tool.
9#[derive(Debug, Clone)]
10pub struct RateLimit {
11    pub max_calls: u32,
12    pub window_ms: u64,
13}
14
15impl Default for RateLimit {
16    fn default() -> Self {
17        Self {
18            max_calls: 60,
19            window_ms: 60_000,
20        }
21    }
22}
23
24/// Sliding-window rate limiter per tool.
25pub struct RateLimiter {
26    windows: HashMap<CompactString, VecDeque<u64>>,
27    limits: HashMap<CompactString, RateLimit>,
28    default_limit: RateLimit,
29    /// Current timestamp in ms — injected by SDK layer (no I/O in kernel).
30    current_time_ms: u64,
31}
32
33impl RateLimiter {
34    pub fn new(default_limit: RateLimit) -> Self {
35        Self {
36            windows: HashMap::new(),
37            limits: HashMap::new(),
38            default_limit,
39            current_time_ms: 0,
40        }
41    }
42
43    pub fn set_limit(&mut self, tool_name: impl Into<CompactString>, limit: RateLimit) {
44        self.limits.insert(tool_name.into(), limit);
45    }
46
47    /// Must be called before each check to provide current time.
48    pub fn set_time(&mut self, now_ms: u64) {
49        self.current_time_ms = now_ms;
50    }
51
52    pub fn check(&mut self, call: &ToolCall) -> Option<GovernanceVerdict> {
53        // current_time_ms defaults to 0; SDK is expected to call set_time() before check.
54        // We don't debug_assert here because 0 is a valid monotonic-clock origin.
55        let limit = self.limits.get(&call.name).unwrap_or(&self.default_limit);
56        let window = self.windows.entry(call.name.clone()).or_default();
57
58        // Evict expired entries
59        let cutoff = self.current_time_ms.saturating_sub(limit.window_ms);
60        while window.front().is_some_and(|&t| t < cutoff) {
61            window.pop_front();
62        }
63
64        if window.len() as u32 >= limit.max_calls {
65            let oldest = window.front().copied().unwrap_or(self.current_time_ms);
66            let retry_after = oldest + limit.window_ms - self.current_time_ms;
67            return Some(GovernanceVerdict::RateLimited {
68                retry_after_ms: retry_after,
69            });
70        }
71
72        window.push_back(self.current_time_ms);
73        None
74    }
75}
76
77impl Default for RateLimiter {
78    fn default() -> Self {
79        Self::new(RateLimit::default())
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    fn make_call(name: &str) -> ToolCall {
88        ToolCall {
89            id: CompactString::new("c"),
90            name: CompactString::new(name),
91            arguments: serde_json::Value::Null,
92        }
93    }
94
95    #[test]
96    fn allows_within_limit() {
97        let mut rl = RateLimiter::new(RateLimit {
98            max_calls: 3,
99            window_ms: 1000,
100        });
101        rl.set_time(100);
102        assert!(rl.check(&make_call("foo")).is_none());
103        assert!(rl.check(&make_call("foo")).is_none());
104        assert!(rl.check(&make_call("foo")).is_none());
105        // 4th call should be limited
106        assert!(rl.check(&make_call("foo")).is_some());
107    }
108
109    #[test]
110    fn expires_old_entries() {
111        let mut rl = RateLimiter::new(RateLimit {
112            max_calls: 1,
113            window_ms: 100,
114        });
115        rl.set_time(0);
116        assert!(rl.check(&make_call("bar")).is_none());
117        assert!(rl.check(&make_call("bar")).is_some());
118
119        rl.set_time(200); // window expired
120        assert!(rl.check(&make_call("bar")).is_none());
121    }
122}
deepstrike_core/governance/rate_limit.rs

deepstrike_core/governance/
rate_limit.rs
