From f60a9cf908b30b75189d32f3fed7d30005df068a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 9 Aug 2025 16:14:51 -0700 Subject: [PATCH] [Bug]: Fix JWTs access not working with model groups (#13474) * fix can_team_access_model * test_find_team_with_model_access_model_group --- litellm/proxy/auth/handle_jwt.py | 3 +- .../proxy/auth/test_handle_jwt.py | 68 ++++++++++++++++--- 2 files changed, 62 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 60529f1e2f4..b8b56833519 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -763,6 +763,7 @@ class JWTAuthManager: proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]: """Find first team with access to the requested model""" + from litellm.proxy.proxy_server import llm_router if not team_ids: if jwt_handler.litellm_jwtauth.enforce_team_based_model_access: @@ -789,7 +790,7 @@ class JWTAuthManager: or can_team_access_model( model=requested_model, team_object=team_object, - llm_router=None, + llm_router=llm_router, team_model_aliases=None, ) ): diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 5627b9aaf90..17efbdcf4b3 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -459,9 +459,9 @@ async def test_nested_jwt_field_access(): 2. Backward compatibility is maintained for flat field names 3. Missing nested paths return appropriate defaults """ - from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy._types import LiteLLM_JWTAuth - + from litellm.proxy.auth.handle_jwt import JWTHandler + # Create JWT handler jwt_handler = JWTHandler() @@ -536,7 +536,7 @@ async def test_nested_jwt_field_access(): assert jwt_handler.get_org_id(flat_token, None) == "org456" # Test 5: object_id_jwt_field with nested access (requires role_mappings) - from litellm.proxy._types import RoleMapping, LitellmUserRoles + from litellm.proxy._types import LitellmUserRoles, RoleMapping jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="profile.object_id", role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)] @@ -588,9 +588,9 @@ async def test_nested_jwt_field_missing_paths(): 2. Partial paths that exist but don't have the final key return defaults 3. team_id_default fallback works with nested fields """ - from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy._types import LiteLLM_JWTAuth - + from litellm.proxy.auth.handle_jwt import JWTHandler + # Create JWT handler jwt_handler = JWTHandler() @@ -626,7 +626,7 @@ async def test_nested_jwt_field_missing_paths(): assert jwt_handler.get_org_id(incomplete_token, "default_org") == "default_org" # Test 5: Missing profile.object_id should return default (requires role_mappings) - from litellm.proxy._types import RoleMapping, LitellmUserRoles + from litellm.proxy._types import LitellmUserRoles, RoleMapping jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="profile.object_id", role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)] @@ -663,9 +663,9 @@ async def test_metadata_prefix_handling_in_nested_fields(): The get_nested_value function should remove metadata. prefix before traversing """ - from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy._types import LiteLLM_JWTAuth - + from litellm.proxy.auth.handle_jwt import JWTHandler + # Create JWT handler jwt_handler = JWTHandler() @@ -685,3 +685,55 @@ async def test_metadata_prefix_handling_in_nested_fields(): # Test 2: user.sub should work normally without metadata prefix jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="sub") assert jwt_handler.get_user_id(token, None) == "u123" + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_model_group(monkeypatch): + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "gpt-4o-mini"}, + "model_info": {"access_groups": ["test-group"]}, + } + ] + ) + import sys + import types + + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + team = LiteLLM_TeamTable(team_id="team-1", models=["test-group"]) + + async def mock_get_team_object(*args, **kwargs): # type: ignore + return team + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids={"team-1"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + assert team_id == "team-1" + assert team_obj.team_id == "team-1" \ No newline at end of file