diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 48fa5e03c..c29661b1d 100644 --- a/headroom/proxy/handlers/openai.py +++ b/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 diff --git a/tests/test_openai_max_completion_tokens.py b/tests/test_openai_max_completion_tokens.py index 91c795712..49740bcaf 100644 --- a/tests/test_openai_max_completion_tokens.py +++ b/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) diff --git a/tests/test_openai_streaming_backend.py b/tests/test_openai_streaming_backend.py index bf1d03ced..5a060a68f 100644 --- a/tests/test_openai_streaming_backend.py +++ b/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"] diff --git a/tests/test_proxy/test_openai_backend_path.py b/tests/test_proxy/test_openai_backend_path.py index ccae9370e..404c699ab 100644 --- a/tests/test_proxy/test_openai_backend_path.py +++ b/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()