From 7b4fd47c6e665e36ee0669ae7d15887a085dd472 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 15:28:14 -0700 Subject: [PATCH] fix(jwt): let x-litellm-team-id select DB membership teams when the token also carries a team claim (#43206) * fix(jwt): let x-litellm-team-id select DB membership teams when the token also carries a team claim Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(jwt): describe header team selection under fallback_to_db_teams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 11 +- litellm/proxy/auth/handle_jwt.py | 58 +++---- .../proxy/auth/test_handle_jwt.py | 141 +++++++++++++++++- 3 files changed, 175 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 89fa0058644..fd13697dab3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5355,11 +5355,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=False, description=( "When True, users whose JWT contains no team claims are authenticated " - "using their database team memberships instead of receiving HTTP 403. " - "Usage is attributed to the user's first resolvable DB team, or to the " - "team specified via the x-litellm-team-id request header (validated " - "against DB membership). Requires user_id_upsert=True so that user " - "records exist before the fallback runs." + "using their database team memberships instead of receiving HTTP 403, " + "with usage attributed to the user's first resolvable DB team. Whether or " + "not the JWT carries team claims, the x-litellm-team-id request header may " + "select any team the user is a member of in the database (validated against " + "DB membership); without the header the JWT team stays the default. Requires " + "user_id_upsert=True so that user records exist before the fallback runs." ), ) issuers: list[JWTIssuerConfig] | None = Field( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index e07b20fd5d5..4f41b283a33 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1930,12 +1930,12 @@ class JWTAuthManager: ) -> HeaderTeam | None: """ 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. + alias. A value that is already an allowed team id never costs a lookup; + under the DB fallback any other value is accepted provisionally, by id + or alias, for the membership check auth_builder runs later. 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: 403 when neither the value nor the team it aliases is @@ -1948,7 +1948,11 @@ class JWTAuthManager: if not header_value: return None - if fallback_to_db_teams and not 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) + + if fallback_to_db_teams: try: await get_team_object( team_id=header_value, @@ -1969,10 +1973,6 @@ class JWTAuthManager: JWTAuthManager._raise_header_team_membership_denial(header_value) return HeaderTeam(header_value=header_value, team_id=header_value) - 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) - 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 ) @@ -2353,9 +2353,9 @@ class JWTAuthManager: 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. The denial + A provisional team_id from the x-litellm-team-id header (accepted under + fallback_to_db_teams because it is outside the JWT's teams) must exist in + the 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 [] @@ -2587,22 +2587,30 @@ class JWTAuthManager: if specific_team_id and not db_team_fallback: all_team_ids.add(specific_team_id) + header_db_fallback: Final = handler.litellm_jwtauth.fallback_to_db_teams and team_id is None + 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, + fallback_to_db_teams=header_db_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, ) + provisional_header_team: Final = ( + header_team + if header_team is not None and header_db_fallback and header_team.team_id not in all_team_ids + else None + ) 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 - # an orphaned team row before that check runs. A genuine membership team - # already exists, so suppressing the upsert in that case costs nothing. + # A provisional header team (accepted because it is outside the + # JWT's teams under fallback_to_db_teams) is validated against DB + # membership further down; never upsert it here or an + # attacker-supplied x-litellm-team-id would create an orphaned team + # row before that check runs. A genuine membership team already + # exists, so suppressing the upsert in that case costs nothing. try: team_object = await get_team_object( team_id=team_id, @@ -2610,10 +2618,10 @@ class JWTAuthManager: user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, - team_id_upsert=(team_id_upsert and not db_team_fallback), + team_id_upsert=(team_id_upsert and provisional_header_team is None), ) except HTTPException: - if not db_team_fallback: + if provisional_header_team is None: raise JWTAuthManager._raise_header_team_membership_denial(header_team.header_value) elif not team_id and not db_team_fallback: @@ -2756,11 +2764,11 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - elif db_team_fallback and header_team is not None and team_id == header_team.team_id: + elif provisional_header_team is not None and team_id == provisional_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, + header_value=provisional_header_team.header_value, ) if not JWTAuthManager._is_team_route_allowed( route=route, @@ -2770,7 +2778,7 @@ class JWTAuthManager: raise HTTPException( status_code=403, detail=( - f"Team '{header_team.header_value}' (from x-litellm-team-id header) " + f"Team '{provisional_header_team.header_value}' (from x-litellm-team-id header) " f"is not allowed to access route '{route}'." ), ) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index f8b9043a23f..b1622e0dff0 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -5324,16 +5324,22 @@ def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deploym @pytest.mark.asyncio -async def test_resolve_team_from_header_defers_to_db_membership_only_without_jwt_claims(): +async def test_resolve_team_from_header_accepts_db_teams_provisionally_under_fallback_even_with_jwt_claims(): """With fallback_to_db_teams=True, an x-litellm-team-id header naming an existing - team is accepted provisionally only when the JWT carries no team claims (allowed - set empty). When the JWT does carry team claims, the header must still be validated - against them, and the flag-off behavior must keep rejecting unknown teams.""" + team is accepted provisionally whether or not the JWT carries team claims; the + union of JWT teams and DB memberships is enforced by auth_builder's later + membership check. Unknown values still 403, and the flag-off behavior keeps + rejecting teams outside the JWT's allowed set.""" known_ids = frozenset({"team-from-db"}) deferred, _, _ = await _resolve_header("team-from-db", set(), True, _teams_by_id(known_ids), _team_alias_lookup_404) assert deferred == HeaderTeam(header_value="team-from-db", team_id="team-from-db") + deferred_with_claims, _, _ = await _resolve_header( + "team-from-db", {"team-1"}, True, _teams_by_id(known_ids), _team_alias_lookup_404 + ) + assert deferred_with_claims == HeaderTeam(header_value="team-from-db", team_id="team-from-db") + with pytest.raises(HTTPException) as exc_info: await _resolve_header("team-x", {"team-1", "team-2"}, True, _teams_by_id(known_ids), _team_alias_lookup_404) assert exc_info.value.status_code == 403 @@ -5849,6 +5855,7 @@ async def _run_auth_builder_with_header_team( allowed_team_ids: set, fake_get_team_by_alias=_team_alias_lookup_404, route: str = "/chat/completions", + send_header: bool = True, ): jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = jwt_auth_config @@ -5909,7 +5916,7 @@ async def _run_auth_builder_with_header_team( user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, - request_headers={"x-litellm-team-id": header_team_id}, + request_headers={"x-litellm-team-id": header_team_id} if send_header else {}, ) @@ -7283,6 +7290,130 @@ async def test_auth_builder_header_alias_under_db_fallback_keeps_the_team_allowe assert allowed["team_id"] == "team_member" +@pytest.mark.asyncio +async def test_auth_builder_header_selects_db_membership_team_when_jwt_also_carries_a_team_claim() -> None: + """Under fallback_to_db_teams, x-litellm-team-id may name a DB-membership + team the JWT does not claim (LIT-8656): the allowed set is the JWT teams + union the user's DB memberships, not the JWT teams alone. The flag-off + path keeps rejecting the same header against the JWT's allowed teams.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member"})) + + by_membership = await _run_auth_builder_with_header_team( + config, token, "team_member", user_object, fake_get_team, {"team_claimed"} + ) + assert by_membership["team_id"] == "team_member" + assert by_membership["team_object"].team_id == "team_member" + + by_claim = await _run_auth_builder_with_header_team( + config, token, "team_claimed", user_object, fake_get_team, {"team_claimed"} + ) + assert by_claim["team_id"] == "team_claimed" + + flag_off = LiteLLM_JWTAuth(fallback_to_db_teams=False, team_id_jwt_field="appid") + with pytest.raises(HTTPException) as exc_info: + await _run_auth_builder_with_header_team( + flag_off, token, "team_member", user_object, fake_get_team, {"team_claimed"} + ) + assert exc_info.value.status_code == 403 + assert "JWT's allowed teams" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_auth_builder_header_non_member_team_is_denied_when_jwt_also_carries_a_team_claim() -> None: + """A header naming a team the user does not belong to stays a membership + denial even when the JWT carries a team claim, and an existing but + non-member team produces the exact same 403 shape as a nonexistent one so + the response is no oracle for which team ids exist.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member", "team_other"})) + + with pytest.raises(HTTPException) as outsider_exc: + await _run_auth_builder_with_header_team( + config, token, "team_other", user_object, fake_get_team, {"team_claimed"} + ) + with pytest.raises(HTTPException) as missing_exc: + await _run_auth_builder_with_header_team( + config, token, "team_ghost", user_object, fake_get_team, {"team_claimed"} + ) + + assert outsider_exc.value.status_code == 403 + assert missing_exc.value.status_code == 403 + assert outsider_exc.value.detail == ( + "x-litellm-team-id 'team_other' does not resolve to a team id or a unique team alias among your " + "team memberships." + ) + assert missing_exc.value.detail.replace("team_ghost", "") == outsider_exc.value.detail.replace( + "team_other", "" + ) + assert "exist" not in missing_exc.value.detail + + +@pytest.mark.asyncio +async def test_auth_builder_no_header_keeps_the_jwt_team_when_fallback_to_db_teams_is_on() -> None: + """With no x-litellm-team-id header, fallback_to_db_teams must not disturb + the claim path: the JWT's own team claim still binds the request.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + + result = await _run_auth_builder_with_header_team( + config, + token, + "team_member", + user_object, + _teams_by_id(frozenset({"team_claimed", "team_member"})), + {"team_claimed"}, + send_header=False, + ) + assert result["team_id"] == "team_claimed" + + +@pytest.mark.asyncio +async def test_auth_builder_team_id_default_does_not_widen_the_header_allowed_set() -> None: + """team_id_default fills in a team for claimless tokens but must not widen + the header's allowed set: a header naming the default team is still held + to DB membership under fallback_to_db_teams.""" + user_object = LiteLLM_UserTable( + user_id="u_default", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_default="team_default") + token = {"sub": "u_default", "scope": ""} + + with pytest.raises(HTTPException) as exc_info: + await _run_auth_builder_with_header_team( + config, + token, + "team_default", + user_object, + _teams_by_id(frozenset({"team_default", "team_member"})), + set(), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == ( + "x-litellm-team-id 'team_default' does not resolve to a team id or a unique team alias among your " + "team memberships." + ) + + @pytest.mark.asyncio async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag(): """Reading the singular team claim during sync is scoped to fallback_to_db_teams.