diff --git a/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py b/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py index fe4b19143fd..6fdab5742d4 100644 --- a/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py +++ b/tests/e2e/mcp/test_mcp_oauth_interactive_e2e.py @@ -26,16 +26,17 @@ 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 documented header placements are covered by two deliberately separate -tests that assert different, achievable contracts. The x-litellm-api-key form -runs the full round-trip, because that header stays free for the LiteLLM key +The documented header placements are covered by separate tests asserting +different, achievable contracts. The x-litellm-api-key form runs the full +interactive round-trip, because that header stays free for the LiteLLM key while the SDK owns `Authorization` for the minted OAuth token. The -Authorization form is a challenge-guard only: putting the LiteLLM key in the -same header the OAuth token must occupy cannot complete the dance (the token -evicts the key on the post-dance retry, so the gateway loses the identity it -needs to resolve the per-user token), so the contract there is that the -gateway fails closed and loudly with the 401 challenge rather than silently -masking the session as an empty tool list. +Authorization form is split by whether a per-user token already exists: with +no stored token, the LiteLLM key in Authorization cannot complete the dance +(the token would evict the key on the post-dance retry, so the gateway loses +the identity it needs), so the contract is that the gateway fails closed and +loudly with the 401 challenge rather than silently masking an empty tool list; +with a stored token, no dance is needed, so a plain Authorization-key session +lists and calls using the token the gateway already holds for that user. """ from __future__ import annotations @@ -62,11 +63,13 @@ GUARDED_STUB_TOOLS = ("echo", "recorded_headers") class TestMcpOauthAuthorizationCode: - """A scoped internal-user key on an authorization_code server: the - x-litellm-api-key form is challenged, completes the interactive dance, and - lists and calls tools with the per-user token the gateway obtained; the - Authorization form is challenged (not masked) but cannot complete the - round-trip, since the OAuth token claims the Authorization slot.""" + """A scoped internal-user key on an authorization_code server, across three + behaviors: the x-litellm-api-key form is challenged, completes the + interactive dance, and lists/calls with the per-user token the gateway + obtained; a key in Authorization with no stored token yet is challenged + (not masked), since the OAuth token would claim the Authorization slot and + the dance cannot complete there; and once a per-user token is stored, a + plain Authorization-key session lists/calls using that stored token.""" @pytest.mark.covers("mcp.list_tools.api_key.succeeds") @pytest.mark.covers("mcp.call_tool.api_key.succeeds") @@ -206,3 +209,79 @@ class TestMcpOauthAuthorizationCode: f"key in Authorization was served a (masked) tool list instead of the OAuth challenge: {challenged}" ) assert challenged.status_code == 401, f"expected the 401 OAuth challenge, got {challenged}" + + @pytest.mark.covers("mcp.list_tools.bearer.succeeds") + @pytest.mark.covers("mcp.call_tool.bearer.succeeds") + def test_stored_token_serves_list_and_call_with_key_in_authorization_header( + self, client: McpClient, resources: ResourceManager + ) -> None: + """A returning user who has already authorized presents the LiteLLM key + in `Authorization: Bearer ` and the gateway lists and calls tools + using the stored per-user token. + + This is the positive counterpart to the challenge-guard above. It works + because no interactive dance is needed once a per-user token exists: the + gateway resolves the user from the LiteLLM key in Authorization and + injects the stored upstream token at egress. (The dance itself cannot be + completed with the key in Authorization, since the SDK would overwrite + that header with the minted token; that is why the token is seeded here + via the x-litellm-api-key dance first, then the plain Authorization-key + session is what serves.) The recorded_headers read-back proves the + upstream received the stored per-user token, not the LiteLLM key, and + the key crosses the gateway boundary in no header.""" + marker = unique_marker() + alias = f"e2emcpauthcodestored{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 + 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", + object_permission=KeyObjectPermission(mcp_servers=[created.server_id]), + ) + ) + resources.defer(lambda: client.gateway.delete_key(key)) + + storage = InMemoryTokenStorage() + seeded = client.poll_oauth_tool_names(alias, {"x-litellm-api-key": f"Bearer {key}"}, storage) + expected = tuple(sorted(f"{alias}-{tool}" for tool in GUARDED_STUB_TOOLS)) + assert seeded == expected, f"the authorizing dance listed {seeded}, expected exactly {expected}" + + authorization_headers = {"Authorization": f"Bearer {key}"} + names = client.poll_tool_names(alias, authorization_headers) + assert names == expected, f"Authorization-key listing was {names}, expected exactly {expected}" + + payload = f"e2e-{marker}" + result = client.call_tool(alias, authorization_headers, f"{alias}-echo", {"text": payload}) + assert result.is_error is False, f"echo with the key in Authorization errored: {result.text[:300]}" + assert result.text == payload + + recorded = client.call_tool(alias, authorization_headers, f"{alias}-recorded_headers", {}) + assert recorded.is_error is False, f"recorded_headers call errored: {recorded.text[:300]}" + upstream_headers = StubRecordedHeaders.model_validate_json(recorded.text).root + assert upstream_headers.get("authorization") == f"Bearer {MCP_STUB_OAUTH_USER_ACCESS_TOKEN}", ( + "upstream must receive the stored per-user token the gateway injects, not the LiteLLM key, " + f"got {upstream_headers.get('authorization')!r}" + ) + 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}"