mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
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:
parent
bd2b93c688
commit
6f8a3b8661
2 changed files with 153 additions and 18 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue