feat(mcp): source the ID-JAG subject from the user's stored SSO assertion - #35147
Conversation
|
|
Greptile SummaryThis PR adds stored SSO assertions as an ID-JAG subject source and updates the corresponding retry, cache, error handling, tests, and dashboard configuration
Confidence Score: 3/5The PR does not yet appear safe to merge because an in-flight ID-JAG exchange can repopulate its cache slot after an upstream-authentication invalidation ExchangedTokenCache serializes cache computation under a per-key lock, but invalidate bypasses that lock and deletes immediately; a computation already awaiting the token endpoint can consequently write into the slot after eviction, undermining the retry path's guarantee that pre-invalidation work cannot survive Files Needing Attention: litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py and its concurrency regression tests
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py | Adds stored-assertion subject resolution, expiry handling, principal-scoped cache slots, fingerprints, and store-independent invalidation |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py | Adds fingerprinted exchanged-token cache entries, but invalidation remains unsynchronized with in-flight cache computations |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py | Introduces an injectable assertion-store protocol and converts storage exceptions into a typed unavailable error |
| litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | Extends upstream authentication retry handling to store-sourced ID-JAG calls |
| ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx | Adds the dashboard field set and validation required to configure both ID-JAG exchange legs |
Reviews (8): Last reviewed commit: "feat(mcp): source the ID-JAG subject fro..." | Re-trigger Greptile
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
d5c4c26 to
9b05c30
Compare
|
@greptileai the assertion-store outage finding was correct and is fixed in
Fixed by mirroring the convention the sibling per-user OAuth store already uses. It is deliberately not mapped to the 412 the other misses use. A missing assertion means sign in again; an unreachable database means nothing the user can act on, and reporting it as an absence would send people through a pointless re-login during an outage. Four regression tests, each verified to fail when the guard is removed: the resolver maps an outage to Two other gaps closed in the same push, both found by measuring which branches my own diff left unexercised: the private-key-JWT arm of the ID-JAG cache-key fingerprint, so rotating a signing key re-mints instead of serving a bearer authorized under the retired key, and a naive stored Note the review is pinned to |
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
|
@greptileai on the remaining uncached-read point, I measured it rather than argue from intuition, and it is more significant than a hand-wave would suggest. Against the live proxy with Postgres So it is one read per call, and it is five of the six statements on that path. Not a rounding error, and worth stating plainly rather than minimizing. I am still proposing to defer the cache rather than add it here, for three reasons. The exchange itself is already cached, so what remains is a single primary-key The blast radius is a brand-new opt-in auth mode. Before this PR Most importantly, a cache here is auth-material caching with a real staleness hazard, and it deserves its own review rather than being appended to this one. The reader judges The shape I would propose for the follow-up: a short-TTL in-memory cache in front of Please re-review the current head |
9b05c30 to
f3e874d
Compare
|
@greptileai you were right, and the finding is a direct consequence of my previous fix. Fixed in The outage guard I added in the last round converted an unhandled error into a silent no-op, and the silence recreated the replay bug the PR set out to fix. Exact sequence: upstream 401, The root cause was that invalidation derived its eviction key by re-reading the store, so it could not do its job precisely when the store was the thing failing. Invalidation no longer depends on the store: A caller that presented its own token is exactly reproducible from the request, so that key is derived directly and no store is consulted. A store-sourced subject records, at resolve time, the cache key it minted under a Worth recording that the alias was too coarse in my first attempt: keying it on the principal alone let an inbound-token-sourced entry and a store-sourced entry for the same user collide, and an existing test ( New regression test reproduces your scenario literally: resolve, take the store down, invalidate, bring the store back, resolve again, and assert the second resolve re-mints instead of replaying. Three mutations verified against it, including deleting the alias write and ignoring the alias at invalidation; all three fail the test. 528 passing in the outbound-credentials suite, On the uncached read, I have nothing to add beyond the measurement in my previous comment and am content to defer to your judgement on whether it should land here or as a follow-up. Please re-review the current head |
f3e874d to
5981d02
Compare
|
@greptileai correct again, and fixed in The alias held a single key, so two store-sourced resolutions for the same principal that mint different keys (a re-login between them changes the subject token; a server edit between them changes the config) would overwrite each other, and the 401 recovery would then evict the wrong entry and leave its own rejected bearer cached for the retry. Rather than try to thread per-request identity through the invalidation call, the alias now holds the set of keys currently live for that principal and invalidation clears all of them. The asymmetry is what makes this the right shape: over-eviction costs a re-mint, under-eviction is the replay bug. The set is bounded per principal so it cannot grow. New regression test walks the exact interleaving: resolve, rotate the stored assertion, resolve again so two distinct keys are live, invalidate once, then point the store back at the first assertion and assert the next resolve re-mints rather than replaying the first bearer. Two mutations verified against it, keeping only the newest key and evicting only one of the set; both fail the test. 529 passing in the outbound-credentials suite, format and strict-rule gate clean. Two CI items on this PR that are not from this branch, for the record. Please re-review the current head |
5981d02 to
940669d
Compare
|
@greptileai this finding splits into two parts, and they need different answers. Head is now On the lost update, I do not think it can occur, and I checked rather than asserting it. On discarding a live key, you are right, and by a different route than locking: the set is bounded, and truncation dropped the oldest key silently. A key the tracking forgets can never be evicted afterwards, so the bound quietly reintroduced the exact replay hole the tracking exists to close. Truncation now evicts the bearer of every key it drops, in the same operation. Forgetting and evicting cannot separate. New test drives one principal past the bound and asserts the pushed-out entry re-mints rather than replaying; removing the eviction-on-truncation line fails it. 530 passing in the outbound-credentials suite, format and strict-rule gate clean. Separately, I verified the dashboard edit path end to end on a live proxy rather than trusting the merge behavior your earlier summary noted: opened the saved ID-JAG server in the dashboard, confirmed leg 1 prefills from its column while leg 2 stays blank because it lives in the encrypted credentials blob the API does not return, saved without re-entering it, and re-ran the tool call. It still returned the correct end user, so omitted credential fields do survive an edit. Please re-review the current head |
940669d to
5d8da16
Compare
|
@greptileai I stopped patching this mechanism and replaced it. Head is Your last finding split into two parts. On the lost update I still believe the read-modify-write is safe: It failed, 52 of 60 survived. 52 is exactly 60 minus the bound of 8, which pinpointed the cause: eviction-on-truncation ran inside the pre-exchange bookkeeping, before So the design is gone rather than patched again. The exchanged-token entry is now addressed by a slot key derived from the principal (plus the caller's own token when it presented one), and the fingerprint of the subject token and config is stored beside the bearer and compared on every read. A mismatch reads as a miss and re-mints. Invalidation is one delete of a key it can always compute, with no store lookup. That deletes the alias map, the remembered-key set, the bound, the truncation-eviction, and the store dependency in invalidation. The whole class of ordering and interleaving question goes with them. Rotation still re-mints, now via the fingerprint instead of the key, and two callers sharing an empty principal cannot receive each other's bearer because a fingerprint mismatch is a miss. The stress harness now passes with all 60 re-minting, and a scaled-down version of it is committed as a regression test alongside a cross-caller fingerprint test. Three mutations verified: ignoring the fingerprint on read, dropping the caller token from the slot key, and disabling eviction; all three fail. 899 passing across the outbound-credentials and manager suites, format and strict-rule gate clean, and I re-ran the live proxy end to end on the restructured cache: both legs still execute and the upstream still resolves the correct end user. Please re-review the current head |
4f53812 to
d1488da
Compare
|
@greptileai the review is still anchored to The first was a defect I introduced while restructuring and want to name rather than let pass quietly. My scripted edit left The second extends the same static-analysis pass across the whole diff instead of the two files I had been editing. Production files now measure 495 basedpyright errors against 495 on clean Also folded in from the previous round: the Local state on Please re-review the current head |
…tion The ID-JAG egress arm could only assert a caller that presented its own IdP identity token on the request, so an agent holding a brokered LiteLLM credential got a 412 and never reached the upstream. The assertion captured at SSO login was already persisted per user for exactly this purpose, but nothing read it back. The arm now falls back to that stored assertion, keyed on the authenticated principal's user_id. The identity is always taken from the credential the gateway authenticated, never from a caller-supplied field, so no caller can select whose identity is asserted upstream. A missing, expired, or unidentified subject stays a 412; ID-JAG exists to assert a specific user and a missing subject has no safe substitute. A store outage is the one exception: it is surfaced as a typed AssertionStoreUnavailable and mapped to 503, so a database blip cannot 500 the egress or the upstream-401 retry, and does not tell the user to sign in again over something they cannot fix. Sourcing a subject from the store rather than the request changed what invalidation can rely on, so the exchanged-token cache changed with it. The entry is now addressed by a slot key derived from the principal, plus the caller's own token when it presented one, with a fingerprint of the subject token and config stored beside the bearer and compared on every read. A mismatch reads as a miss and re-mints, so a rotated assertion or an edited server config cannot be served a bearer authorized under the old inputs, and two callers cannot receive each other's. Invalidation is a single delete of a key it can always compute, needing no store lookup on the recovery path. The upstream-401 invalidate-and-retry path was also gated on a truthy inbound subject token, which skipped recovery entirely for store-sourced calls. The gate is now mode-aware: token_exchange still requires an inbound token because it has nothing else to mint from, id_jag does not. oauth2_id_jag is also now selectable in the admin dashboard with its own field set, instead of being reachable only from config.yaml or the REST API. The auth-type selects drop antd list virtualization: at eleven options the last one no longer mounts, which is a scroll in a browser but makes the option unreachable to anything reading the rendered list.
d1488da to
721e0b3
Compare
TLDR
Problem this solves:
oauth2_id_jagwas not selectable in the admin dashboardHow it solves it:
Relevant issues
Linear ticket
Resolves LIT-4937
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Live proxy against a real Postgres, with a local Okta-shaped authorization server playing three roles at once: the OIDC provider LiteLLM signs into, the ID-JAG org AS for leg 1, and the upstream resource AS for leg 2. Every leg verifies the signature and claims of what it is handed, so a green run cannot be a rubber stamp. The upstream is a real MCP server that rejects any bearer it cannot verify and whose one tool reports the identity it resolved
The end user signs in once through LiteLLM SSO, which is what captures the assertion:
Before, at 440b1bc
The assertion is in the database, and the agent's tool call still cannot reach the upstream. No leg is even attempted:
The authorization server logged nothing for either call
After, at ddb0b89
The transcripts below were first captured at
ddb0b89688and re-run unchanged at9b05c30e88,which adds the assertion-store outage guard described under Changes
Same database, same rig, same request. The agent presents only its brokered LiteLLM key; there is no IdP token anywhere in the request:
The authorization server's own view of the two legs:
The upstream tool's log, which is the point of the whole flow:
A second user whose key is identical in every way except that they have never signed in stays fail-closed:
Assertion expiry, observed live
Re-running the same call about ninety minutes after the login returned a 412 rather than a tool
result, because the captured id_token had passed its one-hour expiry. That is the expiry branch
firing on a live proxy, and it is also the no-refresh limitation in concrete form: a fresh SSO
login restores the flow immediately
Dashboard
The ID-JAG server above was registered through the same REST shape the dashboard posts. To see the UI itself: go to http://localhost:4000/ui/?page=mcp-servers, click Add New MCP Server, pick an HTTP transport, and open the Authentication section. ID-JAG (Okta Cross App Access) is now in the auth-type list, and selecting it reveals the two-leg field set (org token endpoint, resource token endpoint, client id, client secret or private-key JWT, audience, resource indicator, subject token type, scopes). Submitting with neither a client secret nor a private key is refused inline
Type
🆕 New Feature
🐛 Bug Fix
Changes
_id_jagused to hard-fail when the request carried no identity token. It now resolves its subject through one helper: the caller's own inbound token when present, otherwise the assertionsso_assertion_storecaptured for that user at SSO login. That store already had a write side wired into the SSO callback and a retention gate that only persists while anoauth2_id_jagserver is registered;fetch_sso_identity_assertionsimply had no caller. It is injected as a collaborator rather than reached for as a module global, so the arm is testable without patchingThe subject's user always comes from the authenticated principal. There is deliberately no header or
userfield that lets a caller name someone else:end_user_idis documented as arbitrary caller-supplied text unlessvalidate_end_user_id_in_dbis set, so keying off it would let any agent key mint another person's upstream token. Delegated identity, where a shared agent key acts for an arbitrary end user, needs its own permission model and is not in this PRExpiry is judged by the reader, matching the store's stated contract. An expired assertion is a 412 telling the user to sign in again; there is no refresh yet, so the flow currently lives within the id_token's lifetime after a login. That is worth a follow-up but is strictly better than the current state, where it never works at all
Two things had to move with the subject, both of which would have shipped the feature broken:
invalidate_credentialscomputed its cache key from the request's inbound token and returned early without one, so a store-sourced bearer could never be evicted and a token the upstream had already rejected would be replayed until its TTL. It now resolves the subject the same way the arm doesThe upstream-401 invalidate-and-retry branch in
_call_mcp_toolwas gated on a truthy inboundsubject_token, so a store-sourced call skipped recovery entirely. The gate is now mode-aware:token_exchangestill requires an inbound token because it has nothing else to mint from,id_jagdoes notOn the dashboard,
oauth2_id_jagis now an auth type in both the create and edit forms with its own field component. Leg 1 rides the existingtoken_exchange_endpointcolumn; leg 2 and the private-key JWT material live in the credentials blob, matching what the resolver's adapter reads. Both auth-type selects also drop antd's list virtualization. At eleven options the eleventh no longer mounts, which is a scroll in a real browser but makes the last option invisible to anything reading the rendered list, including the existing tests for the two client-forwarded modesMost of the work after the first cut was in the invalidation path, which review rounds walked
through in stages. Recording it because the seam itself is small and this is where the substance is.
fetch_sso_identity_assertionissued a raw Prisma read with no guard, so a database failureescaped credential resolution unhandled. The damaging path was not the egress 500 but
invalidate_credentials, which runs inside the upstream-401 retry branch, so a blip duringrecovery turned a recoverable 401 into a 500.
DbSSOAssertionStore.fetchnow converts any storagefailure into a typed
AssertionStoreUnavailable, mirroring the roleTokenStoreUnavailablealready plays in
oauth_token_store.py, and the resolver maps it toupstream_unavailable(503)as a value. Deliberately not the 412 the other misses use: a missing assertion means sign in again,
an unreachable database means nothing the user can act on, and reporting it as an absence would
send people through a pointless re-login during an outage.
Containing that error then exposed the next layer, and the next, all in the invalidation path. The
short version of a long chain: invalidation derived its eviction key by re-reading the assertion
store, so it could not work precisely when the store was the failing component; keying off the
request alone let a recovered store recompute an identical key and hand a retry the bearer the
upstream had just rejected. Successive attempts to fix that by remembering keys (one per principal,
then a bounded set, then evicting on truncation) each closed one hole and left an adjacent one. A
stress harness over the property that matters, that no bearer minted before an invalidation can be
served after it, failed at 52 of 60 concurrent resolutions and pinpointed why: the truncation
eviction ran before the exchange had written anything.
So that mechanism is deleted rather than patched again. The cached entry is addressed by a slot key
derived from the principal, plus the caller's own token when it presented one, and the fingerprint
of the subject token and config is stored beside the bearer and compared on every read. A mismatch
reads as a miss and re-mints. Invalidation is a single delete of a key it can always compute,
needing no store lookup.
Gone with it: the alias map, the remembered-key set, its bound, the truncation eviction, and the
store dependency during invalidation, along with every ordering question they carried. The two
properties that mattered survive by different means: a rotated assertion or edited config re-mints
via the fingerprint rather than via the key, and two callers cannot receive each other's bearer
because a fingerprint mismatch is a miss rather than a hit. The stress harness now passes at 60 of
60 and a scaled-down version of it is committed.
One deliberate call worth flagging for review:
IdJagFormFields.tsximports antd and therefore needed aneslint-suppressions.jsonentry, the same mechanism every other antd component in that folder already relies on. The section renders inside an antd<Form>and depends onForm.Itemfor binding, edit-form prefill, and the invalidation reset, so going shadcn here would mean hand-wiring all three outside the form store. Adding a suppression does move that ratchet the wrong way; calling it out rather than burying itNot addressed here, by decision rather than oversight: the first-time consent grant in the ticket. Under Okta Cross App Access that authorization is an admin-granted relationship in the IdP, not a runtime prompt the gateway owns, so there is nothing for LiteLLM to build
Final Attestation