feat(handle_jwt.py): map user to team when added via jwt auth (#11108)

* 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 <krrishdholakia@gmail.com> * test: cleanup redundant test --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: JuHyun Bae <jhyun0408@nate.com>
2026-08-05 02:23:42 +00:00 · 2025-05-23 23:23:46 -07:00
parent a52943089b
commit ba2d4d080f
6 changed files with 145 additions and 46 deletions
@@ -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,
@@ -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:
@@ -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(
@@ -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

@@ -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)
@@ -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":