Harden interpreter rebinding validation

Co-authored-by: Isaac Miller <isaac@cmpnd.ai>
Amp-Thread-ID: https://ampcode.com/threads/T-019fcce1-bc8b-753e-886f-9b77c0ccb376

Author: ampagent

dspy/primitives/python_interpreter.py

@@ -22,7 +22,7 @@
 import threading
 from collections.abc import Mapping, Sequence
 from os import PathLike
-from typing import Any, Callable, NoReturn
+from typing import Any, Callable, ClassVar, NoReturn
 
 from pydantic import BaseModel
 from pydantic_core import PydanticSerializationError, to_jsonable_python
@@ -241,6 +241,22 @@ def my_tool(question: str) -> str:
         ```
     """
 
+    _RESERVED_TOOL_NAMES: ClassVar[set[str]] = {
+        "SUBMIT",
+        "FinalOutput",
+        "JsProxy",
+        "_js_tool_call",
+        "buf_stderr",
+        "buf_stdout",
+        "io",
+        "json",
+        "last_exception_args",
+        "old_stderr",
+        "old_stdout",
+        "run_sync",
+        "sys",
+    }
+
     def __init__(
         self,
         deno_command: list[str] | None = None,
@@ -286,6 +302,7 @@ def __init__(
         self.enable_network_access = enable_network_access or []
         self.sync_files = sync_files
         self.tools = {}
+        self._tool_parameters = {}
         self.output_fields = None
         self.callbacks = list(callbacks or [])
         self._tools_registered = True
@@ -302,8 +319,8 @@ def __init__(
         self._session_ended = False
         self._state_lock = threading.Lock()
         self._execution_active = False
-        if tools or output_fields:
-            self.bind(tools=tools or {}, output_fields=output_fields)
+        if tools is not None or output_fields is not None:
+            self.bind(tools=tools if tools is not None else {}, output_fields=output_fields)
 
     @property
     def execution_instructions(self) -> str:
@@ -433,14 +450,25 @@ def bind(
             raise TypeError(f"tools must be a mapping, not {type(tools).__name__}")
 
         copied_tools: dict[str, Callable[..., Any]] = {}
+        copied_tool_parameters = {}
         for name, tool in tools.items():
             if not isinstance(name, str) or not name.isidentifier() or keyword.iskeyword(name):
                 raise CodeInterpreterError(f"Invalid tool name: {name!r}")
-            if name in {"SUBMIT", "FinalOutput"}:
+            if name in self._RESERVED_TOOL_NAMES:
                 raise CodeInterpreterError(f"Tool name {name!r} is reserved by the interpreter")
             if not callable(tool):
                 raise CodeInterpreterError(f"Tool {name!r} must be callable, not {type(tool).__name__}")
+            previous = self._tool_parameters.get(name)
+            if previous is not None and previous[0] is tool:
+                parameters = previous[1]
+            else:
+                try:
+                    parameters = self._extract_parameters(tool)
+                    parameters = json.loads(json.dumps(parameters, allow_nan=False))
+                except (TypeError, ValueError) as e:
+                    raise CodeInterpreterError(f"Tool {name!r} has an unsupported signature: {e}") from e
             copied_tools[name] = tool
+            copied_tool_parameters[name] = (tool, parameters)
 
         copied_output_fields = None
         if output_fields is not None:
@@ -471,6 +499,7 @@ def bind(
             if self._execution_active:
                 raise RuntimeError("Cannot bind tools while PythonInterpreter execution is active.")
             self.tools = copied_tools
+            self._tool_parameters = copied_tool_parameters
             self.output_fields = copied_output_fields
             self._tools_registered = False
 
@@ -565,9 +594,10 @@ def _register_tools(self) -> None:
         # after a previous typed binding.
         tools_info = []
         for name, fn in self.tools.items():
+            cached = self._tool_parameters.get(name)
             tools_info.append({
                 "name": name,
-                "parameters": self._extract_parameters(fn)
+                "parameters": cached[1] if cached is not None and cached[0] is fn else self._extract_parameters(fn)
             })
         params = {"tools": tools_info, "outputs": self.output_fields or []}
 

dspy/primitives/runner.js

@@ -35,10 +35,8 @@ if 'SUBMIT' not in dir():
 // Parameters is an array of {name, type?, default?} objects.
 // Convert a JavaScript/JSON value to Python literal syntax
 const toPythonLiteral = (value) => {
-  if (value === null) return 'None';
-  if (value === true) return 'True';
-  if (value === false) return 'False';
-  return JSON.stringify(value);  // Works for strings, numbers, arrays, objects
+  const serialized = JSON.stringify(value);
+  return `__import__("json").loads(${JSON.stringify(serialized)})`;
 };
 
 const toPythonType = (type) => type === "NoneType" ? "type(None)" : type;
@@ -67,6 +65,30 @@ def ${toolName}(${signature}):
 `;
 };
 
