timeout validation and exception handling

Author: alexzhang13

rlm/environments/ipython_repl.py

@@ -490,6 +490,18 @@ class IPythonREPL(NonIsolatedEnv):
           ``ipykernel`` and is fully isolated from the parent's
           namespace, cwd, and signals.
 
+    Subcall attribution caveat (subprocess mode):
+        Subcall completions and FINAL_VAR answers are tagged with the
+        ``cell_id`` that was active *at the moment the kernel issued the
+        request*. If a cell spawns long-lived kernel-side state (a
+        ``threading.Thread``, an ``asyncio.Task`` left running, a
+        background timer) that calls ``rlm_query`` after the spawning
+        cell has finished, those calls will be tagged with whatever cell
+        is active *at call time*, not the cell that spawned the worker.
+        Practically: avoid leaving background work running across cells,
+        or accept that its rlm_calls will be reported under a later
+        cell's ``REPLResult``.
+
     Args:
         lm_handler_address: (host, port) of the LM handler socket server.
         context_payload: Initial context to load as ``context``/``context_0``.
@@ -540,6 +552,17 @@ def __init__(
             raise ValueError(
                 f"kernel_mode must be 'in_process' or 'subprocess', got {kernel_mode!r}"
             )
+        if startup_timeout <= 0:
+            raise ValueError(f"startup_timeout must be positive, got {startup_timeout!r}")
+        if subcall_timeout is not None and subcall_timeout <= 0:
+            # ``socket.settimeout(0)`` is non-blocking mode (every send/recv
+            # raises immediately) — almost certainly not what a caller
+            # asking for a "0s timeout" actually wants. ``None`` means
+            # "no timeout"; a positive number means "this many seconds".
+            raise ValueError(
+                f"subcall_timeout must be positive or None (got {subcall_timeout!r}); "
+                "use None to disable the timeout."
+            )
 
         super().__init__(
             persistent=persistent,
@@ -582,6 +605,19 @@ def __init__(
         # (cross-thread broker reentry) or corrupt this cell's tracking
         # state (same-thread reentry).
         self._subcall_threads: set[int] = set()
+        # Tracks threads currently inside ``execute_code``. In ``in_process``
+        # mode the LM can reach this REPL via any scaffold bound method's
+        # ``__self__`` (e.g. ``rlm_query.__self__.execute_code(...)``); a
+        # cell that does so would re-enter on the same thread under the
+        # RLock and silently clobber its own ``_pending_llm_calls`` /
+        # ``_last_final_answer`` tracking. Catching same-thread reentry
+        # here turns that into a clear ``RuntimeError``. Subprocess mode
+        # is already shielded by process isolation: the kernel-side
+        # scaffold is plain functions with no ``__self__``.
+        self._executing_threads: set[int] = set()
+        # Single lock guarding both reentry-tracking sets. Touched only
+        # on ``execute_code`` enter/exit and inside ``_tracked_subcall``,
+        # so contention is negligible.
         self._subcall_threads_lock = threading.Lock()
         self._context_count: int = 0
         self._history_count: int = 0
@@ -619,7 +655,13 @@ def __init__(
             if setup_code:
                 self.execute_code(setup_code)
         except BaseException:
-            self.cleanup()
+            # Suppress *cleanup* errors so they don't shadow the original
+            # cause. ``raise`` (with no arg) re-raises the real exception
+            # the user needs to see.
+            try:
+                self.cleanup()
+            except Exception:
+                pass
             raise
 
     # -------------------------------------------------------------------------
@@ -1086,27 +1128,53 @@ def update_handler_address(self, address: tuple[str, int]) -> None:
     # -------------------------------------------------------------------------
 
     def execute_code(self, code: str) -> REPLResult:
-        # Reentry guard: if the caller's thread is currently inside this
-        # REPL's ``subcall_fn`` (in-process synchronously, or subprocess
-        # via the broker handler thread), calling ``execute_code`` would
-        # either deadlock the cell lock (cross-thread broker reentry) or
-        # silently corrupt the in-flight cell's tracking (same-thread
-        # in-process reentry). Fail fast with a clear message.
+        # Reentry guards:
+        #
+        #  1. ``cur in self._subcall_threads`` — the caller is currently
+        #     inside ``subcall_fn`` (in-process synchronously, or via the
+        #     subprocess broker handler thread). Calling ``execute_code``
+        #     would either deadlock the cell lock (cross-thread case) or
+        #     silently corrupt this cell's bookkeeping (same-thread).
+        #
+        #  2. ``cur in self._executing_threads`` — the caller is already
+        #     inside an ``execute_code`` invocation on this same REPL.
+        #     This is the path an LM-generated cell takes when it
+        #     traverses ``rlm_query.__self__.execute_code(...)`` (or any
+        #     scaffold bound method) to call back into the parent. The
+        #     RLock would let it reenter, but the inner call clears
+        #     ``_pending_llm_calls`` / resets ``_last_final_answer`` and
+        #     leaves the outer cell with a corrupt result.
         cur = threading.get_ident()
         with self._subcall_threads_lock:
-            is_reentrant = cur in self._subcall_threads
-        if is_reentrant:
+            in_subcall = cur in self._subcall_threads
+            already_executing = cur in self._executing_threads
+        if in_subcall:
             raise RuntimeError(
                 "Reentrant execute_code on the same instance from inside "
                 "subcall_fn is not supported — it would deadlock or "
                 "clobber the parent cell's bookkeeping. subcall_fn must "
                 "spawn a child REPL (with its own lock) instead of "
                 "calling back into its parent."
             )
+        if already_executing:
+            raise RuntimeError(
+                "Reentrant execute_code on the same instance is not "
+                "supported. The currently-running cell appears to have "
+                "called execute_code on its own REPL (e.g. via "
+                "rlm_query.__self__.execute_code(...)) — that would "
+                "corrupt this cell's tracking state. Use a separate "
+                "REPL instance for nested execution."
+            )
 
-        if self.kernel_mode == "in_process":
-            return self._execute_in_process(code)
-        return self._execute_in_kernel(code, timeout=self.cell_timeout, drain_broker=True)
+        with self._subcall_threads_lock:
+            self._executing_threads.add(cur)
+        try:
+            if self.kernel_mode == "in_process":
+                return self._execute_in_process(code)
+            return self._execute_in_kernel(code, timeout=self.cell_timeout, drain_broker=True)
+        finally:
+            with self._subcall_threads_lock:
+                self._executing_threads.discard(cur)
 
     def _execute_in_process(self, code: str) -> REPLResult:
         start_time = time.perf_counter()

tests/test_ipython_repl.py

@@ -319,6 +319,48 @@ def test_invalid_kernel_mode_rejected():
         IPythonREPL(kernel_mode="nonsense")  # type: ignore[arg-type]
 
 
+@pytest.mark.parametrize("bad", [0, -1, -0.5])
+def test_invalid_startup_timeout_rejected(bad):
+    with pytest.raises(ValueError, match="startup_timeout"):
+        IPythonREPL(kernel_mode="in_process", startup_timeout=bad)
+
+
+@pytest.mark.parametrize("bad", [0, -1, -0.5])
+def test_invalid_subcall_timeout_rejected(bad):
+    with pytest.raises(ValueError, match="subcall_timeout"):
+        IPythonREPL(kernel_mode="in_process", subcall_timeout=bad)
+
+
+def test_subcall_timeout_none_is_allowed():
+    """``None`` is the documented "no timeout" sentinel, distinct from 0."""
+    with IPythonREPL(kernel_mode="in_process", subcall_timeout=None) as repl:
+        assert repl.subcall_timeout is None
+
+
+def test_init_failure_preserves_original_exception_over_cleanup_error(monkeypatch):
+    """If ``setup`` fails *and* cleanup would also fail, ``__init__``
+    must surface the original cause, not the cleanup error."""
+
+    class _Boom(Exception):
+        pass
+
+    # Fail in ``_setup_in_process`` before any IPython state is created
+    # (so the test doesn't leave dangling atexit handlers behind).
+    def _broken_setup(self):
+        raise _Boom("original cause")
+
+    monkeypatch.setattr(IPythonREPL, "_setup_in_process", _broken_setup)
+
+    # Make cleanup itself raise too, so we can verify the original wins.
+    def _broken_cleanup(self):
+        raise RuntimeError("cleanup also failed")
+
+    monkeypatch.setattr(IPythonREPL, "cleanup", _broken_cleanup)
+
+    with pytest.raises(_Boom, match="original cause"):
+        IPythonREPL(kernel_mode="in_process")
+
+
 @BOTH_MODES
 def test_cleanup_is_idempotent(kernel_mode: str):
     repl = IPythonREPL(kernel_mode=kernel_mode)
@@ -673,6 +715,57 @@ def slow_subcall(prompt: str, model: str | None = None) -> RLMChatCompletion:
     )
 
 
+# -----------------------------------------------------------------------------
+# Reentrant cell → execute_code (LM reaches parent via bound-method __self__)
+# -----------------------------------------------------------------------------
+
+
+def test_in_process_cell_cannot_reenter_via_scaffold_self():
+    """An in-process cell that traverses any scaffold bound method's
+    ``__self__`` to call ``execute_code`` on the parent must surface a
+    ``RuntimeError`` rather than silently corrupting the parent's
+    tracking state.
+
+    Subprocess mode is naturally immune: kernel-side scaffolds are
+    plain functions with no ``__self__`` referencing the parent.
+    """
+    with IPythonREPL(kernel_mode="in_process") as repl:
+        # ``rlm_query`` is a bound method; ``rlm_query.__self__`` is the
+        # parent REPL. Any other scaffold (llm_query, FINAL_VAR, ...)
+        # would work just as well as a reentry vector.
+        result = repl.execute_code("rlm_query.__self__.execute_code('print(\"inner\")')")
+    assert (
+        "RuntimeError" in result.stderr
+        or "Reentrant" in result.stderr
+        or "Reentrant" in result.stdout
+    ), result.stderr or result.stdout
+
+
+def test_in_process_scaffold_self_reentry_preserves_outer_state():
+    """If the LM tries to reenter and we raise, the *outer* cell's
+    bookkeeping (``_pending_llm_calls``, ``_last_final_answer``) must
+    still be intact when execution resumes after the caught error."""
+
+    fake = _FakeSubcall(responses=["A"])
+    with IPythonREPL(kernel_mode="in_process", subcall_fn=fake) as repl:
+        # Cell makes a real rlm_query, then attempts reentry. The reentry
+        # raises and is caught locally; the outer rlm_query call should
+        # still be reflected in ``rlm_calls``.
+        result = repl.execute_code(
+            "real = rlm_query('outer')\n"
+            "try:\n"
+            "    rlm_query.__self__.execute_code('inner')\n"
+            "except RuntimeError as e:\n"
+            "    print('caught:', type(e).__name__)\n"
+            "print('outer-result:', real)\n"
+        )
+    assert "caught: RuntimeError" in result.stdout
+    assert "outer-result: A" in result.stdout
+    assert len(result.rlm_calls) == 1, (
+        f"outer rlm_query must remain tracked; got {len(result.rlm_calls)}"
+    )
+
+
 # -----------------------------------------------------------------------------
 # Reentrant subcall_fn → execute_code (deadlock prevention)
 # -----------------------------------------------------------------------------