diff --git a/crates/proxy/src/server/middleware.rs b/crates/proxy/src/server/middleware.rs index d7d41d6..4c81b76 100644 --- a/crates/proxy/src/server/middleware.rs +++ b/crates/proxy/src/server/middleware.rs @@ -45,6 +45,9 @@ fn rate_limit_response(message: &str, retry_after: u64) -> Response { pub struct VirtualKeyContext { /// Database row ID for the virtual key (used for cost accumulation). pub(crate) key_id: i64, + /// Hex-encoded credential hash used as stable distributed rate-limit key. + #[cfg(feature = "redis")] + pub(crate) key_hash_hex: String, pub(crate) rate_state: Arc, /// Optional model allowlist from the virtual key policy. pub(crate) allowed_models: Option>, @@ -343,18 +346,20 @@ pub async fn validate_auth( } } + #[cfg(feature = "redis")] + let key_hash_hex: String = credential_hash.iter().map(|b| format!("{b:02x}")).collect(); + // Enforce TPM limit pre-check if let Some(tpm_limit) = meta.tpm_limit { #[allow(unused_mut, unused_variables)] let mut checked_ext = false; #[cfg(feature = "redis")] { - let hash_hex: String = - credential_hash.iter().map(|b| format!("{b:02x}")).collect(); if let Some(redis_limiter) = crate::ratelimit::get_redis_rate_limiter() { checked_ext = true; - if let Err(retry_after) = - redis_limiter.check_tpm(&hash_hex, tpm_limit, now_ms).await + if let Err(retry_after) = redis_limiter + .check_tpm(&key_hash_hex, tpm_limit, now_ms) + .await { return Err(rate_limit_response( "Token rate limit exceeded for this API key.", @@ -410,6 +415,8 @@ pub async fn validate_auth( // Always insert context for post-response TPM recording and cost tracking. request.extensions_mut().insert(VirtualKeyContext { key_id: meta.id, + #[cfg(feature = "redis")] + key_hash_hex, rate_state: meta.rate_state.clone(), allowed_models: meta.allowed_models.clone(), allowed_routes: meta.allowed_routes.clone(), diff --git a/crates/proxy/src/server/routes.rs b/crates/proxy/src/server/routes.rs index 8af3bee..b4260ba 100644 --- a/crates/proxy/src/server/routes.rs +++ b/crates/proxy/src/server/routes.rs @@ -1153,8 +1153,18 @@ pub(crate) fn record_vk_tpm( output_tokens: u32, ) { if let Some(ctx) = vk_ctx { - ctx.rate_state - .record_tpm(crate::admin::keys::now_ms(), output_tokens); + let now_ms = crate::admin::keys::now_ms(); + ctx.rate_state.record_tpm(now_ms, output_tokens); + + #[cfg(feature = "redis")] + if let Some(redis_limiter) = crate::ratelimit::get_redis_rate_limiter() { + let key_hash_hex = ctx.key_hash_hex.clone(); + tokio::spawn(async move { + redis_limiter + .record_tpm(&key_hash_hex, now_ms, output_tokens) + .await; + }); + } } }