From fea343d849a574cffdfee45ba22be0437024e88a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 17 Aug 2026 11:53:39 +0000 Subject: [PATCH 1/5] fix(proxy): update VerificationToken.spend from SpendLogs api_key SpendLogs can persist api_key from standard_logging_object when the update_database token argument is missing. Key spend now uses that same identifier so LiteLLM_VerificationToken.spend tracks logged spend. Fixes #37144 Co-authored-by: Yuzhong Zhang --- litellm/proxy/db/db_spend_update_writer.py | 29 ++- .../proxy/hooks/proxy_track_cost_callback.py | 4 +- .../proxy/db/test_db_spend_update_writer.py | 208 ++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 52 +++++ 4 files changed, 290 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 0c8c9a853ec..988176b1667 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -39,6 +39,7 @@ from litellm.proxy._types import ( SpendLogsPayload, SpendUpdateQueueItem, ToolDiscoveryQueueItem, + hash_token, ) from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, @@ -166,7 +167,7 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name, prisma_client, ) - from litellm.proxy.utils import ProxyUpdateSpend, hash_token + from litellm.proxy.utils import ProxyUpdateSpend try: verbose_proxy_logger.debug( @@ -195,6 +196,10 @@ class DBSpendUpdateWriter: end_time=end_time, ) payload["spend"] = response_cost or 0.0 + hashed_token = self._resolve_hashed_token_for_key_spend( + hashed_token=hashed_token, + payload_api_key=payload.get("api_key"), + ) if isinstance(payload["startTime"], datetime): payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): @@ -583,6 +588,26 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + @staticmethod + def _resolve_hashed_token_for_key_spend( + hashed_token: str | None, + payload_api_key: str | None, + ) -> str | None: + """ + Use the same key identifier SpendLogs persist when `token` is missing. + + SpendLogs can resolve api_key from standard_logging_object.metadata.user_api_key_hash + even when the update_database `token` argument is None. Key spend must use that + same identifier so LiteLLM_VerificationToken.spend tracks logged spend. + """ + if hashed_token: + return hashed_token + if not payload_api_key: + return hashed_token + if isinstance(payload_api_key, str) and payload_api_key.startswith("sk-"): + return hash_token(payload_api_key) + return payload_api_key + async def _update_key_db( self, response_cost: float | None, @@ -590,7 +615,7 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient | None, ): try: - if hashed_token is None or prisma_client is None: + if not hashed_token or prisma_client is None: return await self.spend_update_queue.add_update( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6abfca1d3a0..ddc1e170106 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -257,7 +257,9 @@ class _ProxyDBLogger(CustomLogger): ) if response_cost is not None: - user_api_key: Final = metadata.get("user_api_key", None) + user_api_key: Final = metadata.get("user_api_key", None) or ( + (sl_object.get("metadata") or {}).get("user_api_key_hash") if sl_object is not None else None + ) if kwargs.get("cache_hit", False) is True: response_cost = 0.0 verbose_proxy_logger.debug("Cache Hit: response_cost %s, for user_id %s", response_cost, user_id) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index ca1827aa38e..9fc86fa21da 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2740,3 +2740,211 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}] assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush PrismaClient.spend_log_flush_requested.clear() + +def _mock_prisma_for_key_spend_commit(): + mock_batcher = MagicMock() + mock_batcher.litellm_verificationtoken = MagicMock() + mock_batcher.litellm_verificationtoken.update_many = MagicMock() + mock_batcher.litellm_usertable = MagicMock() + mock_batcher.litellm_usertable.update_many = MagicMock() + mock_batcher.litellm_teamtable = MagicMock() + mock_batcher.litellm_teamtable.update_many = MagicMock() + mock_batcher.litellm_teammembership = MagicMock() + mock_batcher.litellm_teammembership.update_many = MagicMock() + mock_batcher.litellm_organizationtable = MagicMock() + mock_batcher.litellm_organizationtable.update_many = MagicMock() + mock_batcher.litellm_tagtable = MagicMock() + mock_batcher.litellm_tagtable.update_many = MagicMock() + mock_batcher.litellm_agentstable = MagicMock() + mock_batcher.litellm_agentstable.update_many = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + mock_prisma_client._spend_log_transactions_lock = AsyncMock() + mock_prisma_client.spend_log_transactions = [] + return mock_prisma_client, mock_batcher + + +@pytest.mark.asyncio +async def test_logged_spend_updates_verification_token_spend_when_token_provided(): + """ + Successful logged spend must increment LiteLLM_VerificationToken.spend + for the hashed key written to SpendLogs. + """ + db_writer = DBSpendUpdateWriter() + hashed_token = "a" * 64 + response_cost = 0.05 + mock_prisma_client, mock_batcher = _mock_prisma_for_key_spend_commit() + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + with patch("litellm.proxy.proxy_server.disable_spend_logs", False), patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ), patch( + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch( + "litellm.proxy.proxy_server.master_key", None + ): + await db_writer.update_database( + token=hashed_token, + user_id=None, + end_user_id=None, + start_time=datetime.now(), + end_time=datetime.now(), + team_id=None, + org_id=None, + completion_response=MagicMock(), + response_cost=response_cost, + kwargs={ + "model": "gpt-4", + "custom_llm_provider": "openai", + "litellm_params": {"metadata": {"user_api_key": hashed_token}}, + }, + ) + await asyncio.sleep(0) + + aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert aggregated["key_list_transactions"][hashed_token] == response_cost + + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=aggregated, + ) + + mock_batcher.litellm_verificationtoken.update_many.assert_called() + call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] + assert call_kwargs["where"] == {"token": hashed_token} + assert call_kwargs["data"]["spend"] == {"increment": response_cost} + + +@pytest.mark.asyncio +async def test_logged_spend_updates_verification_token_spend_from_spend_log_api_key(): + """ + Fixes #37144: SpendLogs can persist api_key from standard_logging_object + even when update_database is called with token=None (metadata.user_api_key + missing, but user/team ids still trigger spend tracking). + + LiteLLM_VerificationToken.spend must increment for that same api_key. + """ + db_writer = DBSpendUpdateWriter() + hashed_token = "b" * 64 + response_cost = 0.123 + mock_prisma_client, mock_batcher = _mock_prisma_for_key_spend_commit() + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + kwargs = { + "model": "custom_openai/test-model", + "custom_llm_provider": "openai", + "litellm_params": { + "metadata": { + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + } + }, + "standard_logging_object": { + "response_cost": response_cost, + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + "model_map_information": None, + "hidden_params": {}, + "metadata": {"user_api_key_hash": hashed_token}, + }, + } + + with patch("litellm.proxy.proxy_server.disable_spend_logs", False), patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ), patch( + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch( + "litellm.proxy.proxy_server.master_key", None + ): + await db_writer.update_database( + token=None, + user_id="user-1", + end_user_id=None, + start_time=datetime.now(), + end_time=datetime.now(), + team_id="team-1", + org_id=None, + completion_response=MagicMock(), + response_cost=response_cost, + kwargs=kwargs, + ) + await asyncio.sleep(0) + + spend_log_api_key = mock_prisma_client.spend_log_transactions[0]["api_key"] + assert spend_log_api_key == hashed_token + + aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert hashed_token in aggregated["key_list_transactions"], ( + "VerificationToken.spend was not queued for the SpendLogs api_key" + ) + assert aggregated["key_list_transactions"][hashed_token] == response_cost + + key_only_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": aggregated["key_list_transactions"], + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=key_only_transactions, + ) + + mock_batcher.litellm_verificationtoken.update_many.assert_called() + call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] + assert call_kwargs["where"] == {"token": hashed_token} + assert call_kwargs["data"]["spend"] == {"increment": response_cost} + + +def test_resolve_hashed_token_for_key_spend_uses_payload_when_token_missing(): + hashed = "c" * 64 + assert ( + DBSpendUpdateWriter._resolve_hashed_token_for_key_spend( + hashed_token=None, payload_api_key=hashed + ) + == hashed + ) + assert ( + DBSpendUpdateWriter._resolve_hashed_token_for_key_spend( + hashed_token="", payload_api_key=hashed + ) + == hashed + ) + assert ( + DBSpendUpdateWriter._resolve_hashed_token_for_key_spend( + hashed_token="existing-token", payload_api_key=hashed + ) + == "existing-token" + ) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index ca517474a5c..9e398b256f2 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1875,3 +1875,55 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( 1 if expect_spend_log else 0 ) + + +@pytest.mark.asyncio +async def test_track_cost_callback_uses_standard_logging_key_hash_when_metadata_missing(): + """ + When metadata.user_api_key is missing but standard_logging_object has + user_api_key_hash, the cost callback must still pass that hash as token + so LiteLLM_VerificationToken.spend is updated (issue #37144). + """ + logger = _ProxyDBLogger() + hashed_token = "d" * 64 + kwargs = { + "model": "custom_openai/test-model", + "litellm_params": { + "metadata": { + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + } + }, + "standard_logging_object": { + "response_cost": 0.05, + "metadata": {"user_api_key_hash": hashed_token}, + }, + "stream": False, + } + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, patch( + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment, patch( + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ): + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + call_kwargs = ( + mock_proxy_logging.db_spend_update_writer.update_database.call_args[1] + ) + assert call_kwargs["token"] == hashed_token + mock_increment.assert_awaited_once() + assert mock_increment.call_args[1]["token"] == hashed_token From 83cbcbb51b70f711812c588fb1af3cfc7194fd62 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 20 Aug 2026 04:02:09 +0000 Subject: [PATCH 2/5] fix(proxy): avoid LIT002 empty-dict fallback for key hash Resolve user_api_key_hash from standard_logging metadata without building a throwaway {} so the type-discipline LIT002 total stays within budget. Co-authored-by: Yuzhong Zhang --- litellm/proxy/hooks/proxy_track_cost_callback.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ddc1e170106..88750fef5de 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -257,8 +257,9 @@ class _ProxyDBLogger(CustomLogger): ) if response_cost is not None: + sl_metadata = sl_object.get("metadata") if sl_object is not None else None user_api_key: Final = metadata.get("user_api_key", None) or ( - (sl_object.get("metadata") or {}).get("user_api_key_hash") if sl_object is not None else None + sl_metadata.get("user_api_key_hash") if sl_metadata else None ) if kwargs.get("cache_hit", False) is True: response_cost = 0.0 From 40047b1497a12c8ea177256a2a7231a9fa8e17ba Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 20 Aug 2026 05:13:18 +0000 Subject: [PATCH 3/5] fix(proxy): drop unnecessary isinstance on SpendLogs api_key payload_api_key is already str | None; after the empty check, startswith is enough and basedpyright no longer flags reportUnnecessaryIsInstance. Co-authored-by: Yuzhong Zhang --- litellm/proxy/db/db_spend_update_writer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 988176b1667..84374d652c1 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -604,7 +604,7 @@ class DBSpendUpdateWriter: return hashed_token if not payload_api_key: return hashed_token - if isinstance(payload_api_key, str) and payload_api_key.startswith("sk-"): + if payload_api_key.startswith("sk-"): return hash_token(payload_api_key) return payload_api_key From fd2f0e4c5206cf99dc3de64a23dd46f347e58206 Mon Sep 17 00:00:00 2001 From: Yuzhong Zhang Date: Mon, 24 Aug 2026 23:46:47 +0000 Subject: [PATCH 4/5] fix(proxy): annotate LIT010 sites on verification-token spend path --- litellm/proxy/db/db_spend_update_writer.py | 2 +- litellm/proxy/hooks/proxy_track_cost_callback.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 84374d652c1..23dfa2d9d2d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -196,7 +196,7 @@ class DBSpendUpdateWriter: end_time=end_time, ) payload["spend"] = response_cost or 0.0 - hashed_token = self._resolve_hashed_token_for_key_spend( + hashed_token = self._resolve_hashed_token_for_key_spend( # rebind-ok: SpendLogs api_key fallback hashed_token=hashed_token, payload_api_key=payload.get("api_key"), ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 88750fef5de..aeab38dde56 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -257,7 +257,7 @@ class _ProxyDBLogger(CustomLogger): ) if response_cost is not None: - sl_metadata = sl_object.get("metadata") if sl_object is not None else None + sl_metadata: Final = sl_object.get("metadata") if sl_object is not None else None user_api_key: Final = metadata.get("user_api_key", None) or ( sl_metadata.get("user_api_key_hash") if sl_metadata else None ) From 1825b4176822f477386bbd3f764ec3826026a513 Mon Sep 17 00:00:00 2001 From: Yuzhong Zhang Date: Mon, 24 Aug 2026 23:56:47 +0000 Subject: [PATCH 5/5] test(proxy): avoid TQ008 patch() on verification-token spend tests --- .../proxy/db/test_db_spend_update_writer.py | 169 ++++++++---------- .../hooks/test_proxy_track_cost_callback.py | 48 +++-- 2 files changed, 101 insertions(+), 116 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 9fc86fa21da..46555829a19 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2775,8 +2775,19 @@ def _mock_prisma_for_key_spend_commit(): return mock_prisma_client, mock_batcher +def _stub_proxy_server_for_key_spend(monkeypatch, prisma_client): + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.disable_spend_logs", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + + @pytest.mark.asyncio -async def test_logged_spend_updates_verification_token_spend_when_token_provided(): +async def test_logged_spend_updates_verification_token_spend_when_token_provided(monkeypatch): """ Successful logged spend must increment LiteLLM_VerificationToken.spend for the hashed key written to SpendLogs. @@ -2785,48 +2796,36 @@ async def test_logged_spend_updates_verification_token_spend_when_token_provided hashed_token = "a" * 64 response_cost = 0.05 mock_prisma_client, mock_batcher = _mock_prisma_for_key_spend_commit() - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) + _stub_proxy_server_for_key_spend(monkeypatch, mock_prisma_client) - with patch("litellm.proxy.proxy_server.disable_spend_logs", False), patch( - "litellm.proxy.proxy_server.prisma_client", mock_prisma_client - ), patch( - "litellm.proxy.proxy_server.user_api_key_cache", mock_cache - ), patch( - "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" - ), patch( - "litellm.proxy.proxy_server.general_settings", {} - ), patch( - "litellm.proxy.proxy_server.master_key", None - ): - await db_writer.update_database( - token=hashed_token, - user_id=None, - end_user_id=None, - start_time=datetime.now(), - end_time=datetime.now(), - team_id=None, - org_id=None, - completion_response=MagicMock(), - response_cost=response_cost, - kwargs={ - "model": "gpt-4", - "custom_llm_provider": "openai", - "litellm_params": {"metadata": {"user_api_key": hashed_token}}, - }, - ) - await asyncio.sleep(0) + await db_writer.update_database( + token=hashed_token, + user_id=None, + end_user_id=None, + start_time=datetime.now(), + end_time=datetime.now(), + team_id=None, + org_id=None, + completion_response=MagicMock(), + response_cost=response_cost, + kwargs={ + "model": "gpt-4", + "custom_llm_provider": "openai", + "litellm_params": {"metadata": {"user_api_key": hashed_token}}, + }, + ) + await asyncio.sleep(0) - aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() - assert aggregated["key_list_transactions"][hashed_token] == response_cost + aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert aggregated["key_list_transactions"][hashed_token] == response_cost - with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): - await db_writer._commit_spend_updates_to_db( - prisma_client=mock_prisma_client, - n_retry_times=0, - proxy_logging_obj=MagicMock(), - db_spend_update_transactions=aggregated, - ) + monkeypatch.setattr("litellm.proxy.utils._raise_failed_update_spend_exception", MagicMock()) + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=aggregated, + ) mock_batcher.litellm_verificationtoken.update_many.assert_called() call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] @@ -2835,7 +2834,7 @@ async def test_logged_spend_updates_verification_token_spend_when_token_provided @pytest.mark.asyncio -async def test_logged_spend_updates_verification_token_spend_from_spend_log_api_key(): +async def test_logged_spend_updates_verification_token_spend_from_spend_log_api_key(monkeypatch): """ Fixes #37144: SpendLogs can persist api_key from standard_logging_object even when update_database is called with token=None (metadata.user_api_key @@ -2847,8 +2846,7 @@ async def test_logged_spend_updates_verification_token_spend_from_spend_log_api_ hashed_token = "b" * 64 response_cost = 0.123 mock_prisma_client, mock_batcher = _mock_prisma_for_key_spend_commit() - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) + _stub_proxy_server_for_key_spend(monkeypatch, mock_prisma_client) kwargs = { "model": "custom_openai/test-model", @@ -2870,57 +2868,46 @@ async def test_logged_spend_updates_verification_token_spend_from_spend_log_api_ }, } - with patch("litellm.proxy.proxy_server.disable_spend_logs", False), patch( - "litellm.proxy.proxy_server.prisma_client", mock_prisma_client - ), patch( - "litellm.proxy.proxy_server.user_api_key_cache", mock_cache - ), patch( - "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" - ), patch( - "litellm.proxy.proxy_server.general_settings", {} - ), patch( - "litellm.proxy.proxy_server.master_key", None - ): - await db_writer.update_database( - token=None, - user_id="user-1", - end_user_id=None, - start_time=datetime.now(), - end_time=datetime.now(), - team_id="team-1", - org_id=None, - completion_response=MagicMock(), - response_cost=response_cost, - kwargs=kwargs, - ) - await asyncio.sleep(0) + await db_writer.update_database( + token=None, + user_id="user-1", + end_user_id=None, + start_time=datetime.now(), + end_time=datetime.now(), + team_id="team-1", + org_id=None, + completion_response=MagicMock(), + response_cost=response_cost, + kwargs=kwargs, + ) + await asyncio.sleep(0) - spend_log_api_key = mock_prisma_client.spend_log_transactions[0]["api_key"] - assert spend_log_api_key == hashed_token + spend_log_api_key = mock_prisma_client.spend_log_transactions[0]["api_key"] + assert spend_log_api_key == hashed_token - aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() - assert hashed_token in aggregated["key_list_transactions"], ( - "VerificationToken.spend was not queued for the SpendLogs api_key" - ) - assert aggregated["key_list_transactions"][hashed_token] == response_cost + aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert hashed_token in aggregated["key_list_transactions"], ( + "VerificationToken.spend was not queued for the SpendLogs api_key" + ) + assert aggregated["key_list_transactions"][hashed_token] == response_cost - key_only_transactions = { - "user_list_transactions": {}, - "end_user_list_transactions": {}, - "key_list_transactions": aggregated["key_list_transactions"], - "team_list_transactions": {}, - "team_member_list_transactions": {}, - "org_list_transactions": {}, - "tag_list_transactions": {}, - "agent_list_transactions": {}, - } - with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): - await db_writer._commit_spend_updates_to_db( - prisma_client=mock_prisma_client, - n_retry_times=0, - proxy_logging_obj=MagicMock(), - db_spend_update_transactions=key_only_transactions, - ) + key_only_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": aggregated["key_list_transactions"], + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + monkeypatch.setattr("litellm.proxy.utils._raise_failed_update_spend_exception", MagicMock()) + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=key_only_transactions, + ) mock_batcher.litellm_verificationtoken.update_many.assert_called() call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 9e398b256f2..8a69a51e2c3 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1878,7 +1878,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( @pytest.mark.asyncio -async def test_track_cost_callback_uses_standard_logging_key_hash_when_metadata_missing(): +async def test_track_cost_callback_uses_standard_logging_key_hash_when_metadata_missing(monkeypatch): """ When metadata.user_api_key is missing but standard_logging_object has user_api_key_hash, the cost callback must still pass that hash as token @@ -1901,29 +1901,27 @@ async def test_track_cost_callback_uses_standard_logging_key_hash_when_metadata_ "stream": False, } - with patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ) as mock_proxy_logging, patch( - "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock - ) as mock_increment, patch( - "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock - ): - mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() - mock_proxy_logging.slack_alerting_instance = MagicMock() - mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_increment = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.increment_spend_counters", mock_increment) + monkeypatch.setattr("litellm.proxy.proxy_server.update_cache", AsyncMock()) - await logger._PROXY_track_cost_callback( - kwargs=kwargs, - completion_response=MagicMock(), - start_time=datetime.now(), - end_time=datetime.now(), - ) + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + ) - mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() - call_kwargs = ( - mock_proxy_logging.db_spend_update_writer.update_database.call_args[1] - ) - assert call_kwargs["token"] == hashed_token - mock_increment.assert_awaited_once() - assert mock_increment.call_args[1]["token"] == hashed_token + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + call_kwargs = ( + mock_proxy_logging.db_spend_update_writer.update_database.call_args[1] + ) + assert call_kwargs["token"] == hashed_token + mock_increment.assert_awaited_once() + assert mock_increment.call_args[1]["token"] == hashed_token