fix(proxy): reconcile rust realtime budget reservations

litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs

@@ -14,6 +14,7 @@ use axum::extract::FromRequestParts;
 use axum::http::header::AUTHORIZATION;
 use axum::http::request::Parts;
 use axum::http::StatusCode;
+use serde_json::Value;
 use sha2::{Digest, Sha256};
 use subtle::ConstantTimeEq;
 
@@ -40,6 +41,7 @@ pub struct UserApiKeyAuth {
     pub models: Vec<String>,
     pub tpm_limit: Option<i64>,
     pub rpm_limit: Option<i64>,
+    pub budget_reservation: Option<Value>,
 }
 
 impl UserApiKeyAuth {
@@ -132,17 +134,25 @@ impl FromRequestParts<AppState> for UserApiKeyAuth {
         let model = requested_model(parts.uri.query())
             .map_err(|message| (StatusCode::BAD_REQUEST, message))?;
         let model = model.as_deref();
+        let use_auth_cache = route != "/v1/realtime";
 
         // Virtual key: serve from cache, else verify via the (swappable) backend
         // and cache the result keyed by (route, model, key) hash.
         let hash = key_hash(token, route, model);
-        if let Some(cached) = state.key_cache.get(&hash) {
-            return Ok(cached);
+        if use_auth_cache {
+            if let Some(cached) = state.key_cache.get(&hash) {
+                return Ok(cached);
+            }
         }
 
         match state.authenticator.verify(token, route, model).await {
             Ok(auth) => {
-                state.key_cache.insert(hash, auth.clone());
+                // Budget reservations are request-scoped admission state. Reusing
+                // one from the auth cache would let later sessions bypass a fresh
+                // reservation and then reconcile the same reservation twice.
+                if use_auth_cache && auth.budget_reservation.is_none() {
+                    state.key_cache.insert(hash, auth.clone());
+                }
                 Ok(auth)
             }
             Err(AuthError::Unauthorized) => {
@@ -176,7 +186,9 @@ mod tests {
     use crate::io::realtime_pool::RealtimePool;
     use axum::http::Request;
     use litellm_core::router::Router;
+    use std::sync::atomic::{AtomicUsize, Ordering};
     use std::sync::Arc;
+    use std::time::Duration;
 
     /// Stub authenticator so tests never touch the network. Records nothing; just
     /// returns a fixed identity (or unauthorized) per construction.
@@ -201,18 +213,34 @@ mod tests {
     }
 
     fn test_state(master_key: Option<&str>, authenticator: Arc<dyn KeyAuthenticator>) -> AppState {
+        test_state_with_cache(master_key, authenticator, Arc::new(KeyCache::new()))
+    }
+
+    fn test_state_with_cache(
+        master_key: Option<&str>,
+        authenticator: Arc<dyn KeyAuthenticator>,
+        key_cache: Arc<KeyCache>,
+    ) -> AppState {
         AppState {
             router: Arc::new(Router::new(vec![])),
             master_key: master_key.map(Arc::from),
             loggers: Arc::new(Vec::new()),
             realtime_pool: RealtimePool::disabled(),
             authenticator,
-            key_cache: Arc::new(KeyCache::new()),
+            key_cache,
         }
     }
 
     async fn extract(state: &AppState, header: Option<&str>) -> Result<UserApiKeyAuth, StatusCode> {
-        let mut builder = Request::builder().uri("/");
+        extract_uri(state, header, "/").await
+    }
+
+    async fn extract_uri(
+        state: &AppState,
+        header: Option<&str>,
+        uri: &str,
+    ) -> Result<UserApiKeyAuth, StatusCode> {
+        let mut builder = Request::builder().uri(uri);
         if let Some(value) = header {
             builder = builder.header(AUTHORIZATION, value);
         }
@@ -255,7 +283,11 @@ mod tests {
                 ..UserApiKeyAuth::default()
             }),
         });
-        let state = test_state(Some("sk-master"), stub);
+        let state = test_state_with_cache(
+            Some("sk-master"),
+            stub,
+            Arc::new(KeyCache::with_ttl(Duration::from_secs(60))),
+        );
 
         let auth = extract(&state, Some("Bearer sk-virtual")).await.unwrap();
         assert_eq!(auth.user_id.as_deref(), Some("u-1"));
@@ -264,6 +296,102 @@ mod tests {
         assert_eq!(cached.user_id.as_deref(), Some("u-1"));
     }
 
+    #[tokio::test]
+    async fn auth_with_budget_reservation_is_not_cached() {
+        struct CountingAuthenticator {
+            calls: AtomicUsize,
+        }
+
+        #[axum::async_trait]
+        impl KeyAuthenticator for CountingAuthenticator {
+            async fn verify(
+                &self,
+                _key: &str,
+                _route: &str,
+                _model: Option<&str>,
+            ) -> Result<UserApiKeyAuth, AuthError> {
+                let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
+                Ok(UserApiKeyAuth {
+                    user_id: Some(format!("u-{call}")),
+                    budget_reservation: Some(serde_json::json!({
+                        "reserved_cost": 0.5,
+                        "entries": [{"counter_key": "spend:key:hashed-key"}],
+                        "finalized": false,
+                        "input_cost": 0.1
+                    })),
+                    ..UserApiKeyAuth::default()
+                })
+            }
+        }
+
+        let authenticator = Arc::new(CountingAuthenticator {
+            calls: AtomicUsize::new(0),
+        });
+        let state = test_state_with_cache(
+            Some("sk-master"),
+            authenticator.clone(),
+            Arc::new(KeyCache::with_ttl(Duration::from_secs(60))),
+        );
+
+        let first = extract(&state, Some("Bearer sk-virtual")).await.unwrap();
+        let second = extract(&state, Some("Bearer sk-virtual")).await.unwrap();
+
+        assert_eq!(first.user_id.as_deref(), Some("u-1"));
+        assert_eq!(second.user_id.as_deref(), Some("u-2"));
+        assert_eq!(authenticator.calls.load(Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn realtime_auth_bypasses_cache() {
+        struct CountingAuthenticator {
+            calls: AtomicUsize,
+        }
+
+        #[axum::async_trait]
+        impl KeyAuthenticator for CountingAuthenticator {
+            async fn verify(
+                &self,
+                _key: &str,
+                _route: &str,
+                _model: Option<&str>,
+            ) -> Result<UserApiKeyAuth, AuthError> {
+                let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
+                Ok(UserApiKeyAuth {
+                    user_id: Some(format!("u-{call}")),
+                    ..UserApiKeyAuth::default()
+                })
+            }
+        }
+
+        let authenticator = Arc::new(CountingAuthenticator {
+            calls: AtomicUsize::new(0),
+        });
+        let state = test_state_with_cache(
+            Some("sk-master"),
+            authenticator.clone(),
+            Arc::new(KeyCache::with_ttl(Duration::from_secs(60))),
+        );
+
+        let first = extract_uri(
+            &state,
+            Some("Bearer sk-virtual"),
+            "/v1/realtime?model=gpt-realtime",
+        )
+        .await
+        .unwrap();
+        let second = extract_uri(
+            &state,
+            Some("Bearer sk-virtual"),
+            "/v1/realtime?model=gpt-realtime",
+        )
+        .await
+        .unwrap();
+
+        assert_eq!(first.user_id.as_deref(), Some("u-1"));
+        assert_eq!(second.user_id.as_deref(), Some("u-2"));
+        assert_eq!(authenticator.calls.load(Ordering::SeqCst), 2);
+    }
+
     #[tokio::test]
     async fn unauthorized_virtual_key_maps_to_401() {
         let stub = Arc::new(StubAuthenticator {
@@ -290,7 +418,13 @@ mod tests {
             "blocked": false,
             "models": ["gpt-4o", "gpt-4o-mini"],
             "tpm_limit": 1000,
-            "rpm_limit": 60
+            "rpm_limit": 60,
+            "budget_reservation": {
+                "reserved_cost": 0.5,
+                "entries": [{"counter_key": "spend:key:hashed-abc"}],
+                "finalized": false,
+                "input_cost": 0.1
+            }
         });
 
         let auth: UserApiKeyAuth = serde_json::from_value(body).unwrap();
@@ -306,6 +440,13 @@ mod tests {
         assert_eq!(auth.models, vec!["gpt-4o", "gpt-4o-mini"]);
         assert_eq!(auth.tpm_limit, Some(1000));
         assert_eq!(auth.rpm_limit, Some(60));
+        assert_eq!(
+            auth.budget_reservation
+                .as_ref()
+                .and_then(|reservation| reservation.get("reserved_cost"))
+                .and_then(Value::as_f64),
+            Some(0.5)
+        );
     }
 
     #[test]

litellm-rust/crates/ai-gateway/src/integrations/types.rs

@@ -26,6 +26,7 @@ pub struct RequestMetadata {
     pub user_api_key_hash: Option<String>,
     pub user_api_key_user_id: Option<String>,
     pub user_api_key_team_id: Option<String>,
+    pub user_api_key_budget_reservation: Option<Value>,
 }
 
 /// A logging-callback failure (e.g. a custom logger raised). Mirrors the Python
@@ -141,6 +142,8 @@ pub struct StandardLoggingMetadata {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub user_api_key_end_user_id: Option<String>, // -> SpendLogs.end_user
     #[serde(skip_serializing_if = "Option::is_none")]
+    pub user_api_key_budget_reservation: Option<Value>, // -> budget reservation reconciliation
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub spend_logs_metadata: Option<HashMap<String, Value>>,
 }
 

litellm-rust/crates/ai-gateway/src/realtime/streaming.rs

@@ -174,6 +174,10 @@ impl RealTimeStreaming {
                 user_api_key_hash: self.metadata.user_api_key_hash.clone(),
                 user_api_key_user_id: self.metadata.user_api_key_user_id.clone(),
                 user_api_key_team_id: self.metadata.user_api_key_team_id.clone(),
+                user_api_key_budget_reservation: self
+                    .metadata
+                    .user_api_key_budget_reservation
+                    .clone(),
                 ..Default::default()
             },
             messages: None,
@@ -252,6 +256,12 @@ mod tests {
                 user_api_key_hash: Some("hash123".to_string()),
                 user_api_key_user_id: Some("user-1".to_string()),
                 user_api_key_team_id: Some("team-1".to_string()),
+                user_api_key_budget_reservation: Some(serde_json::json!({
+                    "reserved_cost": 0.5,
+                    "entries": [{"counter_key": "spend:key:hash123"}],
+                    "finalized": false,
+                    "input_cost": 0.1
+                })),
             },
         );
 
@@ -283,6 +293,15 @@ mod tests {
             payload.metadata.user_api_key_hash.as_deref(),
             Some("hash123")
         );
+        assert_eq!(
+            payload
+                .metadata
+                .user_api_key_budget_reservation
+                .as_ref()
+                .and_then(|reservation| reservation.get("reserved_cost"))
+                .and_then(Value::as_f64),
+            Some(0.5)
+        );
 
         streaming.log_messages(SessionStatus::Success);
         assert_eq!(logger.calls.load(Ordering::SeqCst), 1);
@@ -319,6 +338,10 @@ mod tests {
             json.contains("\"response_cost\""),
             "missing response_cost: {json}"
         );
+        assert!(
+            !json.contains("user_api_key_budget_reservation"),
+            "missing reservation should be omitted: {json}"
+        );
         assert_eq!(payload.response_cost, 0.0042);
     }
 

litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs

@@ -63,6 +63,7 @@ fn request_metadata_for_auth(auth: &UserApiKeyAuth, master_key: Option<&str>) ->
         user_api_key_hash,
         user_api_key_user_id: auth.user_id.clone(),
         user_api_key_team_id: auth.team_id.clone(),
+        user_api_key_budget_reservation: auth.budget_reservation.clone(),
     }
 }
 
@@ -185,6 +186,12 @@ mod tests {
                 api_key: Some("hashed-key".to_string()),
                 user_id: Some("user-1".to_string()),
                 team_id: Some("team-1".to_string()),
+                budget_reservation: Some(serde_json::json!({
+                    "reserved_cost": 0.5,
+                    "entries": [{"counter_key": "spend:key:hashed-key"}],
+                    "finalized": false,
+                    "input_cost": 0.1
+                })),
                 ..UserApiKeyAuth::default()
             },
             Some("sk-master"),
@@ -193,6 +200,14 @@ mod tests {
         assert_eq!(metadata.user_api_key_hash.as_deref(), Some("hashed-key"));
         assert_eq!(metadata.user_api_key_user_id.as_deref(), Some("user-1"));
         assert_eq!(metadata.user_api_key_team_id.as_deref(), Some("team-1"));
+        assert_eq!(
+            metadata
+                .user_api_key_budget_reservation
+                .as_ref()
+                .and_then(|reservation| reservation.get("reserved_cost"))
+                .and_then(serde_json::Value::as_f64),
+            Some(0.5)
+        );
     }
 
     #[test]

litellm/proxy/rust_control_plane_endpoints/auth_endpoints.py

@@ -78,12 +78,7 @@ def _synthetic_request(
         "client": ("127.0.0.1", 0),
         "server": ("127.0.0.1", 4000),
     }
-    request = Request(scope, receive)
-    # Admission checks only. Realtime spend is reconciled later through callback
-    # logs, so optimistic reservation here would have no request lifecycle to
-    # release it.
-    request.state.skip_budget_reservation = True
-    return request
+    return Request(scope, receive)
 
 
 @router.post(
@@ -124,4 +119,7 @@ async def verify_key(body: VerifyKeyRequest) -> dict[str, Any]:
             raise
         raise HTTPException(status_code=401, detail="invalid api key")
 
-    return auth.model_dump(exclude_none=True, mode="json")
+    auth_response = auth.model_dump(exclude_none=True, mode="json")
+    if auth.budget_reservation is not None:
+        auth_response["budget_reservation"] = auth.budget_reservation
+    return auth_response

litellm/proxy/rust_control_plane_endpoints/logging_endpoints.py

@@ -94,6 +94,9 @@ class CallbackLogsReplayer:
             "user_api_key_team_id": metadata.get("user_api_key_team_id"),
             "user_api_key_org_id": metadata.get("user_api_key_org_id"),
             "user_api_key_end_user_id": metadata.get("user_api_key_end_user_id"),
+            "user_api_key_budget_reservation": metadata.get(
+                "user_api_key_budget_reservation"
+            ),
             "spend_logs_metadata": metadata.get("spend_logs_metadata"),
         }
 

tests/test_litellm/proxy/rust_control_plane_endpoints/test_auth_endpoints.py

@@ -87,7 +87,7 @@ def test_router_mounts_auth_verify_under_rust_control_plane():
 
 
 @pytest.mark.asyncio
-async def test_synthetic_request_skips_budget_reservation():
+async def test_synthetic_request_allows_budget_reservation():
     request = _synthetic_request(
         route="/v1/realtime",
         authorization_header="Bearer sk-test-key",
@@ -95,7 +95,7 @@ async def test_synthetic_request_skips_budget_reservation():
     )
 
     assert request.url.path == "/v1/realtime"
-    assert request.state.skip_budget_reservation is True
+    assert getattr(request.state, "skip_budget_reservation", False) is False
     assert (await request.json()) == {"model": "gpt-realtime"}
 
 
@@ -133,6 +133,39 @@ async def test_verify_key_returns_model_dump(monkeypatch):
     assert result["user_id"] == "user-123"
 
 
+@pytest.mark.asyncio
+async def test_verify_key_returns_budget_reservation(monkeypatch):
+    budget_reservation = {
+        "reserved_cost": 0.5,
+        "entries": [{"counter_key": "spend:key:hashed-key"}],
+        "finalized": False,
+        "input_cost": 0.1,
+    }
+    expected_auth = UserAPIKeyAuth(
+        api_key="hashed-key",
+        user_id="user-123",
+        budget_reservation=budget_reservation,
+    )
+
+    async def fake_user_api_key_auth(request, api_key):
+        return expected_auth
+
+    monkeypatch.setattr(
+        "litellm.proxy.rust_control_plane_endpoints.auth_endpoints.user_api_key_auth",
+        fake_user_api_key_auth,
+    )
+
+    body = VerifyKeyRequest(
+        api_key="sk-test-key", route="/v1/realtime", model="gpt-realtime"
+    )
+    result = await verify_key(body=body)
+
+    assert "budget_reservation" not in expected_auth.model_dump(
+        exclude_none=True, mode="json"
+    )
+    assert result["budget_reservation"] == budget_reservation
+
+
 @pytest.mark.asyncio
 async def test_verify_key_omits_model_when_absent(monkeypatch):
     captured = {}

tests/test_litellm/proxy/rust_control_plane_endpoints/test_logging_endpoints.py

@@ -17,6 +17,12 @@ from litellm.types.proxy.callback_logs_endpoints import (
 )
 
 REQ_ID = "cb-logs-unit-test-1"
+BUDGET_RESERVATION = {
+    "reserved_cost": 0.5,
+    "entries": [{"counter_key": "spend:key:rust-gateway-test-key"}],
+    "finalized": False,
+    "input_cost": 0.1,
+}
 
 
 def _sample_payload(**overrides):
@@ -37,6 +43,7 @@ def _sample_payload(**overrides):
             "user_api_key_hash": "rust-gateway-test-key",
             "user_api_key_user_id": "user-cb-logs-test",
             "user_api_key_team_id": "team-cb-logs-test",
+            "user_api_key_budget_reservation": BUDGET_RESERVATION,
         },
         "messages": [{"role": "user", "content": "hi"}],
     }
@@ -63,6 +70,7 @@ def test_build_logging_obj_seeds_model_call_details():
     assert md["user_api_key"] == "rust-gateway-test-key"
     assert md["user_api_key_user_id"] == "user-cb-logs-test"
     assert md["user_api_key_team_id"] == "team-cb-logs-test"
+    assert md["user_api_key_budget_reservation"] is BUDGET_RESERVATION
 
 
 def test_response_obj_carries_usage():