Merge pull request #41924 from BerriAI/litellm_role_permissions_normalization

fix(proxy): parse role_permissions where it is read

litellm/proxy/auth/auth_checks.py

@@ -15,10 +15,10 @@ import re
 import time
 from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol
 
 from fastapi import HTTPException, Request, status
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
@@ -2412,22 +2412,22 @@ def _update_last_db_access_time(key: str, value: object | None, last_db_access_t
     last_db_access_time[key] = (value, time.time())
 
 
+ROLE_BASED_PERMISSIONS_ADAPTER: Final[TypeAdapter[list[RoleBasedPermissions]]] = TypeAdapter(list[RoleBasedPermissions])
+
+
 def _get_role_based_permissions(
     rbac_role: RBAC_ROLES,
-    general_settings: dict,
+    general_settings: Mapping[str, object],
     key: Literal["models", "routes"],
 ) -> list[str] | None:
     """
     Get the role based permissions from the general settings.
     """
-    role_based_permissions: Final = cast(
-        list[RoleBasedPermissions] | None,
-        general_settings.get("role_permissions", []),
-    )
-    if role_based_permissions is None:
+    configured: Final = general_settings.get("role_permissions")
+    if configured is None:
         return None
 
-    for role_based_permission in role_based_permissions:
+    for role_based_permission in ROLE_BASED_PERMISSIONS_ADAPTER.validate_python(configured):
         if role_based_permission.role == rbac_role:
             return role_based_permission.models if key == "models" else role_based_permission.routes
 
@@ -2436,7 +2436,7 @@ def _get_role_based_permissions(
 
 def get_role_based_models(
     rbac_role: RBAC_ROLES,
-    general_settings: dict,
+    general_settings: Mapping[str, object],
 ) -> list[str] | None:
     """
     Get the models allowed for a user role.
@@ -2453,7 +2453,7 @@ def get_role_based_models(
 
 def get_role_based_routes(
     rbac_role: RBAC_ROLES,
-    general_settings: dict,
+    general_settings: Mapping[str, object],
 ) -> list[str] | None:
     """
     Get the routes allowed for a user role.

litellm/proxy/proxy_server.py

@@ -111,7 +111,6 @@ from litellm.proxy._types import (
     PassThroughGenericEndpoint,
     ProxyErrorTypes,
     ProxyException,
-    RoleBasedPermissions,
     SpecialModelNames,
     SupportedDBObjectType,
     TeamDefaultSettings,
@@ -317,6 +316,7 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import (
     router as analytics_router,
 )
 from litellm.proxy.auth.auth_checks import (
+    ROLE_BASED_PERMISSIONS_ADAPTER,
     ExperimentalUIJWTToken,
     can_key_call_resolved_model,
     get_team_object,
@@ -6307,9 +6307,7 @@ class ProxyConfig:
             ### RBAC ###
             rbac_role_permissions: Final = general_settings.get("role_permissions", None)
             if rbac_role_permissions is not None:
-                general_settings["role_permissions"] = [  # validate role permissions
-                    RoleBasedPermissions(**role_permission) for role_permission in rbac_role_permissions
-                ]
+                ROLE_BASED_PERMISSIONS_ADAPTER.validate_python(rbac_role_permissions)
 
             ### SSRF URL VALIDATION SETTINGS ###
             _apply_ssrf_general_settings(general_settings)

tests/test_litellm/proxy/test_proxy_server.py

@@ -3237,6 +3237,80 @@ async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeyp
         litellm.max_budget = original_max_budget
 
 
+@pytest.mark.asyncio
+async def test_load_config_role_permissions_usable_by_jwt_auth(tmp_path):
+    from litellm.proxy.auth.auth_checks import get_role_based_models, get_role_based_routes
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    config_file: Final = tmp_path / "config.yaml"
+    config_file.write_text(
+        yaml.dump(
+            {
+                "model_list": [],
+                "general_settings": {
+                    "role_permissions": [
+                        {
+                            "role": "proxy_admin",
+                            "models": ["admin-only-model"],
+                            "routes": ["/v1/embeddings"],
+                        },
+                        {
+                            "role": "internal_user",
+                            "models": ["shared-model"],
+                            "routes": ["/v1/chat/completions"],
+                        },
+                    ]
+                },
+            }
+        )
+    )
+
+    _, _, settings = await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file))
+
+    assert get_role_based_models(rbac_role="internal_user", general_settings=settings) == ["shared-model"]
+    assert get_role_based_routes(rbac_role="internal_user", general_settings=settings) == ["/v1/chat/completions"]
+    assert get_role_based_models(rbac_role="proxy_admin", general_settings=settings) == ["admin-only-model"]
+    assert get_role_based_routes(rbac_role="proxy_admin", general_settings=settings) == ["/v1/embeddings"]
+    assert get_role_based_models(rbac_role="team", general_settings=settings) is None
+
+
+@pytest.mark.asyncio
+async def test_load_config_without_role_permissions_leaves_every_role_unrestricted(tmp_path):
+    from litellm.proxy.auth.auth_checks import get_role_based_models, get_role_based_routes
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    config_file: Final = tmp_path / "config.yaml"
+    config_file.write_text(
+        yaml.dump({"model_list": [], "general_settings": {"max_parallel_requests": 7}})
+    )
+
+    _, _, settings = await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file))
+
+    assert settings["max_parallel_requests"] == 7
+    assert get_role_based_models(rbac_role="internal_user", general_settings=settings) is None
+    assert get_role_based_routes(rbac_role="internal_user", general_settings=settings) is None
+
+
+@pytest.mark.asyncio
+async def test_load_config_rejects_malformed_role_permissions(tmp_path):
+    from pydantic import ValidationError
+
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    config_file: Final = tmp_path / "config.yaml"
+    config_file.write_text(
+        yaml.dump(
+            {
+                "model_list": [],
+                "general_settings": {"role_permissions": [{"role": "not_a_real_role", "models": ["gpt-4o"]}]},
+            }
+        )
+    )
+
+    with pytest.raises(ValidationError):
+        await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file))
+
+
 def test_max_ui_session_budget_default_is_one_dollar():
     """LIT-4662: the dashboard session budget default is a product decision; the
     old 0.25 default locked admins out of auto router Test Connection and the