diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4a6856b6d14..28fbe8a7ddd 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -45,6 +45,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import can_team_access_model +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import PrismaClient, ProxyLogging from .auth_checks import ( @@ -1404,6 +1405,52 @@ class JWTAuthManager: ) return None + @staticmethod + async def _attach_team_from_header_for_admin( + admin_result: JWTAuthBuilderResult, + route: str, + request_headers: Optional[dict], + jwt_handler: JWTHandler, + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + parent_otel_span: Optional[Span], + proxy_logging_obj: ProxyLogging, + ) -> None: + """Attach team context from x-litellm-team-id to an admin result. + + Only applies on LLM API routes so team TPM/RPM limits and attribution + are enforced when admins act on behalf of a team. Admin management + routes ignore the header to preserve pre-existing bypass behavior. + """ + header_team_id = ( + request_headers.get("x-litellm-team-id") if request_headers else None + ) + if not header_team_id or not RouteChecks.is_llm_api_route(route=route): + return + try: + team_object = await get_team_object( + team_id=header_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, + ) + except Exception as e: + # Fall back to pre-PR admin behavior: honor the admin's + # authorization but skip team attribution/limits for this + # request. Log so operators can find the misconfigured caller. + verbose_proxy_logger.warning( + "admin x-litellm-team-id=%r on route=%s could not be resolved (%s); " + "proceeding with admin access, team context NOT attached.", + header_team_id, + route, + e, + ) + return + admin_result["team_id"] = header_team_id + admin_result["team_object"] = team_object + @staticmethod async def auth_builder( api_key: str, @@ -1493,6 +1540,16 @@ class JWTAuthManager: jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token ) if admin_result: + await JWTAuthManager._attach_team_from_header_for_admin( + admin_result=admin_result, + route=route, + request_headers=request_headers, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) return admin_result # Get team with model access diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9779a07b97b..edb1efc9706 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -810,6 +810,19 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_object is not None else None ), + team_tpm_limit=( + team_object.tpm_limit + if team_object is not None + else None + ), + team_rpm_limit=( + team_object.rpm_limit + if team_object is not None + else None + ), + team_models=( + team_object.models if team_object is not None else [] + ), team_metadata=( team_object.metadata if team_object is not None diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index fd293f5c256..cdb9ae4d9ab 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1498,6 +1498,174 @@ async def test_auth_builder_uses_team_from_header_e2e(): assert result["team_object"] == team_object +@pytest.mark.asyncio +async def test_auth_builder_admin_on_llm_route_honors_team_header(): + """JWT proxy_admin + x-litellm-team-id on an LLM API route -> team context is + attached to the admin result so team TPM/RPM limits and attribution apply.""" + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + team_ids_jwt_field="groups", + user_id_jwt_field="sub", + admin_allowed_routes=[ + "management_routes", + "info_routes", + "openai_routes", + ], + ), + ) + + team_object = LiteLLM_TeamTable(team_id="team-low", tpm_limit=100, rpm_limit=2) + + with ( + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), + patch.object(jwt_handler, "is_admin", return_value=True), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock + ) as mock_get_team, + ): + mock_auth_jwt.return_value = { + "sub": "admin-user", + "scope": "", + "groups": [], + } + mock_get_team.return_value = team_object + + result = await JWTAuthManager.auth_builder( + api_key="jwt-token", + jwt_handler=jwt_handler, + request_data={"model": "gpt-4"}, + general_settings={}, + route="/chat/completions", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), + request_headers={"x-litellm-team-id": "team-low"}, + ) + + assert result["is_proxy_admin"] is True + assert result["team_id"] == "team-low" + assert result["team_object"] == team_object + mock_get_team.assert_called_once() + + +@pytest.mark.asyncio +async def test_auth_builder_admin_on_mgmt_route_ignores_team_header(): + """JWT proxy_admin + x-litellm-team-id on an admin management route -> header + is ignored; no team fetch. Preserves pre-existing bypass behavior and avoids + phantom team creation when team_id_upsert is enabled.""" + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + team_ids_jwt_field="groups", + user_id_jwt_field="sub", + team_id_upsert=True, + ), + ) + + with ( + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), + patch.object(jwt_handler, "is_admin", return_value=True), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock + ) as mock_get_team, + ): + mock_auth_jwt.return_value = { + "sub": "admin-user", + "scope": "", + "groups": [], + } + + result = await JWTAuthManager.auth_builder( + api_key="jwt-token", + jwt_handler=jwt_handler, + request_data={}, + general_settings={}, + route="/user/info", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), + request_headers={"x-litellm-team-id": "totally-made-up-team"}, + ) + + assert result["is_proxy_admin"] is True + assert result["team_id"] is None + assert result["team_object"] is None + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_auth_builder_admin_on_llm_route_without_header_unchanged(): + """JWT proxy_admin on an LLM API route without x-litellm-team-id -> no team + context (team limits not applied, admin keeps unrestricted access).""" + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + team_ids_jwt_field="groups", + user_id_jwt_field="sub", + admin_allowed_routes=[ + "management_routes", + "info_routes", + "openai_routes", + ], + ), + ) + + with ( + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), + patch.object(jwt_handler, "is_admin", return_value=True), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock + ) as mock_get_team, + ): + mock_auth_jwt.return_value = { + "sub": "admin-user", + "scope": "", + "groups": [], + } + + result = await JWTAuthManager.auth_builder( + api_key="jwt-token", + jwt_handler=jwt_handler, + request_data={"model": "gpt-4"}, + general_settings={}, + route="/chat/completions", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), + request_headers={}, + ) + + assert result["is_proxy_admin"] is True + assert result["team_id"] is None + assert result["team_object"] is None + mock_get_team.assert_not_called() + + @pytest.mark.asyncio async def test_get_team_alias_with_nested_fields(): """