feat(code-mode): default sandbox limits and per-execution tool-call cap (#4170)

Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>

docs/servers/transforms/code-mode.mdx

@@ -285,7 +285,17 @@ mcp = FastMCP("Server", transforms=[code_mode])
 
 ### Resource Limits
 
-The default `MontySandboxProvider` can enforce execution limits — timeouts, memory caps, recursion depth, and more. Without limits, LLM-generated scripts can run indefinitely.
+The default `MontySandboxProvider` enforces execution limits — timeouts, memory caps, recursion depth, and more.
+
+Constructed with no arguments, it applies a conservative baseline so the out-of-box configuration is not unbounded: `max_duration_secs=30` and `max_memory=100_000_000` (100 MB). Pass an explicit `limits` dict to override it, or `limits=None` to run with no limits at all:
+
+```python
+from fastmcp.experimental.transforms.code_mode import MontySandboxProvider
+
+MontySandboxProvider()                  # baseline: 30s, 100 MB
+MontySandboxProvider(limits={...})      # your own limits
+MontySandboxProvider(limits=None)       # explicitly uncapped
+```
 
 ```python
 from fastmcp.experimental.transforms.code_mode import CodeMode
@@ -308,6 +318,18 @@ All keys are optional — omit any to leave that dimension uncapped:
 | `max_recursion_depth` | `int` | Maximum recursion depth |
 | `gc_interval` | `int` | Garbage collection frequency |
 
+### Tool Call Limits
+
+A single `execute` block can issue many `call_tool()` invocations — a loop in LLM-generated code can fan out into a large number of backend operations from one request. `CodeMode` caps this at `max_tool_calls` (default `50`); exceeding it raises a `ToolError`. Pass `None` for no cap:
+
+```python
+from fastmcp.experimental.transforms.code_mode import CodeMode
+
+CodeMode()                      # default: 50 call_tool() calls per execute()
+CodeMode(max_tool_calls=200)    # raise the cap
+CodeMode(max_tool_calls=None)   # no cap
+```
+
 ### Custom Sandbox Providers
 
 You can replace the default sandbox with any object implementing the `SandboxProvider` protocol:

fastmcp_slim/fastmcp/experimental/transforms/code_mode.py

@@ -10,7 +10,7 @@ if TYPE_CHECKING:
 from mcp.types import TextContent
 from pydantic import Field
 
-from fastmcp.exceptions import NotFoundError
+from fastmcp.exceptions import NotFoundError, ToolError
 from fastmcp.server.context import Context
 from fastmcp.server.transforms import GetToolNext
 from fastmcp.server.transforms.catalog import CatalogTransform
@@ -92,6 +92,25 @@ class SandboxProvider(Protocol):
     ) -> Any: ...
 
 
+class _UnsetType:
+    """Sentinel distinguishing "argument omitted" from an explicit value."""
+
+    def __repr__(self) -> str:
+        return "UNSET"
+
+
+_UNSET = _UnsetType()
+
+
+_DEFAULT_LIMITS: "ResourceLimits" = {
+    "max_duration_secs": 30.0,
+    "max_memory": 100_000_000,  # 100 MB
+}
+"""Baseline limits applied when ``MontySandboxProvider`` is constructed
+without an explicit ``limits`` argument. Pass ``limits=None`` to opt out
+entirely, or a dict to override."""
+
+
 class MontySandboxProvider:
     """Sandbox provider backed by `pydantic-monty`.
 
@@ -101,14 +120,25 @@ class MontySandboxProvider:
             ``max_memory`` (int), ``max_recursion_depth`` (int),
             ``gc_interval`` (int).  All are optional; omit a key to
             leave that limit uncapped.
+
+            When the argument is omitted entirely, a conservative baseline
+            is applied (``max_duration_secs=30``, ``max_memory=100 MB``) so
+            the out-of-box configuration is not unbounded. Pass
+            ``limits=None`` to explicitly run without any limits, or a dict
+            to set your own.
     """
 
     def __init__(
         self,
         *,
-        limits: "ResourceLimits | None" = None,
+        limits: "ResourceLimits | None | _UnsetType" = _UNSET,
     ) -> None:
-        self.limits = limits
+        # Copy the baseline so each provider owns its dict — `limits` is a
+        # mutable public attribute, and sharing the module-level object would
+        # let one provider's edits leak into every other default provider.
+        self.limits: ResourceLimits | None = (
+            _DEFAULT_LIMITS.copy() if isinstance(limits, _UnsetType) else limits
+        )
 
     async def run(
         self,
@@ -484,10 +514,12 @@ class CodeMode(CatalogTransform):
         discovery_tools: list[DiscoveryToolFactory] | None = None,
         execute_tool_name: str = "execute",
         execute_description: str | None = None,
+        max_tool_calls: int | None = 50,
     ) -> None:
         super().__init__()
         self.execute_tool_name = execute_tool_name
         self.execute_description = execute_description
+        self.max_tool_calls = max_tool_calls
         self.sandbox_provider = sandbox_provider or MontySandboxProvider()
 
         self._discovery_factories = (
@@ -556,6 +588,7 @@ class CodeMode(CatalogTransform):
 
     def _make_execute_tool(self) -> Tool:
         transform = self
+        max_tool_calls = self.max_tool_calls
 
         async def execute(
             code: Annotated[
@@ -570,7 +603,18 @@ class CodeMode(CatalogTransform):
         ) -> Any:
             """Execute tool calls using Python code."""
 
+            call_count = 0
+
             async def call_tool(tool_name: str, params: dict[str, Any]) -> Any:
+                nonlocal call_count
+                if max_tool_calls is not None:
+                    call_count += 1
+                    if call_count > max_tool_calls:
+                        raise ToolError(
+                            f"Tool call limit exceeded: at most {max_tool_calls} "
+                            "call_tool() invocations are allowed per execute()."
+                        )
+
                 backend_tools = await transform.get_tool_catalog(ctx)
                 tool = transform._find_tool(tool_name, backend_tools)
                 if tool is None:

tests/experimental/transforms/test_code_mode.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib
 import json
 from typing import Any
@@ -8,6 +9,7 @@ from mcp.types import ImageContent, TextContent
 from fastmcp import Client, FastMCP
 from fastmcp.exceptions import ToolError
 from fastmcp.experimental.transforms.code_mode import (
+    _DEFAULT_LIMITS,
     CodeMode,
     GetSchemas,
     GetToolCatalog,
@@ -711,12 +713,97 @@ async def test_monty_provider_forwards_limits() -> None:
         await provider.run("x = 0\nfor _ in range(10**9):\n    x += 1")
 
 
-async def test_monty_provider_no_limits_by_default() -> None:
+async def test_monty_provider_applies_default_limits() -> None:
     provider = MontySandboxProvider()
+    assert provider.limits == _DEFAULT_LIMITS
+    # Default limits are generous enough for ordinary code.
     result = await provider.run("return 1 + 2")
     assert result == 3
 
 
+async def test_monty_provider_explicit_none_disables_limits() -> None:
+    provider = MontySandboxProvider(limits=None)
+    assert provider.limits is None
+    result = await provider.run("return 1 + 2")
+    assert result == 3
+
+
+async def test_monty_provider_explicit_limits_override_defaults() -> None:
+    provider = MontySandboxProvider(limits={"max_duration_secs": 0.1})
+    assert provider.limits == {"max_duration_secs": 0.1}
+
+
+async def test_monty_provider_default_limits_are_not_shared_between_instances() -> None:
+    """Each default provider must own its limits dict.
+
+    `limits` is a mutable public attribute; if instances shared the
+    module-level baseline, mutating one would silently change the defaults
+    for every other default provider in the process.
+    """
+    a = MontySandboxProvider()
+    b = MontySandboxProvider()
+
+    assert a.limits is not b.limits
+    assert a.limits is not _DEFAULT_LIMITS
+
+    assert a.limits is not None
+    a.limits["max_duration_secs"] = 1
+
+    assert b.limits == {"max_duration_secs": 30.0, "max_memory": 100_000_000}
+    assert _DEFAULT_LIMITS == {"max_duration_secs": 30.0, "max_memory": 100_000_000}
+
+
+async def test_code_mode_max_tool_calls_default_is_50() -> None:
+    assert CodeMode().max_tool_calls == 50
+
+
+async def test_code_mode_max_tool_calls_enforced() -> None:
+    mcp = FastMCP("CodeMode ToolCap")
+
+    @mcp.tool
+    def ping() -> str:
+        return "pong"
+
+    mcp.add_transform(
+        CodeMode(sandbox_provider=_UnsafeTestSandboxProvider(), max_tool_calls=3)
+    )
+
+    code = "\n".join(
+        [
+            "results = []",
+            "for _ in range(5):",
+            "    results.append(await call_tool('ping', {}))",
+            "return results",
+        ]
+    )
+    with pytest.raises(ToolError, match=r"Tool call limit exceeded: at most 3"):
+        await _run_tool(mcp, "execute", {"code": code})
+
+
+async def test_code_mode_max_tool_calls_none_is_unlimited() -> None:
+    mcp = FastMCP("CodeMode ToolCapNone")
+
+    @mcp.tool
+    def ping() -> str:
+        return "pong"
+
+    mcp.add_transform(
+        CodeMode(sandbox_provider=_UnsafeTestSandboxProvider(), max_tool_calls=None)
+    )
+
+    code = "\n".join(
+        [
+            "n = 0",
+            "for _ in range(60):",
+            "    await call_tool('ping', {})",
+            "    n += 1",
+            "return n",
+        ]
+    )
+    result = await _run_tool(mcp, "execute", {"code": code})
+    assert _unwrap_result(result) == 60
+
+
 async def test_monty_provider_cancels_future_when_task_cancelled() -> None:
     """Cancelling the awaiting task must cancel the underlying sandbox future.
 
@@ -725,8 +812,6 @@ async def test_monty_provider_cancels_future_when_task_cancelled() -> None:
     launch seam so the cancellation handling in `run()` is exercised against
     a controllable future rather than a live sandbox thread.
     """
-    import asyncio
-
     loop = asyncio.get_running_loop()
     sandbox_future: asyncio.Future[Any] = loop.create_future()