mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
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
This commit is contained in:
milan-berri
GitHub
parent
7723a54478
commit
3c0d172d4e
2 changed files with 367 additions and 1 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue