fix(sso): pass decoded JWT access token to role mapping during SSO login#24701
Conversation
During SSO login, bearer tokens are stripped from the OAuth response before role mapping runs. Custom role claims encoded inside the JWT access token are lost, so map_jwt_role_to_litellm_role() returns None and the user falls back to internal_user_viewer. process_sso_jwt_access_token() now returns the decoded JWT payload, and a new _sync_user_role_from_jwt_role_map() receives it so jwt_litellm_role_map works correctly during SSO login.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
|
Greptile SummaryThis PR fixes a bug where SSO users configured with Key changes:
Confidence Score: 5/5Safe to merge; the fix is correct, well-tested, and introduces no regressions or security issues. All findings are P2 (style/quality suggestions): a misleading parameter name, a missing TTL on a cache write, and a minor fallback ambiguity. The core logic is sound — the decoded payload is correctly threaded through the call chain, the role is properly persisted to DB and cache, and downstream apply_user_info_values_to_sso_user_defined_values correctly preserves the mapped role. Existing tests were updated only to handle the new tuple arity (no assertions weakened), and the four new tests comprehensively cover the changed behavior with proper mocks. litellm/proxy/management_endpoints/ui_sso.py — specifically the _sync_user_role_from_jwt_role_map function's parameter naming and the missing TTL on the cache update.
|
| Filename | Overview |
|---|---|
| litellm/proxy/management_endpoints/ui_sso.py | Core SSO fix: process_sso_jwt_access_token now returns the decoded JWT payload; new _sync_user_role_from_jwt_role_map applies jwt_litellm_role_map using that payload; get_generic_sso_response extended to a 3-tuple; minor naming and TTL concerns flagged. |
| tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | Existing tests updated to unpack the new 3-tuple return value (no assertions weakened); four new unit tests cover the bug repro, fix verification, DB+cache update, and no-op cases — all fully mocked with no network calls. |
Sequence Diagram
sequenceDiagram
participant Browser
participant Callback as auth_callback
participant GenericSSO as get_generic_sso_response
participant ProcessJWT as process_sso_jwt_access_token
participant SyncRole as _sync_user_role_from_jwt_role_map
participant JWTHandler
participant DB as PrismaClient
participant Cache as DualCache
Browser->>Callback: GET /sso/callback
Callback->>GenericSSO: request, jwt_handler
GenericSSO->>ProcessJWT: token string, result
Note over ProcessJWT: Decode JWT, extract teams + role into result
ProcessJWT-->>GenericSSO: decoded payload dict (new return value)
GenericSSO-->>Callback: 3-tuple result, received_response, decoded payload
Callback->>SSOHandler: handle_callback with decoded payload + jwt_handler
SSOHandler->>DB: get_user_info_from_db
DB-->>SSOHandler: user_info
SSOHandler->>SyncRole: decoded JWT payload
Note over SyncRole: Check sync_user_role_and_teams and jwt_litellm_role_map config
SyncRole->>JWTHandler: map_jwt_role_to_litellm_role
JWTHandler-->>SyncRole: mapped LiteLLM role
SyncRole->>DB: update user_role in DB
SyncRole->>Cache: update user cache entry
SyncRole-->>SSOHandler: user_defined_values updated
SSOHandler->>SSOHandler: apply_user_info_values_to_sso_user_defined_values
SSOHandler->>Browser: RedirectResponse with session key
Reviews (1): Last reviewed commit: "fix(sso): pass decoded JWT access token ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| async def _sync_user_role_from_jwt_role_map( | ||
| jwt_handler: Optional[JWTHandler], | ||
| received_response: Optional[dict], | ||
| user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]], | ||
| prisma_client: PrismaClient, | ||
| user_api_key_cache: DualCache, | ||
| user_defined_values: Optional[SSOUserDefinedValues], | ||
| ) -> None: |
There was a problem hiding this comment.
received_response
The parameter is named received_response but it actually receives the decoded JWT access token payload (or the stripped OAuth response as a fallback). The raw received_response from the OAuth exchange — what the name implies — is the object that intentionally lacks role claims (the bug being fixed). Using the same name for both the stripped OAuth response and the decoded JWT payload will confuse future maintainers.
Consider renaming this parameter to token_claims or jwt_payload to make clear it is the decoded JWT dictionary passed directly to map_jwt_role_to_litellm_role, and update the call site at line 2574 accordingly.
| await user_api_key_cache.async_set_cache( | ||
| key=user_info.user_id, | ||
| value=user_info.model_dump() | ||
| if hasattr(user_info, "model_dump") | ||
| else dict(user_info), | ||
| ) |
There was a problem hiding this comment.
async_set_cache is called without a ttl argument. Depending on DualCache's default, this may write the updated user entry with no expiry (or a different default TTL than other user-cache entries set elsewhere), causing the stale-role scenario this PR fixes to persist until a manual cache flush.
Compare with how the cache is populated in get_user_info_from_db or similar helpers — those likely pass an explicit TTL. Passing the same TTL here keeps cache lifetime consistent across all writes to user_info.user_id.
| # access token, which is stripped from received_response. | ||
| await _sync_user_role_from_jwt_role_map( | ||
| jwt_handler=jwt_handler, | ||
| received_response=access_token_payload or received_response, |
There was a problem hiding this comment.
received_response contradicts fix intent
The PR description explains that received_response lacks custom role claims because bearer tokens were stripped (PR #22923). If access_token_payload is falsy (e.g. None for opaque tokens, or for Google/Microsoft SSO paths where it is never populated), this silently falls back to received_response — the very data source that is missing the claims.
In practice this is harmless today (for Google/MS SSO both will be None so the function exits early), but the guard should be explicit to avoid future confusion. Consider skipping the call entirely when access_token_payload is None rather than using the or fallback, so the two data sources are never silently interchanged.
ryan-crabbe-berri
merged commit a533de0
into
litellm_ryan-march-26
2 participants
Summary
During SSO login, bearer tokens (
access_token,id_token,refresh_token) are stripped from the OAuth response before role mapping runs (introduced in #22923 for security — preventing token leakage in error messages). Custom role claims encoded inside the JWT access token are lost, somap_jwt_role_to_litellm_role()returnsNoneand the user falls back tointernal_user_viewer.Changes
process_sso_jwt_access_token()now returns the decoded JWT payload dict (previously returnedNone)_sync_user_role_from_jwt_role_map()— appliesjwt_litellm_role_mapduring SSO login using the decoded access token payload instead of the stripped OAuth responseget_generic_sso_response()returns a 3-tuple including the decoded payloadTest plan
test_ui_sso.py::TestSyncUserRoleFromJwtRoleMap: