fix(proxy): skip max_tokens rename for backend-routed openai chat (#2401)

## Description

OpenAI-format `POST /v1/chat/completions` requests routed through
`--backend litellm-vertex` fail when the client includes `max_tokens`.
The proxy currently runs its direct-OpenAI compatibility shim before
backend dispatch, renames `max_tokens` to `max_completion_tokens`, then
the LiteLLM path no longer recognizes that field as standard and sweeps
it into `extra_body`. Vertex rejects the resulting request with
`extra_body: Extra inputs are not permitted`.

This change scopes the rename shim to the direct OpenAI path only.
Backend-routed chat requests now keep `max_tokens`, which LiteLLM
already forwards correctly for the Vertex Anthropic path. Direct GPT-5
and o-series compatibility stays unchanged. Closes #2392.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Thread a backend-owned translation flag into
`_normalize_openai_max_tokens`.
- Skip the legacy-to-completion-token rename on backend-routed OpenAI
chat requests.
- Keep the direct OpenAI compatibility path covered with a backend-owned
translation no-op test.
- Add buffered and streaming handler-level regressions for the exact
`litellm-vertex` request shape, proving the request survives the
`/v1/chat/completions` normalization boundary with vendor fields intact.

## Testing

- [x] Unit tests pass (`uv run pytest
tests/test_proxy/test_openai_backend_path.py
tests/test_openai_streaming_backend.py
tests/test_openai_max_completion_tokens.py
tests/test_litellm_openai_passthrough.py -q`)
- [x] Linting passes (`uv run ruff check
headroom/proxy/handlers/openai.py
tests/test_openai_max_completion_tokens.py
tests/test_litellm_openai_passthrough.py
tests/test_proxy/test_openai_backend_path.py
tests/test_openai_streaming_backend.py`)
- [ ] Type checking passes (`uv run mypy headroom`)
- [x] New tests added for new functionality when applicable
- [x] Manual testing performed

### Test Output

```text
$ uv run pytest tests/test_proxy/test_openai_backend_path.py tests/test_openai_streaming_backend.py tests/test_openai_max_completion_tokens.py tests/test_litellm_openai_passthrough.py -q
......sss............                                                    [100%]
20 passed, 3 skipped, 1 warning in 42.13s

$ uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_max_completion_tokens.py tests/test_litellm_openai_passthrough.py tests/test_proxy/test_openai_backend_path.py tests/test_openai_streaming_backend.py
All checks passed!

$ uv run ruff format headroom/proxy/handlers/openai.py tests/test_openai_max_completion_tokens.py tests/test_litellm_openai_passthrough.py tests/test_proxy/test_openai_backend_path.py tests/test_openai_streaming_backend.py --check
5 files already formatted
```

## Real Behavior Proof

- Environment: Windows, synced Headroom development environment, mocked
LiteLLM provider boundary, no paid GCP credentials required
- Exact command / steps: run `uv run pytest
tests/test_proxy/test_openai_backend_path.py
tests/test_openai_streaming_backend.py
tests/test_openai_max_completion_tokens.py
tests/test_litellm_openai_passthrough.py -q`, using the issue payload
shape
`{"model":"claude-sonnet-4-6","max_tokens":32,"messages":[{"role":"user","content":"hi"}],"chat_template_kwargs":{"enable_thinking":false}}`
through `POST /v1/chat/completions`
- Observed result: buffered and streaming `litellm-vertex` requests keep
`max_tokens` as a named backend kwarg, preserve `chat_template_kwargs`
in `extra_body`, omit `max_completion_tokens` from `extra_body`, and
return success through the handler boundary. Direct-path normalization
still renames legacy `max_tokens`.
- Not tested: live Vertex AI request

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [ ] My code follows the project's style guidelines
- [ ] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

- `CHANGELOG.md`: N/A, the release pipeline generates it from the
conventional-commit subject.
- Scope is intentionally narrow: this fixes the exact backend-routed
`max_tokens` failure and does not broaden `extra_body` hardening for
unrelated OpenAI fields.

headroom/proxy/handlers/openai.py

@@ -130,9 +130,14 @@ _OPENAI_BASE_URL_HEADER = "x-headroom-base-url"
 _decode_openai_bearer_payload = decode_openai_bearer_payload
 
 
-def _normalize_openai_max_tokens(body: dict[str, Any]) -> None:
+def _normalize_openai_max_tokens(
+    body: dict[str, Any], *, backend_owns_translation: bool = False
+) -> None:
     """Rename the legacy ``max_tokens`` to ``max_completion_tokens`` in-place.
 
+    This direct-OpenAI compatibility shim leaves provider-specific translation
+    to backend-routed requests.
+
     GPT-5 / o-series chat models reject ``max_tokens`` and require
     ``max_completion_tokens``; gpt-4o/4.1 accept the latter too. So translating
     is a safe, one-way shim for current OpenAI models that lets openai-compatible
@@ -140,7 +145,7 @@ def _normalize_openai_max_tokens(body: dict[str, Any]) -> None:
     No-op when there is no ``max_tokens``; keeps an already-set
     ``max_completion_tokens`` and just drops the rejected legacy key.
     """
-    if not isinstance(body, dict) or "max_tokens" not in body:
+    if backend_owns_translation or not isinstance(body, dict) or "max_tokens" not in body:
         return
     legacy = body.get("max_tokens")
     if legacy is not None and body.get("max_completion_tokens") is None:
@@ -3406,7 +3411,9 @@ class OpenAIHandlerMixin:
         # translate it here — the proxy already owns the outbound body — and
         # those requests work unchanged. No-op when the caller already set
         # `max_completion_tokens`.
-        _normalize_openai_max_tokens(body)
+        _normalize_openai_max_tokens(
+            body, backend_owns_translation=self.anthropic_backend is not None
+        )
 
         # Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity steering
         # on the chat system message. Runs after every other body mutation so the

tests/test_openai_max_completion_tokens.py

@@ -14,11 +14,17 @@ from headroom.proxy.handlers.openai import _normalize_openai_max_tokens
 
 def test_renames_legacy_max_tokens():
     body = {"model": "gpt-5.3-chat-latest", "max_tokens": 256, "messages": []}
-    _normalize_openai_max_tokens(body)
+    _normalize_openai_max_tokens(body, backend_owns_translation=False)
     assert "max_tokens" not in body
     assert body["max_completion_tokens"] == 256
 
 
+def test_backend_owned_translation_preserves_max_tokens():
+    body = {"model": "claude-sonnet-4-6", "max_tokens": 32, "messages": []}
+    _normalize_openai_max_tokens(body, backend_owns_translation=True)
+    assert body == {"model": "claude-sonnet-4-6", "max_tokens": 32, "messages": []}
+
+
 def test_preserves_existing_max_completion_tokens_and_drops_legacy():
     body = {"max_tokens": 256, "max_completion_tokens": 100}
     _normalize_openai_max_tokens(body)

tests/test_openai_streaming_backend.py

@@ -11,6 +11,7 @@ Run with:
 """
 
 import os
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -260,3 +261,49 @@ class TestOpenAIStreamingMock:
                 assert "application/json" in content_type
                 data = response.json()
                 assert data["choices"][0]["message"]["content"] == "Hello!"
+
+    def test_litellm_vertex_streaming_preserves_max_tokens_and_vendor_fields(self):
+        config = ProxyConfig(
+            optimize=False,
+            cache_enabled=False,
+            rate_limit_enabled=False,
+            backend="litellm-vertex",
+        )
+
+        async def fake_stream():
+            yield SimpleNamespace(
+                model_dump=lambda **kwargs: {
+                    "id": "chunk1",
+                    "choices": [{"delta": {"content": "a"}}],
+                }
+            )
+
+        with (
+            patch("headroom.backends.litellm._fetch_bedrock_inference_profiles", return_value={}),
+            patch("headroom.backends.litellm.acompletion", new_callable=AsyncMock) as mock_acomp,
+        ):
+            mock_acomp.return_value = fake_stream()
+            app = create_app(config)
+
+            with TestClient(app) as client:
+                response = client.post(
+                    "/v1/chat/completions",
+                    json={
+                        "model": "claude-sonnet-4-6",
+                        "messages": [{"role": "user", "content": "hi"}],
+                        "max_tokens": 32,
+                        "chat_template_kwargs": {"enable_thinking": False},
+                        "stream": True,
+                    },
+                    headers={"Authorization": "Bearer test-key"},
+                )
+
+        assert response.status_code == 200, response.text
+        assert "text/event-stream" in response.headers.get("content-type", "")
+        assert "data: [DONE]" in response.text
+
+        kwargs = mock_acomp.await_args.kwargs
+        assert kwargs["stream"] is True
+        assert kwargs["max_tokens"] == 32
+        assert kwargs["extra_body"] == {"chat_template_kwargs": {"enable_thinking": False}}
+        assert "max_completion_tokens" not in kwargs["extra_body"]

tests/test_proxy/test_openai_backend_path.py

@@ -16,6 +16,7 @@ don't need a real provider:
 
 from __future__ import annotations
 
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -91,6 +92,21 @@ def _make_mock_backend(response_body: dict, status_code: int = 200) -> MagicMock
     return backend
 
 
+def _make_litellm_response() -> SimpleNamespace:
+    return SimpleNamespace(
+        id="resp_2392",
+        created=123456,
+        choices=[
+            SimpleNamespace(
+                index=0,
+                finish_reason="stop",
+                message=SimpleNamespace(role="assistant", content="ok", tool_calls=None),
+            )
+        ],
+        usage=SimpleNamespace(prompt_tokens=2, completion_tokens=3, total_tokens=5),
+    )
+
+
 def _install_tracker_stub(client: TestClient) -> _RecordingTracker:
     """Force the session_tracker_store to hand out our recording tracker."""
     tracker = _RecordingTracker(provider="openai")
@@ -196,6 +212,40 @@ def test_backend_response_falls_back_to_openai_cached_tokens_when_bedrock_keys_a
     assert call["cache_write_tokens"] == 300
 
 
+def test_litellm_vertex_backend_path_preserves_max_tokens_and_vendor_fields():
+    config = ProxyConfig(
+        optimize=False,
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        backend="litellm-vertex",
+    )
+
+    with (
+        patch("headroom.backends.litellm._fetch_bedrock_inference_profiles", return_value={}),
+        patch("headroom.backends.litellm.acompletion", new_callable=AsyncMock) as mock_acomp,
+    ):
+        mock_acomp.return_value = _make_litellm_response()
+        app = create_app(config)
+        with TestClient(app) as client:
+            response = client.post(
+                "/v1/chat/completions",
+                json={
+                    "model": "claude-sonnet-4-6",
+                    "messages": [{"role": "user", "content": "hi"}],
+                    "max_tokens": 32,
+                    "chat_template_kwargs": {"enable_thinking": False},
+                    "stream": False,
+                },
+                headers={"Authorization": "Bearer test-key"},
+            )
+
+    assert response.status_code == 200, response.text
+    kwargs = mock_acomp.await_args.kwargs
+    assert kwargs["max_tokens"] == 32
+    assert kwargs["extra_body"] == {"chat_template_kwargs": {"enable_thinking": False}}
+    assert "max_completion_tokens" not in kwargs["extra_body"]
+
+
 def test_backend_response_with_ccr_tool_call_is_intercepted_and_resolved():
     """OpenAI-shape response carrying headroom_retrieve → CCR handler resolves it."""
     config = _make_config()