From 65760a0264901796bd7d9fc60c68beba1b26b444 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 24 Jun 2026 18:03:12 -0700 Subject: [PATCH] fix(proxy): reconcile rust realtime budget reservations --- .../ai-gateway/src/auth/user_api_key.rs | 155 +++++++++++++++++- .../ai-gateway/src/integrations/types.rs | 3 + .../ai-gateway/src/realtime/streaming.rs | 23 +++ .../ai-gateway/src/routes/realtime/mod.rs | 15 ++ .../auth_endpoints.py | 12 +- .../logging_endpoints.py | 3 + .../test_auth_endpoints.py | 37 ++++- .../test_logging_endpoints.py | 8 + 8 files changed, 240 insertions(+), 16 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs b/litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs index 15f2b16e91f..44c95bff247 100644 --- a/litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs +++ b/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, pub tpm_limit: Option, pub rpm_limit: Option, + pub budget_reservation: Option, } impl UserApiKeyAuth { @@ -132,17 +134,25 @@ impl FromRequestParts 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) -> AppState { + test_state_with_cache(master_key, authenticator, Arc::new(KeyCache::new())) + } + + fn test_state_with_cache( + master_key: Option<&str>, + authenticator: Arc, + key_cache: Arc, + ) -> 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 { - let mut builder = Request::builder().uri("/"); + extract_uri(state, header, "/").await + } + + async fn extract_uri( + state: &AppState, + header: Option<&str>, + uri: &str, + ) -> Result { + 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 { + 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 { + 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] diff --git a/litellm-rust/crates/ai-gateway/src/integrations/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/types.rs index d61a1f816a7..eb18c7608a5 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/types.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/types.rs @@ -26,6 +26,7 @@ pub struct RequestMetadata { pub user_api_key_hash: Option, pub user_api_key_user_id: Option, pub user_api_key_team_id: Option, + pub user_api_key_budget_reservation: Option, } /// 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, // -> SpendLogs.end_user #[serde(skip_serializing_if = "Option::is_none")] + pub user_api_key_budget_reservation: Option, // -> budget reservation reconciliation + #[serde(skip_serializing_if = "Option::is_none")] pub spend_logs_metadata: Option>, } diff --git a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs index 34c82897808..69390b9f459 100644 --- a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs +++ b/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); } diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs index cf484274268..ea421f2d96c 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/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] diff --git a/litellm/proxy/rust_control_plane_endpoints/auth_endpoints.py b/litellm/proxy/rust_control_plane_endpoints/auth_endpoints.py index b1b38287a74..b6962df0915 100644 --- a/litellm/proxy/rust_control_plane_endpoints/auth_endpoints.py +++ b/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 diff --git a/litellm/proxy/rust_control_plane_endpoints/logging_endpoints.py b/litellm/proxy/rust_control_plane_endpoints/logging_endpoints.py index e39707f9a42..0475a1d6ff0 100644 --- a/litellm/proxy/rust_control_plane_endpoints/logging_endpoints.py +++ b/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"), } diff --git a/tests/test_litellm/proxy/rust_control_plane_endpoints/test_auth_endpoints.py b/tests/test_litellm/proxy/rust_control_plane_endpoints/test_auth_endpoints.py index 57519bf62eb..07bd03d4a25 100644 --- a/tests/test_litellm/proxy/rust_control_plane_endpoints/test_auth_endpoints.py +++ b/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 = {} diff --git a/tests/test_litellm/proxy/rust_control_plane_endpoints/test_logging_endpoints.py b/tests/test_litellm/proxy/rust_control_plane_endpoints/test_logging_endpoints.py index 3eba6ca5dc5..a5da8f4b855 100644 --- a/tests/test_litellm/proxy/rust_control_plane_endpoints/test_logging_endpoints.py +++ b/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():