diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 75d4d13eb71..96d7a379a2c 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 37484 + "limit": 37483 }, "reportArgumentType": { "limit": 2704 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10389 + "limit": 10388 }, "reportFunctionMemberAccess": { "limit": 11 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8a1cb6ae46e..364d8f1b886 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -786,22 +786,55 @@ async def _check_org_team_limits( ) -def _tightest_cap(*caps: Optional[float]) -> Optional[float]: +def _tightest_cap(*caps: float | None) -> float | None: set_caps = tuple(cap for cap in caps if cap is not None) return min(set_caps) if set_caps else None +def _effective_caller_model_scope( + key_models: list[str], + user_models: list[str] | None, +) -> list[str] | None: + """The models a caller may grant to a new team, or None when unrestricted. + + A key with `models=[]` is nominally unrestricted, but on a personal key the + user's own model list still gates every call at request time, so the + caller's true authority is the intersection of both layers. Reading the key + layer alone would let a restricted user mint an all-models team. + `all-proxy-models` (or an empty list) marks a layer unrestricted; + `no-default-models` marks it as granting nothing, yielding an empty scope. + """ + + def _restriction(models: list[str] | None) -> list[str] | None: + if not models: + return None + if SpecialModelNames.no_default_models.value in models: + return [] + if SpecialModelNames.all_proxy_models.value in models: + return None + return models + + key_scope = _restriction(key_models) + user_scope = _restriction(user_models) + if key_scope is None: + return user_scope + if user_scope is None: + return key_scope + return [m for m in key_scope if m in user_scope] + + def _inherit_caller_limits_for_self_served_team( data: NewTeamRequest, user_api_key_dict: UserAPIKeyAuth, + user_obj: LiteLLM_UserTable | None, ) -> NewTeamRequest: """Clamp a self-served team to the creating user's effective limits. A self-served team must never be wider than its creator. The creator's authority spans two layers: the calling key (`tpm_limit`/`rpm_limit`/ `models`) and the underlying user (`user_tpm_limit`/`user_rpm_limit`/ - `user_max_budget`). The primary caller is a UI SSO session whose key - carries no tpm/rpm and only a tiny per-session `max_budget` + `user_max_budget`/`models`). The primary caller is a UI SSO session whose + key carries no tpm/rpm and only a tiny per-session `max_budget` (`max_ui_session_budget`), so reading key limits alone would leave the team uncapped. We take the tightest of both layers per dimension and clamp the team down to it, filling unset fields and shrinking any value (including one @@ -812,9 +845,13 @@ def _inherit_caller_limits_for_self_served_team( effective_tpm = _tightest_cap(user_api_key_dict.tpm_limit, user_api_key_dict.user_tpm_limit) effective_rpm = _tightest_cap(user_api_key_dict.rpm_limit, user_api_key_dict.user_rpm_limit) effective_budget = user_api_key_dict.user_max_budget + effective_models = _effective_caller_model_scope( + key_models=list(user_api_key_dict.models), + user_models=list(user_obj.models) if user_obj is not None else None, + ) return data.model_copy( update={ - "models": data.models if data.models else list(user_api_key_dict.models), + "models": data.models if data.models else list(effective_models or []), "tpm_limit": _tightest_cap(data.tpm_limit, effective_tpm), "rpm_limit": _tightest_cap(data.rpm_limit, effective_rpm), "max_budget": _tightest_cap(data.max_budget, effective_budget), @@ -825,15 +862,17 @@ def _inherit_caller_limits_for_self_served_team( async def _check_user_team_limits( data: Union[NewTeamRequest, UpdateTeamRequest], user_api_key_dict: UserAPIKeyAuth, - prisma_client: PrismaClient, - user_api_key_cache: Any, + user_obj: LiteLLM_UserTable | None, ) -> None: """ Enforce the caller's personal limits when CREATING a standalone team. This validates the requested team budget / models / tpm / rpm against the caller's own limits, so a non-admin user cannot mint a brand-new team that - is richer than themselves. + is richer than themselves. Model scope is the intersection of the key's and + the user's model lists (see _effective_caller_model_scope); when that scope + is restricted, an empty team model list is rejected too, because an empty + list would grant the team every proxy model. Only used by /team/new for standalone teams (organization_id is None). /team/update does NOT call this — an existing team's admin is already @@ -841,14 +880,7 @@ async def _check_user_team_limits( wallet. Org-scoped teams use _check_org_team_limits() instead. """ # Validate team budget against user's max_budget - if data.max_budget is not None and user_api_key_dict.user_id is not None: - user_obj = await get_user_object( - user_id=user_api_key_dict.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - ) - + if data.max_budget is not None: if user_obj is not None and user_obj.max_budget is not None and data.max_budget > user_obj.max_budget: raise HTTPException( status_code=400, @@ -857,14 +889,25 @@ async def _check_user_team_limits( }, ) - # Validate team models against user's allowed models - if data.models is not None and len(user_api_key_dict.models) > 0: + # Validate team models against the caller's effective model scope + effective_models = _effective_caller_model_scope( + key_models=list(user_api_key_dict.models), + user_models=list(user_obj.models) if user_obj is not None else None, + ) + if effective_models is not None: + if not data.models: + raise HTTPException( + status_code=400, + detail={ + "error": f"An empty team model list would grant access to all proxy models, which exceeds the caller's allowed models. Allowed models={effective_models}. User id={user_api_key_dict.user_id}" + }, + ) for m in data.models: - if m not in user_api_key_dict.models: + if m not in effective_models: raise HTTPException( status_code=400, detail={ - "error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}. User id={user_api_key_dict.user_id}" + "error": f"Model not in allowed user models. User allowed models={effective_models}. User id={user_api_key_dict.user_id}" }, ) @@ -1165,6 +1208,15 @@ async def new_team( # Only validate user budget/models/tpm/rpm for standalone teams (not org-scoped) # For org-scoped teams, validation is done by _check_org_team_limits() if data.organization_id is None: + try: + caller_user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except ValueError: + caller_user_obj = None if ( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER and RouteChecks._user_team_creation_enabled() @@ -1172,12 +1224,12 @@ async def new_team( data = _inherit_caller_limits_for_self_served_team( data=data, user_api_key_dict=user_api_key_dict, + user_obj=caller_user_obj, ) await _check_user_team_limits( data=data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + user_obj=caller_user_obj, ) if _should_auto_add_team_creator(user_api_key_dict, general_settings): diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 448a0079674..2d0686b65f5 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 2075 + "limit": 2074 }, "ASYNC230": { "limit": 14 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 5149c6b95ab..b14c678c22b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4485,7 +4485,7 @@ async def test_new_team_standalone_validates_against_user_models(monkeypatch): # Verify exception details assert exc_info.value.code == "400" assert "Model not in allowed user models" in str(exc_info.value.message) - assert "no-default-models" in str(exc_info.value.message) + assert "allowed models=[]" in str(exc_info.value.message) @pytest.mark.asyncio @@ -9729,6 +9729,7 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t", models=[]), user_api_key_dict=caller, + user_obj=None, ) assert result.models == ["gpt-5", "gpt-5-mini"] @@ -9742,6 +9743,7 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t"), user_api_key_dict=caller, + user_obj=None, ) assert result.tpm_limit == 1000 assert result.rpm_limit == 10 @@ -9766,6 +9768,7 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t"), user_api_key_dict=caller, + user_obj=None, ) assert result.tpm_limit == 100 assert result.rpm_limit == 5 @@ -9781,6 +9784,7 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t"), user_api_key_dict=caller, + user_obj=None, ) assert result.tpm_limit == 100 assert result.rpm_limit == 5 @@ -9798,6 +9802,7 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t", tpm_limit=100000, rpm_limit=9999, max_budget=1000.0), user_api_key_dict=caller, + user_obj=None, ) assert result.tpm_limit == 100 assert result.rpm_limit == 5 @@ -9813,6 +9818,7 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t", models=["a"], tpm_limit=1000, rpm_limit=50, max_budget=25.0), user_api_key_dict=caller, + user_obj=None, ) assert result.models == ["a"] assert result.tpm_limit == 1000 @@ -9836,6 +9842,7 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t"), user_api_key_dict=caller, + user_obj=None, ) assert result.models == [] assert result.tpm_limit is None @@ -9854,9 +9861,162 @@ class TestSelfServeTeamLimitInheritance: result = _inherit_caller_limits_for_self_served_team( data=NewTeamRequest(team_alias="t"), user_api_key_dict=caller, + user_obj=None, ) assert result.max_budget is None + def _user(self, **kwargs): + from litellm.proxy._types import LiteLLM_UserTable + + return LiteLLM_UserTable(user_id="u1", **kwargs) + + def test_unrestricted_key_falls_back_to_user_models(self): + """A key with models=[] is nominally unrestricted, but on a personal + key the user's own model list still gates every call, so the team must + inherit the user's models rather than become an all-models team.""" + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _inherit_caller_limits_for_self_served_team, + ) + + caller = self._caller(models=[]) + result = _inherit_caller_limits_for_self_served_team( + data=NewTeamRequest(team_alias="t"), + user_api_key_dict=caller, + user_obj=self._user(models=["gpt-5"]), + ) + assert result.models == ["gpt-5"] + + def test_key_and_user_models_intersect(self): + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _inherit_caller_limits_for_self_served_team, + ) + + caller = self._caller(models=["gpt-5", "gpt-5-mini"]) + result = _inherit_caller_limits_for_self_served_team( + data=NewTeamRequest(team_alias="t"), + user_api_key_dict=caller, + user_obj=self._user(models=["gpt-5-mini", "claude-opus-4-8"]), + ) + assert result.models == ["gpt-5-mini"] + + def test_all_proxy_models_user_sentinel_keeps_key_scope(self): + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _inherit_caller_limits_for_self_served_team, + ) + + caller = self._caller(models=["gpt-5"]) + result = _inherit_caller_limits_for_self_served_team( + data=NewTeamRequest(team_alias="t"), + user_api_key_dict=caller, + user_obj=self._user(models=["all-proxy-models"]), + ) + assert result.models == ["gpt-5"] + + +class TestEffectiveCallerModelScope: + def test_both_layers_unrestricted_is_none(self): + from litellm.proxy.management_endpoints.team_endpoints import _effective_caller_model_scope + + assert _effective_caller_model_scope(key_models=[], user_models=None) is None + assert _effective_caller_model_scope(key_models=[], user_models=[]) is None + assert _effective_caller_model_scope(key_models=["all-proxy-models"], user_models=["all-proxy-models"]) is None + + def test_single_restricted_layer_wins(self): + from litellm.proxy.management_endpoints.team_endpoints import _effective_caller_model_scope + + assert _effective_caller_model_scope(key_models=["m-a"], user_models=None) == ["m-a"] + assert _effective_caller_model_scope(key_models=[], user_models=["m-b"]) == ["m-b"] + + def test_intersection_when_both_restricted(self): + from litellm.proxy.management_endpoints.team_endpoints import _effective_caller_model_scope + + assert _effective_caller_model_scope(key_models=["m-a", "m-b"], user_models=["m-b", "m-c"]) == ["m-b"] + + def test_no_default_models_grants_nothing(self): + from litellm.proxy.management_endpoints.team_endpoints import _effective_caller_model_scope + + assert _effective_caller_model_scope(key_models=[], user_models=["no-default-models"]) == [] + + +class TestCheckUserTeamLimitsModelScope: + def _caller(self, **kwargs): + from litellm.proxy._types import UserAPIKeyAuth + + return UserAPIKeyAuth(api_key="sk-test", user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, **kwargs) + + def _user(self, **kwargs): + from litellm.proxy._types import LiteLLM_UserTable + + return LiteLLM_UserTable(user_id="u1", **kwargs) + + @pytest.mark.asyncio + async def test_rejects_model_beyond_user_scope_despite_unrestricted_key(self): + """Greptile P2 regression: with a models=[] key, the user's own model + restrictions must still bound the team's model list.""" + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _check_user_team_limits + + with pytest.raises(HTTPException) as exc: + await _check_user_team_limits( + data=NewTeamRequest(team_alias="t", models=["gpt-5"]), + user_api_key_dict=self._caller(models=[]), + user_obj=self._user(models=["gpt-5-mini"]), + ) + assert exc.value.status_code == 400 + + @pytest.mark.asyncio + async def test_rejects_empty_team_models_when_scope_restricted(self): + """An empty team model list means all proxy models, so it must be + rejected whenever the caller's effective scope is restricted.""" + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _check_user_team_limits + + with pytest.raises(HTTPException) as exc: + await _check_user_team_limits( + data=NewTeamRequest(team_alias="t", models=[]), + user_api_key_dict=self._caller(models=[]), + user_obj=self._user(models=["gpt-5-mini"]), + ) + assert exc.value.status_code == 400 + + @pytest.mark.asyncio + async def test_allows_models_within_scope(self): + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _check_user_team_limits + + await _check_user_team_limits( + data=NewTeamRequest(team_alias="t", models=["gpt-5-mini"]), + user_api_key_dict=self._caller(models=[]), + user_obj=self._user(models=["gpt-5-mini", "gpt-5"]), + ) + + @pytest.mark.asyncio + async def test_unrestricted_caller_allows_empty_models(self): + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _check_user_team_limits + + await _check_user_team_limits( + data=NewTeamRequest(team_alias="t", models=[]), + user_api_key_dict=self._caller(models=[]), + user_obj=self._user(models=[]), + ) + + @pytest.mark.asyncio + async def test_budget_check_uses_injected_user_obj(self): + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _check_user_team_limits + + with pytest.raises(HTTPException) as exc: + await _check_user_team_limits( + data=NewTeamRequest(team_alias="t", models=["gpt-5"], max_budget=100.0), + user_api_key_dict=self._caller(models=[]), + user_obj=self._user(models=["gpt-5"], max_budget=10.0), + ) + assert exc.value.status_code == 400 + @pytest.mark.parametrize( "flag_enabled,caller_role,expect_inherited", @@ -9870,8 +10030,9 @@ class TestSelfServeTeamLimitInheritance: async def test_new_team_self_serve_inheritance_call_site(flag_enabled, caller_role, expect_inherited): """LIT-3254: new_team applies caller-limit inheritance exactly when allow_user_team_creation is on AND the caller is an INTERNAL_USER. - Org admins and flag-off callers keep today's behavior (an omitted models - list stays empty).""" + For org admins and flag-off callers no inheritance runs, so their omitted + model list (= all proxy models) now fails validation against the caller's + restricted key instead of silently minting an all-models team.""" from fastapi import Request from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth @@ -9935,21 +10096,25 @@ async def test_new_team_self_serve_inheritance_call_site(flag_enabled, caller_ro mock_prisma.db.litellm_teammembership = MagicMock() mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) - await new_team( - data=team_request, - http_request=dummy_request, - user_api_key_dict=caller, - ) - - created_row = mock_prisma.db.litellm_teamtable.create.call_args.kwargs["data"] if expect_inherited: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=caller, + ) + created_row = mock_prisma.db.litellm_teamtable.create.call_args.kwargs["data"] assert created_row["models"] == ["gpt-5"] assert created_row["tpm_limit"] == 1000 assert created_row["rpm_limit"] == 10 else: - assert created_row["models"] == [] - assert created_row.get("tpm_limit") is None - assert created_row.get("rpm_limit") is None + with pytest.raises((HTTPException, ProxyException)) as exc: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=caller, + ) + assert str(getattr(exc.value, "status_code", None) or exc.value.code) == "400" + mock_prisma.db.litellm_teamtable.create.assert_not_called() @pytest.mark.asyncio