@hicaru/pi-rlm 0.2.0 → 0.2.1

package/src/sandbox/worker.py

@@ -12,6 +12,12 @@ Protocol (worker -> parent):  {"id","ok",...result}            # response to a r
 When sandbox code calls llm_query/rlm_query/advance_phase/save_artifact/ask_user_question/todo, the worker writes a request line
 and BLOCKS reading stdin until the matching {"type":"llm_reply","rid",...} arrives. The parent
 services the request in-process (it holds API keys).
+
+Requests and replies are decoupled: `_post` writes a request and returns its rid without
+waiting, and replies are parked in `_inbox` keyed by rid until something asks for them. That
+is what makes `spawn()` / `rlm_await()` / `rlm_await_all()` possible — many requests can be
+in flight at once (the parent already services interrupts concurrently), and a task may be
+awaited in a LATER exec than the one that started it.
 """
 
 from __future__ import annotations
@@ -98,13 +104,42 @@ def _install_read_only_guards():
     return guarded_io_open
 
 
+# _builtin()'s getattr(..., None) fallback would silently inject None for a name this
+# interpreter lacks, surfacing much later as "'NoneType' object is not callable" inside model
+# code. Fail at startup instead. Note "None" is legitimately None, and the block-list below is
+# deliberate — which is why this check runs BEFORE it.
+_MISSING = sorted(name for name, value in _SAFE_BUILTINS.items() if value is None and name != "None")
+if _MISSING:
+    raise RuntimeError(f"unsupported Python interpreter: missing builtins {_MISSING}")
+
+def _blocked_builtin(name: str):
+    """Bind a disabled builtin to a callable that explains itself.
+
+    Binding these to None made `eval(...)` fail with a bare "'NoneType' object is not callable",
+    which reads as a broken sandbox rather than a deliberate block: an audit session spent six
+    execs on it and filed a phantom "namespace corruption" bug. Saying so at the point of failure
+    fixes it for every model without spending native-prompt budget on a rule most runs never hit.
+    """
+    def blocked(*_args, **_kwargs):
+        raise PermissionError(
+            f"{name}() is disabled in the RLM sandbox by design — it is not missing and the "
+            "namespace is not corrupt. Names are already bound, so reference them directly; "
+            "inspect `context` with search() / grep_context() / outline()."
+        )
+    blocked.__name__ = name
+    return blocked
+
+
+# Blocked on purpose (NOT missing) — see _blocked_builtin. The _MISSING check above runs first,
+# so a genuinely absent builtin is still a startup failure rather than a silent None.
 for _blocked in ("eval", "exec", "compile", "input", "globals", "locals"):
-    _SAFE_BUILTINS[_blocked] = None
+    _SAFE_BUILTINS[_blocked] = _blocked_builtin(_blocked)
 
 RESERVED = frozenset(
     {
         "llm_query", "llm_query_batched", "llm_query_chunked",
         "rlm_query", "rlm_query_batched",
+        "spawn", "rlm_await", "rlm_await_all",
         "map_files", "llm_map_reduce",
         "search", "grep_context", "outline",
         "advance_phase", "save_artifact",
@@ -301,14 +336,179 @@ def _send(obj: dict[str, Any]) -> None:
     _REAL_STDOUT.flush()
 
 
+class _StallTimeout(Exception):
+    """No frame from the host while a sub-call was pending."""
+
+
+@contextmanager
+def _stall_alarm(exec_timeout_s: float, stall_timeout_s: float):
+    """Swap the per-cell alarm for a stall alarm while blocked on the parent.
+
+    Sub-LLM latency is network time, not cell compute time, so it must not count against the
+    ```repl``` block timeout — but an unbounded wait is exactly how a lost reply turns into a
+    dead session. The yielded `rearm()` restarts the stall clock on every frame, so a healthy
+    long-running child never trips it.
+    """
+    use = hasattr(signal, "SIGALRM")
+    remaining = signal.getitimer(signal.ITIMER_REAL)[0] if (use and exec_timeout_s > 0) else 0.0
+
+    def _fire(signum, frame):  # noqa: ARG001
+        raise _StallTimeout(
+            f"sub-call stalled — no reply from the host for {stall_timeout_s:g}s "
+            "(the task may still be running; rlm_await it again in a later block)"
+        )
+
+    old = signal.signal(signal.SIGALRM, _fire) if use else None
+
+    def rearm() -> None:
+        if use and stall_timeout_s > 0:
+            signal.setitimer(signal.ITIMER_REAL, stall_timeout_s)
+
+    rearm()
+    try:
+        yield rearm
+    finally:
+        if use:
+            signal.setitimer(signal.ITIMER_REAL, 0)
+            if old is not None:
+                signal.signal(signal.SIGALRM, old)
+            if remaining > 0:                      # restore the cell's remaining budget
+                signal.setitimer(signal.ITIMER_REAL, remaining)
+
+
+def _surfaced_error(message: str) -> str:
+    """The "Error: …" contract value, ALSO written to the cell's stderr.
+
+    A spawn/await misuse whose only trace is the returned value reads to the model as a random
+    string much later — which is exactly how `tasks.items()` blew up on a str.
+    """
+    print(f"[rlm] {message}", file=sys.stderr)
+    return f"Error: {message}"
+
+
+# ---- reply reducers: raw parent replies -> the value the scaffold fn returns ----------------
+# One reducer per result shape, shared by the sync helpers and their spawned equivalents.
+
+
+def _reduce_one(replies: list[dict[str, Any]]) -> str:
+    r = replies[0]
+    return f"Error: {r['error']}" if r.get("error") else r.get("response", "")
+
+
+def _reduce_batch(n: int):
+    """Reducer for a single *_query_batched reply of n prompts."""
+    def reduce(replies: list[dict[str, Any]]) -> list[str]:
+        r = replies[0]
+        if r.get("error"):
+            return [f"Error: {r['error']}"] * n
+        out = r.get("responses")
+        if not isinstance(out, list) or len(out) != n:
+            return ["Error: malformed batched response"] * n
+        return [s if isinstance(s, str) else f"Error: {s}" for s in out]
+    return reduce
+
+
+def _reduce_chunked(sizes: list[int]):
+    """Concatenate several llm_query_batched replies back into one flat chunk list."""
+    per = [_reduce_batch(n) for n in sizes]
+
+    def reduce(replies: list[dict[str, Any]]) -> list[str]:
+        out: list[str] = []
+        for red, rep in zip(per, replies):
+            out.extend(red([rep]))
+        return out
+    return reduce
+
+
+def _reduce_map_files(sizes: list[int], spans: list[tuple[str, int]]):
+    """Flatten the batch replies (same as chunked), then regroup them per path.
+
+    A file larger than the per-prompt budget contributed several requests; its answers rejoin
+    in order, which is what makes map_files a {path: answer} dict rather than a flat list.
+    """
+    flatten = _reduce_chunked(sizes)
+
+    def reduce(replies: list[dict[str, Any]]) -> dict[str, str]:
+        responses = flatten(replies)
+        out: dict[str, str] = {}
+        cursor = 0
+        for path, count in spans:
+            part = responses[cursor:cursor + count]
+            cursor += count
+            out[path] = part[0] if count == 1 and part else "\n\n".join(part)
+        return out
+    return reduce
+
+
+def _spawnable(name: str):
+    """Tag a sync scaffold fn with the request kind spawn() should route it to."""
+    def mark(fn):
+        fn._rlm_name = name
+        return fn
+    return mark
+
+
+class Task:
+    """Handle for parent-side work in flight, returned by spawn().
+
+    Opaque to model code apart from `done` and repr. A Task may be awaited in a later
+    ```repl``` block than the one that created it.
+    """
+
+    __slots__ = ("kind", "label", "_worker", "_rids", "_reduce", "_value", "_settled")
+
+    def __init__(self, worker: "Worker", kind: str, rids, reduce, label: str = ""):
+        self.kind = kind
+        self.label = label
+        self._worker = worker
+        self._rids = tuple(rids)
+        self._reduce = reduce
+        self._value: Any = None
+        self._settled = False
+
+    @staticmethod
+    def resolved(worker: "Worker", kind: str, value: Any, label: str = "") -> "Task":
+        """A Task that never hit the wire — validation errors and empty inputs."""
+        task = Task(worker, kind, (), lambda _replies: value, label)
+        task._value = value
+        task._settled = True
+        return task
+
+    @property
+    def done(self) -> bool:
+        """True once every reply has landed — awaiting will not block."""
+        return self._settled or all(r in self._worker.inbox for r in self._rids)
+
+    def __repr__(self) -> str:
+        return f"<Task {self.kind} {'done' if self.done else 'running'} {self.label}>"
+
+
 class Worker:
-    def __init__(self, depth: int, exec_timeout_s: float, max_prompt_chars: int, read_only: bool = False):
+    def __init__(
+        self,
+        depth: int,
+        exec_timeout_s: float,
+        max_prompt_chars: int,
+        read_only: bool = False,
+        await_timeout_s: float = 600.0,
+    ):
         self.depth = depth
         self.exec_timeout_s = exec_timeout_s
         self.max_prompt_chars = max_prompt_chars
         self.read_only = read_only
+        self.await_timeout_s = await_timeout_s
         self._rid = 0
         self._final_answer: str | None = None
+        # Replies parked by rid until something awaits them. Unbounded by design: a task
+        # the model spawns and never awaits keeps its entry for the life of the process.
+        # Bounded in practice by session length; evicting would silently hang a later
+        # rlm_await, which is strictly worse than the memory.
+        self.inbox: dict[str, dict[str, Any]] = {}
+        self._inflight: set[str] = set()
+        # Requests (exec/snapshot/shutdown) that arrived mid-exec; main() replays them.
+        self._deferred: list[Any] = []
+        # True only while spawn() runs a builder — marks requests that may outlive this exec.
+        self._detached = False
         self.ns: dict[str, Any] = {}
         self._setup()
 
@@ -336,6 +536,9 @@ class Worker:
         ns["llm_query_chunked"] = self._llm_query_chunked
         ns["rlm_query"] = self._rlm_query
         ns["rlm_query_batched"] = self._rlm_query_batched
+        ns["spawn"] = self._spawn
+        ns["rlm_await"] = self._await_task
+        ns["rlm_await_all"] = self._await_all
         ns["map_files"] = self._map_files
         ns["llm_map_reduce"] = self._llm_map_reduce
         ns["search"] = self._search
@@ -391,54 +594,120 @@ class Worker:
 
     # ---- sub-LLM bridge over stdio --------------------------------------------------------
 
-    def _rpc(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
+    def _post(self, kind: str, payload: dict[str, Any]) -> str:
+        """Write one parent request and return its rid WITHOUT waiting for the reply."""
         self._rid += 1
         rid = f"q{self._rid}"
-        _send({"type": kind, "rid": rid, "depth": self.depth, **payload})
-        # The per-cell SIGALRM is wall-clock; it must not count time blocked here waiting for
-        # a sub-LLM reply (network/LLM latency, not local CPU). Pause it across the readline.
-        pause = self.exec_timeout_s > 0 and hasattr(signal, "SIGALRM")
-        if pause:
-            remaining = signal.getitimer(signal.ITIMER_REAL)[0]
-            signal.setitimer(signal.ITIMER_REAL, 0)
+        # Register only after the write succeeds — a broken pipe must not leave an
+        # _inflight entry that nothing will ever settle.
+        _send({"type": kind, "rid": rid, "depth": self.depth,
+               "detached": self._detached, **payload})
+        self._inflight.add(rid)
+        return rid
+
+    def park_reply(self, msg: Any) -> bool:
+        """File an llm_reply against its rid. True when the frame was a reply.
+
+        Public because main() needs it too: a spawned task can settle while the worker
+        sits idle between execs, and that reply must not fall through to "unknown type".
+        """
+        if not isinstance(msg, dict) or msg.get("type") != "llm_reply":
+            return False
+        rid = msg.get("rid")
+        if isinstance(rid, str) and rid in self._inflight:
+            self._inflight.discard(rid)
+            self.inbox[rid] = msg
+        else:
+            # Late reply to an abandoned rid (e.g. a request from a discarded sandbox).
+            print(f"[rlm-sandbox] dropping reply for unknown rid: {rid!r}", file=_REAL_STDERR)
+        return True
+
+    def take_deferred(self) -> Any | None:
+        """Pop a request that arrived mid-exec, for main() to replay. None when empty."""
+        return self._deferred.pop(0) if self._deferred else None
+
+    def _pump(self) -> bool:
+        """Read one frame from the parent into the inbox. False when the pipe closed."""
+        line = _REAL_STDIN.readline()
+        if not line:
+            return False
         try:
-            while True:
-                line = _REAL_STDIN.readline()
-                if not line:
+            msg = json.loads(line)
+        except ValueError:
+            print(f"[rlm-sandbox] skipping non-JSON parent frame: {line[:200]}", file=_REAL_STDERR)
+            return True
+        if self.park_reply(msg):
+            return True
+        # A request (exec/snapshot/shutdown) arriving mid-exec: main() replays it.
+        self._deferred.append(msg)
+        return True
+
+    def _drain_until(self, rids) -> None:
+        """Block until every rid in `rids` has its reply parked in the inbox.
+
+        Bounded: a host that goes silent raises inside the ```repl``` block instead of hanging
+        the session forever.
+        """
+        if all(r in self.inbox for r in rids):
+            return
+        with _stall_alarm(self.exec_timeout_s, self.await_timeout_s) as rearm:
+            while not all(r in self.inbox for r in rids):
+                if not self._pump():
                     raise RuntimeError("parent closed the pipe during a sub-LLM request")
-                msg = json.loads(line)
-                if msg.get("type") == "llm_reply" and msg.get("rid") == rid:
-                    return msg
-                # Stray/late message (e.g. a reply to an earlier timed-out request): skip it.
-                print(
-                    f"[rlm-sandbox] ignoring unexpected message during sub-LLM request: {str(msg)[:200]}",
-                    file=_REAL_STDERR,
-                )
-        finally:
-            if pause and remaining > 0:
-                signal.setitimer(signal.ITIMER_REAL, remaining)
+                rearm()
 
-    def _llm_query(self, prompt: str, model: str | None = None) -> str:
-        r = self._rpc("llm_query", {"prompt": str(prompt), "model": model})
-        return f"Error: {r['error']}" if r.get("error") else r.get("response", "")
+    def _take(self, rids) -> list[dict[str, Any]]:
+        return [self.inbox.pop(r) for r in rids]
 
-    def _llm_query_batched(self, prompts, model: str | None = None) -> list[str]:
+    def _rpc(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
+        """Post one request and block for its reply — the synchronous single-shot path."""
+        rid = self._post(kind, payload)
+        self._drain_until((rid,))
+        return self._take((rid,))[0]
+
+    # ---- spawn / await ---------------------------------------------------------------------
+
+    def _start_prompt(self, kind: str, prompt, model) -> Task:
+        text = str(prompt)
+        # A sub-LLM asked nothing answers something: the confabulation then sits in `answers`
+        # looking exactly like data. Refuse instead of spending a call on it.
+        if not text.strip():
+            return Task.resolved(self, kind, _surfaced_error(
+                f"{kind}() got an empty prompt — a sub-LLM would confabulate an answer to nothing"))
+        rid = self._post(kind, {"prompt": text, "model": model})
+        return Task(self, kind, (rid,), _reduce_one, text[:40])
+
+    def _start_prompts(self, kind: str, prompts, model) -> Task:
         prompts = [str(p) for p in prompts]
         if not prompts:
-            return []
-        r = self._rpc("llm_query_batched", {"prompts": prompts, "model": model})
-        if r.get("error"):
-            return [f"Error: {r['error']}"] * len(prompts)
-        out = r.get("responses")
-        if not isinstance(out, list) or len(out) != len(prompts):
-            return ["Error: malformed batched response"] * len(prompts)
-        return [s if isinstance(s, str) else f"Error: {s}" for s in out]
+            return Task.resolved(self, kind, [])
+        # Only the all-blank case: one blank prompt among twenty is the caller's business.
+        if not any(p.strip() for p in prompts):
+            return Task.resolved(self, kind, [
+                _surfaced_error(f"{kind}() got only empty prompts")
+            ] * len(prompts))
+        rid = self._post(kind, {"prompts": prompts, "model": model})
+        return Task(self, kind, (rid,), _reduce_batch(len(prompts)), f"×{len(prompts)}")
 
-    def _llm_query_chunked(self, text, prompt: str, model: str | None = None) -> list[str]:
-        """Split oversized text into cap-sized chunks and fan out via llm_query_batched.
+    def _start_llm_query(self, prompt, model: str | None = None) -> Task:
+        return self._start_prompt("llm_query", prompt, model)
+
+    def _start_rlm_query(self, prompt, model: str | None = None) -> Task:
+        return self._start_prompt("rlm_query", prompt, model)
+
+    def _start_llm_query_batched(self, prompts, model: str | None = None) -> Task:
+        return self._start_prompts("llm_query_batched", prompts, model)
+
+    def _start_rlm_query_batched(self, prompts, model: str | None = None) -> Task:
+        return self._start_prompts("rlm_query_batched", prompts, model)
+
+    def _start_llm_query_chunked(self, text, prompt: str, model: str | None = None) -> Task:
+        """Split oversized text into cap-sized chunks and post EVERY batch at once.
 
-        Returns one answer per chunk, order preserved. No exceptions escape: errors come
-        back as "Error: ..." strings per chunk (same contract as llm_query_batched).
+        One answer per chunk, order preserved. No exceptions escape: errors come back as
+        "Error: ..." strings per chunk (same contract as llm_query_batched). Because all
+        batches go on the wire together, a large input costs one round-trip of latency
+        rather than one per 20 chunks.
 
         NOTE: budget uses Python code-point length (len) while the parent-side cap check counts
         UTF-16 units (JS string.length); astral/emoji-heavy text may be marginally larger on the
@@ -446,22 +715,93 @@ class Worker:
         """
         text, prompt = str(text), str(prompt)
         if not text:
-            return []
+            return Task.resolved(self, "llm_query_chunked", [])
         budget = self.max_prompt_chars - len(prompt) - _CHUNK_HEADER_OVERHEAD
         if budget < 1_000:
-            return [f"Error: prompt leaves under 1,000 chars per chunk (cap {self.max_prompt_chars:,}) — shorten the instruction"]
+            return Task.resolved(self, "llm_query_chunked", [
+                f"Error: prompt leaves under 1,000 chars per chunk (cap {self.max_prompt_chars:,}) — shorten the instruction"
+            ])
         chunks = _chunk_text(text, budget)
         total = len(chunks)
         if total > _MAX_CHUNKS:
-            return [f"Error: {total} chunks would be needed — filter/slice the text in Python first"]
-        results: list[str] = []
+            return Task.resolved(self, "llm_query_chunked", [
+                f"Error: {total} chunks would be needed — filter/slice the text in Python first"
+            ])
+        rids: list[str] = []
+        sizes: list[int] = []
         for i in range(0, total, _MAX_CHUNK_BATCH):
             batch = [
                 f"{prompt}\n\n[chunk {i + j + 1}/{total} of the input]\n{c}"
                 for j, c in enumerate(chunks[i:i + _MAX_CHUNK_BATCH])
             ]
-            results.extend(self._llm_query_batched(batch, model))
-        return results
+            rids.append(self._post("llm_query_batched", {"prompts": batch, "model": model}))
+            sizes.append(len(batch))
+        return Task(self, "llm_query_chunked", tuple(rids), _reduce_chunked(sizes), f"{total} chunks")
+
+    def _builder_for(self, name: str):
+        # llm_map_reduce is deliberately absent: its reduce step is a SECOND sub-LLM call that
+        # depends on its own map results, so it cannot be one (rids, pure reduce) Task.
+        return {
+            "llm_query": self._start_llm_query,
+            "llm_query_batched": self._start_llm_query_batched,
+            "llm_query_chunked": self._start_llm_query_chunked,
+            "map_files": self._start_map_files,
+            "rlm_query": self._start_rlm_query,
+            "rlm_query_batched": self._start_rlm_query_batched,
+        }.get(name)
+
+    def _spawn(self, fn, *args, **kwargs) -> Task:
+        """Start a sub-call without waiting for it. `fn` is the scaffold function itself.
+
+        Returns a Task for rlm_await / rlm_await_all, possibly in a later ```repl``` block.
+        Misuse returns an already-resolved error Task rather than raising, matching the
+        "Error: ..." contract of the synchronous helpers.
+        """
+        name = getattr(fn, "_rlm_name", None)
+        builder = self._builder_for(name) if isinstance(name, str) else None
+        if builder is None:
+            return Task.resolved(self, "spawn", _surfaced_error(
+                "spawn() takes llm_query, llm_query_batched, llm_query_chunked, map_files, "
+                "rlm_query or rlm_query_batched — not llm_map_reduce, whose reduce step depends "
+                "on its own map results and so cannot be a single Task"))
+        # Mark every request this builder posts as detached: the parent routes them to its
+        # session-scoped registry, since they may outlive the exec that started them.
+        self._detached = True
+        try:
+            return builder(*args, **kwargs)
+        except TypeError as e:
+            return Task.resolved(self, "spawn", _surfaced_error(f"bad spawn arguments — {e}"))
+        finally:
+            self._detached = False
+
+    def _await_task(self, task) -> Any:
+        """Block until `task` has its result. Idempotent — the value is memoized."""
+        if not isinstance(task, Task):
+            return _surfaced_error(
+                f"rlm_await expects a Task from spawn(), got {type(task).__name__}"
+            )
+        if not task._settled:
+            self._drain_until(task._rids)
+            task._value = task._reduce(self._take(task._rids))
+            task._settled = True
+        return task._value
+
+    def _await_all(self, tasks) -> list:
+        """Block until every task has its result. Order matches the input."""
+        tasks = list(tasks)
+        # One union drain so the tasks overlap instead of settling one after another.
+        union: list[str] = []
+        seen: set[str] = set()
+        for t in tasks:
+            if not isinstance(t, Task) or t._settled:
+                continue
+            for rid in t._rids:
+                if rid not in seen:
+                    seen.add(rid)
+                    union.append(rid)
+        if union:
+            self._drain_until(union)
+        return [self._await_task(t) for t in tasks]
 
     # ---- deterministic retrieval (no sub-LLM calls, no root tokens) -----------------------
 
