litellm/tests/e2e/management/test_mcp_lifecycle_e2e.py

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