From 63579f1e355dabb4df2b31f6c654d1d6f7b0ad44 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:12:57 +0000 Subject: [PATCH] fix(spend-tracking): keep batch spend keys joinable after v1.99 provenance gate Batch cost attribution and the legacy queue endpoint already store the VerificationToken hash in user_api_key, but omitted user_api_key_hash. Since v1.99 the spend-log writer re-hashes any key without that provenance flag, so DailyUserSpend.api_key no longer joins VerificationToken and Usage shows key-hash-... rows with null api_key_alias / user_email. Co-authored-by: Mateo Wang --- .../proxy/common_utils/check_batch_cost.py | 1 + litellm/proxy/proxy_server.py | 1 + .../proxy_unit_tests/test_check_batch_cost.py | 43 +++++++++++++++++++ .../test_spend_tracking_utils.py | 36 ++++++++++++++++ 4 files changed, 81 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 354a6ed2fd0..e6f00877a26 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -161,6 +161,7 @@ class CheckBatchCost: metadata: dict[str, object] = { "user_api_key_user_id": job.created_by, "user_api_key": api_key, + "user_api_key_hash": api_key, "user_api_key_team_id": team_id, **(await self._get_user_info(batch_id, job.created_by)), } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..222e79f7ca7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15197,6 +15197,7 @@ async def async_queue_request( # extra_body); see above for the same guard upstream. data["metadata"] = {} data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_hash"] = user_api_key_dict.api_key data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) _headers: Final = _safe_get_request_headers(request).copy() _headers.pop("authorization", None) # do not store the original `sk-..` api key in the db diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff5e8f89d64..9a6ab08e9b6 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2445,6 +2445,7 @@ class TestBatchCostAttribution: metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") assert metadata["user_api_key"] == "hash-alice" + assert metadata["user_api_key_hash"] == "hash-alice" assert metadata["user_api_key_user_id"] == "alice" assert metadata["user_api_key_team_id"] == "team-alpha" assert metadata["user_api_key_alias"] == "prod-key" @@ -2553,6 +2554,48 @@ class TestBatchCostAttribution: assert metadata["user_api_key_alias"] == "prod-key" + @pytest.mark.asyncio + async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self): + """ + CheckBatchCost stores the VerificationToken hash on the managed object. The + spend-log writer must receive matching user_api_key_hash provenance so it + does not re-hash that value; otherwise DailyUserSpend.api_key no longer joins + VerificationToken and Usage shows key-hash-... with a null alias/email. + """ + from datetime import datetime, timezone + from types import SimpleNamespace + + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.proxy.utils import hash_token + + token_hash = hash_token("sk-batch-creator-key") + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key"), + user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None), + ) + metadata = await instance._build_creator_attribution_metadata( + self._job(api_key=token_hash), "batch-1" + ) + + assert metadata["user_api_key"] == token_hash + assert metadata["user_api_key_hash"] == token_hash + + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o", + "call_type": "aretrieve_batch", + "litellm_params": {"metadata": metadata}, + }, + response_obj={ + "id": "batch_123", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + ) + assert payload["api_key"] == token_hash + assert payload["api_key"] != hash_token(token_hash) + class TestPollPageStarvation: """LIT-5462 regression: a row that can never be costed used to keep its slot in the diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9e5917637a8..bea1f8e2d6c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2755,6 +2755,42 @@ def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed(): assert meta["user_api_key"] == hash_token(already_hashed) +def test_get_logging_payload_batch_attribution_keeps_verification_token_hash(): + """ + Batch cost rebuilds metadata with the managed object's already-hashed api_key. + That hash must land in SpendLogs.api_key unchanged so Usage/CloudZero can join + LiteLLM_VerificationToken for api_key_alias and user_email. Regression: without + user_api_key_hash provenance, v1.99+ re-hashed the token and broke the join. + """ + token_hash = hash_token("sk-batch-creator-key") + kwargs = { + "model": "gpt-4o", + "call_type": "aretrieve_batch", + "litellm_params": { + "metadata": { + "user_api_key": token_hash, + "user_api_key_hash": token_hash, + "user_api_key_alias": "batch-creator", + "user_api_key_user_id": "alice", + "user_api_key_user_email": "alice@example.com", + "user_api_key_team_id": "team-1", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj={"id": "batch_123", "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == token_hash + assert payload["api_key"] != hash_token(token_hash) + parsed_meta = json.loads(payload["metadata"]) + assert parsed_meta["user_api_key"] == token_hash + assert parsed_meta["user_api_key_alias"] == "batch-creator" + + def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): already_hashed = hash_token("sk-some-key") different_hash = hash_token("sk-other-key")