scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge 8beaf7d058 into ec94a1f82a

Browse source
This commit is contained in:
mubashir1osmani 2026-08-27 19:08:58 -05:00 • committed by GitHub
commit 59dd2baee2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 53 additions and 81 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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

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

@ -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}"
)

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

@ -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}")

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

@ -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}"