Skip to content

feat(mcp): source the ID-JAG subject from the user's stored SSO assertion - #35147

Merged
yassin-berriai merged 1 commit into
litellm_internal_stagingfrom
litellm_mcp_id_jag_subject_sourcing
Jul 31, 2026
Merged

feat(mcp): source the ID-JAG subject from the user's stored SSO assertion#35147
yassin-berriai merged 1 commit into
litellm_internal_stagingfrom
litellm_mcp_id_jag_subject_sourcing

Conversation

@yassin-berriai

@yassin-berriai yassin-berriai commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • ID-JAG only worked if the caller sent its own IdP token
  • An agent with a brokered LiteLLM key always got a 412
  • The SSO assertion was stored per user but never read back
  • A store-sourced bearer could never be evicted after a 401
  • oauth2_id_jag was not selectable in the admin dashboard

How it solves it:

  • The ID-JAG arm falls back to the user's stored SSO assertion
  • The user comes from the authenticated key, never from a header
  • Invalidation and the 401 retry resolve the same subject
  • ID-JAG gets a dashboard auth type with its own fields

Relevant issues

Linear ticket

Resolves LIT-4937

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to 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:

$ curl -s -L -o /dev/null -w "final status %{http_code} at %{url_effective}\n" \
    "http://localhost:4939/sso/key/generate"
final status 200 at http://localhost:4939/ui/?login=success

$ docker exec lit4937-pg psql -U litellm -d litellm -t \
    -c 'SELECT user_id, length(assertion_b64) FROM "LiteLLM_SSOIdentityAssertion";'
 alice@acme.example |      592

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:

$ curl -s -X POST http://localhost:4939/mcp/ \
    -H "x-litellm-api-key: Bearer $ALICE_KEY" \
    -H "Accept: application/json, text/event-stream" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}'
data: {"jsonrpc":"2.0","id":1,"result":{"_meta":{"litellm.ai/server_outcomes":
      {"jira":{"status":"internal","http_status":412}}},"tools":[]}}

$ curl -s -X POST http://localhost:4939/mcp/ \
    -H "x-litellm-api-key: Bearer $ALICE_KEY" \
    -H "Accept: application/json, text/event-stream" \
    -d '{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"jira-whoami","arguments":{}}}'
data: {"jsonrpc":"2.0","id":2,"result":{"content":[{"type":"text",
      "text":"Error: Tool 'whoami' not found"}],"isError":true}}

The authorization server logged nothing for either call

After, at ddb0b89

The transcripts below were first captured at ddb0b89688 and re-run unchanged at 9b05c30e88,
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:

$ curl -s -X POST http://localhost:4939/mcp/ \
    -H "x-litellm-api-key: Bearer $ALICE_KEY" \
    -H "Accept: application/json, text/event-stream" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"jira-whoami","arguments":{}}}'
data: {"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text",
      "text":"upstream saw acting_user=alice@acme.example"}],"isError":false}}

The authorization server's own view of the two legs:

LEG1 OK: id_token(alice@acme.example) -> ID-JAG assertion  [audience=urn:upstream]
LEG2 OK: ID-JAG(alice@acme.example) -> upstream access token for alice@acme.example

The upstream tool's log, which is the point of the whole flow:

tool=whoami acting_user=alice@acme.example

A second user whose key is identical in every way except that they have never signed in stays fail-closed:

$ curl -s -X POST http://localhost:4939/mcp/ \
    -H "x-litellm-api-key: Bearer $BOB_KEY" \
    -H "Accept: application/json, text/event-stream" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"jira-whoami","arguments":{}}}'
data: {"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text",
      "text":"Error: precondition required: ID-JAG requires an IdP identity assertion for this
       user and none is stored. Sign in through LiteLLM SSO so the gateway captures one."}],
       "isError":true}}

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

$ date '+%H:%M:%S'; # login was at 14:56:30, id_token ttl 3600s
16:32:22

$ curl -s -X POST http://localhost:4939/mcp/ -H "x-litellm-api-key: Bearer $ALICE_KEY" ... tools/list
{"jira":{"status":"internal","http_status":412}}

$ curl -s -L "http://localhost:4939/sso/key/generate"          # sign in again
$ curl -s -X POST http://localhost:4939/mcp/ ... tools/list
{"jira":{"status":"ok","tool_count":1}}
$ curl -s -X POST http://localhost:4939/mcp/ ... tools/call jira-whoami
"upstream saw acting_user=alice@acme.example"

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_jag used 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 assertion sso_assertion_store captured 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 an oauth2_id_jag server is registered; fetch_sso_identity_assertion simply had no caller. It is injected as a collaborator rather than reached for as a module global, so the arm is testable without patching

The subject's user always comes from the authenticated principal. There is deliberately no header or user field that lets a caller name someone else: end_user_id is documented as arbitrary caller-supplied text unless validate_end_user_id_in_db is 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 PR

Expiry 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_credentials computed 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 does

The upstream-401 invalidate-and-retry branch in _call_mcp_tool was gated on a truthy inbound subject_token, so a store-sourced call skipped recovery entirely. 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

On the dashboard, oauth2_id_jag is now an auth type in both the create and edit forms with its own field component. Leg 1 rides the existing token_exchange_endpoint column; 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 modes

Most 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_assertion issued a raw Prisma read with no guard, so a database failure
escaped 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 during
recovery turned a recoverable 401 into a 500. DbSSOAssertionStore.fetch now converts any storage
failure into a typed AssertionStoreUnavailable, mirroring the role TokenStoreUnavailable
already plays in oauth_token_store.py, and the resolver maps it to upstream_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.tsx imports antd and therefore needed an eslint-suppressions.json entry, the same mechanism every other antd component in that folder already relies on. The section renders inside an antd <Form> and depends on Form.Item for 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 it

Not 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

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

Comment thread litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py Outdated
@greptile-apps

greptile-apps Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds stored SSO assertions as an ID-JAG subject source and updates the corresponding retry, cache, error handling, tests, and dashboard configuration

  • Resolves ID-JAG subjects from either an inbound identity token or the authenticated user's stored SSO assertion
  • Reworks exchanged-token cache addressing around principal slots and input fingerprints
  • Adds typed assertion-store outage handling and stored-assertion expiry checks
  • Enables invalidate-and-retry for store-sourced ID-JAG tool calls
  • Adds create and edit dashboard fields for ID-JAG authentication

Confidence Score: 3/5

The 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

Important Files Changed

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

codecov Bot commented Jul 29, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@codspeed-hq

codspeed-hq Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_mcp_id_jag_subject_sourcing (721e0b3) with litellm_internal_staging (551e5d0)1

Open in CodSpeed

Footnotes

  1. No successful run was found on litellm_internal_staging (ae74b06) during the generation of this report, so 551e5d0 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report.

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_id_jag_subject_sourcing branch 3 times, most recently from d5c4c26 to 9b05c30 Compare July 29, 2026 23:29
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai the assertion-store outage finding was correct and is fixed in 9b05c30e88.

fetch_sso_identity_assertion issued a raw Prisma find_unique with no guard, so a database failure escaped credential resolution as an unhandled error. Worse than the 500 on the egress path: invalidate_credentials runs inside the upstream-401 invalidate-and-retry branch, so a blip there turned a recoverable 401 into a 500 during recovery.

Fixed by mirroring the convention the sibling per-user OAuth store already uses. DbSSOAssertionStore.fetch converts any storage failure into a typed AssertionStoreUnavailable, the same role TokenStoreUnavailable plays in oauth_token_store.py, and the resolver maps it to upstream_unavailable (503) as a value rather than letting it raise.

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 upstream_unavailable without calling the IdP, invalidation survives an outage rather than propagating, the live store converts a driver error into the typed failure, and an absent row still reads as None so a genuine never-signed-in user keeps getting the 412.

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 expires_at, which would otherwise raise a TypeError comparing against an aware now.

Note the review is pinned to ddb0b89688, which no longer exists on the branch. Please re-review the current head 9b05c30e88.

Comment thread litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py Outdated
@veria-ai

veria-ai Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@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 log_statement='all', five store-sourced ID-JAG tool calls produced:

assertion-table reads : 5
total SQL statements  : 6
IdP exchange legs     : 0   (exchanged bearer served from cache)

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 find_unique against a one-row-per-user table on a path that then makes an outbound HTTPS call to the upstream MCP server. It dominates the query count, not the latency budget.

The blast radius is a brand-new opt-in auth mode. Before this PR oauth2_id_jag could not complete a store-sourced call at all, so there is no existing traffic whose cost profile changes.

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 expires_at, so a TTL cache that outlives a re-login would keep serving a retired id_token and return 412 to a user who just signed in specifically to clear a 412. Getting that right means invalidating on persist_sso_identity_assertion and reasoning about the cross-pod case where the login lands on a different replica than the read; that is a design with its own live verification, not a two-line wrapper added at the end of a long PR.

The shape I would propose for the follow-up: a short-TTL in-memory cache in front of DbSSOAssertionStore, busted on persist, TTL short enough that cross-pod staleness self-heals within one interval. Happy to do it here instead if you would rather it not ship as a separate change; my read is that it is safer as its own reviewed unit.

Please re-review the current head 9b05c30e88.

Comment thread litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py Outdated
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_id_jag_subject_sourcing branch from 9b05c30 to f3e874d Compare July 29, 2026 23:45
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai you were right, and the finding is a direct consequence of my previous fix. Fixed in f3e874d5a6.

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, invalidate_credentials runs while the store is unreachable so it evicts nothing and returns quietly, the store recovers before the retry resolves, resolution recomputes an identical cache key and hands the retry the bearer the upstream just rejected. One retry burned and the bad token cached until TTL.

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 (tenant, subject, server) alias. Invalidation evicts by that alias. The alias holds no secret; it addresses a cache entry rather than containing one, and it is bounded by the same max-size constant as the token cache. A fresh lookup remains only as the fallback for when nothing was recorded, such as another replica or an aged-out entry.

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 (test_invalidate_credentials_for_id_jag_is_a_noop_without_a_caller_token) caught it rather than my new one. Recording is now restricted to the store-sourced case, which is the only case that needs it.

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, ruff format clean, strict-rule gate passing against the current merge-base.

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 f3e874d5a6.

Comment thread litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py Outdated
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_id_jag_subject_sourcing branch from f3e874d to 5981d02 Compare July 30, 2026 00:00
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai correct again, and fixed in 5981d02daa.

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. misc / Run tests fails at test_openapi_compliance.py::TestRequestCompliance::test_content_schema_uses_discriminator, which I ran on a clean checkout of litellm_internal_staging with none of my code and it fails there identically; my diff touches no interactions or OpenAPI surface. secret-scan hit its 5-minute job timeout during the full-history ggshield scan after its assertion step passed, and nine files already on staging carry the same PEM-header test literals mine does, so it is not a finding on this diff; I re-ran that job.

Please re-review the current head 5981d02daa.

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_id_jag_subject_sourcing branch from 5981d02 to 940669d Compare July 30, 2026 00:09
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai this finding splits into two parts, and they need different answers. Head is now 940669d8e0.

On the lost update, I do not think it can occur, and I checked rather than asserting it. _remember_id_jag_key is a plain def, not a coroutine, and it contains no await; InMemoryCache.get_cache/set_cache/delete_cache are all synchronous (verified with inspect.iscoroutinefunction, all False). The read, the membership check and the write therefore run as one uninterrupted synchronous block on the event loop, and no other resolution can interleave between them. A lock would add contention without changing any observable behavior. If this ever moved to a threaded executor or a cross-process cache the analysis would change, and that is worth a comment if you would like one.

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 940669d8e0.

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_id_jag_subject_sourcing branch from 940669d to 5d8da16 Compare July 30, 2026 00:33
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai I stopped patching this mechanism and replaced it. Head is 5d8da16990.

Your last finding split into two parts. On the lost update I still believe the read-modify-write is safe: _remember_id_jag_key was a plain def with no await and all three InMemoryCache operations are synchronous (checked with inspect.iscoroutinefunction), so nothing could interleave. But rather than argue that, I wrote a stress harness for the property that actually matters, that no bearer minted before an invalidation can ever be served after it, and drove 60 interleaved resolutions at distinct subject tokens through it.

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 get_or_compute had written anything, so it evicted nothing and the bearer landed afterwards under a key nothing was tracking. My previous fix did not hold. Worth noting the first version of that harness was itself vacuous, all 60 coroutines raced onto one shared mutable token and collapsed to a single key through single-flight, and it reported a cheerful pass; the fix was to have the store issue a distinct assertion per call.

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 5d8da16990.

@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_id_jag_subject_sourcing branch 3 times, most recently from 4f53812 to d1488da Compare July 30, 2026 00:51
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai the review is still anchored to 5d8da16990; head is now d1488da97f. Two commits since, neither changing behavior.

The first was a defect I introduced while restructuring and want to name rather than let pass quietly. My scripted edit left invalidate_credentials, _invalidate_id_jag, _authorization_code, _client_credentials and _token_exchange each defined twice. The copies were byte-identical, so Python bound the later one and behavior was unchanged: 899 tests passed and the live proxy worked. Only ruff check and basedpyright's reportRedeclaration saw it, and CI's lint job flagged it on a8a420adba at the same time I was fixing it locally. Duplicates removed, and the surviving invalidate_credentials docstring no longer claims it resolves a subject, which stopped being true with the restructure.

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 litellm_internal_staging, exactly equal, with no redeclarations. In the manager test I had left a parameter unannotated; typing it as MCPAuthType clears three of a four-error delta, and the remaining one is reportPrivateUsage from importing the private helper under test, which the file already does 269 times.

Also folded in from the previous round: the (fingerprint, token) cache entry is now validated through a TypeAdapter rather than hand-rolled isinstance narrowing, matching how this repo asks untyped boundaries to be handled.

Local state on d1488da97f: ruff clean, all changed litellm/ files formatted, strict-rule gate passing against the current merge-base, 899 tests passing, and the concurrency stress harness still at 60 of 60 with no pre-invalidation bearer surviving.

Please re-review the current head d1488da97f.

…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.
@yassin-berriai
yassin-berriai force-pushed the litellm_mcp_id_jag_subject_sourcing branch from d1488da to 721e0b3 Compare July 30, 2026 01:00
@yassin-berriai
yassin-berriai enabled auto-merge (squash) July 31, 2026 15:25
@yassin-berriai
yassin-berriai merged commit 0e9a624 into litellm_internal_staging Jul 31, 2026
79 checks passed
@yassin-berriai
yassin-berriai deleted the litellm_mcp_id_jag_subject_sourcing branch July 31, 2026 17:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants