ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(mcp): preserve omitted fields on PUT /v1/mcp/server partial updates (#29253)

Browse Source 
* fix(mcp): preserve omitted fields on PUT /v1/mcp/server partial updates

Use model_dump(exclude_unset=True) for updates so schema defaults (transport=sse,
allow_all_keys=false, etc.) are not written when callers omit them. Serialize
JSON fields from the filtered dict and only force is_byok on create.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(mcp): allow explicit alias=None on MCP server partial updates

Snapshot caller-provided fields before normalization so omitted alias is not written while an intentional alias=None still clears the stored value.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Shivam Rawat 2026-05-30 12:03:22 -07:00 committed by GitHub
parent ace3c65ab3
commit 2d4c13c00f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 237 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
litellm/proxy/_experimental/mcp_server/db.py
View File

@ -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

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

@ -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:

177
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py Normal file
View File

@ -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"