fix(ai-gateway): close realtime auth review gaps

litellm-rust/Cargo.lock

@@ -571,6 +571,7 @@ dependencies = [
  "reqwest",
  "serde",
  "serde_json",
+ "serde_urlencoded",
  "sha2",
  "subtle",
  "tokio",

litellm-rust/Cargo.toml

@@ -20,6 +20,7 @@ rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+serde_urlencoded = "0.7"
 sha2 = "0.10"
 subtle = "2"
 thiserror = "2.0"

litellm-rust/crates/ai-gateway/Cargo.toml

@@ -22,13 +22,14 @@ futures-util.workspace = true
 serde_json.workspace = true
 axum = { workspace = true, features = ["ws"], optional = true }
 serde = { workspace = true, optional = true }
+serde_urlencoded = { workspace = true, optional = true }
 sha2 = { workspace = true, optional = true }
 subtle = { workspace = true, optional = true }
 pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
 
 [features]
 default = []
-server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"]
+server = ["dep:axum", "dep:subtle", "dep:serde", "dep:serde_urlencoded", "dep:sha2"]
 # Build the gateway's config from the proxy YAML via an embedded Python
 # interpreter (links libpython; requires `litellm` importable at runtime).
 python-config = ["dep:pyo3"]

litellm-rust/crates/ai-gateway/README.md

@@ -17,7 +17,7 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
 Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
 
 - **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
-- **Auth:** `Authorization: Bearer <key>` — the master key (admin) or a virtual key validated by the LiteLLM proxy (see below). Fails closed.
+- **Auth:** `Authorization: Bearer <key>` — proxy-admin only for realtime serving today. Non-admin virtual keys are verified but rejected until realtime usage is reported back to the LiteLLM proxy. Fails closed.
 - **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil`
 
 > **Realtime serving is pure Rust.** Python is used at **load time only** — to
@@ -42,7 +42,9 @@ route/model permissions are enforced in exactly one place. A client presents
   load. A bounded ≤200-entry cache is available for future high-RPS *per-request*
   routes via `LITELLM_AUTH_CACHE_TTL_SECS > 0`, which trades budget/rate-limit
   freshness (bounded by the TTL) for fewer control-plane calls — like the proxy's
-  own ~60s auth cache.
+  own ~60s auth cache. Non-admin virtual-key identities are rejected on the
+  realtime route until realtime usage is reported back to the proxy, so callers
+  cannot consume upstream usage without key/team spend attribution.
 
 Data plane → control plane is itself authenticated with a **dedicated data-plane key**
 (NOT the master key — least privilege): the gateway sends
@@ -73,9 +75,12 @@ export LITELLM_AUTH_VERIFY_URL=https://<proxy-host>/v1/rust_control_plane/authen
 ```
 
 Fails closed: a missing/wrong data-plane key, an unreachable proxy, or a rejected key
-all yield `401`. Revocation and budget changes take effect within the cache TTL
-(if caching is enabled via `LITELLM_AUTH_CACHE_TTL_SECS`; off by default → every connection re-verifies). Keep the proxy on a private network —
-the verify endpoint is internal-only and excluded from the public OpenAPI spec.
+all yield `401`; a valid non-admin virtual key currently yields `403` on
+`/v1/realtime` until realtime spend reporting is available. Revocation and budget
+changes take effect within the cache TTL (if caching is enabled via
+`LITELLM_AUTH_CACHE_TTL_SECS`; off by default → every connection re-verifies).
+Keep the proxy on a private network — the verify endpoint is internal-only and
+excluded from the public OpenAPI spec.
 
 ## Configuration (config.yaml)
 
@@ -117,7 +122,7 @@ overridden at deploy time (e.g. a Render secret file mounted at the same path).
 | `LITELLM_CONFIG_PATH` | yes (config mode) | — | Path to the config.yaml the gateway loads its `model_list` from. The Docker image defaults this to `/app/config.yaml`. |
 | `LITELLM_MASTER_KEY` | yes | — | Admin bearer token (checked locally, no proxy call). Unset ⇒ the master-key path is disabled. |
 | `LITELLM_DATA_PLANE_KEY` | for virtual keys | — | Dedicated secret the gateway sends as `X-LiteLLM-Data-Plane-Key` to authenticate itself to the proxy's verify endpoint. **Must match the proxy's `LITELLM_DATA_PLANE_KEY`.** Not the master key. |
-| `LITELLM_AUTH_VERIFY_URL` | for virtual keys | `http://localhost:4000/v1/rust_control_plane/authentication` | The proxy's verify endpoint the gateway delegates virtual-key auth to. |
+| `LITELLM_AUTH_VERIFY_URL` | for virtual keys | `http://localhost:4000/v1/rust_control_plane/authentication` | The proxy's verify endpoint the gateway delegates virtual-key auth to. Non-admin virtual keys are still rejected by `/v1/realtime` until realtime spend reporting is wired. |
 | `LITELLM_AUTH_CACHE_TTL_SECS` | no | `0` | Verified-key cache TTL. **`0` = off (default)** → every connection re-verifies (budget/rate-limit enforced each time); cheap for realtime since auth is per-connection. Set `> 0` only for high-RPS per-request routes, trading budget/rate-limit freshness for fewer proxy calls. |
 | `OPENAI_API_KEY` | yes | — | Upstream OpenAI key. Referenced by config.yaml as `os.environ/OPENAI_API_KEY` for the gateway→OpenAI dial. |
 | `HOST` | no | `127.0.0.1` | **Set to `0.0.0.0` in any container/deploy** or external traffic is refused. |

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

@@ -1,9 +1,10 @@
 //! Gateway authentication, as an axum **extractor** (the idiomatic pattern —
 //! keeps handlers clean and auth testable).
 //!
-//! For now this is a single **master key**: any caller presenting it as
-//! `Authorization: Bearer <key>` may invoke the gateway. Per-key auth, budgets,
-//! and rate limits are delegated to the Python proxy in a later phase.
+//! For now this supports both the local **master key** and virtual keys resolved
+//! through the Python control plane. Individual routes can still narrow the
+//! accepted identities. For example, realtime rejects non-admin virtual keys
+//! until realtime usage is reported back to the proxy for spend attribution.
 //!
 //! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then
 //! runs during extraction, before the handler body. Routes never re-implement it.

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

@@ -52,6 +52,10 @@ impl UserApiKeyAuth {
             ..Self::default()
         }
     }
