test(e2e): cover MCP gateway auth headers and per-server max concurrency - #33102
Open
tin-berri wants to merge 17 commits into
Open
test(e2e): cover MCP gateway auth headers and per-server max concurrency#33102tin-berri wants to merge 17 commits into
tin-berri wants to merge 17 commits into
Conversation
Contributor
Greptile SummaryThis PR introduces the first live e2e coverage of the MCP gateway, adding
Confidence Score: 5/5Safe to merge; changes are entirely within the e2e test infrastructure and do not touch any production code paths. All changes are additive test scaffolding — a new compose service, a stub server, a test client, and three tests. No production code is modified. The previously flagged issues have all been addressed in follow-up commits. The control-server max_in_flight == burst_size assertion in tests/e2e/mcp/test_mcp_gateway_e2e.py has a timing dependency worth reviewing.
|
| Filename | Overview |
|---|---|
| tests/e2e/mcp/test_mcp_gateway_e2e.py | New e2e suite covering MCP auth headers and per-server concurrency cap; the control server's strict max_in_flight == burst_size assertion is timing-sensitive and may be intermittently fragile on slow CI. |
| tests/e2e/mcp/mcp_client.py | New MCP test client routing management calls via Gateway transport and MCP protocol calls via the official mcp SDK; asyncio.run() usage in ThreadPoolExecutor threads is safe since each thread gets its own loop. |
| tests/e2e/mcp/stub/stub_server.py | Deterministic FastMCP upstream with echo, slow_echo (per-marker in-flight tracking), and stats tools; counter atomicity relies correctly on asyncio single-threaded scheduling with no await between increments. |
| tests/e2e/docker-compose.yml | Adds mcp-stub service with healthcheck and litellm depends_on, addressing the previously flagged startup ordering issue. |
| tests/e2e/models.py | Adds McpServerCreateBody and McpServerInfo Pydantic models; all optional fields correctly default to None for the response shape. |
| tests/e2e/mcp/stub/Dockerfile | Minimal Python 3.13-slim image pinning mcp==1.26.0; intentionally independent of the e2e-dev client range per team decision documented in prior review thread. |
| tests/e2e/mcp/conftest.py | pytest.importorskip guard correctly prevents conftest loading (and downstream mcp imports) when the mcp SDK is absent, matching the playwright pattern. |
Reviews (5): Last reviewed commit: "refactor(e2e): keep the concurrency test..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
Author
Contributor
Contributor
Author
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Contributor
Author
1 similar comment
Contributor
Author
27 tasks
tin-berri
force-pushed
the
litellm_mcp_e2e_tests
branch
from
July 15, 2026 02:34
d3c9670 to
f503c3a
Compare
… on mcp-stub health
…with a parametrized auth-header matrix
… and cover prompt/resource primitives
…he conformance-ready foundation
…appy path it exercises
…Auth flow under both ingress headers
tin-berri
force-pushed
the
litellm_mcp_e2e_tests
branch
from
July 16, 2026 19:51
a03750f to
9b144f1
Compare
…uard, not a full dance The Authorization-header variant cannot complete the interactive round-trip: once the host finishes the OAuth dance the SDK carries the minted upstream token in Authorization, evicting the LiteLLM key the gateway needs to resolve the per-user token. Assert the achievable, correct contract instead: a key in Authorization on a gateway-managed oauth2 server gets the 401 challenge, not a masked empty tool list. Full list-and-call stays on the x-litellm-api-key variant. Red on stock main (masking); flips green with the mode-aware gateway challenge fix.
… and call tools
Positive companion to the Authorization challenge-guard: once a user has
authorized (per-user token stored via the x-litellm-api-key dance), a plain
Authorization: Bearer <key> session needs no dance and the gateway serves
tools using the stored token. Seeds the token via the dance, then lists and
calls over an Authorization-key session, and asserts the upstream received the
stored per-user token (not the key) with no key leak. Claims the previously
unclaimed mcp.{list_tools,call_tool}.bearer.succeeds cells. Green on stock.
| @@ -0,0 +1,310 @@ | |||
| """Deterministic MCP upstreams for the mcp e2e suite. | |||
Collaborator
There was a problem hiding this comment.
use an actual oauth mcp server please (maybe the ones we're having issues with?)
| command: ["--config", "/app/config.yaml", "--port", "4000"] | ||
|
|
||
| # throwaway db | ||
| mcp-stub: |
|
|
||
| def _build_app(upstream: httpx.AsyncClient, litellm_key: str) -> Starlette: | ||
| async def forward(request: Request) -> StreamingResponse: | ||
| url = httpx.URL(path=request.url.path, query=request.url.query.encode()) |
Collaborator
There was a problem hiding this comment.
no httpx, use the tests/e2e/claude.md
Comment on lines
+78
to
+108
| def test_list_and_call_tools_with_x_litellm_api_key_header( | ||
| self, client: McpClient, resources: ResourceManager | ||
| ) -> None: | ||
| marker = unique_marker() | ||
| alias = f"e2emcpauthcode{marker}" | ||
| created = client.create_server( | ||
| McpServerCreateBody( | ||
| alias=alias, | ||
| url=MCP_STUB_OAUTHUSER_URL, | ||
| allow_all_keys=False, | ||
| auth_type="oauth2", | ||
| oauth2_flow="authorization_code", | ||
| authorization_url=MCP_STUB_AUTHORIZE_BROWSER_URL, | ||
| token_url=MCP_STUB_TOKEN_URL, | ||
| credentials=McpServerCredentials( | ||
| client_id=MCP_STUB_OAUTH_USER_CLIENT_ID, | ||
| client_secret=MCP_STUB_OAUTH_USER_CLIENT_SECRET, | ||
| ), | ||
| ) | ||
| ) | ||
| resources.defer(lambda: client.delete_server(created.server_id)) | ||
|
|
||
| stored = client.server_info(created.server_id) | ||
| assert stored.auth_type == "oauth2" | ||
| assert stored.oauth2_flow == "authorization_code" | ||
| assert stored.authorization_url == MCP_STUB_AUTHORIZE_BROWSER_URL | ||
| assert stored.token_url == MCP_STUB_TOKEN_URL | ||
| assert stored.allow_all_keys is False | ||
| assert stored.credentials is None, f"client secret must be redacted on read-back, got {stored.credentials}" | ||
|
|
||
| key = client.gateway.generate_key( |
Collaborator
There was a problem hiding this comment.
this is too big of a function, can you split this into:
- def test_mcp_create_oauth_mcp()
- def test_mcp_fetch_oauth_tools() # assert token was stored in db
- def test_mcp_call_tools() ## check if the token can call tools
- def test_mcp_use_with_key() # try making a completions call
see what i mean?
Removes the auth_forwarder key-stamping relay (unused; raw httpx), the stub's /conformance mount and its fixtures, MCP_STUB_CONFORMANCE_URL, the unused prompt/resource client ops, and the unclaimed passes_official_conformance registry cell. None of it was exercised by a test; the suite stays scoped to the core happy paths against the deterministic stub.
Merged
19 tasks
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
2 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.
Relevant issues
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
Everything below was first captured on 2026-07-13 with the test code at commit dd05a44 and re-executed after each follow-up (the suite run and static gates most recently at head 0f8db78, the negative control at daf31ca; the later commits are readability refactors of the concurrency test with identical assertions), against proxies built from source at exact refs, over the deterministic mcp-stub upstream this PR adds to the compose stack. The suite drives the real MCP protocol through the gateway (initialize, tools/list, tools/call over streamable HTTP with the official mcp SDK); no LLM calls are involved in these flows. Outputs are pasted verbatim except that the pre-existing "spend-log cleanup best-effort failed" session-finish notice (it fires on any local e2e run without a host Postgres at localhost:5432 and does not affect results) is elided where it interleaves with pytest lines. Following the same steps reproduces everything end to end
tests/e2e/.envcarries OPENAI_API_KEY, ANTHROPIC_API_KEY, LITELLM_LICENSE for stack startup parity with the other suites; the mcp suite itself never calls a providerBring the stack up (this builds and starts the new mcp-stub service alongside the proxy) and confirm liveness:
uv sync --inexact --group e2e-dev); output below from the 0f8db78 re-run:The auth tests pass on the old ref (that behavior predates it) and the concurrency test fails at its recorded-state assertion because that proxy neither stores nor enforces the cap; the enforcement half is falsified continuously by the in-test uncapped control server, which must observe all 6 calls overlapping through the identical burst machinery every run (a broken or skipped semaphore makes the capped server record 6 exactly like the control, failing the equality assertion)
x-litellm-api-keyare accepted on LLM routes but the MCP mount only accepts the Bearer-prefixed form both docs document:The tests encode the documented Bearer-prefixed contract; the bare-key inconsistency between the LLM and MCP surfaces is flagged here as a product finding for a separate decision rather than worked around
Teardown both stacks:
docker compose -p <project> -f docker-compose.yml -f <override> down -vStatic gates re-run at 0f8db78:
uv run --no-sync basedpyright tests/e2ereports 0 errors, andpython -m coverage_registry.collector --strict(a full collect-only pass resolving every covers marker against the registry) reports the MCPs module moving from 0/14 on the base branch to 5/15 here, with the caps_concurrency row added by this PRRestructure at f503c3a (the branch is also rebased onto current litellm_internal_staging as of this commit; the uv.lock conflict was resolved by re-locking on top of staging's lock with the exclude-newer snapshot anchor preserved, and the proxy's compose depends_on now gates on both the jaeger and mcp-stub services): the single test_mcp_gateway_e2e.py is split into feature-scoped specs with identical assertions. The two near-duplicate header tests become one parametrized matrix (each parametrization claims its registry cells via param-level covers marks, so the Authorization arm now also runs the allow_all_keys round-trip it previously skipped), and the unrecognized-key rejection becomes its own test run under both headers (previously it only ran under x-litellm-api-key), claiming two new rejects_unknown_key registry rows. Re-run against the compose stack after the split, plus both static gates green at this head:
Readability pass at e5b8967, so each test body reads 1:1 against its QA runbook entry with nothing hidden behind fixtures: tests mint their own virtual key inline (client.gateway.generate_key + a deferred delete) instead of the scoped_key fixture, and build the exact wire header dict inline ({"x-litellm-api-key": f"Bearer {key}"}) instead of the McpAuth wrapper, which is deleted; the suite client's protocol methods now take a plain headers dict. Identical assertions, re-run green (5 passed) with both static gates clean
Explicitness pass at 6d3d1d7: the parametrized header matrix from item 10 is unrolled again into four explicitly named tests (test_list_and_call_tools_with_x_litellm_api_key_header, test_list_and_call_tools_with_authorization_bearer_header, and the two per-header test_unrecognized_key_* rejections), so the literal header dict each test sends sits in its own body instead of behind a parametrize id; covers marks moved back onto each test. Identical assertions; re-run green (5 passed in 9.37s) with both static gates clean at this head
Interactive OAuth coverage at 0ba1fb6: the suite gains test_mcp_oauth_interactive_e2e.py, which drives the full authorization_code flow end to end against a gateway-managed oauth2 server. The stub grows into a deterministic OAuth2 IdP (an auto-approving /oauth/authorize, a token endpoint speaking the authorization_code grant with one-time codes and S256 PKCE verification plus refresh_token, and a /oauthuser/mcp mount that rejects anything but the per-user token that grant hands out; the guard records the headers of the last authorized request, read back through a recorded_headers tool). The MCP-host side is the official mcp SDK's own OAuth machinery (OAuthClientProvider: RFC 9728/8414 discovery against the gateway, RFC 7591 dynamic client registration, PKCE, token exchange), with the browser leg replaced by a redirect chaser following the authorize chain (gateway -> stub IdP -> gateway callback -> host redirect_uri) with plain GETs; the caller's LiteLLM key is injected at the httpx transport layer because the SDK's internally-built OAuth requests bypass client-default headers, and the gateway needs the key on the token exchange to know which user to store the upstream token for. The key must belong to a user: a user-less service-account key completes the client-facing dance, but the gateway logs that it dropped the token and subsequent sessions serve an empty tool list instead of a re-challenge (flagged as a separate product finding). The compose file publishes the stub's port to the host because the browser leg runs from pytest, while token exchange stays proxy-side on the compose network. Two P0 registry rows are added (mcp.list_tools.oauth.completes_authorization_code_flow, mcp.call_tool.oauth.uses_per_user_token); the suite is 6 passed in ~10s at this head with both static gates clean
Foundation pass at 68538f3, adopting how the wider MCP ecosystem tests servers without growing the test surface: the stub gains a /conformance mount implementing the official @modelcontextprotocol/conformance suite's hardcoded fixture contract across all three primitives (test_simple_text / test_error_handling tools, the test://static-text, test://static-binary and test://template/{id}/data resources, the test_simple_prompt / test_prompt_with_arguments prompts), auth_forwarder.py provides an in-process relay that stamps the virtual key onto every request (the official CLI and similar host tools cannot send custom headers), and the suite client speaks prompts/list, prompts/get, resources/list and resources/read. The suite itself stays scoped to the core happy paths (6 tests); running the official suite or covering the prompt/resource registry cells is one small test file away on this foundation, e.g. npx -y @modelcontextprotocol/conformance@0.1.11 server --url http://127.0.0.1://mcp --scenario
A manual sweep of 14 curated official scenarios over this foundation (9 passed) surfaced four real gateway gaps, verified by hand and documented here rather than encoded as tests for now: logging/setLevel and completion/complete both answer -32601 Method not found instead of being relayed upstream; prompts/get resolves alias-prefixed names only, so the unprefixed names a spec-reading host may send are 403'd (tools/call has an unprefixed fallback, prompts/get does not); and binary resource reads drop the base64 blob field on the way through (the stub demonstrably serves it; the gateway's answer has uri and mimeType but no blob). The namespacing contract observed live: prompt names are alias-prefixed like tool names, resource URIs pass through unprefixed. A fifth gap, reproduced at 3b82532 and now pinned as a deliberately failing test: a virtual key in the Authorization header of an MCP session on an oauth2 server skips the pre-session 401 challenge entirely and the session is served a masked, empty tool list, so an OAuth-capable host that engages on 401 never starts the dance
Happy-path relocation at a03750f: the list-and-call tests move out of test_mcp_tool_access_e2e.py (which now holds only the unknown-key ingress gate) into test_mcp_oauth_interactive_e2e.py, so the core happy path runs against a real OAuth server with a scoped internal-user key (object_permission grant, allow_all_keys false) under both documented headers. The x-litellm-api-key form passes end to end. The Authorization form is checked in deliberately failing: it asserts the contract (same 401 challenge, same dance, same tools) and currently fails at the challenge step on the finding-five gap above, making it the regression guard for that fix. Suite at this head: 4 passed, 1 intentional failure, both static gates clean
Rebased onto current litellm_internal_staging at 9b144f1 (128 commits, including the new e2e integration line and the datadog logging suite). Conflicts were confined to the shared harness surfaces where both lines added entries: e2e_config.py (the dd-sink query URL alongside the mcp-stub URLs) and the compose depends_on (now gating on dd-sink and mcp-stub), resolved as unions; uv.lock was re-resolved on staging's own exclude-newer anchor and differs from staging by exactly the two mcp e2e-dev entries. Suite re-run at this head: 4 passed plus the one intentional failure, collector --strict and basedpyright clean, and GitHub reports the PR mergeable again
Type
✅ Test
Changes
Adds
tests/e2e/mcp/, the first live coverage of the MCPs registry module, claiming mcp.list_tools.api_key.succeeds, mcp.call_tool.api_key.succeeds, mcp.list_tools.bearer.succeeds, mcp.call_tool.bearer.succeeds, the new mcp.call_tool.api_key.caps_concurrency cell, and the new mcp.list_tools.{api_key,bearer}.rejects_unknown_key cellsThe suite is laid out as one file per feature, and every test body spells out each QA step inline: register the server over
POST /v1/mcp/server, assert the record round-trips throughGET /v1/mcp/server/{id}, mint a virtual key over /key/generate, build the literal wire header dict, then drive initialize, tools/list, and tools/call through the gateway's per-server URL namespace{PROXY}/{alias}/mcpexactly like a production MCP host, asserting the exact tool listing and an exact echo round-trip. The two documented auth headers are deliberately two explicitly named twin tests (test_list_and_call_tools_with_x_litellm_api_key_header builds {"x-litellm-api-key": f"Bearer {key}"}, test_list_and_call_tools_with_authorization_bearer_header builds {"Authorization": f"Bearer {key}"}) so the header each sends is visible in its own body. The unrecognized-key rejection is likewise its own pair of per-header tests; each settles the server with a real key first, so the 401 it asserts can only be the ingress auth refusing the key rather than record propagationtest_mcp_concurrency_e2e.pyholds the cap spec, unchanged in behavior from earlier revisions of this PRThe concurrency test registers a capped server (max_concurrent_requests 2) and an uncapped control against the same upstream, asserts the cap round-trips (the control echoes null), fires 6 simultaneous slow tool calls at each, and reads the upstream's own per-marker in-flight counters back through the proxy via the stub's stats tool. The proxy-side semaphore queues rather than rejects, so the contract asserted is that all 6 calls succeed while the stub observes exactly 2 overlapping calls on the capped server and exactly 6 on the control
Known failure against current main, on purpose: test_key_in_authorization_header_is_challenged_not_masked pins that a virtual key presented as Authorization: Bearer on a gateway-managed oauth2 server must get the 401 OAuth challenge, not a masked 200+empty tool list. Today any bearer in Authorization skips the pre-session challenge and the session opens masked, so this test fails at the challenge assertion. It is a challenge-guard, not a full round-trip: completing the dance via Authorization is unreachable because the SDK carries the minted OAuth token in Authorization after the dance, evicting the LiteLLM key the gateway needs to resolve the per-user token, so the full list-and-call lives only on the x-litellm-api-key variant. This flips green when the mode-aware gateway challenge fix (PR #33586) reaches main-latest; verified live against a patched gateway
The conformance-ready foundation from proof item 14 (the stub's fixture-complete /conformance mount, auth_forwarder.py, and the client's prompt/resource operations) ships without tests of its own, keeping the suite scoped to the core happy paths
test_mcp_oauth_interactive_e2e.pyholds the interactive OAuth spec described in proof item 13: the 401 challenge contract, the full authorization_code dance run by the mcp SDK's own OAuth client, and the per-user-token egress assertions via the stub's recorded_headers read-backSupporting changes: a deterministic stub MCP upstream (
tests/e2e/mcp/stub/, FastMCP from the official mcp SDK pinned to the version litellm locks) added to the compose stack as the mcp-stub service, reachable by the proxy athttp://mcp-stub:8765/mcpand overridable via E2E_MCP_STUB_URL for deployments where the compose stub is not visible; McpServerCreateBody/McpServerInfo in models.py; the suite client in mcp_client.py routing management calls through the shared Gateway transport and protocol calls through the mcp SDK (added to the e2e-dev dependency group behind the same importorskip pattern as playwright); the caps_concurrency registry row plus grammar note; and the suite line in tests/e2e/CLAUDE.mdQA runbook
Prerequisites: the tests/e2e compose stack (which now includes the mcp-stub service), no provider keys needed for this suite, and the mcp SDK (
uv sync --inexact --group e2e-dev). Every MCP-route request below carries the key Bearer-prefixed in whichever header the step names; an MCP client here means any streamable-http MCP client, e.g. the mcp python SDK, pointed at the per-server URLhttp://localhost:4000/<alias>/mcpFinal Attestation