scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

test(e2e): add MCP bridge, filtering, auth, and OAuth coverage tests

Browse source 
Add e2e tests for the three LLM bridges (chat_completions, responses,
messages) that expand litellm_proxy MCP tool references and auto-execute
tools against the real Datadog MCP server. Add allowed_tools scoping,
namespaced multi-server, upstream_static_auth + transport_http, and
401-not-500 auth regression tests. Unskip the datadog round-trip and
key-access call_tool tests by dropping the rejected telemetry argument.

Add a Datadog OAuth2 PKCE path (dd_oauth.py) that registers the server
with auth_type=oauth2, drives the real authorize dance via DCR + browser
consent + token exchange, stores the per-user token in the gateway vault,
and tests list_tools + call_tool + chat_completion through the stored
token. Requires E2E_DD_STORAGE_STATE (captured via dd_session_capture.py).

Delete the stale guardrail test and the old oauth_chat_client helper.
Extend datadog_mcp.py to return DatadogMcpServer(server_id, alias) and
accept allowed_tools/toolsets kwargs. Add ResponsesMcp models and
bridge methods to mcp_client.py. Add AnthropicMcpTool to models.py.
This commit is contained in:
mubashir1osmani 2026-08-05 15:34:39 -07:00
parent 6d0a71cdb7
commit d016edeea3
17 changed files with 1192 additions and 491 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
tests/e2e/mcp/datadog_mcp.py
View file

@ -3,6 +3,7 @@
from __future__ import annotations
import os
from dataclasses import dataclass
from e2e_config import datadog_mcp_url, unique_marker
from lifecycle import ResourceManager
@ -11,6 +12,12 @@ from mcp_client import McpClient
SEARCH_LOGS_TOOL = "search_datadog_logs"
@dataclass(frozen=True, slots=True)
class DatadogMcpServer:
server_id: str
alias: str
def _dd_api_key() -> str:
return os.environ.get("DD_API_KEY", "").strip()
@ -30,25 +37,31 @@ def assert_dd_mcp_creds() -> None:
)
def _dd_static_headers() -> dict[str, str]:
return {
"DD-API-KEY": _dd_api_key(),
"DD-APPLICATION-KEY": _dd_app_key(),
}
def register_datadog_mcp(
client: McpClient,
resources: ResourceManager,
*,
mcp_access_groups: list[str] | None = None,
) -> str:
allowed_tools: list[str] | None = None,
toolsets: str = "core",
) -> DatadogMcpServer:
assert_dd_mcp_creds()
name = f"e2e_dd_mcp_{unique_marker()}"
server_id = client.register_server(
server_name=name,
alias=name,
url=datadog_mcp_url(toolsets="core"),
url=datadog_mcp_url(toolsets=toolsets),
transport="http",
static_headers={
"DD-API-KEY": _dd_api_key(),
"DD-APPLICATION-KEY": _dd_app_key(),
},
allowed_tools=[SEARCH_LOGS_TOOL],
static_headers=_dd_static_headers(),
allowed_tools=allowed_tools if allowed_tools is not None else [SEARCH_LOGS_TOOL],
mcp_access_groups=mcp_access_groups,
)
resources.defer(lambda: client.delete_server(server_id))
return server_id
return DatadogMcpServer(server_id=server_id, alias=name)

275
tests/e2e/mcp/dd_oauth.py Normal file
View file

@ -0,0 +1,275 @@
"""Shared helpers for e2e tests that exercise the Datadog MCP server through
gateway-managed OAuth2 (authorization_code + PKCE).
Datadog's remote MCP server (mcp.datadoghq.com/v1/mcp) supports OAuth2.1 with
mandatory S256 PKCE. The authorize endpoint (app.datadoghq.com) serves an
interactive consent page, so the browser leg is a headless Chromium primed
with a saved Datadog browser session (E2E_DD_STORAGE_STATE).
The gateway discovers the OAuth endpoints via /.well-known metadata, so the
server is registered with auth_type=oauth2, oauth2_flow=authorization_code
and no explicit authorize/token URLs. The per-user token is stored via
POST /v1/mcp/server/{server_id}/oauth-user-credential.
"""
from __future__ import annotations
import base64
import hashlib
import os
import re
import secrets
import time
import urllib.parse
from dataclasses import dataclass
from typing import TYPE_CHECKING
from urllib.parse import parse_qsl
import httpx
import pytest
from pydantic import BaseModel, TypeAdapter
from e2e_config import REQUEST_TIMEOUT
from e2e_http import AuthHeaders, NoBody, unwrap
from models import McpServerCreateBody, McpServerInfo
from proxy_client import ProxyClient
if TYPE_CHECKING:
from playwright.async_api import Route
DD_MCP_URL = "https://mcp.datadoghq.com/v1/mcp"
DD_AUTHORIZE_URL = "https://app.datadoghq.com/oauth2/v1/authorize"
DD_TOKEN_URL = "https://app.datadoghq.com/api/v2/oauth2/token"
DD_REGISTER_URL = "https://app.datadoghq.com/api/v2/oauth2/register"
OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback"
BROWSER_CONSENT_TIMEOUT = 60.0
@dataclass(frozen=True, slots=True)
class PkceChallenge:
verifier: str
challenge: str
state: str
@dataclass(frozen=True, slots=True)
class DcrClient:
client_id: str
@dataclass(frozen=True, slots=True)
class OAuthToken:
access_token: str
token_type: str
refresh_token: str | None = None
@dataclass(frozen=True, slots=True)
class DatadogMcpOAuthServer:
server_id: str
alias: str
class OAuthCredentialBody(BaseModel):
access_token: str
refresh_token: str | None = None
expires_in: int | None = None
scopes: list[str] | None = None
def assert_dd_oauth_env() -> None:
path = os.environ.get("E2E_DD_STORAGE_STATE", "")
if not path or not os.path.exists(path):
pytest.fail(
"Datadog MCP OAuth e2e requires E2E_DD_STORAGE_STATE to point at a "
"saved Datadog browser session. Capture one with mcp/dd_session_capture.py."
)
def _generate_pkce() -> PkceChallenge:
verifier = secrets.token_urlsafe(64)
challenge = base64.urlsafe_b64encode(
hashlib.sha256(verifier.encode()).digest()
).rstrip(b"=").decode()
state = secrets.token_urlsafe(32)
return PkceChallenge(verifier=verifier, challenge=challenge, state=state)
def _dcr_register() -> DcrClient:
resp = httpx.post(
DD_REGISTER_URL,
json={
"client_name": "e2e-mcp-dd-oauth",
"redirect_uris": [OAUTH_CLIENT_REDIRECT_URI],
"token_endpoint_auth_method": "none",
"grant_types": ["authorization_code", "refresh_token"],
"response_types": ["code"],
},
timeout=REQUEST_TIMEOUT,
)
resp.raise_for_status()
client_id = TypeAdapter(str).validate_python(resp.json()["client_id"])
return DcrClient(client_id=client_id)
async def _browser_authorize(
authorize_url: str, storage_state_path: str
) -> tuple[str, str | None]:
import asyncio
from playwright.async_api import async_playwright
captured: dict[str, str] = {}
trail: list[str] = []
def _note_request(request: object) -> None:
url = getattr(request, "url", "")
if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured:
captured["url"] = url
async def _swallow_redirect(route: "Route") -> None:
await route.fulfill(status=200, content_type="text/plain", body="ok")
async with async_playwright() as playwright:
browser = await playwright.chromium.launch(headless=True)
context = await browser.new_context(storage_state=storage_state_path)
await context.route(
re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect
)
page = await context.new_page()
page.on("request", _note_request)
page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0]))
await page.goto(authorize_url, wait_until="domcontentloaded")
deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT
while "url" not in captured and time.monotonic() < deadline:
try:
await page.wait_for_load_state("networkidle", timeout=8000)
except Exception:
pass
if "url" in captured:
break
control = page.locator(
'button[name="action"][value="approve"], button:has-text("Authorize"), '
'button:has-text("Allow"), button:has-text("@"), a:has-text("@")'
).first
try:
await control.click(timeout=5000)
except Exception:
await asyncio.sleep(0.5)
final_url = page.url
await browser.close()
landing = captured.get("url")
assert landing is not None, (
f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; "
f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}"
)
params = dict(parse_qsl(httpx.URL(landing).query.decode()))
assert "code" in params, f"client redirect_uri carried no code: {landing}"
return params["code"], params.get("state")
def _exchange_code(
code: str, pkce: PkceChallenge, client: DcrClient
) -> OAuthToken:
resp = httpx.post(
DD_TOKEN_URL,
data={
"grant_type": "authorization_code",
"code": code,
"redirect_uri": OAUTH_CLIENT_REDIRECT_URI,
"client_id": client.client_id,
"code_verifier": pkce.verifier,
},
timeout=REQUEST_TIMEOUT,
)
resp.raise_for_status()
data = TypeAdapter(dict[str, object]).validate_python(resp.json())
access_token = str(data["access_token"])
token_type = str(data.get("token_type", "Bearer"))
refresh_token_raw = data.get("refresh_token")
refresh_token = str(refresh_token_raw) if refresh_token_raw is not None else None
return OAuthToken(
access_token=access_token,
token_type=token_type,
refresh_token=refresh_token,
)
def fetch_dd_oauth_token(storage_state_path: str) -> OAuthToken:
"""Drive the full PKCE dance: DCR, authorize (browser), token exchange."""
import asyncio
pkce = _generate_pkce()
dcr = _dcr_register()
params = {
"response_type": "code",
"client_id": dcr.client_id,
"redirect_uri": OAUTH_CLIENT_REDIRECT_URI,
"code_challenge": pkce.challenge,
"code_challenge_method": "S256",
"state": pkce.state,
}
authorize_url = f"{DD_AUTHORIZE_URL}?{urllib.parse.urlencode(params)}"
code, returned_state = asyncio.run(
_browser_authorize(authorize_url, storage_state_path)
)
assert returned_state == pkce.state, (
f"OAuth state mismatch: sent {pkce.state!r}, got {returned_state!r}"
)
return _exchange_code(code, pkce, dcr)
def register_dd_oauth_server(
proxy: ProxyClient, alias: str
) -> DatadogMcpOAuthServer:
"""Register the Datadog MCP server with auth_type=oauth2,
oauth2_flow=authorization_code. The gateway discovers the authorize/token
endpoints via /.well-known metadata."""
resp = unwrap(
proxy.transport.post(
"/v1/mcp/server",
headers=proxy.transport.master,
json=McpServerCreateBody(
alias=alias,
url=DD_MCP_URL,
transport="http",
allow_all_keys=False,
auth_type="oauth2",
oauth2_flow="authorization_code",
),
response_type=McpServerInfo,
)
)
return DatadogMcpOAuthServer(server_id=resp.server_id, alias=alias)
def store_dd_oauth_token(
proxy: ProxyClient,
server_id: str,
key: str,
token: OAuthToken,
) -> None:
"""Store the OAuth access token in the gateway's per-user credential vault
via POST /v1/mcp/server/{server_id}/oauth-user-credential."""
unwrap(
proxy.transport.post(
f"/v1/mcp/server/{server_id}/oauth-user-credential",
headers=AuthHeaders(authorization=f"Bearer {key}"),
json=OAuthCredentialBody(
access_token=token.access_token,
refresh_token=token.refresh_token,
),
response_type=NoBody,
)
)
def delete_dd_oauth_server(proxy: ProxyClient, server_id: str) -> None:
_ = proxy.transport.delete(
f"/v1/mcp/server/{server_id}",
headers=proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)

