From ba2d4d080fda760d616293ab4a56848994824f0e Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 23 May 2025 23:23:46 -0700 Subject: [PATCH] feat(handle_jwt.py): map user to team when added via jwt auth (#11108) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(handle_jwt.py): map user to team when added via jwt auth makes it easy to ensure user belongs to team * test: test_openai_image_edit_litellm_sdk * use n 4 for mapped tests (#11109) * Fix/background health check (#10887) * fix: improve health check logic by deep copying model list on each iteration * test: add async test for background health check reflecting model list changes * fix: validate health check interval before executing background health check * fix: specify type for health check results dictionary * fix(user_api_key_auth.py): handle user custom auth set with no custom settings * bump: version 0.1.21 → 0.2.0 * ci(config.yml): run enterprise and litellm tests separately * fix: fix linting error * docs: add missing docs * [Feat] Add content policy violation error mapping for image editd (#11113) * feat: add image edit mapping for content policy violations * test fix * Expose `/list` and `/info` endpoints for Audit Log events (#11102) * feat(audit_logging_endpoints.py): expose list endpoint to show all audit logs make it easier for user to retrieve individual endpoints * feat(enterprise/): add audit logging endpoint * feat(audit_logging_endpoints.py): expose new GET `/audit/{id}` endpoint make it easier to retrieve view individual audit logs * feat(key_management_event_hooks.py): correctly show the key of the user who initiated the change * fix(key_management_event_hooks.py): add key rotations as an audit log event ' * test(test_audit_logging_endpoints.py): add simple unit testing for audit log endpoint * fix: testing fixes * fix: fix ruff check * [Feat] Use aiohttp transport by default - 97% lower median latency (#11097) * fix: add flag for disabling use_aiohttp_transport * feat: add _create_async_transport * feat: fixes for transport * add httpx-aiohttp * feat: fixes for transport * refactor: fixes for transport * build: fix deps * fixes: test fixes * fix: ensure aiohttp does not auto set content type * test: test fixes * feat: add LiteLLMAiohttpTransport * fix: fixes for responses API handling * test: fixes for responses API handling * test: fixes for responses API handling * feat: fixes for transport * fix: base embedding handler * test: test_async_http_handler_force_ipv4 * test: fix failing deepeval test * fix: add YARL for bedrock urls * fix: issues with transport * fix: comment out linting issues * test fix * test: XAI is unstable * test: fixes for using respx * test: XAI fixes * test: XAI fixes * test: infinity testing fixes * docs(config_settings.md): document param * test: test_openai_image_edit_litellm_sdk * test: remove deprecated test * bump respx==0.22.0 * test: test_xai_message_name_filtering * test: fix anthropic test after bumping httpx * use n 4 for mapped tests (#11109) * fix: use 1 session per event loop * test: test_client_session_helper * fix: linting error * fix: resolving GET requests on httpx 0.28.1 * test fixes proxy unit tests * fix: add ssl verify settings * fix: proxy unit tests * fix: refactor * tests: basic unit tests for aiohttp transports * tests: fixes xai --------- Co-authored-by: Krrish Dholakia * test: cleanup redundant test --------- Co-authored-by: Ishaan Jaff Co-authored-by: JuHyun Bae --- litellm/proxy/auth/handle_jwt.py | 48 ++++++++++++++ .../internal_user_endpoints.py | 3 - .../management_endpoints/team_endpoints.py | 6 +- litellm/proxy/management_endpoints/ui_sso.py | 65 +++++++++---------- tests/litellm/proxy/auth/test_handle_jwt.py | 63 ++++++++++++++++++ tests/proxy_unit_tests/test_proxy_server.py | 6 -- 6 files changed, 145 insertions(+), 46 deletions(-) create mode 100644 tests/litellm/proxy/auth/test_handle_jwt.py 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":