From 288ccb39c019b758b4a059627919f7f1023b6a2b Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 9 Apr 2026 18:20:05 -0700 Subject: [PATCH] resolved greptile comments --- .../spend_management_endpoints.py | 27 ++++++---- .../test_spend_management_endpoints.py | 52 ++++++++++++++----- 2 files changed, 58 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index ef1865a29a9..cd2ccda936e 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -13,11 +13,10 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy._types import ProviderBudgetResponse, ProviderBudgetResponseObject from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import ( - _is_user_team_admin, - _team_member_has_permission, - _user_has_admin_view, -) + +# NOTE: Avoid module-level import from common_utils: proxy_server imports this +# module while common_utils may pull proxy_server during init, which can leave +# those names undefined. Import the helpers locally where they are used. from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team_and_customer, ) @@ -3442,6 +3441,8 @@ 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. """ + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + try: return _user_has_admin_view(user_api_key_dict=user_api_key_dict) except Exception: @@ -3458,6 +3459,11 @@ async def _can_team_member_view_log( Returns True if the team exists and the user is either a team admin or a team member with the ``/spend/logs`` permission. """ + from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_admin, + _team_member_has_permission, + ) + if team_id is None: return False team_row = await prisma_client.db.litellm_teamtable.find_unique( @@ -3509,7 +3515,7 @@ async def _assert_user_can_view_request_id( if row is None: return - if row.user == user_api_key_dict.user_id: + if row.user is not None and row.user == user_api_key_dict.user_id: return if row.team_id: @@ -3539,7 +3545,12 @@ async def _get_permitted_team_ids_for_spend_logs( Return team IDs where the user is either a team admin or has the ``/spend/logs`` permission, allowing them to view team-wide spend logs. """ + # Imported here to avoid circular import: proxy_server imports this module. from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_admin, + _team_member_has_permission, + ) from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache user_obj = await get_user_object( @@ -3559,9 +3570,7 @@ async def _get_permitted_team_ids_for_spend_logs( permitted: List[str] = [] for team_row in team_rows: team_obj = LiteLLM_TeamTable(**team_row.model_dump()) - if _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): permitted.append(team_obj.team_id) elif _team_member_has_permission( user_api_key_dict=user_api_key_dict, 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 f65d4008e0b..01171bc65ae 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 @@ -6,6 +6,7 @@ import sys from datetime import timezone import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient sys.path.insert( @@ -97,15 +98,14 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No from litellm.proxy._types import ( - LiteLLM_TeamTable, - LiteLLM_UserTable, LitellmUserRoles, Member, SpendLogsPayload, UserAPIKeyAuth, ) from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger -from litellm.proxy.proxy_server import app, prisma_client +from litellm.proxy.management_endpoints import common_utils +from litellm.proxy.proxy_server import app from litellm.proxy.spend_tracking import spend_management_endpoints from litellm.router import Router from litellm.types.utils import BudgetConfig @@ -115,7 +115,7 @@ from litellm.types.utils import BudgetConfig async def test_is_admin_view_safe_true(monkeypatch): # Force underlying check to return True monkeypatch.setattr( - spend_management_endpoints, + common_utils, "_user_has_admin_view", lambda user_api_key_dict: True, ) @@ -127,7 +127,7 @@ async def test_is_admin_view_safe_true(monkeypatch): async def test_is_admin_view_safe_false(monkeypatch): # Force underlying check to return False monkeypatch.setattr( - spend_management_endpoints, + common_utils, "_user_has_admin_view", lambda user_api_key_dict: False, ) @@ -141,7 +141,7 @@ async def test_is_admin_view_safe_exception(monkeypatch): def raise_err(*args, **kwargs): raise RuntimeError("boom") - monkeypatch.setattr(spend_management_endpoints, "_user_has_admin_view", raise_err) + monkeypatch.setattr(common_utils, "_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 @@ -187,7 +187,7 @@ async def test_can_team_member_view_log_team_not_found(monkeypatch): prisma = MockPrisma() # Even if admin check would return True, no team means False monkeypatch.setattr( - spend_management_endpoints, + common_utils, "_is_user_team_admin", lambda user_api_key_dict, team_obj: True, ) @@ -227,7 +227,7 @@ async def test_can_team_member_view_log_not_admin(monkeypatch): prisma = MockPrisma() monkeypatch.setattr( - spend_management_endpoints, + common_utils, "_is_user_team_admin", lambda user_api_key_dict, team_obj: False, ) @@ -295,6 +295,37 @@ def test_can_user_view_spend_log_false_for_other_roles(): assert spend_management_endpoints._can_user_view_spend_log(auth) is False +@pytest.mark.asyncio +async def test_assert_user_can_view_request_id_rejects_both_users_none(): + """ + API keys with user_id=None must not be treated as owning a log whose user + field is None (avoid None == None bypass). + """ + + class MockRow: + user = None + team_id = None + + class MockSpendLogs: + async def find_unique(self, where, include=None): + return MockRow() + + class MockDB: + def __init__(self): + self.litellm_spendlogs = MockSpendLogs() + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=None) + with pytest.raises(HTTPException) as exc_info: + await spend_management_endpoints._assert_user_can_view_request_id( + MockPrisma(), auth, "req-none-user" + ) + assert exc_info.value.status_code == 403 + + ignored_keys = [ "request_id", "session_id", @@ -1668,9 +1699,6 @@ class TestSpendLogsPayload: } ) - print(f"payload: {payload}") - print(f"expected_payload: {expected_payload}") - differences = _compare_nested_dicts( payload, expected_payload, ignore_keys=ignored_keys ) @@ -2090,7 +2118,7 @@ async def test_provider_budget_over(disable_budget_sync): ) with pytest.raises(Exception) as e: - response = await router.acompletion( + await router.acompletion( model="azure-gpt-4o", messages=[{"role": "user", "content": "Hello, world!"}], )