Merge pull request #23236 from BerriAI/litellm_/flamboyant-leakey

[Fix] All Proxy Models Not Including Model Access Groups in Key Creation

litellm/proxy/auth/model_checks.py

@@ -108,16 +108,23 @@ def get_key_models(
     """
     all_models: List[str] = []
     if len(user_api_key_dict.models) > 0:
-        all_models = user_api_key_dict.models
+        all_models = list(user_api_key_dict.models)  # copy to avoid mutating cached objects
         if SpecialModelNames.all_team_models.value in all_models:
-            all_models = user_api_key_dict.team_models
+            all_models = list(user_api_key_dict.team_models)  # copy to avoid mutating cached objects
         if SpecialModelNames.all_proxy_models.value in all_models:
-            all_models = proxy_model_list
+            all_models = list(proxy_model_list)  # copy to avoid mutating caller's list
+            if include_model_access_groups:
+                all_models.extend(model_access_groups.keys())
 
     all_models = _get_models_from_access_groups(
-        model_access_groups=model_access_groups, all_models=all_models
+        model_access_groups=model_access_groups,
+        all_models=all_models,
+        include_model_access_groups=include_model_access_groups,
     )
 
+    # deduplicate while preserving order
+    all_models = list(dict.fromkeys(all_models))
+
     verbose_proxy_logger.debug("ALL KEY MODELS - {}".format(len(all_models)))
     return all_models
 
@@ -141,8 +148,8 @@ def get_team_models(
             all_models_set.update(team_models)
         if SpecialModelNames.all_proxy_models.value in all_models_set:
             all_models_set.update(proxy_model_list)
-
-    all_models = list(all_models_set)
+            if include_model_access_groups:
+                all_models_set.update(model_access_groups.keys())
 
     all_models = _get_models_from_access_groups(
         model_access_groups=model_access_groups,
@@ -150,6 +157,9 @@ def get_team_models(
         include_model_access_groups=include_model_access_groups,
     )
 
+    # deduplicate while preserving order
+    all_models = list(dict.fromkeys(all_models))
+
     verbose_proxy_logger.debug("ALL TEAM MODELS - {}".format(len(all_models)))
     return all_models
 

litellm/proxy/management_endpoints/team_endpoints.py

@@ -2827,21 +2827,6 @@ async def validate_membership(
     )
 
 
-def _unfurl_all_proxy_models(
-    team_info: LiteLLM_TeamTable, llm_router: Router
-) -> LiteLLM_TeamTable:
-    if (
-        SpecialModelNames.all_proxy_models.value in team_info.models
-        and llm_router is not None
-    ):
-        team_models: set[str] = set()  # make set to avoid duplicates
-        for model in team_info.models:
-            if model != SpecialModelNames.all_proxy_models.value:
-                team_models.add(model)
-        for model in llm_router.get_model_names():
-            team_models.add(model)
-        team_info.models = list(team_models)
-    return team_info
 
 
 async def _add_team_member_budget_table(
@@ -2972,9 +2957,6 @@ async def team_info(
                 team_info_response_object=_team_info,
             )
 
-        # ## UNFURL 'all-proxy-models' into the team_info.models list ##
-        # if llm_router is not None:
-        #     _team_info = _unfurl_all_proxy_models(_team_info, llm_router)
         response_object = TeamInfoResponseObject(
             team_id=team_id,
             team_info=_team_info,

tests/test_litellm/proxy/auth/test_model_checks.py

@@ -21,6 +21,140 @@ def test_get_team_models_for_all_models_and_team_only_models():
     assert set(result) == set(combined_models)
 
 
+def test_get_team_models_all_proxy_models_includes_access_groups():
+    """
+    When a team has 'all-proxy-models' and include_model_access_groups=True,
+    the result should include model access group names (e.g. 'claude-model-group')
+    in addition to individual model names.
+    """
+    from litellm.proxy.auth.model_checks import get_team_models
+
+    team_models = ["all-proxy-models"]
+    proxy_model_list = ["model1", "model2"]
+    model_access_groups = {
+        "group-a": ["model1"],
+        "group-b": ["model2"],
+    }
+
+    result = get_team_models(
+        team_models, proxy_model_list, model_access_groups, include_model_access_groups=True
+    )
+    assert "group-a" in result
+    assert "group-b" in result
+    assert "model1" in result
+    assert "model2" in result
+    assert len(result) == len(set(result)), "result should have no duplicates"
+
+
+def test_get_team_models_all_proxy_models_without_include_flag():
+    """
+    When include_model_access_groups=False, access group names should NOT
+    appear in the result even with 'all-proxy-models'.
+    """
+    from litellm.proxy.auth.model_checks import get_team_models
+
+    team_models = ["all-proxy-models"]
+    proxy_model_list = ["model1", "model2"]
+    model_access_groups = {
+        "group-a": ["model1"],
+        "group-b": ["model2"],
+    }
+
+    result = get_team_models(
+        team_models, proxy_model_list, model_access_groups, include_model_access_groups=False
+    )
+    assert "group-a" not in result
+    assert "group-b" not in result
+    assert "model1" in result
+    assert "model2" in result
+
+
+def test_get_key_models_all_proxy_models_includes_access_groups():
+    """
+    When a key has 'all-proxy-models' and include_model_access_groups=True,
+    the result should include model access group names.
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.auth.model_checks import get_key_models
+
+    user_api_key_dict = UserAPIKeyAuth(
+        models=["all-proxy-models"],
+        api_key="test-key",
+    )
+    proxy_model_list = ["model1", "model2"]
+    model_access_groups = {
+        "group-a": ["model1"],
+    }
+
+    result = 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=True,
+    )
+    assert "group-a" in result
+    assert "model1" in result
+    assert "model2" in result
+    assert len(result) == len(set(result)), "result should have no duplicates"
+
+
+def test_get_key_models_passes_include_model_access_groups():
+    """
+    When a key explicitly has an access group name in its models list and
+    include_model_access_groups=True, the group name should be retained
+    (not stripped by _get_models_from_access_groups).
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.auth.model_checks import get_key_models
+
+    user_api_key_dict = UserAPIKeyAuth(
+        models=["group-a"],
+        api_key="test-key",
+    )
+    proxy_model_list = ["model1", "model2"]
+    model_access_groups = {
+        "group-a": ["model1", "model2"],
+    }
+
+    result = 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=True,
+    )
+    assert "group-a" in result
+    assert "model1" in result
+    assert "model2" in result
+
+
+def test_get_key_models_does_not_mutate_input():
+    """
+    get_key_models must not mutate user_api_key_dict.models in-place.
+    _get_models_from_access_groups uses .pop()/.extend() which would corrupt
+    cached UserAPIKeyAuth objects if all_models were an alias instead of a copy.
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.auth.model_checks import get_key_models
+
+    original_models = ["group-a", "extra-model"]
+    user_api_key_dict = UserAPIKeyAuth(
+        models=list(original_models),  # give it a list
+        api_key="test-key",
+    )
+    model_access_groups = {
+        "group-a": ["model1", "model2"],
+    }
+
+    _ = get_key_models(
+        user_api_key_dict=user_api_key_dict,
+        proxy_model_list=["model1", "model2"],
+        model_access_groups=model_access_groups,
+        include_model_access_groups=False,
+    )
+    # The original models list on the auth object must be unchanged
+    assert user_api_key_dict.models == original_models
+
+
 @pytest.mark.parametrize(
     "key_models,team_models,proxy_model_list,model_list,expected",
     [