From b01e66dc58ce80a8efcfa5e5ada9762406ab8bae Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Fri, 27 Feb 2026 21:18:44 -0500
Subject: [PATCH] Add resource limits to MontySandboxProvider

---
 .../experimental/transforms/code_mode.py      | 22 +++++++++++++++----
 .../experimental/transforms/test_code_mode.py | 17 +++++++++++++-
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/fastmcp/experimental/transforms/code_mode.py b/src/fastmcp/experimental/transforms/code_mode.py
index 439457099..826c90f8e 100644
--- a/src/fastmcp/experimental/transforms/code_mode.py
+++ b/src/fastmcp/experimental/transforms/code_mode.py
@@ -68,10 +68,22 @@ class SandboxProvider(Protocol):
 
 
 class MontySandboxProvider:
-    """Sandbox provider backed by `pydantic-monty`."""
+    """Sandbox provider backed by `pydantic-monty`.
 
-    def __init__(self, *, install_hint: str = "fastmcp[code-mode]") -> None:
-        self.install_hint = install_hint
+    Args:
+        limits: Resource limits for sandbox execution. Supported keys:
+            ``max_duration_secs`` (float), ``max_allocations`` (int),
+            ``max_memory`` (int), ``max_recursion_depth`` (int),
+            ``gc_interval`` (int).  All are optional; omit a key to
+            leave that limit uncapped.
+    """
+
+    def __init__(
+        self,
+        *,
+        limits: dict[str, Any] | None = None,
+    ) -> None:
+        self.limits = limits
 
     async def run(
         self,
@@ -85,7 +97,7 @@ class MontySandboxProvider:
         except ModuleNotFoundError as exc:
             raise ImportError(
                 "CodeMode requires pydantic-monty for the Monty sandbox provider. "
-                f"Install it with `{self.install_hint}` or pass a custom SandboxProvider."
+                "Install it with `fastmcp[code-mode]` or pass a custom SandboxProvider."
             ) from exc
 
         inputs = inputs or {}
@@ -102,6 +114,8 @@ class MontySandboxProvider:
         run_kwargs: dict[str, Any] = {"external_functions": async_functions}
         if inputs:
             run_kwargs["inputs"] = inputs
+        if self.limits is not None:
+            run_kwargs["limits"] = self.limits
         return await pydantic_monty.run_monty_async(monty, **run_kwargs)
 
 
diff --git a/tests/experimental/transforms/test_code_mode.py b/tests/experimental/transforms/test_code_mode.py
index 1bad81528..7355b5586 100644
--- a/tests/experimental/transforms/test_code_mode.py
+++ b/tests/experimental/transforms/test_code_mode.py
@@ -335,7 +335,7 @@ async def test_code_mode_execute_non_text_content_stringified() -> None:
 async def test_monty_provider_raises_informative_error_when_missing(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    provider = MontySandboxProvider(install_hint="fastmcp[code-mode]")
+    provider = MontySandboxProvider()
     real_import_module = importlib.import_module
 
     def _fake_import_module(name: str, package: str | None = None):
@@ -446,6 +446,21 @@ async def test_code_mode_sandbox_error_surfaces_as_tool_error() -> None:
         await _run_tool(mcp, "execute", {"code": "raise ValueError('boom')"})
 
 
+async def test_monty_provider_forwards_limits() -> None:
+    """MontySandboxProvider passes limits through to pydantic-monty."""
+    provider = MontySandboxProvider(limits={"max_duration_secs": 0.1})
+
+    with pytest.raises(Exception, match="time limit exceeded"):
+        await provider.run("x = 0\nfor _ in range(10**9):\n    x += 1")
+
+
+async def test_monty_provider_no_limits_by_default() -> None:
+    """Without limits, a simple script completes normally."""
+    provider = MontySandboxProvider()
+    result = await provider.run("return 1 + 2")
+    assert result == 3
+
+
 def test_code_mode_rejects_identical_tool_names() -> None:
     """CodeMode raises ValueError when search and execute names collide."""
     with pytest.raises(ValueError, match="must be different"):