@@ -573,12 +913,11 @@ class Worker:
 
     # ---- one-line delegation (structural: orchestrating must be easier than solving) -------
 
-    def _map_files(self, files: Any, prompt: str, model: str | None = None) -> dict[str, str]:
-        """Ask `prompt` of every given file, batched, and return {path: answer}.
+    def _start_map_files(self, files: Any, prompt: str, model: str | None = None) -> Task:
+        """Post every batch map_files needs, WITHOUT waiting. Contract: see _map_files.
 
-        `files` accepts context entries (dicts), paths (strings), or a mix — the whole
-        chunk/batch/collect loop the system prompt used to spell out, as one call.
-        Oversized files are split and their per-chunk answers joined.
+        All batches go on the wire together, so a 100-file map costs one round-trip of latency
+        rather than one per 20 files.
         """
         prompt = str(prompt)
         by_path: list[tuple[str, str]] = []
@@ -595,12 +934,14 @@ class Worker:
                 else:
                     by_path.append((item, ""))
         if not by_path:
-            return {}
+            return Task.resolved(self, "map_files", {})
 
         # Per-file prompt budget; anything larger is chunked and its answers concatenated.
         budget = self.max_prompt_chars - len(prompt) - _CHUNK_HEADER_OVERHEAD - 256
         if budget < 1_000:
-            return {p: "Error: prompt too long to leave room for file content" for p, _ in by_path}
+            return Task.resolved(self, "map_files", {
+                p: "Error: prompt too long to leave room for file content" for p, _ in by_path
+            })
 
         requests: list[str] = []
         spans: list[tuple[str, int]] = []  # (path, number of chunks contributed)
@@ -611,17 +952,24 @@ class Worker:
                 header = f"[file {path}" + (f", part {j + 1}/{len(chunks)}]" if len(chunks) > 1 else "]")
                 requests.append(f"{prompt}\n\n{header}\n{chunk}")
 
-        responses: list[str] = []
+        rids: list[str] = []
+        sizes: list[int] = []
         for i in range(0, len(requests), _MAX_CHUNK_BATCH):
-            responses.extend(self._llm_query_batched(requests[i:i + _MAX_CHUNK_BATCH], model))
+            batch = requests[i:i + _MAX_CHUNK_BATCH]
+            rids.append(self._post("llm_query_batched", {"prompts": batch, "model": model}))
+            sizes.append(len(batch))
+        return Task(self, "map_files", tuple(rids),
+                    _reduce_map_files(sizes, spans), f"{len(by_path)} files")
 
-        out: dict[str, str] = {}
-        cursor = 0
-        for path, count in spans:
-            part = responses[cursor:cursor + count]
-            cursor += count
-            out[path] = part[0] if count == 1 and part else "\n\n".join(part)
-        return out
+    @_spawnable("map_files")
+    def _map_files(self, files: Any, prompt: str, model: str | None = None) -> dict[str, str]:
+        """Ask `prompt` of every given file, batched, and return {path: answer}.
+
+        `files` accepts context entries (dicts), paths (strings), or a mix — the whole
+        chunk/batch/collect loop the system prompt used to spell out, as one call.
+        Oversized files are split and their per-chunk answers joined.
+        """
+        return self._await_task(self._start_map_files(files, prompt, model))
 
     def _llm_map_reduce(
         self,
@@ -657,9 +1005,23 @@ class Worker:
         joined = "\n\n".join(f"[{labels[i]}]\n{a}" for i, a in enumerate(mapped))
         return self._llm_query(f"{reduce_prompt}\n\nPartial answers:\n{joined}", model)
 
+    # ---- sync helpers: await(start(...)), so there is exactly one code path -----------------
+
+    @_spawnable("llm_query")
+    def _llm_query(self, prompt: str, model: str | None = None) -> str:
+        return self._await_task(self._start_llm_query(prompt, model))
+
+    @_spawnable("llm_query_batched")
+    def _llm_query_batched(self, prompts, model: str | None = None) -> list[str]:
+        return self._await_task(self._start_llm_query_batched(prompts, model))
+
+    @_spawnable("llm_query_chunked")
+    def _llm_query_chunked(self, text, prompt: str, model: str | None = None) -> list[str]:
+        return self._await_task(self._start_llm_query_chunked(text, prompt, model))
+
+    @_spawnable("rlm_query")
     def _rlm_query(self, prompt: str, model: str | None = None) -> str:
-        r = self._rpc("rlm_query", {"prompt": str(prompt), "model": model})
-        return f"Error: {r['error']}" if r.get("error") else r.get("response", "")
+        return self._await_task(self._start_rlm_query(prompt, model))
 
     def _ask_user_question(self, questions: list[dict]) -> list[dict]:
         """Present structured questions to the user; blocks until answered.
