diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 783c2f1553..80f05c7b0c 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -30,8 +30,11 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, + Member, ScopeMapping, Span, + TeamMemberAddRequest, + UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -855,6 +858,45 @@ class JWTAuthManager: ) return True + @staticmethod + async def map_user_to_teams( + user_object: Optional[LiteLLM_UserTable], + team_object: Optional[LiteLLM_TeamTable], + ): + """ + Map user to teams. + - If user is not in team, add them to the team + - If user is in team, do nothing + """ + from litellm.proxy.management_endpoints.team_endpoints import team_member_add + + if not user_object: + return None + + if not team_object: + return None + + # check if user is in team + for member in team_object.members_with_roles: + if member.user_id and member.user_id == user_object.user_id: + return None + + data = TeamMemberAddRequest( + member=Member( + user_id=user_object.user_id, + role="user", # [TODO]: allow controlling role within team based on jwt token + ), + team_id=team_object.team_id, + ) + # add user to team + await team_member_add( + data=data, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ), # [TODO]: expose an internal service role, for better tracking + ) + return None + @staticmethod async def auth_builder( api_key: str, @@ -976,6 +1018,12 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, ) + ## MAP USER TO TEAMS + await JWTAuthManager.map_user_to_teams( + user_object=user_object, + team_object=team_object, + ) + # Validate that a valid rbac id is returned for spend tracking JWTAuthManager.validate_object_id( user_id=user_id, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index bbdd8c8133..aeddf8c1e3 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -209,9 +209,6 @@ async def new_user( user_email=data_json.get("user_email", None), ), ), - http_request=Request( - scope={"type": "http", "path": "/user/new"}, - ), user_api_key_dict=user_api_key_dict, ) except HTTPException as e: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 60f0590252..56f6a0042d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -772,7 +772,6 @@ def team_member_add_duplication_check( @management_endpoint_wrapper async def team_member_add( data: TeamMemberAddRequest, - http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1110,7 +1109,7 @@ async def team_member_update( Update team member budgets and team member role """ - from litellm.proxy.proxy_server import prisma_client, premium_user + from litellm.proxy.proxy_server import premium_user, prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -1118,12 +1117,11 @@ async def team_member_update( if data.team_id is None: raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) - if data.role == "admin" and not premium_user: # exactly the same text your proxy throws for add: raise HTTPException( status_code=400, - detail="Assigning team admins is a premium feature. You must be a LiteLLM Enterprise user to use this feature. If you have a license please set `LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/#trial. Pricing: https://www.litellm.ai/#pricing" + detail="Assigning team admins is a premium feature. You must be a LiteLLM Enterprise user to use this feature. If you have a license please set `LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/#trial. Pricing: https://www.litellm.ai/#pricing", ) if data.user_id is None and data.user_email is None: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 163678ab85..a85c43a2a6 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -278,7 +278,6 @@ async def create_team_member_add_task(team_id, user_info): return await team_member_add( data=team_member_add_request, user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - http_request=Request(scope={"type": "http", "path": "/sso/callback"}), ) except Exception as e: verbose_proxy_logger.debug( @@ -356,28 +355,28 @@ async def get_user_info_from_db( user_defined_values: Optional[SSOUserDefinedValues], ) -> Optional[Union[LiteLLM_UserTable, NewUserResponse]]: try: - user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]] = ( - await get_existing_user_info_from_db( - user_id=cast( - Optional[str], - ( - getattr(result, "id", None) - if not isinstance(result, dict) - else result.get("id", None) - ), + user_info: Optional[ + Union[LiteLLM_UserTable, NewUserResponse] + ] = await get_existing_user_info_from_db( + user_id=cast( + Optional[str], + ( + getattr(result, "id", None) + if not isinstance(result, dict) + else result.get("id", None) ), - user_email=cast( - Optional[str], - ( - getattr(result, "email", None) - if not isinstance(result, dict) - else result.get("email", None) - ), + ), + user_email=cast( + Optional[str], + ( + getattr(result, "email", None) + if not isinstance(result, dict) + else result.get("email", None) ), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) verbose_proxy_logger.debug( f"user_info: {user_info}; litellm.default_internal_user_params: {litellm.default_internal_user_params}" @@ -707,9 +706,9 @@ async def insert_sso_user( if user_defined_values.get("max_budget") is None: user_defined_values["max_budget"] = litellm.max_internal_user_budget if user_defined_values.get("budget_duration") is None: - user_defined_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + user_defined_values[ + "budget_duration" + ] = litellm.internal_user_budget_duration if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -904,9 +903,9 @@ class SSOAuthenticationHandler: if state: redirect_params["state"] = state elif "okta" in generic_authorization_endpoint: - redirect_params["state"] = ( - uuid.uuid4().hex - ) # set state param for okta - required + redirect_params[ + "state" + ] = uuid.uuid4().hex # set state param for okta - required return await generic_sso.get_login_redirect(**redirect_params) # type: ignore raise ValueError( "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" @@ -1151,9 +1150,9 @@ class MicrosoftSSOHandler: # if user is trying to get the raw sso response for debugging, return the raw sso response if return_raw_sso_response: - original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = ( - user_team_ids - ) + original_msft_result[ + MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY + ] = user_team_ids return original_msft_result or {} result = MicrosoftSSOHandler.openid_from_response( @@ -1221,9 +1220,9 @@ class MicrosoftSSOHandler: # Fetch user membership from Microsoft Graph API all_group_ids = [] - next_link: Optional[str] = ( - MicrosoftSSOHandler.graph_api_user_groups_endpoint - ) + next_link: Optional[ + str + ] = MicrosoftSSOHandler.graph_api_user_groups_endpoint auth_headers = {"Authorization": f"Bearer {access_token}"} page_count = 0 diff --git a/tests/litellm/proxy/auth/test_handle_jwt.py b/tests/litellm/proxy/auth/test_handle_jwt.py new file mode 100644 index 0000000000..f2e33a5780 --- /dev/null +++ b/tests/litellm/proxy/auth/test_handle_jwt.py @@ -0,0 +1,63 @@ +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, Member +from litellm.proxy.auth.handle_jwt import JWTAuthManager + + +@pytest.mark.asyncio +async def test_map_user_to_teams_user_already_in_team(): + """Test that no action is taken when user is already in team""" + # Setup test data + user = LiteLLM_UserTable(user_id="test_user_1") + team = LiteLLM_TeamTable( + team_id="test_team_1", + members_with_roles=[Member(user_id="test_user_1", role="user")], + ) + + # Mock team_member_add to ensure it's not called + with patch( + "litellm.proxy.management_endpoints.team_endpoints.team_member_add", + new_callable=AsyncMock, + ) as mock_add: + await JWTAuthManager.map_user_to_teams(user_object=user, team_object=team) + mock_add.assert_not_called() + + +@pytest.mark.asyncio +async def test_map_user_to_teams_add_new_user(): + """Test that new user is added to team""" + # Setup test data + user = LiteLLM_UserTable(user_id="test_user_1") + team = LiteLLM_TeamTable(team_id="test_team_1", members_with_roles=[]) + + # Mock team_member_add + with patch( + "litellm.proxy.management_endpoints.team_endpoints.team_member_add", + new_callable=AsyncMock, + ) as mock_add: + await JWTAuthManager.map_user_to_teams(user_object=user, team_object=team) + mock_add.assert_called_once() + # Verify the correct data was passed to team_member_add + call_args = mock_add.call_args[1]["data"] + assert call_args.member.user_id == "test_user_1" + assert call_args.member.role == "user" + assert call_args.team_id == "test_team_1" + + +@pytest.mark.asyncio +async def test_map_user_to_teams_null_inputs(): + """Test that method handles null inputs gracefully""" + # Test with null user + await JWTAuthManager.map_user_to_teams( + user_object=None, team_object=LiteLLM_TeamTable(team_id="test_team_1") + ) + + # Test with null team + await JWTAuthManager.map_user_to_teams( + user_object=LiteLLM_UserTable(user_id="test_user_1"), team_object=None + ) + + # Test with both null + await JWTAuthManager.map_user_to_teams(user_object=None, team_object=None) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 4b7d86b8ca..b503272077 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1040,9 +1040,6 @@ async def test_create_team_member_add(prisma_client, new_member_method): await team_member_add( data=team_member_add_request, user_api_key_dict=UserAPIKeyAuth(user_role="proxy_admin"), - http_request=Request( - scope={"type": "http", "path": "/user/new"}, - ), ) mock_client.assert_called() @@ -1225,9 +1222,6 @@ async def test_create_team_member_add_team_admin( await team_member_add( data=team_member_add_request, user_api_key_dict=valid_token, - http_request=Request( - scope={"type": "http", "path": "/user/new"}, - ), ) except HTTPException as e: if user_role == "user":