fix(proxy): scope model listing direct access to the calling key's grant

Direct access on /model/info was read from the user record alone, so an
unrestricted user calling with a key limited to a few models saw every non-team
deployment, including ones the key gets a 403 on. Resolve the key's grant the
same way and intersect the two.

Resolving a grant now also expands access groups, which the key path needs and
the user path was missing.

Claude-Session: https://claude.ai/code/session_01XL7LBFEew4wi8gphVCDq6n
This commit is contained in:
ryan-crabbe-berri 2026-08-29 16:54:02 -07:00
parent bd2b93c688
commit 6f8a3b8661
2 changed files with 153 additions and 18 deletions

View file

@ -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:

View file

@ -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,