From e1bb542556b5caecf14a41367c18ab6fbc573417 Mon Sep 17 00:00:00 2001
From: Ryan Crabbe <ryan@berri.ai>
Date: Fri, 24 Apr 2026 13:38:28 -0700
Subject: [PATCH] chore: fix linting (ruff PLR0915, black) on admin team-header
 fix

Extract the admin team-header attachment into a helper so
auth_builder stays under the 50-statement lint threshold; apply
black formatting to the two files flagged on the prior commit.
No behavior change.
---
 litellm/proxy/auth/handle_jwt.py              | 60 +++++++++++++------
 litellm/proxy/auth/user_api_key_auth.py       |  4 +-
 .../proxy/auth/test_handle_jwt.py             |  4 +-
 3 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index 44973db8c24..caebbe47e1f 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -1405,6 +1405,39 @@ 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
+        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,
+        )
+        admin_result["team_id"] = header_team_id
+        admin_result["team_object"] = team_object
+
     @staticmethod
     async def auth_builder(
         api_key: str,
@@ -1494,25 +1527,16 @@ class JWTAuthManager:
             jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token
         )
         if admin_result:
-            # When an admin explicitly acts on behalf of a team via
-            # x-litellm-team-id on an LLM API route, fetch the team so
-            # team TPM/RPM limits and attribution apply. For admin
-            # management routes we intentionally ignore the header to
-            # preserve the pre-existing bypass behavior.
-            header_team_id = (
-                request_headers.get("x-litellm-team-id") if request_headers else None
+            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,
             )
-            if header_team_id and RouteChecks.is_llm_api_route(route=route):
-                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,
-                )
-                admin_result["team_id"] = header_team_id
-                admin_result["team_object"] = team_object
             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 9d45f38517c..dca38bf5801 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -821,9 +821,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                                 else None
                             ),
                             team_models=(
-                                team_object.models
-                                if team_object is not None
-                                else []
+                                team_object.models if team_object is not None else []
                             ),
                             team_metadata=(
                                 team_object.metadata
diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py
index 0b7ce48b58c..cdb9ae4d9ab 100644
--- a/tests/test_litellm/proxy/auth/test_handle_jwt.py
+++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py
@@ -1521,9 +1521,7 @@ async def test_auth_builder_admin_on_llm_route_honors_team_header():
         ),
     )
 
-    team_object = LiteLLM_TeamTable(
-        team_id="team-low", tpm_limit=100, rpm_limit=2
-    )
+    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,