From 1a0acaa33bf3abfec5cc939fb2a7a12d41d31de2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 25 Jul 2026 11:37:11 -0700 Subject: [PATCH] fix(auth): route JWT default-team into memberships instead of the create payload (#33082) * fix(auth): route JWT default-team into memberships instead of the create payload JWT auto-provisioning (get_user_object with user_id_upsert) merged litellm.default_internal_user_params verbatim into the Prisma user create, including a teams key. When a default team is configured through the Admin UI it is stored as a list of NewUserRequestTeam objects, but the user table's teams column is String[], so the create raised a Prisma type error and every JWT-authenticated request 401'd with the user never created. Mirror the /user/new path: strip teams (and available_teams) out of the create payload, then route the configured default team through check_if_default_team_set / add_new_user_to_default_team so provisioned users get real membership rows. Reuse the synthetic PROXY_ADMIN UserAPIKeyAuth pattern already used by the team-upsert path to satisfy the membership permission gate, and import the helpers lazily to avoid the auth_checks <-> internal_user_endpoints import cycle. * fix(auth): propagate max_budget_in_team when adding users to default teams * fix: use pipe union instead of Optional for UP045 budget --- litellm/proxy/auth/auth_checks.py | 25 +++++++-- .../internal_user_endpoints.py | 3 + .../proxy/auth/test_auth_checks.py | 55 +++++++++++++++++++ .../test_internal_user_endpoints.py | 55 +++++++++++++++++++ 4 files changed, 134 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ce82ca74267..07dfdc4fb43 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1671,13 +1671,20 @@ async def get_user_object( if response is None: if user_id_upsert: + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + add_new_user_to_default_team, + check_if_default_team_set, + ) + + default_params = litellm.default_internal_user_params or {} + scalar_default_params = { + key: value for key, value in default_params.items() if key not in ("teams", "available_teams") + } new_user_params: Dict[str, Any] = { "user_id": user_id, + **({"user_email": user_email} if user_email is not None else {}), + **scalar_default_params, } - if user_email is not None: - new_user_params["user_email"] = user_email - if litellm.default_internal_user_params is not None: - new_user_params.update(litellm.default_internal_user_params) if ( new_user_params.get("budget_duration") is not None and new_user_params.get("budget_reset_at") is None @@ -1690,6 +1697,16 @@ async def get_user_object( data=new_user_params, include={"organization_memberships": True}, ) + + default_teams = check_if_default_team_set() + if default_teams: + await add_new_user_to_default_team( + user_id=user_id, + user_email=user_email, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + teams=default_teams, + prisma_client=prisma_client, + ) else: if should_check_db: _update_last_db_access_time( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f741783134e..a1592d512f5 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -314,11 +314,13 @@ async def add_new_user_to_default_team( tasks = [] for team in teams: user_role: Literal["user", "admin"] = "user" + max_budget_in_team: float | None = None if isinstance(team, str): team_id = team elif isinstance(team, NewUserRequestTeam): team_id = team.team_id user_role = team.user_role + max_budget_in_team = team.max_budget_in_team else: raise ValueError(f"Invalid team type: {type(team)}") @@ -328,6 +330,7 @@ async def add_new_user_to_default_team( team_id=team_id, user_email=user_email, user_api_key_dict=user_api_key_dict, + max_budget_in_team=max_budget_in_team, user_role=user_role, ) ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5e07d1bcbc5..ccb20976df9 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -882,6 +882,61 @@ async def test_get_user_object_upsert_includes_user_email(): assert creation_args["user_id"] == "new_test_user" +@pytest.mark.asyncio +async def test_get_user_object_upsert_routes_default_team_to_membership(monkeypatch): + """Regression for LIT-4324: a configured default team (list of NewUserRequestTeam + dicts) must not be written into the Prisma create payload (teams is a String[] column + that rejects dicts). Instead it must be routed through add_new_user_to_default_team so + the JWT-provisioned user gets a real team membership.""" + default_params = { + "user_role": "internal_user", + "teams": [{"team_id": "default-team", "user_role": "user"}], + } + monkeypatch.setattr(litellm, "default_internal_user_params", default_params) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + mock_user = MagicMock() + mock_user.organization_memberships = [] + mock_prisma_client.db.litellm_usertable.create = AsyncMock(return_value=mock_user) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + with patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.add_new_user_to_default_team", + new_callable=AsyncMock, + ) as mock_add_to_team: + try: + await get_user_object( + user_id="new_jwt_user", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + user_id_upsert=True, + proxy_logging_obj=None, + ) + except Exception as e: + # mock_user is a MagicMock, so the post-create LiteLLM_UserTable(**dict(...)) + # conversion raises; irrelevant to what we assert. + print(e) + + creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"] + assert "teams" not in creation_args, "teams must be popped before the Prisma create" + assert creation_args["user_role"] == "internal_user" + + mock_add_to_team.assert_awaited_once() + passed_teams = mock_add_to_team.await_args[1]["teams"] + assert [team.team_id for team in passed_teams] == ["default-team"] + assert ( + mock_add_to_team.await_args[1]["user_api_key_dict"].user_role + == LitellmUserRoles.PROXY_ADMIN + ) + + def test_log_budget_lookup_failure_dry_run(): """Dry run: verify _log_budget_lookup_failure logs for schema/DB errors.""" with patch("litellm.proxy.auth.auth_checks.verbose_proxy_logger") as mock_logger: diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index ce2d04f0d26..be0267d69c5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3550,3 +3550,58 @@ async def test_resolve_user_email_metadata_skips_db_when_no_user_ids(mocker): assert result == {} find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_new_user_to_default_team_propagates_max_budget_in_team(mocker): + """A configured per-member budget on a default team must reach the membership + write; dropping it means the member is unlimited within the team budget.""" + from litellm.proxy._types import NewUserRequestTeam + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + add_new_user_to_default_team, + ) + + mock_add = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team", + new_callable=mocker.AsyncMock, + ) + + await add_new_user_to_default_team( + user_id="jwt-user", + user_email="jwt-user@example.com", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + teams=[ + NewUserRequestTeam(team_id="budgeted-team", max_budget_in_team=25.0, user_role="admin"), + NewUserRequestTeam(team_id="uncapped-team"), + ], + prisma_client=mocker.MagicMock(), + ) + + calls = {c.kwargs["team_id"]: c.kwargs for c in mock_add.call_args_list} + assert calls["budgeted-team"]["max_budget_in_team"] == 25.0 + assert calls["budgeted-team"]["user_role"] == "admin" + assert calls["uncapped-team"]["max_budget_in_team"] is None + + +@pytest.mark.asyncio +async def test_add_new_user_to_default_team_string_teams_have_no_member_budget(mocker): + """Bare-string default teams carry no per-member budget.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + add_new_user_to_default_team, + ) + + mock_add = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._add_user_to_team", + new_callable=mocker.AsyncMock, + ) + + await add_new_user_to_default_team( + user_id="jwt-user", + user_email=None, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + teams=["string-team"], + prisma_client=mocker.MagicMock(), + ) + + assert mock_add.call_args.kwargs["max_budget_in_team"] is None + assert mock_add.call_args.kwargs["team_id"] == "string-team"