mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
test(e2e): cover MCP gateway auth headers and per-server max concurrency
This commit is contained in:
parent
fba7ac4428
commit
5edc486928
12 changed files with 534 additions and 2 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
19
tests/e2e/mcp/conftest.py
Normal 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
208
tests/e2e/mcp/mcp_client.py
Normal 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())
|
||||
9
tests/e2e/mcp/stub/Dockerfile
Normal file
9
tests/e2e/mcp/stub/Dockerfile
Normal 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"]
|
||||
71
tests/e2e/mcp/stub/stub_server.py
Normal file
71
tests/e2e/mcp/stub/stub_server.py
Normal 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")
|
||||
174
tests/e2e/mcp/test_mcp_gateway_e2e.py
Normal file
174
tests/e2e/mcp/test_mcp_gateway_e2e.py
Normal 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"
|
||||
)
|
||||
|
|
@ -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
2
uv.lock
generated
|
|
@ -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" },
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue