fix(jwt): accept a team alias in x-litellm-team-id

litellm/proxy/auth/handle_jwt.py

@@ -70,6 +70,7 @@
 from litellm.types.proxy.auth.auth_checks import UserNotFoundError
 
 from .auth_checks import (
+    TeamNotFoundError,
     _allowed_routes_check,
     allowed_routes_check,
     get_actual_routes,
@@ -147,6 +148,12 @@ class _JWTProvisioning:
     team_id_upsert: bool
 
 
+@dataclass(frozen=True, slots=True)
+class HeaderTeam:
+    header_value: str
+    team_id: str
+
+
 class AgentLookup(Protocol):
     """The registered-agent lookups a JWT agent claim is matched against."""
 
@@ -1871,48 +1878,104 @@ def validate_object_id(
         return True
 
     @staticmethod
-    def get_team_id_from_header(
+    def _team_header_value(request_headers: Mapping[str, str] | None) -> str | None:
+        if not request_headers:
+            return None
+        normalized_headers: Final = {k.lower(): v for k, v in request_headers.items()}
+        return normalized_headers.get("x-litellm-team-id")
+
+    @staticmethod
+    def _raise_header_team_not_allowed(header_value: str, allowed_team_ids: set[str]) -> NoReturn:
+        raise HTTPException(
+            status_code=403,
+            detail=f"Team '{header_value}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}",
+        )
+
+    @staticmethod
+    async def _team_id_by_alias(
+        team_alias: str,
+        prisma_client: PrismaClient | None,
+        user_api_key_cache: UserApiKeyCache,
+        parent_otel_span: Span | None,
+        proxy_logging_obj: ProxyLogging,
+    ) -> str | None:
+        if prisma_client is None:
+            return None
+        try:
+            team: Final = await get_team_object_by_alias(
+                team_alias=team_alias,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=parent_otel_span,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+        except HTTPException as exc:
+            if exc.status_code >= 500:
+                raise
+            return None
+        return team.team_id
+
+    @staticmethod
+    async def resolve_team_from_header(
         request_headers: Mapping[str, str] | None,
         allowed_team_ids: set[str],
-        fallback_to_db_teams: bool = False,
-    ) -> str | None:
+        fallback_to_db_teams: bool,
+        prisma_client: PrismaClient | None,
+        user_api_key_cache: UserApiKeyCache,
+        parent_otel_span: Span | None,
+        proxy_logging_obj: ProxyLogging,
+    ) -> HeaderTeam | None:
         """
-        Extract team_id from x-litellm-team-id header if present.
-        Validates that the team is in the user's allowed teams from JWT.
-
-        Args:
-            request_headers: Dictionary of request headers
-            allowed_team_ids: Set of team IDs the user is allowed to access (from JWT)
-            fallback_to_db_teams: When True and the JWT carries no team claims
-                (allowed_team_ids is empty), the header value is returned
-                provisionally and validated against DB memberships later in
-                auth_builder instead of being rejected here.
-
-        Returns:
-            The team_id from header if valid, None otherwise
+        The team named by x-litellm-team-id, which may carry a team id or a team
+        alias. A value that is already an allowed team id (or, under the DB
+        fallback, an existing team id) never costs an alias lookup; an alias is
+        accepted only when the team it names would have been accepted by id.
+        Under the DB fallback only a team row that is provably absent falls
+        through to the alias lookup; a read that failed for any other reason
+        keeps the membership denial the id path already gives.
 
         Raises:
-            HTTPException: If team_id is provided but not in allowed_team_ids
+            HTTPException: 403 when neither the value nor the team it aliases is
+                an allowed team, or the DB fallback's membership denial when the
+                value names no team at all; a 5xx from the alias lookup itself
+                is re-raised rather than reported as a denial
         """
-        if not request_headers:
+        header_value: Final = JWTAuthManager._team_header_value(request_headers)
+        if not header_value:
             return None
 
-        # Normalize headers to lowercase for case-insensitive lookup
-        normalized_headers: Final = {k.lower(): v for k, v in request_headers.items()}
-        header_team_id: Final = normalized_headers.get("x-litellm-team-id")
-
-        if not header_team_id:
-            return None
+        if fallback_to_db_teams and not allowed_team_ids:
+            try:
+                await get_team_object(
+                    team_id=header_value,
+                    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=False,
+                )
+            except TeamNotFoundError:
+                aliased_team_id: Final = await JWTAuthManager._team_id_by_alias(
+                    header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
+                )
+                if aliased_team_id is None:
+                    JWTAuthManager._raise_header_team_membership_denial(header_value)
+                return HeaderTeam(header_value=header_value, team_id=aliased_team_id)
+            except HTTPException:
+                JWTAuthManager._raise_header_team_membership_denial(header_value)
+            return HeaderTeam(header_value=header_value, team_id=header_value)
 
-        defer_to_db_membership: Final = fallback_to_db_teams and not allowed_team_ids
-        if not defer_to_db_membership and header_team_id not in allowed_team_ids:
-            raise HTTPException(
-                status_code=403,
-                detail=f"Team '{header_team_id}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}",
-            )
+        if header_value in allowed_team_ids:
+            verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value)
+            return HeaderTeam(header_value=header_value, team_id=header_value)
 
-        verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_team_id)
-        return header_team_id
+        team_id_by_alias: Final = await JWTAuthManager._team_id_by_alias(
+            header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
+        )
+        if team_id_by_alias is None or team_id_by_alias not in allowed_team_ids:
+            JWTAuthManager._raise_header_team_not_allowed(header_value, allowed_team_ids)
+        verbose_proxy_logger.debug("Using team_id %s for x-litellm-team-id alias: %s", team_id_by_alias, header_value)
+        return HeaderTeam(header_value=header_value, team_id=team_id_by_alias)
 
     @staticmethod
     async def map_user_to_teams(
@@ -2264,31 +2327,34 @@ def _is_team_route_allowed(
         )
 
     @staticmethod
-    def _raise_header_team_membership_denial(team_id: str) -> NoReturn:
+    def _raise_header_team_membership_denial(header_value: str) -> NoReturn:
         """
         The single denial shape for a provisional x-litellm-team-id header,
         raised identically for nonexistent teams and for teams the user is not
-        a member of, so the response does not reveal whether a team id exists.
+        a member of, and naming only the value the caller sent, so the response
+        reveals neither whether a team exists nor which id an alias maps to.
         """
         raise HTTPException(
             status_code=403,
-            detail=(f"Team '{team_id}' (from x-litellm-team-id header) is not in your team memberships."),
+            detail=(f"Team '{header_value}' (from x-litellm-team-id header) is not in your team memberships."),
         )
 
     @staticmethod
     def _validate_header_team_in_db_membership(
         team_id: str,
         user_object: LiteLLM_UserTable | None,
+        header_value: str,
     ) -> None:
         """
         A provisional team_id from the x-litellm-team-id header (accepted without
         JWT-team validation when the JWT carries no team claims) must exist in the
-        user's DB team memberships before it becomes request context.
+        user's DB team memberships before it becomes request context. The denial
+        names `header_value`, the id or alias the caller sent, not `team_id`.
         """
         user_team_ids: Final = user_object.teams if user_object else []
         if team_id in user_team_ids:
             return
-        JWTAuthManager._raise_header_team_membership_denial(team_id)
+        JWTAuthManager._raise_header_team_membership_denial(header_value)
 
     @staticmethod
     async def auth_builder(
@@ -2514,13 +2580,17 @@ async def authorize_jwt(
         if specific_team_id and not db_team_fallback:
             all_team_ids.add(specific_team_id)
 
-        header_team_id: Final = JWTAuthManager.get_team_id_from_header(
+        header_team: Final = await JWTAuthManager.resolve_team_from_header(
             request_headers=request_headers,
             allowed_team_ids=all_team_ids,
             fallback_to_db_teams=db_team_fallback,
+            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:
-            team_id = header_team_id
+        if header_team:
+            team_id = header_team.team_id
             # A provisional header team (accepted only because the JWT carries no
             # team claims) is validated against DB membership further down; never
             # upsert it here or an attacker-supplied x-litellm-team-id would create
@@ -2538,7 +2608,7 @@ async def authorize_jwt(
             except HTTPException:
                 if not db_team_fallback:
                     raise
-                JWTAuthManager._raise_header_team_membership_denial(team_id)
+                JWTAuthManager._raise_header_team_membership_denial(header_team.header_value)
         elif not team_id and not db_team_fallback:
             ## SPECIFIC TEAM ID
             (
@@ -2679,10 +2749,11 @@ async def authorize_jwt(
                 proxy_logging_obj=proxy_logging_obj,
                 team_id_upsert=team_id_upsert,
             )
-        elif db_team_fallback and team_id == header_team_id:
+        elif db_team_fallback and header_team is not None and team_id == header_team.team_id:
             JWTAuthManager._validate_header_team_in_db_membership(
                 team_id=team_id,
                 user_object=user_object,
+                header_value=header_team.header_value,
             )
             if not JWTAuthManager._is_team_route_allowed(
                 route=route,
@@ -2692,7 +2763,8 @@ async def authorize_jwt(
                 raise HTTPException(
                     status_code=403,
                     detail=(
-                        f"Team '{team_id}' (from x-litellm-team-id header) is not allowed to access route '{route}'."
+                        f"Team '{header_team.header_value}' (from x-litellm-team-id header) "
+                        f"is not allowed to access route '{route}'."
                     ),
                 )
 

tests/e2e/coverage_registry/other.yaml

@@ -15,6 +15,8 @@
 - {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"}
 - {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"}
 - {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"}
+- {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"}
+- {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"}
 - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"}
 - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"}
 - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"}

tests/e2e/other/other_client.py

@@ -15,15 +15,25 @@
 
 from dataclasses import dataclass
 
-from e2e_http import NoBody, ProbeResult, Result
+from e2e_http import AuthHeaders, NoBody, ProbeResult, Result
 from idp import Keycloak, keycloak_from_env
 from models import (
+    ChatBody,
+    ChatResponse,
     ReadinessDetailsResponse,
     ReadinessResponse,
     UserListParams,
     UserListResponse,
 )
 from proxy_client import ProxyClient
+from pydantic import Field
+
+
+class TeamHeaders(AuthHeaders):
+    """Bearer auth plus ``x-litellm-team-id``, the header a JWT caller sends to
+    pick one of the teams it belongs to."""
+
+    x_litellm_team_id: str = Field(serialization_alias="x-litellm-team-id")
 
 
 @dataclass(frozen=True, slots=True)
@@ -66,6 +76,18 @@ def readiness_details_unauthenticated(self) -> Result[ReadinessDetailsResponse]:
             response_type=ReadinessDetailsResponse,
         )
 
+    def chat_as_team(self, token: str, team: str, body: ChatBody) -> Result[ChatResponse]:
+        """POST /chat/completions under `token` with `x-litellm-team-id: team`."""
+        return self.proxy.transport.post(
+            "/chat/completions",
+            headers=TeamHeaders(
+                authorization=self.proxy.transport.bearer(token).authorization,
+                x_litellm_team_id=team,
+            ),
+            json=body,
+            response_type=ChatResponse,
+        )
+
     def list_users_as(self, key: str) -> Result[UserListResponse]:
         """GET /user/list under `key`. Admin-only, so it doubles as the master
         key's authorization proof: the master key (proxy admin) reads it, a