From e600f02d2dc05d4298dce9f4ee4f43bd4a55198b Mon Sep 17 00:00:00 2001 From: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Date: Wed, 6 May 2026 22:13:01 +0000 Subject: [PATCH] Handle empty internal tag usage scopes safely Co-authored-by: ishaan-berri --- .../tag_management_endpoints.py | 14 ++- .../test_tag_management_endpoints.py | 92 ++++++++++++++++++- 2 files changed, 97 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 643c5f37105..9b8163e598e 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -59,21 +59,23 @@ async def _get_internal_user_api_keys( if not _is_internal_user_role(user_api_key_dict): return [] + user_api_keys = set() + if user_api_key_dict.api_key: + user_api_keys.add(user_api_key_dict.api_key) + user_id = user_api_key_dict.user_id if user_id is None: - return [] + return sorted(user_api_keys) key_records = await prisma_client.db.litellm_verificationtoken.find_many( where={"user_id": user_id}, select={"token": True}, ) - user_api_keys = { + user_api_keys.update( key_record.token for key_record in key_records if getattr(key_record, "token", None) - } - if user_api_key_dict.api_key: - user_api_keys.add(user_api_key_dict.api_key) + ) return sorted(user_api_keys) @@ -645,6 +647,8 @@ async def get_tag_daily_activity( user_api_key_dict=user_api_key_dict, requested_api_key=api_key, ) + if scoped_api_key_filter == []: + return SpendAnalyticsPaginatedResponse(results=[]) return await get_daily_activity( prisma_client=prisma_client, diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index a4e075069c0..802ae91e9a5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -547,8 +547,6 @@ async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter() mock_db.litellm_verificationtoken.find_many = AsyncMock( return_value=[owned_key_record] ) - mock_get_daily_activity.return_value = "empty-daily-activity-response" - result = await get_tag_daily_activity( start_date="2025-01-01", end_date="2025-01-31", @@ -556,9 +554,95 @@ async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter() user_api_key_dict=mock_user_auth, ) - assert result == "empty-daily-activity-response" + assert result.results == [] + assert result.metadata.total_spend == 0 + assert result.metadata.total_api_requests == 0 + mock_get_daily_activity.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_internal_user_tag_daily_activity_scopes_to_current_key_without_user_id(): + """ + If an internal-user token has no user_id, it should still scope tag usage to + the current request key instead of falling back to proxy-wide tag spend. + """ + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.management_endpoints.tag_management_endpoints import ( + get_tag_daily_activity, + ) + + mock_user_auth = UserAPIKeyAuth( + api_key="current-owned-key", + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.management_endpoints.tag_management_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity, + ): + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_get_daily_activity.return_value = "daily-activity-response" + + result = await get_tag_daily_activity( + start_date="2025-01-01", + end_date="2025-01-31", + user_api_key_dict=mock_user_auth, + ) + + assert result == "daily-activity-response" + mock_db.litellm_verificationtoken.find_many.assert_not_awaited() mock_get_daily_activity.assert_awaited_once() - assert mock_get_daily_activity.await_args.kwargs["api_key"] == [] + assert mock_get_daily_activity.await_args.kwargs["api_key"] == [ + "current-owned-key" + ] + + +@pytest.mark.asyncio +async def test_internal_user_tag_daily_activity_without_any_scoped_keys_returns_empty(): + """ + If an internal-user token has neither user_id nor api_key, the endpoint must + return an empty response instead of dropping the API key filter. + """ + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.management_endpoints.tag_management_endpoints import ( + get_tag_daily_activity, + ) + + mock_user_auth = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.management_endpoints.tag_management_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity, + ): + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + result = await get_tag_daily_activity( + start_date="2025-01-01", + end_date="2025-01-31", + user_api_key_dict=mock_user_auth, + ) + + assert result.results == [] + assert result.metadata.total_spend == 0 + assert result.metadata.total_api_requests == 0 + mock_db.litellm_verificationtoken.find_many.assert_not_awaited() + mock_get_daily_activity.assert_not_awaited() @pytest.mark.asyncio