diff --git a/docs/servers/transforms/code-mode.mdx b/docs/servers/transforms/code-mode.mdx index f09c22b6f..c7ef55bf0 100644 --- a/docs/servers/transforms/code-mode.mdx +++ b/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: diff --git a/fastmcp_slim/fastmcp/experimental/transforms/code_mode.py b/fastmcp_slim/fastmcp/experimental/transforms/code_mode.py index cbdfe6d3d..e9ea26b72 100644 --- a/fastmcp_slim/fastmcp/experimental/transforms/code_mode.py +++ b/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: diff --git a/tests/experimental/transforms/test_code_mode.py b/tests/experimental/transforms/test_code_mode.py index dc1da63b2..952b3a475 100644 --- a/tests/experimental/transforms/test_code_mode.py +++ b/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()