39
tests/e2e/mcp/dd_session_capture.py Normal file
View file

@ -0,0 +1,39 @@
"""Capture a Datadog browser session for the MCP OAuth e2e tests.
Run this once to log into Datadog and save the browser session:
uv run python tests/e2e/mcp/dd_session_capture.py
Then set the env var and run the OAuth tests:
export E2E_DD_STORAGE_STATE=tests/e2e/mcp/.dd_session.json
uv run pytest tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py -v
"""
from __future__ import annotations
import os
from pathlib import Path
DEFAULT_STATE_PATH = Path(__file__).parent / ".dd_session.json"
def capture(state_path: Path) -> None:
from playwright.sync_api import sync_playwright
with sync_playwright() as p:
browser = p.chromium.launch(headless=False)
context = browser.new_context()
page = context.new_page()
page.goto("https://app.datadoghq.com/account/login")
print("Log into Datadog in the browser, then press Enter here.")
input()
context.storage_state(path=str(state_path))
browser.close()
print(f"Session saved to {state_path}")
print(f' export E2E_DD_STORAGE_STATE="{state_path}"')
if __name__ == "__main__":
state = Path(os.environ.get("E2E_DD_STORAGE_STATE", str(DEFAULT_STATE_PATH)))
capture(state)

118
tests/e2e/mcp/mcp_client.py
View file

