From 22563a80d1cb41d954795365c596b7a1fa9e9831 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 21:50:30 -0800 Subject: [PATCH] fix(mcp): keep DeepSeek tool cleanup and the guardrail policy result on lts/1.108.x DeepSeek now sends the params through its own tool cleanup when a hosted MCP tool is dropped, so a tool_choice naming that tool and the tool selection params beside legacy functions are cleared as the line does today. The MCP security guardrail uses the gateway's saved toolset check, so a failed read reports the name as unregistered with the policy result instead of raising. Also fixes three lint errors the line already has, two import orders and one duplicate test import, since make check cannot pass on this line without them. --- litellm/llms/deepseek/chat/transformation.py | 7 ++++- .../redis_update_buffer.py | 1 - .../mcp_security/mcp_security_guardrail.py | 12 ++------ .../mcp/litellm_proxy_mcp_handler.py | 4 +-- litellm/utils.py | 2 +- .../test_streaming_chunk_builder_utils.py | 1 - .../chat/test_deepseek_chat_transformation.py | 30 +++++++++++++++++++ .../guardrail_hooks/test_mcp_security.py | 27 +++++++++++++++++ 8 files changed, 68 insertions(+), 16 deletions(-) diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index e1e9694e2ab..a742cc6c4de 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -42,9 +42,14 @@ class DeepSeekChatConfig(OpenAIGPTConfig): Reference: https://api-docs.deepseek.com/guides/thinking_mode """ + params_without_hosted_mcp: Final = tool_params_without_hosted_mcp( + non_default_params, model=model, route="deepseek", drop_params=drop_params + ) # Let parent handle standard params first optional_params = super().map_openai_params( - tool_params_without_hosted_mcp(non_default_params, model=model, route="deepseek", drop_params=drop_params), + params_without_hosted_mcp + if params_without_hosted_mcp is non_default_params + else self._drop_unsupported_tools(non_default_params), optional_params, model, drop_params, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 22f55ea9153..8fbaa7e3a9a 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -13,7 +13,6 @@ from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import TypeAdapter - from redis.exceptions import RedisError from litellm._internal_context import with_service_target diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 9cd5555fe25..66f493bcea3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -18,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( from litellm.proxy._types import UserAPIKeyAuth from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LITELLM_PROXY_MCP_SERVER_URL_PREFIX, + saved_toolset_exists, ) from litellm.types.guardrails import GuardrailEventHooks @@ -104,7 +105,6 @@ class MCPSecurityGuardrail(CustomGuardrail): from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy.proxy_server import prisma_client async with global_mcp_server_manager.catalog.operation(): not_servers: Final = frozenset( @@ -112,12 +112,4 @@ class MCPSecurityGuardrail(CustomGuardrail): for name in requested_servers if global_mcp_server_manager.get_mcp_server_by_identifier(name) is None ) - if prisma_client is None: - return not_servers - return frozenset( - [ - name - for name in not_servers - if await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is None - ] - ) + return frozenset([name for name in not_servers if not await saved_toolset_exists(name)]) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index b711d051055..1dfa4b2abcb 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -105,7 +105,7 @@ def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: ) -async def _toolset_exists(name: str) -> bool: +async def saved_toolset_exists(name: str) -> bool: try: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -124,7 +124,7 @@ async def _toolset_exists(name: str) -> bool: async def _gateway_served_names( names: Collection[str], servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, - toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists, + toolset_exists: Callable[[str], Awaitable[bool]] = saved_toolset_exists, ) -> frozenset[str]: registered: Final = tuple(servers()) if names else () return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) diff --git a/litellm/utils.py b/litellm/utils.py index fea6b9d2177..8cc0cd3ccce 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -438,11 +438,11 @@ if TYPE_CHECKING: ) from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_llm import CustomLLM from litellm.router_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy, reset_retry_policy, ) - from litellm.llms.custom_llm import CustomLLM from litellm.types.llms.anthropic import ( ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index a46c8bed4ad..716c0af236c 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -18,7 +18,6 @@ from litellm.types.utils import ( ChatCompletionMessageToolCall, Delta, Function, - GenericStreamingChunk, ModelResponseStream, PromptTokensDetails, ServerToolUse, diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index ab38f4fc41f..d4775173fe0 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -952,3 +952,33 @@ def test_deepseek_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(mo assert "tools" not in result assert "tool_choice" not in result assert "parallel_tool_calls" not in result + + +def test_deepseek_clears_a_tool_choice_naming_the_hosted_mcp_tool_drop_params_dropped(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + result: Final = DeepSeekChatConfig().map_openai_params( + { + "tools": [_HOSTED_MCP_TOOL, _function_tool("shell")], + "tool_choice": {"type": "function", "function": {"name": "deepwiki"}}, + }, + {}, + "deepseek-v4-flash", + True, + ) + assert result["tools"] == [_function_tool("shell")] + assert "tool_choice" not in result + + +def test_deepseek_removes_tool_selection_beside_legacy_functions_when_drop_params_drops_every_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + functions: Final = [{"name": "shell", "parameters": {"type": "object", "properties": {}}}] + result: Final = DeepSeekChatConfig().map_openai_params( + {"tools": [_HOSTED_MCP_TOOL], "functions": functions, "tool_choice": "required", "parallel_tool_calls": True}, + {}, + "deepseek-v4-flash", + True, + ) + assert result["functions"] == functions + assert "tools" not in result + assert "tool_choice" not in result + assert "parallel_tool_calls" not in result diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 6501526634c..eeaa105565e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -255,3 +255,30 @@ async def test_allows_saved_toolset_by_name_and_blocks_unknown_name(guardrail, c ): await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) assert exc.value.detail["unregistered_servers"] == ["evil_server"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +async def test_blocks_the_name_with_the_policy_result_when_the_toolset_read_fails(guardrail, call_type): + prisma: Final = MagicMock() + prisma.writer_db = prisma.db + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcptoolsettable.find_first = AsyncMock(side_effect=RuntimeError("toolset table unavailable")) + data: Final = { + "tools": [{"type": "mcp", "server_url": "litellm_proxy/mcp/research"}], + "guardrails": ["test-mcp-security"], + } + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + MCPServerManager(), + ), + pytest.raises(HTTPException) as exc, + ): + await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) + assert exc.value.status_code == 400 + assert exc.value.detail["unregistered_servers"] == ["research"]