From 3c0d172d4ec43080123a157d57e4140f464e18db Mon Sep 17 00:00:00 2001 From: milan-berri Date: Sat, 25 Apr 2026 03:51:42 +0300 Subject: [PATCH] fix(proxy): single-team DB fallback when JWT has no team_id (#26418) * fix(proxy): infer team from DB when JWT has no team and user has one team - When team_id is unset after JWT auth but the user row has exactly one team, set team_id, team_object, and team_membership from DB. - Skip when zero or multiple teams (ambiguous). - Add parametrized unit tests in test_handle_jwt.py. Made-with: Cursor * fix(proxy): JWT single-team DB fallback: catch errors, tests match get_team_object - Wrap get_team_object + get_team_membership in one try/except; log and skip on failure (stale/missing team id no longer fails auth). - Parametrize tests: HTTP 404/500, membership error; use side_effect not return_value=None for missing team row. Made-with: Cursor * refactor(jwt): extract single-team fallback into _resolve_single_team_fallback helper Made-with: Cursor --- litellm/proxy/auth/handle_jwt.py | 70 +++- .../proxy/auth/test_handle_jwt.py | 298 ++++++++++++++++++ 2 files changed, 367 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 28fbe8a7ddd..f50c950d747 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1451,6 +1451,61 @@ class JWTAuthManager: admin_result["team_id"] = header_team_id admin_result["team_object"] = team_object + @staticmethod + async def _resolve_single_team_fallback( + user_object: Optional[LiteLLM_UserTable], + user_id: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + parent_otel_span: Optional[Span], + proxy_logging_obj: ProxyLogging, + team_id_upsert: Optional[bool], + ) -> tuple: + """ + If JWT did not resolve team_id, but the user belongs to exactly one team + in LiteLLM, load that team (and membership when user_id is set) so that + spend / metadata can be attributed correctly. + + Returns (team_id, team_object, team_membership_object). + Any DB error is debug-logged and the tuple is (None, None, None) — no + exception ever propagates from this helper. + """ + if user_object is None or not user_object.teams or len(user_object.teams) != 1: + return None, None, None + + _tid = user_object.teams[0] + try: + team_row = await get_team_object( + team_id=_tid, + 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=team_id_upsert, + ) + if team_row is None: + return None, None, None + + if not user_id: + return _tid, team_row, None + + team_membership = await get_team_membership( + user_id=user_id, + team_id=_tid, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return _tid, team_row, team_membership + except Exception: + verbose_proxy_logger.debug( + "JWT single-team fallback error, skipping. team_id=%s", + _tid, + exc_info=True, + ) + return None, None, None + @staticmethod async def auth_builder( api_key: str, @@ -1515,7 +1570,6 @@ class JWTAuthManager: object_id = jwt_handler.get_object_id(token=jwt_valid_token, default_value=None) # Get basic user info - scopes = jwt_handler.get_scopes(token=jwt_valid_token) user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info( jwt_handler, jwt_valid_token ) @@ -1637,6 +1691,20 @@ class JWTAuthManager: user_api_key_cache=user_api_key_cache, ) + # If JWT did not resolve team_id, attempt single-team DB fallback. + if team_id is None: + team_id, team_object, team_membership_object = ( + await JWTAuthManager._resolve_single_team_fallback( + user_object=user_object, + user_id=user_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, + ) + ) + ## MAP USER TO TEAMS await JWTAuthManager.map_user_to_teams( user_object=user_object, diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index cdb9ae4d9ab..9085469268c 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1,3 +1,4 @@ +from typing import Optional from unittest.mock import AsyncMock, patch import pytest @@ -5,6 +6,7 @@ import pytest from litellm.proxy._types import ( JWTLiteLLMRoleMap, LiteLLM_JWTAuth, + LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, @@ -2267,3 +2269,299 @@ async def test_find_and_validate_specific_team_id_no_hint_for_valid_field(): error_msg = str(exc_info.value) assert "Hint" not in error_msg + + +# --------------------------------------------------------------------------- +# Single-team DB fallback when JWT does not resolve team_id +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ( + "user_id", + "user_teams", + "get_team_object_return", + "expected_team_id", + "expect_get_team_called", + "expect_get_membership_called", + ), + [ + pytest.param( + "user_single_team_fb", + ["team_only_fb"], + "resolved_row", + "team_only_fb", + True, + True, + id="one_db_team_resolves_team_and_membership", + ), + pytest.param( + "user_multi_team_fb", + ["team_a", "team_b"], + "unused", + None, + False, + False, + id="two_db_teams_ambiguous_no_fallback", + ), + pytest.param( + "user_zero_teams_fb", + [], + "unused", + None, + False, + False, + id="zero_db_teams_no_fallback", + ), + pytest.param( + "user_orphan_team_fb", + ["team_missing_in_db"], + "http_404", + None, + True, + False, + id="one_team_id_but_row_missing_in_db", + ), + pytest.param( + "user_orphan_team_non404_fb", + ["team_err"], + "http_500", + None, + True, + False, + id="get_team_object_raises_non404_still_no_raise", + ), + ], +) +@pytest.mark.asyncio +async def test_auth_builder_single_team_db_fallback_when_jwt_has_no_team( + user_id: str, + user_teams: list, + get_team_object_return: Optional[str], + expected_team_id: Optional[str], + expect_get_team_called: bool, + expect_get_membership_called: bool, +) -> None: + """ + JWT does not set team_id (mocks return no team from token/header/routing). Behavior: + - exactly one team on user + get_team_object returns a row -> set team + membership + - two+ teams, or zero teams -> no get_team_object / no membership + - one team id but get_team_object raises (e.g. 404/500) -> skip fallback, no team, no error + """ + if len(user_teams) == 1 and get_team_object_return == "resolved_row": + only = user_teams[0] + team_table = LiteLLM_TeamTable(team_id=only) + membership = LiteLLM_TeamMembership( + user_id=user_id, team_id=only, litellm_budget_table=None + ) + get_team_return_value = team_table + membership_return_value = membership + else: + team_table = None + membership = None + get_team_return_value = None + membership_return_value = None + # "http_404" / "http_500" use get_team_object.side_effect, not return_value + + user_object = LiteLLM_UserTable( + user_id=user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + teams=user_teams, + ) + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + 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, "get_rbac_role", return_value=None), + patch.object(jwt_handler, "get_scopes", return_value=[]), + patch.object(jwt_handler, "get_object_id", return_value=None), + patch.object( + JWTAuthManager, + "get_user_info", + new_callable=AsyncMock, + return_value=(user_id, "u@example.com", True), + ), + patch.object(jwt_handler, "get_org_id", return_value=None), + patch.object(jwt_handler, "get_end_user_id", return_value=None), + patch.object( + JWTAuthManager, + "check_admin_access", + new_callable=AsyncMock, + return_value=None, + ), + patch.object( + JWTAuthManager, + "find_and_validate_specific_team_id", + new_callable=AsyncMock, + return_value=(None, None), + ), + patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()), + patch.object( + JWTAuthManager, + "find_team_with_model_access", + new_callable=AsyncMock, + return_value=(None, None), + ), + patch.object( + JWTAuthManager, + "get_objects", + new_callable=AsyncMock, + return_value=(user_object, None, None, None), + ), + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), + patch.object(JWTAuthManager, "validate_object_id", return_value=True), + patch.object( + JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock + ), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team, + patch( + "litellm.proxy.auth.handle_jwt.get_team_membership", + new_callable=AsyncMock, + ) as mock_get_membership, + ): + mock_auth_jwt.return_value = {"sub": user_id, "scope": ""} + if get_team_object_return in ("http_404", "http_500"): + from fastapi import HTTPException + + code = 404 if get_team_object_return == "http_404" else 500 + mock_get_team.side_effect = HTTPException( + status_code=code, + detail={ + "error": f"Team doesn't exist in db. Team={user_teams[0]}. Create team via `/team/new` call." + }, + ) + else: + mock_get_team.return_value = get_team_return_value + if membership_return_value is not None: + mock_get_membership.return_value = membership_return_value + + result = await JWTAuthManager.auth_builder( + api_key="test_jwt_token", + jwt_handler=jwt_handler, + request_data={"model": "gpt-4"}, + general_settings={"enforce_rbac": False}, + route="/chat/completions", + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert result["team_id"] == expected_team_id + if expected_team_id is not None: + assert result["team_object"] == team_table + assert result["team_membership"] == membership + else: + assert result["team_object"] is None + if not expect_get_membership_called: + assert result["team_membership"] is None + + if expect_get_team_called: + mock_get_team.assert_called() + else: + mock_get_team.assert_not_called() + if expect_get_membership_called: + mock_get_membership.assert_called_once() + else: + mock_get_membership.assert_not_called() + + +@pytest.mark.asyncio +async def test_auth_builder_single_team_fallback_membership_error_skips_no_raise(): + """ + get_team_object succeeds but get_team_membership raises — do not set team; no exception. + """ + from fastapi import HTTPException + + user_id = "u_mem_fail" + team_id_val = "team_mem_fail" + user_object = LiteLLM_UserTable( + user_id=user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + teams=[team_id_val], + ) + team_table = LiteLLM_TeamTable(team_id=team_id_val) + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + 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, "get_rbac_role", return_value=None), + patch.object(jwt_handler, "get_scopes", return_value=[]), + patch.object(jwt_handler, "get_object_id", return_value=None), + patch.object( + JWTAuthManager, + "get_user_info", + new_callable=AsyncMock, + return_value=(user_id, "u@example.com", True), + ), + patch.object(jwt_handler, "get_org_id", return_value=None), + patch.object(jwt_handler, "get_end_user_id", return_value=None), + patch.object( + JWTAuthManager, + "check_admin_access", + new_callable=AsyncMock, + return_value=None, + ), + patch.object( + JWTAuthManager, + "find_and_validate_specific_team_id", + new_callable=AsyncMock, + return_value=(None, None), + ), + patch.object(JWTAuthManager, "get_all_team_ids", return_value=set()), + patch.object( + JWTAuthManager, + "find_team_with_model_access", + new_callable=AsyncMock, + return_value=(None, None), + ), + patch.object( + JWTAuthManager, + "get_objects", + new_callable=AsyncMock, + return_value=(user_object, None, None, None), + ), + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), + patch.object(JWTAuthManager, "validate_object_id", return_value=True), + patch.object( + JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock + ), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team, + patch( + "litellm.proxy.auth.handle_jwt.get_team_membership", + new_callable=AsyncMock, + ) as mock_get_membership, + ): + mock_auth_jwt.return_value = {"sub": user_id, "scope": ""} + mock_get_team.return_value = team_table + mock_get_membership.side_effect = HTTPException( + status_code=500, detail="membership lookup failed" + ) + + result = await JWTAuthManager.auth_builder( + api_key="test_jwt_token", + jwt_handler=jwt_handler, + request_data={"model": "gpt-4"}, + general_settings={"enforce_rbac": False}, + route="/chat/completions", + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert result["team_id"] is None + assert result["team_object"] is None + assert result["team_membership"] is None + mock_get_team.assert_called() + mock_get_membership.assert_called_once()