From 5d80376bd9d2d6f044395054d6ebf3c717c873d5 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 27 May 2026 02:09:06 +0300 Subject: [PATCH] fix(jwt-auth): defer to single-team DB fallback on claim mismatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the single-team DB fallback introduced in #26418 to two more cases where it previously could not run: * `find_and_validate_specific_team_id`: when `team_id_jwt_field` is configured and a claim value is present in the token but the team does not exist in the LiteLLM DB (HTTPException 404 from `get_team_object`), return `(None, None)` instead of raising — the auth_builder fallback then attributes the request to the user's single DB team. Only HTTPException is caught; other errors (e.g. "No DB Connected") still propagate. * `find_team_with_model_access`: when none of the `team_ids_jwt_field` groups resolve to a real LiteLLM team, return `(None, None)` instead of raising 403 so the same fallback path runs. If at least one group DID resolve to a team but none granted the requested model, the original 403 is preserved (legitimate access denial — not a claim mismatch). Tracked via the new `any_claim_team_resolved` flag. The strict `is_required_team_id` raise and `enforce_team_based_model_access` raise remain unchanged. Unit tests cover both new soft-fail paths and guard each preserved path (strict required, enforce_team_based, the preserved 403, and the non-HTTPException propagation). Co-authored-by: Cursor --- litellm/proxy/auth/handle_jwt.py | 35 +++- .../proxy/auth/test_handle_jwt.py | 189 ++++++++++++++++++ 2 files changed, 214 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 5b116c142b2..11f6810aee1 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1037,15 +1037,24 @@ class JWTAuthManager: # First try to get team by team_id if individual_team_id: - team_object = await get_team_object( - team_id=individual_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, - ) - return individual_team_id, team_object + try: + team_object = await get_team_object( + team_id=individual_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, + ) + return individual_team_id, team_object + except HTTPException as e: + # Claim doesn't map to a known team — defer to fallback. + verbose_proxy_logger.debug( + "JWT team_id claim '%s' did not resolve to a team: %s", + individual_team_id, + e.detail, + ) + return None, None # If no team_id found, try to resolve via team_alias_jwt_field team_alias = jwt_handler.get_team_alias( @@ -1132,6 +1141,7 @@ class JWTAuthManager: ) return None, None + any_claim_team_resolved = False for team_id in team_ids: try: team_object = await get_team_object( @@ -1142,6 +1152,9 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, ) + if team_object is not None: + any_claim_team_resolved = True + if team_object and team_object.models is not None: team_models = team_object.models if isinstance(team_models, list) and ( @@ -1166,12 +1179,14 @@ class JWTAuthManager: except Exception: continue - if requested_model: + if requested_model and any_claim_team_resolved: + # Claim teams resolved but none grant the model — deny. raise HTTPException( status_code=403, detail=f"No team has access to the requested model: {requested_model}. Checked teams={team_ids}. Check `/models` to see all available models.", ) + # No claim team resolved — defer to fallback. return None, None @staticmethod diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 90c5d4f4fc5..7179aa510dd 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -2754,3 +2754,192 @@ def test_build_decode_kwargs_no_warning_when_scoped( if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage() ] assert matching == [] + + +# --------------------------------------------------------------------------- +# Defer to single-team DB fallback (PR #26418) when JWT claims are present +# but do not resolve to a LiteLLM team. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_unresolved_claim_returns_none(): + """team_id claim is present in the JWT but the team is missing in the + DB — return (None, None) so the auth_builder single-team fallback can run, + instead of raising and failing auth.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "claim-team-not-in-db"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = HTTPException(status_code=404, detail="missing") + + team_id, team_object = await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert team_id is None + assert team_object is None + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_unresolved_group_claim_returns_none( + monkeypatch, +): + """Group claim resolves to team_ids that don't exist in the DB — return + (None, None) instead of raising 403, so the single-team fallback can run.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}} + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + async def raise_404(*_args, **_kwargs): + raise HTTPException(status_code=404, detail="missing") + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", raise_404) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + team_id, team_object = await JWTAuthManager.find_team_with_model_access( + team_ids={"idp-group-a", "idp-group-b"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert team_id is None + assert team_object is None + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_non_http_exception_still_propagates(): + """Regression guard: only HTTPException (the "team doesn't exist in db" + 404 raised by get_team_object) is softened. Other errors — e.g. "No DB + Connected" — must still propagate so operator-side problems are loud.""" + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "some-claim-team"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = RuntimeError("simulated infrastructure error") + + with pytest.raises(RuntimeError, match="simulated infrastructure error"): + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_enforce_team_based_access_still_raises(): + """Regression guard: when no group claims are present and + `enforce_team_based_model_access` is on, the original 403 still fires — + the new soft-fail only applies to the unresolved-claim path inside the + loop, not to the no-team-claims-at-all path at the top.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(enforce_team_based_model_access=True) + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids=set(), + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + assert "enforce_team_based_model_access" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_resolved_team_without_model_still_raises_403( + monkeypatch, +): + """Regression guard: when the JWT group claim DOES resolve to a real + LiteLLM team but that team does not grant the requested model, keep the + original 403. Only the unresolved-claim case is softened.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"]) + + async def mock_get_team_object(*_args, **_kwargs): + return team + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids={"real-team"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + assert "No team has access to the requested model" in str(exc_info.value.detail)