From 57e831d44510c99c063580594056e5bf67696cce Mon Sep 17 00:00:00 2001 From: jibanez-staticduo Date: Wed, 9 Sep 2026 19:44:19 +0200 Subject: [PATCH] fix(chatgpt): authenticate sideband models and forward image headers --- litellm/images/main.py | 1 + litellm/proxy/auth/user_api_key_auth.py | 22 ++++++---- .../test_litellm/llms/chatgpt/test_images.py | 2 + .../proxy/auth/test_user_api_key_auth.py | 41 +++++++++++++++++++ 4 files changed, 57 insertions(+), 9 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 7290e537369..11df9728ede 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -402,6 +402,7 @@ def image_generation( model=model, prompt=prompt, image_generation_provider_config=image_generation_config, + extra_headers=extra_headers, image_generation_optional_request_params=optional_params, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params_dict, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6b000489d5a..1ab1811fef8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -508,15 +508,6 @@ async def user_api_key_auth_websocket(websocket: WebSocket): request._url = websocket.url - query_params: Final = websocket.query_params - - model: Final = query_params.get("model") - - async def return_body(): - return _realtime_request_body(model) - - request.body = return_body - authorization: Final = websocket.headers.get("authorization") # If no Authorization header, try the api-key header if not authorization: @@ -542,6 +533,19 @@ async def user_api_key_auth_websocket(websocket: WebSocket): # Call user_api_key_auth with the extracted API key # Note: You'll need to modify this to work with WebSocket context if needed try: + from litellm.proxy.realtime_endpoints.call_sessions import decode_call + + call_token: Final = websocket.path_params.get("call_id") or websocket.query_params.get("call_id") + model: Final = ( + decode_call(call_token, f"Bearer {api_key}").alias + if call_token is not None + else websocket.query_params.get("model") + ) + + async def return_body(): + return _realtime_request_body(model) + + request.body = return_body return await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") except Exception as e: if is_invalid_virtual_key_error(e): diff --git a/tests/test_litellm/llms/chatgpt/test_images.py b/tests/test_litellm/llms/chatgpt/test_images.py index 6b69a83839d..0f4eafccfa4 100644 --- a/tests/test_litellm/llms/chatgpt/test_images.py +++ b/tests/test_litellm/llms/chatgpt/test_images.py @@ -25,7 +25,9 @@ def test_generation_routes_with_chatgpt_oauth(chatgpt_tokens): quality="auto", size="auto", background="auto", + extra_headers={"x-gateway-route": "images"}, ) + assert requests[0].headers["x-gateway-route"] == "images" assert result.data[0].b64_json == "aGVsbG8=" assert str(requests[0].url) == "https://chatgpt.com/backend-api/codex/images/generations" assert requests[0].headers["authorization"] == "Bearer test-token-" + "default" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index fd289e33ea6..d4dae5ee150 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -7133,3 +7133,44 @@ def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(t assert report["outcomes"] == ["accepted", "rejected"] auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("attachment", ["path", "query"]) +@pytest.mark.parametrize("credential", ["authorization", "api-key", "subprotocol"]) +@pytest.mark.parametrize("query_model", [b"", b"model=unbudgeted"]) +async def test_sideband_auth_uses_encrypted_model_for_budget_checks(monkeypatch, attachment, credential, query_model): + import hashlib + import importlib + import time + from unittest.mock import AsyncMock + from fastapi import WebSocket + from litellm.llms.chatgpt.codex import CodexRealtimeCall + from litellm.proxy.realtime_endpoints.call_sessions import encode_call + + auth_module = importlib.import_module("litellm.proxy.auth.user_api_key_auth") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-sideband-budget-salt") + token = encode_call(CodexRealtimeCall( + call_id="rtc_test", model="gpt-live-1-codex", alias="budgeted-voice", + owner=hashlib.sha256(b"Bearer owner").hexdigest(), expires_at=time.time() + 300, + )) + seen = [] + + async def authenticate(request, api_key): + seen.append((await request.json(), api_key)) + return "authenticated-with-model" + + monkeypatch.setattr(auth_module, "user_api_key_auth", authenticate) + websocket = WebSocket({ + "type": "websocket", "scheme": "ws", "server": ("localhost", 4000), + "path": "/v1/live/" + token if attachment == "path" else "/v1/realtime", + "path_params": {"call_id": token} if attachment == "path" else {}, + "query_string": query_model + (b"&call_id=" + token.encode() if attachment == "query" else b""), + "headers": { + "authorization": [(b"authorization", b"Bearer owner")], + "api-key": [(b"api-key", b"owner")], + "subprotocol": [(b"sec-websocket-protocol", b"realtime, openai-insecure-api-key.owner")], + }[credential], + }, AsyncMock(), AsyncMock()) + assert await auth_module.user_api_key_auth_websocket(websocket) == "authenticated-with-model" + assert seen == [({"model": "budgeted-voice"}, "Bearer owner")]