fix(images): honor max_retries and prevent Router retry amplification (#48705)

* test(images): cover SDK retry budgets and router ownership

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(images): forward max_retries to the OpenAI image generation client

Per-call, deployment and default_litellm_params max_retries were filtered
out before reaching the OpenAI-compatible image handler, which then always
used 2 SDK retries. Forward the caller's value, and pin it to 0 for routed
calls so the Router stays the only retry owner, matching completion()

Co-authored-by: Ericz <yonglzh@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(images): read image max_retries through a typed kwargs view

Co-authored-by: Ericz <yonglzh@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(images): assert the public invalid retry error mapping

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(images): preserve the invalid retry status code

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(images): group the proxy test imports

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(images): discard proxy retry test routers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Ericz <yonglzh@gmail.com>

litellm/images/main.py

@@ -284,6 +284,12 @@ def image_generation(
         extra_headers: Final = kwargs.get("extra_headers", None)
         headers: Final[dict] = kwargs.get("headers", None) or {}
         base_model: Final = kwargs.get("base_model", None)
+        call_kwargs: Final = cast(Mapping[str, object], kwargs)  # cast-ok: image_generation's **kwargs is untyped
+        max_retries: Final = call_kwargs.get("max_retries")
+        is_router_call: Final = any(
+            isinstance(call_metadata := call_kwargs.get(k), Mapping) and "model_group" in call_metadata
+            for k in ("metadata", "litellm_metadata")
+        )
         if extra_headers is not None:
             headers.update(extra_headers)
         model_response: ImageResponse = litellm.utils.ImageResponse()
@@ -550,6 +556,7 @@ def image_generation(
                 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)
+            sdk_max_retries: Final = 0 if is_router_call else max_retries
             model_response = openai_chat_completions.image_generation(
                 model=model,
                 prompt=prompt,
@@ -557,7 +564,9 @@ def image_generation(
                 api_key=api_key or dynamic_api_key,
                 api_base=api_base,
                 logging_obj=litellm_logging_obj,
-                optional_params=optional_params,
+                optional_params=(
+                    optional_params if sdk_max_retries is None else {**optional_params, "max_retries": sdk_max_retries}
+                ),
                 model_response=model_response,
                 organization=organization,
                 aimg_generation=aimg_generation,

tests/integration/providers/test_image_gen_drop_params_wire.py

@@ -1,16 +1,101 @@
 import json
-from typing import Final
+import threading
+from types import MappingProxyType
+from typing import Final, Literal
 
+import httpx
 import openai
 import pytest
 from integration._support.client import Gateway
 from integration._support.wire import Reply, Request, wire_server
 from pydantic import JsonValue, TypeAdapter
 
+import litellm
+
 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
 _PROMPT: Final = "a transparent lighthouse at dawn"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("router_retries", (0, 1))
+@pytest.mark.parametrize("sdk_retries", (0, 5))
+@pytest.mark.parametrize("outcome", ("503", "timeout", "success"))
+async def test_proxy_image_retry_attempts_are_owned_only_by_router(
+    monkeypatch: pytest.MonkeyPatch,
+    router_retries: int,
+    sdk_retries: int,
+    outcome: Literal["503", "timeout", "success"],
+) -> None:
+    from litellm.proxy import proxy_server
+
+    release: Final = threading.Event()
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.target == "/v1/images/generations"
+        assert request.headers["authorization"] == "Bearer synthetic-image-key"
+        assert _JSON_OBJECT.validate_json(request.body) == {"model": "gpt-image-1-mini", "prompt": "a red circle"}
+        if outcome == "timeout":
+            return Reply(chunks=(b"{}",), gate_before_first=release)
+        if outcome == "503":
+            return Reply(
+                status=503,
+                headers=MappingProxyType({"retry-after-ms": "1"}),
+                body=b'{"error":{"message":"scripted unavailable","type":"server_error"}}',
+            )
+        return Reply(body=b'{"created":1700000000,"data":[{"b64_json":"aW1n"}]}')
+
+    monkeypatch.setattr(litellm, "callbacks", [])
+    monkeypatch.setattr(litellm, "success_callback", [])
+    monkeypatch.setattr(litellm, "failure_callback", [])
+    monkeypatch.setattr(litellm, "num_retries", None)
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    litellm.in_memory_llm_clients_cache.flush_cache()
+    with wire_server(respond) as wire:
+        router: Final = litellm.Router(
+            model_list=[
+                {
+                    "model_name": "image-retries",
+                    "litellm_params": {
+                        "model": "openai/gpt-image-1-mini",
+                        "api_base": f"{wire.url}/v1",
+                        "api_key": "synthetic-image-key",
+                        "max_retries": sdk_retries,
+                        "timeout": 0.1,
+                    },
+                }
+            ],
+            num_retries=router_retries,
+            disable_cooldowns=True,
+        )
+        monkeypatch.setattr(proxy_server, "llm_router", router)
+        monkeypatch.setattr(proxy_server, "master_key", None)
+        monkeypatch.setattr(proxy_server, "prisma_client", None)
+        monkeypatch.setattr(proxy_server, "general_settings", {})
+        try:
+            async with httpx.AsyncClient(
+                transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://proxy.test"
+            ) as client:
+                response: Final = await client.post(
+                    "/v1/images/generations",
+                    json={"model": "image-retries", "prompt": "a red circle", "max_retries": sdk_retries},
+                )
+        finally:
+            release.set()
+            router.discard()
+            litellm.in_memory_llm_clients_cache.flush_cache()
+        requests: Final = wire.drain()
+    status: Final = {"503": 503, "timeout": 408, "success": 200}[outcome]
+    attempts: Final = 1 if outcome == "success" else router_retries + 1
+    assert response.status_code == status, response.text
+    assert len(requests) == attempts, tuple(request.headers.get("x-stainless-retry-count") for request in requests)
+    if outcome == "success":
+        data: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(
+            _JSON_OBJECT.validate_json(response.content)["data"]
+        )
+        assert data[0]["b64_json"] == "aW1n"
+
+
 def test_image_generation_additional_drop_params_reaches_provider_body(gateway: Gateway) -> None:
     def respond(request: Request) -> Reply:
         assert request.method == "POST"

tests/unit/images/test_main.py

@@ -3,14 +3,18 @@ import base64
 import json
 from datetime import datetime
 from pathlib import Path
-from typing import Final
+from collections.abc import Generator
+from queue import SimpleQueue
+from typing import Final, Literal, TypedDict
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
 import respx
-from openai import AsyncOpenAI
+from openai import AsyncOpenAI, OpenAI
 from openai.types.image import Image
+from pydantic import JsonValue, TypeAdapter
+from typing_extensions import ReadOnly
 
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
@@ -19,6 +23,212 @@ from litellm.types.utils import CallTypes, StandardLoggingPayload
 import os
 
 
+class ImageRetryParameters(TypedDict, total=False):
+    max_retries: ReadOnly[int | None]
+
+
+@pytest.fixture
+def image_retry_delays() -> Generator[None, None, None]:
+    with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock):
+        yield
+
+
+def _image_retry_transport(
+    requests: SimpleQueue[httpx.Request], outcome: Literal["503", "timeout", "success"]
+) -> httpx.MockTransport:
+    def respond(request: httpx.Request) -> httpx.Response:
+        requests.put(request)
+        assert request.method == "POST"
+        assert request.url.path == "/v1/images/generations"
+        body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content)
+        assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"}
+        if outcome == "timeout":
+            raise httpx.ReadTimeout("scripted image timeout", request=request)
+        if outcome == "503":
+            return httpx.Response(503, json={"error": {"message": "scripted unavailable", "type": "server_error"}})
+        return httpx.Response(200, json={"created": 1700000000, "data": [{"b64_json": "aW1n"}]})
+
+    return httpx.MockTransport(respond)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("image_retry_delays")
+@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
+@pytest.mark.parametrize("outcome", ("503", "timeout"))
+@pytest.mark.parametrize(
+    ("retries", "attempts"),
+    ((0, 1), (1, 2), (None, 3), ("unset", 3), ("global", 3)),
+    ids=("zero", "one", "none", "unset", "global"),
+)
+async def test_image_generation_honors_sdk_retry_budget(
+    monkeypatch: pytest.MonkeyPatch,
+    is_async: bool,
+    outcome: Literal["503", "timeout"],
+    retries: int | None | Literal["unset", "global"],
+    attempts: int,
+) -> None:
+    requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
+    transport: Final = _image_retry_transport(requests, outcome)
+    if retries == "global":
+        monkeypatch.setattr(litellm, "max_retries", 0, raising=False)
+    retry_params: Final[ImageRetryParameters] = {} if isinstance(retries, str) else {"max_retries": retries}
+    error: Final = litellm.Timeout if outcome == "timeout" else litellm.ServiceUnavailableError
+    if is_async:
+        async with AsyncOpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            max_retries=5,
+            http_client=httpx.AsyncClient(transport=transport),
+        ) as client:
+            with pytest.raises(error):
+                await litellm.aimage_generation(
+                    model="openai/gpt-image-1-mini",
+                    prompt="a red circle",
+                    api_key="synthetic-image-key",
+                    api_base="https://image.test/v1",
+                    client=client,
+                    **retry_params,
+                )
+    else:
+        with OpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            max_retries=5,
+            http_client=httpx.Client(transport=transport),
+        ) as client:
+            with pytest.raises(error):
+                litellm.image_generation(
+                    model="openai/gpt-image-1-mini",
+                    prompt="a red circle",
+                    api_key="synthetic-image-key",
+                    api_base="https://image.test/v1",
+                    client=client,
+                    **retry_params,
+                )
+    assert requests.qsize() == attempts
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("image_retry_delays")
+@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
+@pytest.mark.parametrize("provider", ("litellm_proxy", "hosted_vllm"))
+async def test_compatible_image_generation_honors_zero_retries(provider: str, is_async: bool) -> None:
+    requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
+    transport: Final = _image_retry_transport(requests, "503")
+    if is_async:
+        async with AsyncOpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            http_client=httpx.AsyncClient(transport=transport),
+        ) as client:
+            with pytest.raises(litellm.ServiceUnavailableError):
+                await litellm.aimage_generation(
+                    model=f"{provider}/gpt-image-1-mini",
+                    prompt="a red circle",
+                    api_key="synthetic-image-key",
+                    api_base="https://image.test/v1",
+                    client=client,
+                    max_retries=0,
+                )
+    else:
+        with OpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            http_client=httpx.Client(transport=transport),
+        ) as client:
+            with pytest.raises(litellm.ServiceUnavailableError):
+                litellm.image_generation(
+                    model=f"{provider}/gpt-image-1-mini",
+                    prompt="a red circle",
+                    api_key="synthetic-image-key",
+                    api_base="https://image.test/v1",
+                    client=client,
+                    max_retries=0,
+                )
+    assert requests.qsize() == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("image_retry_delays")
+@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
+async def test_image_generation_rejects_non_integer_retries_before_sending(is_async: bool) -> None:
+    requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
+    transport: Final = _image_retry_transport(requests, "success")
+    if is_async:
+        async with AsyncOpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            http_client=httpx.AsyncClient(transport=transport),
+        ) as client:
+            with pytest.raises(litellm.BadRequestError) as raised:
+                await litellm.aimage_generation(
+                    model="openai/gpt-image-1-mini",
+                    prompt="a red circle",
+                    api_key="synthetic-image-key",
+                    api_base="https://image.test/v1",
+                    client=client,
+                    max_retries="invalid",
+                )
+    else:
+        with OpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            http_client=httpx.Client(transport=transport),
+        ) as client:
+            with pytest.raises(litellm.BadRequestError) as raised:
+                litellm.image_generation(
+                    model="openai/gpt-image-1-mini",
+                    prompt="a red circle",
+                    api_key="synthetic-image-key",
+                    api_base="https://image.test/v1",
+                    client=client,
+                    max_retries="invalid",
+                )
+    assert raised.value.status_code == 422
+    assert "max retries must be an int" in raised.value.message
+    assert requests.qsize() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
+async def test_image_generation_returns_success_without_forwarding_retry_control(is_async: bool) -> None:
+    requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
+    transport: Final = _image_retry_transport(requests, "success")
+    if is_async:
+        async with AsyncOpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            http_client=httpx.AsyncClient(transport=transport),
+        ) as client:
+            async_response: Final = await litellm.aimage_generation(
+                model="openai/gpt-image-1-mini",
+                prompt="a red circle",
+                api_key="synthetic-image-key",
+                api_base="https://image.test/v1",
+                client=client,
+                max_retries=0,
+            )
+            async_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(async_response.data)
+            assert async_data[0].b64_json == "aW1n"
+    else:
+        with OpenAI(
+            api_key="synthetic-image-key",
+            base_url="https://image.test/v1",
+            http_client=httpx.Client(transport=transport),
+        ) as client:
+            sync_response: Final = litellm.image_generation(
+                model="openai/gpt-image-1-mini",
+                prompt="a red circle",
+                api_key="synthetic-image-key",
+                api_base="https://image.test/v1",
+                client=client,
+                max_retries=0,
+            )
+            sync_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(sync_response.data)
+            assert sync_data[0].b64_json == "aW1n"
+    assert requests.qsize() == 1
+
+
 def test_image_generation_keeps_an_internal_prefixed_kwarg_out_of_the_provider_request(
     respx_mock: respx.MockRouter,
 ) -> None:

