Implement M3.5.7: Rate limiting + idempotency (20 tests)

2026-08-26 13:35:50 -07:00
parent 1b0bc29027
commit f05565edd0
8 changed files with 1083 additions and 125 deletions
@@ -0,0 +1,243 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
+
+/// Rate limit error with retry guidance
+#[derive(Debug, Clone)]
+pub struct RateLimitError {
+    pub retry_after_seconds: u64,
+    pub limit_window_secs: u64,
+    pub reason: String,
+}
+
+impl RateLimitError {
+    pub fn reason(&self) -> String {
+        format!(
+            "{} (retry after {} seconds, window: {} seconds)",
+            self.reason, self.retry_after_seconds, self.limit_window_secs
+        )
+    }
+}
+
+/// Token bucket for a single endpoint
+#[derive(Debug, Clone)]
+struct TokenBucket {
+    tokens: f64,
+    last_refill: Instant,
+    capacity: f64,    // max tokens (per hour)
+    refill_rate: f64, // tokens per second
+}
+
+impl TokenBucket {
+    fn new(capacity: f64, refill_rate: f64) -> Self {
+        Self {
+            tokens: capacity,
+            last_refill: Instant::now(),
+            capacity,
+            refill_rate,
+        }
+    }
+
+    /// Refill tokens based on elapsed time
+    fn refill(&mut self) {
+        let now = Instant::now();
+        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
+        let refilled = elapsed * self.refill_rate;
+
+        self.tokens = (self.tokens + refilled).min(self.capacity);
+        self.last_refill = now;
+    }
+
+    /// Try to consume 1 token. Returns Ok if successful, Err(retry_after_secs) if rate limited.
+    fn try_consume(&mut self) -> Result<(), u64> {
+        self.refill();
+
+        if self.tokens >= 1.0 {
+            self.tokens -= 1.0;
+            return Ok(());
+        }
+
+        // Rate limited: estimate time until next token available
+        let tokens_needed = 1.0 - self.tokens;
+        let retry_after = (tokens_needed / self.refill_rate).ceil() as u64;
+        Err(retry_after.max(1))
+    }
+}
+
+/// Rate limiter with per-apikey, per-endpoint buckets
+pub struct RateLimiter {
+    buckets: Arc<Mutex<HashMap<String, Arc<Mutex<TokenBucket>>>>>,
+    limit_config: LimitConfig,
+}
+
+#[derive(Clone, Debug)]
+pub struct LimitConfig {
+    pub ingest_per_hour: f64,
+    pub query_per_hour: f64,
+    pub projects_per_hour: f64,
+    pub burst_per_second: f64, // Currently unused but kept for API compatibility
+}
+
+impl Default for LimitConfig {
+    fn default() -> Self {
+        Self {
+            ingest_per_hour: 100.0,
+            query_per_hour: 1000.0,
+            projects_per_hour: 100.0,
+            burst_per_second: 10.0,
+        }
+    }
+}
+
+impl RateLimiter {
+    pub fn new(config: LimitConfig) -> Self {
+        Self {
+            buckets: Arc::new(Mutex::new(HashMap::new())),
+            limit_config: config,
+        }
+    }
+
+    /// Get or create bucket for apikey + endpoint
+    fn get_or_create_bucket(&self, apikey_endpoint: &str) -> Arc<Mutex<TokenBucket>> {
+        let mut buckets = self.buckets.lock().unwrap();
+        let config = &self.limit_config;
+
+        if !buckets.contains_key(apikey_endpoint) {
+            // Determine limit based on endpoint
+            let capacity = if apikey_endpoint.contains("/memory/ingest") {
+                config.ingest_per_hour
+            } else if apikey_endpoint.contains("/memory/query") {
+                config.query_per_hour
+            } else if apikey_endpoint.contains("/memory/projects") {
+                config.projects_per_hour
+            } else {
+                // Unlimited for unknown endpoints
+                f64::INFINITY
+            };
+
+            let refill_rate = if capacity.is_infinite() {
+                f64::INFINITY
+            } else {
+                capacity / 3600.0 // per second
+            };
+
+            let bucket = TokenBucket::new(capacity, refill_rate);
+            buckets.insert(apikey_endpoint.to_string(), Arc::new(Mutex::new(bucket)));
+        }
+
+        buckets[apikey_endpoint].clone()
+    }
+
+    /// Check rate limit for apikey + endpoint. Returns Ok or Err with retry guidance.
+    pub fn check(&self, apikey: &str, endpoint: &str) -> Result<(), RateLimitError> {
+        let key = format!("{}::{}", apikey, endpoint);
+        let bucket = self.get_or_create_bucket(&key);
+        let mut b = bucket.lock().unwrap();
+
+        match b.try_consume() {
+            Ok(_) => Ok(()),
+            Err(retry_after) => {
+                let window_secs = if endpoint.contains("/memory/ingest") {
+                    3600
+                } else if endpoint.contains("/memory/query") {
+                    3600
+                } else if endpoint.contains("/memory/projects") {
+                    3600
+                } else {
+                    3600
+                };
+
+                Err(RateLimitError {
+                    retry_after_seconds: retry_after,
+                    limit_window_secs: window_secs,
+                    reason: format!(
+                        "rate_limit_exceeded for {}",
+                        endpoint
+                    ),
+                })
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_token_bucket_refill() {
+        let mut bucket = TokenBucket::new(100.0, 100.0 / 3600.0);
+        assert!(bucket.try_consume().is_ok());
+        // After one consumption, should have 99 tokens
+        assert_eq!((bucket.tokens * 1.0) as i64, 99);
+    }
+
+    #[test]
+    fn test_rate_limit_within_capacity() {
+        let config = LimitConfig {
+            ingest_per_hour: 5.0,
+            query_per_hour: 10.0,
+            projects_per_hour: 10.0,
+            burst_per_second: 10.0,
+        };
+        let limiter = RateLimiter::new(config);
+
+        // First 5 should succeed
+        for _ in 0..5 {
+            assert!(limiter.check("apikey1", "/memory/ingest").is_ok());
+        }
+
+        // 6th should fail
+        let err = limiter.check("apikey1", "/memory/ingest");
+        assert!(err.is_err());
+        if let Err(e) = err {
+            assert!(e.retry_after_seconds > 0);
+        }
+    }
+
+    #[test]
+    fn test_per_apikey_isolation() {
+        let config = LimitConfig {
+            ingest_per_hour: 5.0,
+            query_per_hour: 10.0,
+            projects_per_hour: 10.0,
+            burst_per_second: 10.0,
+        };
+        let limiter = RateLimiter::new(config);
+
+        // apikey1 uses up 5 ingest requests
+        for _ in 0..5 {
+            assert!(limiter.check("apikey1", "/memory/ingest").is_ok());
+        }
+        assert!(limiter.check("apikey1", "/memory/ingest").is_err());
+
+        // apikey2 should have its own 5
+        for _ in 0..5 {
+            assert!(limiter.check("apikey2", "/memory/ingest").is_ok());
+        }
+        assert!(limiter.check("apikey2", "/memory/ingest").is_err());
+    }
+
+    #[test]
+    fn test_per_endpoint_isolation() {
+        let config = LimitConfig {
+            ingest_per_hour: 5.0,
+            query_per_hour: 10.0,
+            projects_per_hour: 10.0,
+            burst_per_second: 10.0,
+        };
+        let limiter = RateLimiter::new(config);
+
+        // Use up 5 ingest
+        for _ in 0..5 {
+            assert!(limiter.check("apikey1", "/memory/ingest").is_ok());
+        }
+        assert!(limiter.check("apikey1", "/memory/ingest").is_err());
+
+        // Query should have separate 10 limit
+        for _ in 0..10 {
+            assert!(limiter.check("apikey1", "/memory/query").is_ok());
+        }
+        assert!(limiter.check("apikey1", "/memory/query").is_err());
+    }
+}