fix(health-check): keep max_tokens for completion mode, add unit tests

Greptile P1: litellm.atext_completion (completion mode) accepts max_tokens,
so unconditionally stripping it via _filter_model_params silently dropped
the BACKGROUND_HEALTH_CHECK_MAX_TOKENS cost cap for completion-mode
deployments. Make _filter_model_params mode-aware via a keep_max_tokens
flag and use it in the completion-mode lambda. Strict non-chat handlers
(image_generation, embedding, etc.) still strip max_tokens as before.

Adds tests/litellm/litellm_core_utils/test_health_check_utils.py covering
both code paths and addressing the codecov/patch failure.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

litellm/litellm_core_utils/health_check_helpers.py

@@ -157,7 +157,9 @@ class HealthCheckHelpers:
                 **model_params,
             ),
             "completion": lambda: litellm.atext_completion(
-                **_filter_model_params(model_params=model_params),
+                **_filter_model_params(
+                    model_params=model_params, keep_max_tokens=True
+                ),
                 prompt=prompt or "test",
             ),
             "embedding": lambda: litellm.aembedding(

litellm/litellm_core_utils/health_check_utils.py

@@ -2,29 +2,39 @@
 Utils used for litellm.ahealth_check()
 """
 
-#: litellm params that ``_update_litellm_params_for_health_check`` injects for
-#: chat-completion health checks but that are invalid (or rejected) on every
-#: other handler in :class:`HealthCheckHelpers.get_mode_handlers` — image and
-#: video generation, embeddings, audio speech / transcription, rerank, ocr,
-#: responses, batch, etc. ``messages`` is always chat-only; ``max_tokens`` is
-#: chat/completion-only and is rejected by strict providers (e.g. OpenAI's
-#: image-generation endpoints return 400 ``Unknown parameter: 'max_tokens'``).
-_NON_CHAT_HEALTH_CHECK_STRIP_KEYS = {"messages", "max_tokens"}
+#: ``messages`` is injected by ``_update_litellm_params_for_health_check`` for
+#: chat-style health checks; every non-chat handler in
+#: :class:`HealthCheckHelpers.get_mode_handlers` rejects it, and
+#: ``litellm.atext_completion`` takes ``prompt`` instead.
+_COMPLETION_HEALTH_CHECK_STRIP_KEYS = {"messages"}
+
+#: Strict non-chat handlers (OpenAI image generation, embeddings, rerank, etc.)
+#: also reject ``max_tokens`` — e.g. ``dall-e-3`` returns
+#: ``400 Unknown parameter: 'max_tokens'``. ``atext_completion`` does accept
+#: ``max_tokens``, so the ``completion`` mode keeps it (preserves the
+#: ``BACKGROUND_HEALTH_CHECK_MAX_TOKENS`` cost cap).
+_NON_CHAT_HEALTH_CHECK_STRIP_KEYS = _COMPLETION_HEALTH_CHECK_STRIP_KEYS | {
+    "max_tokens"
+}
 
 
-def _filter_model_params(model_params: dict) -> dict:
+def _filter_model_params(
+    model_params: dict, *, keep_max_tokens: bool = False
+) -> dict:
     """Strip chat-only params before invoking a non-chat health check handler.
 
     ``litellm.acompletion`` is the only mode handler that consumes
     ``model_params`` unfiltered; every other handler routes through this
-    helper, so removing chat-completion-only keys here keeps strict providers
-    (OpenAI image generation, etc.) from rejecting the request.
+    helper. ``litellm.atext_completion`` (the ``completion`` mode) accepts
+    ``max_tokens`` and should pass ``keep_max_tokens=True`` to preserve the
+    cost-control cap.
     """
-    return {
-        k: v
-        for k, v in model_params.items()
-        if k not in _NON_CHAT_HEALTH_CHECK_STRIP_KEYS
-    }
+    strip_keys = (
+        _COMPLETION_HEALTH_CHECK_STRIP_KEYS
+        if keep_max_tokens
+        else _NON_CHAT_HEALTH_CHECK_STRIP_KEYS
+    )
+    return {k: v for k, v in model_params.items() if k not in strip_keys}
 
 
 def _create_health_check_response(response_headers: dict) -> dict:

tests/litellm/litellm_core_utils/test_health_check_utils.py

@@ -0,0 +1,71 @@
+"""
+Tests for ``litellm.litellm_core_utils.health_check_utils._filter_model_params``.
+
+Regression coverage for #26406 / #26604: chat-only params (``messages``,
+``max_tokens``) must be stripped before non-chat health check handlers, but
+``max_tokens`` must be preserved for the ``completion`` mode handler
+(``litellm.atext_completion``) so the ``BACKGROUND_HEALTH_CHECK_MAX_TOKENS``
+cost cap still applies.
+"""
+
+from litellm.litellm_core_utils.health_check_utils import (
+    _COMPLETION_HEALTH_CHECK_STRIP_KEYS,
+    _NON_CHAT_HEALTH_CHECK_STRIP_KEYS,
+    _filter_model_params,
+)
+
+
+def _sample_params() -> dict:
+    return {
+        "model": "openai/dall-e-3",
+        "messages": [{"role": "user", "content": "test"}],
+        "max_tokens": 5,
+        "api_key": "sk-fake",
+    }
+
+
+def test_filter_strips_messages_and_max_tokens_for_non_chat_handlers():
+    """Default behavior: strip both ``messages`` and ``max_tokens``.
+
+    Reproduces the original failure (OpenAI image-generation 400 on
+    ``max_tokens``) — these keys must not reach non-chat handlers.
+    """
+    filtered = _filter_model_params(model_params=_sample_params())
+
+    assert "messages" not in filtered
+    assert "max_tokens" not in filtered
+    assert filtered == {"model": "openai/dall-e-3", "api_key": "sk-fake"}
+
+
+def test_filter_keeps_max_tokens_for_completion_mode():
+    """``completion`` mode passes ``keep_max_tokens=True``.
+
+    ``litellm.atext_completion`` accepts ``max_tokens``; stripping it would
+    silently remove the ``BACKGROUND_HEALTH_CHECK_MAX_TOKENS`` cost cap.
+    """
+    filtered = _filter_model_params(
+        model_params=_sample_params(), keep_max_tokens=True
+    )
+
+    assert "messages" not in filtered
+    assert filtered.get("max_tokens") == 5
+    assert filtered["model"] == "openai/dall-e-3"
+    assert filtered["api_key"] == "sk-fake"
+
+
+def test_filter_does_not_mutate_input():
+    params = _sample_params()
+    snapshot = dict(params)
+
+    _filter_model_params(model_params=params)
+    _filter_model_params(model_params=params, keep_max_tokens=True)
+
+    assert params == snapshot
+
+
+def test_strip_key_sets_are_consistent():
+    """``_NON_CHAT`` is a strict superset of ``_COMPLETION`` — only
+    ``max_tokens`` differs between the two modes."""
+    assert _COMPLETION_HEALTH_CHECK_STRIP_KEYS == {"messages"}
+    assert _NON_CHAT_HEALTH_CHECK_STRIP_KEYS == {"messages", "max_tokens"}
+    assert _COMPLETION_HEALTH_CHECK_STRIP_KEYS < _NON_CHAT_HEALTH_CHECK_STRIP_KEYS