diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f536f521d57..a01cafab4ff 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12373,25 +12373,53 @@ async def get_all_team_models( return returned_team_models +def _resolve_model_grant_to_deployment_ids( + models: Sequence[str], + llm_router: Router, +) -> tuple[str, ...]: + """ + Resolve a `models` grant (a user's or a key's) to the deployment ids it can call. + + An empty grant and the 'all-proxy-models' sentinel both mean unrestricted at call + time (see `_check_model_access_helper`), so both expand to every non-team deployment. + A grant entry naming an access group also grants that group's members, and naming a + deployed model that shares the name grants the model itself, matching the union the + call-time check applies. + """ + if not models or SpecialModelNames.all_proxy_models.value in models: + return tuple(llm_router.get_model_ids(exclude_team_models=True)) + + access_groups: Final = llm_router.get_model_access_groups() + granted_model_names: Final = tuple(name for model in models for name in (model, *access_groups.get(model, ()))) + return tuple( + model_id + for name in granted_model_names + for deployment in (llm_router.get_model_list(model_name=name) or ()) + if (model_id := deployment.get("model_info", {}).get("id", None)) is not None + ) + + def get_direct_access_models( user_db_object: LiteLLM_UserTable, llm_router: Router, -) -> list[str]: + key_models: Sequence[str] = (), +) -> tuple[str, ...]: """ - Get all models that user has direct access to. + Get all models the caller has direct (non-team) access to. - The 'all-proxy-models' sentinel grants direct access to every non-team - deployment, mirroring how get_key_models expands it for the key/team path. + Both the user record and the calling key are enforced at call time, so direct access + is the intersection of the two grants. An unrestricted key (empty grant, or the + 'all-proxy-models' sentinel) leaves the user's grant untouched. """ - if not user_db_object.models or SpecialModelNames.all_proxy_models.value in user_db_object.models: - return llm_router.get_model_ids(exclude_team_models=True) + user_model_ids: Final = _resolve_model_grant_to_deployment_ids( + cast(Sequence[str], user_db_object.models), # cast-ok: user.models is a String[] column + llm_router, + ) + if not key_models or SpecialModelNames.all_proxy_models.value in key_models: + return user_model_ids - return [ - model_id - for model in user_db_object.models - for deployment in (llm_router.get_model_list(model_name=model) or []) - if (model_id := deployment.get("model_info", {}).get("id", None)) is not None - ] + key_model_ids: Final = frozenset(_resolve_model_grant_to_deployment_ids(key_models, llm_router)) + return tuple(model_id for model_id in user_model_ids if model_id in key_model_ids) def _filter_models_to_user_accessible(all_models: list[dict]) -> list[dict]: @@ -12415,10 +12443,10 @@ async def _populate_team_access_on_models( without filtering the model list. """ user_teams: list[str] | Literal["*"] | None = None - direct_access_models: list[str] = [] + direct_access_models: Sequence[str] = () if _user_has_admin_view(user_api_key_dict): user_teams = "*" - direct_access_models = llm_router.get_model_ids(exclude_team_models=True) # has access to all models + direct_access_models = tuple(llm_router.get_model_ids(exclude_team_models=True)) # access to all models elif user_api_key_dict.user_id is not None: user_db_object: Final[SupportsModelDump | None] = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} @@ -12429,6 +12457,7 @@ async def _populate_team_access_on_models( direct_access_models = get_direct_access_models( user_db_object=user_object, llm_router=llm_router, + key_models=cast(Sequence[str], user_api_key_dict.models), # cast-ok: key.models is a String[] column ) if user_teams is not None: team_models: Final = await get_all_team_models( @@ -12450,7 +12479,7 @@ async def _populate_team_access_on_models( if can_use_model: _model["model_info"]["access_via_team_ids"] = team_models.get(model_id, []) - direct_access_model_ids: Final = set(direct_access_models) + direct_access_model_ids: Final = frozenset(direct_access_models) for _model in all_models: model_id = _model.get("model_info", {}).get("id", None) if model_id is not None: diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index cb1521f10ae..9a8dde3c512 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -1487,7 +1487,7 @@ def test_get_direct_access_models_expands_all_proxy_models_sentinel(): result = ps.get_direct_access_models(user_db_object=user, llm_router=router) - assert result == ["global-id-1", "global-id-2"] + assert result == ("global-id-1", "global-id-2") router.get_model_ids.assert_called_once_with(exclude_team_models=True) router.get_model_list.assert_not_called() @@ -1502,7 +1502,7 @@ def test_get_direct_access_models_resolves_explicit_model_names(): result = ps.get_direct_access_models(user_db_object=user, llm_router=router) - assert result == ["gpt4o-id"] + assert result == ("gpt4o-id",) router.get_model_ids.assert_not_called() router.get_model_list.assert_called_once_with(model_name="gpt-4o") @@ -1519,7 +1519,7 @@ def test_get_direct_access_models_empty_models_grants_all_non_team_models(): result = ps.get_direct_access_models(user_db_object=user, llm_router=router) - assert result == ["global-id-1", "global-id-2"] + assert result == ("global-id-1", "global-id-2") router.get_model_ids.assert_called_once_with(exclude_team_models=True) router.get_model_list.assert_not_called() @@ -1562,6 +1562,112 @@ async def test_populate_team_access_grants_empty_models_user_direct_access(monke assert visible[0]["model_info"]["direct_access"] is True +def test_get_direct_access_models_restricted_key_narrows_unrestricted_user(): + """A key scoped to one model cannot call the rest, so the listing must not show + every non-team model just because the user record is unrestricted.""" + router = MagicMock() + router.get_model_ids.return_value = ["gpt4o-id", "sonnet-id"] + router.get_model_access_groups.return_value = {} + router.get_model_list.side_effect = lambda model_name: ( + [{"model_info": {"id": "gpt4o-id"}}] if model_name == "gpt-4o" else [] + ) + + user = LiteLLM_UserTable(user_id="u", models=[], teams=[]) + + result = ps.get_direct_access_models(user_db_object=user, llm_router=router, key_models=("gpt-4o",)) + + assert result == ("gpt4o-id",) + + +def test_get_direct_access_models_all_proxy_models_key_keeps_team_scoped_user_grant(): + """'all-proxy-models' on the key means unrestricted, so it must leave the user's + grant alone rather than clipping it to the non-team deployment set.""" + router = MagicMock() + router.get_model_ids.return_value = ["global-id"] + router.get_model_access_groups.return_value = {} + router.get_model_list.side_effect = lambda model_name: ( + [{"model_info": {"id": "byok-id"}}] if model_name == "byok-model" else [] + ) + + user = LiteLLM_UserTable(user_id="u", models=["byok-model"], teams=[]) + + result = ps.get_direct_access_models( + user_db_object=user, + llm_router=router, + key_models=(ps.SpecialModelNames.all_proxy_models.value,), + ) + + assert result == ("byok-id",) + + +def test_get_direct_access_models_expands_access_group_grant(): + """A grant naming an access group can call the group's members at call time, so the + listing must resolve the members instead of looking up the group name as a model.""" + router = MagicMock() + router.get_model_access_groups.return_value = {"beta-models": ["gpt-4o", "sonnet"]} + router.get_model_list.side_effect = lambda model_name: { + "gpt-4o": [{"model_info": {"id": "gpt4o-id"}}], + "sonnet": [{"model_info": {"id": "sonnet-id"}}], + }.get(model_name, []) + + user = LiteLLM_UserTable(user_id="u", models=["beta-models"], teams=[]) + + result = ps.get_direct_access_models(user_db_object=user, llm_router=router) + + assert result == ("gpt4o-id", "sonnet-id") + + +@pytest.mark.asyncio +async def test_populate_team_access_hides_models_the_calling_key_cannot_call(monkeypatch): + """An unrestricted user calling with a key scoped to one model must only see that + model as direct access; the others 403 at the key check, so listing them over-promises.""" + allowed_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "gpt4o-id", "db_model": False}, + } + blocked_row = { + "model_name": "sonnet", + "litellm_params": {"model": "sonnet"}, + "model_info": {"id": "sonnet-id", "db_model": False}, + } + + router = MagicMock() + router.get_model_ids.return_value = ["gpt4o-id", "sonnet-id"] + router.get_model_access_groups.return_value = {} + router.get_model_list.side_effect = lambda model_name: ( + [{"model_info": {"id": "gpt4o-id"}}] if model_name == "gpt-4o" else [] + ) + + user_row = LiteLLM_UserTable( + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=[], + teams=[], + ) + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + + monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={})) + + caller = UserAPIKeyAuth( + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["gpt-4o"], + team_models=[], + ) + + populated = await ps._populate_team_access_on_models( + user_api_key_dict=caller, + prisma_client=prisma_client, + llm_router=router, + all_models=[allowed_row, blocked_row], + ) + visible = ps._filter_models_to_user_accessible(populated) + + assert [m["model_info"]["id"] for m in visible] == ["gpt4o-id"] + + @pytest.mark.asyncio async def test_populate_team_access_grants_all_proxy_models_user_direct_access( monkeypatch,