scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(vertex): keep Messages authentication off the event loop (#48523)

Browse source
This commit is contained in:
tin-berri 2026-10-26 22:43:23 -04:00 • committed by GitHub
parent cc09b52616
commit f2b59be6ea
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 172 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

92
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py
View file

@ -1,4 +1,4 @@
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import Final
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
@ -26,34 +26,92 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
def should_strip_billing_metadata(self) -> bool:
return True
@staticmethod
def _vertex_auth_parameters(
litellm_params: dict[str, object], # mutable-ok: delegates to the existing Vertex parameter readers
) -> tuple[str | None, str | None, str | None]:
return (
VertexBase.safe_get_vertex_ai_project(litellm_params),
VertexBase.safe_get_vertex_ai_location(litellm_params),
VertexBase.safe_get_vertex_ai_credentials(litellm_params),
)
def validate_anthropic_messages_environment(
self,
headers: dict,
headers: dict[str, str],
model: str,
messages: Sequence[object],
optional_params: dict,
litellm_params: dict,
optional_params: dict[str, object],
litellm_params: dict[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict, str | None]:
) -> tuple[dict[str, str], str | None]:
"""
OPTIONAL
Validate the environment for the request
"""
# Work on a local copy — router shallow-copies litellm_params so the caller's
# headers dict may be the shared deployment extra_headers object.
headers = dict(headers)
vertex_ai_project: Final = VertexBase.safe_get_vertex_ai_project(litellm_params)
vertex_ai_location: Final = VertexBase.safe_get_vertex_ai_location(litellm_params)
vertex_credentials: Final = VertexBase.safe_get_vertex_ai_credentials(litellm_params)
vertex_ai_project, vertex_ai_location, vertex_credentials = self._vertex_auth_parameters(litellm_params)
access_token, project_id = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_ai_project,
custom_llm_provider="vertex_ai",
)
headers["Authorization"] = f"Bearer {access_token}"
return self._finalize_vertex_messages_environment(
headers,
model,
messages,
optional_params,
api_base,
access_token,
project_id,
vertex_ai_project,
vertex_ai_location,
)
async def avalidate_anthropic_messages_environment(
self,
headers: dict[str, str], # mutable-ok: mirrors the provider validation contract
model: str,
messages: Sequence[object],
optional_params: dict[str, object], # mutable-ok: mirrors the provider validation contract
litellm_params: dict[str, object], # mutable-ok: mirrors the provider validation contract
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers
vertex_ai_project, vertex_ai_location, vertex_credentials = self._vertex_auth_parameters(litellm_params)
access_token, project_id = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_ai_project,
custom_llm_provider="vertex_ai",
)
return self._finalize_vertex_messages_environment(
headers,
model,
messages,
optional_params,
api_base,
access_token,
project_id,
vertex_ai_project,
vertex_ai_location,
)
def _finalize_vertex_messages_environment(
self,
headers: Mapping[str, str],
model: str,
messages: Sequence[object],
optional_params: dict, # mutable-ok: preserves the existing provider optional-parameter contract
api_base: str | None,
access_token: str,
project_id: str,
vertex_ai_project: str | None,
vertex_ai_location: str | None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers
# Router shallow-copies litellm_params, so headers may be shared between requests.
request_headers: Final = dict(headers)
request_headers["Authorization"] = f"Bearer {access_token}"
api_base = self.get_complete_vertex_url(
custom_api_base=api_base,
@ -65,14 +123,14 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
model=model,
)
headers["content-type"] = "application/json"
request_headers["content-type"] = "application/json"
# Add beta headers for Vertex AI
tools: Final = optional_params.get("tools", [])
beta_values: Final[set[str]] = set()
# Get existing beta headers if any
existing_beta: Final = headers.get("anthropic-beta")
existing_beta: Final = request_headers.get("anthropic-beta")
if existing_beta:
beta_values.update(b.strip() for b in existing_beta.split(","))
@ -117,9 +175,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value)
if beta_values:
headers["anthropic-beta"] = ",".join(beta_values)
request_headers["anthropic-beta"] = ",".join(beta_values)
return headers, api_base
return request_headers, api_base
def get_complete_url(
self,

22
tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py
View file

@ -6,6 +6,7 @@ import uuid
from typing import Any, Dict, Final, List
import httpx
import responses
import respx
import pytest
from fastapi.testclient import TestClient
@ -1709,16 +1710,20 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet
@pytest.mark.asyncio
@pytest.mark.parametrize("client_headers", _CLIENT_BETA_HEADERS)
async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex(
local_beta_headers_config, client_headers
):
local_beta_headers_config: None, client_headers: dict[str, str]
) -> None:
"""Vertex rawPredict takes the beta as the `anthropic-beta` header and 400s on `safeguards` without it, so the beta rides along with the field."""
from litellm.llms.anthropic.pass_through.messages import handler
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
safeguards, safeguard_results = _claude_code_auto_mode_request()
captured: dict[str, object] = {}
with patch.object(VertexBase, "_ensure_access_token", return_value=("test-token", "test-project")):
with responses.RequestsMock(assert_all_requests_are_fired=False) as oauth:
oauth.add(
responses.POST,
"https://oauth2.googleapis.com/token",
json={"access_token": "test-token", "expires_in": 3600},
)
response = await handler.anthropic_messages(
max_tokens=16,
messages=[{"role": "user", "content": "hi"}],
@ -1726,7 +1731,14 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet
custom_llm_provider="vertex_ai",
vertex_project="test-project",
vertex_location="global",
vertex_credentials="{}",
vertex_credentials=json.dumps(
{
"type": "authorized_user",
"client_id": "synthetic-client",
"client_secret": "synthetic-secret",
"refresh_token": "synthetic-refresh",
}
),
client=_upstream_answering_with(safeguard_results, captured),
safeguards=safeguards,
extra_headers=client_headers,

80
tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py
View file

@ -2,18 +2,98 @@ import copy
import json
import os
from collections.abc import Mapping
from datetime import datetime, timedelta
from queue import SimpleQueue
from threading import get_ident
from typing import Final
from unittest.mock import MagicMock, patch
import pytest
import requests
import responses
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import (
VertexAIPartnerModelsAnthropicMessagesConfig,
)
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS
from litellm.types.router import GenericLiteLLMParams
@pytest.mark.parametrize("stream", (False, True))
@pytest.mark.parametrize("infer_project", (False, True), ids=("explicit-project", "credential-project"))
@pytest.mark.asyncio
async def test_messages_auth_load_and_refresh_leave_the_event_loop(stream: bool, infer_project: bool) -> None:
config: Final = VertexAIPartnerModelsAnthropicMessagesConfig()
loop_thread: Final = get_ident()
token_threads: Final[SimpleQueue[int]] = SimpleQueue()
now: Final = datetime(2025, 1, 1)
model: Final = "claude-sonnet-5"
headers: Final = {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"}
optional_params: Final = {
"stream": stream,
"tools": [{"type": ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value, "name": "web_search"}],
}
litellm_params: Final = {
"vertex_ai_project": None if infer_project else "auth-project",
"vertex_ai_location": "global",
"vertex_credentials": json.dumps(
{
"type": "authorized_user",
"client_id": "synthetic-client",
"client_secret": "synthetic-secret",
"refresh_token": "synthetic-refresh",
"quota_project_id": "auth-project",
}
),
}
async def validate() -> tuple[Mapping[str, object], str | None]:
return await config.avalidate_anthropic_messages_environment(
headers=headers, model=model, messages=[], optional_params=optional_params, litellm_params=litellm_params
)
with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth:
def token_response(_request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
token_threads.put(get_ident())
return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600})
oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
cold_headers, cold_url = await validate()
assert token_threads.get_nowait() != loop_thread
assert cold_headers["Authorization"] == "Bearer token-1"
assert await validate() == (cold_headers, cold_url)
assert len(oauth.calls) == 1
clock.return_value = now + timedelta(hours=2)
refreshed_headers, refreshed_url = await validate()
assert token_threads.get_nowait() != loop_thread
assert refreshed_headers["Authorization"] == "Bearer token-2"
assert refreshed_url == cold_url
assert len(oauth.calls) == 2
assert config.validate_anthropic_messages_environment(
headers=headers, model=model, messages=[], optional_params=optional_params, litellm_params=litellm_params
) == (refreshed_headers, refreshed_url)
assert len(oauth.calls) == 2
endpoint: Final = "streamRawPredict" if stream else "rawPredict"
assert cold_url == (
"https://aiplatform.googleapis.com/v1/projects/auth-project/locations/global/"
f"publishers/anthropic/models/{model}:{endpoint}"
)
assert refreshed_headers == {
"Authorization": "Bearer token-2",
"anthropic-beta": cold_headers["anthropic-beta"],
"x-request-id": "req-1",
"content-type": "application/json",
}
beta_header: Final = cold_headers["anthropic-beta"]
assert isinstance(beta_header, str)
assert set(beta_header.split(",")) == {"retained-beta", ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value}
assert headers == {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"}
def test_validate_environment_uses_vertex_ai_location():
config = VertexAIPartnerModelsAnthropicMessagesConfig()
headers = {}