@@ -866,17 +1228,9 @@ class Worker:
             return response
         return response if isinstance(response, str) else "ok"
 
+    @_spawnable("rlm_query_batched")
     def _rlm_query_batched(self, prompts, model: str | None = None) -> list[str]:
-        prompts = [str(p) for p in prompts]
-        if not prompts:
-            return []
-        r = self._rpc("rlm_query_batched", {"prompts": prompts, "model": model})
-        if r.get("error"):
-            return [f"Error: {r['error']}"] * len(prompts)
-        out = r.get("responses")
-        if not isinstance(out, list) or len(out) != len(prompts):
-            return ["Error: malformed batched response"] * len(prompts)
-        return [s if isinstance(s, str) else f"Error: {s}" for s in out]
+        return self._await_task(self._start_rlm_query_batched(prompts, model))
 
     # ---- context + execution --------------------------------------------------------------
 
@@ -997,6 +1351,13 @@ class Worker:
         for k, v in self.ns.items():
             if k.startswith("_") or _CONTEXT_NAME.match(k) or k in RESERVED or k == "__builtins__":
                 continue
+            # A Task holds a back-reference to this Worker, so dill would happily pickle the
+            # whole process. Top-level guard only: a Task nested inside a list/dict still
+            # falls to the generic `except` below and skips the variable — which is why this
+            # guard is explicit rather than left to that fallback.
+            if isinstance(v, Task):
+                skipped.append(k)
+                continue
             try:
                 blob = s.dumps(v)
                 if len(blob) > MAX_VAR_BYTES:
