From 6ea95a6379021208f7e9cc3c16e3ecfae865bf79 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 24 Apr 2026 09:40:59 -0700 Subject: [PATCH 1/3] fix(jwt-auth): apply team TPM/RPM + attribution for admins using x-litellm-team-id Scope the header-driven team fetch to LLM API routes so admin management routes keep the pre-existing bypass behavior (no phantom teams, no 404s on mgmt calls). Team context is threaded onto UserAPIKeyAuth so spend logs, rate limits, and team_models attribution are correctly applied when admins act on behalf of a team via x-litellm-team-id. --- litellm/proxy/auth/handle_jwt.py | 20 +++ litellm/proxy/auth/user_api_key_auth.py | 15 ++ .../proxy/auth/test_handle_jwt.py | 170 ++++++++++++++++++ 3 files changed, 205 insertions(+) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4a6856b6d14..44973db8c24 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 ( @@ -1493,6 +1494,25 @@ 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 + ) + 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 ff4f63593cb..9d45f38517c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -810,6 +810,21 @@ 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..0b7ce48b58c 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1498,6 +1498,176 @@ 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(): """ From e1bb542556b5caecf14a41367c18ab6fbc573417 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 24 Apr 2026 13:38:28 -0700 Subject: [PATCH 2/3] 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, From a0bba43cea0cd88684902f1cf33b953554bb53a4 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 24 Apr 2026 14:31:36 -0700 Subject: [PATCH 3/3] fix(jwt-auth): soft-fail unresolvable x-litellm-team-id for admins Previously, an admin JWT sending a stale/typo'd/missing x-litellm-team-id on an LLM API route received a hard 404 from get_team_object, blocking the request. Restore pre-PR admin behavior: if the header can't be resolved, skip team attribution and proceed with admin access, logging a warning with the header value and route so the misconfigured caller is diagnosable. --- litellm/proxy/auth/handle_jwt.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index caebbe47e1f..28fbe8a7ddd 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1427,14 +1427,27 @@ class JWTAuthManager: ) 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, - ) + 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