From ce9437b8d75f7991874859b5039760259764acfd Mon Sep 17 00:00:00 2001
From: Isaac Miller <17116851+isaacbmiller@users.noreply.github.com>
Date: Thu, 30 Jul 2026 16:18:53 +0000
Subject: [PATCH] fix(flex): keep execution internals replaceable

Exclude the private bridge from whole-program persistence and rebuild it from canonical source after loading.

Clarify interpreter and predictor-budget contracts so future source lowering and alternate CodeInterpreter implementations remain compatible.
---
 docs/docs/api/modules/Flex.md       | 10 +++++-----
 docs/docs/diving-deeper/flex.md     |  6 +++---
 dspy/flex/flex.py                   | 26 ++++++++++++++++++++++----
 dspy/flex/primitives_doc.py         |  4 +++-
 dspy/primitives/code_interpreter.py |  6 ++++--
 tests/flex/test_flex_binding.py     |  3 +++
 tests/flex/test_flex_gepa.py        |  7 +++++++
 tests/flex/test_flex_interpreter.py |  4 ++--
 8 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/docs/docs/api/modules/Flex.md b/docs/docs/api/modules/Flex.md
index 6e4447b0d3..f5409337a5 100644
--- a/docs/docs/api/modules/Flex.md
+++ b/docs/docs/api/modules/Flex.md
@@ -90,9 +90,9 @@ The `program_trace` parameter is opt-in *by declaration*: only metrics that name
 
 ## Sandboxed Execution
 
-`Flex` always runs its generated code in a sandbox — like `dspy.RLM`, it never runs it in the host Python process. `interpreter_factory` defaults to `dspy.PythonInterpreter` (Deno/Pyodide) and must be a **zero-argument factory** returning a fresh `CodeInterpreter`; a bare instance is not accepted, so each parallel evaluation during optimization gets its own session. The code is authored by the reflection model, so isolating it keeps it from running with your host's full permissions. The optimizer-authored glue — control flow, string work, arithmetic, imports — runs inside the sandbox, and only provided-tool calls, predictor construction, and predictor calls bridge back to the host, which makes the real LM calls.
+`Flex` always runs its generated code in a sandbox — like `dspy.RLM`, it never runs it in the host Python process. `interpreter_factory` defaults to `dspy.PythonInterpreter` (Deno/Pyodide) and must be a **zero-argument factory** returning a fresh `CodeInterpreter`; a bare instance is not accepted, so parallel evaluations receive isolated sessions. The factory is called once per sandbox session, including separate sessions requested by nested code-executing modules. The code is authored by the reflection model, so isolating it keeps it from running with your host's full permissions. With the default interpreter, optimizer-authored control flow, string work, arithmetic, and supported imports run inside the sandbox, and only provided-tool calls, predictor construction, and predictor calls bridge back to the host, which makes the real LM calls.
 
