From ef030235fd6267f685d879cfbfaa89f4bc671895 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 2 Jul 2026 17:32:46 -0700 Subject: [PATCH] test(e2e): add vertex_ai passthrough spend-log coverage (#31781) * test(e2e): add vertex_ai passthrough spend-log coverage Port the de-flake of the SDK-based vertex spend test (#31689) into the tests/e2e/llm_translation harness. The vertexai SDK intermittently ignored the proxy api_endpoint override and billed Vertex directly, so the request never reached LiteLLM and no spend was logged; driving native generateContent over the shared transport always reaches the proxy, which the harness already guarantees. The vertex deployment is added at runtime through /model/new with use_in_pass_through rather than declared in the gateway config, and deleted on teardown. That registers the deployment's service account for the /vertex_ai route, so the passthrough call sends only its litellm virtual key in x-litellm-api-key and no upstream bearer, and the proxy mints the Vertex token itself. The credential is the one the proxy already holds, read from the same VERTEXAI_CREDENTIALS/VERTEXAI_PROJECT env; the test never mints a token. Asserts both that the forward succeeds and that a costed SpendLogs row lands (vertex_ai provider, a gemini model, spend > 0, call_type pass_through_endpoint), correlated by the x-litellm-call-id header. * Update tests/e2e/llm_translation/test_vertex_passthrough_e2e.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update tests/e2e/llm_translation/test_vertex_passthrough_e2e.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 16 +- .../e2e/llm_translation/passthrough_client.py | 27 +++ .../test_vertex_passthrough_e2e.py | 166 ++++++++++++++++++ 3 files changed, 206 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/llm_translation/test_vertex_passthrough_e2e.py diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index 5e4a448857f..44d6e79122e 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -28,7 +28,7 @@ Status: `covered` / `partial` / `gap`. |----------|---------------|-----------|------------|-------------|--------| | Gemini (`/gemini/v1beta/models/{m}:generateContent` / `:streamGenerateContent`) | live | live | live | live | **covered** | | Anthropic (`/anthropic/v1/messages`) | live | live | live | live | **covered** | -| Vertex AI (`/vertex_ai/...`) | - | - | - | - | gap (gcloud auth) | +| Vertex AI (`/vertex_ai/v1/projects/{p}/locations/{loc}/.../models/{m}:generateContent`) | live | - | - | live | **partial** | | OpenAI / Bedrock / Cohere / Mistral / VLLM | - | - | - | - | gap | Each covered cell asserts: `call_type == "pass_through_endpoint"`, `spend > 0`, @@ -60,11 +60,21 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_nonstreaming_logs_cost` | anthropic native, non-stream, cost | | `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | +| `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | + +Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is +added at runtime instead of declared in the gateway config: the test POSTs `/model/new` +with `use_in_pass_through`, so the proxy registers that deployment's service account for +the `/vertex_ai` route, then deletes it on teardown. The passthrough call sends only its +litellm virtual key (`x-litellm-api-key`), no upstream bearer, and the proxy mints the +Vertex token itself. Credentials (`VERTEXAI_PROJECT` / `VERTEXAI_CREDENTIALS`) are read +from the same env the proxy uses, so the test never mints a token. ## Gaps -- Vertex / OpenAI / Bedrock / Cohere passthrough (same shape; add once the - provider credential is configured; Vertex is closest - route exists, auth stale). +- Vertex streaming / tool-call passthrough (non-streaming + cost now covered). +- OpenAI / Bedrock / Cohere passthrough (same shape; add once the provider + credential is configured). - Non-passthrough tool calls over `/chat/completions` end to end with cost. - Image / audio / rerank / responses / realtime translation + cost. - Streaming cost-injection (`include_cost_in_streaming_usage`); passthrough on diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py index fff4064a328..77dcf68a1e3 100644 --- a/tests/e2e/llm_translation/passthrough_client.py +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -48,6 +48,16 @@ class AnthropicHeaders(Headers): tags: str | None = None +class VertexHeaders(Headers): + # Only the litellm virtual key; the /vertex_ai passthrough mints the Vertex token + # from the proxy's own service account (the deployment marked use_in_pass_through), + # so no upstream Authorization bearer is sent from the client. + x_litellm_api_key: str = Field(serialization_alias="x-litellm-api-key") + content_type: str = Field( + default="application/json", serialization_alias="Content-Type" + ) + + class AltSseParams(BaseModel): alt: str = "sse" @@ -132,6 +142,23 @@ class PassthroughClient: stream=True, ) + # ---- Vertex AI native passthrough (/vertex_ai/v1/projects/...) ------- + + def vertex_generate( + self, key: str, project: str, location: str, model: str, text: str + ) -> StreamingResponse: + path = ( + f"/vertex_ai/v1/projects/{project}/locations/{location}" + f"/publishers/google/models/{model}:generateContent" + ) + return self.gateway.transport.send( + path, + headers=VertexHeaders(x_litellm_api_key=key), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=text)])] + ), + ) + # ---- Anthropic native passthrough (/anthropic/v1/messages) ---------- def anthropic_message( diff --git a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py new file mode 100644 index 00000000000..78d2bb358d2 --- /dev/null +++ b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py @@ -0,0 +1,166 @@ +"""Live e2e: a native Vertex AI generateContent call over the proxy's /vertex_ai +passthrough is forwarded to Vertex and still logged as a costed SpendLogs row. + +Ports the de-flake of the SDK-based spend test (#31689). That test configured the +vertexai SDK with an api_endpoint override pointing at the proxy, but the SDK +intermittently ignored the override and billed the public Vertex endpoint directly, +so the request never reached LiteLLM and no spend was recorded; the bypass, not +logging lag, was the flake. Driving raw HTTP through the shared transport always +reaches the proxy, which the harness already guarantees, so the only residual +nondeterminism is the ~60s async spend flush the poll absorbs. + +The vertex deployment is added at runtime through the management endpoint rather than +declared in the gateway config: the test POSTs /model/new with use_in_pass_through so +the proxy registers that deployment's service account for the /vertex_ai route, then +deletes it on teardown. The credential is the one the proxy already holds (read from +the same VERTEXAI_CREDENTIALS the deployment uses), so the passthrough call sends only +its litellm virtual key in x-litellm-api-key and no upstream bearer, and the proxy +mints the Vertex token itself. The test never mints a token. + +Asserts both sides of the promise: the forward succeeds (2xx with a candidate) and +the costed row lands (call_type pass_through_endpoint, vertex_ai provider, a gemini +model, spend > 0), correlated by the x-litellm-call-id header. +""" + +import os + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import NoBody, require_successful_call, unwrap +from lifecycle import ResourceManager +from models import SpendLogRow +from passthrough_client import PassthroughClient + +pytestmark = pytest.mark.e2e + +VERTEX_MODEL = "gemini-2.5-flash" +# The added deployment's region and the passthrough URL's region are the same constant, +# so they always agree; the proxy registers passthrough credentials per project+region. +VERTEX_LOCATION = os.environ.get("VERTEXAI_LOCATION", "us-central1") + + +@pytest.fixture(scope="session") +def vertex_project() -> str: + """The Vertex project to bill, read from the same VERTEXAI_PROJECT the proxy uses. + Skip when unset, since that is an environment gap rather than a behavior failure.""" + project = os.environ.get("VERTEXAI_PROJECT") + if not project: + pytest.skip("set VERTEXAI_PROJECT (the project the vertex deployment bills)") + return project + + +@pytest.fixture(scope="session") +def vertex_credentials() -> str: + """The service-account JSON the added deployment authenticates with, the same + VERTEXAI_CREDENTIALS the proxy holds. Skip when unset.""" + credentials = os.environ.get("VERTEXAI_CREDENTIALS") + if not credentials: + pytest.skip("set VERTEXAI_CREDENTIALS (the vertex service-account JSON)") + return credentials + + +class _VertexDeploymentParams(BaseModel): + model: str + vertex_project: str + vertex_location: str + vertex_credentials: str + use_in_pass_through: bool + + +class _ModelInfoId(BaseModel): + id: str + + +class _ModelNewBody(BaseModel): + model_name: str + litellm_params: _VertexDeploymentParams + model_info: _ModelInfoId + + +class _ModelNewResponse(BaseModel): + model_id: str + + +class _ModelDeleteBody(BaseModel): + id: str + + +def _add_vertex_passthrough_model( + client: PassthroughClient, model_name: str, project: str, credentials: str +) -> str: + return unwrap( + client.gateway.transport.post( + "/model/new", + headers=client.gateway.transport.master, + json=_ModelNewBody( + model_name=model_name, + litellm_params=_VertexDeploymentParams( + model=f"vertex_ai/{VERTEX_MODEL}", + vertex_project=project, + vertex_location=VERTEX_LOCATION, + vertex_credentials=credentials, + use_in_pass_through=True, + ), + model_info=_ModelInfoId(id=model_name), + ), + response_type=_ModelNewResponse, + ) + ).model_id + + +def _delete_model(client: PassthroughClient, model_id: str) -> None: + _ = client.gateway.transport.post( + "/model/delete", + headers=client.gateway.transport.master, + json=_ModelDeleteBody(id=model_id), + response_type=NoBody, + ) + + +def _costed_row(client: PassthroughClient, call_id: str | None) -> SpendLogRow: + """The passthrough call's SpendLogs row, polled until it carries a cost. + + A 2xx passthrough call that produced no costed row is a hard failure, not a skip: + a billed Vertex call that LiteLLM did not track is the exact regression #31689 + guards against.""" + assert call_id, "vertex passthrough response had no x-litellm-call-id header" + rows = client.gateway.poll_logs_for_request_id( + call_id, + predicate=lambda rs: (rs[0].spend or 0) > 0, + ) + assert rows, f"no SpendLogs row for vertex passthrough call_id {call_id}" + row = rows[0] + assert row.call_type == "pass_through_endpoint", f"unexpected call_type: {row}" + assert (row.spend or 0) > 0, f"vertex passthrough call was not costed: {row}" + assert row.status == "success", f"unexpected status: {row}" + return row + + +class TestVertexPassthroughSpendTracking: + def test_vertex_passthrough_via_managed_model_logs_cost( + self, + client: PassthroughClient, + scoped_key: str, + resources: ResourceManager, + vertex_project: str, + vertex_credentials: str, + ) -> None: + model_name = f"e2e-vertex-pt-{unique_marker()}" + model_id = _add_vertex_passthrough_model(client, model_name, vertex_project, vertex_credentials) + resources.defer(lambda: _delete_model(client, model_id)) + + result = client.vertex_generate( + key=scoped_key, + project=vertex_project, + location=VERTEX_LOCATION, + model=VERTEX_MODEL, + text=f"reply with one word {unique_marker()}", + ) + require_successful_call(result) + assert '"candidates"' in result.body, f"vertex passthrough returned no candidates: {result.body[:300]}" + + row = _costed_row(client, result.call_id) + assert row.custom_llm_provider == "vertex_ai", f"passthrough spend logged under the wrong provider: {row}" + assert "gemini" in (row.model or ""), f"unexpected model in spend log: {row}"