move max_concurrent_subcalls as a general REPL property. add more tests

Author: alexzhang13

rlm/core/rlm.py

@@ -242,8 +242,7 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
                 env_kwargs["custom_sub_tools"] = self.custom_sub_tools
             if self.compaction and self.environment_type == "local":
                 env_kwargs["compaction"] = True
-            if self.environment_type == "local":
-                env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
+            env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
             environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
 
             if self.persistent:

rlm/environments/base_env.py

@@ -212,9 +212,12 @@ class BaseEnv(ABC):
             }
     """
 
-    def __init__(self, persistent: bool = False, depth: int = 1, **kwargs):
+    def __init__(
+        self, persistent: bool = False, depth: int = 1, max_concurrent_subcalls: int = 4, **kwargs
+    ):
         self.persistent = persistent
         self.depth = depth
+        self.max_concurrent_subcalls = max_concurrent_subcalls
         self.kwargs = kwargs
 
     @abstractmethod

rlm/environments/local_repl.py

@@ -139,11 +139,15 @@ def __init__(
         max_concurrent_subcalls: int = 4,
         **kwargs,
     ):
-        super().__init__(persistent=persistent, depth=depth, **kwargs)
+        super().__init__(
+            persistent=persistent,
+            depth=depth,
+            max_concurrent_subcalls=max_concurrent_subcalls,
+            **kwargs,
+        )
 
         self.lm_handler_address = lm_handler_address
         self.subcall_fn = subcall_fn  # Callback for recursive RLM calls (depth > 1 support)
-        self.max_concurrent_subcalls = max_concurrent_subcalls
         self.original_cwd = os.getcwd()
         self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_")
         self._lock = threading.Lock()
@@ -366,8 +370,7 @@ def _run_subcall(index: int, prompt: str) -> None:
 
             with ThreadPoolExecutor(max_workers=max_workers) as executor:
                 futures = [
-                    executor.submit(_run_subcall, i, prompt)
-                    for i, prompt in enumerate(prompts)
+                    executor.submit(_run_subcall, i, prompt) for i, prompt in enumerate(prompts)
                 ]
                 # Wait for all futures to complete; exceptions are captured inside _run_subcall
                 for future in as_completed(futures):

tests/test_rlm_query.py

@@ -281,6 +281,7 @@ def sometimes_fail(prompt, model):
 
     def test_batched_pending_calls_ordered(self):
         """Pending LLM calls should be appended in prompt order for deterministic metadata."""
+
         def delayed_subcall(prompt, model):
             # Reverse delay so later prompts finish first
             delay = 0.1 if prompt == "first" else 0.01
@@ -299,12 +300,175 @@ def test_single_prompt_skips_thread_pool(self):
         """Single prompt should not use ThreadPoolExecutor (no overhead)."""
         subcall_fn = MagicMock(return_value=_make_completion("solo"))
         repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4)
-        result = repl.execute_code("answers = rlm_query_batched(['only'])")
+        repl.execute_code("answers = rlm_query_batched(['only'])")
         assert repl.locals["answers"] == ["solo"]
         subcall_fn.assert_called_once_with("only", None)
         repl.cleanup()
 
 
+class TestMaxConcurrentSubcallsBounds:
+    """Tests for edge cases and boundary conditions of max_concurrent_subcalls."""
+
+    def test_max_concurrent_subcalls_default_value(self):
+        """Default max_concurrent_subcalls should be 4."""
+        repl = LocalREPL()
+        assert repl.max_concurrent_subcalls == 4
+        repl.cleanup()
+
+    def test_max_concurrent_subcalls_custom_value(self):
+        """Should accept custom max_concurrent_subcalls values."""
+        repl = LocalREPL(max_concurrent_subcalls=8)
+        assert repl.max_concurrent_subcalls == 8
+        repl.cleanup()
+
+    def test_max_concurrent_subcalls_one(self):
+        """max_concurrent_subcalls=1 should force sequential execution."""
+        order = []
+        lock = threading.Lock()
+
+        def sequential_subcall(prompt, model):
+            with lock:
+                order.append(f"start-{prompt}")
+            time.sleep(0.05)
+            with lock:
+                order.append(f"end-{prompt}")
+            return _make_completion(f"done-{prompt}")
+
+        repl = LocalREPL(subcall_fn=sequential_subcall, max_concurrent_subcalls=1)
+        result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c'])")
+        assert result.stderr == ""
+        answers = repl.locals["answers"]
+        assert answers == ["done-a", "done-b", "done-c"]
+        # With max_workers=1, tasks run one at a time so each start
+        # should be followed by its end before the next start
+        for i in range(0, len(order) - 1, 2):
+            prompt_id = order[i].split("-", 1)[1]
+            assert order[i] == f"start-{prompt_id}"
+            assert order[i + 1] == f"end-{prompt_id}"
+        repl.cleanup()
+
+    def test_max_concurrent_subcalls_larger_than_prompts(self):
+        """max_concurrent_subcalls larger than prompt count should work fine."""
+        subcall_fn = MagicMock(return_value=_make_completion("ok"))
+        repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=100)
+        result = repl.execute_code("answers = rlm_query_batched(['a', 'b'])")
+        assert result.stderr == ""
+        assert repl.locals["answers"] == ["ok", "ok"]
+        assert subcall_fn.call_count == 2
+        repl.cleanup()
+
+    def test_max_concurrent_subcalls_exact_match(self):
+        """max_concurrent_subcalls == prompt count should allow full parallelism."""
+        active_count = []
+        active_lock = threading.Lock()
+        active = [0]
+
+        def tracked_subcall(prompt, model):
+            with active_lock:
+                active[0] += 1
+                active_count.append(active[0])
+            time.sleep(0.1)
+            with active_lock:
+                active[0] -= 1
+            return _make_completion(f"ok {prompt}")
+
+        repl = LocalREPL(subcall_fn=tracked_subcall, max_concurrent_subcalls=3)
+        result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c'])")
+        assert result.stderr == ""
+        # All 3 should be able to run at once
+        assert max(active_count) <= 3
+        repl.cleanup()
+
+    def test_batched_all_failures_parallel(self):
+        """All subcalls failing in parallel should return all error strings."""
+
+        def always_fail(prompt, model):
+            raise ValueError(f"fail-{prompt}")
+
+        repl = LocalREPL(subcall_fn=always_fail, max_concurrent_subcalls=4)
+        result = repl.execute_code("answers = rlm_query_batched(['x', 'y', 'z'])")
+        assert result.stderr == ""
+        answers = repl.locals["answers"]
+        assert len(answers) == 3
+        for i, prompt in enumerate(["x", "y", "z"]):
+            assert "Error" in answers[i]
+            assert f"fail-{prompt}" in answers[i]
+        repl.cleanup()
+
+    def test_batched_large_batch_with_low_concurrency(self):
+        """Many prompts with low concurrency should still complete correctly."""
+        max_concurrent = 2
+        active_count = []
+        active_lock = threading.Lock()
+        active = [0]
+
+        def tracked_subcall(prompt, model):
+            with active_lock:
+                active[0] += 1
+                active_count.append(active[0])
+            time.sleep(0.02)
+            with active_lock:
+                active[0] -= 1
+            return _make_completion(f"r-{prompt}")
+
+        prompts = [f"p{i}" for i in range(10)]
+        repl = LocalREPL(subcall_fn=tracked_subcall, max_concurrent_subcalls=max_concurrent)
+        result = repl.execute_code(f"answers = rlm_query_batched({prompts!r})")
+        assert result.stderr == ""
+        answers = repl.locals["answers"]
+        assert len(answers) == 10
+        assert answers == [f"r-p{i}" for i in range(10)]
+        # Concurrency cap should be respected
+        assert max(active_count) <= max_concurrent
+        repl.cleanup()
+
+    def test_batched_empty_prompts_with_parallel(self):
+        """Empty prompt list should return empty list regardless of concurrency setting."""
+        subcall_fn = MagicMock()
+        repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4)
+        repl.execute_code("answers = rlm_query_batched([])")
+        assert repl.locals["answers"] == []
+        subcall_fn.assert_not_called()
+        repl.cleanup()
+
+    def test_pending_calls_exclude_failures(self):
+        """Failed subcalls should not appear in pending_llm_calls."""
+
+        def fail_second(prompt, model):
+            if prompt == "bad":
+                raise RuntimeError("boom")
+            return _make_completion(f"ok-{prompt}")
+
+        repl = LocalREPL(subcall_fn=fail_second, max_concurrent_subcalls=4)
+        result = repl.execute_code("answers = rlm_query_batched(['good1', 'bad', 'good2'])")
+        assert result.stderr == ""
+        # Only 2 successful completions should be in rlm_calls
+        assert len(result.rlm_calls) == 2
+        assert result.rlm_calls[0].response == "ok-good1"
+        assert result.rlm_calls[1].response == "ok-good2"
+        repl.cleanup()
+
+
+class TestMaxConcurrentSubcallsOnBaseEnv:
+    """Tests that max_concurrent_subcalls is a property of BaseEnv, not just LocalREPL."""
+
+    def test_base_env_has_max_concurrent_subcalls(self):
+        """BaseEnv should accept and store max_concurrent_subcalls."""
+
+        # BaseEnv is abstract, so we test via LocalREPL which inherits from it
+        repl = LocalREPL(max_concurrent_subcalls=16)
+        assert repl.max_concurrent_subcalls == 16
+        # Verify it's set via the BaseEnv.__init__ chain
+        assert hasattr(repl, "max_concurrent_subcalls")
+        repl.cleanup()
+
+    def test_base_env_default_max_concurrent_subcalls(self):
+        """BaseEnv default should be 4."""
+        repl = LocalREPL()
+        assert repl.max_concurrent_subcalls == 4
+        repl.cleanup()
+
+
 class TestRlmQueryScaffoldRestoration:
     """Test that rlm_query and rlm_query_batched are restored after overwrite."""