-Because the default builds a `PythonInterpreter`, *running* a `Flex` needs [Deno](https://deno.land/) installed and raises with install instructions otherwise — construction, save, and load are interpreter-free; each call spins up a fresh sandbox and tears it down on return. To customize the sandbox — grant filesystem or network access, or use another `CodeInterpreter` backend — pass your own factory:
+Because the default builds a `PythonInterpreter`, *running* a `Flex` needs [Deno](https://deno.land/) installed and raises with install instructions otherwise — construction, save, and load are interpreter-free; each call owns a fresh outer sandbox and tears it down on return. To customize execution, pass your own factory. This is a low-level hook: Flex may validate or lower source and install its guest shim before execution, while the interpreter determines the Python and standard-library subset actually available. Source optimized for one interpreter is not guaranteed to run unchanged on another.
 
 ```python
 solve = dspy.Flex(
@@ -101,7 +101,7 @@ solve = dspy.Flex(
 )
 ```
 
-There is nothing to clean up: each call creates its own interpreter and shuts it down on return, so a `Flex` holds no live sessions between calls.
+There is nothing to clean up: each call owns and shuts down every interpreter session it creates, so a `Flex` holds no live sessions between calls.
 
 ## Tools
 
@@ -136,8 +136,8 @@ The interpreter, like the LM, is a **runtime dependency and is not serialized**.
 |-----------|------|---------|-------------|
 | `signature` | `str \| Signature` | required | Declares the module's inputs and outputs (e.g. `"invoice -> total_cents: int"`). |
 | `tools` | `list[Callable \| dspy.Tool]` | `None` | Tools the generated code may call. With tools, the baseline is a `dspy.RLM`; without, a `dspy.Predict`. |
-| `interpreter_factory` | `Callable[[], CodeInterpreter]` | `PythonInterpreter` | Zero-arg factory returning the sandbox that runs the generated code; defaults to `dspy.PythonInterpreter` (needs Deno), like `dspy.RLM`. A bare interpreter instance is not accepted. |
-| `max_predictor_calls` | `int` | `100` | Cap on bridged LM calls per `forward` (a runaway guard). `None` disables it. |
+| `interpreter_factory` | `Callable[[], CodeInterpreter]` | `PythonInterpreter` | Zero-arg factory returning a fresh sandbox session; defaults to `dspy.PythonInterpreter` (needs Deno), like `dspy.RLM`. A bare interpreter instance is not accepted. Supported Python and libraries are interpreter-dependent. |
+| `max_predictor_calls` | `int` | `100` | Cap on predictor invocations admitted by the Flex bridge per `forward` (a runaway guard). This is not a count of every internal LM call made by compound modules. `None` disables it. |
 
 ## Notes
 
diff --git a/docs/docs/diving-deeper/flex.md b/docs/docs/diving-deeper/flex.md
index c62cd2e7b3..75c7e10f60 100644
--- a/docs/docs/diving-deeper/flex.md
+++ b/docs/docs/diving-deeper/flex.md
@@ -38,7 +38,7 @@ A code candidate can bind cleanly and still raise mid-`forward` on some inputs 
 
 ### 8. Generated code always runs in an interpreter, never in-process
 
-`Flex` runs `module_src` in a sandbox: the `interpreter_factory` defaults to `dspy.PythonInterpreter` (Deno/Pyodide), matching `dspy.RLM`, and — like `dspy.RLM` — must be a *zero-argument factory* (a bare instance or `None` is rejected). Since the code is authored by the reflection model, isolating it keeps it from running with the host's full permissions: the optimizer-authored glue runs isolated, and only provided-tool calls, predictor construction, and predictor calls bridge back to the host to make real LM calls. The factory is called per `forward` — each call gets a fresh interpreter, shut down on return — so parallel evaluations are isolated by construction; pass your own factory to customize the sandbox (grant filesystem/network access, or use another backend). `max_predictor_calls` caps bridged LM calls per `forward` as a runaway guard.
+`Flex` runs `module_src` in a sandbox: the `interpreter_factory` defaults to `dspy.PythonInterpreter` (Deno/Pyodide), matching `dspy.RLM`, and — like `dspy.RLM` — must be a *zero-argument factory* (a bare instance or `None` is rejected). Since the code is authored by the reflection model, isolating it keeps it from running with the host's full permissions: the optimizer-authored glue runs isolated, and only provided-tool calls, predictor construction, and predictor calls bridge back to the host. The factory creates a fresh interpreter for each sandbox session; a forward owns an outer session and nested code-executing modules may request separate sessions. Flex may validate or lower source and install its guest shim before execution, while each interpreter defines its supported Python and standard-library subset, so source portability between custom interpreters is not automatic. `max_predictor_calls` caps predictor invocations admitted by the Flex bridge per `forward`; compound modules may make additional internal LM calls.
 
 ### 9. The declared output types are enforced at the sandbox boundary
 
@@ -82,10 +82,10 @@ Add `program_trace=None` as a sixth parameter to your metric and GEPA passes the
 Plain functions or `dspy.Tool` instances, referenced by name in the generated code, so each name must be a valid Python identifier. Providing tools makes the baseline a `dspy.RLM` and tells the code proposer the tools are in scope — to wire into `dspy.RLM`/`dspy.ReAct`, call directly, or supplement with its own inline helpers.
 
 **`interpreter_factory=...`**
-Defaults to `dspy.PythonInterpreter` (sandboxed, needs Deno), like `dspy.RLM`. Must be a zero-argument callable returning a fresh `CodeInterpreter`; each parallel evaluation gets its own session. As in `dspy.RLM`, a bare interpreter instance is not accepted — pass a factory.
+Defaults to `dspy.PythonInterpreter` (sandboxed, needs Deno), like `dspy.RLM`. Must be a zero-argument callable returning a fresh `CodeInterpreter` for each sandbox session; parallel evaluations and nested code-executing modules can therefore receive isolated sessions. As in `dspy.RLM`, a bare interpreter instance is not accepted. This low-level hook does not guarantee source or standard-library portability between different interpreters.
 
 **`max_predictor_calls`**
-Caps bridged LM calls per `forward` as a runaway guard. `None` disables it.
+Caps predictor invocations admitted by the Flex bridge per `forward` as a runaway guard. It does not count every internal LM call made by a compound module. `None` disables it.
 
 ### What the generated code can use
 
diff --git a/dspy/flex/flex.py b/dspy/flex/flex.py
index 9553ef57b7..4e600bf85a 100644
--- a/dspy/flex/flex.py
+++ b/dspy/flex/flex.py
@@ -24,15 +24,18 @@ class Flex(Module, Parameter):
     The optimizer-authored code runs inside an interpreter. ``Flex`` never
     runs it in the host Python process. ``interpreter_factory`` defaults to ``dspy.PythonInterpreter``
     (Deno/Pyodide) and must be a zero-argument callable returning a new ``CodeInterpreter``.
+    Flex may validate or lower source and install its guest shim before execution; a custom
+    interpreter therefore defines the Python and standard-library subset available to that source.
     The optimizer-authored glue runs isolated; only provided-tool calls, predictor construction,
     and predictor calls bridge back to the host, which makes the real LM calls.
 
     Args:
         signature: A ``dspy.Signature`` class or string declaring inputs/outputs.
         tools: ``dspy.Tool`` instances or named callables.
-        interpreter_factory: Zero-argument callable returning a ``CodeInterpreter`` for each forward
-            pass. Defaults to ``dspy.PythonInterpreter`` (sandbox, requires Deno).
-        max_predictor_calls: Cap on bridged LM calls per ``forward``; ``None`` disables it.
+        interpreter_factory: Zero-argument callable returning a fresh ``CodeInterpreter`` for each
+            sandbox session. Defaults to ``dspy.PythonInterpreter`` (sandbox, requires Deno).
+        max_predictor_calls: Cap on predictor invocations admitted by the Flex bridge per
+            ``forward``; ``None`` disables it.
     """
 
     def __init__(
@@ -58,10 +61,25 @@ def __init__(
         self._interpreter_factory = interpreter_factory
         self._max_predictor_calls = max_predictor_calls
 
+        self._rebuild_bridge()
+        self._bind_code(self._baseline_src())
+
+    def _rebuild_bridge(self) -> None:
         from dspy.flex.bridge import BridgeRuntime
 
         self._bridge: BridgeRuntime = BridgeRuntime(self, self._interpreter_factory, self._max_predictor_calls)
-        self._bind_code(self._baseline_src())
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = super().__getstate__()
+        state.pop("_bridge", None)
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        state.pop("_bridge", None)  # Ignore bridges persisted by prerelease versions.
+        super().__setstate__(state)
+        self._rebuild_bridge()
+        if self._module_src is not None:
+            self._bridge.bind(self._module_src)
 
     @property
     def signature(self) -> Any:
diff --git a/dspy/flex/primitives_doc.py b/dspy/flex/primitives_doc.py
index 64adcb620a..17d8f79566 100644
--- a/dspy/flex/primitives_doc.py
+++ b/dspy/flex/primitives_doc.py
@@ -34,7 +34,9 @@ def forward(self, **inputs):
 ```
 
 `dspy` is already in scope — do NOT add `import` statements at module scope (import
-stdlib modules like `re` *inside* `forward` if you need them).
+stdlib modules like `re` *inside* `forward` if you need them). Python and standard-library
+availability is defined by this Flex's configured interpreter; candidates are evaluated in that
+interpreter and should not assume host-Python access.
 
 Available DSPy primitives:
 
diff --git a/dspy/primitives/code_interpreter.py b/dspy/primitives/code_interpreter.py
index ce43810aef..6e44a19eba 100644
--- a/dspy/primitives/code_interpreter.py
+++ b/dspy/primitives/code_interpreter.py
@@ -74,11 +74,13 @@ class CodeInterpreter(Protocol):
     """
 
     @property
-    def tools(self) -> dict[str, Callable[..., str]]:
+    def tools(self) -> dict[str, Callable[..., Any]]:
         """Tools available for interpreter code to call.
 
         Tools are host-side functions that can be invoked from within the
-        interpreter. Each tool accepts keyword arguments and returns a string.
+        interpreter. Each tool accepts keyword arguments. Return values must
+        satisfy the boundary supported by the interpreter; Flex tools must
+        return JSON-compatible values.
 
         Implementations should accept tools via constructor and expose them
         through this property.
diff --git a/tests/flex/test_flex_binding.py b/tests/flex/test_flex_binding.py
index a945ee0b66..f35f50e4b9 100644
--- a/tests/flex/test_flex_binding.py
+++ b/tests/flex/test_flex_binding.py
@@ -283,10 +283,13 @@ def forward(self, **kwargs):
 
     program = Program()
     program.flex._bind_code(ECHO_MODULE)
+    assert "_bridge" not in program.flex.__getstate__()
     program.save(tmp_path, save_program=True)
 
     loaded = dspy.load(str(tmp_path), allow_pickle=True)
     assert loaded.flex.module_src == ECHO_MODULE
+    assert loaded.flex._bridge._flex is loaded.flex
+    assert loaded.flex._bridge._module_src == ECHO_MODULE
 
 
 @deno_required
diff --git a/tests/flex/test_flex_gepa.py b/tests/flex/test_flex_gepa.py
index ac108d4b9e..49fc0e9689 100644
--- a/tests/flex/test_flex_gepa.py
+++ b/tests/flex/test_flex_gepa.py
@@ -13,6 +13,7 @@
 
 from __future__ import annotations
 
+import shutil
 import textwrap
 
 import pytest
@@ -24,6 +25,8 @@
 from dspy.utils.dummies import DummyLM
 from dspy.utils.exceptions import LMRateLimitError
 
+deno_required = pytest.mark.skipif(shutil.which("deno") is None, reason="Deno is not installed")
+
 # A plain dspy.Predict module class that binds without an LM (no RLM interpreter needed).
 SIMPLE_MODULE = textwrap.dedent(
     """
@@ -145,6 +148,7 @@ def test_build_program_rebinds_flex_code() -> None:
 # --- adapter: selection eval passes the trace to a flex metric ---------------
 
 
+@deno_required
 def test_selection_eval_passes_program_trace_to_declaring_metric() -> None:
     """The selection eval (capture_traces=False) must pass the execution trace to a metric that
     declares `program_trace`, so a trace-dependent score (e.g. an LLM-call penalty that rewards
@@ -168,6 +172,7 @@ def trace_aware_metric(gold, pred, trace=None, pred_name=None, pred_trace=None,
     assert seen.get("trace") is None  # eval-mode semantics of the `trace` argument are preserved
 
 
+@deno_required
 def test_selection_eval_keeps_vanilla_semantics_for_legacy_metric() -> None:
     seen: dict[str, object] = {}
 
@@ -189,6 +194,7 @@ def legacy_metric(gold, pred, trace=None, pred_name=None, pred_trace=None):
     assert batch.scores == [0.75]
 
 
+@deno_required
 def test_selection_eval_binds_metric_with_required_contract_params() -> None:
     """A metric written to the full GEPAFeedbackMetric signature with `trace`/`pred_name`/
     `pred_trace` REQUIRED (no defaults) must still bind at flex scoring time: flex passes those
@@ -231,6 +237,7 @@ def forward(self, **inputs):
 ).strip()
 
 
+@deno_required
 def test_evaluate_scores_stay_aligned_when_an_example_crashes() -> None:
     """A code candidate that binds fine but raises at runtime on one input must still score
     one-per-example, aligned by position — the gepa engine pairs scores with example ids
diff --git a/tests/flex/test_flex_interpreter.py b/tests/flex/test_flex_interpreter.py
index be912d02fa..26a87c34bb 100644
--- a/tests/flex/test_flex_interpreter.py
+++ b/tests/flex/test_flex_interpreter.py
@@ -745,7 +745,7 @@ class ProbeSig(dspy.Signature):
         )
     )
 
-    MODULE = textwrap.dedent(
+    module = textwrap.dedent(
         """
         class ProbeModule(dspy.Module):
             def __init__(self):
@@ -765,7 +765,7 @@ def forward(self, **inputs):
     ).strip()
 
     flex = Flex(ProbeSig, tools=[probe, spawn_inner], interpreter_factory=lambda: dspy.PythonInterpreter())
-    flex._bind_code(MODULE)
+    flex._bind_code(module)
     flex(task="go")
 
     # Every layer ran inside a Pyodide sandbox, not the host process: