fix(h2o): gate the rename to chat completions and to max_tokens-discarding routes

Two review findings, both confirmed by execution on the 1.93.0 tree this
targets.

Blocker: the hook broke /v1/messages on Azure-routed models. The
deployment-hook dispatch is not chat-specific, it runs for every
@client-decorated async entrypoint, and litellm.anthropic_messages
declares max_tokens as a REQUIRED parameter. Popping it made litellm's
own wrapper raise TypeError on the following await
original_function(...), outside this hook's try/except, with a traceback
that never names the hook. Reproduced against a real install: identical
call, OK without the hook registered, TypeError with it. Gate on
call_type so only completion/acompletion are rewritten. This also covers
atext_completion, whose endpoint has no max_completion_tokens, and
aembedding.

Non-blocking: a configured max_completion_tokens is no longer a trigger
on its own. convert_model_to_litellm_config sets it for every reasoning
model, Azure or not, and only the Azure branch adds the drop, so the old
predicate would strip max_tokens from non-Azure providers and leave those
requests with no ceiling at all. Trigger on the drop list or an Azure
route instead, which keeps the drop-less reasoning-Azure case firing and
closes the previously documented caller-supplied-ceiling edge. Zero
non-Azure entries reach this today, so the change is latent-only.

Verified after the change: anthropic_messages OK, acompletion still
renames max_tokens=50 to max_completion_tokens=50. Both new tests fail
on the previous head.

Tests: 22.

Author: achraf-mer

litellm/integrations/h2o/litellm_max_tokens_rename_hook.py

@@ -46,33 +46,44 @@
 -----
 Fires only when BOTH hold for the deployment actually selected:
 
+  * the call is a chat completion, and
   * the request carries a positive integer `max_tokens`, and
-  * that deployment expects `max_completion_tokens` (it has one configured,
-    or `max_tokens` in its `additional_drop_params`).
+  * the selected deployment is one that discards or rejects `max_tokens`,
+    i.e. it lists `max_tokens` in `additional_drop_params`, or it is an Azure
+    route.
+
+A configured `max_completion_tokens` is deliberately NOT a trigger on its own.
+`convert_model_to_litellm_config` sets it for every reasoning model, Azure or
+not (`use_completion_tokens = is_reasoning_model or is_azure_provider`), and
+only the Azure branch adds the drop. Treating it as a trigger would strip
+`max_tokens` from non-Azure providers that need it, leaving those requests with
+no ceiling at all, which is worse than the bug being fixed. It also cannot be
+distinguished from a caller-supplied value once kwargs are merged.
 
 Deployments that natively accept `max_tokens` (Anthropic, Bedrock, vLLM,
 non-2025 Azure) are left untouched, including when they share a model group
 with an Azure deployment.
 
 Two scope limits worth knowing:
 
-  * ASYNC PATH ONLY. `async_pre_call_deployment_hook` is dispatched by the
-    @client decorator's ASYNC wrapper (litellm/utils.py, in wrapper_async).
-    The sync wrapper does not dispatch it, so a direct sync
+  * CHAT COMPLETIONS ONLY, enforced via `call_type`. The dispatch is NOT
+    chat-specific: `wrapper_async` runs it for every @client-decorated async
+    entrypoint. `litellm.anthropic_messages` (which backs /v1/messages, and
+    which bridges Azure and other non-Anthropic providers) declares
+    `max_tokens: int` as a REQUIRED parameter, so popping it there makes
+    litellm's own wrapper raise
+    `TypeError: anthropic_messages() missing 1 required positional argument`
+    on the following `await original_function(*args, **kwargs)`. That is
+    outside this hook's try/except and cannot be caught here, and the
+    traceback never names this hook. `atext_completion` (/v1/completions has
+    no `max_completion_tokens`) and `aembedding` reach the same dispatch.
+
+  * ASYNC PATH ONLY. The dispatch lives in the @client decorator's ASYNC
+    wrapper; the sync wrapper does not dispatch it, so a direct sync
     `litellm.completion()` or `Router._completion()` bypasses this hook. That
     is fine for how the hook is deployed: it is registered only in the proxy
     config, and the proxy maps /chat/completions to `acompletion`
-    (proxy/route_llm_request.py), so all proxied traffic takes the async path.
-    Code embedding litellm in-process and calling sync `completion()` would
-    not get the rename.
-
-  * The predicate reads the MERGED kwargs, which cannot distinguish a
-    deployment-configured `max_completion_tokens` from a caller-supplied one.
-    A request sending BOTH `max_tokens` and `max_completion_tokens` to a
-    max_tokens-native deployment is therefore rewritten. The reported case,
-    and everything h2oGPTe core emits, sends `max_tokens` alone. Keying only
-    on `additional_drop_params` would remove this edge, at the cost of not
-    firing for an Azure deployment configured with a ceiling but no drop.
+    (proxy/route_llm_request.py), so all proxied traffic is async.
 
 INTERACTION WITH THE CAP HOOK
 -----------------------------
@@ -122,25 +133,45 @@ def _positive_int(value: Any) -> Optional[int]:
         return value if value > 0 else None
 
     @staticmethod
