test(e2e): cover MCP gateway auth headers and per-server max concurrency

This commit is contained in:
Tin Chi Lo 2026-07-13 11:09:01 -07:00
parent fba7ac4428
commit 5edc486928
12 changed files with 534 additions and 2 deletions

View file

@ -189,6 +189,7 @@ dev = [
e2e-dev = [
"playwright==1.61.0",
"websockets>=15.0.1,<16.0",
"mcp>=1.26.0,<2.0",
]
proxy-dev = [
"prisma==0.11.0",

View file

@ -13,6 +13,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `realtime/` - realtime websocket sessions, including the pipecat audio path
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip)
- `mcp/` - the MCP gateway: server registration via `/v1/mcp/server`, tool listing/calling over the streamable-http protocol under both auth headers, and per-server enforcement such as `max_concurrent_requests`; `stub/` holds the deterministic upstream MCP server the compose stack runs for it
- `logging/` - logging-integration delivery (datadog and friends)
- `security/` - secret handling and log-leak protection
- `router/` - routing and reliability behavior (fallbacks, cooldowns)
@ -107,7 +108,7 @@ MCPs - endpoint features with the protocol op as the variant
mcp.<operation>.<auth_family>.<assertion>
operation : list_tools | call_tool | list_resources | read_resource | list_prompts | get_prompt
auth_family : none | api_key | bearer | oauth
assertion : succeeds | denied_without_permission
assertion : succeeds | denied_without_permission | caps_concurrency
e.g. mcp.call_tool.oauth.succeeds
```

View file

@ -79,6 +79,14 @@
assertions: [succeeds]
source: "rest_endpoints.py:305-334"
rationale: No upstream auth required; demo servers
- id: mcp.call_tool.api_key.caps_concurrency
module: mcp
tier: P1
operation: call_tool
auth_family: api_key
assertions: [caps_concurrency]
source: "mcp_server_manager.py:3751"
rationale: Per-server max_concurrent_requests semaphore bounds in-flight upstream tool calls; excess calls queue
- id: mcp.get_prompt.api_key.succeeds
module: mcp
tier: P1

View file

@ -166,7 +166,15 @@ services:
target: /app/config.yaml
command: ["--config", "/app/config.yaml", "--port", "4000"]
# throwaway db
mcp-stub:
build: ./mcp/stub
healthcheck:
test: ["CMD", "python", "-c", "import socket; socket.create_connection(('localhost', 8765), timeout=2)"]
interval: 3s
timeout: 3s
retries: 20
# throwaway db
db:
image: postgres:16
environment:

View file

@ -36,6 +36,12 @@ OTEL_QUERY_URL = os.environ.get("E2E_OTEL_QUERY_URL", "http://localhost:16686").
# service records every intake POST and replays them on GET /requests).
DD_SINK_URL = os.environ.get("E2E_DD_SINK_URL", "http://localhost:9915").rstrip("/")
# The MCP upstream the mcp suite registers on the proxy: the mcp-stub compose
# service, addressed by service name on the compose network. It must be
# reachable from the proxy, not from pytest; override when the proxy under
# test runs somewhere the compose stub is not visible from.
MCP_STUB_URL = os.environ.get("E2E_MCP_STUB_URL", "http://mcp-stub:8765/mcp")
# Writes on the proxy are eventually consistent (e.g. spend rows flush on
# proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once.
POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120"))

19
tests/e2e/mcp/conftest.py Normal file
View file

@ -0,0 +1,19 @@
"""Fixtures for the mcp e2e suite.
The shared lifecycle (`resources`, `scoped_key`), the liveness gate, and the
`e2e` marker come from the parent tests/e2e/conftest.py; this adds the suite
client. The mcp SDK lives in the `e2e-dev` dependency group (like playwright
for management/), so its import is guarded to skip the suite rather than
error at collection when it is not installed.
"""
import pytest
pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`")
from mcp_client import McpClient, build_client # noqa: E402 # import must follow the importorskip guard above
@pytest.fixture(scope="session")
def client() -> McpClient:
return build_client()

208
tests/e2e/mcp/mcp_client.py Normal file
View file

@ -0,0 +1,208 @@
"""Client for the mcp e2e suite.
Management routes (/v1/mcp/server CRUD) go through the shared Gateway
transport. The MCP protocol itself (initialize, tools/list, tools/call over
streamable HTTP) goes through the official mcp SDK, the same client library
production MCP hosts use, aimed at the gateway's per-server URL namespace
{PROXY}/{alias}/mcp. The gateway accepts the LiteLLM virtual key as either
`x-litellm-api-key: Bearer sk-...` or `Authorization: Bearer sk-...` (both
Bearer-prefixed on the MCP routes, matching the docs); `McpAuth` models the
two header styles so each test states which one it drives.
"""
from __future__ import annotations
import asyncio
import time
from dataclasses import dataclass
from typing import Literal, Mapping, cast
import httpx
import pytest
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.types import CallToolResult, TextContent
from pydantic import BaseModel
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
from e2e_gateway import Gateway, build_gateway
from e2e_http import NoBody, unwrap
from models import McpServerCreateBody, McpServerInfo
McpHeaderName = Literal["x-litellm-api-key", "Authorization"]
ToolArguments = Mapping[str, str | float]
@dataclass(frozen=True, slots=True)
class McpAuth:
"""One of the two documented ways to present a LiteLLM key to the MCP gateway."""
header_name: McpHeaderName
key: str
def headers(self) -> dict[str, str]:
return {self.header_name: f"Bearer {self.key}"}
@dataclass(frozen=True, slots=True)
class McpToolNames:
names: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class McpDenied:
"""A tools/list attempt the gateway rejected before serving the session."""
status_code: int | None
message: str
type ListToolsOutcome = McpToolNames | McpDenied
@dataclass(frozen=True, slots=True)
class McpToolText:
"""The text content of one tools/call result."""
text: str
is_error: bool
class StubToolStats(BaseModel):
"""JSON the stub's `stats` tool returns for one marker (see stub/stub_server.py)."""
marker: str
max_in_flight: int
completed: int
def _mcp_url(alias: str) -> str:
return f"{PROXY_BASE_URL}/{alias}/mcp"
def _first_text(result: CallToolResult) -> str:
first = result.content[0] if result.content else None
if isinstance(first, TextContent):
return first.text
return f"<non-text content: {type(first).__name__}>"
def _http_status(root: BaseException) -> int | None:
"""The HTTP status behind an SDK failure; the SDK wraps transport errors in
(possibly nested) ExceptionGroups, so walk them without recursing."""
pending: list[BaseException] = [root] # mutable-ok: bounded worklist over an exception tree
while pending:
exc = pending.pop()
if isinstance(exc, httpx.HTTPStatusError):
return exc.response.status_code
if isinstance(exc, BaseExceptionGroup):
pending.extend(cast("tuple[BaseException, ...]", exc.exceptions))
return None
def _http_client(headers: dict[str, str]) -> httpx.AsyncClient:
return httpx.AsyncClient(headers=headers, timeout=httpx.Timeout(REQUEST_TIMEOUT), follow_redirects=True)
async def _list_tool_names(url: str, headers: dict[str, str]) -> tuple[str, ...]:
async with _http_client(headers) as http_client:
async with streamable_http_client(url, http_client=http_client) as (read, write, _):
async with ClientSession(read, write) as session:
await session.initialize()
listed = await session.list_tools()
return tuple(sorted(tool.name for tool in listed.tools))
async def _call_tool(url: str, headers: dict[str, str], tool: str, arguments: ToolArguments) -> McpToolText:
async with _http_client(headers) as http_client:
async with streamable_http_client(url, http_client=http_client) as (read, write, _):
async with ClientSession(read, write) as session:
await session.initialize()
result = await session.call_tool(tool, dict(arguments))
return McpToolText(text=_first_text(result), is_error=bool(result.isError))
@dataclass(frozen=True, slots=True)
class McpClient:
gateway: Gateway
def create_server(self, body: McpServerCreateBody) -> McpServerInfo:
return unwrap(
self.gateway.transport.post(
"/v1/mcp/server",
headers=self.gateway.transport.master,
json=body,
response_type=McpServerInfo,
)
)
def server_info(self, server_id: str) -> McpServerInfo:
return unwrap(
self.gateway.transport.get(
f"/v1/mcp/server/{server_id}",
headers=self.gateway.transport.master,
params=NoBody(),
response_type=McpServerInfo,
)
)
def delete_server(self, server_id: str) -> None:
_ = self.gateway.transport.delete(
f"/v1/mcp/server/{server_id}",
headers=self.gateway.transport.master,
json=NoBody(),
response_type=NoBody,
)
def list_tools_once(self, alias: str, auth: McpAuth) -> ListToolsOutcome:
try:
return McpToolNames(names=asyncio.run(_list_tool_names(_mcp_url(alias), auth.headers())))
except Exception as exc: # noqa: BLE001 - the SDK raises ExceptionGroup-wrapped transport errors; modelled as a value
return McpDenied(status_code=_http_status(exc), message=str(exc))
def poll_tool_names(self, alias: str, auth: McpAuth) -> tuple[str, ...]:
"""tools/list to the shared deadline: a just-created server record
propagates to the gateway asynchronously and a just-created key can lag
the auth cache, so the first attempts may 401 or list nothing."""
deadline = time.monotonic() + self.gateway.poll_timeout
outcome: ListToolsOutcome = McpDenied(status_code=None, message="never attempted")
while time.monotonic() < deadline:
outcome = self.list_tools_once(alias, auth)
match outcome:
case McpToolNames(names=names) if names:
return names
case _:
time.sleep(self.gateway.poll_interval)
pytest.fail(
f"MCP tools for {alias!r} never listed within {self.gateway.poll_timeout}s; last outcome: {outcome}"
)
def call_tool(self, alias: str, auth: McpAuth, tool: str, arguments: ToolArguments) -> McpToolText:
return asyncio.run(_call_tool(_mcp_url(alias), auth.headers(), tool, arguments))
def burst_call_tool(
self,
alias: str,
auth: McpAuth,
tool: str,
arguments_per_call: tuple[ToolArguments, ...],
) -> tuple[McpToolText, ...]:
"""Fire every call at once, each over its own MCP session, like N
independent clients hitting the gateway simultaneously."""
async def _burst() -> tuple[McpToolText, ...]:
url = _mcp_url(alias)
headers = auth.headers()
results = await asyncio.gather(*(_call_tool(url, headers, tool, args) for args in arguments_per_call))
return tuple(results)
return asyncio.run(_burst())
def stub_stats(self, alias: str, auth: McpAuth, stats_tool: str, marker: str) -> StubToolStats:
outcome = self.call_tool(alias, auth, stats_tool, {"marker": marker})
return StubToolStats.model_validate_json(outcome.text)
def build_client() -> McpClient:
return McpClient(gateway=build_gateway())

View file

@ -0,0 +1,9 @@
FROM python:3.13-slim
RUN pip install --no-cache-dir mcp==1.26.0
WORKDIR /app
COPY stub_server.py /app/stub_server.py
EXPOSE 8765
CMD ["python", "stub_server.py"]

View file

@ -0,0 +1,71 @@
"""Deterministic MCP upstream for the mcp e2e suite.
Serves the streamable-http MCP protocol with three tools. `echo` answers
immediately so auth tests can assert an exact round-trip. `slow_echo` holds the
request open for `sleep_seconds` while per-`marker` in-flight and max-in-flight
counters track how many calls the proxy let through simultaneously; that is the
observable a per-server `max_concurrent_requests` cap must bound. `stats` reads
those counters back, so tests observe upstream concurrency through the proxy
itself and the stub needs no side-channel port.
Counter updates are plain attribute mutations between awaits, so asyncio's
single-threaded scheduling makes them atomic; markers come from
`unique_marker()` so concurrent test runs never share a counter.
"""
from __future__ import annotations
import asyncio
import json
from dataclasses import dataclass
from mcp.server.fastmcp import FastMCP
mcp = FastMCP("e2e-stub", host="0.0.0.0", port=8765, stateless_http=True)
@dataclass
class _MarkerStats:
in_flight: int = 0
max_in_flight: int = 0
completed: int = 0
_stats: dict[str, _MarkerStats] = {}
@mcp.tool()
def echo(text: str) -> str:
"""Return `text` unchanged."""
return text
@mcp.tool()
async def slow_echo(text: str, marker: str, sleep_seconds: float) -> str:
"""Return `text` after `sleep_seconds`, recording concurrency under `marker`."""
stats = _stats.setdefault(marker, _MarkerStats())
stats.in_flight += 1
stats.max_in_flight = max(stats.max_in_flight, stats.in_flight)
try:
await asyncio.sleep(sleep_seconds)
finally:
stats.in_flight -= 1
stats.completed += 1
return text
@mcp.tool()
def stats(marker: str) -> str:
"""Return the JSON stats recorded for `marker`."""
recorded = _stats.get(marker, _MarkerStats())
return json.dumps(
{
"marker": marker,
"max_in_flight": recorded.max_in_flight,
"completed": recorded.completed,
}
)
if __name__ == "__main__":
mcp.run(transport="streamable-http")

View file

@ -0,0 +1,174 @@
"""Live e2e: MCP gateway tool access and the per-server concurrency cap.
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.call_tool.api_key.caps_concurrency (`max_concurrent_requests`), all
against the deterministic mcp-stub compose service (tests/e2e/mcp/stub/).
Each test walks the full lifecycle: register the server over the management
API and defer its deletion, assert the recorded state round-trips
(GET /v1/mcp/server/{id} echoes what was configured), 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. Both auth headers carry `Bearer <key>`; that is
the documented contract on the MCP routes (a bare key in `x-litellm-api-key`
is accepted on LLM routes but 401s here).
The concurrency test's observable is the stub's own per-marker in-flight
counter, read back through the proxy via the stub's `stats` tool. The
proxy-side semaphore queues excess tool calls instead of rejecting them, so
with a cap of MAX_CONCURRENT a burst of BURST simultaneous slow_echo calls
must all succeed while the stub never sees more than MAX_CONCURRENT of them
overlap. The uncapped control server drives the identical burst through the
identical machinery and must see all BURST calls overlap; that is the
fail-before-fix evidence built into the test, since a broken (or skipped)
semaphore makes the capped server record BURST in-flight exactly like the
control and the equality assertion fail.
"""
from __future__ import annotations
import pytest
from e2e_config import MCP_STUB_URL, unique_marker
from lifecycle import ResourceManager
from mcp_client import McpAuth, McpClient, McpDenied
from models import McpServerCreateBody
pytestmark = pytest.mark.e2e
STUB_TOOLS = ("echo", "slow_echo", "stats")
MAX_CONCURRENT = 2
BURST = 6
SLOW_CALL_SECONDS = 2.0
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, scoped_key: str
) -> 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
auth = McpAuth(header_name="x-litellm-api-key", key=scoped_key)
names = client.poll_tool_names(alias, auth)
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, auth, f"{alias}-echo", {"text": payload})
assert result.is_error is False, f"echo call errored: {result.text[:300]}"
assert result.text == payload
denied = client.list_tools_once(alias, McpAuth(header_name="x-litellm-api-key", key="sk-not-a-real-key"))
assert isinstance(denied, McpDenied), f"unrecognized key was served tools: {denied}"
assert denied.status_code == 401, f"expected 401 for an unrecognized key, got {denied}"
@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, scoped_key: str
) -> 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
auth = McpAuth(header_name="Authorization", key=scoped_key)
names = client.poll_tool_names(alias, auth)
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, auth, f"{alias}-echo", {"text": payload})
assert result.is_error is False, f"echo call errored: {result.text[:300]}"
assert result.text == payload
class TestMcpServerMaxConcurrency:
"""`max_concurrent_requests` bounds how many tool calls the gateway lets
reach one server at a time; excess calls queue and still succeed."""
@pytest.mark.covers("mcp.call_tool.api_key.caps_concurrency")
def test_max_concurrent_requests_caps_in_flight_upstream_calls(
self, client: McpClient, resources: ResourceManager, scoped_key: str
) -> None:
auth = McpAuth(header_name="x-litellm-api-key", key=scoped_key)
capped_alias = f"e2emcpcap{unique_marker()}"
capped = client.create_server(
McpServerCreateBody(
alias=capped_alias,
url=MCP_STUB_URL,
allow_all_keys=True,
max_concurrent_requests=MAX_CONCURRENT,
)
)
resources.defer(lambda: client.delete_server(capped.server_id))
control_alias = f"e2emcpfree{unique_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))
assert client.server_info(capped.server_id).max_concurrent_requests == MAX_CONCURRENT
assert client.server_info(control.server_id).max_concurrent_requests is None
_ = client.poll_tool_names(capped_alias, auth)
_ = client.poll_tool_names(control_alias, auth)
capped_marker = unique_marker()
capped_results = client.burst_call_tool(
capped_alias,
auth,
f"{capped_alias}-slow_echo",
tuple(
{"text": "capped", "marker": capped_marker, "sleep_seconds": SLOW_CALL_SECONDS}
for _ in range(BURST)
),
)
assert len(capped_results) == BURST
assert all(result.is_error is False and result.text == "capped" for result in capped_results), (
f"queued calls must all still succeed under the cap: {capped_results}"
)
capped_stats = client.stub_stats(capped_alias, auth, f"{capped_alias}-stats", capped_marker)
assert capped_stats.completed == BURST
assert capped_stats.max_in_flight == MAX_CONCURRENT, (
f"stub saw {capped_stats.max_in_flight} overlapping calls; the cap of {MAX_CONCURRENT} "
f"must be the exact ceiling AND be reached ({BURST} queued calls saturate it)"
)
control_marker = unique_marker()
control_results = client.burst_call_tool(
control_alias,
auth,
f"{control_alias}-slow_echo",
tuple(
{"text": "control", "marker": control_marker, "sleep_seconds": SLOW_CALL_SECONDS}
for _ in range(BURST)
),
)
assert all(result.is_error is False and result.text == "control" for result in control_results)
control_stats = client.stub_stats(control_alias, auth, f"{control_alias}-stats", control_marker)
assert control_stats.completed == BURST
assert control_stats.max_in_flight == BURST, (
f"uncapped control saw {control_stats.max_in_flight} overlapping calls, expected all {BURST}; "
f"if this fails the instrument cannot detect over-cap concurrency and the capped assertion is vacuous"
)

