scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

test: e2e coverage for MCP toolset enforcement at team, org, and user levels

Browse source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Yucheng Low 2026-08-28 10:26:32 -07:00
parent 987e7623c9
commit ceea12e434
5 changed files with 382 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
tests/e2e/coverage_registry/mcp.yaml
View file

@ -119,3 +119,57 @@
assertions: [succeeds]
source: "server.py:1089"
rationale: Smoke; rarely used; same auth model as tools
- id: mcp.list_tools.api_key.team_toolset_scoped
module: mcp
tier: P0
operation: list_tools
auth_family: api_key
assertions: [toolset_scoped]
source: "auth/user_api_key_auth_mcp.py:2453"
rationale: "A toolset attached to a team's object_permission narrows the team key's tools/list to exactly the tools it names (LIT-5749: stored but inert before the fix)"
fail_before_fix: proven
- id: mcp.call_tool.api_key.team_toolset_denied_outside
module: mcp
tier: P0
operation: call_tool
auth_family: api_key
assertions: [denied_outside_toolset]
source: "mcp_server_manager.py:4925"
rationale: "A team key calling a tool on the granted server but outside the team's toolset is refused; the named tool stays callable (LIT-5749)"
fail_before_fix: proven
- id: mcp.list_tools.api_key.org_toolset_scoped
module: mcp
tier: P1
operation: list_tools
auth_family: api_key
assertions: [toolset_scoped]
source: "auth/user_api_key_auth_mcp.py:1449"
rationale: "A toolset on an org's object_permission caps a member team's key that declares no MCP grants of its own (LIT-5749: the org substitution amplifier)"
fail_before_fix: proven
- id: mcp.call_tool.api_key.org_toolset_denied_outside
module: mcp
tier: P1
operation: call_tool
auth_family: api_key
assertions: [denied_outside_toolset]
source: "mcp_server_manager.py:4925"
rationale: "An org-inherited key calling outside the org's toolset is refused while the named tool stays callable (LIT-5749)"
fail_before_fix: proven
- id: mcp.list_tools.api_key.user_toolset_scoped
module: mcp
tier: P1
operation: list_tools
auth_family: api_key
assertions: [toolset_scoped]
source: "auth/user_api_key_auth_mcp.py:2859"
rationale: "A toolset on an internal user's row ceils the servers the user's own key grants; user-level toolsets narrow, never grant (LIT-5749)"
fail_before_fix: proven
- id: mcp.call_tool.api_key.user_toolset_denied_outside
module: mcp
tier: P1
operation: call_tool
auth_family: api_key
assertions: [denied_outside_toolset]
source: "mcp_server_manager.py:4925"
rationale: "A key whose user row holds a toolset is refused calling outside it even though the key itself grants the whole server (LIT-5749)"
fail_before_fix: proven

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

@ -3,6 +3,7 @@
from __future__ import annotations
import os
from collections.abc import Sequence
from e2e_config import datadog_mcp_url, unique_marker
from lifecycle import ResourceManager
@ -35,7 +36,11 @@ def register_datadog_mcp(
resources: ResourceManager,
*,
mcp_access_groups: list[str] | None = None,
allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,),
) -> str:
"""Register the Datadog remote MCP server. `allowed_tools` defaults to the
search-logs slice; pass None to expose the full core toolset (needed when a
test must prove narrowing, so a second tool has to exist to be denied)."""
assert_dd_mcp_creds()
name = f"e2e_dd_mcp_{unique_marker()}"
server_id = client.register_server(
@ -47,7 +52,7 @@ def register_datadog_mcp(
"DD-API-KEY": _dd_api_key(),
"DD-APPLICATION-KEY": _dd_app_key(),
},
allowed_tools=[SEARCH_LOGS_TOOL],
allowed_tools=list(allowed_tools) if allowed_tools is not None else None,
mcp_access_groups=mcp_access_groups,
)
resources.defer(lambda: client.delete_server(server_id))

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

