diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index ff9211742f3f..1625198892fa 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -52,20 +52,20 @@ def _get_models_from_access_groups( model_access_groups: dict[str, list[str]], all_models: list[str], include_model_access_groups: bool | None = False, + proxy_model_list: Sequence[str] | None = None, ) -> list[str]: - idx_to_remove: Final = [] - new_models: Final = [] - for idx, model in enumerate(all_models): - if model in model_access_groups: - if not include_model_access_groups: # remove access group, unless requested - e.g. when creating a key - idx_to_remove.append(idx) - new_models.extend(model_access_groups[model]) - - for idx in sorted(idx_to_remove, reverse=True): - all_models.pop(idx) - - all_models.extend(new_models) - return all_models + # a grant naming both a deployed model and an access group means both at runtime + # (_check_model_access_helper unions them), so listings must keep the literal too + deployed_model_names: Final = frozenset(proxy_model_list or ()) + kept_models: Final = [ + model + for model in all_models + if model not in model_access_groups or include_model_access_groups or model in deployed_model_names + ] + member_models: Final = [ + member for model in all_models if model in model_access_groups for member in model_access_groups[model] + ] + return kept_models + member_models async def get_mcp_server_ids( @@ -128,6 +128,7 @@ def get_key_models( model_access_groups=model_access_groups, all_models=all_models, include_model_access_groups=include_model_access_groups, + proxy_model_list=proxy_model_list, ) # deduplicate while preserving order @@ -169,6 +170,7 @@ def get_team_models( model_access_groups=model_access_groups, all_models=list(all_models_set), include_model_access_groups=include_model_access_groups, + proxy_model_list=proxy_model_list, ) # deduplicate while preserving order diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0f0079c542f5..c3dce158bfb6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11858,6 +11858,8 @@ def _add_team_models_to_all_models( Add team models to all models """ team_models: Final[dict[str, set[str]]] = {} + proxy_model_list: Final = llm_router.get_model_names() + model_access_groups: Final = llm_router.get_model_access_groups() for team_object in team_db_objects_typed: if ( @@ -11879,7 +11881,12 @@ def _add_team_models_to_all_models( if can_add_model: team_models.setdefault(model_id, set()).add(team_object.team_id) else: - for model_name in team_object.models: + resolved_model_names = get_team_models( + team_models=team_object.models, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + for model_name in resolved_model_names: _models = llm_router.get_model_list(model_name=model_name, team_id=team_object.team_id) if _models is not None: for model in _models: diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index e6c0eaee3c4b..5161554b9693 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -133,6 +133,50 @@ def test_get_key_models_passes_include_model_access_groups(): assert "model2" in result +def test_get_key_models_keeps_literal_model_colliding_with_group_name(): + """A name that is BOTH a deployed model and an access group grants both at + runtime (_check_model_access_helper unions them), so the listing must keep + the literal model alongside the group members instead of dropping it.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.model_checks import get_key_models + + user_api_key_dict = UserAPIKeyAuth(models=["beta-models"], api_key="test-key") + + result = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=["beta-models", "member-a", "unrelated"], + model_access_groups={"beta-models": ["member-a"]}, + include_model_access_groups=False, + ) + assert sorted(result) == ["beta-models", "member-a"] + + +def test_get_team_models_keeps_literal_model_colliding_with_group_name(): + """Team flavor of the collision case: literal deployment survives group expansion.""" + from litellm.proxy.auth.model_checks import get_team_models + + result = get_team_models( + team_models=["beta-models"], + proxy_model_list=["beta-models", "member-a", "unrelated"], + model_access_groups={"beta-models": ["member-a"]}, + include_model_access_groups=False, + ) + assert sorted(result) == ["beta-models", "member-a"] + + +def test_get_team_models_drops_group_name_that_is_not_a_deployed_model(): + """No collision: a pure access-group name is still replaced by its members.""" + from litellm.proxy.auth.model_checks import get_team_models + + result = get_team_models( + team_models=["beta-models"], + proxy_model_list=["member-a", "unrelated"], + model_access_groups={"beta-models": ["member-a"]}, + include_model_access_groups=False, + ) + assert result == ["member-a"] + + def test_get_key_models_does_not_mutate_input(): """ get_key_models must not mutate user_api_key_dict.models in-place. 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 e73f1d08cb5e..2f4018b55aba 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 @@ -580,6 +580,63 @@ async def test_populate_team_access_gives_view_only_admin_full_admin_scope(monke assert by_id["global-id-1"]["model_info"]["direct_access"] is True +@pytest.mark.asyncio +async def test_populate_team_access_grants_config_access_group_model(): + """LIT-4433: a team whose only model grant is a CONFIG-defined access group + (model_info.access_groups) must have that group's member deployments listed in + access_via_team_ids. Before the fix _add_team_models_to_all_models passed the + access-group name straight to get_model_list, which never matched, leaving the + team's /v2/model/info?include_team_models=true result empty.""" + team_id = "team-access-group-only" + access_group_model = { + "model_name": "team-allowed-model-a", + "litellm_params": {"model": "gpt-4"}, + "model_info": { + "id": "model-a-id", + "access_groups": ["test-access-group"], + "db_model": False, + }, + } + + router = MagicMock() + router.get_model_names.return_value = ["team-allowed-model-a"] + router.get_model_access_groups.return_value = {"test-access-group": ["team-allowed-model-a"]} + router.get_model_ids.return_value = [] + + def get_model_list(model_name=None, team_id=None): + if model_name == "team-allowed-model-a": + return [access_group_model] + return None + + router.get_model_list.side_effect = get_model_list + + team_db_object = MagicMock() + team_db_object.model_dump.return_value = { + "team_id": team_id, + "models": ["test-access-group"], + "access_group_ids": [], + } + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_db_object]) + + admin = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]) + result = await ps._populate_team_access_on_models( + user_api_key_dict=admin, + prisma_client=prisma_client, + llm_router=router, + all_models=[ + { + "model_name": "team-allowed-model-a", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "model-a-id", "access_groups": ["test-access-group"], "db_model": False}, + } + ], + ) + + by_id = {m["model_info"]["id"]: m for m in result} + assert by_id["model-a-id"]["model_info"]["access_via_team_ids"] == [team_id] + + @pytest.mark.asyncio async def test_model_info_v1_team_id_without_db_fails_fast(monkeypatch): """`teamId` without a connected DB raises 500 before any enrichment work runs.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 3acb9fcafd32..57b2c874962a 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1708,6 +1708,149 @@ def test_add_team_models_to_all_models(): assert result == {"gpt-4-model-2": {"team1"}} +def _make_router_with_access_groups(model_names, model_access_groups, deployments): + llm_router = MagicMock() + llm_router.get_model_names.return_value = model_names + llm_router.get_model_access_groups.return_value = model_access_groups + + def get_model_list(model_name=None, team_id=None): + matched = [ + deployment + for deployment in deployments + if deployment["model_name"] == model_name + and ( + team_id is None + or deployment.get("model_info", {}).get("team_id") is None + or deployment.get("model_info", {}).get("team_id") == team_id + ) + ] + return matched or None + + llm_router.get_model_list.side_effect = get_model_list + return llm_router + + +def test_add_team_models_to_all_models_resolves_config_access_group(): + """ + LIT-4433: a CONFIG-defined access group (model_info.access_groups) named in + team.models must resolve to its member deployments' ids. The pre-fix code + passed the group name straight to get_model_list, which never matched, so the + team's /v2/model/info?include_team_models=true result was empty. + """ + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.proxy_server import _add_team_models_to_all_models + + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "team-a" + team.models = ["test-access-group"] + + llm_router = _make_router_with_access_groups( + model_names=["team-allowed-model-a"], + model_access_groups={"test-access-group": ["team-allowed-model-a"]}, + deployments=[{"model_name": "team-allowed-model-a", "model_info": {"id": "model-a-id"}}], + ) + + result = _add_team_models_to_all_models(team_db_objects_typed=[team], llm_router=llm_router) + assert result == {"model-a-id": {"team-a"}} + + +def test_add_team_models_to_all_models_resolves_mixed_literal_and_access_group(): + """A team.models list mixing a literal model name and a config access-group + name must resolve both to their deployment ids.""" + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.proxy_server import _add_team_models_to_all_models + + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "team-a" + team.models = ["team-allowed-model-b", "test-access-group"] + + llm_router = _make_router_with_access_groups( + model_names=["team-allowed-model-a", "team-allowed-model-b"], + model_access_groups={"test-access-group": ["team-allowed-model-a"]}, + deployments=[ + {"model_name": "team-allowed-model-a", "model_info": {"id": "model-a-id"}}, + {"model_name": "team-allowed-model-b", "model_info": {"id": "model-b-id"}}, + ], + ) + + result = _add_team_models_to_all_models(team_db_objects_typed=[team], llm_router=llm_router) + assert result == {"model-a-id": {"team-a"}, "model-b-id": {"team-a"}} + + +def test_add_team_models_to_all_models_keeps_literal_model_colliding_with_group_name(): + """A team.models entry that names BOTH a deployed model and an access group + grants both at runtime, so the /v2 team map must contain the literal + deployment's id alongside the group members' ids.""" + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.proxy_server import _add_team_models_to_all_models + + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "team-a" + team.models = ["beta-models"] + + llm_router = _make_router_with_access_groups( + model_names=["beta-models", "member-a"], + model_access_groups={"beta-models": ["member-a"]}, + deployments=[ + {"model_name": "beta-models", "model_info": {"id": "collision-id"}}, + {"model_name": "member-a", "model_info": {"id": "member-a-id"}}, + ], + ) + + result = _add_team_models_to_all_models(team_db_objects_typed=[team], llm_router=llm_router) + assert result == {"collision-id": {"team-a"}, "member-a-id": {"team-a"}} + + +def test_add_team_models_to_all_models_excludes_other_access_group(): + """Only the access group named in team.models is expanded; deployments that + belong solely to a different access group must not leak into the team map.""" + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.proxy_server import _add_team_models_to_all_models + + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "team-a" + team.models = ["test-access-group"] + + llm_router = _make_router_with_access_groups( + model_names=["team-allowed-model-a", "forbidden-model"], + model_access_groups={ + "test-access-group": ["team-allowed-model-a"], + "other-access-group": ["forbidden-model"], + }, + deployments=[ + {"model_name": "team-allowed-model-a", "model_info": {"id": "model-a-id"}}, + {"model_name": "forbidden-model", "model_info": {"id": "forbidden-id"}}, + ], + ) + + result = _add_team_models_to_all_models(team_db_objects_typed=[team], llm_router=llm_router) + assert result == {"model-a-id": {"team-a"}} + + +def test_add_team_models_to_all_models_excludes_other_teams_byok_with_shared_name(): + """A BYOK deployment owned by a DIFFERENT team but sharing the resolved model + name must not be added for this team. Guards the team_id filter passed to + get_model_list: dropping it would leak the other team's private deployment.""" + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.proxy_server import _add_team_models_to_all_models + + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "team-a" + team.models = ["test-access-group"] + + llm_router = _make_router_with_access_groups( + model_names=["team-allowed-model-a"], + model_access_groups={"test-access-group": ["team-allowed-model-a"]}, + deployments=[ + {"model_name": "team-allowed-model-a", "model_info": {"id": "model-a-id", "team_id": "team-a"}}, + {"model_name": "team-allowed-model-a", "model_info": {"id": "other-team-byok-id", "team_id": "team-b"}}, + ], + ) + + result = _add_team_models_to_all_models(team_db_objects_typed=[team], llm_router=llm_router) + assert result == {"model-a-id": {"team-a"}} + + @pytest.mark.asyncio async def test_apply_search_filter_matches_team_public_model_name(): """