scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 138

fix(proxy): list team access group models in /v1/model/info and keep key model lists narrow (#32599)

Browse source 
* 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>
This commit is contained in:
L 2026-10-16 01:23:03 +08:00 • committed by GitHub
parent 3740d48844
commit bdd80781f2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 297 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
litellm/proxy/proxy_server.py
View file

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

33
litellm/proxy/utils.py
View file

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

168
tests/unit/proxy/proxy_server/test_routes_model_info.py
View file

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

58
tests/unit/proxy/utils/helpers/test_model_access.py
View file

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