+const registeredToolNames = new Set(), shadowedGlobals = new Map();
+
+const replaceToolWrappers = (tools) => {
+  const nextNames = new Set(tools.map(tool => typeof tool === 'string' ? tool : tool.name));
+  for (const name of registeredToolNames) {
+    if (!nextNames.has(name)) {
+      const shadowed = shadowedGlobals.get(name);
+      if (shadowed.exists) pyodide.globals.set(name, shadowed.value);
+      else pyodide.globals.delete(name);
+      shadowedGlobals.delete(name);
+    }
+  }
+  for (const tool of tools) {
+    const name = typeof tool === 'string' ? tool : tool.name;
+    if (!registeredToolNames.has(name)) {
+      const exists = pyodide.globals.has(name);
+      shadowedGlobals.set(name, { exists, value: exists ? pyodide.globals.get(name) : undefined });
+    }
+    pyodide.runPython(makeToolWrapper(name, typeof tool === 'string' ? [] : tool.parameters || []));
+  }
+  registeredToolNames.clear();
+  for (const name of nextNames) registeredToolNames.add(name);
+};
+
 // Generate SUBMIT function with output field signature.
 // Outputs is an array of {name, type?} objects.
 const makeSubmitWrapper = (outputs) => {
@@ -290,21 +312,9 @@ while (true) {
   }
 
   if (method === "register") {
-    const toolNames = [];
-
-    // Register tools with typed signatures
-    if (params.tools) {
-      for (const tool of params.tools) {
-        // Support both old format (string) and new format (object with parameters)
-        if (typeof tool === 'string') {
-          pyodide.runPython(makeToolWrapper(tool, []));
-          toolNames.push(tool);
-        } else {
-          pyodide.runPython(makeToolWrapper(tool.name, tool.parameters || []));
-          toolNames.push(tool.name);
-        }
-      }
-    }
+    const tools = params.tools || [];
+    replaceToolWrappers(tools);
+    const toolNames = tools.map(tool => typeof tool === 'string' ? tool : tool.name);
 
     // Register SUBMIT with output signature
     if (params.outputs) {

tests/predict/test_rlm.py

@@ -7,6 +7,7 @@
 """
 
 import base64
+import inspect
 from contextlib import contextmanager
 
 import pytest
@@ -1492,18 +1493,53 @@ def accepts_none(value: type(None)) -> str:
         assert pooled_interpreter.execute("state + 1") == 42
         assert pooled_interpreter.deno_process is process
 
+    def test_restores_caller_namespace_shadowed_by_invocation_tools(self, pooled_interpreter):
+        pooled_interpreter.execute("llm_query = 'original query'\nllm_query_batched = 'original batch'")
+        process = pooled_interpreter.deno_process
+        rlm = RLM("query -> answer", max_iters=1)
+        rlm.generate_action = make_mock_predictor([
+            {"reasoning": "Submit", "code": "SUBMIT(answer='done')"},
+        ])
+
+        assert rlm(pooled_interpreter, query="test").answer == "done"
+        assert pooled_interpreter.execute("(llm_query, llm_query_batched)") == ["original query", "original batch"]
+        assert pooled_interpreter.deno_process is process
+
+    def test_reuses_validated_metadata_when_restoring_caller_tools(self, pooled_interpreter):
+        class StatefulTool:
+            reads = 0
+
+            @property
+            def __signature__(self):
+                self.reads += 1
+                return inspect.Signature() if self.reads == 1 else "malformed"
+
+            def __call__(self):
+                return "baseline"
+
+        baseline = StatefulTool()
+        pooled_interpreter.bind(tools={"baseline": baseline})
+        rlm = RLM("query -> answer", max_iters=1)
+        rlm.generate_action = make_mock_predictor([
+            {"reasoning": "Submit", "code": "SUBMIT(answer='done')"},
+        ])
+
+        assert rlm(pooled_interpreter, query="test").answer == "done"
+        assert pooled_interpreter.execute("baseline()") == "baseline"
+        assert baseline.reads == 1
+
     def test_reserved_tool_rejection_preserves_caller_owned_interpreter(self, pooled_interpreter):
         def baseline():
             return "baseline"
 
-        def FinalOutput():
+        def reserved_tool():
             return "unreachable"
 
         previous_outputs = [{"name": "previous", "type": "str"}]
         pooled_interpreter.bind(tools={"baseline": baseline}, output_fields=previous_outputs)
         pooled_interpreter.execute("state = 41")
         process = pooled_interpreter.deno_process
-        rlm = RLM("query -> answer", tools=[FinalOutput], max_iters=1)
+        rlm = RLM("query -> answer", tools=[Tool(reserved_tool, name="FinalOutput")], max_iters=1)
 
         with pytest.raises(CodeInterpreterError, match="'FinalOutput' is reserved"):
             rlm(pooled_interpreter, query="test")

tests/primitives/test_python_interpreter.py

@@ -1,5 +1,6 @@
 import asyncio
 import dataclasses
+import inspect
 import io
 import json
 import os
@@ -736,6 +737,72 @@ def accepts_none(value: type(None)) -> str:
     assert pooled_interpreter.deno_process is process
 
 
+def test_tool_with_json_default_preserves_session(pooled_interpreter):
+    def tool(value: list = [None]):  # noqa: B006 - exercises JSON-compatible mutable defaults
+        return value
+
+    pooled_interpreter.execute("state = 41")
+    process = pooled_interpreter.deno_process
+    pooled_interpreter.bind(tools={"tool": tool})
+
+    assert pooled_interpreter.execute("tool()") == [None]
+    assert pooled_interpreter.execute("state + 1") == 42
+    assert pooled_interpreter.deno_process is process
+
+
+def test_unsupported_tool_signature_is_rejected_without_changing_session(pooled_interpreter):
+    unsupported_default = object()
+
+    def baseline():
+        return "baseline"
+
+    def unsupported(value=unsupported_default):
+        return value
+
+    pooled_interpreter.bind(tools={"baseline": baseline})
+    pooled_interpreter.execute("state = 41")
+    process = pooled_interpreter.deno_process
+
+    with pytest.raises(CodeInterpreterError, match="unsupported signature"):
+        pooled_interpreter.bind(tools={"unsupported": unsupported})
+
+    assert pooled_interpreter.tools == {"baseline": baseline}
+    assert pooled_interpreter.execute("baseline()") == "baseline"
+    assert pooled_interpreter.execute("state + 1") == 42
+    assert pooled_interpreter.deno_process is process
+
+
+def test_bind_snapshots_validated_tool_defaults(pooled_interpreter):
+    default = [None]
+
+    def tool(value=default):
+        return value
+
+    pooled_interpreter.bind(tools={"tool": tool})
+    default.append(object())
+
+    assert pooled_interpreter.execute("tool()") == [None]
+
+
+def test_bind_does_not_reinspect_tool_signature(pooled_interpreter):
+    class StatefulTool:
+        reads = 0
+
+        @property
+        def __signature__(self):
+            self.reads += 1
+            return inspect.Signature() if self.reads == 1 else "malformed"
+
+        def __call__(self):
+            return "ok"
+
+    tool = StatefulTool()
+    pooled_interpreter.bind(tools={"tool": tool})
+
+    assert pooled_interpreter.execute("tool()") == "ok"
+    assert tool.reads == 1
+
+
 def test_bind_after_start_replaces_tools_and_preserves_namespace(pooled_interpreter):
     old_calls = []
 
@@ -755,6 +822,18 @@ def old_tool():
     assert old_calls == [True]
 
 
+def test_bind_restores_shadowed_globals_and_removes_stale_wrappers(pooled_interpreter):
+    pooled_interpreter.execute("tool = 'original'")
+    pooled_interpreter.bind(tools={"tool": lambda: "host", "ephemeral": lambda: "temporary"})
+    assert pooled_interpreter.execute("tool()") == "host"
+    assert pooled_interpreter.execute("ephemeral()") == "temporary"
+
+    pooled_interpreter.bind(tools={})
+
+    assert pooled_interpreter.execute("tool") == "original"
+    assert pooled_interpreter.execute("'ephemeral' in globals()") is False
+
+
 def test_bind_none_output_fields_restores_default_submit(pooled_interpreter):
     pooled_interpreter.bind(tools={}, output_fields=[{"name": "answer", "type": "str"}])
     assert pooled_interpreter.execute("SUBMIT(answer='typed')") == FinalOutput({"answer": "typed"})

tests/primitives/test_python_interpreter_config.py

@@ -72,7 +72,7 @@ def second():
     assert interpreter.output_fields is None
 
 
-@pytest.mark.parametrize("name", ["not valid", "class", "SUBMIT", "FinalOutput"])
+@pytest.mark.parametrize("name", ["not valid", "class", "SUBMIT", "FinalOutput", "json", "run_sync"])
 def test_bind_rejects_invalid_tool_names_without_changing_binding(name):
     def original():
         return "original"
@@ -85,6 +85,23 @@ def original():
     assert interpreter.tools == {"original": original}
 
 
+def test_bind_rejects_unsupported_tool_signature_without_changing_binding():
+    unsupported_default = object()
+
+    def original():
+        return "original"
+
+    def unsupported(value=unsupported_default):
+        return value
+
+    interpreter = PythonInterpreter(deno_command=["deno", "run"], tools={"original": original})
+
+    with pytest.raises(CodeInterpreterError, match="unsupported signature"):
+        interpreter.bind(tools={"unsupported": unsupported})
+
+    assert interpreter.tools == {"original": original}
+
+
 def test_bind_rejects_invalid_output_fields_without_changing_binding():
     original_fields = [{"name": "answer", "type": "str"}]
     interpreter = PythonInterpreter(deno_command=["deno", "run"], output_fields=original_fields)
@@ -110,6 +127,20 @@ def test_constructor_rejects_reserved_tool_name():
         PythonInterpreter(deno_command=["deno", "run"], tools={"FinalOutput": lambda: None})
 
 
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"tools": []}, "tools must be a mapping"),
+        ({"tools": ""}, "tools must be a mapping"),
+        ({"output_fields": ""}, "output_fields must be a sequence"),
+        ({"output_fields": b""}, "output_fields must be a sequence"),
+    ],
+)
+def test_constructor_validates_falsy_bindings(kwargs, message):
+    with pytest.raises(TypeError, match=message):
+        PythonInterpreter(deno_command=["deno", "run"], **kwargs)
+
+
 def test_bind_after_shutdown_fails():
     interpreter = PythonInterpreter(deno_command=["deno", "run"])
     interpreter.shutdown()