poimen-memory/crates/mem-cli/src/rate_limiter.rs

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());
    }
}