diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index a6f0d145e9..e30667776c 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -30,6 +30,8 @@ from litellm.types.mcp import MCPCredentials def _prepare_mcp_server_data( data: Union[NewMCPServerRequest, UpdateMCPServerRequest], + exclude_unset: bool = False, + fields_set: Optional[Set[str]] = None, ) -> Dict[str, Any]: """ Helper function to prepare MCP server data for database operations. @@ -37,17 +39,39 @@ def _prepare_mcp_server_data( Args: data: NewMCPServerRequest or UpdateMCPServerRequest object + exclude_unset: When True, only fields the caller explicitly provided are + included. Used for partial updates (PUT /v1/mcp/server) so omitted + fields keep their existing DB value instead of being silently reset + to a Pydantic schema default. ``exclude_none`` is not enough here: + non-Optional fields (e.g. ``transport=MCPTransport.sse``, + ``mcp_access_groups=[]``, ``allow_all_keys=False``) are backfilled + with their default when omitted, and a non-None default survives the + ``exclude_none`` filter and overwrites the row. Returns: Dict with properly serialized JSON fields """ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - # Convert model to dict - data_dict = data.model_dump(exclude_none=True) - # Ensure alias is always present in the dict (even if None) - if "alias" not in data_dict: - data_dict["alias"] = getattr(data, "alias", None) + # Convert model to dict. + # - Partial update (exclude_unset): only caller-provided keys are emitted, so + # omitted fields are never written and keep their existing DB value. + # - Create (exclude_none): drop None-valued fields and let DB defaults apply. + if exclude_unset: + if fields_set is None: + fields_set = data.fields_set() + data_dict = data.model_dump(exclude_unset=True) + # ``validate_and_normalize_mcp_server_payload`` always assigns ``alias`` + # on the payload, which marks it as set even when the caller omitted it. + # Drop it only when the original request omitted alias; an explicit + # ``alias=None`` is a valid request to clear the stored alias. + if data_dict.get("alias") is None and "alias" not in fields_set: + data_dict.pop("alias", None) + else: + data_dict = data.model_dump(exclude_none=True) + # Ensure alias is always present in the dict (even if None) + if "alias" not in data_dict: + data_dict["alias"] = getattr(data, "alias", None) # Handle credentials serialization credentials = data_dict.get("credentials") @@ -57,33 +81,33 @@ def _prepare_mcp_server_data( ) data_dict["credentials"] = safe_dumps(data_dict["credentials"]) - # Handle static_headers serialization - if data.static_headers is not None: - data_dict["static_headers"] = safe_dumps(data.static_headers) + # Serialize JSON fields from ``data_dict`` (not ``data``) so the + # exclude_unset filter is respected. Reading back from ``data`` would + # reintroduce defaults (e.g. ``env={}``) for fields the caller never set. + if data_dict.get("static_headers") is not None: + data_dict["static_headers"] = safe_dumps(data_dict["static_headers"]) - # Handle mcp_info serialization - if data.mcp_info is not None: - data_dict["mcp_info"] = safe_dumps(data.mcp_info) + if data_dict.get("mcp_info") is not None: + data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"]) - # Handle env serialization - if data.env is not None: - data_dict["env"] = safe_dumps(data.env) + if data_dict.get("env") is not None: + data_dict["env"] = safe_dumps(data_dict["env"]) - # Handle tool name override serialization - if data.tool_name_to_display_name is not None: + if data_dict.get("tool_name_to_display_name") is not None: data_dict["tool_name_to_display_name"] = safe_dumps( - data.tool_name_to_display_name + data_dict["tool_name_to_display_name"] ) - if data.tool_name_to_description is not None: + if data_dict.get("tool_name_to_description") is not None: data_dict["tool_name_to_description"] = safe_dumps( - data.tool_name_to_description + data_dict["tool_name_to_description"] ) # mcp_access_groups is already List[str], no serialization needed - # Force include is_byok even when False (exclude_none=True would not drop it, - # but be explicit to ensure a False value is always written to the DB). - data_dict["is_byok"] = getattr(data, "is_byok", False) + # On create, force is_byok so a False value is always written to the DB. On + # partial update, only write it when the caller explicitly provided it. + if not exclude_unset: + data_dict["is_byok"] = getattr(data, "is_byok", False) return data_dict @@ -398,7 +422,10 @@ async def create_mcp_server( async def update_mcp_server( - prisma_client: PrismaClient, data: UpdateMCPServerRequest, touched_by: str + prisma_client: PrismaClient, + data: UpdateMCPServerRequest, + touched_by: str, + fields_set: Optional[Set[str]] = None, ) -> LiteLLM_MCPServerTable: """ Update a new mcp server record in the db @@ -407,8 +434,13 @@ async def update_mcp_server( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - # Use helper to prepare data with proper JSON serialization - data_dict = _prepare_mcp_server_data(data) + # Use helper to prepare data with proper JSON serialization. + # exclude_unset=True makes this a true partial update: fields the caller did + # not provide are not written, so they keep their existing DB value instead + # of being reset to a schema default (transport=sse, allow_all_keys=False...). + data_dict = _prepare_mcp_server_data( + data, exclude_unset=True, fields_set=fields_set + ) # Pre-fetch existing record once if we need it for auth_type or credential logic existing = None diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 431ff49c7c..b35e2b6e3f 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2135,6 +2135,8 @@ if MCP_AVAILABLE: "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) + payload_fields_set = set(payload.fields_set()) + # Validate and normalize payload fields validate_and_normalize_mcp_server_payload(payload) @@ -2154,6 +2156,7 @@ if MCP_AVAILABLE: prisma_client, payload, touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + fields_set=payload_fields_set, ) if mcp_server_record_updated is None: 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 new file mode 100644 index 0000000000..b5e0f20f66 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -0,0 +1,177 @@ +""" +Tests for partial-update semantics of PUT /v1/mcp/server. + +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. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy._experimental.mcp_server.db import ( + create_mcp_server, + update_mcp_server, +) +from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest + + +def _mock_prisma(): + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable = AsyncMock() + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=MagicMock()) + return mock_prisma + + +async def _run_update(data: UpdateMCPServerRequest, fields_set=None) -> dict: + mock_prisma = _mock_prisma() + await update_mcp_server(mock_prisma, data, "test-user", fields_set=fields_set) + return mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + +@pytest.mark.asyncio +async def test_partial_update_omits_unset_defaultful_fields(): + """ + A PUT touching only allowed_tools must not write transport, + mcp_access_groups, allow_all_keys, available_on_public_internet, + delegate_auth_to_upstream, is_byok, args, env or byok_description. + """ + data = UpdateMCPServerRequest( + server_id="my-test-server", + allowed_tools=["foo"], + ) + + data_dict = await _run_update(data) + + # The intended change is present. + assert data_dict["allowed_tools"] == ["foo"] + + # Fields the caller did not provide must not be in the write payload, so the + # existing DB value is preserved. + for trapped_field in ( + "transport", + "mcp_access_groups", + "allow_all_keys", + "available_on_public_internet", + "delegate_auth_to_upstream", + "is_byok", + "args", + "env", + "byok_description", + ): + assert trapped_field not in data_dict, ( + f"{trapped_field} should not be written on a partial update that " + f"omitted it (would reset the row to a schema default)" + ) + + +@pytest.mark.asyncio +async def test_partial_update_preserves_http_transport(): + """The reported prod incident: a PUT without transport must not flip http->sse.""" + data = UpdateMCPServerRequest( + server_id="atlassian_url", + allowed_tools=[], + ) + + data_dict = await _run_update(data) + + assert "transport" not in data_dict + assert data_dict["allowed_tools"] == [] + + +@pytest.mark.asyncio +async def test_partial_update_writes_explicitly_provided_fields(): + """Explicitly provided fields are written, including falsy/default-equal values.""" + data = UpdateMCPServerRequest( + server_id="my-test-server", + url="https://example.com/mcp", + transport="http", + allow_all_keys=False, + mcp_access_groups=["mcp-dev-sandbox"], + available_on_public_internet=True, + ) + + data_dict = await _run_update(data) + + assert data_dict["transport"] == "http" + # Explicitly provided False must still be written. + assert data_dict["allow_all_keys"] is False + assert data_dict["mcp_access_groups"] == ["mcp-dev-sandbox"] + assert data_dict["available_on_public_internet"] is True + + +@pytest.mark.asyncio +async def test_partial_update_can_explicitly_reset_allow_all_keys(): + """Caller can still reset a field to its default by sending it explicitly.""" + enabled = await _run_update( + UpdateMCPServerRequest(server_id="s", allow_all_keys=True) + ) + assert enabled["allow_all_keys"] is True + + disabled = await _run_update( + UpdateMCPServerRequest(server_id="s", allow_all_keys=False) + ) + assert disabled["allow_all_keys"] is False + + +@pytest.mark.asyncio +async def test_partial_update_does_not_clear_alias_when_unset(): + """alias is force-normalized on the payload; an unset/None alias must not be written.""" + data = UpdateMCPServerRequest( + server_id="my-test-server", + allowed_tools=["foo"], + ) + fields_set = set(data.fields_set()) + # Simulate validate_and_normalize_mcp_server_payload assigning alias=None. + data.alias = None + + data_dict = await _run_update(data, fields_set=fields_set) + + assert "alias" not in data_dict + + +@pytest.mark.asyncio +async def test_partial_update_can_explicitly_clear_alias(): + """Caller can clear an existing alias by explicitly sending alias=None.""" + data = UpdateMCPServerRequest( + server_id="my-test-server", + alias=None, + ) + fields_set = set(data.fields_set()) + # Simulate validate_and_normalize_mcp_server_payload preserving alias=None. + data.alias = None + + data_dict = await _run_update(data, fields_set=fields_set) + + assert "alias" in data_dict + assert data_dict["alias"] is None + + +@pytest.mark.asyncio +async def test_create_still_writes_defaults(): + """ + Regression guard: create (POST) must keep writing defaults so DB columns + without a default get populated. exclude_unset is update-only. + """ + mock_prisma = _mock_prisma() + data = NewMCPServerRequest( + server_id="new-server", + url="https://example.com/mcp", + transport="http", + ) + + await create_mcp_server(mock_prisma, data, "test-user") + + data_dict = mock_prisma.db.litellm_mcpservertable.create.call_args[1]["data"] + + assert data_dict["transport"] == "http" + # is_byok is force-written on create. + assert data_dict["is_byok"] is False + # alias key is always present on create (even if None). + assert "alias" in data_dict + # audit fields set by create_mcp_server. + assert data_dict["created_by"] == "test-user" + assert data_dict["updated_by"] == "test-user"