@ -18,8 +18,24 @@ from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict, Field, RootModel
from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap
from models import KeyGenerateBody, ObjectPermission
from e2e_http import (
AnthropicHeaders,
Headers,
NoBody,
Result,
StreamingResponse,
Success,
UnknownApiError,
unwrap,
)
from models import (
AnthropicMessagesBody,
AnthropicMessagesResponse,
ChatBody,
ChatResponse,
KeyGenerateBody,
ObjectPermission,
)
from proxy_client import ProxyClient
McpToolArg = str | int | float | bool | list[str] | dict[str, str]
@ -144,6 +160,67 @@ class McpCallToolResponse(BaseModel):
return "\n".join(part.text for part in self.content if part.text)
class ResponsesMcpTool(BaseModel):
type: str = "mcp"
server_label: str
server_url: str
require_approval: str = "never"
allowed_tools: list[str] | None = None
class ResponsesMcpInputMessage(BaseModel):
role: str = "user"
type: str = "message"
content: str
class ResponsesMcpBody(BaseModel):
model: str
input: list[ResponsesMcpInputMessage]
instructions: str | None = None
stream: bool = False
tools: list[ResponsesMcpTool]
class ResponsesMcpOutputContent(BaseModel):
type: str | None = None
text: str | None = None
class ResponsesMcpOutputItem(BaseModel):
model_config = ConfigDict(extra="allow")
type: str | None = None
content: list[ResponsesMcpOutputContent] = []
name: str | None = None
arguments: str | None = None
class ResponsesMcpResult(BaseModel):
model_config = ConfigDict(extra="allow")
id: str | None = None
status: str | None = None
model: str | None = None
output: list[ResponsesMcpOutputItem] = []
@property
def text(self) -> str:
return "".join(
content.text or "" for item in self.output for content in item.content
)
@property
def mcp_tools_fetched(self) -> ResponsesMcpOutputItem | None:
return next(
(item for item in self.output if item.type == "mcp_tools_fetched"), None
)
@property
def tool_execution_results(self) -> ResponsesMcpOutputItem | None:
return next(
(item for item in self.output if item.type == "tool_execution_results"), None
)
@dataclass(frozen=True, slots=True)
class McpClient:
proxy: ProxyClient
@ -371,6 +448,43 @@ class McpClient:
response_type=McpCallToolResponse,
)
def chat_with_mcp(self, key: str, body: ChatBody) -> Result[ChatResponse]:
return self.proxy.transport.post(
"/chat/completions",
headers=self.proxy.transport.bearer(key),
json=body,
response_type=ChatResponse,
)
def responses_with_mcp(
self, key: str, body: ResponsesMcpBody
) -> Result[ResponsesMcpResult]:
return self.proxy.transport.post(
"/v1/responses",
headers=self.proxy.transport.bearer(key),
json=body,
response_type=ResponsesMcpResult,
)
def messages_with_mcp(
self, key: str, body: AnthropicMessagesBody
) -> Result[AnthropicMessagesResponse]:
return self.proxy.transport.post(
"/v1/messages",
headers=AnthropicHeaders(authorization=self.proxy.transport.bearer(key).authorization),
json=body,
response_type=AnthropicMessagesResponse,
)
def messages_stream_with_mcp(
self, key: str, body: AnthropicMessagesBody
) -> StreamingResponse:
return self.proxy.transport.stream(
"/v1/messages",
headers=AnthropicHeaders(authorization=self.proxy.transport.bearer(key).authorization),
json=body,
)
def _is_mcp_not_synced(
result: Result[McpCallToolResponse],

270
tests/e2e/mcp/oauth_chat_client.py
View file

@ -1,270 +0,0 @@
"""Client for the mcp chat-completion OAuth e2e suite.
Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the
per-user upstream token by driving the interactive authorize dance with the
official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium
primed with a human's saved browser session for the OAuth MCP under test), then exercises the server through
/chat/completions, where the gateway lists and executes its tools with the
stored per-user token.
Management routes (/v1/mcp/server CRUD, /chat/completions) go through the
shared ProxyClient transport. The MCP protocol used to seed the token goes through
the mcp SDK, the same library production MCP hosts run.
"""
from __future__ import annotations
import asyncio
import re
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING
from urllib.parse import parse_qsl
import httpx
import pytest
from mcp import ClientSession
from mcp.client.auth import OAuthClientProvider
from mcp.client.streamable_http import streamable_http_client
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
from proxy_client import ProxyClient
from e2e_http import AuthHeaders, NoBody, unwrap
from models import ChatBody, ChatResponse, McpServerCreateBody, McpServerInfo
if TYPE_CHECKING:
from playwright.async_api import Route
# Where the "browser" lands at the end of the authorize dance. Nothing listens
# here: the route interceptor short-circuits the final redirect and reads the
# code/state off its query string, exactly like a desktop MCP host intercepting
# its loopback redirect.
OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback"
BROWSER_CONSENT_TIMEOUT = 60.0
def _mcp_url(alias: str) -> str:
return f"{PROXY_BASE_URL}/{alias}/mcp"
class InMemoryTokenStorage:
"""The mcp SDK's TokenStorage protocol, in memory for one dance: the
DCR-registered client and the gateway tokens minted for it."""
def __init__(self) -> None:
self._tokens: OAuthToken | None = None
self._client_info: OAuthClientInformationFull | None = None
async def get_tokens(self) -> OAuthToken | None:
return self._tokens
async def set_tokens(self, tokens: OAuthToken) -> None:
self._tokens = tokens
async def get_client_info(self) -> OAuthClientInformationFull | None:
return self._client_info
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
self._client_info = client_info
async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> tuple[str, str | None]:
"""Play the browser's role for a real upstream whose authorize endpoint
serves an interactive consent page. A headless Chromium primed with a saved
browser session opens the gateway authorize URL, advances common consent
controls (Approve / Authorize / Allow), and rides the chain through the
gateway callback to the host redirect_uri. The final hop is intercepted and
short-circuited, since nothing listens there, and its code/state are read off
the query string."""
from playwright.async_api import async_playwright
captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener
trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance
def _note_request(request: object) -> None:
url = getattr(request, "url", "")
if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured:
captured["url"] = url
async def _swallow_redirect(route: "Route") -> None:
await route.fulfill(status=200, content_type="text/plain", body="ok")
async with async_playwright() as playwright:
browser = await playwright.chromium.launch(headless=True)
context = await browser.new_context(storage_state=storage_state_path)
await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect)
page = await context.new_page()
page.on("request", _note_request)
page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0]))
await page.goto(start_url, wait_until="domcontentloaded")
deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT
while "url" not in captured and time.monotonic() < deadline:
try:
await page.wait_for_load_state("networkidle", timeout=8000)
except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it
pass
if "url" in captured:
break
control = page.locator(
'button[name="action"][value="approve"], button:has-text("Authorize"), '
'button:has-text("Allow"), button:has-text("@"), a:has-text("@")'
).first
try:
await control.click(timeout=5000)
except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check
await asyncio.sleep(0.5)
final_url = page.url
await browser.close()
landing = captured.get("url")
assert landing is not None, (
f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; "
f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}"
)
params = dict(parse_qsl(httpx.URL(landing).query.decode()))
assert "code" in params, f"client redirect_uri carried no code: {landing}"
return params["code"], params.get("state")
def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path: str) -> OAuthClientProvider:
"""The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR,
PKCE, token exchange) with the browser leg driven by Playwright against the
upstream's consent screen."""
code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks
async def redirect_handler(authorize_url: str) -> None:
code, state = await _browser_follow_authorize(authorize_url, storage_state_path)
code_holder["code"] = code
code_holder["state"] = state
async def callback_handler() -> tuple[str, str | None]:
code = code_holder.get("code")
assert code is not None, "callback_handler ran before the authorize redirect completed"
return code, code_holder.get("state")
return OAuthClientProvider(
server_url=url,
client_metadata=OAuthClientMetadata.model_validate(
{
"redirect_uris": [OAUTH_CLIENT_REDIRECT_URI],
"token_endpoint_auth_method": "none",
"grant_types": ["authorization_code", "refresh_token"],
"response_types": ["code"],
"client_name": "e2e-mcp-host",
}
),
storage=storage,
redirect_handler=redirect_handler,
callback_handler=callback_handler,
)
class _HeaderInjectingTransport(httpx.AsyncBaseTransport):
"""Adds the caller's LiteLLM key header to every outgoing SDK request
(discovery, DCR, token exchange), so the gateway resolves which user to
store the upstream token for from the key on the token exchange, exactly
like a production MCP host configured with a LiteLLM key header."""
def __init__(self, inner: httpx.AsyncBaseTransport, headers: dict[str, str]) -> None:
self._inner = inner
self._headers = headers
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
for name, value in self._headers.items():
if name not in request.headers:
request.headers[name] = value
return await self._inner.handle_async_request(request)
def _oauth_http_client(headers: dict[str, str], auth: OAuthClientProvider) -> httpx.AsyncClient:
return httpx.AsyncClient(
headers=headers,
auth=auth,
timeout=httpx.Timeout(REQUEST_TIMEOUT),
follow_redirects=True,
transport=_HeaderInjectingTransport(httpx.AsyncHTTPTransport(), headers),
)
async def _seed_via_dance(
url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str
) -> tuple[str, ...]:
async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) 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))
@dataclass(frozen=True, slots=True)
class ChatMcpClient:
proxy: ProxyClient
def create_server(self, body: McpServerCreateBody) -> McpServerInfo:
return unwrap(
self.proxy.transport.post(
"/v1/mcp/server",
headers=self.proxy.transport.master,
json=body,
response_type=McpServerInfo,
)
)
def server_info(self, server_id: str) -> McpServerInfo:
return unwrap(
self.proxy.transport.get(
f"/v1/mcp/server/{server_id}",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpServerInfo,
)
)
def delete_server(self, server_id: str) -> None:
_ = self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]:
"""Drive the interactive authorize dance for `key`'s user so the gateway
stores their upstream token, retried to the shared deadline since the
just-created server and key propagate asynchronously. The LiteLLM key
rides x-litellm-api-key so the gateway binds the token to that user.
Returns the upstream tool names the dance listed, proof the token works."""
headers = {"x-litellm-api-key": f"Bearer {key}"}
storage = InMemoryTokenStorage()
deadline = time.monotonic() + self.proxy.poll_timeout
last_error: Exception | None = None
while time.monotonic() < deadline:
try:
return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path))
except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below
last_error = exc
time.sleep(self.proxy.poll_interval)
pytest.fail(
f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; "
f"last error: {last_error!r}"
)
def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse:
"""POST /chat/completions carrying the LiteLLM key in `headers` (either
ingress form) with an MCP server attached in `body.tools`. The gateway
resolves the user from the key and lists/executes the server's tools
with that user's stored upstream token."""
return unwrap(
self.proxy.transport.post(
"/chat/completions",
headers=headers,
json=body,
response_type=ChatResponse,
)
)
def build_chat_client(proxy: ProxyClient) -> ChatMcpClient:
return ChatMcpClient(proxy=proxy)

