fix(health_check): skip max_tokens for image_generation mode (#26417)

* fix(health_check): skip max_tokens for image_generation mode

`_update_litellm_params_for_health_check` injected `max_tokens` for
every deployment. OpenAI `/v1/images/generations` strictly rejects
unknown fields, so health checks for dall-e-* and gpt-image-1 always
failed with `400 "Unknown parameter: 'max_tokens'"` even though the
actual image endpoint calls succeed. Skip the `max_tokens` injection
when `model_info.mode == "image_generation"`. `messages` still gets
injected (downstream `_filter_model_params` already strips it for
non-chat handlers).

* Switch to allow-list with per-deployment override

Per @krrishdholakia review: deny-listing image_generation only re-introduces
the same bug for every other non-chat mode (embedding, audio_*, rerank,
video_generation, ocr, search, moderation, ...).

Replace the single image_generation skip with `_MAX_TOKEN_SUPPORT_MODES =
{chat, completion, responses}`. Missing `mode` is treated as chat for
backward compatibility. New modes are safe by default.

Add `model_info.health_check_supports_max_tokens` as an operator escape
hatch — True forces injection on a non-listed deployment (operator wants
to bound probe tokens), False suppresses it on a chat-style deployment
behind a strict-schema provider.

Tests: parametrize over 3 chat-style + 10 non-chat modes, plus override
on/off and the no-mode legacy path.

litellm/proxy/health_check.py

@@ -32,6 +32,30 @@ ILLEGAL_DISPLAY_PARAMS = [
 
 MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"]
 
+# Modes whose health-check probe is a chat-style completion call and
+# therefore accept `max_tokens`. Other modes (embedding, image_generation,
+# audio_*, rerank, video_generation, ocr, search, moderation, ...) hit
+# endpoints that reject unknown fields with 400 "Unknown parameter:
+# 'max_tokens'". Allow-list so new modes are safe by default.
+# Per-deployment override: `model_info.health_check_supports_max_tokens`.
+_MAX_TOKEN_SUPPORT_MODES: frozenset = frozenset({"chat", "completion", "responses"})
+
+
+def _should_inject_health_check_max_tokens(model_info: dict) -> bool:
+    """
+    Whether the health-check probe should include `max_tokens`.
+
+    Order:
+      1. `model_info.health_check_supports_max_tokens` (operator override).
+      2. `_MAX_TOKEN_SUPPORT_MODES`. Missing `mode` is treated as `chat`
+         for backward compatibility.
+    """
+    explicit = model_info.get("health_check_supports_max_tokens")
+    if explicit is not None:
+        return bool(explicit)
+    mode = model_info.get("mode") or "chat"
+    return mode in _MAX_TOKEN_SUPPORT_MODES
+
 
 def _get_process_rss_mb() -> Optional[float]:
     """
@@ -362,14 +386,22 @@ def _update_litellm_params_for_health_check(
     Update the litellm params for health check.
 
     - gets a short `messages` param for health check
+    - adds a bounded `max_tokens` when the deployment is a chat-style mode
+      (`chat`, `completion`, `responses`) or the operator explicitly opts in
+      via `model_info.health_check_supports_max_tokens`. Non-chat endpoints
+      (image, embedding, audio_*, rerank, video, ocr, search, moderation, ...)
+      reject unknown fields with 400 "Unknown parameter: 'max_tokens'".
     - updates the `model` param with the `health_check_model` if it exists Doc: https://docs.litellm.ai/docs/proxy/health#wildcard-routes
     - updates the `voice` param with the `health_check_voice` for `audio_speech` mode if it exists Doc: https://docs.litellm.ai/docs/proxy/health#text-to-speech-models
     - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID
     """
     litellm_params["messages"] = _get_random_llm_message()
-    _resolved_max_tokens = _resolve_health_check_max_tokens(model_info, litellm_params)
-    if _resolved_max_tokens is not None:
-        litellm_params["max_tokens"] = _resolved_max_tokens
+    if _should_inject_health_check_max_tokens(model_info):
+        _resolved_max_tokens = _resolve_health_check_max_tokens(
+            model_info, litellm_params
+        )
+        if _resolved_max_tokens is not None:
+            litellm_params["max_tokens"] = _resolved_max_tokens
 
     _health_check_model = model_info.get("health_check_model", None)
     if _health_check_model is not None:

tests/test_litellm/proxy/test_health_check_max_tokens.py

@@ -225,3 +225,129 @@ def test_wildcard_ignores_reasoning_split_model_info(monkeypatch):
     litellm_params = {"model": "openai/*"}
 
     assert _resolve_health_check_max_tokens(model_info, litellm_params) is None
+
+
+# ---------------------------------------------------------------------------
+# image_generation must not receive max_tokens.
+#
+# _update_litellm_params_for_health_check injected `max_tokens` for every
+# deployment. For `mode: image_generation` that leaked into OpenAI
+# `/v1/images/generations`, which strictly rejects unknown fields with
+# `400 "Unknown parameter: 'max_tokens'"`, marking dall-e-* and
+# gpt-image-1 as permanently unhealthy even though their actual image
+# calls succeed. `messages` still gets injected (downstream
+# `_filter_model_params` already strips it for non-chat handlers).
+# ---------------------------------------------------------------------------
+
+
+def test_image_generation_mode_skips_max_tokens():
+    """image_generation must not receive max_tokens."""
+    model_info = {"mode": "image_generation"}
+    litellm_params = {"model": "openai/dall-e-3", "api_key": "sk-test"}
+
+    updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert "max_tokens" not in updated
+    # connection-level params must still pass through unchanged
+    assert updated["api_key"] == "sk-test"
+
+
+def test_health_check_max_tokens_value_is_ignored_for_non_chat_modes():
+    """A configured `health_check_max_tokens` *value* (the int that controls
+    how many tokens to inject) is still skipped when the mode is outside the
+    allow-list — the inject decision runs before value resolution, so the
+    value never reaches `_resolve_health_check_max_tokens`. Note this is
+    distinct from `health_check_supports_max_tokens` (the bool that toggles
+    injection on/off per deployment)."""
+    model_info = {"mode": "image_generation", "health_check_max_tokens": 50}
+    litellm_params = {"model": "openai/dall-e-3"}
+
+    updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert "max_tokens" not in updated
+
+
+def test_chat_mode_still_injects_max_tokens():
+    """Regression guard: the chat-style probe payload is unchanged."""
+    model_info = {"mode": "chat"}
+    litellm_params = {"model": "gpt-4"}
+
+    updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert updated["max_tokens"] == 5
+
+
+def test_no_mode_still_injects_max_tokens():
+    """Regression guard: model_info without `mode` keeps the legacy path."""
+    model_info: dict = {}
+    litellm_params = {"model": "gpt-4"}
+
+    updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert updated["max_tokens"] == 5
+
+
+# ---------------------------------------------------------------------------
+# Allow-list behavior: only chat-style modes (chat / completion / responses)
+# receive max_tokens. Every other mode is skipped by default.
+#
+# Per-deployment override via `health_check_supports_max_tokens` lets the
+# operator force injection on (e.g. a non-listed but max_tokens-capable
+# endpoint where they want to bound probe token usage) or off (e.g. a
+# chat-style provider with a strict schema that rejects unknown fields).
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("mode", ["chat", "completion", "responses"])
+def test_chat_style_modes_inject_max_tokens(mode):
+    updated = _update_litellm_params_for_health_check(
+        {"mode": mode}, {"model": f"openai/dummy-{mode}"}
+    )
+
+    assert updated["max_tokens"] == 5
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        "embedding",
+        "image_generation",
+        "image_edit",
+        "audio_speech",
+        "audio_transcription",
+        "rerank",
+        "video_generation",
+        "ocr",
+        "search",
+        "moderation",
+    ],
+)
+def test_non_chat_modes_skip_max_tokens(mode):
+    updated = _update_litellm_params_for_health_check(
+        {"mode": mode}, {"model": f"openai/dummy-{mode}"}
+    )
+
+    assert "max_tokens" not in updated
+
+
+def test_explicit_override_true_forces_injection_outside_allowlist():
+    """Operator opts a non-listed deployment in to bound probe token usage."""
+    model_info = {
+        "mode": "image_generation",
+        "health_check_supports_max_tokens": True,
+    }
+    litellm_params = {"model": "openai/some-future-image-model"}
+
+    updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert updated["max_tokens"] == 5
+
+
+def test_explicit_override_false_suppresses_injection_inside_allowlist():
+    """Operator opts a chat-style deployment out (strict-schema provider)."""
+    model_info = {"mode": "chat", "health_check_supports_max_tokens": False}
+    litellm_params = {"model": "openai/strict-schema-chat"}
+
+    updated = _update_litellm_params_for_health_check(model_info, litellm_params)
+
+    assert "max_tokens" not in updated