diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7b7e03e05b3..9b906a4b96e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/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) diff --git a/litellm/proxy/mcp_server/README.md b/litellm/proxy/mcp_server/README.md index f2485cea549..e57231eee74 100644 --- a/litellm/proxy/mcp_server/README.md +++ b/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. diff --git a/litellm/proxy/mcp_server/code_worker.py b/litellm/proxy/mcp_server/code_worker.py index bedf5a37f6a..78b3f3b9d4e 100644 --- a/litellm/proxy/mcp_server/code_worker.py +++ b/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 diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index a7e71540198..95c0d6162e4 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/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 diff --git a/tests/integration/mcp/test_mcp_code_dispatch.py b/tests/integration/mcp/test_mcp_code_dispatch.py index 369e02f8edc..332a2211e39 100644 --- a/tests/integration/mcp/test_mcp_code_dispatch.py +++ b/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] diff --git a/tests/unit/proxy/hooks/test_mcp_code_mode_admission.py b/tests/unit/proxy/hooks/test_mcp_code_mode_admission.py index cce9244d0c5..22502aa844b 100644 --- a/tests/unit/proxy/hooks/test_mcp_code_mode_admission.py +++ b/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() diff --git a/tests/unit/proxy/mcp_server/test_mcp_tool_search.py b/tests/unit/proxy/mcp_server/test_mcp_tool_search.py index b21202ba227..b3907c2981e 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_tool_search.py +++ b/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"