@ -20,7 +20,19 @@ from pydantic import BaseModel, ConfigDict, Field, RootModel
from e2e_config import settle_propagation
from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap
from models import KeyGenerateBody, ObjectPermission
from models import (
KeyGenerateBody,
ObjectPermission,
OrgDeleteBody,
OrgNewBody,
OrgNewResponse,
TeamDeleteBody,
TeamNewBody,
TeamNewResponse,
UserDeleteBody,
UserNewBody,
UserNewResponse,
)
from proxy_client import ProxyClient
McpToolArg = str | int | float | bool | list[str] | dict[str, str]
@ -46,6 +58,20 @@ class McpServerNewResponse(BaseModel):
server_id: str
class McpToolsetToolSpec(BaseModel):
server_id: str
tool_name: str
class McpToolsetNewBody(BaseModel):
toolset_name: str
tools: list[McpToolsetToolSpec]
class McpToolsetNewResponse(BaseModel):
toolset_id: str
class McpServerRow(BaseModel):
server_id: str
alias: str | None = None
@ -74,9 +100,7 @@ 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:
@ -225,6 +249,7 @@ class McpClient:
mcp_servers: list[str] | None,
mcp_access_groups: list[str] | None = None,
models: list[str] | None = None,
team_id: str | None = None,
) -> str:
object_permission = (
ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups)
@ -235,10 +260,114 @@ class McpClient:
KeyGenerateBody(
models=models if models is not None else [],
user_id=user_id,
team_id=team_id,
object_permission=object_permission,
)
)
def create_toolset(self, *, name: str, server_id: str, tool_names: list[str]) -> str:
"""Create a named toolset over `server_id` (admin), returning its id."""
return unwrap(
self.proxy.transport.post(
"/v1/mcp/toolset",
headers=self.proxy.transport.master,
json=McpToolsetNewBody(
toolset_name=name,
tools=[McpToolsetToolSpec(server_id=server_id, tool_name=tool) for tool in tool_names],
),
response_type=McpToolsetNewResponse,
)
).toolset_id
def delete_toolset(self, toolset_id: str) -> None:
_ = self.proxy.transport.delete(
f"/v1/mcp/toolset/{toolset_id}",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
def create_team(
self,
*,
alias: str,
object_permission: ObjectPermission | None = None,
organization_id: str | None = None,
) -> str:
return unwrap(
self.proxy.transport.post(
"/team/new",
headers=self.proxy.transport.master,
json=TeamNewBody(
team_alias=alias,
organization_id=organization_id,
object_permission=object_permission,
),
response_type=TeamNewResponse,
)
).team_id
def delete_team(self, team_id: str) -> None:
_ = self.proxy.transport.post(
"/team/delete",
headers=self.proxy.transport.master,
json=TeamDeleteBody(team_ids=[team_id]),
response_type=NoBody,
)
def create_org(
self,
*,
alias: str,
object_permission: ObjectPermission | None = None,
) -> str:
return unwrap(
self.proxy.transport.post(
"/organization/new",
headers=self.proxy.transport.master,
json=OrgNewBody(
organization_alias=alias,
object_permission=object_permission,
),
response_type=OrgNewResponse,
)
).organization_id
def delete_org(self, organization_id: str) -> None:
_ = self.proxy.transport.delete(
"/organization/delete",
headers=self.proxy.transport.master,
json=OrgDeleteBody(organization_ids=[organization_id]),
response_type=NoBody,
)
def create_user(
self,
*,
user_email: str,
object_permission: ObjectPermission | None = None,
) -> str:
return unwrap(
self.proxy.transport.post(
"/user/new",
headers=self.proxy.transport.master,
json=UserNewBody(
user_email=user_email,
user_role="internal_user",
object_permission=object_permission,
),
response_type=UserNewResponse,
)
).user_id
def delete_user(self, user_id: str) -> None:
_ = self.proxy.transport.post(
"/user/delete",
headers=self.proxy.transport.master,
json=UserDeleteBody(user_ids=[user_id]),
response_type=NoBody,
)
def list_tools(self, key: str) -> Result[McpToolsListResponse]:
return self.proxy.transport.get(
"/mcp-rest/tools/list",
@ -316,13 +445,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 +494,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 +527,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

181
tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py Normal file
View file

@ -0,0 +1,181 @@
"""Live e2e: an MCP toolset attached to a team, org, or internal user narrows a
real MCP server to exactly the tools the toolset names (LIT-5749).
An admin registers the Datadog remote MCP server with its full core toolset (no
`allowed_tools` cap, so more than one tool exists to be denied), creates a
toolset naming only `search_datadog_logs` on it, and attaches that toolset at
one principal level per test. A control key granted the whole server first
proves the upstream is alive and serves a second tool, so a later denial is an
authorization decision rather than a dead server. The scoped key must then see
exactly the toolset's one tool on `tools/list`, be refused (403) calling any
other tool on the same server, and still successfully call the granted tool.
Levels covered, one per test: a team holding server + toolset, a team that
inherits its grant from an org holding server + toolset, and an internal user
whose toolset ceils the servers their own key grants.
"""
from __future__ import annotations
import pytest
from datadog_mcp import SEARCH_LOGS_TOOL, 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
from models import ObjectPermission
pytestmark = pytest.mark.e2e
def _open_server_and_toolset(client: McpClient, resources: ResourceManager) -> tuple[str, str]:
"""Register the DD server uncapped and a toolset naming only the search-logs
tool on it, returning (server_id, toolset_id)."""
server_id = register_datadog_mcp(client, resources, allowed_tools=None)
client.await_registered(server_id)
toolset_id = client.create_toolset(
name=f"e2e-search-only-{unique_marker()}",
server_id=server_id,
tool_names=[SEARCH_LOGS_TOOL],
)
resources.defer(lambda: client.delete_toolset(toolset_id))
return server_id, toolset_id
def _control_tools(client: McpClient, resources: ResourceManager, server_id: str) -> tuple[str, str]:
"""A fully-granted control key's view of the server: the fully-qualified
search-logs tool name and one tool outside the toolset. Proves the upstream
is alive and actually serves something a toolset can deny."""
control_key = client.generate_key(user_id=f"e2e-ts-control-{unique_marker()}", mcp_servers=[server_id])
resources.defer(lambda: client.proxy.delete_key(control_key))
granted_tool = client.await_tool(control_key, server_id, SEARCH_LOGS_TOOL)
all_tools = unwrap(client.list_tools(control_key)).tool_names_for_server(server_id)
outside = sorted(all_tools - {granted_tool})
assert outside, (
f"the uncapped Datadog core toolset served only {all_tools}; a toolset "
f"cannot be proven to narrow a one-tool server"
)
return granted_tool, outside[0]
def _assert_toolset_ceiling(
client: McpClient,
scoped_key: str,
*,
server_id: str,
granted_tool: str,
outside_tool: str,
) -> None:
"""The scoped key sees exactly the toolset's tool, is 403-refused on a tool
outside it, and can still execute the granted one."""
_ = client.await_tool(scoped_key, server_id, SEARCH_LOGS_TOOL)
listed = unwrap(client.list_tools(scoped_key)).tool_names_for_server(server_id)
assert listed == frozenset({granted_tool}), (
f"toolset-scoped key must list exactly {granted_tool!r}; the toolset ceiling leaked: {sorted(listed)}"
)
denied = client.await_call_tool_denied(scoped_key, server_id=server_id, name=outside_tool, arguments={})
assert denied.status_code == 403
result = client.await_call_tool(
scoped_key,
server_id=server_id,
name=granted_tool,
arguments={"query": f"service:e2e-toolset-{unique_marker()}", "from": DD_SEARCH_FROM},
)
assert result.is_error is not True, f"the toolset-granted tool must stay callable; upstream said: {result.all_text}"
class TestMcpToolsetEnforcementPerLevel:
@pytest.mark.covers("mcp.list_tools.api_key.team_toolset_scoped")
@pytest.mark.covers("mcp.call_tool.api_key.team_toolset_denied_outside")
def test_team_toolset_narrows_team_key(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
server_id, toolset_id = _open_server_and_toolset(client, resources)
granted_tool, outside_tool = _control_tools(client, resources, server_id)
team_id = client.create_team(
alias=f"e2e-ts-team-{unique_marker()}",
object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]),
)
resources.defer(lambda: client.delete_team(team_id))
team_key = client.generate_key(
user_id=f"e2e-ts-team-user-{unique_marker()}",
mcp_servers=None,
team_id=team_id,
)
resources.defer(lambda: client.proxy.delete_key(team_key))
_assert_toolset_ceiling(
client,
team_key,
server_id=server_id,
granted_tool=granted_tool,
outside_tool=outside_tool,
)
@pytest.mark.covers("mcp.list_tools.api_key.org_toolset_scoped")
@pytest.mark.covers("mcp.call_tool.api_key.org_toolset_denied_outside")
def test_org_toolset_caps_inherited_team_key(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
server_id, toolset_id = _open_server_and_toolset(client, resources)
granted_tool, outside_tool = _control_tools(client, resources, server_id)
org_id = client.create_org(
alias=f"e2e-ts-org-{unique_marker()}",
object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]),
)
resources.defer(lambda: client.delete_org(org_id))
# The team declares no MCP grants of its own; whatever its key can reach
# comes from the org, so the org's toolset must cap it.
team_id = client.create_team(alias=f"e2e-ts-org-team-{unique_marker()}", organization_id=org_id)
resources.defer(lambda: client.delete_team(team_id))
team_key = client.generate_key(
user_id=f"e2e-ts-org-user-{unique_marker()}",
mcp_servers=None,
team_id=team_id,
)
resources.defer(lambda: client.proxy.delete_key(team_key))
_assert_toolset_ceiling(
client,
team_key,
server_id=server_id,
granted_tool=granted_tool,
outside_tool=outside_tool,
)
@pytest.mark.covers("mcp.list_tools.api_key.user_toolset_scoped")
@pytest.mark.covers("mcp.call_tool.api_key.user_toolset_denied_outside")
def test_user_toolset_ceils_own_key_grant(
self,
client: McpClient,
resources: ResourceManager,
) -> None:
server_id, toolset_id = _open_server_and_toolset(client, resources)
granted_tool, outside_tool = _control_tools(client, resources, server_id)
# The user's row holds only the toolset; their key grants the whole
# server. The user-level toolset is a ceiling over the key's grant.
user_id = client.create_user(
user_email=f"e2e-ts-user-{unique_marker()}@example.com",
object_permission=ObjectPermission(mcp_toolsets=[toolset_id]),
)
resources.defer(lambda: client.delete_user(user_id))
user_key = client.generate_key(user_id=user_id, mcp_servers=[server_id])
resources.defer(lambda: client.proxy.delete_key(user_key))
_assert_toolset_ceiling(
client,
user_key,
server_id=server_id,
granted_tool=granted_tool,
outside_tool=outside_tool,
)

4
tests/e2e/models.py
View file

@ -52,6 +52,7 @@ class KeyMetadata(BaseModel):
class ObjectPermission(BaseModel):
mcp_servers: list[str] | None = None
mcp_access_groups: list[str] | None = None
mcp_toolsets: list[str] | None = None
class KeyGenerateBody(BaseModel):
@ -922,6 +923,7 @@ class TeamNewBody(BaseModel):
team_id: str | None = None
organization_id: str | None = None
metadata: TeamMetadata | None = None
object_permission: ObjectPermission | None = None
class TeamNewResponse(BaseModel):
@ -979,6 +981,7 @@ class UserNewBody(BaseModel):
user_email: str
user_role: UserRole
user_id: str | None = None
object_permission: ObjectPermission | None = None
class UserNewResponse(BaseModel):
@ -1029,6 +1032,7 @@ class UserListResponse(BaseModel):
class OrgNewBody(BaseModel):
organization_alias: str
models: list[str] = []
object_permission: ObjectPermission | None = None
class OrgNewResponse(BaseModel):