+
+    pub fn is_proxy_admin(&self) -> bool {
+        self.user_role.as_deref() == Some("proxy_admin")
+    }
 }
 
 /// SHA-256 over `(route, model, key)`, used as the cache key. Hashing keeps
@@ -79,15 +83,25 @@ fn bearer_token(parts: &Parts) -> Option<&str> {
         .filter(|token| !token.is_empty())
 }
 
-/// Pull a query parameter's value out of a raw query string, e.g. `model` from
-/// `model=gpt-realtime&foo=bar`. Returns the raw (un-percent-decoded) value;
-/// model names use only query-safe characters (alphanumerics, `-`, `_`, `.`, `/`),
-/// so no decoding is needed for the values we authorize on.
-fn query_param<'a>(query: Option<&'a str>, name: &str) -> Option<&'a str> {
-    query?.split('&').find_map(|pair| {
-        let (k, v) = pair.split_once('=')?;
-        (k == name).then_some(v)
-    })
+#[derive(serde::Deserialize)]
+struct RequestedModelQuery {
+    model: Option<String>,
+}
+
+/// Decode the requested model with the same query parser Axum's `Query` extractor
+/// uses later in the realtime handler. This keeps auth and serving aligned: an
+/// encoded key like `mo%64el=restricted` is still treated as `model=restricted`
+/// during authorization.
+fn requested_model(query: Option<&str>) -> Result<Option<String>, String> {
+    let Some(query) = query else {
+        return Ok(None);
+    };
+    let parsed: RequestedModelQuery =
+        serde_urlencoded::from_str(query).map_err(|err| format!("invalid query string: {err}"))?;
+    Ok(parsed
+        .model
+        .map(|model| model.trim().to_string())
+        .filter(|model| !model.is_empty()))
 }
 
 #[axum::async_trait]
