diff --git a/CHANGELOG.md b/CHANGELOG.md index eb9b8cb..21c78c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,9 @@ ## [Unreleased] +### Changed +- **Rate limiter replaced with `governor` (GCRA, lock-free)** (`src/middleware/rate_limit.rs`): replaced the custom `Mutex>>` sliding-window implementation with the `governor` crate (GCRA algorithm). Enforcement is now lock-free (atomics) and O(1) per check — no more O(n) `Vec::retain` on every request, no more background cleanup task, no more TOCTOU race (closes #82 root cause). Adds optional `rate_limit_burst` field on `AgentPolicy` (defaults to `rate_limit` for full backward compatibility). Closes #98. + ### Added - **Kubernetes ConfigMap watcher** (`src/kubernetes.rs`, `src/reload.rs`, `src/config.rs`): Arbitus can now watch a Kubernetes ConfigMap via the K8s API and hot-reload its configuration on every `Apply` event — no `SIGUSR1` or 30-second polling required. Set `kubernetes.configmap_name` in `gateway.yml` to enable. Namespace defaults to the pod's own namespace (read from the projected service-account token). Built as an optional feature (`kubernetes`, enabled by default in the official image) using `kube 0.98` + `kube-runtime`. The Helm chart gains a `kubernetesWatcher` section that automatically provisions the Role + RoleBinding and injects the `kubernetes:` config block. Closes #132. diff --git a/Cargo.lock b/Cargo.lock index 0d63c0d..d607099 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -108,6 +108,7 @@ dependencies = [ "chrono", "clap", "futures-util", + "governor", "hex", "jsonschema 0.18.3", "jsonwebtoken 10.3.0", @@ -763,6 +764,19 @@ dependencies = [ "syn", ] +[[package]] +name = "dashmap" +version = "5.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +dependencies = [ + "cfg-if", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "data-encoding" version = "2.10.0" @@ -1187,6 +1201,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + [[package]] name = "futures-util" version = "0.3.32" @@ -1262,6 +1282,26 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +[[package]] +name = "governor" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68a7f542ee6b35af73b06abc0dad1c1bae89964e4e253bc4b587b91c9637867b" +dependencies = [ + "cfg-if", + "dashmap", + "futures", + "futures-timer", + "no-std-compat", + "nonzero_ext", + "parking_lot", + "portable-atomic", + "quanta", + "rand 0.8.5", + "smallvec", + "spinning_top", +] + [[package]] name = "group" version = "0.13.0" @@ -2234,6 +2274,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "no-std-compat" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" + [[package]] name = "nom" version = "7.1.3" @@ -2253,6 +2299,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2690,6 +2742,12 @@ dependencies = [ "nom 7.1.3", ] +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + [[package]] name = "potential_utf" version = "0.1.5" @@ -2800,6 +2858,21 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + [[package]] name = "quinn" version = "0.11.9" @@ -2972,6 +3045,15 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3677,6 +3759,15 @@ version = "0.9.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +[[package]] +name = "spinning_top" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96d2d1d716fb500937168cc09353ffdc7a012be8475ac7308e1bdf0e3923300" +dependencies = [ + "lock_api", +] + [[package]] name = "spki" version = "0.7.3" @@ -4482,6 +4573,22 @@ dependencies = [ "winsafe", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + [[package]] name = "winapi-util" version = "0.1.11" @@ -4491,6 +4598,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-core" version = "0.62.2" diff --git a/Cargo.toml b/Cargo.toml index 2a08abe..7590feb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,6 +60,7 @@ sha2 = "0.10" hex = "0.4" which = "7" lru = "0.16.3" +governor = "0.6" kube = { version = "0.98", features = ["client", "runtime"], optional = true } k8s-openapi = { version = "0.24", features = ["v1_31"], optional = true } regorus = "0.2" diff --git a/README.md b/README.md index 4740738..de6862f 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,7 @@ Agent (Cursor, Claude, etc.) - **Multiple upstreams** — route different agents to different MCP servers - **Circuit breaker** — automatic upstream failure isolation with half-open recovery - **Config hot-reload** — reload on `SIGUSR1`, automatically every 30 seconds, or event-driven via Kubernetes ConfigMap watcher +- **Lock-free rate limiting** — GCRA algorithm via `governor`; O(1) per check, configurable burst allowance - **Metrics** — Prometheus-compatible `/metrics` endpoint with cost/token estimation - **OpenTelemetry** — export traces to any OTLP backend (Jaeger, Tempo, Honeycomb, Datadog) - **Dashboard** — `/dashboard` audit viewer with per-agent filtering diff --git a/src/a2a/interceptor.rs b/src/a2a/interceptor.rs index ebb23d4..36d8e0d 100644 --- a/src/a2a/interceptor.rs +++ b/src/a2a/interceptor.rs @@ -201,6 +201,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, diff --git a/src/bin/arbitus.rs b/src/bin/arbitus.rs index cbb6f1c..ee827e7 100644 --- a/src/bin/arbitus.rs +++ b/src/bin/arbitus.rs @@ -1087,6 +1087,7 @@ audits: [] allowed_tools: None, denied_tools: vec![], rate_limit: 60, + rate_limit_burst: None, tool_rate_limits: std::collections::HashMap::new(), upstream: None, api_key: None, diff --git a/src/config.rs b/src/config.rs index 12d06ef..416f0b9 100644 --- a/src/config.rs +++ b/src/config.rs @@ -322,6 +322,10 @@ pub struct AgentPolicy { pub denied_tools: Vec, #[serde(default = "default_rate_limit")] pub rate_limit: usize, + /// Maximum burst size for the agent rate limiter. + /// Defaults to `rate_limit` (entire quota usable as burst, matching the old sliding-window behaviour). + #[serde(default)] + pub rate_limit_burst: Option, /// Per-tool rate limits (calls/min). Checked in addition to the global rate_limit. #[serde(default)] pub tool_rate_limits: HashMap, @@ -757,6 +761,7 @@ pub(crate) fn make_agent( allowed_tools: allowed.map(|v| v.into_iter().map(String::from).collect()), denied_tools: denied.into_iter().map(String::from).collect(), rate_limit, + rate_limit_burst: None, tool_rate_limits: std::collections::HashMap::new(), upstream: None, api_key: None, diff --git a/src/gateway.rs b/src/gateway.rs index 7ce2bf5..0cdd172 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -1137,6 +1137,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 100, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, @@ -1231,6 +1232,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 100, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: Some(upstream_name.to_string()), api_key: None, @@ -1381,6 +1383,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 100, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, @@ -1608,6 +1611,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 100, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, diff --git a/src/live_config.rs b/src/live_config.rs index ae6dfef..99d55a4 100644 --- a/src/live_config.rs +++ b/src/live_config.rs @@ -45,6 +45,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 60, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: Some(key.to_string()), @@ -66,6 +67,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 60, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, @@ -197,6 +199,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 60, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, diff --git a/src/middleware/hitl.rs b/src/middleware/hitl.rs index ef560ef..da3082f 100644 --- a/src/middleware/hitl.rs +++ b/src/middleware/hitl.rs @@ -19,6 +19,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 60, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, diff --git a/src/middleware/rate_limit.rs b/src/middleware/rate_limit.rs index bfcffb9..ec19802 100644 --- a/src/middleware/rate_limit.rs +++ b/src/middleware/rate_limit.rs @@ -1,14 +1,225 @@ -use super::{Decision, McpContext, Middleware}; +use super::{Decision, McpContext, Middleware, RateLimitInfo}; use crate::live_config::LiveConfig; use async_trait::async_trait; +use governor::{DefaultDirectRateLimiter, Quota, RateLimiter}; use std::{ collections::HashMap, - sync::Arc, - time::{Duration, Instant}, + num::NonZeroU32, + sync::{ + Arc, RwLock, + atomic::{AtomicI64, AtomicU64, Ordering}, + }, + time::{SystemTime, UNIX_EPOCH}, }; -use tokio::sync::{Mutex, watch}; +use tokio::sync::watch; + +// ── Per-key limiter entry ───────────────────────────────────────────────────── + +/// One rate-limiter entry per (agent | tool | IP) key. +/// +/// `governor` handles enforcement via lock-free atomics (GCRA). +/// The separate `remaining` / `window_start_ms` atomics give us an O(1) +/// approximation of burst-capacity remaining for `X-RateLimit-Remaining` +/// response headers — no Vec, no Mutex. +struct LimiterEntry { + limiter: DefaultDirectRateLimiter, + /// Configured quota (requests / min). + limit: usize, + /// Approximate remaining capacity in the current 60-second window. + remaining: AtomicI64, + /// Unix-epoch milliseconds when the current window started. + window_start_ms: AtomicU64, +} + +impl LimiterEntry { + fn new(limit: usize, burst: usize) -> Self { + let nz_limit = NonZeroU32::new(limit.max(1) as u32).expect("rate_limit > 0"); + let nz_burst = NonZeroU32::new(burst.max(1) as u32).expect("burst > 0"); + let quota = Quota::per_minute(nz_limit).allow_burst(nz_burst); + Self { + limiter: RateLimiter::direct(quota), + limit, + remaining: AtomicI64::new(burst as i64), + window_start_ms: AtomicU64::new(now_ms()), + } + } + + /// Check whether a cell is available. + /// Returns `(allowed, remaining, reset_after_secs)`. + fn check(&self) -> (bool, usize, u64) { + // Reset the remaining counter once per 60-second window. + let ws = self.window_start_ms.load(Ordering::Relaxed); + let now = now_ms(); + if now.saturating_sub(ws) >= 60_000 { + // CAS prevents two threads from both resetting at once. + if self + .window_start_ms + .compare_exchange(ws, now, Ordering::SeqCst, Ordering::Relaxed) + .is_ok() + { + self.remaining.store(self.limit as i64, Ordering::Relaxed); + } + } + + let elapsed_secs = now.saturating_sub(self.window_start_ms.load(Ordering::Relaxed)) / 1000; + let reset_after = 60u64.saturating_sub(elapsed_secs); + + match self.limiter.check() { + Ok(_) => { + let r = (self.remaining.fetch_sub(1, Ordering::Relaxed) - 1).max(0) as usize; + (true, r, reset_after) + } + Err(_) => (false, 0, reset_after), + } + } +} + +fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} + +// ── Limiter store ───────────────────────────────────────────────────────────── + +type EntryMap = Arc>>>; + +/// Look up an existing entry or create a new one with the given `limit`/`burst`. +/// If the stored entry's limit differs from the current config (after a hot-reload), +/// replace it so the new rate takes effect immediately. +fn get_or_create(map: &EntryMap, key: K, limit: usize, burst: usize) -> Arc +where + K: std::hash::Hash + Eq + Clone, +{ + // Fast path: entry exists and has the right limit. + { + let m = map.read().unwrap(); + if let Some(e) = m.get(&key) + && e.limit == limit + { + return Arc::clone(e); + } + } + // Slow path: create (or replace) the entry. + let mut m = map.write().unwrap(); + let entry = Arc::new(LimiterEntry::new(limit, burst)); + m.insert(key, Arc::clone(&entry)); + entry +} + +// ── Middleware ───────────────────────────────────────────────────────────────── + +pub struct RateLimitMiddleware { + config: watch::Receiver>, + /// One entry per agent_id. + agent_limiters: EntryMap, + /// One entry per (agent_id, tool_name) pair. + tool_limiters: EntryMap<(String, String)>, + /// One entry per client IP string. + ip_limiters: EntryMap, +} + +impl RateLimitMiddleware { + pub fn new(config: watch::Receiver>) -> Self { + Self { + config, + agent_limiters: Arc::new(RwLock::new(HashMap::new())), + tool_limiters: Arc::new(RwLock::new(HashMap::new())), + ip_limiters: Arc::new(RwLock::new(HashMap::new())), + } + } +} + +#[async_trait] +impl Middleware for RateLimitMiddleware { + fn name(&self) -> &'static str { + "rate_limit" + } + + async fn check(&self, ctx: &McpContext) -> Decision { + if !matches!( + ctx.method.as_str(), + "tools/call" | "resources/read" | "resources/subscribe" | "prompts/get" + ) { + return Decision::Allow { rl: None }; + } + + let (global_limit, global_burst, tool_limit, ip_limit) = { + let cfg = self.config.borrow(); + let Some(policy) = cfg.agents.get(&ctx.agent_id) else { + return Decision::Allow { rl: None }; // unknown agents blocked by AuthMiddleware + }; + let burst = policy.rate_limit_burst.unwrap_or(policy.rate_limit); + let tool_limit = ctx + .tool_name + .as_ref() + .and_then(|t| policy.tool_rate_limits.get(t).copied()); + (policy.rate_limit, burst, tool_limit, cfg.ip_rate_limit) + }; + + // ── IP rate limit (cheapest rejection — checked first) ──────────────── + if let (Some(limit), Some(ip)) = (ip_limit, ctx.client_ip.as_ref()) { + let entry = get_or_create(&self.ip_limiters, ip.clone(), limit, limit); + let (allowed, remaining, reset_after) = entry.check(); + if !allowed { + return Decision::Block { + reason: format!("IP rate limit exceeded ({limit}/min)"), + rl: Some(RateLimitInfo { + limit, + remaining, + reset_after_secs: reset_after, + }), + }; + } + } -type ToolCounts = Arc>>>; + // ── Global agent rate limit ─────────────────────────────────────────── + let agent_entry = get_or_create( + &self.agent_limiters, + ctx.agent_id.clone(), + global_limit, + global_burst, + ); + let (allowed, remaining, reset_after) = agent_entry.check(); + if !allowed { + return Decision::Block { + reason: format!("rate limit exceeded ({global_limit}/min)"), + rl: Some(RateLimitInfo { + limit: global_limit, + remaining: 0, + reset_after_secs: reset_after, + }), + }; + } + let agent_rl = RateLimitInfo { + limit: global_limit, + remaining, + reset_after_secs: reset_after, + }; + + // ── Per-tool rate limit ─────────────────────────────────────────────── + if let (Some(limit), Some(tool)) = (tool_limit, ctx.tool_name.as_ref()) { + let key = (ctx.agent_id.clone(), tool.clone()); + let entry = get_or_create(&self.tool_limiters, key, limit, limit); + let (allowed, _, reset_after) = entry.check(); + if !allowed { + return Decision::Block { + reason: format!("tool '{tool}' rate limit exceeded ({limit}/min)"), + rl: Some(RateLimitInfo { + limit, + remaining: 0, + reset_after_secs: reset_after, + }), + }; + } + } + + Decision::Allow { rl: Some(agent_rl) } + } +} + +// ── Tests ───────────────────────────────────────────────────────────────────── #[cfg(test)] mod tests { @@ -20,6 +231,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit, + rate_limit_burst: None, tool_rate_limits: HashMap::new(), upstream: None, api_key: None, @@ -77,7 +289,6 @@ mod tests { #[tokio::test] async fn unknown_agent_passes_to_auth_middleware() { - // Rate limit doesn't block unknown agents — that's auth's job let mw = make_mw(HashMap::new(), None); assert!(matches!( mw.check(&ctx("ghost", "echo", None)).await, @@ -128,6 +339,7 @@ mod tests { allowed_tools: None, denied_tools: vec![], rate_limit: 100, + rate_limit_burst: None, tool_rate_limits: tool_limits, upstream: None, api_key: None, @@ -256,174 +468,30 @@ mod tests { Decision::Block { .. } )); } -} -pub struct RateLimitMiddleware { - config: watch::Receiver>, - /// Per-agent sliding window counters — keyed by agent_id. - counts: Arc>>>, - /// Per-(agent, tool) sliding window counters for tool_rate_limits. - tool_counts: ToolCounts, - /// Per-IP sliding window counters (HTTP mode). Keyed by client IP string. - ip_counts: Arc>>>, -} - -impl RateLimitMiddleware { - pub fn new(config: watch::Receiver>) -> Self { - let counts = Arc::new(Mutex::new(HashMap::new())); - let tool_counts = Arc::new(Mutex::new(HashMap::new())); - let ip_counts = Arc::new(Mutex::new(HashMap::new())); - - // Background task: purge inactive entries every 5 minutes to prevent - // unbounded HashMap growth when many distinct agents/IPs are seen. - { - let counts = Arc::clone(&counts); - let tool_counts = Arc::clone(&tool_counts); - let ip_counts = Arc::clone(&ip_counts); - tokio::spawn(async move { - let mut interval = tokio::time::interval(Duration::from_secs(300)); - interval.tick().await; // skip immediate tick - loop { - interval.tick().await; - let window = Duration::from_secs(60); - let now = Instant::now(); - { - let mut m = counts.lock().await; - m.retain(|_, ts: &mut Vec| { - ts.retain(|t| now.duration_since(*t) < window); - !ts.is_empty() - }); - } - { - let mut m = tool_counts.lock().await; - m.retain(|_, ts: &mut Vec| { - ts.retain(|t| now.duration_since(*t) < window); - !ts.is_empty() - }); - } - { - let mut m = ip_counts.lock().await; - m.retain(|_, ts: &mut Vec| { - ts.retain(|t| now.duration_since(*t) < window); - !ts.is_empty() - }); - } - } - }); - } - - Self { - config, - counts, - tool_counts, - ip_counts, - } - } -} - -/// Seconds until the oldest timestamp in `ts` ages out of the 60s window. -fn window_reset_secs(ts: &[Instant], now: Instant) -> u64 { - ts.first() - .map(|oldest| { - let elapsed = now.duration_since(*oldest).as_secs(); - 60u64.saturating_sub(elapsed) - }) - .unwrap_or(60) -} - -#[async_trait] -impl Middleware for RateLimitMiddleware { - fn name(&self) -> &'static str { - "rate_limit" - } - - async fn check(&self, ctx: &McpContext) -> Decision { - use super::RateLimitInfo; - - if !matches!( - ctx.method.as_str(), - "tools/call" | "resources/read" | "resources/subscribe" | "prompts/get" - ) { - return Decision::Allow { rl: None }; - } - - let (global_limit, tool_limit, ip_limit) = { - let cfg = self.config.borrow(); - let Some(policy) = cfg.agents.get(&ctx.agent_id) else { - return Decision::Allow { rl: None }; // unknown agents are blocked by AuthMiddleware - }; - let tool_limit = ctx - .tool_name - .as_ref() - .and_then(|t| policy.tool_rate_limits.get(t).copied()); - (policy.rate_limit, tool_limit, cfg.ip_rate_limit) - }; - - let now = Instant::now(); - let window = Duration::from_secs(60); - - // ── IP rate limit (checked first — cheapest rejection) ───────────────── - if let (Some(limit), Some(ip)) = (ip_limit, ctx.client_ip.as_ref()) { - let mut ip_counts = self.ip_counts.lock().await; - let ts = ip_counts.entry(ip.clone()).or_default(); - ts.retain(|t| now.duration_since(*t) < window); - if ts.len() >= limit { - return Decision::Block { - reason: format!("IP rate limit exceeded ({limit}/min)"), - rl: Some(RateLimitInfo { - limit, - remaining: 0, - reset_after_secs: window_reset_secs(ts, now), - }), - }; - } - ts.push(now); - } - - // ── Global agent rate limit ──────────────────────────────────────────── - let agent_rl = { - let mut counts = self.counts.lock().await; - let ts = counts.entry(ctx.agent_id.clone()).or_default(); - ts.retain(|t| now.duration_since(*t) < window); - - if ts.len() >= global_limit { - return Decision::Block { - reason: format!("rate limit exceeded ({global_limit}/min)"), - rl: Some(RateLimitInfo { - limit: global_limit, - remaining: 0, - reset_after_secs: window_reset_secs(ts, now), - }), - }; - } - ts.push(now); - RateLimitInfo { - limit: global_limit, - remaining: global_limit.saturating_sub(ts.len()), - reset_after_secs: window_reset_secs(ts, now), - } - }; - - // ── Per-tool rate limit ──────────────────────────────────────────────── - if let (Some(limit), Some(tool)) = (tool_limit, ctx.tool_name.as_ref()) { - let key = (ctx.agent_id.clone(), tool.clone()); - let mut tool_counts = self.tool_counts.lock().await; - let ts = tool_counts.entry(key.clone()).or_default(); - ts.retain(|t| now.duration_since(*t) < window); - - if ts.len() >= limit { - return Decision::Block { - reason: format!("tool '{tool}' rate limit exceeded ({limit}/min)"), - rl: Some(RateLimitInfo { - limit, - remaining: 0, - reset_after_secs: window_reset_secs(ts, now), - }), - }; - } - ts.push(now); - } - - Decision::Allow { rl: Some(agent_rl) } + #[tokio::test] + async fn burst_config_respected() { + let mut agents = HashMap::new(); + agents.insert( + "a".to_string(), + AgentPolicy { + rate_limit: 60, + rate_limit_burst: Some(2), // only 2 rapid requests allowed + ..policy(60) + }, + ); + let mw = make_mw(agents, None); + assert!(matches!( + mw.check(&ctx("a", "echo", None)).await, + Decision::Allow { .. } + )); + assert!(matches!( + mw.check(&ctx("a", "echo", None)).await, + Decision::Allow { .. } + )); + assert!(matches!( + mw.check(&ctx("a", "echo", None)).await, + Decision::Block { .. } + )); } }