scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 90

fix(handle_jwt): use canonical UserTable.user_id instead of raw JWT claim

Browse source
This commit is contained in:
Varshith 2026-05-22 01:15:09 -05:00
parent d04373f4ce
commit b73ad40963
2 changed files with 238 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/auth/handle_jwt.py
View file

@ -1754,6 +1754,10 @@ class JWTAuthManager:
org_alias=org_alias,
)
# Reconcile to canonical UserTable.user_id if get_objects resolved via sso_user_id/email fuzzy match.
if user_object is not None and user_object.user_id and user_id != user_object.user_id:
user_id = user_object.user_id
# Derive org_id from org_object if resolved by alias
resolved_org_id = org_object.organization_id if org_object else org_id

234
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -2754,3 +2754,237 @@ def test_build_decode_kwargs_no_warning_when_scoped(
if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage()
]
assert matching == []
@pytest.mark.asyncio
async def test_auth_builder_reconciles_jwt_email_to_canonical_uuid():
"""JWT user_id is email; get_objects resolves the canonical UUID row — result must carry the UUID."""
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
canonical_uuid = "uuid-aaaa-1234-5678-canonical"
user_object = LiteLLM_UserTable(
user_id=canonical_uuid, user_role=LitellmUserRoles.INTERNAL_USER
)
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
) as mock_check_rbac,
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@example.com", "user@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),
):
mock_auth_jwt.return_value = {"sub": "user@example.com", "scope": ""}
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result["user_id"] == canonical_uuid
assert result["user_object"] == user_object
@pytest.mark.asyncio
async def test_auth_builder_reconciles_sso_external_id_to_canonical_uuid():
"""JWT user_id is an opaque SSO external ID; get_objects fuzzy-matches to the canonical UUID row."""
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
canonical_uuid = "uuid-bbbb-5678-canonical"
user_object = LiteLLM_UserTable(
user_id=canonical_uuid, user_role=LitellmUserRoles.INTERNAL_USER
)
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
) as mock_check_rbac,
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=("ext-sso-id-xyz", "user@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),
):
mock_auth_jwt.return_value = {"sub": "ext-sso-id-xyz", "scope": ""}
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result["user_id"] == canonical_uuid
assert result["user_object"] == user_object
@pytest.mark.asyncio
async def test_auth_builder_no_reconciliation_when_jwt_user_id_already_canonical():
"""When JWT user_id already matches user_object.user_id, result is unchanged."""
api_key = "test_jwt_token"
request_data = {"model": "gpt-4"}
general_settings = {"enforce_rbac": False}
route = "/chat/completions"
canonical_uuid = "uuid-cccc-9999-canonical"
user_object = LiteLLM_UserTable(
user_id=canonical_uuid, user_role=LitellmUserRoles.INTERNAL_USER
)
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
) as mock_check_rbac,
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=(canonical_uuid, "user@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),
):
mock_auth_jwt.return_value = {"sub": canonical_uuid, "scope": ""}
result = await JWTAuthManager.auth_builder(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result["user_id"] == canonical_uuid
assert result["user_object"] == user_object