fix(chatgpt): forward client credentials for image generation (#48522)

* fix(chatgpt): forward client credentials for image generation

* fix(chatgpt): validate image headers at the typed boundary

* refactor(chatgpt): use read-only image header types

* fix(chatgpt): preserve image header credentials and workspace

litellm/images/main.py

@@ -1,7 +1,7 @@
 import asyncio
 import contextvars
 import importlib
-from collections.abc import Coroutine
+from collections.abc import Coroutine, Mapping
 from functools import partial
 from typing import TYPE_CHECKING, Final, Literal, cast, overload
 
@@ -85,6 +85,47 @@ def _get_ImageEditRequestUtils() -> type["ImageEditRequestUtils"]:
     return _ImageEditRequestUtils_cache
 
 
+def _chatgpt_image_params(kwargs: Mapping[str, object]) -> GenericLiteLLMParams:
+    from litellm.responses.utils import ResponsesAPIRequestUtils
+
+    headers: Final = ResponsesAPIRequestUtils.merge_client_forwarded_headers(
+        TypeAdapter[dict[str, object] | None](dict[str, object] | None).validate_python(kwargs.get("extra_headers")),
+        TypeAdapter[dict[str, str] | None](dict[str, str] | None).validate_python(kwargs.get("headers")),
+    )
+    return GenericLiteLLMParams.model_validate({**kwargs, "extra_headers": headers})
+
+
+def _chatgpt_image_headers(model: str, params: GenericLiteLLMParams, headers: Mapping[str, str]) -> Mapping[str, str]:
+    from litellm.llms.chatgpt.common_utils import get_chatgpt_request_credential
+    from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig
+
+    credential: Final = get_chatgpt_request_credential(params)
+    credential_headers: Final = (
+        {
+            "Authorization": f"Bearer {credential.access_token}",
+            **({"ChatGPT-Account-Id": credential.account_id} if credential.account_id else {}),
+        }
+        if credential is not None
+        else {}
+    )
+    explicit_account_id: Final = next(
+        (value for name, value in headers.items() if name.lower() == "chatgpt-account-id"), None
+    )
+    auth_headers: Final = ChatGPTResponsesAPIConfig().validate_environment(
+        headers=credential_headers, model=model, litellm_params=params
+    )
+    return {
+        **{k: v for k, v in auth_headers.items() if k.lower() not in ("content-type", "accept")},
+        **{
+            k: v
+            for k, v in headers.items()
+            if k.lower() not in ("authorization", "chatgpt-account-id", "cookie", "x-litellm-api-key")
+        },
+        **({"ChatGPT-Account-Id": explicit_account_id} if credential is None and explicit_account_id else {}),
+        "Accept": "application/json",
+    }
+
+
 ##### Image Generation #######################
 @client
 async def aimage_generation(*args, **kwargs) -> ImageResponse:
@@ -121,6 +162,12 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
             model=model,
             custom_llm_provider=provider_hints.get("custom_llm_provider"),
             api_base=provider_hints.get("api_base"),
+            litellm_params=(
+                _chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python(kwargs))
+                if provider_hints.get("custom_llm_provider") == "chatgpt"
+                or (isinstance(model, str) and model.startswith("chatgpt/"))
+                else None
+            ),
         )
 
         # Await normally
@@ -240,12 +287,18 @@ def image_generation(
         if extra_headers is not None:
             headers.update(extra_headers)
         model_response: ImageResponse = litellm.utils.ImageResponse()
+        chatgpt_params: Final = (
+            _chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python({**kwargs, "api_key": api_key}))
+            if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/"))
+            else None
+        )
         dynamic_api_key: str | None = None
         if model is not None or custom_llm_provider is not None:
             model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
                 model=model,
                 custom_llm_provider=custom_llm_provider,
                 api_base=api_base,
+                litellm_params=chatgpt_params,
             )
         else:
             model = "dall-e-2"
@@ -268,7 +321,10 @@ def image_generation(
             "size",
             "style",
         ]