@@ -1034,6 +1395,8 @@ def main() -> None:
     ap = argparse.ArgumentParser()
     ap.add_argument("--depth", type=int, default=int(os.environ.get("RLM_DEPTH", "1")))
     ap.add_argument("--timeout", type=float, default=float(os.environ.get("RLM_EXEC_TIMEOUT_S", "600")))
+    ap.add_argument("--await-timeout", type=float,
+                    default=float(os.environ.get("RLM_AWAIT_TIMEOUT_S", "600")))
     ap.add_argument("--max-prompt-chars", type=int,
                     default=int(os.environ.get("RLM_MAX_PROMPT_CHARS", "400000")))
     ap.add_argument("--read-only", action="store_true",
@@ -1042,17 +1405,32 @@ def main() -> None:
     args = ap.parse_args()
 
     worker = Worker(depth=args.depth, exec_timeout_s=args.timeout,
-                    max_prompt_chars=args.max_prompt_chars, read_only=args.read_only)
+                    max_prompt_chars=args.max_prompt_chars, read_only=args.read_only,
+                    await_timeout_s=args.await_timeout)
     _send({"id": "_init", "ok": True})
 
-    for raw in _REAL_STDIN:
-        raw = raw.strip()
-        if not raw:
+    while True:
+        # Requests that arrived mid-exec were parked by _pump; replay them before reading.
+        req = worker.take_deferred()
+        if req is None:
+            line = _REAL_STDIN.readline()
+            if not line:
+                return
+            raw = line.strip()
+            if not raw:
+                continue
+            try:
+                req = json.loads(raw)
+            except json.JSONDecodeError as e:
+                _send({"id": "?", "ok": False, "error": f"bad json: {e}"})
+                continue
+        # A task spawned in an earlier exec settling while the worker is idle. Park it for
+        # a later rlm_await; without this it would fall through to "unknown type" and the
+        # result would be lost.
+        if worker.park_reply(req):
             continue
-        try:
-            req = json.loads(raw)
-        except json.JSONDecodeError as e:
-            _send({"id": "?", "ok": False, "error": f"bad json: {e}"})
+        if not isinstance(req, dict):
+            _send({"id": "?", "ok": False, "error": f"expected an object, got {type(req).__name__}"})
             continue
         rid, kind = req.get("id", "?"), req.get("type")
         try: