Merge pull request #26438 from BerriAI/litellm_fix-jwt-admin-bypass

fix(jwt-auth): apply team TPM/RPM + attribution for admins using x-litellm-team-id

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

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

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():
     """