Skip to content

perf(mcp): cache SSO identity assertion reads on the ID-JAG path - #39348

Merged
yassin-berriai merged 7 commits into
litellm_internal_stagingfrom
litellm_lit_5060_cache_sso_assertion_reads
Sep 4, 2026
Merged

yassin-berriai merged 7 commits into
litellm_internal_stagingfrom
litellm_lit_5060_cache_sso_assertion_reads

Conversation

@devin-ai-integration

@devin-ai-integration devin-ai-integration Bot commented Sep 2, 2026 •

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • Every store-sourced ID-JAG tool call reads the SSO assertion row from Postgres
  • The exchanged bearer is already cached, so that read was the only remaining DB work per call

How it solves it:

  • Short-TTL in-memory cache in front of the assertion read (MCP_SSO_ASSERTION_CACHE_TTL_SECONDS, default 60s)
  • The SSO login write busts the cached entry, so a re-login is visible immediately on that pod
  • A fetch that was mid-read during any login skips caching, so it cannot repopulate the previous assertion
  • Misses (no row, undecryptable, unparseable) are not cached, so a first login is never masked

User Flow

Before: an agent calling an oauth2_id_jag MCP server with a LiteLLM key gets a DB read on every call

  1. The user signs in at https://litellm-domain/sso/key/generate and the gateway captures their IdP id_token
  2. The agent sends POST https://litellm-domain/mcp/ with tools/call for a tool on that server, five times in a row
  3. Every call succeeds, and the operator's Postgres statement log shows five reads of the assertion table for those five calls

After: the same five calls hit Postgres once

  1. The user signs in at https://litellm-domain/sso/key/generate and the gateway captures their IdP id_token
  2. The agent sends POST https://litellm-domain/mcp/ with tools/call for a tool on that server, five times in a row
  3. Every call succeeds, and the Postgres statement log shows one read of the assertion table; the other four calls are served from memory
  4. If the user signs in again, the very next call already uses the new assertion (same pod), or within 60 seconds (login landed on another pod)

Relevant issues

Follow-up to #35147, where this was measured and deferred

Linear ticket

Resolves LIT-5060

Pre-Submission checklist

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

  • I have added meaningful tests
  • The handful of test files covering my change pass locally, e.g. uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*, make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
  • My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
  • 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

A live run needs an enterprise IdP that supports RFC 8693 ID-JAG exchange plus an EMA-protected MCP upstream, which this environment does not have, so the section below is the runbook a reviewer with that setup can follow. It is the same measurement used in #35147 and the ticket. Statement counts below are the expected values, not a captured run

Shared setup: a proxy with SSO configured and one MCP server with auth_type: oauth2_id_jag, Postgres started with log_statement='all'. Sign in once through /sso/key/generate so an assertion is stored, then run

for i in $(seq 1 5); do
  curl -s -X POST http://localhost:4000/mcp/ \
    -H "Authorization: Bearer $LITELLM_KEY" -H "Content-Type: application/json" \
    -H "Accept: application/json, text/event-stream" \
    -d '{"jsonrpc":"2.0","id":'$i',"method":"tools/call","params":{"name":"<server>-<tool>","arguments":{}}}' \
    -o /dev/null -w "%{http_code}\n"
done
grep -c 'FROM "public"."LiteLLM_SSOIdentityAssertion"' /var/log/postgresql/postgresql-*.log

Before (2ce4e3f)

  1. Five 200 lines from the loop
  2. grep -c prints 5: one assertion-table read per call

After (205e943)

  1. Five 200 lines from the loop
  2. grep -c prints 1: the first call reads, the other four are served from the in-process cache
  3. Sign in again through /sso/key/generate, run one more call: the count goes to 2, since the login busted the cached entry

Type

🚄 Infrastructure

Caveats (if any)

Medium

  • Cross-pod staleness of up to MCP_SSO_ASSERTION_CACHE_TTL_SECONDS (60s) after a login on another replica
    • Only bites a user who re-logs in to clear a 412 and immediately calls through a different pod; it self-heals within one TTL

Low

  • Deletion on user removal (LIT-4669) and the refresh write (LIT-5059) are not implemented yet; when they land they must call invalidate(user_id) on the same SSOAssertionCache

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

Link to Devin session: https://app.devin.ai/sessions/4ce3dfeb377f4a91bcba8ea4c24ca5c6
Open in Devin Desktop: https://app.devin.ai/desktop/session/4ce3dfeb377f4a91bcba8ea4c24ca5c6?variant=devin
Requested by: @yassin-berriai

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

  • Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
  • Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

  • Disable automatic comment, CI, and merge conflict monitoring

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

@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

PR #39348 (BerriAI/litellm, author devin-ai-integration[bot]) has no enterprise label (labels: none), so it is out of scope. No changes made; no risk label applied and no Linear routing performed.

@codspeed

codspeed Bot commented Sep 2, 2026 •

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_lit_5060_cache_sso_assertion_reads (4630307) with litellm_internal_staging (a53c550)

Open in CodSpeed

@greptile-apps

greptile-apps Bot commented Sep 2, 2026 •

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR adds a short-lived process-local cache for SSO identity assertions used by ID-JAG credential resolution.

  • Caches successfully decoded assertion reads while leaving misses uncached.
  • Invalidates the cached user after login persistence.
  • Uses a process-wide epoch to prevent pre-login reads from repopulating stale entries.
  • Adds cache-hit, invalidation-race, miss, TTL, and dependency-injection tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py Adds assertion caching and replaces finite per-user generation markers with a non-expiring process-wide epoch guard.
litellm/constants.py Defines the configurable SSO assertion cache TTL with a 60-second default.
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py Adds coverage for cache hits, login invalidation, overlapping reads, misses, expiration, and injected caches.
tests/test_litellm/test_router.py Renames an existing router test without changing its behavior.

Reviews (4): Last reviewed commit: "fix(mcp): use a process-wide epoch for s..." | Re-trigger Greptile

@codecov

codecov Bot commented Sep 2, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

yassin-berriai and others added 2 commits September 3, 2026 22:14
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

@greptileai please re-review after 709ee4c, which adds per-user generation guards against stale repopulation plus a regression test

@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

Fixed in 709ee4c: fetch records a per-user generation before the DB read and only caches if login did not bump it

@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

TTL test now persists through a second cache instead of mutating stored. The mock line is exactly 120 chars, ruff format keeps it

…separate namespaces

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

@greptileai please re-review after bca1822, which moves generation markers into their own cache so user ids cannot collide with them

…passes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

@greptileai please re-review after 205e943, which replaces the per-user markers with a process-wide epoch counter

…_assertion_reads

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@yassin-berriai
yassin-berriai merged commit 1add1b4 into litellm_internal_staging Sep 4, 2026
79 of 80 checks passed
@yassin-berriai
yassin-berriai deleted the litellm_lit_5060_cache_sso_assertion_reads branch September 4, 2026 00: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.

2 participants