fix(vertex): label passthrough requests with the key identity and drop forged identity labels

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -69,6 +69,14 @@ from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_
 from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError
 from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base
 from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request
+from litellm.llms.vertex_ai.request_labels import (
+    VERTEX_LABELS_HEADER,
+    VERTEX_REQUEST_LABELS_MAX,
+    is_vertex_labels_header,
+    resolve_vertex_request_labels,
+    vertex_partner_labels_header,
+    vertex_request_label_fields,
+)
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error
 from litellm.proxy._types import *
@@ -2633,6 +2641,42 @@ def _without_dropped_anthropic_betas(
     return MappingProxyType({name: value for name, value in upstream_values if value is not None})
 
 
+_VERTEX_BODY_LABELED_METHODS: Final = (":generateContent", ":streamGenerateContent", ":predict")
+_VERTEX_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object])
+
+
+def _is_vertex_anthropic_route(endpoint: str) -> bool:
+    return "/publishers/anthropic/" in f"/{endpoint}"
+
+
+def _vertex_body_with_resolved_labels(
+    label_fields: tuple[str, ...], identity: Mapping[str, object], body: Mapping[str, object]
+) -> Mapping[str, object]:
+    resolved: Final = resolve_vertex_request_labels(
+        label_fields, body.get("labels"), {"litellm_metadata": identity}, VERTEX_REQUEST_LABELS_MAX
+    )
+    unlabeled: Final = {name: value for name, value in body.items() if name != "labels"}
+    return {**unlabeled, "labels": dict(resolved)} if resolved else unlabeled
+
+
+def _upstream_vertex_headers(
+    dropped_anthropic_betas: frozenset[str],
+    label_fields: tuple[str, ...],
+    identity: Mapping[str, object],
+    headers: Mapping[str, str],
+) -> Mapping[str, str]:
+    without_dropped_betas: Final = _without_dropped_anthropic_betas(dropped_anthropic_betas, headers)
+    if not label_fields:
+        return without_dropped_betas
+    labels_header: Final = vertex_partner_labels_header(
+        label_fields, None, without_dropped_betas, {"litellm_metadata": identity}
+    )
+    unlabeled: Final = {
+        name: value for name, value in without_dropped_betas.items() if not is_vertex_labels_header(name)
+    }
+    return MappingProxyType(unlabeled if labels_header is None else {**unlabeled, VERTEX_LABELS_HEADER: labels_header})
+
+
 def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]:
     if not _is_vertex_anthropic_count_tokens_route(endpoint):
         return headers
