fix(proxy): enforce virtual-key TPM when Redis rate limiter is enabled (#14)

* fix(proxy): record virtual-key TPM usage in redis limiter

* fix: keep redis TPM key only in redis builds

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<RateLimitState>,
     /// Optional model allowlist from the virtual key policy.
     pub(crate) allowed_models: Option<Vec<String>>,
@@ -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(),

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