use std::time::{Duration, Instant};
pub struct RateLimiter {
tokens: u32,
max_tokens: u32,
last_refill: Instant,
refill_interval: Duration,
}
impl RateLimiter {
pub fn new(max_tokens: u32, refill_interval: Duration) -> Self {
Self {
tokens: max_tokens,
max_tokens,
last_refill: Instant::now(),
refill_interval,
}
}
pub fn default_write_limiter() -> Self {
Self::new(30, Duration::from_secs(60))
}
pub fn check(&mut self) -> bool {
self.refill();
if self.tokens > 0 {
self.tokens -= 1;
true
} else {
false
}
}
fn refill(&mut self) {
let elapsed = self.last_refill.elapsed();
if elapsed >= self.refill_interval {
self.tokens = self.max_tokens;
self.last_refill = Instant::now();
}
}
pub fn is_write_method(method: &str) -> bool {
matches!(
method,
"sessions.send"
| "sessions.create"
| "sessions.patch"
| "sessions.compact"
| "sessions.reset"
| "sessions.delete"
| "chat.send"
| "chat.abort"
| "agents.create"
| "agents.update"
| "agents.delete"
| "agent.send"
| "config.set"
| "config.patch"
| "config.apply"
| "cron.add"
| "cron.remove"
| "cron.run"
| "cron.update"
| "cron.delete"
| "memory.store"
| "exec.approval.set"
| "exec.approval.resolve"
| "system.shutdown"
| "system.stop"
| "system.restart"
| "system.update.run"
| "node.pair.request"
| "node.pair.approve"
| "node.pair.reject"
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn allows_up_to_max() {
let mut rl = RateLimiter::new(3, Duration::from_secs(60));
assert!(rl.check());
assert!(rl.check());
assert!(rl.check());
assert!(!rl.check());
}
#[test]
fn write_methods_classified() {
assert!(RateLimiter::is_write_method("sessions.send"));
assert!(RateLimiter::is_write_method("chat.send"));
assert!(RateLimiter::is_write_method("config.set"));
assert!(!RateLimiter::is_write_method("health"));
assert!(!RateLimiter::is_write_method("agents.list"));
assert!(!RateLimiter::is_write_method("models.list"));
}
}