-    def _deployment_uses_completion_tokens(kwargs: Dict[str, Any]) -> bool:
-        """True when the SELECTED deployment expects `max_completion_tokens`.
+    def _is_chat_completion(call_type: Any) -> bool:
+        """True only for the chat-completion call types.
+
+        `call_type` is a `CallTypes` enum member (or None for an unrecognised
+        entrypoint), so compare on its value and tolerate a bare string.
+        """
+        return getattr(call_type, "value", call_type) in ("completion", "acompletion")
+
+    @staticmethod
+    def _deployment_discards_max_tokens(kwargs: Dict[str, Any]) -> bool:
+        """True when the SELECTED deployment would discard or reject
+        `max_tokens`, so the caller's limit only survives as
+        `max_completion_tokens`.
 
         Read straight off the merged kwargs rather than the router, so a mixed
         model group is judged per selected deployment instead of per group.
         """
-        if kwargs.get("max_completion_tokens") is not None:
-            return True
         drop = kwargs.get("additional_drop_params") or []
-        return isinstance(drop, (list, tuple)) and "max_tokens" in drop
+        if isinstance(drop, (list, tuple)) and "max_tokens" in drop:
+            return True
+        # Azure 2025+ rejects max_tokens outright. Reasoning Azure entries are
+        # exempt from the drop above but still need the rename, so recognise
+        # the route itself. custom_llm_provider is not always populated at this
+        # point, so the prefixed model string is the primary signal.
+        if kwargs.get("custom_llm_provider") == "azure":
+            return True
+        model = kwargs.get("model")
+        return isinstance(model, str) and model.startswith("azure/")
 
     async def async_pre_call_deployment_hook(
         self, kwargs: Dict[str, Any], call_type: Any
     ) -> Optional[dict]:
         try:
+            if not self._is_chat_completion(call_type):
+                return None
             requested = self._positive_int(kwargs.get("max_tokens"))
             if requested is None:
                 return None
-            if not self._deployment_uses_completion_tokens(kwargs):
+            if not self._deployment_discards_max_tokens(kwargs):
                 return None
 
             # A configured deployment ceiling (or a caller-supplied value) must

tests/test_litellm/integrations/h2o/test_max_tokens_rename_hook.py

@@ -130,12 +130,42 @@ async def test_plain_deployment_with_no_params_is_a_noop(hook):
 
 @pytest.mark.asyncio
 async def test_malformed_drop_params_never_propagates(hook):
-    """A non-list additional_drop_params must not fail the request."""
-    out = await _run(hook, {"model": "azure/gpt-4o-mini", "max_tokens": 50,
+    """A non-list additional_drop_params must not raise. On a non-Azure model
+    it is the only possible signal, so the request is left alone."""
+    out = await _run(hook, {"model": "openrouter/some-model", "max_tokens": 50,
                             "additional_drop_params": "max_tokens"})
     assert out is None
 
 
+@pytest.mark.asyncio
+async def test_azure_route_alone_triggers_the_rename(hook):
+    """Reasoning Azure entries get max_completion_tokens but are exempt from
+    the drop, so the route itself has to be recognised."""
+    out = await _run(hook, {"model": "azure/gpt-5-mini", "max_tokens": 50,
+                            "max_completion_tokens": 16384})
+    assert out["max_completion_tokens"] == 50
+    assert "max_tokens" not in out
+
+
+@pytest.mark.asyncio
+async def test_non_azure_ceiling_alone_does_not_trigger_the_rename(hook):
+    """convert_model_to_litellm_config sets max_completion_tokens for EVERY
+    reasoning model, Azure or not, and only the Azure branch adds the drop.
+    Renaming on a non-Azure provider would strip max_tokens and leave the
+    request with no ceiling at all, which is worse than the bug being fixed."""
+    out = await _run(hook, {"model": "openrouter/grok-4", "max_tokens": 50,
+                            "max_completion_tokens": 16384})
+    assert out is None
+
+
+@pytest.mark.asyncio
+async def test_native_deployment_with_caller_supplied_ceiling_is_untouched(hook):
+    """A caller sending BOTH fields to a max_tokens-native deployment must not
+    have max_tokens stripped."""
+    out = await _run(hook, dict(BEDROCK_KWARGS, max_completion_tokens=4096))
+    assert out is None
+
+
 @pytest.mark.asyncio
 async def test_other_kwargs_are_preserved(hook):
     """The hook returns the FULL kwargs dict; the dispatcher replaces kwargs
@@ -235,3 +265,44 @@ async def test_litellm_dispatch_leaves_a_max_tokens_native_deployment_alone():
     assert seen is not None
     assert seen["max_tokens"] == 50, seen
     assert seen.get("max_completion_tokens") is None, seen
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_survives_the_rename_hook():
+    """litellm.anthropic_messages backs /v1/messages and bridges Azure, and it
+    declares max_tokens as a REQUIRED parameter. The deployment-hook dispatch
+    is not chat-specific, so popping max_tokens there made litellm's own
+    @client wrapper raise
+    `TypeError: anthropic_messages() missing 1 required positional argument`
+    on the next `await original_function(...)`. That is outside this hook's
+    try/except, and the traceback never names the hook.
+    """
+    import litellm
+
+    previous = litellm.callbacks
+    litellm.callbacks = [MaxTokensRenameHook()]
+    try:
+        await litellm.anthropic_messages(
+            model="azure/gpt-4o-mini",
+            messages=[{"role": "user", "content": "hi"}],
+            max_tokens=50,
+            max_completion_tokens=16384,        # merged in by the router
+            additional_drop_params=["max_tokens"],
+            api_key="dummy",
+            api_version="2025-04-01-preview",
+            api_base="https://example.openai.azure.com",
+            mock_response="ok",
+        )
+    finally:
+        litellm.callbacks = previous
+
+
+@pytest.mark.asyncio
+async def test_non_chat_call_types_are_skipped():
+    """Same guard at the unit level, for the other entrypoints that reach this
+    dispatch: /v1/completions has no max_completion_tokens, and embeddings
+    have no notion of one."""
+    hook = MaxTokensRenameHook()
+    for call_type in ("anthropic_messages", "atext_completion", "aembedding", None):
+        out = await hook.async_pre_call_deployment_hook(dict(AZURE_KWARGS), call_type)
+        assert out is None, f"{call_type} must not be rewritten"