@@ -3081,6 +3125,19 @@ async def _base_vertex_proxy_route(
     if is_streaming_request:
         target += "?alt=sse"
 
+    label_fields: Final = vertex_request_label_fields()
+    identity: Final[Mapping[str, object]] = (
+        get_passthrough_router_request_metadata(user_api_key_dict) if label_fields else MappingProxyType({})
+    )
+    if label_fields and encoded_endpoint.endswith(_VERTEX_BODY_LABELED_METHODS):
+        body: Final = _VERTEX_REQUEST_BODY.validate_python(await read_request_body(request))
+        if body:
+            setattr(
+                request.state,
+                LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
+                _vertex_body_with_resolved_labels(label_fields, identity, body),
+            )
+
     ## CREATE PASS-THROUGH
     endpoint_func: Final = create_pass_through_route(
         endpoint=endpoint,
@@ -3088,7 +3145,10 @@ async def _base_vertex_proxy_route(
         custom_headers=_upstream_headers_for_vertex_route(endpoint, headers),
         is_streaming_request=is_streaming_request,
         upstream_headers_filter=partial(
-            _without_dropped_anthropic_betas, _vertex_dropped_anthropic_betas(general_settings, request.headers)
+            _upstream_vertex_headers,
+            _vertex_dropped_anthropic_betas(general_settings, request.headers),
+            label_fields if _is_vertex_anthropic_route(encoded_endpoint) else (),
+            identity,
         ),
     )  # dynamically construct pass-through endpoint based on incoming path
 

tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 import importlib
 import json
 from collections.abc import Iterator
@@ -555,3 +556,133 @@ def test_vertex_pass_through_forwards_a_request_without_betas_despite_a_malforme
 
     assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY)
     assert "anthropic-beta" not in upstream.calls.last.request.headers
+
+
+_VERTEX_GEMINI_PATH: Final = "v1/projects/caller-project/locations/global/publishers/google/models/gemini-3.8-flash"
+_VERTEX_GEMINI_BODY: Final = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}
+_VERTEX_LABELED_KEY: Final = "sk-billing-svc"
+_VERTEX_CALLER_LABELS: Final = {"user_api_key_team_alias": "forged", "env": "prod"}
+_VERTEX_RESOLVED_LABELS: Final = {
+    "user_api_key_alias": "billing-svc",
+    "user_api_key_team_alias": "ai_platform",
+    "env": "prod",
+}
+
+
+def _labels_header(labels: dict[str, str]) -> str:
+    return base64.b64encode(json.dumps(labels).encode()).decode()
+
+
+@pytest.fixture
+def labeled_key(monkeypatch: pytest.MonkeyPatch) -> str:
+    from fastapi import Request
+
+    from litellm.proxy import proxy_server
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    async def billing_service_auth(request: Request, api_key: str) -> UserAPIKeyAuth:
+        assert api_key == _VERTEX_LABELED_KEY
+        return UserAPIKeyAuth(api_key=api_key, key_alias="billing-svc", team_alias="AI Platform")
+
+    monkeypatch.setattr(proxy_server, "user_custom_auth", billing_service_auth)
+    return _VERTEX_LABELED_KEY
+
+
+@pytest.mark.parametrize("method", [":generateContent", ":streamGenerateContent", ":predict"])
+@pytest.mark.parametrize(
+    ("label_fields", "sent_labels", "upstream_labels"),
+    [
+        pytest.param(
+            ["user_api_key_alias", "user_api_key_team_alias"],
+            _VERTEX_CALLER_LABELS,
+            _VERTEX_RESOLVED_LABELS,
+            id="identity-replaces-forged-label",
+        ),
+        pytest.param(["user_api_key_org_id"], {"user_api_key_team_alias": "forged"}, None, id="only-forged-label"),
+        pytest.param(None, _VERTEX_CALLER_LABELS, _VERTEX_CALLER_LABELS, id="setting-off"),
+    ],
+)
+def test_vertex_pass_through_labels_the_request_body_from_the_key_identity(
+    vertex_proxy: TestClient,
+    respx_mock: "respx.MockRouter",
+    monkeypatch: pytest.MonkeyPatch,
+    labeled_key: str,
+    method: str,
+    label_fields: list[str] | None,
+    sent_labels: dict[str, str],
+    upstream_labels: dict[str, str] | None,
+) -> None:
+    monkeypatch.setattr(litellm, "vertex_request_label_fields", label_fields)
+    upstream: Final = respx_mock.post(
+        url__startswith=f"https://aiplatform.googleapis.com/{_VERTEX_GEMINI_PATH}{method}"
+    ).mock(return_value=httpx.Response(200, json={}))
+
+    response: Final = vertex_proxy.post(
+        f"/vertex_ai/{_VERTEX_GEMINI_PATH}{method}",
+        json={**_VERTEX_GEMINI_BODY, "labels": sent_labels},
+        headers={"x-litellm-api-key": labeled_key, "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}"},
+    )
+
+    assert response.status_code == 200
+    expected_labels: Final = {} if upstream_labels is None else {"labels": upstream_labels}
+    assert json.loads(upstream.calls.last.request.content) == {**_VERTEX_GEMINI_BODY, **expected_labels}
+
+
+@pytest.mark.parametrize(
+    ("label_fields", "caller_headers", "upstream_labels"),
+    [
+        pytest.param(
+            ["user_api_key_alias", "user_api_key_team_alias"],
+            {"x-pass-x-vertex-ai-labels": _labels_header(_VERTEX_CALLER_LABELS)},
+            _VERTEX_RESOLVED_LABELS,
+            id="identity-replaces-forged-header",
+        ),
+        pytest.param(
+            ["user_api_key_alias", "user_api_key_team_alias"],
+            {},
+            {"user_api_key_alias": "billing-svc", "user_api_key_team_alias": "ai_platform"},
+            id="identity-without-caller-header",
+        ),
+        pytest.param(
+            ["user_api_key_org_id"],
+            {"x-pass-x-vertex-ai-labels": _labels_header({"user_api_key_team_alias": "forged"})},
+            None,
+            id="only-forged-header",
+        ),
+        pytest.param(
+            None,
+            {"x-pass-x-vertex-ai-labels": _labels_header(_VERTEX_CALLER_LABELS)},
+            _VERTEX_CALLER_LABELS,
+            id="setting-off",
+        ),
+    ],
+)
+def test_vertex_pass_through_labels_the_claude_request_header_from_the_key_identity(
+    vertex_proxy: TestClient,
+    respx_mock: "respx.MockRouter",
+    monkeypatch: pytest.MonkeyPatch,
+    labeled_key: str,
+    label_fields: list[str] | None,
+    caller_headers: dict[str, str],
+    upstream_labels: dict[str, str] | None,
+) -> None:
+    monkeypatch.setattr(litellm, "vertex_request_label_fields", label_fields)
+    upstream: Final = respx_mock.post(f"https://aiplatform.googleapis.com/{_VERTEX_CALLER_PATH}").mock(
+        return_value=httpx.Response(200, json=_VERTEX_CLAUDE_REPLY)
+    )
+
+    response: Final = vertex_proxy.post(
+        f"/vertex_ai/{_VERTEX_CALLER_PATH}",
+        json=_VERTEX_CLAUDE_BODY,
+        headers={
+            "x-litellm-api-key": labeled_key,
+            "Authorization": f"Bearer {_VERTEX_CALLER_GOOGLE_TOKEN}",
+            **caller_headers,
+        },
+    )
+
+    assert (response.status_code, response.json()) == (200, _VERTEX_CLAUDE_REPLY)
+    outbound: Final = upstream.calls.last.request
+    sent_headers: Final = tuple(json.loads(base64.b64decode(value)) for value in outbound.headers.get_list("x-vertex-ai-labels"))
+    assert sent_headers == (() if upstream_labels is None else (upstream_labels,))
+    assert json.loads(outbound.content) == _VERTEX_CLAUDE_BODY