diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..34f90d2218d 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -61,11 +61,25 @@ class McpToolMcpInfo(BaseModel): alias: str | None = None +class McpToolInputSchema(BaseModel): + properties: dict[str, object] = {} + + class McpToolEntry(BaseModel): + model_config = ConfigDict(populate_by_name=True) + name: str description: str | None = None + input_schema: McpToolInputSchema = Field(alias="inputSchema") mcp_info: McpToolMcpInfo | None = None + def assert_arguments_are_documented(self, arguments: McpToolArguments) -> None: + undocumented = frozenset(arguments).difference(self.input_schema.properties) + assert not undocumented, ( + f"tool {self.name!r} arguments are absent from its advertised input schema: " + f"{sorted(undocumented)}; documented arguments: {sorted(self.input_schema.properties)}" + ) + class McpToolsListResponse(BaseModel): tools: list[McpToolEntry] = [] @@ -74,18 +88,20 @@ class McpToolsListResponse(BaseModel): def tool_names_for_server(self, server_id: str) -> frozenset[str]: return frozenset( - tool.name - for tool in self.tools - if tool.mcp_info is not None and tool.mcp_info.server_id == server_id + tool.name for tool in self.tools if tool.mcp_info is not None and tool.mcp_info.server_id == server_id ) def tool_name_containing(self, server_id: str, needle: str) -> str | None: + tool = self.tool_containing(server_id, needle) + return tool.name if tool is not None else None + + def tool_containing(self, server_id: str, needle: str) -> McpToolEntry | None: needle_l = needle.lower() for tool in self.tools: if tool.mcp_info is None or tool.mcp_info.server_id != server_id: continue if needle_l in tool.name.lower() or tool.name.lower().endswith(needle_l): - return tool.name + return tool return None @@ -248,8 +264,11 @@ class McpClient: ) def await_tool(self, key: str, server_id: str, needle: str) -> str: + return self.await_tool_entry(key, server_id, needle).name + + def await_tool_entry(self, key: str, server_id: str, needle: str) -> McpToolEntry: """Poll tools/list until `server_id` serves a tool matching `needle`, and - return its fully-qualified name. Fails at poll_timeout. + return its entry. Fails at poll_timeout. /v1/mcp/server returns as soon as the DB row is written, but the gateway runs the initialize + tools/list handshake against the upstream lazily on @@ -261,9 +280,9 @@ class McpClient: while True: result = self.list_tools(key) if isinstance(result, Success): - tool_name = result.data.tool_name_containing(server_id, needle) - if tool_name is not None: - return tool_name + tool = result.data.tool_containing(server_id, needle) + if tool is not None: + return tool if time.monotonic() >= deadline: raise AssertionError( f"server {server_id} never served a tool matching {needle!r} within " @@ -316,13 +335,10 @@ class McpClient: if isinstance(last, UnknownApiError) and last.status_code == 403: return last if not _is_mcp_not_synced(last, tool_name=name): - raise AssertionError( - f"ungranted key's tools/call was not 403 access_denied: {last}" - ) + raise AssertionError(f"ungranted key's tools/call was not 403 access_denied: {last}") if time.monotonic() >= deadline: raise AssertionError( - f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; " - f"last result: {last}" + f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; last result: {last}" ) time.sleep(self.proxy.poll_interval) @@ -368,9 +384,7 @@ class McpClient: return self.proxy.transport.post( "/mcp-rest/tools/call", headers=ApiKeyHeaders(x_litellm_api_key=key), - json=McpCallToolBody( - name=name, arguments=dict(arguments), server_id=server_id - ), + json=McpCallToolBody(name=name, arguments=dict(arguments), server_id=server_id), response_type=McpCallToolResponse, ) @@ -403,9 +417,7 @@ def _is_mcp_not_synced( # Gateway: "Tool search_datadog_logs not found" (optionally inside a longer message) if tool_name is not None: - return ( - re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None - ) + return re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None return re.search(r"\btool\s+\S+\s+not found\b", body_l) is not None diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index 138f654272d..bb351454648 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -34,8 +34,7 @@ def _assert_datadog_logger_active(proxy: ProxyClient) -> None: f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" ) assert DD_LOGGER_NAME in result.body, ( - f"the proxy must report the {DD_LOGGER_NAME} callback active " - f"(callbacks + DD_* env); got: {result.body[:400]}" + f"the proxy must report the {DD_LOGGER_NAME} callback active (callbacks + DD_* env); got: {result.body[:400]}" ) @@ -49,16 +48,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, @@ -88,24 +77,22 @@ 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 = client.await_tool_entry(key, server_id, SEARCH_LOGS_TOOL) + arguments = { + "query": marker, + "from": DD_SEARCH_FROM, + "to": "now", + "max_tokens": 5000, + } + tool.assert_arguments_are_documented(arguments) call = client.await_call_tool( key, server_id=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" - }, - }, + name=tool.name, + arguments=arguments, ) assert call.is_error is not True, f"search_datadog_logs errored: {call}" body = call.all_text assert marker in body, ( - f"search_datadog_logs response must include the seeded marker {marker!r}; " - f"got: {body[:800]!r}" + f"search_datadog_logs response must include the seeded marker {marker!r}; got: {body[:800]!r}" ) diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py index 60a349ddc5e..683530db0cb 100644 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ b/tests/e2e/mcp/test_mcp_guardrail_e2e.py @@ -78,16 +78,6 @@ def _search_on_synced_pod( 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"], @@ -99,9 +89,7 @@ class TestMcpToolCallGuardrail: 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_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)) @@ -110,7 +98,7 @@ class TestMcpToolCallGuardrail: 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) + tool = client.await_tool_entry(key, server_id, SEARCH_LOGS_TOOL) def search(query: str) -> Result[McpCallToolResponse]: arguments: McpToolArguments = { @@ -118,9 +106,9 @@ class TestMcpToolCallGuardrail: "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) + tool.assert_arguments_are_documented(arguments) + 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 @@ -173,6 +161,4 @@ class TestMcpToolCallGuardrail: 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}" - ) + 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..a11d9760c5c 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -47,20 +47,9 @@ class TestMcpKeyWithoutAccessIsDenied: denied_tools = unwrap(client.list_tools(denied_key)).tool_names_for_server(server_id) assert denied_tools == frozenset(), ( - f"ungranted key saw the server's tools; tools/list leaked across the permission " - f"boundary: {denied_tools}" + f"ungranted key saw the server's tools; tools/list leaked across the permission 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, @@ -73,21 +62,19 @@ class TestMcpKeyWithoutAccessIsDenied: permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) - tool_name = client.await_tool(permitted_key, server_id, SEARCH_LOGS_TOOL) + tool = client.await_tool_entry(permitted_key, server_id, SEARCH_LOGS_TOOL) search_args = { "query": "service:litellm", "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000, - "telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"}, } + tool.assert_arguments_are_documented(search_args) permitted_call = client.await_call_tool( - permitted_key, server_id=server_id, name=tool_name, arguments=search_args + permitted_key, server_id=server_id, name=tool.name, arguments=search_args ) assert permitted_call.is_error is not True, f"granted key's tool call errored: {permitted_call}" - denied = client.await_call_tool_denied( - denied_key, server_id=server_id, name=tool_name, arguments=search_args - ) + denied = client.await_call_tool_denied(denied_key, server_id=server_id, name=tool.name, arguments=search_args) assert "access_denied" in denied.body, f"403 was not an MCP access denial: {denied.body}"