@@ -115,7 +129,9 @@ impl FromRequestParts<AppState> for UserApiKeyAuth {
         // the key's route/model restrictions for what's actually being requested
         // (not just that the key exists — closes the model-authorization bypass).
         let route = parts.uri.path();
-        let model = query_param(parts.uri.query(), "model");
+        let model = requested_model(parts.uri.query())
+            .map_err(|message| (StatusCode::BAD_REQUEST, message))?;
+        let model = model.as_deref();
 
         // Virtual key: serve from cache, else verify via the (swappable) backend
         // and cache the result keyed by (route, model, key) hash.
@@ -303,9 +319,19 @@ mod tests {
 
     #[test]
     fn admin_is_proxy_admin() {
-        assert_eq!(
-            UserApiKeyAuth::admin().user_role.as_deref(),
-            Some("proxy_admin")
-        );
+        assert!(UserApiKeyAuth::admin().is_proxy_admin());
+        assert!(!UserApiKeyAuth::default().is_proxy_admin());
+    }
+
+    #[test]
+    fn requested_model_decodes_query_parameter_names_like_axum() {
+        let model = requested_model(Some("mo%64el=gpt-realtime&foo=bar")).unwrap();
+        assert_eq!(model.as_deref(), Some("gpt-realtime"));
+    }
+
+    #[test]
+    fn requested_model_trims_empty_values() {
+        let model = requested_model(Some("model=%20%20%20")).unwrap();
+        assert_eq!(model, None);
     }
 }

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

@@ -1,8 +1,7 @@
 //! `GET /v1/realtime` (WebSocket).
 //!
 //! This file is the **axum surface**: `router()`, the handler, and the small
-//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. Auth is
-//! the `RequireMasterKey` extractor, so the handler stays thin.
+//! socket↔events adapter. The pure logic (no axum) lives in [`service`].
 
 mod service;
 
@@ -33,16 +32,31 @@ struct RealtimeQuery {
     model: String,
 }
 
+fn require_realtime_billing_safe_auth(auth: &UserApiKeyAuth) -> Result<(), (StatusCode, String)> {
+    if auth.is_proxy_admin() {
+        return Ok(());
+    }
+    Err((
+        StatusCode::FORBIDDEN,
+        "non-admin virtual-key realtime is disabled until realtime usage is reported to the control plane".to_string(),
+    ))
+}
+
 /// Auth runs via the `UserApiKeyAuth` extractor (master key → admin, otherwise
-/// cache then the swappable authenticator). We validate the model BEFORE the
-/// upgrade so failures are clean HTTP (400/404), not a socket that opens then
-/// closes, then hand the socket to `bridge`.
+/// cache then the swappable authenticator). Non-admin virtual keys are rejected
+/// until realtime usage is reported back to the proxy; otherwise callers could
+/// consume upstream realtime spend without that spend being charged to their key
+/// or team. We validate the model BEFORE the upgrade so failures are clean HTTP
+/// (400/404), not a socket that opens then closes, then hand the socket to
+/// `bridge`.
 async fn handle(
-    _auth: UserApiKeyAuth,
+    auth: UserApiKeyAuth,
     ws: WebSocketUpgrade,
     State(state): State<AppState>,
     Query(query): Query<RealtimeQuery>,
 ) -> Result<Response, (StatusCode, String)> {
+    require_realtime_billing_safe_auth(&auth)?;
+
     if query.model.trim().is_empty() {
         return Err((
             StatusCode::BAD_REQUEST,
@@ -87,3 +101,25 @@ async fn bridge(
     futures_util::pin_mut!(client_in, client_out);
     let _ = service::run(&router, &pool, &model, None, client_in, client_out).await;
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn realtime_allows_proxy_admin_identity() {
+        assert!(require_realtime_billing_safe_auth(&UserApiKeyAuth::admin()).is_ok());
+    }
+
+    #[test]
+    fn realtime_rejects_non_admin_virtual_keys_until_spend_is_reported() {
+        let err = require_realtime_billing_safe_auth(&UserApiKeyAuth {
+            user_role: Some("internal_user".to_string()),
+            ..UserApiKeyAuth::default()
+        })
+        .unwrap_err();
+
+        assert_eq!(err.0, StatusCode::FORBIDDEN);
+        assert!(err.1.contains("usage is reported"));
+    }
+}