2
tests/e2e/mcp/test_mcp_access_group_e2e.py
View file

@ -28,7 +28,7 @@ class TestMcpAccessGroupToolSelection:
self, client: McpClient, resources: ResourceManager
) -> None:
group = f"e2e-mcp-grp-{unique_marker()}"
server_id = register_datadog_mcp(client, resources, mcp_access_groups=[group])
server_id = register_datadog_mcp(client, resources, mcp_access_groups=[group]).server_id
client.await_registered(server_id)
granted = client.generate_key(

35
tests/e2e/mcp/test_mcp_auth_errors_e2e.py Normal file
View file

@ -0,0 +1,35 @@
"""Live e2e: the MCP REST endpoints return 401 for an invalid key, not a
flattened 500.
Before PR #31011 the MCP protocol path flattened auth errors to 500; the REST
path shares the same user_api_key_auth dependency as /chat/completions, so a
401 here proves the gateway's auth error mapping is intact. Budget enforcement
(429) is the same dependency as chat and is covered by the quota_management
suite.
"""
from __future__ import annotations
import pytest
from e2e_http import UnauthorizedError
from lifecycle import ResourceManager
from mcp_client import McpClient
pytestmark = pytest.mark.e2e
GARBAGE_KEY = "sk-deadbeef-not-a-real-key"
class TestMcpAuthStatusCodes:
@pytest.mark.covers("mcp.auth.api_key.returns_401_not_500")
def test_invalid_key_returns_401_not_500(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
result = client.list_tools(GARBAGE_KEY)
assert isinstance(result, UnauthorizedError), (
f"invalid key on /mcp-rest/tools/list must return 401, not 500; "
f"got: {result}"
)

99
tests/e2e/mcp/test_mcp_chat_completions_e2e.py Normal file
View file

@ -0,0 +1,99 @@
"""Live e2e: /chat/completions expands a gateway-registered MCP server and
auto-executes its tools in one agentic turn.
Registers the real Datadog remote MCP server, grants a key access to it, then
sends a /chat/completions request whose ``tools`` array carries an
``{type: "mcp", server_url: "litellm_proxy/mcp/<alias>"}`` reference. The gateway
lists the server's tools, feeds them to the model, the model calls
search_datadog_logs, the gateway executes the call upstream, and folds the
result back into a follow-up completion. The response must carry
provider_specific_fields.mcp_list_tools (the gateway listed tools),
mcp_tool_calls (the model called one), and mcp_call_results (the gateway
executed it), proving the full bridge loop ran end to end against a real MCP
server and a real LLM.
"""
from __future__ import annotations
import pytest
from datadog_mcp import assert_dd_mcp_creds, register_datadog_mcp
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient
from models import ChatBody, ChatMessage, McpChatTool
pytestmark = pytest.mark.e2e
class TestChatCompletionMcpAutoExecute:
@pytest.mark.covers("mcp.chat_completion.api_key.auto_executes_tools")
def test_chat_completion_lists_calls_and_executes_mcp_tools(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd = register_datadog_mcp(client, resources)
client.await_registered(dd.server_id)
key = client.generate_key(
user_id=f"e2e-mcp-chat-{unique_marker()}",
mcp_servers=[dd.server_id],
models=[CHEAP_ANTHROPIC_MODEL],
)
resources.defer(lambda: client.proxy.delete_key(key))
body = ChatBody(
model=CHEAP_ANTHROPIC_MODEL,
messages=[
ChatMessage(
role="user",
content=(
"Use the search_datadog_logs tool to search for logs "
"with query 'service:litellm' from now-30m to now with "
"max_tokens 500. After you get results, summarize what you found."
),
)
],
max_tokens=1024,
tools=[
McpChatTool(
type="mcp",
server_url=f"litellm_proxy/mcp/{dd.alias}",
server_label="datadog",
require_approval="never",
)
],
)
response = unwrap(client.chat_with_mcp(key, body))
assert response.choices, f"chat completion returned no choices: {response}"
message = response.choices[0].message
assert message is not None, f"choice had no message: {response}"
psf = message.provider_specific_fields
assert psf is not None, (
"provider_specific_fields missing; the gateway did not attach MCP metadata "
f"(mcp_list_tools / mcp_tool_calls / mcp_call_results): {message}"
)
assert psf.mcp_list_tools, (
"mcp_list_tools is empty; the gateway never listed the Datadog server's tools "
"through the chat bridge"
)
assert psf.mcp_tool_calls, (
"mcp_tool_calls is empty; the model did not call any MCP tool "
"(it may not have seen the expanded tools)"
)
assert psf.mcp_call_results, (
"mcp_call_results is empty; the gateway did not execute the tool call upstream"
)
result_text = next(
(r.result for r in psf.mcp_call_results if r.result), None
)
assert result_text, (
"mcp_call_results has no result text; the tool call returned nothing"
)

23
tests/e2e/mcp/test_mcp_datadog_e2e.py
View file

@ -49,16 +49,6 @@ def _seed_completion(proxy: ProxyClient, *, key: str, marker: str) -> None:
class TestDatadogMcpRoundTrip:
@pytest.mark.skip(
reason=(
"LIT-5052: this test sends a `telemetry` argument that Datadog's "
"search_datadog_logs tool now rejects, so every tool call fails validation with "
"'unexpected additional properties [\"telemetry\"]' before the round-trip "
"assertion is reached. `telemetry` was never a documented Datadog parameter; the "
"test relied on the server ignoring unknown properties. Unskip once the argument "
"is dropped."
)
)
@pytest.mark.covers("mcp.list_tools.api_key.succeeds", "mcp.call_tool.api_key.succeeds")
def test_search_logs_finds_seeded_completion(
self,
@ -69,13 +59,13 @@ class TestDatadogMcpRoundTrip:
assert_dd_mcp_creds()
_assert_datadog_logger_active(client.proxy)
server_id = register_datadog_mcp(client, resources)
client.await_registered(server_id)
dd = register_datadog_mcp(client, resources)
client.await_registered(dd.server_id)
marker = f"{MARKER_PREFIX}{unique_marker()}"
key = client.generate_key(
user_id=f"e2e-dd-mcp-{unique_marker()}",
mcp_servers=[server_id],
mcp_servers=[dd.server_id],
models=[CHEAP_ANTHROPIC_MODEL],
)
resources.defer(lambda: client.proxy.delete_key(key))
@ -88,19 +78,16 @@ class TestDatadogMcpRoundTrip:
"within the poll deadline; MCP search would have nothing to find"
)
tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL)
tool_name = client.await_tool(key, dd.server_id, SEARCH_LOGS_TOOL)
call = client.await_call_tool(
key,
server_id=server_id,
server_id=dd.server_id,
name=tool_name,
arguments={
"query": marker,
"from": DD_SEARCH_FROM,
"to": "now",
"max_tokens": 5000,
"telemetry": {
"intent": "e2e assert seeded litellm completion log is searchable via MCP"
},
},
)
assert call.is_error is not True, f"search_datadog_logs errored: {call}"

159
tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py Normal file
View file

@ -0,0 +1,159 @@
"""Live e2e: the Datadog MCP server through gateway-managed OAuth2 (PKCE).
Registers the Datadog MCP server with auth_type=oauth2,
oauth2_flow=authorization_code. The gateway discovers the OAuth endpoints via
/.well-known metadata. A real PKCE authorize dance (DCR, browser consent,
token exchange) produces an access token, which is stored in the gateway's
per-user credential vault. Then the key lists tools, calls one, and drives a
chat completion through the MCP bridge, all using the stored per-user token.
Requires E2E_DD_STORAGE_STATE pointing at a saved Datadog browser session.
"""
from __future__ import annotations
import os
import pytest
from dd_oauth import (
assert_dd_oauth_env,
delete_dd_oauth_server,
fetch_dd_oauth_token,
register_dd_oauth_server,
store_dd_oauth_token,
)
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient
from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, ObjectPermission
pytestmark = [
pytest.mark.e2e,
pytest.mark.skipif(
not os.environ.get("E2E_DD_STORAGE_STATE"),
reason="set E2E_DD_STORAGE_STATE to a Datadog session captured via mcp/dd_session_capture.py",
),
]
pytest.importorskip("mcp", reason="mcp SDK not installed")
pytest.importorskip("playwright.async_api", reason="playwright not installed")
SEARCH_LOGS_TOOL = "search_datadog_logs"
class TestDatadogMcpOAuth:
@pytest.mark.covers(
"mcp.list_tools.oauth.succeeds",
"mcp.call_tool.oauth.succeeds",
)
def test_oauth_list_and_call_tools(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_oauth_env()
marker = unique_marker()
alias = f"e2e_dd_oauth_{marker}"
dd = register_dd_oauth_server(client.proxy, alias)
resources.defer(lambda: delete_dd_oauth_server(client.proxy, dd.server_id))
client.await_registered(dd.server_id)
key = client.proxy.generate_key(
KeyGenerateBody(
models=[CHEAP_ANTHROPIC_MODEL],
user_id=f"e2e-dd-oauth-{marker}",
object_permission=ObjectPermission(mcp_servers=[dd.server_id]),
)
)
resources.defer(lambda: client.proxy.delete_key(key))
token = fetch_dd_oauth_token(os.environ["E2E_DD_STORAGE_STATE"])
store_dd_oauth_token(client.proxy, dd.server_id, key, token)
tool_name = client.await_tool(key, dd.server_id, SEARCH_LOGS_TOOL)
assert tool_name, f"OAuth token did not surface any tools for server {dd.server_id}"
call = client.await_call_tool(
key,
server_id=dd.server_id,
name=tool_name,
arguments={
"query": "service:litellm",
"from": "now-30m",
"to": "now",
"max_tokens": 500,
},
)
assert call.is_error is not True, f"search_datadog_logs errored via OAuth: {call}"
assert call.all_text, f"OAuth tool call returned empty text: {call}"
@pytest.mark.covers("mcp.chat_completion.oauth.auto_executes_tools")
def test_oauth_chat_completion_auto_executes_tools(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_oauth_env()
marker = unique_marker()
alias = f"e2e_dd_oauth_chat_{marker}"
dd = register_dd_oauth_server(client.proxy, alias)
resources.defer(lambda: delete_dd_oauth_server(client.proxy, dd.server_id))
client.await_registered(dd.server_id)
key = client.proxy.generate_key(
KeyGenerateBody(
models=[CHEAP_ANTHROPIC_MODEL],
user_id=f"e2e-dd-oauth-chat-{marker}",
object_permission=ObjectPermission(mcp_servers=[dd.server_id]),
)
)
resources.defer(lambda: client.proxy.delete_key(key))
token = fetch_dd_oauth_token(os.environ["E2E_DD_STORAGE_STATE"])
store_dd_oauth_token(client.proxy, dd.server_id, key, token)
response = unwrap(
client.chat_with_mcp(
key,
ChatBody(
model=CHEAP_ANTHROPIC_MODEL,
messages=[
ChatMessage(
role="user",
content=(
"Use the search_datadog_logs tool to search for logs "
"with query 'service:litellm' from now-30m to now with "
"max_tokens 500. After you get results, summarize what you found."
),
)
],
max_tokens=1024,
tools=[
McpChatTool(
type="mcp",
server_url=f"litellm_proxy/mcp/{dd.alias}",
server_label="datadog",
require_approval="never",
)
],
),
)
)
assert response.choices, f"chat completion returned no choices: {response}"
message = response.choices[0].message
assert message is not None, f"choice had no message: {response}"
psf = message.provider_specific_fields
assert psf is not None, (
f"provider_specific_fields missing; the gateway did not attach MCP metadata: {message}"
)
assert psf.mcp_list_tools, (
f"mcp_list_tools is empty; the gateway never listed tools via the stored OAuth token: {psf}"
)
assert psf.mcp_call_results, (
f"mcp_call_results is empty; the gateway did not execute any tool via OAuth: {psf}"
)

178
tests/e2e/mcp/test_mcp_guardrail_e2e.py
View file

@ -1,178 +0,0 @@
"""Live e2e: a guardrail on the MCP tool-call path blocks banned content in the
tool arguments before the call reaches the upstream MCP server.
A general litellm_content_filter guardrail is configured with mode=pre_mcp_call
(the event type the proxy rewrites pre_call to for a call_mcp_tool) and default_on
(per-key/request guardrail selection is dropped from the synthetic MCP request the
hook sees, so default_on is how it attaches to tools/call). The banned keyword is
unique per run, so default_on only ever intercepts this test's own banned call.
Against the real Datadog MCP server, calling search_datadog_logs with the banned
keyword in the query is blocked with HTTP 400 attributed to the pre_mcp_call hook,
and the tool never runs; the same guardrail lets a clean query through to Datadog.
This is the enforced half (the block) plus the pass-through half in one spec.
"""
from __future__ import annotations
import time
from collections.abc import Callable
import pytest
from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp
from e2e_config import DD_SEARCH_FROM, unique_marker
from e2e_http import Result, Success, UnknownApiError
from lifecycle import ResourceManager
from mcp_client import McpCallToolResponse, McpClient, McpToolArguments
pytestmark = pytest.mark.e2e
# Stage runs several data-plane pods behind the shared key, and each picks up a
# newly registered guardrail or MCP server only on its next periodic DB sync (~30s in
# proxy_server.py). Every pod is guaranteed to have refreshed only once a full sync
# interval has elapsed since the later of those two writes; before then a banned call
# routed to a lagging pod passes through as legitimate in-flight propagation, not a leak.
FULL_SYNC_SECONDS = 40.0
POST_SYNC_VERIFICATION_CALLS = 4
def _poll_until_blocked(
search: Callable[[str], Result[McpCallToolResponse]], banned_keyword: str, client: McpClient
) -> Result[McpCallToolResponse]:
"""Retry a banned tool call until the guardrail blocks it (400) or the deadline
passes, returning the last result. Absorbs the control-plane -> data-plane
guardrail-sync delay so the check waits for enforcement instead of racing it."""
deadline = time.monotonic() + client.proxy.poll_timeout
last: Result[McpCallToolResponse] = search(f"tell me about {banned_keyword}")
while time.monotonic() < deadline:
if isinstance(last, UnknownApiError) and last.status_code == 400:
return last
time.sleep(client.proxy.poll_interval)
last = search(f"tell me about {banned_keyword}")
return last
def _pod_lacks_mcp_server(result: Result[McpCallToolResponse]) -> bool:
"""True when the pod that served the call answered as though the MCP server or its
tool does not exist (500 "Tool ... not found"), i.e. its MCP registry has not synced
yet and the request never reached the guardrail at all."""
if not isinstance(result, UnknownApiError) or result.status_code != 500:
return False
body = result.body.lower()
return "not found" in body and ("tool" in body or "server" in body)
def _search_on_synced_pod(
search: Callable[[str], Result[McpCallToolResponse]], query: str, client: McpClient
) -> Result[McpCallToolResponse]:
"""Issue `query`, retrying to the poll deadline only while the serving pod does not
know the MCP server yet. Every other outcome, guardrail block or pass-through, comes
back untouched so the caller's assertion still decides it."""
deadline = time.monotonic() + client.proxy.poll_timeout
last = search(query)
while _pod_lacks_mcp_server(last) and time.monotonic() < deadline:
time.sleep(client.proxy.poll_interval)
last = search(query)
return last
class TestMcpToolCallGuardrail:
@pytest.mark.skip(
reason=(
"LIT-5052: the control call sends a `telemetry` argument that Datadog's "
"search_datadog_logs tool now rejects, so the clean-argument half of this test "
"errors with 'unexpected additional properties [\"telemetry\"]' and the guardrail "
"block it exists to prove is never exercised. `telemetry` was never a documented "
"Datadog parameter; the test relied on the server ignoring unknown properties. "
"Unskip once the argument is dropped."
)
)
@pytest.mark.covers(
"guardrail.litellm_content_filter.pre_mcp_call.blocks",
exercised_on=["mcp_operations"],
)
def test_content_filter_blocks_banned_keyword_in_tool_args(
self, client: McpClient, resources: ResourceManager
) -> None:
assert_dd_mcp_creds()
marker = unique_marker()
banned_keyword = f"e2eblocked{marker}"
guardrail_id = client.register_mcp_content_filter(
name=f"e2e-mcp-cf-{marker}", blocked_keyword=banned_keyword
)
guardrail_created_at = time.monotonic()
resources.defer(lambda: client.delete_guardrail(guardrail_id))
server_id = register_datadog_mcp(client, resources)
server_registered_at = time.monotonic()
key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id])
resources.defer(lambda: client.proxy.delete_key(key))
tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL)
def search(query: str) -> Result[McpCallToolResponse]:
arguments: McpToolArguments = {
"query": query,
"from": DD_SEARCH_FROM,
"to": "now",
"max_tokens": 500,
"telemetry": {"intent": "e2e mcp guardrail check"},
}
return client.call_tool(key, server_id=server_id, name=tool_name, arguments=arguments)
# Registering the guardrail is a control-plane write; the data-plane worker
# that serves tools/call picks it up on its next guardrail sync, so an
# immediate call can race the propagation and slip through. Poll the banned
# call to the deadline and require a block, so the check proves enforcement
# rather than catching a pre-sync pass-through. The keyword is unique per
# run, so this only ever intercepts this test's own call.
blocked = _poll_until_blocked(search, banned_keyword, client)
match blocked:
case UnknownApiError(status_code=400, body=body):
assert banned_keyword in body or "content blocked" in body.lower(), (
f"the block must name the content-filter reason, got: {body[:300]}"
)
assert "pre_mcp_call" in body, (
f"the block must be attributed to the MCP tool-call hook (pre_mcp_call), got: {body[:300]}"
)
case _:
pytest.fail(
"content_filter never blocked the banned keyword on the MCP tool call within "
f"{client.proxy.poll_timeout}s (the guardrail or the MCP server never synced to "
f"the data plane); last result: {blocked}"
)
# The block above only proves the one pod that served it has synced; another
# pod could still lack the guardrail and let the banned call reach Datadog.
# Wait out the full sync interval from the later of the guardrail create and the
# MCP server registration (each syncs on its own clock, so the earlier write's
# deadline can elapse while a pod still lacks the other) so every pod has
# refreshed from the DB, then require the banned call to stay blocked across
# several attempts. A pass-through now is a genuine partial-propagation leak, not
# a race. Client load balancing still can't guarantee every pod is hit, so this
# samples several worker selections rather than proving all pods synced.
sync_remaining = max(guardrail_created_at, server_registered_at) + FULL_SYNC_SECONDS - time.monotonic()
if sync_remaining > 0:
time.sleep(sync_remaining)
for attempt in range(1, POST_SYNC_VERIFICATION_CALLS + 1):
reblocked = _search_on_synced_pod(search, f"still about {banned_keyword} #{attempt}", client)
assert isinstance(reblocked, UnknownApiError) and reblocked.status_code == 400, (
"after the sync interval every data-plane pod must block the banned keyword, but "
f"attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was not blocked (a pod still "
f"lacks the guardrail, or never synced the MCP server): {reblocked}"
)
if attempt < POST_SYNC_VERIFICATION_CALLS:
time.sleep(client.proxy.poll_interval)
allowed = _search_on_synced_pod(search, f"e2e-clean-{marker}", client)
match allowed:
case Success(data=result):
assert result.is_error is not True, (
f"a clean MCP tool call must reach the server and not error, got: {result}"
)
case _:
pytest.fail(
f"a clean MCP tool call must pass the guardrail and reach the server; got {allowed}"
)

15
tests/e2e/mcp/test_mcp_key_access_e2e.py
View file

@ -37,7 +37,7 @@ class TestMcpKeyWithoutAccessIsDenied:
client: McpClient,
resources: ResourceManager,
) -> None:
server_id = register_datadog_mcp(client, resources)
server_id = register_datadog_mcp(client, resources).server_id
client.await_registered(server_id)
permitted_key = _key(client, resources, mcp_servers=[server_id])
@ -51,23 +51,13 @@ class TestMcpKeyWithoutAccessIsDenied:
f"boundary: {denied_tools}"
)
@pytest.mark.skip(
reason=(
"LIT-5052: the control call proving a granted key CAN invoke the tool sends a "
"`telemetry` argument that Datadog's search_datadog_logs tool now rejects, so it "
"errors with 'unexpected additional properties [\"telemetry\"]' and the denial "
"assertion is never reached. `telemetry` was never a documented Datadog "
"parameter; the test relied on the server ignoring unknown properties. Unskip "
"once the argument is dropped."
)
)
@pytest.mark.covers("mcp.call_tool.api_key.denied_without_permission")
def test_call_tool_denied_without_permission(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
server_id = register_datadog_mcp(client, resources)
server_id = register_datadog_mcp(client, resources).server_id
client.await_registered(server_id)
permitted_key = _key(client, resources, mcp_servers=[server_id])
@ -80,7 +70,6 @@ class TestMcpKeyWithoutAccessIsDenied:
"from": DD_SEARCH_FROM,
"to": "now",
"max_tokens": 1000,
"telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"},
}
permitted_call = client.await_call_tool(
permitted_key, server_id=server_id, name=tool_name, arguments=search_args

114
tests/e2e/mcp/test_mcp_messages_e2e.py Normal file
View file

@ -0,0 +1,114 @@
"""Live e2e: /v1/messages expands a gateway-registered MCP server and
auto-executes its tools, running the tool_use loop internally and returning a
final Anthropic response (the Claude Code path).
Registers the real Datadog remote MCP server, grants a key access to it, then
sends a /v1/messages request whose ``tools`` array carries an
``{type: "mcp", server_url: "litellm_proxy/mcp/<alias>"}`` reference. The
gateway intercepts the litellm_proxy reference (which Anthropic cannot reach),
expands it into native Anthropic custom tools under the caller's credentials,
runs the tool_use loop (model calls search_datadog_logs, gateway executes it,
feeds the result back as a tool_result), and returns the final answer. The
response has no MCP-specific metadata; the proof is the final text answer,
meaning the loop completed and the model used the tool result.
The streaming variant exercises the same loop but with stream=True, where the
gateway runs the tool_use loop non-streaming internally and then fakes a stream
of the final answer as Anthropic SSE events.
"""
from __future__ import annotations
import pytest
from datadog_mcp import assert_dd_mcp_creds, register_datadog_mcp
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_http import require_successful_call, unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient
from models import AnthropicMcpTool, AnthropicMessagesBody, ChatMessage
pytestmark = pytest.mark.e2e
TOOL_PROMPT = (
"Use the search_datadog_logs tool to search for logs "
"with query 'service:litellm' from now-30m to now with "
"max_tokens 500. After you get results, summarize what you found in one sentence."
)
def _messages_body(model: str, alias: str) -> AnthropicMessagesBody:
return AnthropicMessagesBody(
model=model,
max_tokens=1024,
messages=[ChatMessage(role="user", content=TOOL_PROMPT)],
tools=[
AnthropicMcpTool(
server_label="datadog",
server_url=f"litellm_proxy/mcp/{alias}",
require_approval="never",
)
],
)
class TestMessagesMcpAutoExecute:
@pytest.mark.covers("mcp.messages.api_key.auto_executes_tools")
def test_messages_runs_tool_loop_and_returns_final_answer(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd = register_datadog_mcp(client, resources)
client.await_registered(dd.server_id)
key = client.generate_key(
user_id=f"e2e-mcp-msg-{unique_marker()}",
mcp_servers=[dd.server_id],
models=[CHEAP_ANTHROPIC_MODEL],
)
resources.defer(lambda: client.proxy.delete_key(key))
response = unwrap(
client.messages_with_mcp(key, _messages_body(CHEAP_ANTHROPIC_MODEL, dd.alias))
)
assert response.content, f"/v1/messages returned no content blocks: {response}"
text = "".join(block.text or "" for block in response.content)
assert text.strip(), (
f"/v1/messages returned no text after the MCP tool loop; the gateway "
f"may not have completed the tool_use loop: {response}"
)
@pytest.mark.covers("mcp.messages.api_key.stream_auto_executes_tools")
def test_messages_stream_runs_tool_loop_and_returns_final_answer(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd = register_datadog_mcp(client, resources)
client.await_registered(dd.server_id)
key = client.generate_key(
user_id=f"e2e-mcp-msg-stream-{unique_marker()}",
mcp_servers=[dd.server_id],
models=[CHEAP_ANTHROPIC_MODEL],
)
resources.defer(lambda: client.proxy.delete_key(key))
body = _messages_body(CHEAP_ANTHROPIC_MODEL, dd.alias)
body.stream = True
result = client.messages_stream_with_mcp(key, body)
require_successful_call(result)
assert result.is_streaming, f"response was not streamed: {result.headers}"
assert not result.stream_error, f"stream errored: {result.stream_error}"
assert result.stream_events, "stream produced no SSE events"
assert any("content_block_delta" in event for event in result.stream_events), (
"stream carried no content deltas"
)
assert any("message_stop" in event for event in result.stream_events), (
"stream never reached message_stop"
)

83
tests/e2e/mcp/test_mcp_responses_e2e.py Normal file
View file

@ -0,0 +1,83 @@
"""Live e2e: /v1/responses expands a gateway-registered MCP server and
auto-executes its tools, surfacing the results as response output items.
Registers the real Datadog remote MCP server, grants a key access to it, then
sends a /v1/responses request whose ``tools`` array carries an
``{type: "mcp", server_url: "litellm_proxy/mcp/<alias>"}`` reference. The
gateway lists the server's tools, feeds them to the model in Responses API
format, the model calls search_datadog_logs, the gateway executes the call
upstream, and appends ``mcp_tools_fetched`` and ``tool_execution_results``
output items to the response. Their presence proves the full Responses API
bridge loop ran end to end against a real MCP server and a real LLM.
"""
from __future__ import annotations
import pytest
from datadog_mcp import assert_dd_mcp_creds, register_datadog_mcp
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient, ResponsesMcpBody, ResponsesMcpInputMessage, ResponsesMcpTool
pytestmark = pytest.mark.e2e
class TestResponsesMcpAutoExecute:
@pytest.mark.covers("mcp.responses.api_key.auto_executes_tools")
def test_responses_lists_calls_and_executes_mcp_tools(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd = register_datadog_mcp(client, resources)
client.await_registered(dd.server_id)
key = client.generate_key(
user_id=f"e2e-mcp-resp-{unique_marker()}",
mcp_servers=[dd.server_id],
models=[CHEAP_ANTHROPIC_MODEL],
)
resources.defer(lambda: client.proxy.delete_key(key))
body = ResponsesMcpBody(
model=CHEAP_ANTHROPIC_MODEL,
input=[
ResponsesMcpInputMessage(
content=(
"Use the search_datadog_logs tool to search for logs "
"with query 'service:litellm' from now-30m to now with "
"max_tokens 500. After you get results, summarize what you found."
)
)
],
instructions="You are a helpful assistant.",
tools=[
ResponsesMcpTool(
server_label="datadog",
server_url=f"litellm_proxy/mcp/{dd.alias}",
require_approval="never",
)
],
)
result = unwrap(client.responses_with_mcp(key, body))
fetched = result.mcp_tools_fetched
assert fetched is not None, (
"response.output has no mcp_tools_fetched item; the gateway did not list "
f"the Datadog server's tools through the responses bridge: {result.output}"
)
assert fetched.content, (
"mcp_tools_fetched item has no content; the tool list was empty"
)
executed = result.tool_execution_results
assert executed is not None, (
"response.output has no tool_execution_results item; the gateway did not "
f"execute any MCP tool through the responses bridge: {result.output}"
)
assert executed.content, (
"tool_execution_results item has no content; the tool call returned nothing"
)

134
tests/e2e/mcp/test_mcp_server_config_e2e.py Normal file
View file

@ -0,0 +1,134 @@
"""Live e2e: the gateway brokers Datadog's static-header HTTP transport and
multi-server tool namespacing.
Three cells in one spec:
1. upstream_static_auth: a server registered with static_headers (Datadog's
DD-API-KEY / DD-APPLICATION-KEY) injects them on every upstream call, so
search_datadog_logs succeeds. The tool result must be non-empty, proving
the static credentials reached Datadog.
2. transport_http: the Datadog server uses the streamable HTTP transport, and
a successful tools/list + tools/call round-trip proves that transport path
works end to end. This is the same call as upstream_static_auth but asserts
the transport-specific cell.
3. namespaced_multi_server: two registered servers' tools remain
distinguishable on the aggregate tools/list (each tool carries its own
mcp_info.server_id), so a multi-server tenant never sees tools collide.
"""
from __future__ import annotations
import pytest
from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp
from e2e_config import DD_SEARCH_FROM, unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient, McpToolArguments
pytestmark = pytest.mark.e2e
def _search_args(query: str) -> McpToolArguments:
return {
"query": query,
"from": DD_SEARCH_FROM,
"to": "now",
"max_tokens": 500,
}
class TestUpstreamStaticAuthAndTransport:
@pytest.mark.covers(
"mcp.call_tool.api_key.upstream_static_auth",
"mcp.call_tool.api_key.transport_http",
)
def test_static_header_http_transport_call_succeeds(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd = register_datadog_mcp(client, resources)
client.await_registered(dd.server_id)
key = client.generate_key(
user_id=f"e2e-mcp-static-{unique_marker()}",
mcp_servers=[dd.server_id],
)
resources.defer(lambda: client.proxy.delete_key(key))
tool_name = client.await_tool(key, dd.server_id, SEARCH_LOGS_TOOL)
call = client.await_call_tool(
key,
server_id=dd.server_id,
name=tool_name,
arguments=_search_args("service:litellm"),
)
assert call.is_error is not True, (
f"search_datadog_logs errored with static headers + http transport: {call}"
)
assert call.all_text, (
"search_datadog_logs returned empty text; the static DD-API-KEY / "
"DD-APPLICATION-KEY headers may not have reached the upstream"
)
class TestNamespacedMultiServer:
@pytest.mark.covers("mcp.list_tools.api_key.namespaced_multi_server")
def test_two_servers_tools_remain_distinguishable(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd_a = register_datadog_mcp(client, resources)
dd_b = register_datadog_mcp(client, resources)
client.await_registered(dd_a.server_id)
client.await_registered(dd_b.server_id)
key_a = client.generate_key(
user_id=f"e2e-mcp-ns-a-{unique_marker()}",
mcp_servers=[dd_a.server_id],
)
resources.defer(lambda: client.proxy.delete_key(key_a))
key_b = client.generate_key(
user_id=f"e2e-mcp-ns-b-{unique_marker()}",
mcp_servers=[dd_b.server_id],
)
resources.defer(lambda: client.proxy.delete_key(key_b))
_ = client.await_tool(key_a, dd_a.server_id, SEARCH_LOGS_TOOL)
_ = client.await_tool(key_b, dd_b.server_id, SEARCH_LOGS_TOOL)
key_both = client.generate_key(
user_id=f"e2e-mcp-ns-both-{unique_marker()}",
mcp_servers=[dd_a.server_id, dd_b.server_id],
)
resources.defer(lambda: client.proxy.delete_key(key_both))
tools = unwrap(client.list_tools(key_both))
a_tools = tools.tool_names_for_server(dd_a.server_id)
b_tools = tools.tool_names_for_server(dd_b.server_id)
assert a_tools, (
f"server A's tools are missing from the aggregate list; "
f"the multi-server namespace collapsed: {tools.tools}"
)
assert b_tools, (
f"server B's tools are missing from the aggregate list; "
f"the multi-server namespace collapsed: {tools.tools}"
)
a_entries = tuple(
t for t in tools.tools
if t.mcp_info and t.mcp_info.server_id == dd_a.server_id
)
b_entries = tuple(
t for t in tools.tools
if t.mcp_info and t.mcp_info.server_id == dd_b.server_id
)
assert len(a_entries) == len(b_tools) and len(b_entries) == len(b_tools), (
f"each server's tools must carry its own mcp_info.server_id so a "
f"multi-server tenant can tell them apart; "
f"A entries={a_entries}, B entries={b_entries}"
)

100
tests/e2e/mcp/test_mcp_tool_filtering_e2e.py Normal file
View file

@ -0,0 +1,100 @@
"""Live e2e: the gateway's allowed_tools filtering on the Datadog MCP server.
A server registered with a narrow allowed_tools list hides every other tool
from tools/list, and tools/call on a hidden tool is blocked (403 or 404,
depending on whether the tool was in the gateway's resolved tool map).
Note: disallowed_tools and allowed_params are config-only fields today, they
have no DB column in the Prisma schema and are silently dropped on the
management API path. Those cells are product gaps, not test gaps.
All against the real Datadog remote MCP server.
"""
from __future__ import annotations
import pytest
from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp
from e2e_config import DD_SEARCH_FROM, unique_marker
from e2e_http import UnknownApiError, unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient, McpToolArguments
pytestmark = pytest.mark.e2e
BOGUS_TOOL = "nonexistent_e2e_tool"
def _search_args(query: str) -> McpToolArguments:
return {
"query": query,
"from": DD_SEARCH_FROM,
"to": "now",
"max_tokens": 500,
}
def _key(
client: McpClient, resources: ResourceManager, server_id: str, label: str
) -> str:
key = client.generate_key(
user_id=f"e2e-mcp-{label}-{unique_marker()}",
mcp_servers=[server_id],
)
resources.defer(lambda: client.proxy.delete_key(key))
return key
class TestAllowedToolsScoping:
@pytest.mark.covers("mcp.list_tools.api_key.allowed_tools_scoped")
def test_list_tools_hides_non_allowed_tools(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd = register_datadog_mcp(client, resources, allowed_tools=[BOGUS_TOOL])
client.await_registered(dd.server_id)
key = _key(client, resources, dd.server_id, "allow-list")
tools = unwrap(client.list_tools(key)).tool_names_for_server(dd.server_id)
assert SEARCH_LOGS_TOOL not in tools, (
f"search_datadog_logs must be hidden by the allowed_tools filter "
f"(only {BOGUS_TOOL!r} is allowed), but it appeared in tools/list: {tools}"
)
@pytest.mark.covers("mcp.call_tool.api_key.allowed_tools_scoped")
def test_call_tool_denied_outside_allowed_tools(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
assert_dd_mcp_creds()
dd_all = register_datadog_mcp(client, resources)
client.await_registered(dd_all.server_id)
dd_narrow = register_datadog_mcp(client, resources, allowed_tools=[BOGUS_TOOL])
client.await_registered(dd_narrow.server_id)
key_all = _key(client, resources, dd_all.server_id, "allow-call-control")
_ = client.await_tool(key_all, dd_all.server_id, SEARCH_LOGS_TOOL)
key_narrow = _key(client, resources, dd_narrow.server_id, "allow-call-narrow")
result = client.call_tool(
key_narrow,
server_id=dd_narrow.server_id,
name=SEARCH_LOGS_TOOL,
arguments=_search_args("service:litellm"),
)
match result:
case UnknownApiError(status_code=403):
pass
case UnknownApiError(status_code=404):
pass
case _:
pytest.fail(
f"calling a tool outside allowed_tools must be blocked (403 or 404), "
f"got: {result}"
)

10
tests/e2e/models.py
View file

@ -363,7 +363,15 @@ class AnthropicCustomTool(BaseModel):
input_schema: ToolInputSchema
type AnthropicTool = AnthropicToolSearchTool | AnthropicCustomTool
class AnthropicMcpTool(BaseModel):
type: Literal["mcp"] = "mcp"
server_label: str
server_url: str
require_approval: str = "never"
allowed_tools: list[str] | None = None
type AnthropicTool = AnthropicToolSearchTool | AnthropicCustomTool | AnthropicMcpTool
class AnthropicMessagesBody(BaseModel):