fix(code-mode): cancel Monty sandbox future on task cancellation (#4169)

fastmcp_slim/fastmcp/experimental/transforms/code_mode.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib
 import json
 from collections.abc import Awaitable, Callable, Sequence
@@ -131,9 +132,38 @@ class MontySandboxProvider:
         }
 
         monty = pydantic_monty.Monty(code, inputs=list(inputs))
-        return await monty.run_async(
-            inputs=inputs or None,
-            external_functions=async_functions or None,
+        future = asyncio.ensure_future(
+            self._run_monty(
+                monty,
+                inputs=inputs or None,
+                external_functions=async_functions or None,
+            )
+        )
+        try:
+            return await future
+        except asyncio.CancelledError:
+            # Awaiting alone does not stop the native sandbox thread when the
+            # surrounding task is cancelled (e.g. an HTTP client disconnects
+            # mid-execution). Explicitly cancel so the Monty runtime tears the
+            # thread down instead of leaving it running to completion.
+            future.cancel()
+            raise
+
+    def _run_monty(
+        self,
+        monty: Any,
+        *,
+        inputs: dict[str, Any] | None,
+        external_functions: dict[str, Callable[..., Any]] | None,
+    ) -> Any:
+        """Launch the sandbox and return its awaitable.
+
+        Isolated so the cancellation handling in `run()` can be exercised
+        without a live `pydantic-monty` runtime.
+        """
+        return monty.run_async(
+            inputs=inputs,
+            external_functions=external_functions,
             limits=self.limits,
         )
 

tests/experimental/transforms/test_code_mode.py

@@ -715,3 +715,36 @@ async def test_monty_provider_no_limits_by_default() -> None:
     provider = MontySandboxProvider()
     result = await provider.run("return 1 + 2")
     assert result == 3
+
+
+async def test_monty_provider_cancels_future_when_task_cancelled() -> None:
+    """Cancelling the awaiting task must cancel the underlying sandbox future.
+
+    Otherwise the native Monty thread keeps running to completion after a
+    client disconnects or the request times out. A subclass overrides the
+    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()
+
+    class _NeverFinishingProvider(MontySandboxProvider):
+        def _run_monty(self, monty: Any, *, inputs: Any, external_functions: Any):
+            return sandbox_future
+
+    provider = _NeverFinishingProvider()
+    task = asyncio.create_task(provider.run("return 1"))
+
+    # Advance the task to `await future` (no suspension point before it).
+    for _ in range(3):
+        await asyncio.sleep(0)
+        if not task.done():
+            break
+
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert sandbox_future.cancelled()