car_engine/

rate_limit.rs

//! Token bucket rate limiter for tool calls with backpressure support.

use std::collections::HashMap;
use tokio::sync::Mutex;
use tokio::time::Instant;

/// Per-tool rate limit configuration.
#[derive(Debug, Clone)]
pub struct RateLimit {
    pub max_calls: u32,
    pub interval_secs: f64,
}

/// Token bucket rate limiter for tool calls.
///
/// Each tool can have an independent rate limit. When a tool's bucket is empty,
/// `acquire()` blocks until a token becomes available (backpressure).
pub struct RateLimiter {
    limits: Mutex<HashMap<String, RateLimit>>,
    buckets: Mutex<HashMap<String, TokenBucket>>,
}

struct TokenBucket {
    tokens: f64,
    max_tokens: f64,
    refill_rate: f64, // tokens per second
    last_refill: Instant,
}

impl TokenBucket {
    fn new(max_tokens: f64, refill_rate: f64) -> Self {
        Self {
            tokens: max_tokens,
            max_tokens,
            refill_rate,
            last_refill: Instant::now(),
        }
    }

    /// Refill tokens based on elapsed time since last refill.
    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.max_tokens);
        self.last_refill = now;
    }

    /// Try to consume one token. Returns true if successful.
    fn try_consume(&mut self) -> bool {
        self.refill();
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }

    /// Seconds until one token is available (0.0 if already available).
    fn time_until_available(&mut self) -> f64 {
        self.refill();
        if self.tokens >= 1.0 {
            return 0.0;
        }
        let deficit = 1.0 - self.tokens;
        deficit / self.refill_rate
    }
}

impl RateLimiter {
    /// Create an empty rate limiter with no limits configured.
    pub fn new() -> Self {
        Self {
            limits: Mutex::new(HashMap::new()),
            buckets: Mutex::new(HashMap::new()),
        }
    }

    /// Configure a rate limit for a specific tool.
    ///
    /// `max_calls` tokens over `interval_secs` seconds. The refill rate is
    /// `max_calls / interval_secs` tokens per second.
    pub async fn set_limit(&self, tool: &str, limit: RateLimit) {
        let max_tokens = limit.max_calls as f64;
        let refill_rate = max_tokens / limit.interval_secs;

        self.limits.lock().await.insert(tool.to_string(), limit);

        self.buckets
            .lock()
            .await
            .insert(tool.to_string(), TokenBucket::new(max_tokens, refill_rate));
    }

    /// Configure a rate limit from a NON-async context. Returns false if
    /// either lock was held and no limit was installed.
    ///
    /// Exists for the synchronous `Runtime::with_*` builders, which cannot
    /// `.await` [`Self::set_limit`] but must still install a tool's declared
    /// rate limit — a limit the builder skipped would leave the tool
    /// unbounded, which is the opposite of what declaring one meant. The locks
    /// are uncontended during construction, so `try_lock` succeeds.
    pub fn try_set_limit(&self, tool: &str, limit: RateLimit) -> bool {
        let max_tokens = limit.max_calls as f64;
        let refill_rate = max_tokens / limit.interval_secs;

        let (Ok(mut limits), Ok(mut buckets)) = (self.limits.try_lock(), self.buckets.try_lock())
        else {
            return false;
        };
        limits.insert(tool.to_string(), limit);
        buckets.insert(tool.to_string(), TokenBucket::new(max_tokens, refill_rate));
        true
    }

    /// Wait until a token is available for the given tool, then consume it.
    ///
    /// If no rate limit is configured for the tool, returns immediately.
    /// This provides backpressure: callers block until capacity is available.
    pub async fn acquire(&self, tool: &str) {
        loop {
            let wait_time = {
                let mut buckets = self.buckets.lock().await;
                let bucket = match buckets.get_mut(tool) {
                    Some(b) => b,
                    None => return, // no limit configured
                };

                if bucket.try_consume() {
                    return;
                }

                bucket.time_until_available()
            };

            // Sleep outside the lock to allow other tasks to proceed.
            tokio::time::sleep(std::time::Duration::from_secs_f64(wait_time)).await;
        }
    }

