@hicaru/pi-rlm 0.1.9 → 0.2.1

package/src/sandbox/worker.py

@@ -12,13 +12,22 @@ 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
 
 import argparse
+import fnmatch
+import heapq
 import io
 import json
+import math
 import os
 import pickle
 import re
@@ -95,19 +104,53 @@ 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",
         "ask_user_question", "todo",
         "load_library",
         "SHOW_VARS", "answer", "context",
     }
 )
+# NOTE: `answers` and `plan` are deliberately NOT reserved. They are seeded by the scaffold but
+# owned by the model, so they must appear in SHOW_VARS and be captured by snapshots — losing a
+# memoized answer across a resume is exactly the failure the memo exists to prevent.
 # Only the single name `context` is the packed world. Legacy context_N names are filtered out.
 _CONTEXT_NAME = re.compile(r"context(_\d+)?\Z")
 
@@ -134,6 +177,145 @@ def _chunk_text(text: str, chunk_chars: int) -> list[str]:
     return chunks
 
 
+# ---- deterministic retrieval over `context` -----------------------------------------------
+#
+# The RLM paper's trajectories retrieve by having the root model hand-write regex over the
+# context (App. E.1). Frontier models do that well; small/fast models guess keywords badly and
+# the first decomposition attempt disproportionately decides the outcome (paper §5, Fig. 4a).
+# These primitives make retrieval deterministic and token-free: no sub-LLM call, no root tokens
+# spent on printed file bodies — the model gets ranked pointers and decides what to delegate.
+
+_INDEX_WINDOW_LINES = 40       # a window is the retrieval unit: big enough to carry meaning
+_INDEX_MAX_WINDOWS = 20_000    # ceiling so a huge load_library() cannot exhaust worker memory
+_SNIPPET_CHARS = 400
+_GREP_HARD_CAP = 200           # absolute ceiling on returned grep hits, whatever k asks for
+_BM25_K1 = 1.2
+_BM25_B = 0.75
+
+_TOKEN_SPLIT = re.compile(r"[^0-9A-Za-z]+")     # also splits snake_case and paths
+_CAMEL_SPLIT = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")
+
+# Definition-ish lines across the languages this plugin is likely to meet. Deliberately
+# lexical: an outline is an orientation aid, not a parse tree.
+_OUTLINE_LINE = re.compile(
+    r"^\s*(?:"
+    r"(?:export\s+)?(?:default\s+)?(?:async\s+)?(?:function|class|interface|type|enum|struct|impl|trait|namespace)\s+\w+"
+    r"|(?:export\s+)?(?:const|let|var)\s+\w+\s*[:=]\s*(?:async\s*)?(?:function|\(|<)"
+    r"|(?:pub\s+)?(?:async\s+)?fn\s+\w+"
+    r"|def\s+\w+|class\s+\w+"
+    r"|func\s+\w+"
+    r"|#{1,4}\s+\S"
+    r")"
+)
+
+
+def _tokenize(text: str) -> list[str]:
+    """Lowercased alphanumeric runs, plus camelCase parts so `resolveModelId` matches `model id`."""
+    out: list[str] = []
+    for raw in _TOKEN_SPLIT.split(text):
+        if not raw:
+            continue
+        lowered = raw.lower()
+        out.append(lowered)
+        if len(raw) > 3:
+            parts = _CAMEL_SPLIT.split(raw)
+            if len(parts) > 1:
+                for part in parts:
+                    piece = part.lower()
+                    if piece and piece != lowered:
+                        out.append(piece)
+    return out
+
+
+def _context_entries(context: Any) -> list[tuple[str, str]]:
+    """(path, content) pairs for either context shape: list[dict] bundles or a raw string."""
+    if isinstance(context, str):
+        return [("<context>", context)]
+    if not isinstance(context, list):
+        return []
+    out: list[tuple[str, str]] = []
+    for i, item in enumerate(context):
+        if isinstance(item, dict):
+            content = item.get("content", "")
+            out.append((
+                str(item.get("path", f"<context[{i}]>")),
+                content if isinstance(content, str) else str(content),
+            ))
+        elif isinstance(item, str):
+            out.append((f"<context[{i}]>", item))
+    return out
+
+
+class _Bm25Index:
+    """Okapi BM25 over fixed-line windows of `context`. Built lazily, discarded on change."""
+
+    __slots__ = ("paths", "starts", "texts", "postings", "doc_len", "avg_len", "truncated")
+
+    def __init__(self, entries: list[tuple[str, str]]) -> None:
+        self.paths: list[str] = []
+        self.starts: list[int] = []
+        self.texts: list[str] = []
+        self.postings: dict[str, list[tuple[int, int]]] = {}
+        self.doc_len: list[int] = []
+        self.truncated = False
+
+        for path, content in entries:
+            if not content:
+                continue
+            lines = content.split("\n")
+            for start in range(0, len(lines), _INDEX_WINDOW_LINES):
+                if len(self.texts) >= _INDEX_MAX_WINDOWS:
+                    self.truncated = True
+                    break
+                window = "\n".join(lines[start:start + _INDEX_WINDOW_LINES])
+                idx = len(self.texts)
+                self.paths.append(path)
+                self.starts.append(start + 1)
+                self.texts.append(window)
+                terms = _tokenize(window)
+                self.doc_len.append(len(terms))
+                freq: dict[str, int] = {}
+                for term in terms:
+                    freq[term] = freq.get(term, 0) + 1
+                for term, tf in freq.items():
+                    self.postings.setdefault(term, []).append((idx, tf))
+            if self.truncated:
+                break
+
+        total = len(self.doc_len)
+        self.avg_len = (sum(self.doc_len) / total) if total else 1.0
+
+    def query(self, terms: list[str], k: int, path_glob: str | None) -> list[dict[str, Any]]:
+        total = len(self.texts)
+        if total == 0:
+            return []
+        scores: dict[int, float] = {}
+        for term in set(terms):
+            posting = self.postings.get(term)
+            if not posting:
+                continue
+            df = len(posting)
+            idf = math.log(1.0 + (total - df + 0.5) / (df + 0.5))
+            for idx, tf in posting:
+                norm = _BM25_K1 * (1.0 - _BM25_B + _BM25_B * self.doc_len[idx] / self.avg_len)
+                scores[idx] = scores.get(idx, 0.0) + idf * (tf * (_BM25_K1 + 1.0)) / (tf + norm)
+        if path_glob:
+            scores = {i: s for i, s in scores.items() if fnmatch.fnmatch(self.paths[i], path_glob)}
+        if not scores:
+            return []
+        top = heapq.nlargest(k, scores.items(), key=lambda kv: kv[1])
+        out: list[dict[str, Any]] = [None] * len(top)  # type: ignore[list-item]
+        for i, (idx, score) in enumerate(top):
+            text = self.texts[idx]
+            out[i] = {
+                "path": self.paths[idx],
+                "line": self.starts[idx],
+                "score": round(score, 3),
+                "snippet": text[:_SNIPPET_CHARS],
+            }
+        return out
+
+
 class _AnswerDict(dict):
     """`answer` dict; flipping `ready` True captures the final answer for the parent."""
 
