feat: configurable Redis rate limit failure policy (open/closed)

Add RATE_LIMIT_FAIL_POLICY env var (open/closed) so operators can choose
whether to allow or reject requests when Redis is unavailable. Defaults
to fail-open for backward compatibility.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/proxy/src/main.rs

@@ -52,6 +52,12 @@ async fn main() {
     let (multi_config, model_router) = config::MultiConfig::load();
     let listen_port = multi_config.listen_port;
 
+    // Wire up WEBHOOK_URLS env var (if LiteLLM callbacks weren't already set).
+    if let Some(cb) = anyllm_proxy::callbacks::CallbackConfig::from_env() {
+        anyllm_proxy::server::routes::set_callbacks(cb);
+        tracing::info!("webhook callbacks configured from WEBHOOK_URLS env var");
+    }
+
     tracing::info!(
         backends = ?multi_config.backends.keys().collect::<Vec<_>>(),
         default = %multi_config.default_backend,
@@ -98,10 +104,11 @@ async fn main() {
     // multiple proxy instances share rate limit state.
     #[cfg(feature = "redis")]
     if let Ok(redis_url) = std::env::var("REDIS_URL") {
-        match anyllm_proxy::ratelimit::RedisRateLimiter::new(&redis_url).await {
+        let fail_policy = anyllm_proxy::ratelimit::RateLimitFailPolicy::from_env();
+        match anyllm_proxy::ratelimit::RedisRateLimiter::new(&redis_url, fail_policy).await {
             Ok(limiter) => {
                 anyllm_proxy::ratelimit::set_redis_rate_limiter(limiter);
-                tracing::info!("Redis distributed rate limiting enabled");
+                tracing::info!(?fail_policy, "Redis distributed rate limiting enabled");
             }
             Err(e) => {
                 tracing::error!("Redis connection failed: {e}. Using local-only rate limiting.");
@@ -256,6 +263,20 @@ async fn main() {
         // Make virtual keys available to the auth middleware.
         anyllm_proxy::server::middleware::set_virtual_keys(virtual_keys.clone());
 
+        let virtual_keys_pruner = virtual_keys.clone();
+        tokio::spawn(async move {
+            let mut interval = tokio::time::interval(std::time::Duration::from_secs(300));
+            loop {
+                interval.tick().await;
+                let now = anyllm_proxy::admin::keys::now_ms();
+                // Check and prune old rate limit states
+                for entry in virtual_keys_pruner.iter() {
+                    let _ = entry.rate_state.check_rpm(0, now);
+                    let _ = entry.rate_state.check_tpm(0, now);
+                }
+            }
+        });
+
         let shared = admin::state::SharedState {
             db: db.clone(),
             events_tx: events_tx.clone(),
@@ -265,6 +286,7 @@ async fn main() {
             log_reload: Some(log_reload),
             config_write_lock: Arc::new(tokio::sync::Mutex::new(())),
             virtual_keys,
+            model_router: model_router.clone(),
         };
 
         // Admin token: use env var or generate random UUID written to a file.

crates/proxy/src/ratelimit.rs

@@ -0,0 +1,270 @@
+//! Distributed rate limiting via Redis sorted sets.
+//!
+//! When `REDIS_URL` is set and the `redis` feature is enabled, RPM/TPM
+//! checks are performed against Redis so multiple proxy instances share
+//! rate limit state. On Redis failure, the proxy falls back to local
+//! in-memory rate limiting (each instance limits independently).
+
+/// Policy for handling Redis rate limiter errors.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum RateLimitFailPolicy {
+    /// Allow requests when Redis is unavailable (default).
+    Open,
+    /// Reject requests when Redis is unavailable.
+    Closed,
+}
+
+impl RateLimitFailPolicy {
+    pub fn from_env_str(s: &str) -> Self {
+        match s.to_lowercase().as_str() {
+            "closed" => Self::Closed,
+            _ => Self::Open,
+        }
+    }
+
+    pub fn from_env() -> Self {
+        std::env::var("RATE_LIMIT_FAIL_POLICY")
+            .map(|v| Self::from_env_str(&v))
+            .unwrap_or(Self::Open)
+    }
+}
+
+#[cfg(feature = "redis")]
+use redis::aio::ConnectionManager;
+#[cfg(feature = "redis")]
+use std::sync::{LazyLock, OnceLock};
+
+#[cfg(feature = "redis")]
+static REDIS_RATE_LIMITER: OnceLock<RedisRateLimiter> = OnceLock::new();
+
+/// Initialize the global Redis rate limiter. Called once from main.
+#[cfg(feature = "redis")]
+pub fn set_redis_rate_limiter(limiter: RedisRateLimiter) {
+    let _ = REDIS_RATE_LIMITER.set(limiter);
+}
+
+/// Get the global Redis rate limiter, if initialized.
+#[cfg(feature = "redis")]
+pub fn get_redis_rate_limiter() -> Option<&'static RedisRateLimiter> {
+    REDIS_RATE_LIMITER.get()
+}
+
+/// Stub when redis feature is not enabled.
+#[cfg(not(feature = "redis"))]
+pub fn get_redis_rate_limiter() -> Option<&'static ()> {
+    None
+}
+
+/// Redis-backed distributed rate limiter using sorted sets.
+///
+/// Keys use the format `anyllm:rl:{key_hash_hex}:rpm` and `anyllm:rl:{key_hash_hex}:tpm`.
+/// Each request is a member scored by its timestamp in milliseconds.
+/// A Lua script atomically trims expired entries, checks the count/sum,
+/// and adds the new entry if within limits.
+#[cfg(feature = "redis")]
+pub struct RedisRateLimiter {
+    conn: ConnectionManager,
+    fail_policy: RateLimitFailPolicy,
+}
+
+#[cfg(feature = "redis")]
+impl RedisRateLimiter {
+    /// Connect to Redis and create a rate limiter.
+    pub async fn new(redis_url: &str, fail_policy: RateLimitFailPolicy) -> Result<Self, redis::RedisError> {
+        let client = redis::Client::open(redis_url)?;
+        let conn = ConnectionManager::new(client).await?;
+        Ok(Self { conn, fail_policy })
+    }
+
+    /// Get the underlying connection manager for reuse (e.g., by cache layer).
+    pub fn connection(&self) -> &ConnectionManager {
+        &self.conn
+    }
+
+    /// Check RPM limit. Returns Ok(()) if allowed, Err(retry_after_secs) if exceeded.
+    /// On Redis error, returns Ok(()) to allow the request (fail-open).
+    pub async fn check_rpm(&self, key_hash_hex: &str, limit: u32, now_ms: u64) -> Result<(), u64> {
+        let redis_key = format!("anyllm:rl:{key_hash_hex}:rpm");
+        match self.check_rpm_inner(&redis_key, limit, now_ms).await {
+            Ok(result) => result,
+            Err(e) => match self.fail_policy {
+                RateLimitFailPolicy::Open => {
+                    tracing::warn!(error = %e, "Redis RPM check failed, allowing request (fail-open)");
+                    Ok(())
+                }
+                RateLimitFailPolicy::Closed => {
+                    tracing::error!(error = %e, "Redis RPM check failed, rejecting request (fail-closed)");
+                    Err(1)
+                }
+            },
+        }
+    }
+
+    async fn check_rpm_inner(
+        &self,
+        redis_key: &str,
+        limit: u32,
+        now_ms: u64,
+    ) -> Result<Result<(), u64>, redis::RedisError> {
+        let mut conn = self.conn.clone();
+        let cutoff = now_ms.saturating_sub(60_000);
+        let member_id = format!("{now_ms}:{}", uuid::Uuid::new_v4().as_simple());
+
+        // Hashed once at first use; avoids re-computing SHA1 per request.
+        static RPM_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| {
+            redis::Script::new(
+                r#"
+                redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
+                local count = redis.call('ZCARD', KEYS[1])
+                if count >= tonumber(ARGV[2]) then
+                    local oldest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
+                    if oldest and #oldest >= 2 then
+                        return oldest[2]
+                    end
+                    return tostring(ARGV[3])
+                end
+                redis.call('ZADD', KEYS[1], ARGV[3], ARGV[4])
+                redis.call('EXPIRE', KEYS[1], 120)
+                return 0
+                "#,
+            )
+        });
+
+        let result: i64 = RPM_SCRIPT
+            .key(redis_key)
+            .arg(cutoff)
+            .arg(limit)
+            .arg(now_ms)
+            .arg(&member_id)
+            .invoke_async(&mut conn)
+            .await?;
+
+        if result == 0 {
+            Ok(Ok(()))
+        } else {
+            let oldest_ms = result as u64;
+            let retry_after_ms = (oldest_ms + 60_000).saturating_sub(now_ms);
+            Ok(Err((retry_after_ms / 1000).max(1)))
+        }
+    }
+
+    /// Check TPM limit. Returns Ok(()) if allowed, Err(retry_after_secs) if exceeded.
+    /// On Redis error, returns Ok(()) (fail-open).
+    pub async fn check_tpm(&self, key_hash_hex: &str, limit: u32, now_ms: u64) -> Result<(), u64> {
+        let redis_key = format!("anyllm:rl:{key_hash_hex}:tpm");
+        match self.check_tpm_inner(&redis_key, limit, now_ms).await {
+            Ok(result) => result,
+            Err(e) => match self.fail_policy {
+                RateLimitFailPolicy::Open => {
+                    tracing::warn!(error = %e, "Redis TPM check failed, allowing request (fail-open)");
+                    Ok(())
+                }
+                RateLimitFailPolicy::Closed => {
+                    tracing::error!(error = %e, "Redis TPM check failed, rejecting request (fail-closed)");
+                    Err(1)
+                }
+            },
+        }
+    }
+
+    async fn check_tpm_inner(
+        &self,
+        redis_key: &str,
+        limit: u32,
+        now_ms: u64,
+    ) -> Result<Result<(), u64>, redis::RedisError> {
+        let mut conn = self.conn.clone();
+        let cutoff = now_ms.saturating_sub(60_000);
+
+        // For TPM, members are scored by timestamp and the member value encodes the token count.
+        // We sum member names (which are "{tokens}:{uuid}") to get total tokens.
+        static TPM_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| {
+            redis::Script::new(
+                r#"
+                redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
+                local members = redis.call('ZRANGE', KEYS[1], 0, -1)
+                local total = 0
+                for _, m in ipairs(members) do
+                    local tokens = tonumber(string.match(m, '^(%d+):'))
+                    if tokens then total = total + tokens end
+                end
+                if total >= tonumber(ARGV[2]) then
+                    local oldest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
+                    if oldest and #oldest >= 2 then
+                        return oldest[2]
+                    end
+                    return tostring(ARGV[3])
+                end
+                return 0
+                "#,
+            )
+        });
+
+        let result: i64 = TPM_SCRIPT
+            .key(redis_key)
+            .arg(cutoff)
+            .arg(limit)
+            .arg(now_ms)
+            .invoke_async(&mut conn)
+            .await?;
+
+        if result == 0 {
+            Ok(Ok(()))
+        } else {
+            let oldest_ms = result as u64;
+            let retry_after_ms = (oldest_ms + 60_000).saturating_sub(now_ms);
+            Ok(Err((retry_after_ms / 1000).max(1)))
+        }
+    }
+
+    /// Record TPM tokens after a response is received.
+    pub async fn record_tpm(&self, key_hash_hex: &str, now_ms: u64, tokens: u32) {
+        let redis_key = format!("anyllm:rl:{key_hash_hex}:tpm");
+        let member = format!("{tokens}:{}", uuid::Uuid::new_v4().as_simple());
+        let mut conn = self.conn.clone();
+        let result: Result<(), redis::RedisError> = redis::pipe()
+            .zadd(&redis_key, member, now_ms as f64)
+            .expire(&redis_key, 120)
+            .query_async(&mut conn)
+            .await;
+        if let Err(e) = result {
+            tracing::warn!(error = %e, "Redis TPM record failed");
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::RateLimitFailPolicy;
+
+    #[test]
+    fn get_redis_rate_limiter_returns_none_without_init() {
+        // When redis feature is not enabled, or when not initialized,
+        // the function should return None.
+        assert!(super::get_redis_rate_limiter().is_none());
+    }
+
+    #[test]
+    fn parse_rate_limit_fail_policy() {
+        assert!(matches!(
+            RateLimitFailPolicy::from_env_str("open"),
+            RateLimitFailPolicy::Open
+        ));
+        assert!(matches!(
+            RateLimitFailPolicy::from_env_str("closed"),
+            RateLimitFailPolicy::Closed
+        ));
+        assert!(matches!(
+            RateLimitFailPolicy::from_env_str("OPEN"),
+            RateLimitFailPolicy::Open
+        ));
+        assert!(matches!(
+            RateLimitFailPolicy::from_env_str("CLOSED"),
+            RateLimitFailPolicy::Closed
+        ));
+        assert!(matches!(
+            RateLimitFailPolicy::from_env_str("unknown"),
+            RateLimitFailPolicy::Open
+        ));
+    }
+}