    /// Non-blocking attempt to acquire a token for the given tool.
    ///
    /// Returns `true` if a token was consumed, `false` if the bucket is empty.
    /// Returns `true` if no rate limit is configured for the tool.
    pub async fn try_acquire(&self, tool: &str) -> bool {
        let mut buckets = self.buckets.lock().await;
        match buckets.get_mut(tool) {
            Some(bucket) => bucket.try_consume(),
            None => true, // no limit configured
        }
    }
}

impl Default for RateLimiter {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test(start_paused = true)]
    async fn test_token_bucket_refills_correctly() {
        let limiter = RateLimiter::new();
        limiter
            .set_limit(
                "tool_a",
                RateLimit {
                    max_calls: 2,
                    interval_secs: 1.0,
                },
            )
            .await;

        // Consume both tokens.
        assert!(limiter.try_acquire("tool_a").await);
        assert!(limiter.try_acquire("tool_a").await);
        // Bucket is empty.
        assert!(!limiter.try_acquire("tool_a").await);

        // Advancing 0.5s refills exactly one token at 2 tokens/sec.
        tokio::time::advance(std::time::Duration::from_millis(500)).await;
        assert!(limiter.try_acquire("tool_a").await);
        assert!(!limiter.try_acquire("tool_a").await);
    }

    #[tokio::test(start_paused = true)]
    async fn test_acquire_blocks_when_empty() {
        let limiter = Arc::new(RateLimiter::new());
        limiter
            .set_limit(
                "tool_b",
                RateLimit {
                    max_calls: 1,
                    interval_secs: 0.2,
                },
            )
            .await;

        // Drain the single token.
        assert!(limiter.try_acquire("tool_b").await);
        assert!(!limiter.try_acquire("tool_b").await);

        // acquire() remains pending until virtual time reaches the refill.
        let waiting = {
            let limiter = limiter.clone();
            tokio::spawn(async move { limiter.acquire("tool_b").await })
        };
        tokio::task::yield_now().await;
        assert!(
            !waiting.is_finished(),
            "empty bucket must block acquisition"
        );
        tokio::time::advance(std::time::Duration::from_millis(200)).await;
        waiting
            .await
            .expect("acquire task must complete after refill");
    }

    #[tokio::test]
    async fn test_independent_tool_limits() {
        let limiter = RateLimiter::new();
        limiter
            .set_limit(
                "fast",
                RateLimit {
                    max_calls: 10,
                    interval_secs: 1.0,
                },
            )
            .await;
        limiter
            .set_limit(
                "slow",
                RateLimit {
                    max_calls: 1,
                    interval_secs: 1.0,
                },
            )
            .await;

        // Drain the slow bucket.
        assert!(limiter.try_acquire("slow").await);
        assert!(!limiter.try_acquire("slow").await);

        // fast bucket should still have tokens.
        for _ in 0..10 {
            assert!(limiter.try_acquire("fast").await);
        }
        assert!(!limiter.try_acquire("fast").await);
    }

    #[tokio::test]
    async fn test_no_limit_always_passes() {
        let limiter = RateLimiter::new();
        // No limit set for "unconfigured".
        assert!(limiter.try_acquire("unconfigured").await);
        limiter.acquire("unconfigured").await; // should return immediately
    }

    use std::sync::Arc;

    #[tokio::test(start_paused = true)]
    async fn test_max_tokens_cap() {
        let limiter = RateLimiter::new();
        limiter
            .set_limit(
                "capped",
                RateLimit {
                    max_calls: 2,
                    interval_secs: 1.0,
                },
            )
            .await;

        // Drain the initial bucket, then advance through many refill intervals.
        assert!(limiter.try_acquire("capped").await);
        assert!(limiter.try_acquire("capped").await);
        assert!(!limiter.try_acquire("capped").await);
        tokio::time::advance(std::time::Duration::from_secs(10)).await;

        // Refill remains capped at exactly two tokens.
        assert!(limiter.try_acquire("capped").await);
        assert!(limiter.try_acquire("capped").await);
        assert!(!limiter.try_acquire("capped").await);
    }
}