From d016edeea3aa2a67c0b30ba56be21004a03f5a3a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 5 Aug 2026 15:34:39 -0700 Subject: [PATCH] test(e2e): add MCP bridge, filtering, auth, and OAuth coverage tests 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. --- tests/e2e/mcp/datadog_mcp.py | 29 +- tests/e2e/mcp/dd_oauth.py | 275 ++++++++++++++++++ tests/e2e/mcp/dd_session_capture.py | 39 +++ tests/e2e/mcp/mcp_client.py | 118 +++++++- tests/e2e/mcp/oauth_chat_client.py | 270 ----------------- tests/e2e/mcp/test_mcp_access_group_e2e.py | 2 +- tests/e2e/mcp/test_mcp_auth_errors_e2e.py | 35 +++ .../e2e/mcp/test_mcp_chat_completions_e2e.py | 99 +++++++ tests/e2e/mcp/test_mcp_datadog_e2e.py | 23 +- tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py | 159 ++++++++++ tests/e2e/mcp/test_mcp_guardrail_e2e.py | 178 ------------ tests/e2e/mcp/test_mcp_key_access_e2e.py | 15 +- tests/e2e/mcp/test_mcp_messages_e2e.py | 114 ++++++++ tests/e2e/mcp/test_mcp_responses_e2e.py | 83 ++++++ tests/e2e/mcp/test_mcp_server_config_e2e.py | 134 +++++++++ tests/e2e/mcp/test_mcp_tool_filtering_e2e.py | 100 +++++++ tests/e2e/models.py | 10 +- 17 files changed, 1192 insertions(+), 491 deletions(-) create mode 100644 tests/e2e/mcp/dd_oauth.py create mode 100644 tests/e2e/mcp/dd_session_capture.py delete mode 100644 tests/e2e/mcp/oauth_chat_client.py create mode 100644 tests/e2e/mcp/test_mcp_auth_errors_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_chat_completions_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py delete mode 100644 tests/e2e/mcp/test_mcp_guardrail_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_messages_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_responses_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_server_config_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_tool_filtering_e2e.py diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..2843c73c91b 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -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) diff --git a/tests/e2e/mcp/dd_oauth.py b/tests/e2e/mcp/dd_oauth.py new file mode 100644 index 00000000000..5e14e2692c2 --- /dev/null +++ b/tests/e2e/mcp/dd_oauth.py @@ -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, + ) diff --git a/tests/e2e/mcp/dd_session_capture.py b/tests/e2e/mcp/dd_session_capture.py new file mode 100644 index 00000000000..b2ed612b010 --- /dev/null +++ b/tests/e2e/mcp/dd_session_capture.py @@ -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) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 33ec557c339..87eb8f4fee5 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -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], diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py deleted file mode 100644 index 4378a5869d0..00000000000 --- a/tests/e2e/mcp/oauth_chat_client.py +++ /dev/null @@ -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) diff --git a/tests/e2e/mcp/test_mcp_access_group_e2e.py b/tests/e2e/mcp/test_mcp_access_group_e2e.py index f72b75fd43d..08527cc2c61 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/tests/e2e/mcp/test_mcp_access_group_e2e.py @@ -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( diff --git a/tests/e2e/mcp/test_mcp_auth_errors_e2e.py b/tests/e2e/mcp/test_mcp_auth_errors_e2e.py new file mode 100644 index 00000000000..cc0f21f07b4 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_auth_errors_e2e.py @@ -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}" + ) diff --git a/tests/e2e/mcp/test_mcp_chat_completions_e2e.py b/tests/e2e/mcp/test_mcp_chat_completions_e2e.py new file mode 100644 index 00000000000..9f26e257318 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_chat_completions_e2e.py @@ -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/"}`` 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" + ) diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index 138f654272d..03ed26f18d5 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -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}" diff --git a/tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py b/tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py new file mode 100644 index 00000000000..0401c12d361 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py @@ -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}" + ) diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py deleted file mode 100644 index 60a349ddc5e..00000000000 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ /dev/null @@ -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}" - ) diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 788a0a3f45c..429f0840e4b 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -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 diff --git a/tests/e2e/mcp/test_mcp_messages_e2e.py b/tests/e2e/mcp/test_mcp_messages_e2e.py new file mode 100644 index 00000000000..da91505b0f9 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_messages_e2e.py @@ -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/"}`` 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" + ) diff --git a/tests/e2e/mcp/test_mcp_responses_e2e.py b/tests/e2e/mcp/test_mcp_responses_e2e.py new file mode 100644 index 00000000000..7991dd4d080 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_responses_e2e.py @@ -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/"}`` 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" + ) diff --git a/tests/e2e/mcp/test_mcp_server_config_e2e.py b/tests/e2e/mcp/test_mcp_server_config_e2e.py new file mode 100644 index 00000000000..fad4a27827a --- /dev/null +++ b/tests/e2e/mcp/test_mcp_server_config_e2e.py @@ -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}" + ) diff --git a/tests/e2e/mcp/test_mcp_tool_filtering_e2e.py b/tests/e2e/mcp/test_mcp_tool_filtering_e2e.py new file mode 100644 index 00000000000..3b08478f566 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_tool_filtering_e2e.py @@ -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}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index f1c0ede0e85..bb938f84253 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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):