From 8bf491c9391ca517052c3371bb032c95de14ba77 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 12 Nov 2025 18:44:21 -0800 Subject: [PATCH] [Fix] /spend/logs/ui Access Control (#16446) * RBAC for /spend/logs/ui * Addressing comments --- .../spend_management_endpoints.py | 68 ++++ .../test_spend_management_endpoints.py | 291 ++++++++++++++++++ 2 files changed, 359 insertions(+) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 51ece5dd84a..b60acd14281 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -18,6 +18,10 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( ) from litellm.proxy.utils import handle_exception_on_proxy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_admin, + _user_has_admin_view, +) if TYPE_CHECKING: from litellm.proxy.proxy_server import PrismaClient @@ -1749,6 +1753,28 @@ async def ui_view_spend_logs( # noqa: PLR0915 where_conditions["spend"]["gte"] = min_spend if max_spend is not None: where_conditions["spend"]["lte"] = max_spend + is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) + if not is_admin_view: + if team_id is not None: + can_view_team = await _can_team_member_view_log( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + team_id=team_id, + ) + if not can_view_team: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Not authorized to view team spend for team_id={}".format( + team_id + ) + }, + ) + where_conditions["team_id"] = team_id + else: + if _can_user_view_spend_log(user_api_key_dict=user_api_key_dict): + where_conditions["user"] = user_api_key_dict.user_id + where_conditions.pop("team_id", None) # Calculate skip value for pagination skip = (page - 1) * page_size @@ -2990,3 +3016,45 @@ def _build_status_filter_condition(status_filter: Optional[str]) -> Dict[str, An return {"OR": [{"status": {"equals": "success"}}, {"status": None}]} else: return {"status": {"equals": status_filter}} + + +def _is_admin_view_safe(user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Safely determine if the current user has admin view permissions. + Wraps the underlying check and defaults to False on any exception. + """ + try: + return _user_has_admin_view(user_api_key_dict=user_api_key_dict) + except Exception: + return False + + +async def _can_team_member_view_log( + prisma_client, + user_api_key_dict: UserAPIKeyAuth, + team_id: Optional[str], +) -> bool: + """ + Check if the requesting user can view spend logs for the given team. + Returns True only if the team exists and the user is a team admin. + """ + if team_id is None: + return False + team_obj = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + if team_obj is None: + return False + return _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj) + + +def _can_user_view_spend_log(user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Check if the requesting user can view their own spend logs. + """ + user_role = user_api_key_dict.user_role + user_id = user_api_key_dict.user_id + return user_role in ( + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ) and user_id is not None diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index c9c602e1cc2..3a9229b3c45 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -21,6 +21,169 @@ from litellm.proxy.proxy_server import app, prisma_client from litellm.proxy.spend_tracking import spend_management_endpoints from litellm.router import Router from litellm.types.utils import BudgetConfig +from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, Member +from litellm.proxy.spend_tracking import spend_management_endpoints +import litellm.proxy.proxy_server as ps + +@pytest.mark.asyncio +async def test_is_admin_view_safe_true(monkeypatch): + # Force underlying check to return True + monkeypatch.setattr( + spend_management_endpoints, "_user_has_admin_view", lambda user_api_key_dict: True + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user") + assert spend_management_endpoints._is_admin_view_safe(auth) is True + + +@pytest.mark.asyncio +async def test_is_admin_view_safe_false(monkeypatch): + # Force underlying check to return False + monkeypatch.setattr( + spend_management_endpoints, "_user_has_admin_view", lambda user_api_key_dict: False + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + assert spend_management_endpoints._is_admin_view_safe(auth) is False + + +@pytest.mark.asyncio +async def test_is_admin_view_safe_exception(monkeypatch): + # Ensure exceptions are swallowed and return False + def raise_err(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(spend_management_endpoints, "_user_has_admin_view", raise_err) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + assert spend_management_endpoints._is_admin_view_safe(auth) is False + + +@pytest.mark.asyncio +async def test_can_team_member_view_log_none_team_id(): + # team_id=None should immediately return False + class MockPrisma: + class DB: + class TeamTable: + async def find_unique(self, where: dict): + return None + + def __init__(self): + self.litellm_teamtable = self.TeamTable() + + def __init__(self): + self.db = self.DB() + + prisma = MockPrisma() + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + allowed = await spend_management_endpoints._can_team_member_view_log( + prisma, auth, None + ) + assert allowed is False + + +@pytest.mark.asyncio +async def test_can_team_member_view_log_team_not_found(monkeypatch): + # Non-existent team should return False + class MockPrisma: + class DB: + class TeamTable: + async def find_unique(self, where: dict): + return None + + def __init__(self): + self.litellm_teamtable = self.TeamTable() + + def __init__(self): + self.db = self.DB() + + prisma = MockPrisma() + # Even if admin check would return True, no team means False + monkeypatch.setattr( + spend_management_endpoints, "_is_user_team_admin", lambda user_api_key_dict, team_obj: True + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + allowed = await spend_management_endpoints._can_team_member_view_log( + prisma, auth, "team_x" + ) + assert allowed is False + + +@pytest.mark.asyncio +async def test_can_team_member_view_log_not_admin(monkeypatch): + # Existing team but caller is not a team admin -> False + class MockTeam: + pass + + class MockPrisma: + class DB: + class TeamTable: + async def find_unique(self, where: dict): + return MockTeam() + + def __init__(self): + self.litellm_teamtable = self.TeamTable() + + def __init__(self): + self.db = self.DB() + + prisma = MockPrisma() + monkeypatch.setattr( + spend_management_endpoints, "_is_user_team_admin", lambda user_api_key_dict, team_obj: False + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + allowed = await spend_management_endpoints._can_team_member_view_log( + prisma, auth, "team_x" + ) + assert allowed is False + + +@pytest.mark.asyncio +async def test_can_team_member_view_log_admin(monkeypatch): + # Existing team and caller is team admin -> True + class MockTeam: + pass + + class MockPrisma: + class DB: + class TeamTable: + async def find_unique(self, where: dict): + return MockTeam() + + def __init__(self): + self.litellm_teamtable = self.TeamTable() + + def __init__(self): + self.db = self.DB() + + prisma = MockPrisma() + monkeypatch.setattr( + spend_management_endpoints, "_is_user_team_admin", lambda user_api_key_dict, team_obj: True + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + allowed = await spend_management_endpoints._can_team_member_view_log( + prisma, auth, "team_x" + ) + assert allowed is True + + +def test_can_user_view_spend_log_true_for_internal_user(): + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="u1") + assert spend_management_endpoints._can_user_view_spend_log(auth) is True + + +def test_can_user_view_spend_log_true_for_internal_view_only(): + auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, user_id="u1" + ) + assert spend_management_endpoints._can_user_view_spend_log(auth) is True + + +def test_can_user_view_spend_log_false_without_user_id(): + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=None) + assert spend_management_endpoints._can_user_view_spend_log(auth) is False + + +def test_can_user_view_spend_log_false_for_other_roles(): + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + assert spend_management_endpoints._can_user_view_spend_log(auth) is False ignored_keys = [ "request_id", @@ -255,6 +418,134 @@ async def test_ui_view_spend_logs_with_team_id(client, monkeypatch): assert data["data"][0]["team_id"] == "team1" +@pytest.mark.asyncio +async def test_ui_view_spend_logs_internal_user_scoped_without_user_id(client, monkeypatch): + """ + Internal users should only be able to view their own spend even if user_id is not provided. + """ + # Mock spend logs for 2 users + mock_spend_logs = [ + {"id": "log1", "request_id": "req1", "api_key": "sk-test-key", "user": "internal_user_1", "team_id": "team1", "spend": 0.05, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-3.5-turbo"}, + {"id": "log2", "request_id": "req2", "api_key": "sk-test-key", "user": "internal_user_2", "team_id": "team1", "spend": 0.10, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4"}, + ] + + # Prisma client mock that filters by "user" where condition + class MockDB: + async def find_many(self, *args, **kwargs): + where = kwargs.get("where", {}) + if "user" in where and where["user"] == "internal_user_1": + return [mock_spend_logs[0]] + return mock_spend_logs + + async def count(self, *args, **kwargs): + where = kwargs.get("where", {}) + if "user" in where and where["user"] == "internal_user_1": + return 1 + return len(mock_spend_logs) + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + self.db.litellm_spendlogs = self.db + + mock_prisma_client = MockPrismaClient() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Override auth dependency to return INTERNAL_USER with specific user_id + # Override using the function reference attached to the running app module + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1" + ) + + try: + start_date = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)).strftime("%Y-%m-%d %H:%M:%S") + end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + + # No user_id provided; should auto-scope to authenticated internal user's own id + response = client.get( + "/spend/logs/ui", + params={"start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["user"] == "internal_user_1" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeypatch): + """ + Team admins should be able to view team-wide spend when team_id is provided. + """ + # Mock spend logs for two teams + mock_spend_logs = [ + {"id": "log1", "request_id": "req1", "api_key": "sk-test-key", "user": "member1", "team_id": "team_admin_team", "spend": 0.05, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-3.5-turbo"}, + {"id": "log2", "request_id": "req2", "api_key": "sk-test-key", "user": "member2", "team_id": "team_other", "spend": 0.10, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4"}, + ] + + class MockDB: + async def find_many(self, *args, **kwargs): + where = kwargs.get("where", {}) + if "team_id" in where and where["team_id"] == "team_admin_team": + return [mock_spend_logs[0]] + return mock_spend_logs + + async def count(self, *args, **kwargs): + where = kwargs.get("where", {}) + if "team_id" in where and where["team_id"] == "team_admin_team": + return 1 + return len(mock_spend_logs) + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + self.db.litellm_spendlogs = self.db + # Team lookup for RBAC check + class TeamTable: + def __init__(self): + # user "admin_user" is team admin + self.members_with_roles = [Member(user_id="admin_user", role="admin")] + + async def find_unique(where: dict): + if where == {"team_id": "team_admin_team"}: + return TeamTable() + return None + + self.db.litellm_teamtable = self + self.litellm_teamtable = self + self.find_unique = find_unique + + mock_prisma_client = MockPrismaClient() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Override auth dependency to return INTERNAL_USER (who is a team admin via team.members_with_roles) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="admin_user" + ) + + try: + start_date = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)).strftime("%Y-%m-%d %H:%M:%S") + end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + + response = client.get( + "/spend/logs/ui", + params={"team_id": "team_admin_team", "start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["team_id"] == "team_admin_team" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + @pytest.mark.asyncio async def test_ui_view_spend_logs_pagination(client, monkeypatch): # Create a larger set of mock data for pagination testing