tests/unit/test_router_per_deployment_num_retries.py

@@ -6,7 +6,10 @@ GitHub Issue: #18968 - Per-deployment max_retries/num_retries in litellm_params
 import httpx
 import pytest
 import pytest_asyncio
-from unittest.mock import patch
+from queue import SimpleQueue
+from typing import Final, Literal
+from unittest.mock import AsyncMock, patch
+from pydantic import JsonValue, TypeAdapter
 
 import litellm
 from litellm import Router
@@ -14,6 +17,70 @@ from litellm.types.router import RetryPolicy
 from litellm.integrations.custom_logger import CustomLogger
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
+@pytest.mark.parametrize("source", ("deployment", "default", "request"))
+@pytest.mark.parametrize(("sdk_retries", "router_retries"), ((0, 0), (0, 1), (5, 1)))
+async def test_routed_images_do_not_multiply_sdk_and_router_retries(
+    monkeypatch: pytest.MonkeyPatch,
+    is_async: bool,
+    source: Literal["deployment", "default", "request"],
+    sdk_retries: int,
+    router_retries: int,
+) -> None:
+    requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        requests.put(request)
+        assert request.url.path == "/v1/images/generations"
+        body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content)
+        assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"}
+        return httpx.Response(
+            503,
+            headers={"retry-after": "0"},
+            json={"error": {"message": "scripted unavailable", "type": "server_error"}},
+        )
+
+    transport: Final = httpx.MockTransport(respond)
+    litellm.in_memory_llm_clients_cache.flush_cache()
+    with httpx.Client(transport=transport) as sync_session:
+        async with httpx.AsyncClient(transport=transport) as async_session:
+            monkeypatch.setattr(litellm, "client_session", sync_session)
+            monkeypatch.setattr(litellm, "aclient_session", async_session)
+            monkeypatch.setattr(litellm, "callbacks", [])
+            monkeypatch.setattr(litellm, "success_callback", [])
+            monkeypatch.setattr(litellm, "failure_callback", [])
+            monkeypatch.setattr(litellm, "num_retries", None)
+            router: Final = Router(
+                model_list=[
+                    {
+                        "model_name": "image-retries",
+                        "litellm_params": {
+                            "model": "openai/gpt-image-1-mini",
+                            "api_key": "synthetic-image-key",
+                            "api_base": "https://router-image.test/v1",
+                            **({"max_retries": sdk_retries} if source == "deployment" else {}),
+                        },
+                    }
+                ],
+                default_litellm_params={"max_retries": sdk_retries} if source == "default" else {},
+                num_retries=router_retries,
+                disable_cooldowns=True,
+            )
+            call_params: Final = {"max_retries": sdk_retries} if source == "request" else {}
+            try:
+                with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock):
+                    if is_async:
+                        with pytest.raises(litellm.ServiceUnavailableError):
+                            await router.aimage_generation(model="image-retries", prompt="a red circle", **call_params)
+                    else:
+                        with pytest.raises(litellm.ServiceUnavailableError):
+                            router.image_generation(model="image-retries", prompt="a red circle", **call_params)
+            finally:
+                litellm.in_memory_llm_clients_cache.flush_cache()
+    assert requests.qsize() == router_retries + 1
+
+
 class TestPerDeploymentNumRetries:
     """Test that per-deployment num_retries in litellm_params is correctly used."""