View file

@ -95,6 +95,31 @@ class KeyInfoResponse(BaseModel):
info: KeyInfo
# ---------- mcp servers ----------
class McpServerCreateBody(BaseModel):
"""POST /v1/mcp/server. `allow_all_keys` opts the server out of per-key
object_permission grants so any virtual key on the proxy may use it."""
alias: str
url: str
transport: str = "http"
allow_all_keys: bool = True
max_concurrent_requests: int | None = None
class McpServerInfo(BaseModel):
"""Response of POST /v1/mcp/server and GET /v1/mcp/server/{server_id}."""
server_id: str
alias: str | None = None
url: str | None = None
transport: str | None = None
allow_all_keys: bool | None = None
max_concurrent_requests: int | None = None
# ---------- customers ----------

2
uv.lock generated
View file

@ -3926,6 +3926,7 @@ dev = [
{ name = "vcrpy" },
]
e2e-dev = [
{ name = "mcp" },
{ name = "playwright" },
{ name = "websockets" },
]
@ -4101,6 +4102,7 @@ dev = [
{ name = "vcrpy", specifier = "==8.2.1" },
]
e2e-dev = [
{ name = "mcp", specifier = ">=1.26.0,<2.0" },
{ name = "playwright", specifier = "==1.61.0" },
{ name = "websockets", specifier = ">=15.0.1,<16.0" },
]