diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1efc9611fe6..0bead9fb2f4 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3959,7 +3959,13 @@ def _check_model_access_helper( # Filter out models that are access_groups filtered_models: Final = [m for m in models if m not in access_groups] - if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases): + team_alias_target: Final = team_model_aliases.get(model) if team_model_aliases else None + if team_alias_target is not None and _check_model_access_helper( + model=team_alias_target, + llm_router=llm_router, + models=models, + team_id=team_id, + ): return True if _model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models): @@ -4045,24 +4051,6 @@ def _can_object_call_model( ) -def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool: - """ - Returns True if `model` being accessed is an alias of a team model - - - `model=gpt-4o` - - `team_model_aliases={"gpt-4o": "gpt-4o-team-1"}` - - returns True - - - `model=gp-4o` - - `team_model_aliases={"o-3": "o3-preview"}` - - returns False - """ - if team_model_aliases: - if model in team_model_aliases: - return True - return False - - def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]: """ Expand key model sentinels before auth checks. diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8777e24e209..c2a66eece87 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -403,6 +403,122 @@ async def test_can_key_call_model_all_team_models_no_team_id_is_unrestricted(): ) +_TEAM_ALIASES: Final = { + "fast": "openai/gpt-4.1-mini", + "smart": "openai/gpt-4.1", + "vision": "openai/gpt-4o", +} + + +def _aliased_key(models: list[str]) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-key", + team_id="team-1", + models=models, + team_models=["openai/gpt-4.1-mini", "openai/gpt-4.1"], + team_model_aliases=_TEAM_ALIASES, + ) + + +@pytest.mark.asyncio +async def test_can_key_call_model_team_alias_target_outside_key_models_denied(): + """A team alias whose target is not in key.models is denied.""" + from litellm.proxy.auth.auth_checks import can_key_call_model + + valid_token: Final = _aliased_key(models=["openai/gpt-4.1-mini"]) + + for alias in ("smart", "vision"): + with pytest.raises(ProxyException) as exc_info: + await can_key_call_model(model=alias, llm_model_list=None, valid_token=valid_token, llm_router=None) + assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied + assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN) + + +@pytest.mark.asyncio +async def test_can_key_call_model_team_alias_target_in_key_models_allowed(): + """A team alias whose target is in key.models is allowed, and a plain model name still resolves as before.""" + from litellm.proxy.auth.auth_checks import can_key_call_model + + valid_token: Final = _aliased_key(models=["openai/gpt-4.1-mini"]) + + for model in ("fast", "openai/gpt-4.1-mini"): + assert await can_key_call_model(model=model, llm_model_list=None, valid_token=valid_token, llm_router=None) + + with pytest.raises(ProxyException) as exc_info: + await can_key_call_model(model="openai/gpt-4.1", llm_model_list=None, valid_token=valid_token, llm_router=None) + assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied + + +@pytest.mark.asyncio +async def test_can_key_call_model_team_alias_unrestricted_key_allowed(): + """A key with no model restriction, or the all-proxy-models sentinel, can call any team alias.""" + from litellm.proxy._types import SpecialModelNames + from litellm.proxy.auth.auth_checks import can_key_call_model + + for models in ([], [SpecialModelNames.all_proxy_models.value]): + valid_token: Final = _aliased_key(models=models) + for alias in _TEAM_ALIASES: + assert await can_key_call_model(model=alias, llm_model_list=None, valid_token=valid_token, llm_router=None) + + +@pytest.mark.asyncio +async def test_can_key_call_model_team_alias_all_team_models_sentinel(): + """With the all-team-models sentinel the alias target must be in team_models.""" + from litellm.proxy._types import SpecialModelNames + from litellm.proxy.auth.auth_checks import can_key_call_model + + valid_token: Final = _aliased_key(models=[SpecialModelNames.all_team_models.value]) + + for alias in ("fast", "smart"): + assert await can_key_call_model(model=alias, llm_model_list=None, valid_token=valid_token, llm_router=None) + + with pytest.raises(ProxyException) as exc_info: + await can_key_call_model(model="vision", llm_model_list=None, valid_token=valid_token, llm_router=None) + assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied + + +@pytest.mark.asyncio +async def test_can_team_access_model_team_alias_checks_target(): + """A team alias is allowed only when its target is in team.models.""" + from litellm.proxy.auth.auth_checks import can_team_access_model + + team_object: Final = LiteLLM_TeamTable(team_id="team-1", models=["openai/gpt-4.1-mini"]) + + assert await can_team_access_model( + model="fast", team_object=team_object, llm_router=None, team_model_aliases=_TEAM_ALIASES + ) + + with pytest.raises(ProxyException) as exc_info: + await can_team_access_model( + model="smart", team_object=team_object, llm_router=None, team_model_aliases=_TEAM_ALIASES + ) + assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied + + +def test_check_model_access_helper_team_alias_target_matches_wildcard(): + """A team alias target is matched against wildcard entries in the allowlist.""" + from litellm.proxy.auth.auth_checks import _check_model_access_helper + + assert _check_model_access_helper( + model="fast", llm_router=None, models=["openai/*"], team_model_aliases=_TEAM_ALIASES, team_id="team-1" + ) + assert not _check_model_access_helper( + model="fast", llm_router=None, models=["anthropic/*"], team_model_aliases=_TEAM_ALIASES, team_id="team-1" + ) + + +def test_check_model_access_helper_self_referential_alias_terminates(): + """An alias that points at its own name is decided by the allowlist alone.""" + from litellm.proxy.auth.auth_checks import _check_model_access_helper + + self_alias: Final = {"same": "same"} + + assert _check_model_access_helper(model="same", llm_router=None, models=["same"], team_model_aliases=self_alias) + assert not _check_model_access_helper( + model="same", llm_router=None, models=["other"], team_model_aliases=self_alias + ) + + def test_resolve_key_models_teamless_all_team_models_returns_empty(): """_resolve_key_models_for_auth_check must return [] for a teamless key with all-team-models, making it equivalent to an unscoped key (unrestricted