diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index c6d0478..c2c2371 100644 --- a/crates/proxy/src/main.rs +++ b/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::>(), 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. diff --git a/crates/proxy/src/ratelimit.rs b/crates/proxy/src/ratelimit.rs new file mode 100644 index 0000000..9889d1d --- /dev/null +++ b/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 = 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 { + 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, 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 = 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, 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 = 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 + )); + } +}