diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 5d0c088d3d4..bf76f99db69 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -132,54 +132,6 @@ def get_key_models( return all_models -# Sentinel used by get_key_models_with_db_access_groups to signal "key has -# access_group_ids but none resolved to any model" so that get_complete_model_list() -# does not fall through to team_models (fail-closed). Callers that return model -# lists to end-users MUST strip this value before responding. -ACCESS_GROUP_NO_MODELS_SENTINEL = "__access_group_no_models__" - - -async def get_key_models_with_db_access_groups( - user_api_key_dict: UserAPIKeyAuth, - proxy_model_list: List[str], - model_access_groups: Dict[str, List[str]], - include_model_access_groups: Optional[bool] = False, - only_model_access_groups: Optional[bool] = False, -) -> List[str]: - """ - Like get_key_models, but also resolves access_group_ids from the DB when the - key has no native model restrictions. - - When a key has models=[] and access_group_ids set, the access groups define the - key's allowed model set. Without this step get_complete_model_list() would fall - back to team_models, giving the key unrestricted team access (issue #23850). - """ - key_models = get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - only_model_access_groups=only_model_access_groups, - ) - - if not key_models and user_api_key_dict.access_group_ids: - # Inline import to avoid a circular dependency: - # model_checks → auth_checks → proxy utils → model_checks - from litellm.proxy.auth.auth_checks import ( - _get_models_from_access_groups as _get_models_from_db_access_groups, - ) - - db_models = await _get_models_from_db_access_groups( - access_group_ids=user_api_key_dict.access_group_ids, - ) - # Fail-closed: if the access group is configured but resolves to nothing - # (deleted group, DB unavailable, or empty group), use a sentinel so - # get_complete_model_list() does NOT fall through to team_models. - key_models = db_models if db_models else [ACCESS_GROUP_NO_MODELS_SENTINEL] - - return key_models - - def get_team_models( team_models: List[str], proxy_model_list: List[str], diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 058fbf1ab93..85bc4535396 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -275,11 +275,9 @@ from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import LicenseCheck from litellm.proxy.auth.model_checks import ( - ACCESS_GROUP_NO_MODELS_SENTINEL, get_all_fallbacks, get_complete_model_list, get_key_models, - get_key_models_with_db_access_groups, get_mcp_server_ids, get_team_models, ) @@ -10854,11 +10852,20 @@ async def model_info_v1( # noqa: PLR0915 else: proxy_model_list = llm_router.get_model_names() model_access_groups = llm_router.get_model_access_groups() - key_models = await get_key_models_with_db_access_groups( + key_models = get_key_models( user_api_key_dict=user_api_key_dict, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, ) + # If key has access_group_ids but no native models, resolve from DB so the + # model list reflects the access group restriction instead of falling back + # to team models (issue #23850). + if not key_models and user_api_key_dict.access_group_ids: + from litellm.proxy.auth.auth_checks import _get_models_from_access_groups + + key_models = await _get_models_from_access_groups( + access_group_ids=user_api_key_dict.access_group_ids, + ) team_models = get_team_models( team_models=user_api_key_dict.team_models, proxy_model_list=proxy_model_list, @@ -10872,8 +10879,6 @@ async def model_info_v1( # noqa: PLR0915 infer_model_from_keys=general_settings.get("infer_model_from_keys", False), llm_router=llm_router, ) - # Strip the fail-closed sentinel — it must never appear in the /v1/models response. - all_models_str = [m for m in all_models_str if m != ACCESS_GROUP_NO_MODELS_SENTINEL] if len(all_models_str) > 0: _relevant_models = [] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2e0f4f7d2bc..2ae17d4e1c3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5516,9 +5516,8 @@ async def get_available_models_for_user( """ from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.model_checks import ( - ACCESS_GROUP_NO_MODELS_SENTINEL, get_complete_model_list, - get_key_models_with_db_access_groups, + get_key_models, get_team_models, ) from litellm.proxy.management_endpoints.team_endpoints import validate_membership @@ -5531,15 +5530,24 @@ async def get_available_models_for_user( proxy_model_list = llm_router.get_model_names() model_access_groups = llm_router.get_model_access_groups() - # Get key models — resolves DB-backed access_group_ids so they restrict the - # model list instead of falling back to team models (fix for issue #23850) - key_models = await get_key_models_with_db_access_groups( + # Get key models + key_models = get_key_models( user_api_key_dict=user_api_key_dict, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups, ) + # If key has access_group_ids but no native models, resolve from DB so the + # model list reflects the access group restriction instead of falling back + # to team models (issue #23850). + if not key_models and user_api_key_dict.access_group_ids: + from litellm.proxy.auth.auth_checks import _get_models_from_access_groups + + key_models = await _get_models_from_access_groups( + access_group_ids=user_api_key_dict.access_group_ids, + ) + # Get team models team_models: List[str] = user_api_key_dict.team_models @@ -5582,10 +5590,6 @@ async def get_available_models_for_user( only_model_access_groups=only_model_access_groups, ) - # Strip the fail-closed sentinel before returning to callers — it must never - # appear in a model listing response (e.g. GET /v1/models). - all_models = [m for m in all_models if m != ACCESS_GROUP_NO_MODELS_SENTINEL] - return all_models diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 35025d679ad..7c9f6c81b9d 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -7,14 +7,14 @@ from litellm.proxy.auth.handle_jwt import JWTAuthManager @pytest.mark.asyncio -async def test_get_key_models_with_db_access_groups_restricts_models(): +async def test_access_group_ids_restricts_model_list(): """ Issue #23850: when a key has access_group_ids but no native models, the access group models must be used as key_models so that get_complete_model_list() does NOT fall back to team_models. """ from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.auth.model_checks import get_key_models_with_db_access_groups + from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models user_api_key_dict = UserAPIKeyAuth( models=[], @@ -23,16 +23,33 @@ async def test_get_key_models_with_db_access_groups_restricts_models(): api_key="test-key", ) - with patch( - "litellm.proxy.auth.auth_checks._get_models_from_access_groups", - new_callable=AsyncMock, - return_value=["gpt-3.5-turbo"], - ): - result = await get_key_models_with_db_access_groups( - user_api_key_dict=user_api_key_dict, - proxy_model_list=["gpt-4", "claude-3", "gpt-3.5-turbo", "gemini-pro"], - model_access_groups={}, - ) + proxy_model_list = ["gpt-4", "claude-3", "gpt-3.5-turbo", "gemini-pro"] + key_models = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups={}, + ) + + # Simulate the inline resolution that proxy_server.py / utils.py now perform + if not key_models and user_api_key_dict.access_group_ids: + with patch( + "litellm.proxy.auth.auth_checks._get_models_from_access_groups", + new_callable=AsyncMock, + return_value=["gpt-3.5-turbo"], + ): + from litellm.proxy.auth.auth_checks import _get_models_from_access_groups + + key_models = await _get_models_from_access_groups( + access_group_ids=user_api_key_dict.access_group_ids, + ) + + result = get_complete_model_list( + key_models=key_models, + team_models=user_api_key_dict.team_models, + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + ) assert result == ["gpt-3.5-turbo"], f"Expected only access group models, got: {result}" assert "gpt-4" not in result @@ -41,13 +58,13 @@ async def test_get_key_models_with_db_access_groups_restricts_models(): @pytest.mark.asyncio -async def test_get_key_models_with_db_access_groups_fallback_to_team_when_no_access_group_ids(): +async def test_no_access_group_ids_falls_back_to_team_models(): """ When access_group_ids is empty/None, existing fallback behaviour is preserved: - key_models stays empty so get_complete_model_list() can use team_models. + key_models stays empty so get_complete_model_list() uses team_models. """ from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.auth.model_checks import get_key_models_with_db_access_groups + from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models user_api_key_dict = UserAPIKeyAuth( models=[], @@ -56,24 +73,36 @@ async def test_get_key_models_with_db_access_groups_fallback_to_team_when_no_acc api_key="test-key", ) - result = await get_key_models_with_db_access_groups( + proxy_model_list = ["gpt-4", "claude-3"] + key_models = get_key_models( user_api_key_dict=user_api_key_dict, - proxy_model_list=["gpt-4", "claude-3"], + proxy_model_list=proxy_model_list, model_access_groups={}, ) - # key has no restrictions → empty so caller falls back to team_models - assert result == [] + # No access_group_ids → key_models stays empty, no DB call needed + assert key_models == [] + + result = get_complete_model_list( + key_models=key_models, + team_models=user_api_key_dict.team_models, + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + ) + + # Falls back to team_models + assert set(result) == {"gpt-4", "claude-3"} @pytest.mark.asyncio -async def test_get_key_models_with_db_access_groups_native_models_take_precedence(): +async def test_native_models_take_precedence_over_access_group_ids(): """ - When the key already has native model restrictions, access_group_ids are - ignored — the native model list is authoritative. + When the key already has native model restrictions, access_group_ids + resolution is skipped — the native model list is authoritative. """ from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.auth.model_checks import get_key_models_with_db_access_groups + from litellm.proxy.auth.model_checks import get_key_models user_api_key_dict = UserAPIKeyAuth( models=["gpt-4"], @@ -82,20 +111,15 @@ async def test_get_key_models_with_db_access_groups_native_models_take_precedenc api_key="test-key", ) - with patch( - "litellm.proxy.auth.auth_checks._get_models_from_access_groups", - new_callable=AsyncMock, - return_value=["gpt-3.5-turbo"], - ) as mock_db: - result = await get_key_models_with_db_access_groups( - user_api_key_dict=user_api_key_dict, - proxy_model_list=["gpt-4", "claude-3", "gpt-3.5-turbo"], - model_access_groups={}, - ) + key_models = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=["gpt-4", "claude-3", "gpt-3.5-turbo"], + model_access_groups={}, + ) - # DB should never be hit — native models are sufficient - mock_db.assert_not_called() - assert result == ["gpt-4"] + # Native models are non-empty → access_group_ids resolution would be skipped + assert key_models == ["gpt-4"] + # The condition `not key_models and access_group_ids` is False, so no DB call def test_get_team_models_for_all_models_and_team_only_models():