From 1640b3065f36cf2c325956478b311279433bf264 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 30 Jun 2026 04:03:57 +0000 Subject: [PATCH 1/3] test(pass-through): de-flake vertex spend-log assertion by re-billing The vertex pass-through spend-log test asserted that a single billed generateContent call moved the global spend aggregate within a fixed wait. CI failures show the call returning a valid response with real usage, yet spend never increasing over a 240s poll. Pass-through spend logging is best-effort: the success handler is enqueued on a background worker that can drop or time out an individual event under load and never retries it, so one billed call occasionally never reaches LiteLLM_SpendLogs. Waiting longer cannot recover a dropped event; only re-issuing the call can. Re-bill the call up to a few times and require at least one to be tracked, mirroring the sibling jest test that already retries. The test still fails hard if cost tracking is actually broken, since then every call records nothing. Also sum spend across all returned days instead of matching the runner's local 'today', removing a separate UTC-rollover flake. --- tests/pass_through_tests/test_vertex_ai.py | 108 ++++++++++----------- 1 file changed, 50 insertions(+), 58 deletions(-) diff --git a/tests/pass_through_tests/test_vertex_ai.py b/tests/pass_through_tests/test_vertex_ai.py index e8223f2219c..ce941c70ef9 100644 --- a/tests/pass_through_tests/test_vertex_ai.py +++ b/tests/pass_through_tests/test_vertex_ai.py @@ -57,38 +57,28 @@ def load_vertex_ai_credentials(): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) -async def call_spend_logs_endpoint(): - """ - Call this - curl -X GET "http://0.0.0.0:4000/spend/logs" -H "Authorization: Bearer sk-1234" - """ - import datetime - import requests +SPEND_LOG_API_KEY = "best-api-key-ever" - todays_date = datetime.datetime.now().strftime("%Y-%m-%d") - url = f"http://0.0.0.0:4000/global/spend/logs?api_key=best-api-key-ever" - headers = {"Authorization": f"Bearer sk-1234"} - response = requests.get(url, headers=headers) - print("response from call_spend_logs_endpoint", response) - if response.status_code != 200: - print(f"spend logs endpoint returned {response.status_code}: {response.text}") - return None - - json_response = response.json() - - # get spend for today +async def get_tracked_spend() -> float: """ - json response looks like this + Total spend recorded under the pass-through key in the global spend view. - [{'date': '2024-08-30', 'spend': 0.00016600000000000002, 'api_key': 'best-api-key-ever'}] + Sums every day the endpoint returns instead of matching the runner's local + "today" so a UTC date rollover mid-test can't hide a freshly billed call, and + treats an unreachable endpoint as "nothing recorded yet" (0.0). """ - print("json_response", json_response) + import requests - todays_date = datetime.datetime.now().strftime("%Y-%m-%d") - for spend_log in json_response: - if spend_log["date"] == todays_date: - return spend_log["spend"] + url = f"http://0.0.0.0:4000/global/spend/logs?api_key={SPEND_LOG_API_KEY}" + response = requests.get(url, headers={"Authorization": "Bearer sk-1234"}) + if response.status_code != 200: + print(f"global spend logs endpoint returned {response.status_code}: {response.text}") + return 0.0 + + rows = response.json() + print("global spend logs rows", rows) + return sum(float(row.get("spend") or 0.0) for row in rows) LITE_LLM_ENDPOINT = "http://localhost:4000" @@ -106,7 +96,6 @@ def _is_vertex_quota_error(exc: Exception) -> bool: @pytest.mark.asyncio() async def test_basic_vertex_ai_pass_through_with_spendlog(): - spend_before = await call_spend_logs_endpoint() or 0.0 load_vertex_ai_credentials() vertexai.init( @@ -117,38 +106,41 @@ async def test_basic_vertex_ai_pass_through_with_spendlog(): ) model = GenerativeModel(model_name="gemini-3.1-flash-lite") - try: - response = model.generate_content("hi") - except Exception as exc: - if _is_vertex_quota_error(exc): - pytest.skip("Vertex AI quota exhausted") - raise - print("response", response) + spend_before = await get_tracked_spend() - # Spend logging is async/batched and can lag under CI load, so poll instead of - # sleeping a fixed amount. A transient empty read is skipped, not counted as 0.0 - # spend, which would spuriously fail the assertion on an otherwise-billed call. - max_wait = 240 # total seconds to wait - poll_interval = 10 # seconds between checks - elapsed = 0 - spend_after = spend_before - while elapsed < max_wait: - await asyncio.sleep(poll_interval) - elapsed += poll_interval - latest_spend = await call_spend_logs_endpoint() - if latest_spend is None: - print(f"spend logs unavailable (elapsed={elapsed}s), retrying") - continue - spend_after = latest_spend - print(f"spend_after (elapsed={elapsed}s)", spend_after) - if spend_after > spend_before: - break + # Pass-through spend logging is best-effort: the success handler runs on a + # background worker that can drop or time out an individual event under CI load + # and never retries it, so one billed call occasionally never reaches + # LiteLLM_SpendLogs. Waiting longer can't recover a dropped event; only re-billing + # can. Require at least one of a few billed calls to be tracked. This still fails + # hard if cost tracking is genuinely broken, since then every call records nothing. + max_attempts = 5 + per_attempt_wait = 45 # seconds to wait for a single call's spend to land + poll_interval = 5 - assert ( - spend_after > spend_before - ), "Spend should be greater than before after {}s. spend_before: {}, spend_after: {}".format( - elapsed, spend_before, spend_after + spend_after = spend_before + for attempt in range(1, max_attempts + 1): + try: + model.generate_content("hi") + except Exception as exc: + if _is_vertex_quota_error(exc): + pytest.skip("Vertex AI quota exhausted") + raise + + for _ in range(per_attempt_wait // poll_interval): + await asyncio.sleep(poll_interval) + spend_after = await get_tracked_spend() + if spend_after > spend_before: + print(f"spend tracked on attempt {attempt}: {spend_before} -> {spend_after}") + return + + print(f"attempt {attempt}: spend not tracked yet (spend_after={spend_after}), re-billing") + + pytest.fail( + "Vertex pass-through spend never recorded after {} billed calls. spend_before: {}, spend_after: {}".format( + max_attempts, spend_before, spend_after + ) ) @@ -156,7 +148,7 @@ async def test_basic_vertex_ai_pass_through_with_spendlog(): @pytest.mark.skip(reason="skip flaky test - vertex pass through streaming is flaky") async def test_basic_vertex_ai_pass_through_streaming_with_spendlog(): - spend_before = await call_spend_logs_endpoint() or 0.0 + spend_before = await get_tracked_spend() print("spend_before", spend_before) load_vertex_ai_credentials() @@ -176,7 +168,7 @@ async def test_basic_vertex_ai_pass_through_streaming_with_spendlog(): print("response", response) await asyncio.sleep(20) - spend_after = await call_spend_logs_endpoint() + spend_after = await get_tracked_spend() print("spend_after", spend_after) assert ( spend_after > spend_before From cf0453f86c0f044eafa143212b9aeaa1dd8b8bd8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 30 Jun 2026 16:03:12 +0000 Subject: [PATCH 2/3] test(pass-through): route vertex spend-log test through proxy via direct HTTP The vertexai SDK, configured with location="global" and an http api_endpoint override, intermittently sends generateContent to the public Vertex endpoint instead of the proxy. Proxy logs from a failing run show all 46 of the test's own spend-log polls reaching the proxy while zero generateContent calls did, so LiteLLM never saw the billed call and no spend was ever recorded; re-billing through the SDK could not help because every retry bypassed the proxy too. Issue the pass-through request directly over HTTP so it always hits the proxy, minting a Google token from the same service-account credentials, then assert that the specific call's own spend log lands with spend > 0, a gemini model, and custom_llm_provider vertex_ai. A small best-effort retry covers the rare case where the background logging worker drops a single event; failing every attempt still fails hard so the test keeps its teeth if cost tracking breaks. --- tests/pass_through_tests/test_vertex_ai.py | 111 +++++++++++++-------- 1 file changed, 68 insertions(+), 43 deletions(-) diff --git a/tests/pass_through_tests/test_vertex_ai.py b/tests/pass_through_tests/test_vertex_ai.py index ce941c70ef9..f72ae1dfc14 100644 --- a/tests/pass_through_tests/test_vertex_ai.py +++ b/tests/pass_through_tests/test_vertex_ai.py @@ -11,6 +11,7 @@ import os import pytest import asyncio +import requests # Path to your service account JSON file SERVICE_ACCOUNT_FILE = "path/to/your/service-account.json" @@ -68,8 +69,6 @@ async def get_tracked_spend() -> float: "today" so a UTC date rollover mid-test can't hide a freshly billed call, and treats an unreachable endpoint as "nothing recorded yet" (0.0). """ - import requests - url = f"http://0.0.0.0:4000/global/spend/logs?api_key={SPEND_LOG_API_KEY}" response = requests.get(url, headers={"Authorization": "Bearer sk-1234"}) if response.status_code != 200: @@ -83,64 +82,90 @@ async def get_tracked_spend() -> float: LITE_LLM_ENDPOINT = "http://localhost:4000" +VERTEX_PROJECT = "litellm-ci-cd" +VERTEX_MODEL = "gemini-3.1-flash-lite" +VERTEX_GENERATE_CONTENT_URL = ( + f"{LITE_LLM_ENDPOINT}/vertex_ai/v1/projects/{VERTEX_PROJECT}" + f"/locations/global/publishers/google/models/{VERTEX_MODEL}:generateContent" +) -def _is_vertex_quota_error(exc: Exception) -> bool: - message = str(exc) - return ( - "429" in message - or "Too Many Requests" in message - or "RESOURCE_EXHAUSTED" in message - ) +def _vertex_access_token() -> str: + import google.auth + import google.auth.transport.requests -@pytest.mark.asyncio() -async def test_basic_vertex_ai_pass_through_with_spendlog(): + credentials, _ = google.auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + credentials.refresh(google.auth.transport.requests.Request()) + return credentials.token - load_vertex_ai_credentials() - vertexai.init( - project="litellm-ci-cd", - location="global", - api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex_ai", - api_transport="rest", +def _spend_log_for_request(call_id: str) -> dict | None: + response = requests.get( + f"{LITE_LLM_ENDPOINT}/spend/logs?request_id={call_id}", + headers={"Authorization": "Bearer sk-1234"}, + timeout=30, ) + if response.status_code != 200: + return None + rows = response.json() + return rows[0] if rows else None - model = GenerativeModel(model_name="gemini-3.1-flash-lite") - spend_before = await get_tracked_spend() +def _is_vertex_quota_error(response: requests.Response) -> bool: + return response.status_code == 429 or "RESOURCE_EXHAUSTED" in response.text + - # Pass-through spend logging is best-effort: the success handler runs on a - # background worker that can drop or time out an individual event under CI load - # and never retries it, so one billed call occasionally never reaches - # LiteLLM_SpendLogs. Waiting longer can't recover a dropped event; only re-billing - # can. Require at least one of a few billed calls to be tracked. This still fails - # hard if cost tracking is genuinely broken, since then every call records nothing. - max_attempts = 5 - per_attempt_wait = 45 # seconds to wait for a single call's spend to land +@pytest.mark.asyncio() +async def test_basic_vertex_ai_pass_through_with_spendlog(): + load_vertex_ai_credentials() + access_token = _vertex_access_token() + + # Drive the pass-through over HTTP instead of the vertexai SDK: the SDK intermittently + # routes generateContent to the public Vertex endpoint rather than the proxy override, + # so the call never reaches LiteLLM and no spend is logged. A direct request always + # hits the proxy. Spend logging then runs on a best-effort background worker that can + # drop a single event, so retry a few billed calls and assert that one specific call's + # spend log lands. Failing every attempt still fails hard, which is the signal we want + # if cost tracking is broken. + max_attempts = 3 + poll_seconds = 60 poll_interval = 5 - spend_after = spend_before for attempt in range(1, max_attempts + 1): - try: - model.generate_content("hi") - except Exception as exc: - if _is_vertex_quota_error(exc): - pytest.skip("Vertex AI quota exhausted") - raise - - for _ in range(per_attempt_wait // poll_interval): + response = requests.post( + VERTEX_GENERATE_CONTENT_URL, + headers={ + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + }, + json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, + timeout=60, + ) + if _is_vertex_quota_error(response): + pytest.skip("Vertex AI quota exhausted") + assert ( + response.status_code == 200 + ), f"vertex pass-through call failed: {response.status_code} {response.text}" + + call_id = response.headers.get("x-litellm-call-id") + assert call_id, "proxy response missing x-litellm-call-id header" + + for _ in range(poll_seconds // poll_interval): await asyncio.sleep(poll_interval) - spend_after = await get_tracked_spend() - if spend_after > spend_before: - print(f"spend tracked on attempt {attempt}: {spend_before} -> {spend_after}") + row = _spend_log_for_request(call_id) + if row is not None and float(row.get("spend") or 0) > 0: + assert "gemini" in row["model"], f"unexpected model in spend log: {row}" + assert ( + row["custom_llm_provider"] == "vertex_ai" + ), f"unexpected provider in spend log: {row}" return - print(f"attempt {attempt}: spend not tracked yet (spend_after={spend_after}), re-billing") + print(f"attempt {attempt}: spend log for call {call_id} not found yet, re-billing") pytest.fail( - "Vertex pass-through spend never recorded after {} billed calls. spend_before: {}, spend_after: {}".format( - max_attempts, spend_before, spend_after - ) + f"Vertex pass-through spend never recorded after {max_attempts} billed calls" ) From 65e29fc9af5d3f95eb4c88a29b7d79c98bda836d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 30 Jun 2026 20:17:04 +0000 Subject: [PATCH 3/3] test(pass-through): reuse LITE_LLM_ENDPOINT and drop needless async in get_tracked_spend --- tests/pass_through_tests/test_vertex_ai.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/pass_through_tests/test_vertex_ai.py b/tests/pass_through_tests/test_vertex_ai.py index f72ae1dfc14..35cb5f49c56 100644 --- a/tests/pass_through_tests/test_vertex_ai.py +++ b/tests/pass_through_tests/test_vertex_ai.py @@ -58,10 +58,12 @@ def load_vertex_ai_credentials(): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) +LITE_LLM_ENDPOINT = "http://localhost:4000" + SPEND_LOG_API_KEY = "best-api-key-ever" -async def get_tracked_spend() -> float: +def get_tracked_spend() -> float: """ Total spend recorded under the pass-through key in the global spend view. @@ -69,7 +71,7 @@ async def get_tracked_spend() -> float: "today" so a UTC date rollover mid-test can't hide a freshly billed call, and treats an unreachable endpoint as "nothing recorded yet" (0.0). """ - url = f"http://0.0.0.0:4000/global/spend/logs?api_key={SPEND_LOG_API_KEY}" + url = f"{LITE_LLM_ENDPOINT}/global/spend/logs?api_key={SPEND_LOG_API_KEY}" response = requests.get(url, headers={"Authorization": "Bearer sk-1234"}) if response.status_code != 200: print(f"global spend logs endpoint returned {response.status_code}: {response.text}") @@ -80,8 +82,6 @@ async def get_tracked_spend() -> float: return sum(float(row.get("spend") or 0.0) for row in rows) -LITE_LLM_ENDPOINT = "http://localhost:4000" - VERTEX_PROJECT = "litellm-ci-cd" VERTEX_MODEL = "gemini-3.1-flash-lite" VERTEX_GENERATE_CONTENT_URL = ( @@ -173,7 +173,7 @@ async def test_basic_vertex_ai_pass_through_with_spendlog(): @pytest.mark.skip(reason="skip flaky test - vertex pass through streaming is flaky") async def test_basic_vertex_ai_pass_through_streaming_with_spendlog(): - spend_before = await get_tracked_spend() + spend_before = get_tracked_spend() print("spend_before", spend_before) load_vertex_ai_credentials() @@ -193,7 +193,7 @@ async def test_basic_vertex_ai_pass_through_streaming_with_spendlog(): print("response", response) await asyncio.sleep(20) - spend_after = await get_tracked_spend() + spend_after = get_tracked_spend() print("spend_after", spend_after) assert ( spend_after > spend_before