From 17ed4c4f7a31a6b1d0d64fb91264dc68ecbe3031 Mon Sep 17 00:00:00 2001 From: AxelRay <110029405+axelray-dev@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:24:51 +0800 Subject: [PATCH 1/7] fix(proxy): scope user activity by owned API keys --- .../common_daily_activity.py | 35 +++++++ .../internal_user_endpoints.py | 32 ++++--- .../test_common_daily_activity.py | 92 +++++++++++++++++++ .../test_internal_user_endpoints.py | 68 ++++++++++++-- 4 files changed, 203 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index bf8e7bc15fb..c4db06e8982 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -560,6 +560,41 @@ async def get_api_key_metadata( return await attach_user_emails(prisma_client, combined) +async def _get_deleted_keys_for_user( + prisma_client: PrismaClient, + user_id: str, +) -> Sequence["PrismaDeletedVerificationToken"]: + try: + return await DeletedVerificationTokenRepository(prisma_client).table.find_many( + where={"user_id": user_id}, + order={"deleted_at": "desc"}, + ) + except Exception as e: + verbose_proxy_logger.warning("Failed to fetch deleted key metadata for user %s: %s", user_id, e) + return () + + +async def get_user_api_key_filter( + prisma_client: PrismaClient, + user_id: str, + api_key: str | None, +) -> list[str]: + """Return the key digests that should scope a user's activity query.""" + active_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(where={"user_id": user_id}) + deleted_keys: Final = await _get_deleted_keys_for_user(prisma_client, user_id) + + user_api_keys: Final = list( + dict.fromkeys( + key.token + for key in (*active_keys, *deleted_keys) + if getattr(key, "token", None) and getattr(key, "user_id", None) == user_id + ) + ) + if api_key is None: + return user_api_keys + return [api_key] if api_key in user_api_keys else [] + + def _adjust_dates_for_timezone( start_date: str, end_date: str, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6509977f7ff..da228f8b9c5 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -54,6 +54,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( DailySpendRecord, get_daily_activity, get_daily_activity_aggregated, + get_user_api_key_filter, ) from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, @@ -3025,29 +3026,32 @@ async def get_user_daily_activity_aggregated( try: is_admin: Final = _user_has_admin_view(user_api_key_dict) - if is_admin: - entity_id = user_id # None means global view, otherwise filter by user - else: - caller_user_id: Final = require_caller_user_id_for_non_admin(user_api_key_dict) - if user_id is None: - user_id = caller_user_id - if user_id != caller_user_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "Non-admin users can only view their own spend data."}, - ) - entity_id = user_id + caller_user_id: Final[str | None] = ( + None if is_admin else require_caller_user_id_for_non_admin(user_api_key_dict) + ) + requested_user_id: Final[str | None] = user_id if is_admin or user_id is not None else caller_user_id + if caller_user_id is not None and requested_user_id != caller_user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Non-admin users can only view their own spend data."}, + ) + + api_key_filter: Final = ( + api_key + if requested_user_id is None + else await get_user_api_key_filter(prisma_client, requested_user_id, api_key) + ) return await get_daily_activity_aggregated( prisma_client=prisma_client, table_name="litellm_dailyuserspend", entity_id_field="user_id", - entity_id=entity_id, + entity_id=None, entity_metadata_field=None, start_date=start_date, end_date=end_date, model=model, - api_key=api_key, + api_key=api_key_filter, timezone_offset_minutes=timezone, include_current_utc_day=include_current_utc_day, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index baaf3f4ba2f..7157d871afd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -25,6 +25,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, + get_user_api_key_filter, global_rollup_reconciled_through, update_metrics, ) @@ -331,6 +332,20 @@ async def test_get_api_key_metadata_returns_active_key_metadata(): assert result["active-key-hash-123"]["team_id"] == "team-abc" +@pytest.mark.asyncio +async def test_get_user_api_key_filter_scopes_to_active_and_deleted_user_keys(): + mock_prisma = MagicMock() + + active_key = SimpleNamespace(token="active-key", user_id="target-user") + unrelated_key = SimpleNamespace(token="unrelated-key", user_id="other-user") + deleted_key = SimpleNamespace(token="deleted-key", user_id="target-user") + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[active_key, unrelated_key]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[deleted_key]) + + assert await get_user_api_key_filter(mock_prisma, "target-user", None) == ["active-key", "deleted-key"] + assert await get_user_api_key_filter(mock_prisma, "target-user", "unrelated-key") == [] + + @pytest.mark.asyncio async def test_get_api_key_metadata_falls_back_to_deleted_keys(): """Test that get_api_key_metadata should fall back to deleted keys table for missing keys.""" @@ -1663,6 +1678,83 @@ async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == {"key-1"} +@pytest.mark.asyncio +async def test_get_daily_activity_aggregated_user_key_filter_keeps_usage_with_legacy_user_id( + _aggregated_postgresql: psycopg.Connection, +): + rows: Final = [ + ( + "owned-row", + "legacy-user-id", + "2026-09-22", + "owned-key", + "gpt-5", + "", + "openai", + "/v1/chat/completions", + 10, + 7.0, + 1, + 1, + ), + ( + "deleted-row", + "another-legacy-user-id", + "2026-09-22", + "deleted-key", + "gpt-5", + "", + "openai", + "/v1/chat/completions", + 10, + 3.0, + 1, + 1, + ), + ( + "unrelated-row", + "other-user-id", + "2026-09-22", + "unrelated-key", + "gpt-5", + "", + "openai", + "/v1/chat/completions", + 10, + 100.0, + 1, + 1, + ), + ] + _seed_daily_user_spend(_aggregated_postgresql, rows) + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[ + SimpleNamespace(token="owned-key", key_alias="owned", team_id=None, user_id="target-user"), + SimpleNamespace(token="deleted-key", key_alias="deleted", team_id=None, user_id="target-user"), + ] + ) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + + result = await get_daily_activity_aggregated( + prisma_client=mock_prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2026-09-22", + end_date="2026-09-22", + model=None, + api_key=["owned-key", "deleted-key"], + ) + + assert result.metadata.total_spend == 10.0 + assert set(result.results[0].breakdown.api_keys) == {"owned-key", "deleted-key"} + assert "unrelated-key" not in result.results[0].breakdown.api_keys + + def _prisma_with_marker(marker: str | None) -> MagicMock: prisma = MagicMock() prisma.db = MagicMock() diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index f684e2040bd..9b13cb2a37d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2523,6 +2523,48 @@ async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch, ) +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_filters_by_user_owned_keys(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity_aggregated, + ) + + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_response = MagicMock() + mock_get_user_api_key_filter = AsyncMock(return_value=["owned-key"]) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_user_api_key_filter", + mock_get_user_api_key_filter, + ) + mock_get_daily_agg = AsyncMock(return_value=mock_response) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + result = await get_user_daily_activity_aggregated( + start_date="2026-09-22", + end_date="2026-09-22", + model=None, + api_key=None, + user_id="target-user", + timezone=None, + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) + + assert result is mock_response + mock_get_user_api_key_filter.assert_awaited_once_with(mock_prisma_client, "target-user", None) + assert mock_get_daily_agg.call_args.kwargs["entity_id"] is None + assert mock_get_daily_agg.call_args.kwargs["api_key"] == ["owned-key"] + + @pytest.mark.asyncio async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_users( monkeypatch, @@ -2577,19 +2619,25 @@ async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_us new_callable=AsyncMock, return_value=mock_response, ) as mock_get_daily_agg: - result = await get_user_daily_activity_aggregated( - start_date="2025-01-01", - end_date="2025-01-31", - model=None, - api_key=None, - user_id=None, - timezone=None, - user_api_key_dict=non_admin_key_dict, - ) + with patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_user_api_key_filter", + new_callable=AsyncMock, + return_value=["regular-user-key"], + ): + result = await get_user_daily_activity_aggregated( + start_date="2025-01-01", + end_date="2025-01-31", + model=None, + api_key=None, + user_id=None, + timezone=None, + user_api_key_dict=non_admin_key_dict, + ) assert result is mock_response mock_get_daily_agg.assert_called_once() - assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "regular-user-123" + assert mock_get_daily_agg.call_args.kwargs["entity_id"] is None + assert mock_get_daily_agg.call_args.kwargs["api_key"] == ["regular-user-key"] @pytest.mark.asyncio From 2a14059c11f04819ce5aee3d698de9fda92689b9 Mon Sep 17 00:00:00 2001 From: AxelRay <110029405+axelray-dev@users.noreply.github.com> Date: Wed, 23 Sep 2026 15:52:11 +0800 Subject: [PATCH 2/7] fix(proxy): satisfy type discipline gate --- .../management_endpoints/common_daily_activity.py | 12 +++++++----- .../management_endpoints/internal_user_endpoints.py | 4 +++- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index c4db06e8982..329fef4260d 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -566,8 +566,8 @@ async def _get_deleted_keys_for_user( ) -> Sequence["PrismaDeletedVerificationToken"]: try: return await DeletedVerificationTokenRepository(prisma_client).table.find_many( - where={"user_id": user_id}, - order={"deleted_at": "desc"}, + where={"user_id": user_id}, # mutable-ok: Prisma query payload is consumed as a mutable mapping + order={"deleted_at": "desc"}, # mutable-ok: Prisma query payload is consumed as a mutable mapping ) except Exception as e: verbose_proxy_logger.warning("Failed to fetch deleted key metadata for user %s: %s", user_id, e) @@ -580,10 +580,12 @@ async def get_user_api_key_filter( api_key: str | None, ) -> list[str]: """Return the key digests that should scope a user's activity query.""" - active_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(where={"user_id": user_id}) + active_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many( + where={"user_id": user_id}, # mutable-ok: Prisma query payload is consumed as a mutable mapping + ) deleted_keys: Final = await _get_deleted_keys_for_user(prisma_client, user_id) - user_api_keys: Final = list( + user_api_keys: Final = list( # mutable-ok: helper returns a list for the downstream SQL filter dict.fromkeys( key.token for key in (*active_keys, *deleted_keys) @@ -592,7 +594,7 @@ async def get_user_api_key_filter( ) if api_key is None: return user_api_keys - return [api_key] if api_key in user_api_keys else [] + return [api_key] if api_key in user_api_keys else [] # mutable-ok: preserve the list return shape def _adjust_dates_for_timezone( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index da228f8b9c5..13218b2b8b9 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -3033,7 +3033,9 @@ async def get_user_daily_activity_aggregated( if caller_user_id is not None and requested_user_id != caller_user_id: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "Non-admin users can only view their own spend data."}, + detail={ + "error": "Non-admin users can only view their own spend data." + }, # mutable-ok: FastAPI requires a mutable error payload ) api_key_filter: Final = ( From ed3e209501b7f759f693592340f8a7429408a22e Mon Sep 17 00:00:00 2001 From: AxelRay <110029405+axelray-dev@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:02:36 +0800 Subject: [PATCH 3/7] test(proxy): cover deleted key lookup fallback --- .../test_common_daily_activity.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 7157d871afd..84f9cc202de 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -346,6 +346,18 @@ async def test_get_user_api_key_filter_scopes_to_active_and_deleted_user_keys(): assert await get_user_api_key_filter(mock_prisma, "target-user", "unrelated-key") == [] +@pytest.mark.asyncio +async def test_get_user_api_key_filter_handles_deleted_lookup_failure_and_owned_key(): + mock_prisma = MagicMock() + active_key = SimpleNamespace(token="active-key", user_id="target-user") + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[active_key]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( + side_effect=RuntimeError("deleted-key table unavailable") + ) + + assert await get_user_api_key_filter(mock_prisma, "target-user", "active-key") == ["active-key"] + + @pytest.mark.asyncio async def test_get_api_key_metadata_falls_back_to_deleted_keys(): """Test that get_api_key_metadata should fall back to deleted keys table for missing keys.""" From 003909d325afc09c446a96b2730871c7f3868f33 Mon Sep 17 00:00:00 2001 From: AxelRay <110029405+axelray-dev@users.noreply.github.com> Date: Fri, 25 Sep 2026 04:14:21 +0800 Subject: [PATCH 4/7] fix(activity): fail closed on deleted key lookup errors --- .../management_endpoints/common_daily_activity.py | 12 ++++-------- .../test_common_daily_activity.py | 3 ++- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 329fef4260d..3a808327b8d 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -564,14 +564,10 @@ async def _get_deleted_keys_for_user( prisma_client: PrismaClient, user_id: str, ) -> Sequence["PrismaDeletedVerificationToken"]: - try: - return await DeletedVerificationTokenRepository(prisma_client).table.find_many( - where={"user_id": user_id}, # mutable-ok: Prisma query payload is consumed as a mutable mapping - order={"deleted_at": "desc"}, # mutable-ok: Prisma query payload is consumed as a mutable mapping - ) - except Exception as e: - verbose_proxy_logger.warning("Failed to fetch deleted key metadata for user %s: %s", user_id, e) - return () + return await DeletedVerificationTokenRepository(prisma_client).table.find_many( + where={"user_id": user_id}, # mutable-ok: Prisma query payload is consumed as a mutable mapping + order={"deleted_at": "desc"}, # mutable-ok: Prisma query payload is consumed as a mutable mapping + ) async def get_user_api_key_filter( diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 84f9cc202de..9a61c7ac5dd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -355,7 +355,8 @@ async def test_get_user_api_key_filter_handles_deleted_lookup_failure_and_owned_ side_effect=RuntimeError("deleted-key table unavailable") ) - assert await get_user_api_key_filter(mock_prisma, "target-user", "active-key") == ["active-key"] + with pytest.raises(RuntimeError, match="deleted-key table unavailable"): + await get_user_api_key_filter(mock_prisma, "target-user", "active-key") @pytest.mark.asyncio From 47cee6e62010ab3ecffa9f5169f2e5eebc57be9e Mon Sep 17 00:00:00 2001 From: AxelRay <110029405+axelray-dev@users.noreply.github.com> Date: Sat, 3 Oct 2026 23:30:50 +0800 Subject: [PATCH 5/7] style: format aggregated activity endpoint --- .../proxy/management_endpoints/internal_user_endpoints.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 5fb534490e1..2970a6880ed 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -3007,6 +3007,7 @@ async def get_user_daily_activity( detail={"error": f"Failed to fetch analytics: {e}"}, ) + @router.get( "/user/daily/activity/aggregated", tags=["Budget & Spend Tracking", "Internal User management"], @@ -3055,9 +3056,7 @@ async def get_user_daily_activity_aggregated( ) if isinstance(resolved_entity_ids, ScopeDenied): raise_public(resolved_entity_ids) - requested_user_id: Final[str | None] = ( - resolved_entity_ids[0] if resolved_entity_ids is not None else None - ) + requested_user_id: Final[str | None] = resolved_entity_ids[0] if resolved_entity_ids is not None else None api_key_filter: Final = ( api_key if requested_user_id is None From f580e4842223522b46e9e2243faef9be4ed9fdb9 Mon Sep 17 00:00:00 2001 From: AxelRay <110029405+axelray-dev@users.noreply.github.com> Date: Sat, 3 Oct 2026 23:35:10 +0800 Subject: [PATCH 6/7] lint: document FastAPI dependency injection --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 2970a6880ed..6d10caf2cad 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -3033,7 +3033,7 @@ async def get_user_daily_activity_aggregated( default=False, description="When the range ends on the caller's current local day, include today's UTC bucket.", ), - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection ) -> SpendAnalyticsPaginatedResponse: from litellm.proxy.proxy_server import prisma_client From db8e525c4319cbdc935229637b7e2aede62f487d Mon Sep 17 00:00:00 2001 From: AxelRay <110029405+axelray-dev@users.noreply.github.com> Date: Sun, 4 Oct 2026 00:13:23 +0800 Subject: [PATCH 7/7] test: cover aggregated user activity endpoint --- .../internal_user_endpoints.py | 6 +- .../test_internal_user_endpoints.py | 292 +++++++++++++++--- 2 files changed, 261 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6d10caf2cad..da2e213b871 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -3049,6 +3049,8 @@ async def get_user_daily_activity_aggregated( detail={"error": "Please provide start_date and end_date"}, ) + typed_prisma_client: Final["PrismaClient"] = prisma_client + try: resolved_entity_ids: Final = resolve_user_daily_activity_entity_ids( user_id=user_id, @@ -3060,9 +3062,9 @@ async def get_user_daily_activity_aggregated( api_key_filter: Final = ( api_key if requested_user_id is None - else await get_user_api_key_filter(prisma_client, requested_user_id, api_key) + else await get_user_api_key_filter(typed_prisma_client, requested_user_id, api_key) ) - repository = daily_activity_repository(prisma_client) + repository = daily_activity_repository(typed_prisma_client) scope = daily_activity_scope( table="litellm_dailyuserspend", entity_id_field="user_id", diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 799fe59147c..4a952c4356a 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2143,7 +2143,11 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo "user_alias": "Spruce", } - replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0} + replacement: Final = { + "model-spruce": {"budget_limit": 0, "time_period": "1d"}, + "model-birch": 5.0, + "model-cedar": 0, + } request: Final = UpdateUserRequest( user_id="user-spruce", model_max_budget=replacement, @@ -2185,10 +2189,14 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2224,10 +2232,14 @@ async def test_user_status_update_refreshes_cached_user( prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2261,10 +2273,14 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke prisma_client: Final = mocker.MagicMock() prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2296,7 +2312,9 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke @pytest.mark.asyncio @pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"]) -async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None: +async def test_user_max_budget_update_evicts_cached_user_on_every_worker( + mocker: MockerFixture, all_users: bool +) -> None: from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update @@ -2309,10 +2327,14 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2477,6 +2499,204 @@ async def test_get_user_daily_activity_rejects_service_account_caller(monkeypatc mock_get_daily.assert_not_called() +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_requires_database(monkeypatch): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity_aggregated, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + with pytest.raises(HTTPException) as exc_info: + await get_user_daily_activity_aggregated( + start_date="2026-01-01", + end_date="2026-01-01", + model=None, + api_key=None, + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.status_code == 500 + assert "database" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_requires_dates(monkeypatch): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity_aggregated, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + + with pytest.raises(HTTPException) as exc_info: + await get_user_daily_activity_aggregated( + start_date=None, + end_date="2026-01-01", + model=None, + api_key=None, + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.status_code == 400 + assert "start_date" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_scopes_non_admin(monkeypatch): + import litellm.proxy.management_endpoints.internal_user_endpoints as endpoints + + prisma_client = MagicMock() + repository = object() + expected = SimpleNamespace(result="aggregated") + calls = {} + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + + async def fake_key_filter(client, user_id, api_key): + calls["filter"] = (client, user_id, api_key) + return ["owned-digest"] + + def fake_repository(client): + calls["repository"] = client + return repository + + async def fake_aggregate(repository_arg, scope): + calls["aggregate"] = (repository_arg, scope) + return expected + + monkeypatch.setattr(endpoints, "get_user_api_key_filter", fake_key_filter) + monkeypatch.setattr(endpoints, "daily_activity_repository", fake_repository) + monkeypatch.setattr(endpoints, "get_daily_activity_aggregated", fake_aggregate) + + result = await endpoints.get_user_daily_activity_aggregated( + start_date="2026-01-01", + end_date="2026-01-02", + model="gpt-4o", + api_key="requested-digest", + user_id=None, + timezone=480, + include_current_utc_day=True, + user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert result is expected + assert calls["filter"] == (prisma_client, "caller", "requested-digest") + assert calls["repository"] is prisma_client + assert calls["aggregate"][0] is repository + scope = calls["aggregate"][1] + assert scope.table.value == "litellm_dailyuserspend" + assert scope.entity_ids is None + assert scope.api_keys == ("owned-digest",) + assert scope.start_date == "2026-01-01" + assert scope.end_date == "2026-01-02" + assert scope.model == "gpt-4o" + assert scope.timezone_offset_minutes == 480 + assert scope.include_current_utc_day is True + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_admin_keeps_requested_key(monkeypatch): + import litellm.proxy.management_endpoints.internal_user_endpoints as endpoints + + prisma_client = MagicMock() + repository = object() + expected = SimpleNamespace(result="admin-aggregated") + captured_scope = None + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + + async def unexpected_key_filter(*args): + raise AssertionError("admin global scope must not resolve a user key filter") + + def fake_repository(client): + assert client is prisma_client + return repository + + async def fake_aggregate(repository_arg, scope): + nonlocal captured_scope + assert repository_arg is repository + captured_scope = scope + return expected + + monkeypatch.setattr(endpoints, "get_user_api_key_filter", unexpected_key_filter) + monkeypatch.setattr(endpoints, "daily_activity_repository", fake_repository) + monkeypatch.setattr(endpoints, "get_daily_activity_aggregated", fake_aggregate) + + result = await endpoints.get_user_daily_activity_aggregated( + start_date="2026-01-01", + end_date="2026-01-01", + model=None, + api_key="admin-digest", + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result is expected + assert captured_scope is not None + assert captured_scope.entity_ids is None + assert captured_scope.api_keys == ("admin-digest",) + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_rejects_other_non_admin(monkeypatch): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity_aggregated, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + + with pytest.raises(HTTPException) as exc_info: + await get_user_daily_activity_aggregated( + start_date="2026-01-01", + end_date="2026-01-01", + model=None, + api_key=None, + user_id="other-user", + timezone=None, + include_current_utc_day=False, + user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.status_code == 403 + assert "only view their own" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_maps_failures(monkeypatch): + import litellm.proxy.management_endpoints.internal_user_endpoints as endpoints + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr(endpoints, "daily_activity_repository", lambda client: object()) + + async def fail_aggregate(repository, scope): + raise RuntimeError("aggregate unavailable") + + monkeypatch.setattr(endpoints, "get_daily_activity_aggregated", fail_aggregate) + + with pytest.raises(HTTPException) as exc_info: + await endpoints.get_user_daily_activity_aggregated( + start_date="2026-01-01", + end_date="2026-01-01", + model=None, + api_key=None, + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert exc_info.value.status_code == 500 + assert "aggregate unavailable" in str(exc_info.value.detail) + + @pytest.mark.asyncio async def test_delete_user_cleans_up_created_by_invitation_links(mocker): """ @@ -2511,24 +2731,12 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) # Mock all delete_many calls - mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock( - return_value=[] - ) - mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock( - return_value=1 - ) - mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1) + mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -2619,9 +2827,15 @@ async def test_delete_user_evicts_jwt_key_mapping_cache_of_its_keys(mocker): mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: substitute the database dependency - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache - mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj", None + ) # test-quality-ok: delete_user reads it off proxy_server at call time await delete_user( data=DeleteUserRequest(user_ids=["jwt-user"]), @@ -4530,12 +4744,18 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() for row in (deleted, survivor): await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache - mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj", None + ) # test-quality-ok: delete_user reads it off proxy_server at call time broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -4641,7 +4861,9 @@ async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_ assert raised.value.code == "503" assert raised.value.type == ProxyErrorTypes.no_db_connection assert isinstance(raised.value.__cause__, httpx.ConnectError) - outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()] + outage_logs: Final = [ + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage() + ] assert outage_logs == ["Database unavailable during user search: ConnectError"]