diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index ecaaf35e817..f13b1c72558 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -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( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d5c3427f29a..6f98717400e 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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( diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..f853e9ff8d6 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -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" diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d571fb36546..bebca73ef0c 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -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"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 9d5f1658e91..415c72bbb3c 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -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, ) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 2b897f5f07f..1ef0d89a8f9 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -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", diff --git a/tests/e2e/management/test_mcp_lifecycle_e2e.py b/tests/e2e/management/test_mcp_lifecycle_e2e.py new file mode 100644 index 00000000000..d96b6c509ec --- /dev/null +++ b/tests/e2e/management/test_mcp_lifecycle_e2e.py @@ -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}" + ) diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..352b4446cfd 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from 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)) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..210fc7a1e98 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -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, diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py new file mode 100644 index 00000000000..d164dd2a48b --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -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)}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2c6c0e9bbd4..b618275243a 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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 diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 3b15e57dfcf..47c5d0abaff 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -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, ) diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 66841725d1d..6ab35dd6cd5 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -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""), NoBody) + assert isinstance(result, ValidationError) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index a508c97b9fb..88931525c12 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -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"} diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 44fdbaa3e41..be32304e682 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -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", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index f6bd79c5d2d..ce5fd76b7ad 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -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" + }