perf(mcp): cache _get_allowed_mcp_servers per request (PR #27847 follow-up) - #27862
Closed
mateo-berri wants to merge 17 commits into
Closed
perf(mcp): cache _get_allowed_mcp_servers per request (PR #27847 follow-up)#27862mateo-berri wants to merge 17 commits into
mateo-berri wants to merge 17 commits into
Conversation
For MCP servers configured with extra_headers: [Authorization], the gateway forwards the client token directly to the upstream. When that token is rejected (expired or invalid) the upstream returns 401, but the MCP SDK starts the SSE stream with 200 OK before calling handlers, so the 401 can't be returned mid-stream. Fix: add a pre-flight httpx probe in handle_streamable_http_mcp — before the SDK opens the session — so the gateway can still return HTTP 401 with WWW-Authenticate: Bearer authorization_uri=<gateway-discovery-url> when the upstream rejects the token. The probe fails-open (returns 200) on network errors so a transient hiccup does not block valid requests. Co-authored-by: Cursor <cursoragent@cursor.com>
…de effects - Extract forwarded_auth outside the pass-through server loop (was called N times for the same scope value) - Gather all upstream auth probes concurrently with asyncio.gather instead of sequentially; eliminates N×5 s worst-case latency - Switch probe from POST+initialize JSON-RPC body to HEAD request; HEAD carries the Authorization header so the upstream rejects invalid tokens with 401 but never allocates a session or writes an audit entry Co-authored-by: Cursor <cursoragent@cursor.com>
Replaces bare httpx.AsyncClient with the project-standard get_async_httpx_client(httpxSpecialProvider.MCP) to satisfy the ensure_async_clients_test code coverage check and avoid the +500 ms per-request overhead of creating a new client on every probe call. Co-authored-by: Cursor <cursoragent@cursor.com>
…eam_auth Moves the parallel upstream auth probe logic out of handle_streamable_http_mcp into a dedicated helper to satisfy Ruff PLR0915 (Too many statements > 50). Co-authored-by: Cursor <cursoragent@cursor.com>
…bypass _check_passthrough_upstream_auth was resolving user-supplied server names directly before authorization ran, letting any permitted LiteLLM key trigger an upstream HEAD probe to a server it was not allowed to use. Changes: - Call _get_allowed_mcp_servers inside the helper so only servers the caller's key is authorized for are probed. - Move the call site to after toolset scoping so the auth context is fully resolved before the probe list is built. - Thread user_api_key_auth into the helper signature (replaces the raw mcp_servers name list). Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
… probe _prepare_mcp_server_headers skips caller Authorization when the server uses OAuth client-credentials (M2M), but the pre-flight probe was still selecting those servers and forwarding the caller's raw token in the HEAD request. Exclude servers with has_client_credentials from the probe list to match the actual downstream header-preparation logic. Co-authored-by: Cursor <cursoragent@cursor.com>
Per RFC 9110, 401 means "go get new credentials." Mapping an upstream 403 to a gateway 401 causes OAuth clients to restart the authorization flow, obtain a fresh token with identical scopes, hit 403 again, and loop indefinitely. 401 from upstream → gateway 401 + WWW-Authenticate (re-authorize) 403 from upstream → gateway 403 (no WWW-Authenticate hint) Co-authored-by: Cursor <cursoragent@cursor.com>
… key The pre-flight upstream probe must not forward the caller's Authorization header when it could itself be the LiteLLM proxy API key. Restrict the probe to requests that supply x-litellm-api-key explicitly — only then is the Authorization header unambiguously the upstream OAuth token the caller wants forwarded.
Co-authored-by: Yassin Kortam <yassin@berri.ai>
Use AsyncHTTPHandler.post() and catch httpx.HTTPStatusError explicitly so the 401/403 we want to surface is not silently swallowed by the broad fail-open except Exception block. Avoids reaching into the handler's private client attribute, which would silently regress to fail-open if AsyncHTTPHandler is ever refactored.
Co-authored-by: Yassin Kortam <yassin@berri.ai>
AsyncHTTPHandler.post() calls raise_for_status() internally, so a real upstream 401/403 lands as httpx.HTTPStatusError. Add a test that exercises that specific exception path so a regression that swallows the error in the broad fail-open except Exception would be caught.
The pass-through auth probe added in PR #27847 introduced a second call to _get_allowed_mcp_servers per request — the SDK-dispatched list_tools / call_tool handlers already run the full key/team/end_user/agent/org permission chain, so probing doubled the work for every pass-through auth request. Prime a per-request ContextVar at the entry of handle_streamable_http_mcp and memoize _get_allowed_mcp_servers against (mcp_servers, client_ip). The probe and the SDK handlers then share a single resolution. Outside the request scope the cache is None and the function behaves exactly as before.
|
claude-bot seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Contributor
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Cache key omits
user_api_key_auth, risking stale results- Added id(user_api_key_auth) to the per-request cache key so the pre-merge probe entry no longer shadows the post-merge _list_mcp_tools resolution.
You can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit 7aba8af. Configure here.
| client_ip, | ||
| ) | ||
| if request_cache is not None and cache_key in request_cache: | ||
| return request_cache[cache_key] |
Contributor
There was a problem hiding this comment.
Cache key omits user_api_key_auth, risking stale results
Medium Severity
The cache_key is (mcp_servers, client_ip) but the function result also depends on user_api_key_auth, which is not part of the key. Within a single request the probe in _check_passthrough_upstream_auth populates the cache with the toolset-scoped auth, but the SDK handler path (_list_mcp_tools) calls _merge_toolset_permissions first, which can expand object_permission.mcp_servers with additional server IDs from resolved toolsets. The subsequent _get_allowed_mcp_servers call matches the same cache key and returns the stale result computed from the narrower pre-merge auth, silently dropping servers that the merged permissions would have granted.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 7aba8af. Configure here.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses the remaining Greptile concern on #27847: the pass-through auth probe was doubling the work of
_get_allowed_mcp_serversbecause the SDK-dispatchedlist_tools/call_toolhandlers run the same key/team/end_user/agent/org permission chain on every request.Prime a per-request
ContextVar(_mcp_request_allowed_servers_cache) at the entry ofhandle_streamable_http_mcpand memoize_get_allowed_mcp_serversagainst(mcp_servers, client_ip). Probe + handlers share one resolution. Outside the request scope the cache isNoneand behavior is unchanged.Tests cover both the dedupe-in-request and no-cache-outside-request paths.
Based on top of #27847.
Note
Medium Risk
Changes MCP request-time authorization/probing behavior and adds per-request caching; mistakes could alter access control decisions or how 401/403 are surfaced to clients.
Overview
Adds a per-request
ContextVarcache (primed inhandle_streamable_http_mcp) to memoize_get_allowed_mcp_serversby(mcp_servers, client_ip), so the pass-through upstream auth probe and subsequentlist_tools/call_tooldispatch reuse one permission resolution.Tightens StreamableHTTP pre-flight behavior by probing upstream
Authorizationonly whenx-litellm-api-keyis present (avoiding proxy-key leakage) and by propagating pre-header handler exceptions in_stream_mcp_asgi_responseso upstream 401/403 can be returned promptly. Adds tests for the probe helpers, cache dedupe/no-cache paths, and exception propagation.Reviewed by Cursor Bugbot for commit 7aba8af. Bugbot is set up for automated code reviews on this repo. Configure here.