fix(mcp): preserve disabled behavior and return code timeouts

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -721,7 +721,6 @@ class RequestRateLimiterStash:
     """
 
     owner_litellm_call_id: str | None = None
-    borrowed_parallel_keys: frozenset[str] = frozenset()
     rate_limit_response: RateLimitResponse | None = None
     parallel_slot: ParallelSlotAcquisition | None = None
     parallel_slot_reused_at: datetime | None = None
@@ -746,6 +745,7 @@ class RequestRateLimiterStash:
     batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = ()
     reservation_released: bool = False
     tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)
+    borrowed_parallel_keys: frozenset[str] = frozenset()
 
 
 @dataclass(frozen=True, slots=True)

litellm/proxy/mcp_server/README.md

@@ -49,7 +49,7 @@ Keep large identifiers as strings. Values crossing the JavaScript bridge must be
 
 ## Failures and cancellation
 
-Each script has a 120-second deadline, a 64 MiB JavaScript heap limit, and a maximum of 100 tool calls with eight in flight. Each proxy process admits four scripts at once; additional requests receive a busy error. Code is limited to 128 KiB, individual bridge messages to 2 MiB, and returned output to 256 KiB. Use direct calls for tools whose results or duration exceed these limits.
+Each script has a 25-second deadline, a 64 MiB JavaScript heap limit, and a maximum of 100 tool calls with eight in flight. Each proxy process admits four scripts at once; additional requests receive a busy error. Code is limited to 128 KiB, individual bridge messages to 2 MiB, and returned output to 256 KiB. The deadline leaves room for a structured error before the existing MCP transport's 30-second response limit. Authentication or hooks can still consume that transport budget. Use direct calls for tools whose results or duration exceed the code-mode limits.
 
 A script that exceeds a limit fails explicitly. The result includes `output`, `error`, and `calls_started`; the count records attempted tool calls, not confirmed upstream actions. When cancellation reaches the running request, it stops further dispatch and terminates the worker, but cannot undo a request already sent upstream. A client timeout or disconnect does not guarantee server-side cancellation; work can continue until the script deadline. A failed script can therefore have partial effects. Check the upstream state before retrying a write; the gateway does not automatically retry whole scripts.
 

litellm/proxy/mcp_server/code_worker.py

@@ -13,7 +13,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
 
 @dataclass(frozen=True, slots=True)
 class CodeExecutionLimits:
-    timeout_seconds: float = 120
+    timeout_seconds: float = 25  # Leave response time inside the MCP bridge's 30-second header deadline.
     memory_bytes: int = 64 * 1024 * 1024
     script_bytes: int = 128 * 1024
     message_bytes: int = 2 * 1024 * 1024

litellm/proxy/mcp_server/operations.py

@@ -315,7 +315,7 @@ async def _mcp_server_rate_limit_rejection(
 
 async def prepare_virtual_mcp_call(
     name: str,
-    arguments: dict[str, object],
+    arguments: Mapping[str, object],
     user_api_key_auth: UserAPIKeyAuth,
     logging_auth: UserAPIKeyAuth | None = None,
     raw_headers: Mapping[str, str] | None = None,
@@ -365,7 +365,7 @@ async def _build_virtual_call_logging_obj(
     client_ip: str | None = None,
 ) -> LiteLLMLoggingObj | None:
     _, logging_obj = await prepare_virtual_mcp_call(
-        name, dict(arguments), user_api_key_auth, logging_auth, raw_headers, client_ip
+        name, arguments, user_api_key_auth, logging_auth, raw_headers, client_ip
     )
     return logging_obj
 

tests/integration/mcp/test_mcp_code_dispatch.py

@@ -6,7 +6,7 @@ from typing import Final
 import pytest
 import yaml
 from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value
-from integration._support.mcp import McpCaller, mcp_peer, register_mcp, tool_calls
+from integration._support.mcp import McpCaller, mcp_peer, official_client_outcomes, register_mcp, tool_calls
 from integration._support.process import owned_proxy
 
 
@@ -65,3 +65,19 @@ def test_code_calls_chain_with_caller_permissions(code_gateway: Gateway) -> None
         assert refused["isError"] is True
         assert len(tool_calls(allowed.drain())) == 2
         assert tool_calls(denied.drain()) == ()
+
+
+def test_code_timeout_returns_error_and_releases_caller_slot(code_gateway: Gateway) -> None:
+    with code_gateway.scenario() as scenario:
+        key: Final = scenario.key(max_parallel_requests=1)
+        _, result = official_client_outcomes(
+            code_gateway, key, "/mcp", "litellm_execute_code", {"code": "while (true) {}"}
+        )
+        assert result.status == 200 and not result.ok and result.text is not None, result
+        body: Final = JSON_OBJECT.validate_json(result.text)
+        error: Final = body["error"]
+        assert isinstance(error, str) and "timed out" in error, body
+        assert body["calls_started"] == 0
+        recovery: Final = McpCaller(code_gateway, key, "mcp").call("litellm_execute_code", {"code": "text(42)"})
+        assert recovery.ok and recovery.text is not None, recovery
+        assert JSON_OBJECT.validate_json(recovery.text)["output"] == [42]

tests/unit/proxy/hooks/test_mcp_code_mode_admission.py

@@ -187,3 +187,13 @@ async def test_cancelled_outer_does_not_refund_completed_child_usage() -> None:
         with pytest.raises(HTTPException, match="tokens"):
             with isolated_mcp_tool_call():
                 await admit(handler, cache, auth, "completed-child-charge-still-counts")
+
+
+def test_request_stash_preserves_existing_positional_constructor() -> None:
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse
+
+    response = RateLimitResponse(overall_code="OK", statuses=[])
+    stash = RequestRateLimiterStash("original-owner", response)
+    assert stash.owner_litellm_call_id == "original-owner"
+    assert stash.rate_limit_response is response
+    assert stash.borrowed_parallel_keys == frozenset()

tests/unit/proxy/mcp_server/test_mcp_tool_search.py

@@ -1468,4 +1468,73 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No
         if original_affinity is not None:
             os.sched_setaffinity(0, original_affinity)
 
+@pytest.mark.asyncio
+async def test_disabled_code_mode_preserves_legacy_pre_hook_argument_mutation(monkeypatch: pytest.MonkeyPatch) -> None:
+    from mcp.types import TextContent
 
+    from litellm.caching.caching import DualCache
+    from litellm.integrations.custom_logger import CustomLogger
+    from litellm.proxy import proxy_server
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+    from litellm.proxy.mcp_server.tool_registry import global_mcp_tool_registry
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.types.mcp import MCPTransport
+    from litellm.types.mcp_server.mcp_server_manager import MCPServer
+    from litellm.types.utils import CallTypesLiteral
+
+    tool_name = "legacy-rewrite-echo"
+
+    class RewriteVirtualArguments(CustomLogger):
+        async def async_pre_call_hook(
+            self,
+            user_api_key_dict: UserAPIKeyAuth,
+            cache: DualCache,
+            data: dict[str, object],
+            call_type: CallTypesLiteral,
+        ) -> dict[str, object]:
+            # Only the virtual pre-hook has `name`; the later per-tool hook has
+            # `mcp_tool_name`. A second rewrite there would hide a broken handoff.
+            if data.get("name") == tool_name and "mcp_tool_name" not in data:
+                arguments = data.get("arguments")
+                assert isinstance(arguments, dict)
+                arguments["message"] = "rewritten by pre-hook"
+            return data
+
+    async def echo(message: str) -> str:
+        return message
+
+    server = MCPServer(
+        server_id="legacy-rewrite-server",
+        name="legacy-rewrite",
+        alias="legacy-rewrite",
+        transport=MCPTransport.http,
+        allow_all_keys=True,
+    )
+    manager = mcp_operations.global_mcp_server_manager
+    monkeypatch.setattr(litellm, "mcp_code_mode", False, raising=False)
+    monkeypatch.setattr(litellm, "callbacks", [RewriteVirtualArguments()])
+    monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=UserApiKeyCache()))
+    monkeypatch.setitem(manager.registry, server.server_id, server)
+    monkeypatch.setitem(manager.tool_name_to_mcp_server_name_mapping, tool_name, server.name)
+    monkeypatch.setattr(global_mcp_tool_registry, "tools", dict(global_mcp_tool_registry.tools))
+    global_mcp_tool_registry.register_tool(
+        name=tool_name,
+        description="Echo the supplied message",
+        input_schema={"type": "object", "properties": {"message": {"type": "string"}}},
+        handler=echo,
+        server_id=server.server_id,
+    )
+    auth = UserAPIKeyAuth(
+        api_key="legacy-rewrite-test",
+        object_permission=_make_perm(mcp_tool_search_enabled=True),
+    )
+    result = await mcp_operations._dispatch_virtual_mcp_tool(
+        name=MCP_TOOL_CALL_TOOL_NAME,
+        arguments={"tool_name": tool_name, "arguments": {"message": "original message"}},
+        user_api_key_auth=auth,
+        client_ip=None,
+    )
+    assert result is not None
+    assert not result.is_error
+    assert isinstance(result.content[0], TextContent)
+    assert result.content[0].text == "rewritten by pre-hook"