From 9c2f380a99491cf05d77148dc466930a1c4c0be9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 13 Jan 2026 17:01:51 -0800 Subject: [PATCH] only show own internal user usage --- .../common_daily_activity.py | 9 +- .../management_endpoints/team_endpoints.py | 34 +- .../test_team_endpoints.py | 363 ++++++++++++++++++ 3 files changed, 401 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index f52abf86b9..c52491efc7 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -343,7 +343,7 @@ def _build_where_conditions( start_date: str, end_date: str, model: Optional[str], - api_key: Optional[str], + api_key: Optional[Union[str, List[str]]], exclude_entity_ids: Optional[List[str]] = None, ) -> Dict[str, Any]: """Build prisma where clause for daily activity queries.""" @@ -357,7 +357,10 @@ def _build_where_conditions( if model: where_conditions["model"] = model if api_key: - where_conditions["api_key"] = api_key + if isinstance(api_key, list): + where_conditions["api_key"] = {"in": api_key} + else: + where_conditions["api_key"] = api_key if entity_id is not None: if isinstance(entity_id, list): @@ -445,7 +448,7 @@ async def get_daily_activity( start_date: Optional[str], end_date: Optional[str], model: Optional[str], - api_key: Optional[str], + api_key: Optional[Union[str, List[str]]], page: int, page_size: int, exclude_entity_ids: Optional[List[str]] = None, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 78caa86db7..d1549b5116 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3601,7 +3601,7 @@ async def get_team_daily_activity( }, ) - ## Fetch team aliases + ## Fetch team aliases and check team admin status where_condition = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} @@ -3612,6 +3612,36 @@ async def get_team_daily_activity( t.team_id: {"team_alias": t.team_alias} for t in team_aliases } + # Check if user is team admin for any requested teams + # If not, filter by user's API keys + user_api_keys: Optional[List[str]] = None + if not _user_has_admin_view(user_api_key_dict) and team_ids_list and team_aliases: + # Check if user is team admin for any of the teams + is_team_admin_for_any = False + for team_alias in team_aliases: + team_obj = LiteLLM_TeamTable(**team_alias.model_dump()) + if _is_user_team_admin( + user_api_key_dict=user_api_key_dict, team_obj=team_obj + ): + is_team_admin_for_any = True + break + + # If user is not a team admin for any team, filter by their API keys + if not is_team_admin_for_any: + # Get all API keys for this user + user_keys = await prisma_client.db.litellm_verificationtoken.find_many( + where={"user_id": user_api_key_dict.user_id} + ) + user_api_keys = [key.token for key in user_keys if key.token] + # If user has no API keys, return empty result + if not user_api_keys: + user_api_keys = [""] # Use empty string to ensure no matches + + # If api_key parameter is provided, use it; otherwise use user_api_keys if set + final_api_key_filter: Optional[Union[str, List[str]]] = api_key + if final_api_key_filter is None and user_api_keys is not None: + final_api_key_filter = user_api_keys + return await get_daily_activity( prisma_client=prisma_client, table_name="litellm_dailyteamspend", @@ -3622,7 +3652,7 @@ async def get_team_daily_activity( start_date=start_date, end_date=end_date, model=model, - api_key=api_key, + api_key=final_api_key_filter, page=page, page_size=page_size, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index e296066b99..bbff7448e1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( LiteLLM_OrganizationTable, LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, Member, ProxyErrorTypes, @@ -4476,6 +4477,187 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): assert deserialized_settings == router_settings_data +@pytest.mark.asyncio +async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( + mock_db_client, +): + """ + Test that non-team-admin users only see their own spend (filtered by their API keys) + when calling /team/daily/activity endpoint. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + get_team_daily_activity, + ) + + # Create a non-admin user + user_id = "test_user_123" + team_id = "test_team_456" + user_api_key_dict = UserAPIKeyAuth( + user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER + ) + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="test@example.com", + user_role="internal_user", + ) + + # Mock team with user as non-admin member + mock_team_member = Member(user_id=user_id, role="user") + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [mock_team_member] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "user"}], + } + + # Mock user's API keys + user_api_key_1 = MagicMock() + user_api_key_1.token = "user_key_1" + user_api_key_2 = MagicMock() + user_api_key_2.token = "user_key_2" + + # Setup mocks + mock_db_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[mock_team] + ) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[user_api_key_1, user_api_key_2] + ) + + # Mock get_user_object + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + # Mock get_daily_activity to capture the api_key parameter + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity: + mock_get_daily_activity.return_value = MagicMock() + + # Call the endpoint + await get_team_daily_activity( + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-02", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=user_api_key_dict, + ) + + # Verify get_daily_activity was called with user's API keys as filter + mock_get_daily_activity.assert_called_once() + call_kwargs = mock_get_daily_activity.call_args[1] + assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"] + assert call_kwargs["entity_id"] == [team_id] + + # Verify user's API keys were fetched + mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() + api_key_call_kwargs = ( + mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] + ) + assert api_key_call_kwargs["where"] == {"user_id": user_id} + + +@pytest.mark.asyncio +async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client): + """ + Test that team admin users see all team spend (no API key filtering) + when calling /team/daily/activity endpoint. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + get_team_daily_activity, + ) + + # Create a team admin user + user_id = "test_admin_123" + team_id = "test_team_456" + user_api_key_dict = UserAPIKeyAuth( + user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER + ) + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="admin@example.com", + user_role="internal_user", + ) + + # Mock team with user as admin member + mock_team_member = Member(user_id=user_id, role="admin") + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [mock_team_member] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "admin"}], + } + + # Setup mocks + mock_db_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[mock_team] + ) + + # Mock get_user_object + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + # Mock get_daily_activity to capture the api_key parameter + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity: + mock_get_daily_activity.return_value = MagicMock() + + # Call the endpoint + await get_team_daily_activity( + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-02", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=user_api_key_dict, + ) + + # Verify get_daily_activity was called WITHOUT API key filtering + mock_get_daily_activity.assert_called_once() + call_kwargs = mock_get_daily_activity.call_args[1] + assert call_kwargs["api_key"] is None + assert call_kwargs["entity_id"] == [team_id] + + # Verify user's API keys were NOT fetched (since they're admin) + if hasattr( + mock_db_client.db.litellm_verificationtoken, "find_many" + ) and mock_db_client.db.litellm_verificationtoken.find_many.called: + # If it was called, that's unexpected for admin users + assert False, "API keys should not be fetched for team admin users" + + @pytest.mark.asyncio async def test_update_team_with_router_settings(mock_db_client, mock_admin_auth): """ @@ -4552,3 +4734,184 @@ async def test_update_team_with_router_settings(mock_db_client, mock_admin_auth) # Verify router_settings can be deserialized and matches input deserialized_settings = json.loads(team_data["router_settings"]) assert deserialized_settings == router_settings_data + + +@pytest.mark.asyncio +async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( + mock_db_client, +): + """ + Test that non-team-admin users only see their own spend (filtered by their API keys) + when calling /team/daily/activity endpoint. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + get_team_daily_activity, + ) + + # Create a non-admin user + user_id = "test_user_123" + team_id = "test_team_456" + user_api_key_dict = UserAPIKeyAuth( + user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER + ) + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="test@example.com", + user_role="internal_user", + ) + + # Mock team with user as non-admin member + mock_team_member = Member(user_id=user_id, role="user") + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [mock_team_member] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "user"}], + } + + # Mock user's API keys + user_api_key_1 = MagicMock() + user_api_key_1.token = "user_key_1" + user_api_key_2 = MagicMock() + user_api_key_2.token = "user_key_2" + + # Setup mocks + mock_db_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[mock_team] + ) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[user_api_key_1, user_api_key_2] + ) + + # Mock get_user_object + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + # Mock get_daily_activity to capture the api_key parameter + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity: + mock_get_daily_activity.return_value = MagicMock() + + # Call the endpoint + await get_team_daily_activity( + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-02", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=user_api_key_dict, + ) + + # Verify get_daily_activity was called with user's API keys as filter + mock_get_daily_activity.assert_called_once() + call_kwargs = mock_get_daily_activity.call_args[1] + assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"] + assert call_kwargs["entity_id"] == [team_id] + + # Verify user's API keys were fetched + mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() + api_key_call_kwargs = ( + mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] + ) + assert api_key_call_kwargs["where"] == {"user_id": user_id} + + +@pytest.mark.asyncio +async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client): + """ + Test that team admin users see all team spend (no API key filtering) + when calling /team/daily/activity endpoint. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + get_team_daily_activity, + ) + + # Create a team admin user + user_id = "test_admin_123" + team_id = "test_team_456" + user_api_key_dict = UserAPIKeyAuth( + user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER + ) + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="admin@example.com", + user_role="internal_user", + ) + + # Mock team with user as admin member + mock_team_member = Member(user_id=user_id, role="admin") + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [mock_team_member] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "admin"}], + } + + # Setup mocks + mock_db_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[mock_team] + ) + + # Mock get_user_object + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + # Mock get_daily_activity to capture the api_key parameter + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity: + mock_get_daily_activity.return_value = MagicMock() + + # Call the endpoint + await get_team_daily_activity( + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-02", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=user_api_key_dict, + ) + + # Verify get_daily_activity was called WITHOUT API key filtering + mock_get_daily_activity.assert_called_once() + call_kwargs = mock_get_daily_activity.call_args[1] + assert call_kwargs["api_key"] is None + assert call_kwargs["entity_id"] == [team_id] + + # Verify user's API keys were NOT fetched (since they're admin) + if hasattr( + mock_db_client.db.litellm_verificationtoken, "find_many" + ) and mock_db_client.db.litellm_verificationtoken.find_many.called: + # If it was called, that's unexpected for admin users + assert False, "API keys should not be fetched for team admin users"