fix(mcp): drop the cached per-user OAuth token when the credential row changes - #32302
Conversation
…w changes The v2 authorization_code chain Cached(Refreshing(V2PerUserTokenStore)) caches a positive token until its expires_at (or 300s without one), and CachedOAuthTokenStore.invalidate had no callers, so a re-authorization or revocation wrote the DB while egress kept serving the replaced token from the in-process cache until its TTL. LazyPerUserOAuthTokenStore now exposes invalidate, MCPServerManager threads it to the write side, and the three credential write sites (the OAuth callback, the Tools-tab persist endpoint, and the revoke endpoint) drop the cache entry after the row changes. The v2 refresher's own persist stays untouched; RefreshingTokenStore already feeds the rotated token back into the cache in the same fetch
Greptile SummaryThis PR fixes a cache-staleness bug where a per-user OAuth token stored under the v2
Confidence Score: 5/5Safe to merge; the change is additive and narrowly scoped to cache invalidation on credential writes, with no modifications to the read path or auth logic. All three write sites are correctly instrumented, invalidation is fail-open (cache errors never abort a DB write), the RecordNotFoundError branch is explicitly tested, and no existing test assertions were modified. The DualCache.async_delete_cache path was confirmed to clear both in-memory and Redis layers. No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py | Adds InvalidatableOAuthTokenStore protocol and CachedOAuthTokenStore.invalidate that delegates to backend.delete(); clean and minimal addition |
| litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py | Adds LazyPerUserOAuthTokenStore.invalidate mirroring the fetch fast-path and the _store_for_fetch slow-path; correctly handles cold-process Redis entry invalidation and in-flight local-fetch concurrency control |
| litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | Exposes invalidate_user_oauth_token_cache (fail-open, logs on error) and makes the per-user store injectable for tests; constructor change is backward-compatible |
| litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py | Inserts cache invalidation after a successful DB write in _store_per_user_token_server_side, correctly placed between the DB write and the v1 Redis warm-up |
| litellm/proxy/management_endpoints/mcp_management_endpoints.py | Adds cache invalidation to store_mcp_oauth_user_credential (re-auth) and delete_mcp_oauth_user_credential (revoke); revoke invalidates inside the if cred_to_delete is not None block but after both the try/except and RecordNotFoundError pass-through, so it fires even on concurrent deletes |
| tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py | Three new async tests cover invalidate on a cold process, after a prior fetch, and after a Redis chain is built; mock stores gain invalidate methods matching the new protocol |
| tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py | Two new tests verify the manager seam delegates to the injected store and swallows store errors; no modifications to existing tests |
| tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py | Two new tests assert that a successful DB write triggers the invalidate and a failed DB write skips it; new tests only, no modifications to existing ones |
| tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py | Three new tests cover store, revoke (happy path), and revoke-on-already-gone (RecordNotFoundError) invalidation; all new, no modifications to existing test assertions |
Reviews (4): Last reviewed commit: "test(mcp): cover invalidate on the redis..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
This comment was marked as outdated.
This comment was marked as outdated.
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile's review flagged that only the happy-path delete asserted the invalidate; a refactor moving the call inside the try block would silently skip the cache drop when the row was already deleted by a concurrent request while the cache still held the revoked token. The new test fails on exactly that mutation
Merging this PR will improve performance by 31.51%
|
| Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|
| ⚡ | test_completion_with_tools |
4.2 ms | 3.2 ms | +31.51% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing litellm_mcp_oauth_cache_invalidate_on_reauth (9e331d8) with litellm_internal_staging (db24027)
Codecov flagged the redis fast path of LazyPerUserOAuthTokenStore.invalidate as unexercised; the existing invalidate tests only ran the no-redis chain. The new test builds the redis chain via a fetch and asserts a subsequent invalidate reaches the same store instance without a rebuild
|
bugbot run |
|
@greptileai rereview |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 9e331d8. Configure here.
mateo-berri
left a comment
mateo-berri
left a comment
There was a problem hiding this comment.
codecov/patch — 73.33% of diff hit (target 74.11%)
Is the coverage not high enough or is this a stale metric?
Stale metric, ran locally and fully covered |
|
Note on the codecov/patch red: it is a reporting artifact, not missing coverage. A local coverage run over the four mapped test files shows zero missing executable lines in every diff region, and the tests that exercise the lines codecov flags (for example the revoke-endpoint invalidate) ran and passed inside the exact CI shard whose upload codecov ingested (proxy-endpoints, see the run log for test_delete_mcp_oauth_user_credential_invalidates_cached_token). Codecov's own PR comment and its compare API also disagree with each other about which files are missing, and the compare view marks a module-top import and a docstring line as missed while function bodies in the same module read as hit, which no coherent coverage report can produce. A shard rerun and remerge did not change the number, so this is stuck on codecov's side |
…w changes (BerriAI#32302) * fix(mcp): drop the cached per-user OAuth token when the credential row changes The v2 authorization_code chain Cached(Refreshing(V2PerUserTokenStore)) caches a positive token until its expires_at (or 300s without one), and CachedOAuthTokenStore.invalidate had no callers, so a re-authorization or revocation wrote the DB while egress kept serving the replaced token from the in-process cache until its TTL. LazyPerUserOAuthTokenStore now exposes invalidate, MCPServerManager threads it to the write side, and the three credential write sites (the OAuth callback, the Tools-tab persist endpoint, and the revoke endpoint) drop the cache entry after the row changes. The v2 refresher's own persist stays untouched; RefreshingTokenStore already feeds the rotated token back into the cache in the same fetch * test(mcp): pin cache invalidation on the revoke already-gone branch Greptile's review flagged that only the happy-path delete asserted the invalidate; a refactor moving the call inside the try block would silently skip the cache drop when the row was already deleted by a concurrent request while the cache still held the revoked token. The new test fails on exactly that mutation * test(mcp): cover invalidate on the redis-backed lazy store path Codecov flagged the redis fast path of LazyPerUserOAuthTokenStore.invalidate as unexercised; the existing invalidate tests only ran the no-redis chain. The new test builds the redis chain via a fetch and asserts a subsequent invalidate reaches the same store instance without a rebuild
2 participants
Relevant issues
Found while live-verifying #31362: overwriting a stored per-user OAuth token did not change what the gateway sent upstream until the proxy was restarted. Sibling of the preflight work in #31362; this PR fixes the write side, that one the connect-time probe
Linear ticket
Pre-Submission checklist
make test-unit@greptileaiand received a Confidence Score of at least 4/5Screenshots / Proof of Fix
Reproduced on latest
litellm_internal_staging(2f0cdb3): the v2 per-user token chainCached(Refreshing(V2PerUserTokenStore))caches a positive token until itsexpires_at(300s without one), andCachedOAuthTokenStore.invalidatehas zero callers, so after a user re-authorizes or revokes, egress keeps serving the replaced token from the in-process cache. The revocation case is the bad one: a credential the user deleted keeps flowing to the upstream until the old token's expiryLive before/after on a Postgres-backed proxy (
python litellm/proxy/proxy_cli.py --config qa_config.yaml --port 33514 --detailed_debug --use_v2_migration_resolver). There is no third-party provider we can revoke tokens against on demand, so the upstream is a minimal local MCP server that acceptsBearer fresh-tokenand rejects everything else with an RFC 6750 challenge, and it logs theAuthorizationheader of every request; everything on the client side of the proxy is exactly what an end user would runSetup:
curl http://localhost:33514/user/newfor usersqa-cache-user-aandqa-cache-user-b, thencurl http://localhost:33514/key/generateper user with"object_permission": {"mcp_servers": ["dd5ff6fc6cab4acc3b53e1bc24410e2a"]}. Tokens are stored and revoked through the proxy's own endpoints (the same ones the Tools tab uses), so the write happens in the proxy process where the cache livesScenario A, re-authorization. Store a stale token, list tools once (this warms the cache; the upstream rejects the stale token so the list is empty), then store a fresh token the way a re-auth would and list again
Before (unfixed
litellm_internal_staging): the fresh token is in the DB but the list is still empty, and the upstream log shows the proxy sent the stale token again after the re-authAfter (this PR): the same fourth step returns the tool immediately
Scenario B, revocation (second user). Store a valid token, list tools once (tool visible, cache warm), revoke it via the DELETE endpoint, list again
Before: the revoked credential keeps working; the tool is still listed and the upstream log shows the deleted token still being sent
After: the revoked credential stops flowing on the next request; the list is empty and the upstream receives no authorized call for that user at all
First-time auth is unaffected in both runs (a
Noneread is never cached, so the first stored token is picked up immediately; scenario A step 2 and scenario B step 2 show it)Type
🐛 Bug Fix
Changes
oauth_token_store.pyadds anInvalidatableOAuthTokenStoreprotocol (fetch plus invalidate) andper_user_oauth_store.pyimplements it onLazyPerUserOAuthTokenStore, mirroringfetch's build-and-bookkeeping path so an invalidate on a cold process still drops a shared Redis entry written by another replica, and aDualCacheTokenCacheBackend.deleteclears both the in-process layer and RedisMCPServerManagernow holds the per-user store it hands the resolver (injectable for tests) and exposesinvalidate_user_oauth_token_cache(user_id, server_id); a cache-drop failure is logged and never raised, since the DB write already succeeded and the TTL remains the backstopThe three write sites that change the credential row outside the chain now call it:
_store_per_user_token_server_side(OAuth callback, code exchange and refresh),store_mcp_oauth_user_credential(Tools-tab persist), anddelete_mcp_oauth_user_credential(revoke), matching the invalidation the BYOK twin endpoints already do via_invalidate_byok_cred_cache. The v2 refresher's persist is deliberately not a call site:RefreshingTokenStorereturns the rotated token into the samefetchthat caches it, so that path is already self-consistentRegression tests: the lazy store's invalidate threading (
test_per_user_oauth_store.py), the manager seam and its fail-open error handling (test_mcp_server_manager.py), and one test per write site asserting the invalidate fires with the right(user_id, server_id)(test_mcp_management_endpoints.py,test_discoverable_endpoints.py), plus one asserting a failed DB write does not invalidate, and one asserting the revoke path still invalidates when the row is already gone (a concurrent delete) so the cache drop cannot be refactored inside the try block. All fail on unfixed codeNote
Medium Risk
Touches MCP OAuth egress credential caching and token persistence paths; incorrect invalidation could cause extra upstream auth or brief stale tokens, but changes are scoped to cache drops after known writes with TTL as backstop.
Overview
Fixes a bug where re-authorizing or revoking a per-user MCP OAuth credential updated the DB but egress kept sending the old token from the v2
CachedOAuthTokenStorechain until cache TTL or restart.Adds an
InvalidatableOAuthTokenStoreprotocol andinvalidateonLazyPerUserOAuthTokenStore(including cold-start / cross-replica Redis cases).MCPServerManager.invalidate_user_oauth_token_cachedelegates to that store; failures are logged only (DB write already succeeded).Write paths now call invalidate after a successful credential change: OAuth callback
_store_per_user_token_server_side(not on DB failure), Tools-tab store, and delete revoke (including when the row is already gone). Then the existing v1 Redis warm in the callback still runs as before.Reviewed by Cursor Bugbot for commit 9e331d8. Bugbot is set up for automated code reviews on this repo. Configure here.