From 9b144f131142abe89389ea09e9e72c447d631244 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 16:00:09 -0700 Subject: [PATCH] test(e2e): run the list-and-call happy path through the interactive OAuth flow under both ingress headers --- .../e2e/mcp/test_mcp_oauth_interactive_e2e.py | 131 ++++++++++++++---- tests/e2e/mcp/test_mcp_tool_access_e2e.py | 98 +++---------- tests/e2e/models.py | 8 ++ 3 files changed, 128 insertions(+), 109 deletions(-) diff --git a/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py b/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py index 00b855c3676..9dc28a7a48d 100644 --- a/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py +++ b/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py @@ -1,6 +1,8 @@ -"""Live e2e: the interactive (authorization_code) OAuth MCP flow. +"""Live e2e: list-and-call through the interactive (authorization_code) OAuth flow. -Covers mcp.list_tools.oauth.completes_authorization_code_flow and +Covers the core list_tools/call_tool happy path against a real OAuth MCP +server, once per documented ingress header, plus +mcp.list_tools.oauth.completes_authorization_code_flow and mcp.call_tool.oauth.uses_per_user_token: a gateway-managed oauth2 server in the authorization_code flow must challenge an unauthorized MCP session with 401 + WWW-Authenticate, let a real MCP host complete the whole interactive @@ -10,16 +12,27 @@ PKCE-verified token exchange), and then serve tools/list and tools/call using the per-user upstream token the gateway obtained; the caller's virtual key must never leave the gateway. -The MCP-host side is the official mcp SDK's own OAuth machinery -(OAuthClientProvider), the same code path desktop MCP hosts run; only the -browser leg is replaced by a redirect chaser that follows the authorize chain -(gateway -> stub IdP -> gateway callback -> host redirect_uri) with plain -GETs. The upstream is the /oauthuser stub mount, which 401s anything but the -exact access token the stub IdP's authorization_code grant hands out, so a -served call proves the gateway completed the code exchange (the stub verifies -client credentials, redirect_uri, one-time code, and the S256 code_verifier) -rather than forwarding anything it already had. The recorded_headers -read-back then makes the injected token explicit and adds the leak check. +The caller is a real scoped internal user, not a wildcard: the virtual key is +user-bound and granted the oauth server through object_permission (the server +is not allow_all_keys), so the flow tested is exactly what a permissioned +customer key experiences. The MCP-host side is the official mcp SDK's own +OAuth machinery (OAuthClientProvider), the same code path desktop MCP hosts +run; only the browser leg is replaced by a redirect chaser that follows the +authorize chain (gateway -> stub IdP -> gateway callback -> host +redirect_uri) with plain GETs. The upstream is the /oauthuser stub mount, +which 401s anything but the exact access token the stub IdP's +authorization_code grant hands out, so a served call proves the gateway +completed the code exchange rather than forwarding anything it already had; +the recorded_headers read-back makes the injected token explicit and adds the +leak check. + +The two header styles are deliberately separate, explicitly named tests, and +they assert the same contract: which header carries the key must not change +what the OAuth flow does. Once the dance completes, the SDK carries the +gateway-minted OAuth token in `Authorization` on MCP traffic (overwriting a +key placed there), which is why the x-litellm-api-key form is what production +hosts configure; the Authorization form pins that a key presented the other +documented way is still challenged into the same working flow. """ from __future__ import annotations @@ -38,7 +51,7 @@ from e2e_config import ( ) from lifecycle import ResourceManager from mcp_client import InMemoryTokenStorage, McpClient, McpDenied, StubRecordedHeaders -from models import KeyGenerateBody, McpServerCreateBody, McpServerCredentials +from models import KeyGenerateBody, KeyObjectPermission, McpServerCreateBody, McpServerCredentials pytestmark = pytest.mark.e2e @@ -46,19 +59,12 @@ GUARDED_STUB_TOOLS = ("echo", "recorded_headers") class TestMcpOauthAuthorizationCode: - """An authorization_code server challenges unauthenticated sessions, runs - the full interactive dance with a real MCP host, and serves tools with the - per-user token the gateway obtained from the IdP. - - There is deliberately no authorization-bearer twin of this test: the OAuth - token owns the `Authorization` header on MCP traffic, so the virtual key - must ride `x-litellm-api-key`. A key presented in `Authorization` instead - never even reaches the dance on current gateways: any bearer there skips - the pre-session 401 challenge and the session is served a masked, empty - tool list, so an OAuth-capable host (which engages on 401) never starts - authorizing. That gap is a known product issue, reproduced live while - building this test; when it is fixed, the twin belongs here.""" + """A scoped internal-user key on an authorization_code server is + challenged, completes the interactive dance, and lists and calls tools + with the per-user token the gateway obtained from the IdP.""" + @pytest.mark.covers("mcp.list_tools.api_key.succeeds") + @pytest.mark.covers("mcp.call_tool.api_key.succeeds") @pytest.mark.covers("mcp.list_tools.oauth.completes_authorization_code_flow") @pytest.mark.covers("mcp.call_tool.oauth.uses_per_user_token") def test_list_and_call_tools_with_x_litellm_api_key_header( @@ -70,7 +76,7 @@ class TestMcpOauthAuthorizationCode: McpServerCreateBody( alias=alias, url=MCP_STUB_OAUTHUSER_URL, - allow_all_keys=True, + allow_all_keys=False, auth_type="oauth2", oauth2_flow="authorization_code", authorization_url=MCP_STUB_AUTHORIZE_BROWSER_URL, @@ -88,9 +94,15 @@ class TestMcpOauthAuthorizationCode: 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(KeyGenerateBody(user_id="e2e-test-user")) + key = client.gateway.generate_key( + KeyGenerateBody( + user_id="e2e-test-user", + object_permission=KeyObjectPermission(mcp_servers=[created.server_id]), + ) + ) resources.defer(lambda: client.gateway.delete_key(key)) headers = {"x-litellm-api-key": f"Bearer {key}"} @@ -124,3 +136,68 @@ class TestMcpOauthAuthorizationCode: ) leaked = sorted(name for name, value in upstream_headers.items() if key in value) assert leaked == [], f"caller's virtual key crossed the gateway boundary in header(s) {leaked}" + + @pytest.mark.covers("mcp.list_tools.bearer.succeeds") + @pytest.mark.covers("mcp.call_tool.bearer.succeeds") + def test_list_and_call_tools_with_authorization_bearer_header( + self, client: McpClient, resources: ResourceManager + ) -> None: + """The identical contract with the key presented as + `Authorization: Bearer ` instead: the gateway must challenge the + pre-dance session with 401 exactly like the x-litellm-api-key form, + and the completed dance must serve the same tools.""" + marker = unique_marker() + alias = f"e2emcpauthcodez{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.allow_all_keys is False + + key = client.gateway.generate_key( + KeyGenerateBody( + user_id="e2e-test-user", + object_permission=KeyObjectPermission(mcp_servers=[created.server_id]), + ) + ) + resources.defer(lambda: client.gateway.delete_key(key)) + headers = {"Authorization": f"Bearer {key}"} + + control_alias = f"e2emcpauthcodezctl{marker}" + control = client.create_server( + McpServerCreateBody(alias=control_alias, url=MCP_STUB_URL, allow_all_keys=True) + ) + resources.defer(lambda: client.delete_server(control.server_id)) + _ = client.poll_tool_names(control_alias, headers) + + challenged = client.list_tools_once(alias, headers) + assert isinstance(challenged, McpDenied), ( + f"session without a user token was served tools instead of the OAuth challenge: {challenged}" + ) + assert challenged.status_code == 401, f"expected the 401 OAuth challenge, got {challenged}" + + storage = InMemoryTokenStorage() + names = client.poll_oauth_tool_names(alias, headers, storage) + expected = tuple(sorted(f"{alias}-{tool}" for tool in GUARDED_STUB_TOOLS)) + assert names == expected, f"post-dance listing was {names}, expected exactly {expected}" + + payload = f"e2e-{marker}" + result = client.oauth_call_tool(alias, headers, storage, f"{alias}-echo", {"text": payload}) + assert result.is_error is False, f"echo through the user-token upstream errored: {result.text[:300]}" + assert result.text == payload diff --git a/tests/e2e/mcp/test_mcp_tool_access_e2e.py b/tests/e2e/mcp/test_mcp_tool_access_e2e.py index 2602425d218..44bb7cf7741 100644 --- a/tests/e2e/mcp/test_mcp_tool_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_tool_access_e2e.py @@ -1,28 +1,20 @@ -"""Live e2e: MCP gateway tool access under both documented auth headers. +"""Live e2e: the MCP gateway's ingress key gate under both documented headers. -Covers mcp.list_tools.api_key.succeeds + mcp.call_tool.api_key.succeeds (the -`x-litellm-api-key` header), mcp.list_tools.bearer.succeeds + -mcp.call_tool.bearer.succeeds (the `Authorization` header), and -mcp.list_tools.{api_key,bearer}.rejects_unknown_key (the ingress gate), -all against the deterministic mcp-stub compose service (tests/e2e/mcp/stub/). +Covers mcp.list_tools.{api_key,bearer}.rejects_unknown_key against the +deterministic mcp-stub compose service (tests/e2e/mcp/stub/): a key the proxy +does not recognize must be refused with an explicit 401 at session +establishment, whichever documented header carries it. The list-and-call +happy path itself lives in test_mcp_oauth_interactive_e2e.py, where it runs +against a real OAuth server; these tests keep the plain ingress gate pinned +with a deliberately minimal setup. Each test body spells out every step a human QA run would take, in order: -create the server over the management API, read the record back, mint a -virtual key over /key/generate, build the exact wire header, then drive the -real MCP protocol through the gateway the way a production MCP host does -(initialize, tools/list, tools/call over streamable HTTP at -{PROXY}/{alias}/mcp) and assert the enforced behavior. Teardown is the only -thing delegated (resources.defer), because it must run even when an -assertion fails. - -The two documented header styles are deliberately separate, explicitly named -tests rather than one parametrized body, so the literal header each one sends -(`x-litellm-api-key: Bearer sk-...` / `Authorization: Bearer sk-...`) is -visible in its own body. Both are Bearer-prefixed on the MCP routes, matching -the docs; a bare key in `x-litellm-api-key` is accepted on LLM routes but -401s here. The unknown-key rejection tests settle the server with a real key -first, so their 401 can only be the gateway's ingress auth refusing the key, -never record propagation. +create the server over the management API, mint a virtual key over +/key/generate, build the exact wire header, settle the server with the real +key (so the later refusal can only be the ingress auth, never record +propagation), then present an unknown key over the same header and require +the explicit 401. Teardown is the only thing delegated (resources.defer), +because it must run even when an assertion fails. """ from __future__ import annotations @@ -36,67 +28,9 @@ from models import KeyGenerateBody, McpServerCreateBody pytestmark = pytest.mark.e2e -STUB_TOOLS = ("echo", "slow_echo", "stats") - - class TestMcpToolAccess: - """Any virtual key reaches an allow_all_keys server through either - documented auth header; a key the proxy does not recognize is turned away - at the door.""" - - @pytest.mark.covers("mcp.list_tools.api_key.succeeds") - @pytest.mark.covers("mcp.call_tool.api_key.succeeds") - def test_list_and_call_tools_with_x_litellm_api_key_header( - self, client: McpClient, resources: ResourceManager - ) -> None: - alias = f"e2emcp{unique_marker()}" - created = client.create_server(McpServerCreateBody(alias=alias, url=MCP_STUB_URL, allow_all_keys=True)) - resources.defer(lambda: client.delete_server(created.server_id)) - - stored = client.server_info(created.server_id) - assert stored.alias == alias - assert stored.url == MCP_STUB_URL - assert stored.allow_all_keys is True - - key = client.gateway.generate_key(KeyGenerateBody()) - resources.defer(lambda: client.gateway.delete_key(key)) - - headers = {"x-litellm-api-key": f"Bearer {key}"} - names = client.poll_tool_names(alias, headers) - expected = tuple(sorted(f"{alias}-{tool}" for tool in STUB_TOOLS)) - assert names == expected, f"gateway listed {names}, expected exactly {expected}" - - payload = f"e2e-{unique_marker()}" - result = client.call_tool(alias, headers, f"{alias}-echo", {"text": payload}) - assert result.is_error is False, f"echo call errored: {result.text[:300]}" - assert result.text == payload - - @pytest.mark.covers("mcp.list_tools.bearer.succeeds") - @pytest.mark.covers("mcp.call_tool.bearer.succeeds") - def test_list_and_call_tools_with_authorization_bearer_header( - self, client: McpClient, resources: ResourceManager - ) -> None: - alias = f"e2emcp{unique_marker()}" - created = client.create_server(McpServerCreateBody(alias=alias, url=MCP_STUB_URL, allow_all_keys=True)) - resources.defer(lambda: client.delete_server(created.server_id)) - - stored = client.server_info(created.server_id) - assert stored.alias == alias - assert stored.url == MCP_STUB_URL - assert stored.allow_all_keys is True - - key = client.gateway.generate_key(KeyGenerateBody()) - resources.defer(lambda: client.gateway.delete_key(key)) - - headers = {"Authorization": f"Bearer {key}"} - names = client.poll_tool_names(alias, headers) - expected = tuple(sorted(f"{alias}-{tool}" for tool in STUB_TOOLS)) - assert names == expected, f"gateway listed {names}, expected exactly {expected}" - - payload = f"e2e-{unique_marker()}" - result = client.call_tool(alias, headers, f"{alias}-echo", {"text": payload}) - assert result.is_error is False, f"echo call errored: {result.text[:300]}" - assert result.text == payload + """A key the proxy does not recognize is turned away at the door, under + either documented auth header.""" @pytest.mark.covers("mcp.list_tools.api_key.rejects_unknown_key") def test_unrecognized_key_in_x_litellm_api_key_header_is_turned_away( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index f54674515f1..40b667add05 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -39,6 +39,13 @@ class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None +class KeyObjectPermission(BaseModel): + """object_permission on /key/generate: the MCP servers (by server_id) a key + may list and call. A key without one only reaches allow_all_keys servers.""" + + mcp_servers: list[str] | None = None + + class KeyGenerateBody(BaseModel): models: list[str] = [] duration: str | None = None @@ -57,6 +64,7 @@ class KeyGenerateBody(BaseModel): rpm_limit: int | None = None allowed_routes: list[str] | None = None metadata: KeyMetadata | None = None + object_permission: KeyObjectPermission | None = None class KeyGenerateResponse(BaseModel):