-        non_default_params: Final = filter_out_litellm_params(kwargs, excluding=openai_params)
+        non_default_params: Final = filter_out_litellm_params(
+            kwargs,
+            excluding=[*openai_params, "extra_headers", "headers"] if chatgpt_params is not None else openai_params,
+        )
 
         image_generation_config: BaseImageGenerationConfig | None = None
         if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values():
@@ -490,7 +546,7 @@ def image_generation(
             or custom_llm_provider == LlmProviders.LITELLM_PROXY.value
             or custom_llm_provider in litellm.openai_compatible_providers
         ):
-            if extra_headers is not None:
+            if extra_headers is not None and chatgpt_params is None:
                 optional_params["extra_headers"] = extra_headers
             # Forward OpenAI organization if present (set by proxy pre-call utils)
             organization: Final[str | None] = kwargs.get("organization", None)
@@ -506,7 +562,17 @@ def image_generation(
                 organization=organization,
                 aimg_generation=aimg_generation,
                 client=client,
-                headers=headers,
+                headers=(
+                    dict(
+                        _chatgpt_image_headers(
+                            model,
+                            chatgpt_params.model_copy(update={"api_key": dynamic_api_key or api_key}),
+                            TypeAdapter(dict[str, str]).validate_python(headers),
+                        )
+                    )
+                    if chatgpt_params is not None
+                    else headers
+                ),
                 litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "bedrock":

tests/unit/images/test_main.py

@@ -1,6 +1,8 @@
 import asyncio
+import base64
 import json
 from datetime import datetime
+from pathlib import Path
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -405,3 +407,56 @@ async def test_azure_image_generation_request_body():
         call_args = mock_post.call_args
         request_json = call_args.kwargs.get("json", {})
         assert request_json == expected_body
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("transport", ["headers", "extra_headers", "scoped", "api_key"])
+@pytest.mark.parametrize("saved_key", [None, "fake-saved-gateway-key"])
+async def test_chatgpt_images_forward_client_login_and_workspace(
+    asynchronous: bool,
+    transport: str,
+    saved_key: str | None,
+    respx_mock: respx.MockRouter,
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
+    claims: Final = {"https://api.openai.com/auth": {"chatgpt_account_id": "client-account"}}
+    payload: Final = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=")
+    token: Final = f"test.{payload}.test"
+    headers: Final = {"aUtHoRiZaTiOn": f"Bearer {token}", "ChatGPT-Account-Id": "selected-workspace"}
+    credentials: Final = (
+        {"provider_specific_header": {"custom_llm_provider": "chatgpt", "extra_headers": headers}}
+        if transport == "scoped"
+        else {"api_key": token, "extra_headers": {"cHaTgPt-AcCoUnT-Id": "selected-workspace"}}
+        if transport == "api_key"
+        else {transport: {**headers, "cookie": "fake-cookie", "x-litellm-api-key": "fake-gateway-key"}}
+    )
+    route: Final = respx_mock.post("https://chatgpt.com/backend-api/codex/images/generations").mock(
+        return_value=httpx.Response(200, json={"created": 1700000000, "data": [{"b64_json": "aW1n"}]})
+    )
+    params: Final = {
+        "model": "chatgpt/gpt-image-2",
+        "prompt": "A blue circle",
+        "api_key": saved_key,
+        **(
+            {
+                "client": AsyncOpenAI(
+                    api_key="fake", base_url="https://chatgpt.com/backend-api/codex", http_client=httpx.AsyncClient()
+                )
+            }
+            if asynchronous
+            else {}
+        ),
+        **credentials,
+    }
+    response: Final = await litellm.aimage_generation(**params) if asynchronous else litellm.image_generation(**params)
+    assert response.data[0].b64_json == "aW1n"
+    request: Final = route.calls.last.request
+    assert request.headers["authorization"] == f"Bearer {token}"
+    assert request.headers["chatgpt-account-id"] == "selected-workspace"
+    assert request.headers.get_list("content-type") == ["application/json"]
+    assert "cookie" not in request.headers
+    assert "x-litellm-api-key" not in request.headers
+    assert json.loads(request.content) == {"model": "gpt-image-2", "prompt": "A blue circle"}