@@ -154,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()
 
@@ -174,6 +521,8 @@ class Worker:
         self.ns = {"__builtins__": builtins, "__name__": "__main__"}
         self._context_payload: Any | None = None  # pristine restore for the single `context` var
         self._nudged: set[str] = set()
+        self._index: _Bm25Index | None = None
+        self._index_stamp: tuple[int, int] | None = None  # (id(context), len(context))
         self._restore_scaffold()
 
     def _capture_answer(self, content: Any) -> None:
@@ -187,6 +536,20 @@ 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
+        ns["grep_context"] = self._grep_context
+        ns["outline"] = self._outline
+        # env_tips memo (paper App. C.3): "If a value isn't in `answers`, it doesn't exist."
+        # Re-created only when deleted — contents must survive every turn.
+        if not isinstance(ns.get("answers"), dict):
+            ns["answers"] = {}
+        if not isinstance(ns.get("plan"), dict):
+            ns["plan"] = {}
         ns["advance_phase"] = self._advance_phase
         ns["save_artifact"] = self._save_artifact
         ns["ask_user_question"] = self._ask_user_question
@@ -231,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)
 
-        Returns one answer per chunk, order preserved. No exceptions escape: errors come
-        back as "Error: ..." strings per chunk (same contract as llm_query_batched).
+    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.
+
+        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
@@ -286,26 +715,313 @@ 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) -----------------------
+
+    def _entries(self) -> list[tuple[str, str]]:
+        return _context_entries(self.ns.get("context"))
+
+    def _get_index(self) -> _Bm25Index:
+        """Build the BM25 index on first use; rebuild when `context` was replaced or resized.
+
+        Identity+length is a cheap stamp that catches the two ways context actually changes:
+        load_library() extending the list, and the model re-binding the name. In-place edits
+        that preserve length are not detected — documented, and rare in practice.
+        """
+        ctx = self.ns.get("context")
+        stamp = (id(ctx), len(ctx) if isinstance(ctx, (list, str)) else 0)
+        if self._index is None or self._index_stamp != stamp:
+            self._index = _Bm25Index(self._entries())
+            self._index_stamp = stamp
+        return self._index
 
