mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
test(e2e): run the list-and-call happy path through the interactive OAuth flow under both ingress headers
This commit is contained in:
Tin Chi Lo
parent
9fb7271ac6
commit
9b144f1311
3 changed files with 128 additions and 109 deletions
|
|
@ -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 <key>` 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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue