diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8887d86570a..480daf448d8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5488,6 +5488,31 @@ async def can_key_call_model( - Exception: If token not allowed to call model """ key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) + key_access_group_ids: Final = valid_token.access_group_ids or [] + + if key_access_group_ids and not key_models: + models_from_groups: Final = await get_models_from_access_groups( + access_group_ids=key_access_group_ids, + prisma_client=prisma_client, + ) + if not models_from_groups: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=model), + internal_message=f"key access groups grant no models. Tried to access {model}", + type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type="key"), + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + return can_object_call_model( + model=model, + llm_router=llm_router, + models=models_from_groups, + team_model_aliases=valid_token.team_model_aliases, + team_id=valid_token.team_id, + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + object_type="key", + ) + try: return can_object_call_model( model=model, @@ -5500,17 +5525,16 @@ async def can_key_call_model( ) except ProxyException: # Fallback: check key's access_group_ids - key_access_group_ids: Final = valid_token.access_group_ids or [] if key_access_group_ids: - models_from_groups: Final = await get_models_from_access_groups( + models_from_groups_fallback: Final = await get_models_from_access_groups( access_group_ids=key_access_group_ids, prisma_client=prisma_client, ) - if models_from_groups: + if models_from_groups_fallback: return can_object_call_model( model=model, llm_router=llm_router, - models=models_from_groups, + models=models_from_groups_fallback, team_model_aliases=valid_token.team_model_aliases, team_id=valid_token.team_id, key_model_aliases=key_model_aliases_for_auth_check(valid_token), @@ -5710,11 +5734,37 @@ async def can_team_access_model( 1. First checks native team-level model permissions (current implementation) 2. If not allowed natively, falls back to access_group_ids on the team """ + team_models: Final = team_object.models if team_object else [] + team_access_group_ids: Final = (team_object.access_group_ids or []) if team_object else [] + + if team_access_group_ids and not team_models: + models_from_groups: Final = await get_models_from_access_groups( + access_group_ids=team_access_group_ids, + prisma_client=prisma_client, + ) + if not models_from_groups: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=model), + internal_message=f"team access groups grant no models. Tried to access {model}", + type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type="team"), + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + return can_object_call_model( + model=model, + llm_router=llm_router, + models=models_from_groups, + team_model_aliases=team_model_aliases, + team_id=team_object.team_id if team_object else None, + key_model_aliases=key_model_aliases, + object_type="team", + ) + try: return can_object_call_model( model=model, llm_router=llm_router, - models=team_object.models if team_object else [], + models=team_models, team_model_aliases=team_model_aliases, team_id=team_object.team_id if team_object else None, key_model_aliases=key_model_aliases, @@ -5722,17 +5772,16 @@ async def can_team_access_model( ) except ProxyException: # Fallback: check team's access_group_ids - team_access_group_ids: Final = (team_object.access_group_ids or []) if team_object else [] if team_access_group_ids: - models_from_groups: Final = await get_models_from_access_groups( + models_from_groups_fallback: Final = await get_models_from_access_groups( access_group_ids=team_access_group_ids, prisma_client=prisma_client, ) - if models_from_groups: + if models_from_groups_fallback: return can_object_call_model( model=model, llm_router=llm_router, - models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])), + models=list(dict.fromkeys([*team_models, *models_from_groups_fallback])), team_model_aliases=team_model_aliases, team_id=team_object.team_id if team_object else None, key_model_aliases=key_model_aliases, diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..c60f141e2d2 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -1565,6 +1565,66 @@ async def test_can_team_access_model_access_group_ids_denied(): ) +@pytest.mark.asyncio +async def test_can_team_access_model_empty_models_access_group_denied(): + """Test can_team_access_model denies access when team has empty models and model is not in access groups.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import ModelAccessDeniedProxyException + from litellm.proxy.auth.auth_checks import can_team_access_model + + team_object = LiteLLM_TeamTable( + team_id="test-team", + models=[], + access_group_ids=["ag-with-gpt4"], + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.get_models_from_access_groups", + new_callable=AsyncMock, + return_value=["gpt-4"], + ), + pytest.raises(ModelAccessDeniedProxyException), + ): + await can_team_access_model( + model="claude-3", + team_object=team_object, + llm_router=None, + team_model_aliases=None, + ) + + +@pytest.mark.asyncio +async def test_can_team_access_model_empty_models_empty_access_group_denied(): + """Test can_team_access_model denies access when team has empty models and access groups have no models.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import ModelAccessDeniedProxyException + from litellm.proxy.auth.auth_checks import can_team_access_model + + team_object = LiteLLM_TeamTable( + team_id="test-team", + models=[], + access_group_ids=["ag-empty"], + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.get_models_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + pytest.raises(ModelAccessDeniedProxyException), + ): + await can_team_access_model( + model="gpt-4", + team_object=team_object, + llm_router=None, + team_model_aliases=None, + ) + + # --------------------------------------------------------------------------- # can_key_call_model with access_group_ids fallback # --------------------------------------------------------------------------- @@ -1604,6 +1664,86 @@ async def test_can_key_call_model_via_access_group_ids(): ) +@pytest.mark.asyncio +async def test_can_key_call_model_access_group_ids_denied_for_unlisted_model(): + """Test can_key_call_model denies access when key has access_group_ids but requested model is not in access group.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import ModelAccessDeniedProxyException + from litellm.proxy.auth.auth_checks import can_key_call_model + + user_api_key_object = UserAPIKeyAuth( + token="test-token", + models=[], + access_group_ids=["ag-with-gpt4"], + ) + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4", "api_key": "test"}, + }, + { + "model_name": "claude-3", + "litellm_params": {"model": "anthropic/claude-3", "api_key": "test"}, + }, + ] + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.get_models_from_access_groups", + new_callable=AsyncMock, + return_value=["gpt-4"], + ), + pytest.raises(ModelAccessDeniedProxyException), + ): + await can_key_call_model( + model="claude-3", + llm_model_list=[], + valid_token=user_api_key_object, + llm_router=router, + ) + + +@pytest.mark.asyncio +async def test_can_key_call_model_access_group_ids_empty_models_denied(): + """Test can_key_call_model denies access when key has access_group_ids but access groups contain no models.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import ModelAccessDeniedProxyException + from litellm.proxy.auth.auth_checks import can_key_call_model + + user_api_key_object = UserAPIKeyAuth( + token="test-token", + models=[], + access_group_ids=["ag-empty"], + ) + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4", "api_key": "test"}, + } + ] + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.get_models_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + pytest.raises(ModelAccessDeniedProxyException), + ): + await can_key_call_model( + model="gpt-4", + llm_model_list=[], + valid_token=user_api_key_object, + llm_router=router, + ) + + # --------------------------------------------------------------------------- # _key_access_group_grants_model (key access group overriding team restriction) # ---------------------------------------------------------------------------