+    def _search(self, query: str, k: int = 10, path_glob: str | None = None) -> list[dict[str, Any]]:
+        """Rank `context` windows against a natural-language query (BM25).
+
+        Returns [{path, line, score, snippet}] — pointers, not bodies. Follow up by slicing the
+        named files out of `context` and delegating them to llm_query / map_files.
+        """
+        terms = _tokenize(str(query))
+        if not terms:
+            return []
+        try:
+            limit = max(1, min(int(k), 100))
+        except (TypeError, ValueError):
+            limit = 10
+        return self._get_index().query(terms, limit, path_glob)
+
+    def _grep_context(
+        self,
+        pattern: str,
+        k: int = 50,
+        path_glob: str | None = None,
+        before: int = 0,
+        after: int = 0,
+    ) -> dict[str, Any]:
+        """Regex over `context`, capped and shaped.
+
+        Returns {"hits": [{path, line, text}], "counts": {path: n}, "total": n, "truncated": bool}.
+        `counts` is complete even when `hits` is capped, so a wide pattern reports its shape
+        instead of flooding stdout.
+        """
+        try:
+            rx = re.compile(pattern)
+        except re.error as e:
+            return {"hits": [], "counts": {}, "total": 0, "truncated": False, "error": f"bad regex: {e}"}
+        try:
+            limit = max(1, min(int(k), _GREP_HARD_CAP))
+        except (TypeError, ValueError):
+            limit = 50
+        pad_before = max(0, min(int(before or 0), 10))
+        pad_after = max(0, min(int(after or 0), 10))
+
+        hits: list[dict[str, Any]] = []
+        counts: dict[str, int] = {}
+        total = 0
+        for path, content in self._entries():
+            if path_glob and not fnmatch.fnmatch(path, path_glob):
+                continue
+            if not rx.search(content):
+                continue
+            lines = content.split("\n")
+            for i, line in enumerate(lines):
+                if not rx.search(line):
+                    continue
+                total += 1
+                counts[path] = counts.get(path, 0) + 1
+                if len(hits) >= limit:
+                    continue
+                lo = max(0, i - pad_before)
+                hi = min(len(lines), i + pad_after + 1)
+                hits.append({"path": path, "line": i + 1, "text": "\n".join(lines[lo:hi])[:_SNIPPET_CHARS]})
+        return {"hits": hits, "counts": counts, "total": total, "truncated": total > len(hits)}
+
+    def _outline(self, path: str) -> str:
+        """Definition/heading skeleton of one context file — orient in ~200 chars, not 20K.
+
+        `path` matches exactly, then by suffix, then as a glob.
+        """
+        target = str(path)
+        entries = self._entries()
+        content: str | None = None
+        for p, c in entries:
+            if p == target:
+                content = c
+                break
+        if content is None:
+            for p, c in entries:
+                if p.endswith(target) or fnmatch.fnmatch(p, target):
+                    content = c
+                    target = p
+                    break
+        if content is None:
+            return f"Error: no context file matching {path!r} — use search() or list paths from context"
+        out: list[str] = [f"# {target}"]
+        for i, line in enumerate(content.split("\n")):
+            if _OUTLINE_LINE.match(line):
+                out.append(f"{i + 1}: {line.strip()[:160]}")
+        if len(out) == 1:
+            return f"# {target}\n(no definition-like lines found)"
+        return "\n".join(out)
+
+    # ---- one-line delegation (structural: orchestrating must be easier than solving) -------
+
+    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.
+
+        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]] = []
+        lookup: dict[str, str] | None = None
+        for item in files if isinstance(files, (list, tuple)) else [files]:
+            if isinstance(item, dict):
+                content = item.get("content", "")
+                by_path.append((str(item.get("path", "?")), content if isinstance(content, str) else str(content)))
+            elif isinstance(item, str):
+                if lookup is None:
+                    lookup = {p: c for p, c in self._entries()}
+                if item in lookup:
+                    by_path.append((item, lookup[item]))
+                else:
+                    by_path.append((item, ""))
+        if not by_path:
+            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 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)
+        for path, content in by_path:
+            chunks = _chunk_text(content, budget) if len(content) > budget else [content]
+            spans.append((path, len(chunks)))
+            for j, chunk in enumerate(chunks):
+                header = f"[file {path}" + (f", part {j + 1}/{len(chunks)}]" if len(chunks) > 1 else "]")
+                requests.append(f"{prompt}\n\n{header}\n{chunk}")
+
+        rids: list[str] = []
+        sizes: list[int] = []
+        for i in range(0, len(requests), _MAX_CHUNK_BATCH):
+            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")
+
+    @_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,
+        items: Any,
+        map_prompt: str,
+        reduce_prompt: str,
+        model: str | None = None,
+    ) -> str:
+        """Map `map_prompt` over `items` in one batch, then reduce the answers with one call.
+
+        The paper's canonical strategy ("query an LLM per chunk ... then query an LLM with all
+        the buffers") as a single call, so the root never hand-rolls the loop.
+        """
+        map_prompt, reduce_prompt = str(map_prompt), str(reduce_prompt)
+        seq = list(items) if isinstance(items, (list, tuple)) else [items]
+        if not seq:
+            return "Error: llm_map_reduce got no items"
+        texts = [
+            (str(it.get("content", "")) if isinstance(it, dict) else str(it))
+            for it in seq
+        ]
+        labels = [
+            (str(it.get("path", f"item {i + 1}")) if isinstance(it, dict) else f"item {i + 1}")
+            for i, it in enumerate(seq)
+        ]
+        mapped: list[str] = []
+        for i in range(0, len(texts), _MAX_CHUNK_BATCH):
+            batch = [
+                f"{map_prompt}\n\n[{labels[i + j]}]\n{t}"
+                for j, t in enumerate(texts[i:i + _MAX_CHUNK_BATCH])
+            ]
+            mapped.extend(self._llm_query_batched(batch, model))
+        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.
@@ -512,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 --------------------------------------------------------------
 
@@ -643,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:
@@ -680,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",
@@ -688,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: