From b30a70f34caf8a450708a14081a7bec70bdc578d Mon Sep 17 00:00:00 2001 From: Abhishek Acharya Date: Thu, 22 Oct 2026 21:12:16 +1100 Subject: [PATCH] fix(vertex_ai): send the caller's headers to Vertex Model Garden endpoints VertexAIModelGardenModels.completion received the caller's headers (including extra_headers) but never passed them to the OpenAI like handler, which then built its own headers. Calls to vertex_ai/openai/ endpoints, sync, async and streaming, went out without them. Pass them on in a new dict, keeping the Vertex access token as Authorization like the partner models route does. --- .../vertex_ai/vertex_model_garden/main.py | 7 ++ .../test_vertex_model_garden_openapi.py | 110 +++++++++++++++++- 2 files changed, 116 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index 84907f01685..2f1eb108ca5 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -136,6 +136,13 @@ class VertexAIModelGardenModels(VertexBase): timeout=timeout, encoding=encoding, custom_llm_provider="vertex_ai", + # Send the caller's headers (extra_headers included) in a new dict, and keep + # the Vertex access token as Authorization, like the partner models route + headers={ + "Content-Type": "application/json", + **(headers or {}), + "Authorization": f"Bearer {access_token}", + }, ) except Exception as e: diff --git a/tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py b/tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py index b80d4714253..ac58a2ad56b 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py +++ b/tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py @@ -5,13 +5,17 @@ import os import sys from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm from litellm.llms.vertex_ai.common_utils import ( vertex_model_garden_model_id_in_json_body, ) -from litellm.llms.vertex_ai.vertex_model_garden.main import create_vertex_url +from litellm.llms.vertex_ai.vertex_model_garden.main import ( + VertexAIModelGardenModels, + create_vertex_url, +) @pytest.mark.parametrize( @@ -235,3 +239,107 @@ async def test_default_api_base_when_none_provided_publisher_model( "test-project/locations/us-central1/endpoints/openapi/chat/completions" ) assert request_body["model"] == "xai/grok-4.1-fast-reasoning" + + +USER_API_BASE = "https://my-endpoint.example.com/v1" +CHAT_COMPLETION = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1234567890, + "model": "5464397967697903616", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +STREAM_CHUNK = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "5464397967697903616", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": "stop"}], +} + + +@pytest.fixture +def model_garden_chat_route(respx_mock, monkeypatch, clean_vertex_env, _reset_litellm_http_client_cache): + """The Model Garden route with only the network, the token mint and the optional + vertexai import replaced""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + vertexai_stub = MagicMock() + monkeypatch.setitem(sys.modules, "vertexai", vertexai_stub) + monkeypatch.setitem(sys.modules, "vertexai.preview", vertexai_stub.preview) + monkeypatch.setattr( + VertexAIModelGardenModels, + "_ensure_access_token", + lambda self, **kwargs: ("fake-token", "test-project"), + ) + return respx_mock.post(f"{USER_API_BASE}/chat/completions") + + +def _model_garden_kwargs(**overrides): + return { + "model": "vertex_ai/openai/5464397967697903616", + "messages": [{"role": "user", "content": "hello"}], + "api_base": USER_API_BASE, + "vertex_ai_location": "us-central1", + "vertex_ai_project": "test-project", + **overrides, + } + + +@pytest.mark.parametrize("header_kwarg", ["extra_headers", "headers"]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_model_garden_sends_the_callers_headers_with_the_vertex_token( + model_garden_chat_route, header_kwarg, sync_mode +): + model_garden_chat_route.mock(return_value=httpx.Response(200, json=CHAT_COMPLETION)) + kwargs = _model_garden_kwargs(**{header_kwarg: {"X-Test": "from-caller"}}) + + response = litellm.completion(**kwargs) if sync_mode else await litellm.acompletion(**kwargs) + + assert response.choices[0].message.content == "hi" + sent_headers = model_garden_chat_route.calls[0].request.headers + assert sent_headers["x-test"] == "from-caller" + assert sent_headers["authorization"] == "Bearer fake-token" + assert sent_headers["content-type"] == "application/json" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_model_garden_streaming_sends_the_callers_headers(model_garden_chat_route, sync_mode): + model_garden_chat_route.mock( + return_value=httpx.Response(200, content=f"data: {json.dumps(STREAM_CHUNK)}\n\ndata: [DONE]\n\n") + ) + kwargs = _model_garden_kwargs(stream=True, extra_headers={"X-Test": "from-caller"}) + + if sync_mode: + text = "".join(chunk.choices[0].delta.content or "" for chunk in litellm.completion(**kwargs)) + else: + stream = await litellm.acompletion(**kwargs) + text = "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) + + assert text == "hi" + sent_headers = model_garden_chat_route.calls[0].request.headers + assert sent_headers["x-test"] == "from-caller" + assert sent_headers["authorization"] == "Bearer fake-token" + + +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_model_garden_keeps_the_vertex_token_over_a_caller_authorization_header( + model_garden_chat_route, +): + model_garden_chat_route.mock(return_value=httpx.Response(200, json=CHAT_COMPLETION)) + + await litellm.acompletion(**_model_garden_kwargs(extra_headers={"Authorization": "Bearer from-caller"})) + + sent_headers = model_garden_chat_route.calls[0].request.headers + assert sent_headers["authorization"] == "Bearer fake-token"