diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 0c8c9a853ec..23dfa2d9d2d 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( # rebind-ok: SpendLogs api_key fallback + 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 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..aeab38dde56 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -257,7 +257,10 @@ class _ProxyDBLogger(CustomLogger): ) if response_cost is not None: - user_api_key: Final = metadata.get("user_api_key", 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 + ) 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..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 @@ -2740,3 +2740,198 @@ 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 + + +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(monkeypatch): + """ + 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() + _stub_proxy_server_for_key_spend(monkeypatch, mock_prisma_client) + + 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 + + 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] + 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(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 + 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() + _stub_proxy_server_for_key_spend(monkeypatch, mock_prisma_client) + + 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}, + }, + } + + 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": {}, + } + 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] + 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..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 @@ -1875,3 +1875,53 @@ 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(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 + 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, + } + + 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(), + ) + + 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