From bdd80781f2725932d313627e825c79bd34fadcd7 Mon Sep 17 00:00:00 2001 From: L <13051960278@163.com> Date: Fri, 16 Oct 2026 01:23:03 +0800 Subject: [PATCH] fix(proxy): list team access group models in /v1/model/info and keep key model lists narrow (#32599) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(proxy): list models from team access groups * 修复 PR #32599 的类型检查并补齐模型列表回退测试 * fix(proxy): keep a key's own model list narrow when adding team access group models A key with its own model list was shown every model from its team's access groups on /v1/models, /v1/model/info and /model_group/info, though auth denies those calls. Team access group models now widen a key's list only when the key inherits the team's models. /v1/model/info reuses the shared listing helper without a catch-all, and route-level tests replace the mock-only ones * fix(proxy): skip the list build in the unrestricted model info check and test every listing route The unrestricted check for /v1/model/info only needs to know whether the key and team lists are empty, so it no longer builds and expands the model list a second time. The route test now covers /v1/models and /model_group/info next to both model info routes * fix(proxy): scope model info for personal and agent keys with no model list * fix(proxy): keep config keys' team access groups and own-team team_id scope in model listing --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 64 +++++-- litellm/proxy/utils.py | 33 +++- .../proxy_server/test_routes_model_info.py | 168 +++++++++++++++++- .../proxy/utils/helpers/test_model_access.py | 58 ++++++ 4 files changed, 297 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 03ab1752a92..504775c2a91 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -347,6 +347,7 @@ from litellm.proxy._experimental.mcp_server.catalog import public_catalog_operat from litellm.proxy._experimental.mcp_server.stdio_gate import MCP_STDIO_ENABLED_ENV_VAR, is_mcp_stdio_flag_key from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot from litellm.proxy._types import * +from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) @@ -357,6 +358,7 @@ from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, get_team_object, log_db_metrics, + personal_key_team_model_access_applies, ) from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET from litellm.proxy.auth.auth_utils import ( @@ -16493,41 +16495,64 @@ def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed return isinstance(team_public_model_name, str) and team_public_model_name in allowed_model_names -def _get_v1_model_info_allowed_model_names( +async def _v1_model_info_lists_every_deployment( user_api_key_dict: UserAPIKeyAuth, llm_router: Router, -) -> set[str] | None: - """Return key/team allowlisted public model names, or None if unrestricted.""" + resolve_agent_ceiling: CeilingResolver, +) -> bool: model_access_groups: Final = llm_router.get_model_access_groups() proxy_model_list: Final = llm_router.get_model_names() - key_models: Final = get_key_models( + if get_key_models( user_api_key_dict=user_api_key_dict, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, - ) - team_models: Final = get_team_models( + ) or get_team_models( team_models=user_api_key_dict.team_models, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, - ) - if not key_models and not team_models: + ): + return False + if personal_key_team_model_access_applies(user_api_key_dict): + return False + return not user_api_key_dict.agent_id or await resolve_agent_ceiling(user_api_key_dict.agent_id) is None + + +async def _get_v1_model_info_allowed_model_names_with_access_groups( + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + resolve_agent_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> frozenset[str] | None: + if await _v1_model_info_lists_every_deployment( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + resolve_agent_ceiling=resolve_agent_ceiling, + ): return None - return set( - get_complete_model_list( - key_models=key_models, - team_models=team_models, - proxy_model_list=proxy_model_list, - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + + from litellm.proxy.utils import ( + get_available_models_for_user, # pyright: ignore[reportUnknownVariableType] # untyped settings parameter + ) + + return frozenset( + await get_available_models_for_user( + user_api_key_dict=user_api_key_dict, llm_router=llm_router, - return_wildcard_routes=False, + general_settings=general_settings, + user_model=cast(str | None, user_model), # cast-ok: legacy untyped proxy setting + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + resolve_agent_ceiling=resolve_agent_ceiling, ) ) def _filter_v1_model_info_deployments( all_models: list[dict], - allowed_model_names: set[str] | None, + allowed_model_names: Collection[str] | None, ) -> list[dict]: if allowed_model_names is None: return all_models @@ -16730,9 +16755,12 @@ async def model_info_v1( all_models = expand_wildcard_deployments_for_model_info(all_models) - allowed_model_names: Final = _get_v1_model_info_allowed_model_names( + allowed_model_names: Final = await _get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=user_api_key_dict, llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) all_models = _filter_v1_model_info_deployments( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7b4182e5a6f..5b5956eba12 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -167,6 +167,7 @@ from litellm.proxy._types import ( CallInfo, LiteLLM_VerificationTokenView, Member, + SpecialModelNames, UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling @@ -8873,13 +8874,19 @@ async def _get_team_object_for_access_groups( return None +@dataclass(frozen=True, slots=True) +class _AccessGroupModels: + team: tuple[str, ...] = () + key: tuple[str, ...] = () + + async def _get_access_group_models( user_api_key_dict: "UserAPIKeyAuth", team_object: Optional["LiteLLM_TeamTableCachedObj"], prisma_client: Optional["PrismaClient"], user_api_key_cache: Optional["UserApiKeyCache"], proxy_logging_obj: Optional["ProxyLogging"], -) -> tuple[str, ...]: +) -> _AccessGroupModels: from litellm.proxy.auth.auth_checks import ( get_authorized_resources_from_key_access_groups, get_models_from_access_groups, @@ -8896,7 +8903,7 @@ async def _get_access_group_models( team_object=team_object, resource_field="access_model_names", ) - return tuple(dict.fromkeys((*team_group_models, *key_group_models))) + return _AccessGroupModels(team=tuple(team_group_models), key=tuple(key_group_models)) async def _agent_access_group_visible_models( @@ -9008,7 +9015,7 @@ async def get_available_models_for_user( key_models: Final[Sequence[str]] = ( () - if requested_team_object is not None + if requested_team_object is not None and requested_team_object.team_id != user_api_key_dict.team_id else get_key_models( user_api_key_dict=user_api_key_dict, proxy_model_list=proxy_model_list, @@ -9043,11 +9050,25 @@ async def get_available_models_for_user( proxy_logging_obj=proxy_logging_obj, ) if key_models or team_models - else () + else _AccessGroupModels() ) - granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models - granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models + key_model_names: Final = cast("list[str]", user_api_key_dict.models) # cast-ok: key models are a string array + key_model_check_skipped: Final = ( + bool(user_api_key_dict.config) or SpecialModelNames.all_team_models.value in key_model_names + ) + granted_key_models: Final = ( + ( + *key_models, + *access_group_models.key, + *(access_group_models.team if key_model_check_skipped else ()), + ) + if key_models + else key_models + ) + granted_team_models: Final = ( + (*team_models, *access_group_models.team, *access_group_models.key) if team_models else team_models + ) # Get complete model list complete_models: Final = get_complete_model_list( diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 37e02059179..cb7b5b13fa3 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -463,7 +463,7 @@ def test_model_info_id_lookup_reports_the_same_cost_as_the_list( _invalidate_model_cost_lowercase_map() -def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch): +async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch): from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.proxy.auth import model_checks @@ -483,12 +483,15 @@ def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment]) - allowed_model_names = proxy_server._get_v1_model_info_allowed_model_names( + allowed_model_names = await proxy_server._get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=UserAPIKeyAuth( api_key="sk-test", models=[SpecialModelNames.all_proxy_models.value], ), llm_router=router, + prisma_client=None, + user_api_key_cache=None, + proxy_logging_obj=None, ) result = proxy_server._filter_v1_model_info_deployments( @@ -597,6 +600,167 @@ def test_model_info_team_key_cannot_see_other_teams_byok_model( assert _BYOK_INTERNAL_NAME not in surfaced_names +_ACCESS_GROUP_TEAM_ID = "team-with-access-group" +_ACCESS_GROUP_ID = "ag-group-models" + + +@pytest.fixture +def access_group_team(monkeypatch, mock_prisma): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + router: Final = litellm.Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-unused"}} + for name in ("model-direct", "model-group", "model-other") + ] + ) + cache: Final = UserApiKeyCache() + + def seed_team(models: list[str]) -> None: + cache.set_cache( + key=f"team_id:{_ACCESS_GROUP_TEAM_ID}", + value=LiteLLM_TeamTableCachedObj( + team_id=_ACCESS_GROUP_TEAM_ID, models=models, access_group_ids=[_ACCESS_GROUP_ID] + ), + model_type=LiteLLM_TeamTableCachedObj, + ) + + cache.set_cache( + key=f"access_group_id:{_ACCESS_GROUP_ID}", + value=LiteLLM_AccessGroupTable( + access_group_id=_ACCESS_GROUP_ID, + access_group_name="group-models", + access_model_names=["model-group"], + assigned_team_ids=[_ACCESS_GROUP_TEAM_ID], + ), + model_type=LiteLLM_AccessGroupTable, + ) + mock_prisma.db.litellm_usertable.find_unique.return_value = None + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + return seed_team + + +@pytest.mark.parametrize( + ("path", "name_field"), + [ + ("/v1/model/info", "model_name"), + ("/model/info", "model_name"), + ("/v1/models", "id"), + ("/model_group/info", "model_group"), + ], +) +@pytest.mark.parametrize( + ("key_models", "team_models", "expected_names"), + [ + (["all-team-models"], ["no-default-models"], ["model-group"]), + ([], ["model-direct"], ["model-direct", "model-group"]), + (["model-direct"], ["model-direct"], ["model-direct"]), + ], + ids=["team-key-gets-access-group-models", "team-models-plus-access-group", "key-own-list-stays-narrow"], +) +def test_model_listing_routes_list_team_access_group_models_the_key_can_call( + client, auth_as, access_group_team, path, name_field, key_models, team_models, expected_names +): + from litellm.proxy._types import LitellmUserRoles + + access_group_team(team_models) + with auth_as( + role=LitellmUserRoles.INTERNAL_USER, + user_id=None, + team_id=_ACCESS_GROUP_TEAM_ID, + models=key_models, + team_models=team_models, + ): + response: Final = client.get(path) + + assert response.status_code == 200, response.text + assert sorted(model[name_field] for model in response.json()["data"]) == expected_names + + +def test_v1_model_info_unrestricted_admin_key_still_lists_team_byok_rows( + client, auth_as, byok_team_router, mock_prisma, monkeypatch +): + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + with auth_as(): + response: Final = client.get("/v1/model/info") + + assert response.status_code == 200, response.text + assert [model["model_name"] for model in response.json()["data"]] == [_BYOK_PUBLIC_NAME] + + +@pytest.mark.parametrize( + ("key_models", "expected_names"), + [([], ["model-direct"]), (["model-direct", "model-other"], ["model-direct"])], + ids=["no-key-list", "own-key-list"], +) +def test_v1_model_info_scopes_a_personal_key_to_its_owner_teams( + client, auth_as, access_group_team, monkeypatch, key_models, expected_names +): + from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles + from litellm.proxy.auth import auth_checks + + async def owner(**_: object) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id="owner", teams=[_ACCESS_GROUP_TEAM_ID]) + + async def owner_team(team_id: str, **_: object) -> LiteLLM_TeamTableCachedObj: + return LiteLLM_TeamTableCachedObj(team_id=team_id, models=["model-direct"]) + + async def no_membership(**_: object) -> None: + return None + + access_group_team(["model-direct"]) + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(auth_checks, "get_user_object", owner) + monkeypatch.setattr(auth_checks, "get_team_object", owner_team) + monkeypatch.setattr(auth_checks, "get_team_membership", no_membership) + with auth_as(role=LitellmUserRoles.INTERNAL_USER, user_id="owner", models=key_models): + response: Final = client.get("/v1/model/info") + + assert response.status_code == 200, response.text + assert sorted(model["model_name"] for model in response.json()["data"]) == expected_names + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("ceiling_models", "expected"), + [(frozenset({"model-direct"}), frozenset({"model-direct"})), (None, None)], + ids=["capped-agent", "agent-without-access-groups"], +) +async def test_v1_model_info_caps_an_agent_key_with_no_model_list_to_its_access_groups( + access_group_team, ceiling_models, expected +): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + async def ceiling(agent_id: str) -> AgentAccessGroupCeiling | None: + assert agent_id == "agent-1" + if ceiling_models is None: + return None + return AgentAccessGroupCeiling( + access_group_ids=("agent-group",), + models=ceiling_models, + mcp_server_ids=frozenset(), + agent_ids=frozenset(), + ) + + allowed: Final = await proxy_server._get_v1_model_info_allowed_model_names_with_access_groups( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-agent", agent_id="agent-1", models=[]), + llm_router=proxy_server.llm_router, + prisma_client=None, + user_api_key_cache=None, + proxy_logging_obj=None, + resolve_agent_ceiling=ceiling, + ) + + assert allowed == expected + + # --------------------------------------------------------------------------- # GET /model_group/info # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/utils/helpers/test_model_access.py b/tests/unit/proxy/utils/helpers/test_model_access.py index 7f77f938323..ecf1810f34a 100644 --- a/tests/unit/proxy/utils/helpers/test_model_access.py +++ b/tests/unit/proxy/utils/helpers/test_model_access.py @@ -448,6 +448,64 @@ async def test_get_available_models_for_user_resolves_team_access_group_models( assert sorted(result) == ["model-a", "model-b"] +@pytest.mark.asyncio +async def test_get_available_models_for_user_lists_team_access_group_models_only_where_auth_grants_them( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy._types import Member + + teams = { + "team-1": LiteLLM_TeamTableCachedObj(team_id="team-1", models=["model-a"], access_group_ids=["ag-1"]), + "team-2": LiteLLM_TeamTableCachedObj( + team_id="team-2", models=["model-c"], members_with_roles=[Member(user_id="user-1", role="user")] + ), + } + + async def _get_team_object(team_id: str, **_kwargs): + return teams[team_id] + + async def _get_access_object(**_kwargs): + return LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="team-group", + access_model_names=["model-b"], + assigned_team_ids=["team-1"], + ) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + + async def _list_for(key_models: list[str], config: dict[str, list[str]] | None = None, team_id: str | None = None): + return sorted( + await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=key_models, + team_models=["model-a"], + config=config or {}, + ), + llm_router=_router_with_models(["model-a", "model-b", "model-c"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + team_id=team_id, + ) + ) + + assert await _list_for(["model-a"]) == ["model-a"] + assert await _list_for(["model-a"], team_id="team-1") == ["model-a"] + assert await _list_for(["model-a"], team_id="team-2") == ["model-c"] + assert await _list_for(["model-a"], config={"model_list": []}) == ["model-a", "model-b"] + assert await _list_for(["all-team-models"]) == ["model-a", "model-b"] + assert await _list_for([]) == ["model-a", "model-b"] + + @pytest.mark.asyncio async def test_get_available_models_for_user_without_access_groups_grants_nothing( monkeypatch,