scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

fix(mcp): honor an explicit null on toolset update, cover MCP lifecycle e2e

Browse source 
PUT /v1/mcp/toolset dumped its payload with exclude_none, so a field sent as
null looked exactly like one the caller left out and the stored value
survived. An admin could not clear a toolset's description: the save reported
success and the old text came straight back. It now dumps with exclude_unset,
so absent keeps and null clears, which is what PUT /v1/mcp/server already did.
A null tools list clears the selection to empty, and a null toolset_name is
ignored because a toolset always has a name.

Adds create, read, partial-update, clear and delete e2e coverage for MCP
servers and toolsets, with every read-back polled on every replica so an edit
that lands on one replica and not another fails the test, plus an enforcement
test proving a key granted a toolset lists exactly that toolset's tools
against the real Datadog upstream.
This commit is contained in:
Yuneng Jiang 2026-09-06 07:23:35 +00:00
parent 0318b4acdc
commit 856abb30c6
No known key found for this signature in database
16 changed files with 997 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/_experimental/mcp_server/toolset_db.py
View file

@ -132,9 +132,14 @@ async def update_mcp_toolset(
data: UpdateMCPToolsetRequest,
touched_by: str,
) -> MCPToolset | None:
data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"})
"""A partial update: a field the caller left out keeps its stored value and a
field sent as null is cleared, where a null ``tools`` is an empty list and a null
``toolset_name`` is ignored because a toolset always has a name."""
data_dict: Final = data.model_dump(exclude_unset=True, exclude={"toolset_id"})
if "tools" in data_dict:
data_dict["tools"] = json.dumps(data_dict["tools"])
data_dict["tools"] = json.dumps(data_dict["tools"] or ())
if data_dict.get("toolset_name", "") is None:
_ = data_dict.pop("toolset_name")
data_dict["updated_by"] = touched_by
try:
row: Final = await _toolset_table(prisma_client).update(

3
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -2673,6 +2673,8 @@ if MCP_AVAILABLE:
"""
Updates the MCP Server in the db.
Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared.
Parameters:
- payload: UpdateMCPServerRequest - Required. The updated mcp server data.
```
@ -3098,6 +3100,7 @@ if MCP_AVAILABLE:
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
litellm_changed_by: str | None = Header(None),
):
"""Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
raise HTTPException(

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

@ -119,3 +119,11 @@
assertions: [succeeds]
source: "server.py:1089"
rationale: Smoke; rarely used; same auth model as tools
- id: mcp.list_tools.api_key.toolset_scoped
module: mcp
tier: P0
operation: list_tools
auth_family: api_key
assertions: [toolset_scoped]
source: "user_api_key_auth_mcp.py:2137"
rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves"

10
tests/e2e/coverage_registry/mgmt.yaml
View file

@ -74,3 +74,13 @@
- {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"}
- {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"}
- {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven}
- {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"}
- {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"}
- {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"}
- {id: mgmt.mcp_server.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:2665", rationale: "An explicit null clears the stored field (absent keeps, null clears)"}
- {id: mgmt.mcp_server.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:2139", rationale: "A deleted server is gone by id and from the list on every replica"}
- {id: mgmt.mcp_toolset.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3009", rationale: "Toolset tools read back under the exact server_id and tool_name written; a toolset stored under one name and read under another granted nothing"}
- {id: mgmt.mcp_toolset.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:3098", rationale: "Editing the description leaves the tools and name intact"}
- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"}
- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"}
- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"}

64
tests/e2e/e2e_http.py
View file

@ -49,6 +49,11 @@ class AnthropicHeaders(AuthHeaders):
anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version")
class PartialBody(BaseModel):
"""A body for a partial-update route (absent = keep, null = clear): a field left
unset is omitted from the wire, and a field set to None is sent as JSON null."""
class NoBody(BaseModel):
"""Empty body/query for routes that take none."""
@ -252,6 +257,13 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None:
f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}"
)
def wire_body(json: BaseModel) -> dict[str, object]:
if isinstance(json, PartialBody):
return json.model_dump(by_alias=True, exclude_unset=True)
return json.model_dump(by_alias=True, exclude_none=True)
def _headers(headers: BaseModel) -> dict[str, str]:
dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True)
return {key: str(value) for key, value in dumped.items()}
@ -307,9 +319,26 @@ def request_with_retry[T: RetryableResponse](
return issue()
def _classify[R: BaseModel](
resp: requests.Response, response_type: type[R]
) -> Result[R]:
class ClassifiableResponse(Protocol):
"""What classifying an outcome reads off a response. requests.Response satisfies
it, and so does a fake, so the classification rules are testable on their own."""
@property
def status_code(self) -> int: ...
@property
def ok(self) -> bool: ...
@property
def text(self) -> str: ...
@property
def content(self) -> bytes: ...
def json(self) -> object: ...
def classify[R: BaseModel](resp: ClassifiableResponse, response_type: type[R]) -> Result[R]:
if resp.status_code == 401:
return UnauthorizedError(body=resp.text)
if resp.status_code == 429:
@ -317,7 +346,8 @@ def _classify[R: BaseModel](
if not resp.ok:
return UnknownApiError(status_code=resp.status_code, body=resp.text)
try:
return Success(status_code=resp.status_code, data=response_type.model_validate(resp.json()))
payload: Final[object] = resp.json() if resp.content else {}
return Success(status_code=resp.status_code, data=response_type.model_validate(payload))
except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value
return ValidationError(message=str(exc))
@ -335,13 +365,13 @@ def post[R: BaseModel](
lambda: requests.post(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def get[R: BaseModel](
@ -363,7 +393,7 @@ def get[R: BaseModel](
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def get_external[R: BaseModel](
@ -383,7 +413,7 @@ def get_external[R: BaseModel](
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def delete[R: BaseModel](
@ -400,14 +430,14 @@ def delete[R: BaseModel](
lambda: requests.delete(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
params=_params(params),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def patch[R: BaseModel](
@ -423,13 +453,13 @@ def patch[R: BaseModel](
lambda: requests.patch(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def put[R: BaseModel](
@ -445,13 +475,13 @@ def put[R: BaseModel](
lambda: requests.put(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def probe(
@ -555,7 +585,7 @@ def send(
str(url),
headers=_headers(headers),
params=_params(params),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
stream=stream,
timeout=timeout,
)
@ -605,7 +635,7 @@ def upload[R: BaseModel](
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def stream_binary(
@ -623,7 +653,7 @@ def stream_binary(
resp = requests.post(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
stream=True,
timeout=timeout,
)

35
tests/e2e/management/management_client.py
View file

@ -43,6 +43,9 @@ from models import (
KeyResetSpendBody,
KeyResetSpendResponse,
KeyUpdateBody,
McpServerCreateBody,
McpServerRow,
McpServerUpdateBody,
ModelDeleteBody,
OrgDeleteBody,
OrgInfoParams,
@ -537,6 +540,38 @@ class ManagementClient:
).root
)
def create_mcp_server(self, body: McpServerCreateBody) -> McpServerRow:
return unwrap(
self.proxy.transport.post(
"/v1/mcp/server",
headers=self.proxy.transport.master,
json=body,
response_type=McpServerRow,
)
)
def update_mcp_server(self, body: McpServerUpdateBody) -> McpServerRow:
"""PUT /v1/mcp/server, the call behind the dashboard's Save Changes: a partial
update where a field left unset keeps its stored value and None clears it."""
return unwrap(
self.proxy.transport.put(
"/v1/mcp/server",
headers=self.proxy.transport.master,
json=body,
response_type=McpServerRow,
)
)
def delete_mcp_server(self, server_id: str) -> Result[NoBody]:
"""DELETE /v1/mcp/server/{server_id}. Returns the outcome so the act phase can
unwrap it while a deferred teardown can ignore an already-deleted server."""
return self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
return self.proxy.transport.send(
"/chat/completions",

294
tests/e2e/management/test_mcp_lifecycle_e2e.py Normal file
View file

@ -0,0 +1,294 @@
"""Live e2e: the MCP server and toolset management routes' lifecycle contract.
Two customer defects sit on these routes, and each step here is the read-back that
would have caught one of them: a dashboard edit that took several saves to stick
because the read landed on a replica the write had not reached, and a toolset whose
tools were stored under one name and read back under another, so it granted
nothing. Every read-back therefore polls every replica that serves the route
(ProxyClient.read_back_everywhere) and asserts the exact values written, and both
update routes are held to the same partial-update contract: a field left out of the
payload keeps its stored value, a field sent as null is cleared. The server URL is
unreachable on purpose; only persistence is under test, never a tool call.
"""
from __future__ import annotations
from collections.abc import Callable, Mapping
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from management_client import ManagementClient
from models import (
McpInfo,
McpServerCreateBody,
McpServerListResponse,
McpServerRow,
McpServerUpdateBody,
ToolsetCreateBody,
ToolsetListResponse,
ToolsetRow,
ToolsetTool,
ToolsetUpdateBody,
)
pytestmark = pytest.mark.e2e
UNREACHABLE_URL: Final = "https://e2e-fake-mcp.test.local/mcp"
def _create_server(client: ManagementClient, resources: ResourceManager) -> tuple[McpServerCreateBody, str]:
name: Final = f"e2e_mcp_lifecycle_{unique_marker()}"
body: Final = McpServerCreateBody(
server_name=name,
alias=name,
url=UNREACHABLE_URL,
transport="http",
description="e2e lifecycle server",
mcp_info=McpInfo(
server_name=f"{name} (display)",
description="shown on the MCP page",
logo_url="https://e2e.test.local/logo.png",
),
)
server_id: Final = client.create_mcp_server(body).server_id
resources.defer(lambda: client.delete_mcp_server(server_id))
return body, server_id
def _assert_server_matches(row: McpServerRow, written: McpServerCreateBody, *, where: str) -> None:
stored: Final = (row.server_name, row.alias, row.url, row.transport, row.description, row.mcp_info)
expected: Final = (
written.server_name,
written.alias,
written.url,
written.transport,
written.description,
written.mcp_info,
)
assert stored == expected, f"{where}: stored {stored}, expected {expected}"
def _server_everywhere(
client: ManagementClient, server_id: str, *, settled: Callable[[McpServerRow], bool]
) -> Mapping[str, McpServerRow]:
return client.proxy.read_back_everywhere(f"/v1/mcp/server/{server_id}", McpServerRow, settled=settled)
def _listed_server_everywhere(client: ManagementClient, server_id: str) -> Mapping[str, McpServerRow]:
listings: Final = client.proxy.read_back_everywhere(
"/v1/mcp/server",
McpServerListResponse,
settled=lambda rows: any(row.server_id == server_id for row in rows.root),
)
return {replica: next(row for row in rows.root if row.server_id == server_id) for replica, rows in listings.items()}
class TestMcpServerLifecycle:
@pytest.mark.covers("mgmt.mcp_server.new.persists")
def test_create_persists_every_field_on_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
by_id: Final = _server_everywhere(client, server_id, settled=lambda row: row.server_id == server_id)
for replica, row in by_id.items():
_assert_server_matches(row, body, where=f"GET /v1/mcp/server/{server_id} on {replica}")
@pytest.mark.skip(
reason=(
"product gap: GET /v1/mcp/server builds each row from the in-memory registry, whose "
"_build_mcp_server_table sets description from mcp_info['description'], so the list "
"reports the mcp_info description while GET /v1/mcp/server/{server_id} reports the "
"stored description column. A server created with both set to different text reads "
"back with two different descriptions depending on the route"
)
)
@pytest.mark.covers("mgmt.mcp_server.list.persists")
def test_created_server_is_listed_with_every_field(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
for replica, row in _listed_server_everywhere(client, server_id).items():
_assert_server_matches(row, body, where=f"GET /v1/mcp/server on {replica}")
@pytest.mark.covers("mgmt.mcp_server.update.preserves_unrelated_fields")
def test_updating_only_the_alias_keeps_every_other_field_on_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
renamed: Final = f"{body.alias}_renamed"
_ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, alias=renamed))
after_one_put: Final = _server_everywhere(client, server_id, settled=lambda row: row.alias == renamed)
for replica, row in after_one_put.items():
_assert_server_matches(
row,
body.model_copy(update={"alias": renamed}),
where=f"GET /v1/mcp/server/{server_id} on {replica} after one PUT of alias",
)
@pytest.mark.covers("mgmt.mcp_server.update.clear_persists")
def test_clearing_the_description_with_null_reads_back_null(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
_ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, description=None))
cleared: Final = _server_everywhere(client, server_id, settled=lambda row: row.description is None)
for replica, row in cleared.items():
_assert_server_matches(
row,
body.model_copy(update={"description": None}),
where=f"GET /v1/mcp/server/{server_id} on {replica} after PUT description=null",
)
@pytest.mark.covers("mgmt.mcp_server.delete.persists")
def test_delete_removes_the_server_from_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
_ = unwrap(client.delete_mcp_server(server_id))
gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/server/{server_id}")
assert set(gone.values()) == {404}, f"a deleted server must 404 on every replica; got {dict(gone)}"
listings: Final = client.proxy.read_back_everywhere(
"/v1/mcp/server",
McpServerListResponse,
settled=lambda rows: all(row.server_id != server_id for row in rows.root),
)
for replica, rows in listings.items():
assert all(row.server_id != server_id for row in rows.root), (
f"GET /v1/mcp/server on {replica} still lists the deleted server {server_id}"
)
def _create_toolset(
client: ManagementClient, resources: ResourceManager, server_id: str
) -> tuple[ToolsetCreateBody, str]:
body: Final = ToolsetCreateBody(
toolset_name=f"e2e_toolset_{unique_marker()}",
description="e2e lifecycle toolset",
tools=[
ToolsetTool(server_id=server_id, tool_name="search_datadog_logs"),
ToolsetTool(server_id=server_id, tool_name="get_datadog_metric"),
],
)
toolset_id: Final = client.proxy.create_toolset(body).toolset_id
resources.defer(lambda: client.proxy.delete_toolset(toolset_id))
return body, toolset_id
def _assert_toolset_matches(row: ToolsetRow, written: ToolsetCreateBody, *, where: str) -> None:
stored: Final = (row.toolset_name, row.description, row.tools)
expected: Final = (written.toolset_name, written.description, written.tools)
assert stored == expected, f"{where}: stored {stored}, expected {expected}"
def _toolset_everywhere(
client: ManagementClient, toolset_id: str, *, settled: Callable[[ToolsetRow], bool]
) -> Mapping[str, ToolsetRow]:
return client.proxy.read_back_everywhere(f"/v1/mcp/toolset/{toolset_id}", ToolsetRow, settled=settled)
class TestMcpToolsetLifecycle:
@pytest.mark.covers("mgmt.mcp_toolset.new.persists")
def test_create_persists_both_tools_under_the_exact_names_written(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
by_id: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.toolset_id == toolset_id)
for replica, row in by_id.items():
_assert_toolset_matches(row, body, where=f"GET /v1/mcp/toolset/{toolset_id} on {replica}")
listings: Final = client.proxy.read_back_everywhere(
"/v1/mcp/toolset",
ToolsetListResponse,
settled=lambda rows: any(row.toolset_id == toolset_id for row in rows.root),
)
for replica, rows in listings.items():
_assert_toolset_matches(
next(row for row in rows.root if row.toolset_id == toolset_id),
body,
where=f"GET /v1/mcp/toolset on {replica}",
)
@pytest.mark.covers("mgmt.mcp_toolset.update.preserves_unrelated_fields")
def test_updating_only_the_description_keeps_the_tools_and_name(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
_ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description="edited"))
edited: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description == "edited")
for replica, row in edited.items():
_assert_toolset_matches(
row,
body.model_copy(update={"description": "edited"}),
where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of description",
)
@pytest.mark.covers("mgmt.mcp_toolset.update.persists")
def test_updating_the_tools_to_one_entry_reads_back_exactly_that_entry(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
kept: Final = body.tools[:1]
_ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, tools=kept))
narrowed: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.tools == kept)
for replica, row in narrowed.items():
_assert_toolset_matches(
row,
body.model_copy(update={"tools": kept}),
where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of one tool",
)
@pytest.mark.covers("mgmt.mcp_toolset.update.clear_persists")
def test_clearing_the_description_with_null_reads_back_null(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
_ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description=None))
cleared: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description is None)
for replica, row in cleared.items():
_assert_toolset_matches(
row,
body.model_copy(update={"description": None}),
where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT description=null",
)
@pytest.mark.covers("mgmt.mcp_toolset.delete.persists")
def test_delete_removes_the_toolset_from_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
_, toolset_id = _create_toolset(client, resources, server_id)
_ = unwrap(client.proxy.delete_toolset(toolset_id))
gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/toolset/{toolset_id}")
assert set(gone.values()) == {404}, f"a deleted toolset must 404 on every replica; got {dict(gone)}"
listings: Final = client.proxy.read_back_everywhere(
"/v1/mcp/toolset",
ToolsetListResponse,
settled=lambda rows: all(row.toolset_id != toolset_id for row in rows.root),
)
for replica, rows in listings.items():
assert all(row.toolset_id != toolset_id for row in rows.root), (
f"GET /v1/mcp/toolset on {replica} still lists the deleted toolset {toolset_id}"
)

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 core Datadog toolset with its credentials from the env. By default
the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose
every tool the core toolset serves."""
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=None if allowed_tools is None else list(allowed_tools),
mcp_access_groups=mcp_access_groups,
)
resources.defer(lambda: client.delete_server(server_id))

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

@ -16,11 +16,11 @@ import time
from collections.abc import Mapping
from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict, Field, RootModel
from pydantic import BaseModel, ConfigDict, Field
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, McpServerListResponse, McpServerRow, ObjectPermission
from proxy_client import ProxyClient
McpToolArg = str | int | float | bool | list[str] | dict[str, str]
@ -46,16 +46,6 @@ class McpServerNewResponse(BaseModel):
server_id: str
class McpServerRow(BaseModel):
server_id: str
alias: str | None = None
url: str | None = None
class McpServersListResponse(RootModel[list[McpServerRow]]):
pass
class McpToolMcpInfo(BaseModel):
server_id: str | None = None
alias: str | None = None
@ -193,7 +183,7 @@ class McpClient:
"/v1/mcp/server",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpServersListResponse,
response_type=McpServerListResponse,
)
).root
@ -224,11 +214,16 @@ class McpClient:
user_id: str,
mcp_servers: list[str] | None,
mcp_access_groups: list[str] | None = None,
mcp_toolsets: list[str] | None = None,
models: list[str] | None = None,
) -> str:
object_permission = (
ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups)
if mcp_servers is not None or mcp_access_groups is not None
ObjectPermission(
mcp_servers=mcp_servers,
mcp_access_groups=mcp_access_groups,
mcp_toolsets=mcp_toolsets,
)
if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None
else None
)
return self.proxy.generate_key(
@ -272,6 +267,20 @@ class McpClient:
)
time.sleep(self.proxy.poll_interval)
def await_tools(self, key: str, server_id: str, *, expected: frozenset[str]) -> frozenset[str]:
"""Poll tools/list until `server_id`'s tools as `key` sees them are exactly
`expected`, and return the last listing either way, so the caller's equality
assertion names the difference. Fails at poll_timeout only when the read
itself never succeeded."""
deadline = time.monotonic() + self.proxy.poll_timeout
while True:
result = self.list_tools(key)
if isinstance(result, Success) and result.data.tool_names_for_server(server_id) == expected:
return expected
if time.monotonic() >= deadline:
return unwrap(result).tool_names_for_server(server_id)
time.sleep(self.proxy.poll_interval)
def await_call_tool(
self,
key: str,

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

@ -0,0 +1,77 @@
"""Live e2e: a key granted a toolset lists exactly the toolset's tools.
An admin registers the real Datadog remote MCP server with its whole core toolset
exposed, discovers two of its tool names through a key granted the server outright,
and curates a toolset naming exactly those two. A second key is granted the server
plus that toolset, and its tools/list must come back as exactly those two names: no
more, so the rest of the server's catalog stays hidden behind the toolset, and no
fewer, so a tool stored under one name and read under another (which granted
nothing) fails here first. Requires DD_API_KEY + DD_APP_KEY (the suite's real MCP
upstream).
"""
from __future__ import annotations
from typing import Final
import pytest
from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp
from e2e_config import unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient
from models import ToolsetCreateBody, ToolsetTool
pytestmark = pytest.mark.e2e
def _key(
client: McpClient,
resources: ResourceManager,
label: str,
*,
server_id: str,
toolset_id: str | None = None,
) -> str:
key: Final = client.generate_key(
user_id=f"e2e-mcp-{label}-{unique_marker()}",
mcp_servers=[server_id],
mcp_toolsets=None if toolset_id is None else [toolset_id],
)
resources.defer(lambda: client.proxy.delete_key(key))
return key
class TestMcpToolsetEnforcement:
@pytest.mark.covers("mcp.list_tools.api_key.toolset_scoped")
def test_key_granted_a_toolset_lists_exactly_its_tools(self, client: McpClient, resources: ResourceManager) -> None:
server_id: Final = register_datadog_mcp(client, resources, allowed_tools=None)
client.await_registered(server_id)
catalog_key: Final = _key(client, resources, "catalog", server_id=server_id)
_ = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL)
catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id)
assert len(catalog) > 2, (
f"the Datadog core toolset must serve more tools than the toolset names, or the "
f"restriction has nothing to hide; got {sorted(catalog)}"
)
chosen: Final = frozenset(sorted(catalog)[:2])
toolset: Final = client.proxy.create_toolset(
ToolsetCreateBody(
toolset_name=f"e2e_toolset_{unique_marker()}",
description="two Datadog tools",
tools=[ToolsetTool(server_id=server_id, tool_name=name) for name in sorted(chosen)],
)
)
resources.defer(lambda: client.proxy.delete_toolset(toolset.toolset_id))
assert frozenset(tool.tool_name for tool in toolset.tools) == chosen, (
f"toolset stored {toolset.tools}, expected the two names {sorted(chosen)} verbatim"
)
scoped_key: Final = _key(client, resources, "toolset", server_id=server_id, toolset_id=toolset.toolset_id)
listed: Final = client.await_tools(scoped_key, server_id, expected=chosen)
assert listed == chosen, (
f"a key granted the toolset must list exactly its two tools; "
f"got {sorted(listed)}, expected {sorted(chosen)}"
)

73
tests/e2e/models.py
View file

@ -10,6 +10,7 @@ from collections.abc import Sequence
from datetime import datetime
from typing import Literal
from e2e_http import PartialBody
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator
# ---------- keys ----------
@ -54,6 +55,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):
@ -76,7 +78,7 @@ class KeyGenerateBody(BaseModel):
allowed_passthrough_routes: list[str] | None = None
metadata: KeyMetadata | None = None
object_permission: ObjectPermission | None = None
router_settings: "RouterSettingsOverride | None" = None
router_settings: RouterSettingsOverride | None = None
class KeyGenerateResponse(BaseModel):
@ -505,6 +507,15 @@ class CountTokensResponse(BaseModel):
# ---------- mcp servers ----------
class McpInfo(BaseModel):
"""The `mcp_info` display block stored on an MCP server; only the fields the
lifecycle test writes and reads back."""
server_name: str | None = None
description: str | None = None
logo_url: str | None = None
class McpServerCreateBody(BaseModel):
"""POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
`oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
@ -519,6 +530,18 @@ class McpServerCreateBody(BaseModel):
oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None
authorization_url: str | None = None
token_url: str | None = None
server_name: str | None = None
description: str | None = None
mcp_info: McpInfo | None = None
class McpServerUpdateBody(PartialBody):
"""PUT /v1/mcp/server: a field left unset keeps its stored value, a field set
to None is cleared."""
server_id: str
alias: str | None = None
description: str | None = None
class McpServerInfo(BaseModel):
@ -532,6 +555,54 @@ class McpServerInfo(BaseModel):
allow_all_keys: bool | None = None
class McpServerRow(McpServerInfo):
"""A stored MCP server as the create, get, and list routes return it: the
fields the lifecycle test asserts survive the round trip."""
server_name: str | None = None
transport: str | None = None
description: str | None = None
mcp_info: McpInfo | None = None
class McpServerListResponse(RootModel[list[McpServerRow]]):
"""GET /v1/mcp/server answers with a bare array of servers."""
class ToolsetTool(BaseModel):
server_id: str
tool_name: str
class ToolsetCreateBody(BaseModel):
toolset_name: str
description: str | None = None
tools: list[ToolsetTool]
class ToolsetUpdateBody(PartialBody):
"""PUT /v1/mcp/toolset: a field left unset keeps its stored value, a field set
to None is cleared."""
toolset_id: str
description: str | None = None
tools: list[ToolsetTool] | None = None
class ToolsetRow(BaseModel):
"""A stored toolset as POST /v1/mcp/toolset, GET /v1/mcp/toolset/{toolset_id},
and each row of GET /v1/mcp/toolset return it."""
toolset_id: str
toolset_name: str
description: str | None = None
tools: list[ToolsetTool] = []
class ToolsetListResponse(RootModel[list[ToolsetRow]]):
"""GET /v1/mcp/toolset answers with a bare array of toolsets."""
class EmbedBody(BaseModel):
model: str
input: str

203
tests/e2e/proxy_client.py
View file

@ -16,6 +16,8 @@ from datetime import datetime
from types import MappingProxyType
from typing import Final
from pydantic import BaseModel
from e2e_http import (
AnthropicHeaders,
AuthHeaders,
@ -24,6 +26,7 @@ from e2e_http import (
Result,
StreamingResponse,
Success,
UnknownApiError,
is_ok,
unwrap,
)
@ -68,6 +71,9 @@ from models import (
SpendLogsPage,
SpendLogsPageParams,
SpendLogsParams,
ToolsetCreateBody,
ToolsetRow,
ToolsetUpdateBody,
)
from e2e_config import (
CONTROL_PLANE_BASE_URL,
@ -80,7 +86,7 @@ from e2e_config import (
SLOW_PROVIDER_TIMEOUT_SECONDS,
settle_propagation,
)
from transport import HttpTransport, SplitTransport, Transport
from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path
RowsPredicate = Callable[[list[SpendLogRow]], bool]
@ -239,10 +245,97 @@ def servable_timeout_message(
)
type ReplicaRead[T] = Callable[[float], T]
@dataclass(frozen=True, slots=True)
class Converged[T]:
"""Every replica answered with something `settled` accepts, keyed by replica."""
answers: Mapping[str, T]
@dataclass(frozen=True, slots=True)
class NeverConvergedOn[T]:
"""`replica` ran out its budget without an answer `settled` accepts; `last` is
its final answer, so the failure can say what that replica still serves."""
replica: str
last: T
def _last_answer[T](
read: ReplicaRead[T],
*,
settled: Callable[[T], bool],
timeout: float,
interval: float,
request_timeout: float,
now: Callable[[], float],
sleep: Callable[[float], None],
) -> T:
"""Poll `read` until `settled` accepts its answer or `timeout` runs out, and
return the last answer either way. Each read's request timeout is clamped to
the budget left, and the final poll runs even when less than an interval
remains, so a deadline never skips the read that would have settled."""
deadline: Final = now() + timeout
answer = read(min(request_timeout, timeout))
while not settled(answer):
remaining = deadline - now()
if remaining <= 0:
return answer
sleep(min(interval, remaining))
answer = read(min(request_timeout, remaining))
return answer
def await_everywhere[T](
reads: Mapping[str, ReplicaRead[T]],
*,
settled: Callable[[T], bool],
timeout: float,
interval: float,
request_timeout: float,
now: Callable[[], float],
sleep: Callable[[float], None],
) -> Converged[T] | NeverConvergedOn[T]:
"""`_last_answer` against every replica in turn, each with the full budget, so a
write counts as visible only once the last replica reflects it, and stop at the
first replica that never converges. Clock and sleep are injected."""
answers: dict[str, T] = {}
for replica, read in reads.items():
answer = _last_answer(
read,
settled=settled,
timeout=timeout,
interval=interval,
request_timeout=request_timeout,
now=now,
sleep=sleep,
)
if not settled(answer):
return NeverConvergedOn(replica=replica, last=answer)
answers[replica] = answer
return Converged(answers=MappingProxyType(answers))
def _is_not_found[R: BaseModel](result: Result[R]) -> bool:
return isinstance(result, UnknownApiError) and result.status_code == 404
def _status_of[R: BaseModel](result: Result[R]) -> int:
match result:
case Success(status_code=status_code) | UnknownApiError(status_code=status_code):
return status_code
case _:
return -1
@dataclass(frozen=True, slots=True)
class ProxyClient:
transport: Transport
replicas: Mapping[str, Transport]
control_replicas: Mapping[str, Transport]
poll_timeout: float = 120.0
poll_interval: float = 5.0
model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT
@ -443,6 +536,105 @@ class ProxyClient:
if not is_ok(result):
warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2)
# ---- replica read-back ----------------------------------------------
def replicas_for(self, path: str) -> Mapping[str, Transport]:
"""The replicas that serve `path`: every data-plane replica for an LLM route,
and for a management route the control-plane replicas, which in a split
deployment is the one service that serves it (the data-plane replicas trim
management routes) and in a monolith is every replica."""
return self.control_replicas if is_control_plane_path(path) else self.replicas
def read_back_everywhere[R: BaseModel](
self, path: str, response_type: type[R], *, settled: Callable[[R], bool]
) -> Mapping[str, R]:
"""GET `path` on every replica that serves it, polling each to poll_timeout
until `settled` accepts its body, and fail naming the first replica that
never converged. Returns each replica's settled body, keyed by replica, so
the caller can assert the rest of it."""
outcome: Final = await_everywhere(
{url: self._reader(transport, path, response_type) for url, transport in self.replicas_for(path).items()},
settled=lambda result: isinstance(result, Success) and settled(result.data),
timeout=self.poll_timeout,
interval=self.poll_interval,
request_timeout=REQUEST_TIMEOUT,
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case Converged(answers=answers):
return MappingProxyType({url: unwrap(result) for url, result in answers.items()})
case NeverConvergedOn(replica=replica, last=last):
raise AssertionError(
f"GET {path} on {replica} never converged within {self.poll_timeout}s of the write; "
f"last read: {last}"
)
def gone_everywhere(self, path: str) -> Mapping[str, int]:
"""Poll GET `path` on every replica that serves it until each stops serving
it, and fail naming the first replica that still does at poll_timeout.
Returns each replica's final status, so the caller asserts the 404 itself."""
outcome: Final = await_everywhere(
{url: self._reader(transport, path, NoBody) for url, transport in self.replicas_for(path).items()},
settled=_is_not_found,
timeout=self.poll_timeout,
interval=self.poll_interval,
request_timeout=REQUEST_TIMEOUT,
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case Converged(answers=answers):
return MappingProxyType({url: _status_of(result) for url, result in answers.items()})
case NeverConvergedOn(replica=replica, last=last):
raise AssertionError(
f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}"
)
@staticmethod
def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]:
return lambda request_timeout: transport.get(
path,
headers=transport.master,
params=NoBody(),
response_type=response_type,
timeout=request_timeout,
)
# ---- mcp toolsets ---------------------------------------------------
def create_toolset(self, body: ToolsetCreateBody) -> ToolsetRow:
return unwrap(
self.transport.post(
"/v1/mcp/toolset",
headers=self.transport.master,
json=body,
response_type=ToolsetRow,
)
)
def update_toolset(self, body: ToolsetUpdateBody) -> ToolsetRow:
"""PUT /v1/mcp/toolset: a partial update where a field left unset keeps its
stored value and None clears it."""
return unwrap(
self.transport.put(
"/v1/mcp/toolset",
headers=self.transport.master,
json=body,
response_type=ToolsetRow,
)
)
def delete_toolset(self, toolset_id: str) -> Result[NoBody]:
"""DELETE /v1/mcp/toolset/{toolset_id}. Returns the outcome so the act phase
can unwrap it while a deferred teardown can ignore an already-deleted row."""
return self.transport.delete(
f"/v1/mcp/toolset/{toolset_id}",
headers=self.transport.master,
json=NoBody(),
response_type=NoBody,
)
def create_credential(self, body: CredentialCreateBody) -> None:
unwrap(
self.transport.post(
@ -610,7 +802,10 @@ def build_proxy_client(
base URLs are the same for a monolithic proxy, so routing is then a no-op.
``replica_urls`` (PROXY_REPLICA_URLS) names every data-plane replica the model
barrier polls directly; it is the data-plane URL itself unless the stack
exports each gateway's own address.
exports each gateway's own address. Management read-backs poll those same
replicas when the two planes share a base URL (a monolith, where every replica
serves every route) and the control plane alone when they differ (a split
deployment, where the data-plane replicas do not serve management routes).
The endpoints are injectable for callers that resolve the proxy some other
way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must
@ -638,9 +833,13 @@ def build_proxy_client(
for url in replica_urls
}
)
control_replicas: Final = (
replicas if control_plane_base_url == base_url else MappingProxyType({control_plane_base_url: split.control})
)
return ProxyClient(
transport=split,
replicas=replicas,
control_replicas=control_replicas,
poll_timeout=POLL_TIMEOUT,
poll_interval=POLL_INTERVAL,
)

77
tests/e2e/test_e2e_http.py
View file

@ -18,8 +18,19 @@ from types import MappingProxyType
from typing import Final
import pytest
from e2e_http import RETRY_ATTEMPTS, TRANSIENT_STATUSES, request_with_retry, streaming_outcome
from e2e_http import (
RETRY_ATTEMPTS,
TRANSIENT_STATUSES,
NoBody,
PartialBody,
Success,
ValidationError,
classify,
request_with_retry,
streaming_outcome,
wire_body,
)
from pydantic import BaseModel, TypeAdapter
@dataclass
@ -134,3 +145,65 @@ class TestStreamEventArrivals:
assert result.stream_events == []
assert result.stream_event_arrivals == []
assert result.body == "bad request"
class _ServerUpdate(PartialBody):
server_id: str
alias: str | None = None
description: str | None = None
class _ServerCreate(BaseModel):
alias: str
description: str | None = None
class TestWireBody:
"""A partial-update body must put exactly the caller's choice on the wire: an
omitted field stays off it so the route keeps the stored value, and an explicit
None goes out as JSON null so the route clears it. Plain bodies keep dropping
None, which is what every create route expects."""
def test_partial_body_omits_unset_fields_and_sends_explicit_none_as_null(self) -> None:
assert wire_body(_ServerUpdate(server_id="s1", description=None)) == {"server_id": "s1", "description": None}
assert wire_body(_ServerUpdate(server_id="s1", alias="renamed")) == {"server_id": "s1", "alias": "renamed"}
def test_plain_body_drops_none_fields(self) -> None:
assert wire_body(_ServerCreate(alias="a", description=None)) == {"alias": "a"}
_JSON: Final[TypeAdapter[object]] = TypeAdapter(object)
@dataclass
class FakeJsonResponse:
"""The `_classify` view of a response: a status, the raw body bytes, and the
parse that would raise on an empty one."""
status_code: int
content: bytes
@property
def ok(self) -> bool:
return self.status_code < 400
@property
def text(self) -> str:
return self.content.decode()
def json(self) -> object:
return _JSON.validate_json(self.content)
class TestClassifyEmptyBody:
"""A delete that answers 202 with no body is a success, not a parse failure:
the MCP server and toolset delete routes both answer that way, and reading it
as a failure would hide a delete that did not happen behind one that did."""
def test_empty_2xx_body_is_a_success(self) -> None:
result: Final = classify(FakeJsonResponse(status_code=202, content=b""), NoBody)
assert isinstance(result, Success) and result.status_code == 202
def test_body_that_is_not_json_is_still_a_validation_failure(self) -> None:
result: Final = classify(FakeJsonResponse(status_code=200, content=b"<html/>"), NoBody)
assert isinstance(result, ValidationError)

81
tests/e2e/test_proxy_client.py
View file

@ -1,10 +1,11 @@
"""Harness coverage for the model barrier that gates on every replica.
"""Harness coverage for the barriers that gate on every replica.
No proxy needed and no ``e2e`` marker: this pins that a model registered through
the control plane only counts as servable once every configured replica lists it
on /v1/models, which is what keeps a two-gateway stack from handing a test a
model that one gateway has not reloaded yet. The fakes are plain pollers and an
injected clock, so nothing here monkeypatches anything.
on /v1/models, and that a management write only counts as read back once every
replica that serves the route reflects it, which is what keeps a multi-replica
stack from handing a test a replica the write has not reached yet. The fakes are
plain pollers and an injected clock, so nothing here monkeypatches anything.
"""
from __future__ import annotations
@ -19,7 +20,17 @@ import pytest
from e2e_config import parse_replica_urls
from e2e_http import Success
from models import ModelListEntry, ModelsListResponse
from proxy_client import ModelsPoller, NotServableOn, Servable, await_servable_everywhere
from proxy_client import (
Converged,
ModelsPoller,
NeverConvergedOn,
NotServableOn,
ReplicaRead,
Servable,
await_everywhere,
await_servable_everywhere,
build_proxy_client,
)
MODEL: Final = "gpt-under-test"
TIMEOUT: Final = 10.0
@ -85,3 +96,63 @@ class TestParseReplicaUrls:
def test_falls_back_to_the_data_plane_address_when_unset(self) -> None:
assert parse_replica_urls("", "http://lb") == ("http://lb",)
def _answers(answers: Iterable[str]) -> ReplicaRead[str]:
it: Final = iter(answers)
return lambda _timeout: next(it)
def _await_everywhere(reads: Mapping[str, ReplicaRead[str]]) -> Converged[str] | NeverConvergedOn[str]:
clock: Final = FakeClock()
return await_everywhere(
reads,
settled=lambda answer: answer == "renamed",
timeout=TIMEOUT,
interval=INTERVAL,
request_timeout=5.0,
now=clock.now,
sleep=clock.sleep,
)
class TestAwaitEverywhere:
def test_waits_for_the_lagging_replica_and_returns_every_settled_answer(self) -> None:
reads: Final = {
"gateway-1": _answers(repeat("renamed")),
"gateway-2": _answers(chain(repeat("stale", 2), repeat("renamed"))),
}
outcome: Final = _await_everywhere(reads)
assert isinstance(outcome, Converged)
assert dict(outcome.answers) == {"gateway-1": "renamed", "gateway-2": "renamed"}
def test_names_the_replica_that_never_converges_with_what_it_last_served(self) -> None:
reads: Final = {
"gateway-1": _answers(repeat("renamed")),
"gateway-2": _answers(repeat("stale")),
}
assert _await_everywhere(reads) == NeverConvergedOn(replica="gateway-2", last="stale")
def test_polls_until_the_deadline_before_giving_up(self) -> None:
lagging: Final = chain(repeat("stale", int(TIMEOUT / INTERVAL)), repeat("renamed"))
outcome: Final = _await_everywhere({"gateway-1": _answers(lagging)})
assert isinstance(outcome, Converged), outcome
class TestReplicasFor:
def test_split_deployment_reads_management_routes_back_from_the_control_plane(self) -> None:
client: Final = build_proxy_client(
base_url="http://lb",
control_plane_base_url="http://backend",
replica_urls=("http://gateway-1", "http://gateway-2"),
)
assert set(client.replicas_for("/v1/mcp/server/abc")) == {"http://backend"}
assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"}
def test_monolith_reads_management_routes_back_from_every_replica(self) -> None:
client: Final = build_proxy_client(
base_url="http://lb",
control_plane_base_url="http://lb",
replica_urls=("http://pod-1", "http://pod-2"),
)
assert set(client.replicas_for("/v1/mcp/server/abc")) == {"http://pod-1", "http://pod-2"}

1
tests/e2e/transport.py
View file

@ -291,6 +291,7 @@ class HttpTransport:
# (/chat, /embeddings, and native passthrough like /gemini, /anthropic) are NOT
# here and fall through to the data plane. Matched as path prefixes.
CONTROL_PLANE_PREFIXES: tuple[str, ...] = (
"/v1/mcp/",
"/key",
"/user",
"/team",

63
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py
View file

@ -1,10 +1,11 @@
"""
Tests for partial-update semantics of PUT /v1/mcp/server.
Tests for partial-update semantics of PUT /v1/mcp/server and PUT /v1/mcp/toolset.
A partial update must only write the fields the caller explicitly provided.
Omitting a field must NOT reset it to its Pydantic schema default (e.g.
``transport=sse``, ``mcp_access_groups=[]``, ``allow_all_keys=False``), which
would silently overwrite the existing DB row.
would silently overwrite the existing DB row, and a field the caller sent as null
must be cleared rather than left at its stored value.
"""
import json
@ -847,3 +848,61 @@ async def test_cf_pair_switch_does_not_clear_dcr_bridge():
data = UpdateMCPServerRequest(server_id="s", auth_type="oauth_delegate")
data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough")
assert "dcr_bridge" not in data_dict
def _mock_toolset_prisma():
"""A prisma double whose update answers with a row the reader can expand, so the
call under test returns instead of failing inside the row mapper."""
updated_row = MagicMock()
updated_row.model_dump.return_value = {
"toolset_id": "ts-1",
"toolset_name": "ops",
"description": None,
"tools": "[]",
}
mock_prisma = MagicMock()
mock_prisma.db.litellm_mcptoolsettable = AsyncMock()
mock_prisma.db.litellm_mcptoolsettable.update = AsyncMock(return_value=updated_row)
return mock_prisma
async def _run_toolset_update(payload: dict) -> dict:
"""The columns PUT /v1/mcp/toolset writes for this payload, minus the audit stamp
every write carries. The prisma double is injected, so nothing is patched."""
from litellm.proxy._experimental.mcp_server.toolset_db import update_mcp_toolset
from litellm.types.mcp_server.mcp_toolset import UpdateMCPToolsetRequest
mock_prisma = _mock_toolset_prisma()
await update_mcp_toolset(mock_prisma, UpdateMCPToolsetRequest.model_validate(payload), "test-user")
written = dict(mock_prisma.db.litellm_mcptoolsettable.update.call_args[1]["data"])
assert written.pop("updated_by") == "test-user"
return written
@pytest.mark.asyncio
async def test_toolset_partial_update_clears_description_on_explicit_null():
"""The dump used to drop None, so a null description could never clear the stored
one: the toolset kept a description its owner had deleted."""
assert await _run_toolset_update({"toolset_id": "ts-1", "description": None}) == {"description": None}
@pytest.mark.asyncio
async def test_toolset_partial_update_omits_the_fields_the_caller_left_out():
tools = [{"server_id": "s1", "tool_name": "alpha"}]
assert await _run_toolset_update({"toolset_id": "ts-1", "tools": tools}) == {"tools": json.dumps(tools)}
@pytest.mark.asyncio
async def test_toolset_partial_update_writes_null_tools_as_an_empty_list():
"""Prisma ``tools`` is a required Json column defaulting to [], so a null clears
the selection to none rather than writing SQL null."""
assert await _run_toolset_update({"toolset_id": "ts-1", "tools": None}) == {"tools": "[]"}
@pytest.mark.asyncio
async def test_toolset_partial_update_ignores_a_null_name():
"""A toolset always has a name, so a null toolset_name is a no-op, not a clear
that would write a NOT NULL column to null."""
assert await _run_toolset_update({"toolset_id": "ts-1", "toolset_name": None, "description": "kept"}) == {
"description": "kept"
}