@hicaru/pi-rlm 0.2.1 → 0.2.2

package/src/sandbox/context-file.ts

@@ -0,0 +1,154 @@
+/**
+ * Temp-file transport for sandbox context payloads, with refcounted sharing.
+ *
+ * Two callers, two ownership models, one writer:
+ *  - `writeContextTempFile` — non-owning. `load_library` uses it because the WORKER unlinks
+ *    that file after reading it (see sandbox.ts serviceInterrupt / worker.py `_load_library`).
+ *  - `pinContext` — refcounted. Every child RLM of one node inherits the SAME payload, so an
+ *    18-way fan-out would otherwise cost 18 serializations and 18 files. Pins are keyed by
+ *    payload identity, which is a free version key: `mergeLibraryIntoContext` always returns a
+ *    NEW array, so loading a library mints a new key and old holders keep their own file.
+ *
+ * Serialization is chunked with an await between chunks so the event loop is never blocked for
+ * more than ~SERIALIZE_CHUNK entries. A Worker Thread was considered and rejected: posting the
+ * payload structured-clones the whole array, which costs about what the stringify costs and
+ * doubles peak RSS.
+ */
+
+import { open, unlink, type FileHandle } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+/** Entries serialized per await. Bounds the longest synchronous span on the event loop. */
+const SERIALIZE_CHUNK = 64;
+
+/** A temp file on disk holding a serialized context payload. */
+export interface ContextTempFile {
+  readonly path: string;
+  /** True when the file holds JSON; false when the payload was a raw string. */
+  readonly json: boolean;
+}
+
+/** A shared, refcounted context file. Every holder must `release()` exactly once. */
+export interface PinnedContext extends ContextTempFile {
+  /** Drop this holder's reference; unlinks once the last holder releases. Idempotent. */
+  release(): Promise<void>;
+}
+
+interface PinEntry extends ContextTempFile {
+  refs: number;
+}
+
+/**
+ * Live pins keyed by payload identity, storing the in-flight PROMISE rather than the settled
+ * entry. Children of one node race here (each drives its own sandbox, so nothing else
+ * serializes them); inserting the promise before the first await makes them join one write
+ * instead of each starting their own and orphaning the loser's file.
+ */
+const pins = new Map<unknown, Promise<PinEntry>>();
+
+function tempPath(isJson: boolean): string {
+  const suffix = isJson ? "json" : "txt";
+  return join(
+    tmpdir(),
+    `rlm-ctx-${process.pid}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}.${suffix}`,
+  );
+}
+
+/** Write a JSON array incrementally, yielding to the event loop between chunks. */
+async function writeChunkedArray(handle: FileHandle, items: readonly unknown[]): Promise<void> {
+  await handle.write("[");
+  const buf = new Array<string>(SERIALIZE_CHUNK);
+  let n = 0;
+  for (let i = 0; i < items.length; i++) {
+    // Comma prefix beats trimming a trailing one; no `+=` accumulation anywhere.
+    buf[n++] = i === 0 ? JSON.stringify(items[i]) : `,${JSON.stringify(items[i])}`;
+    if (n === SERIALIZE_CHUNK) {
+      await handle.write(buf.join("")); // the await is the yield point
+      n = 0;
+    }
+  }
+  if (n > 0) await handle.write(buf.slice(0, n).join(""));
+  await handle.write("]");
+}
+
+/**
+ * Serialize a payload to a fresh temp file. The caller owns the file and decides when (or
+ * whether) to unlink it.
+ */
+export async function writeContextTempFile(payload: unknown): Promise<ContextTempFile> {
+  const json = typeof payload !== "string";
+  const path = tempPath(json);
+  try {
+    const handle = await open(path, "w");
+    try {
+      if (typeof payload === "string") await handle.write(payload);
+      else if (Array.isArray(payload)) await writeChunkedArray(handle, payload);
+      else await handle.write(JSON.stringify(payload));
+    } finally {
+      await handle.close();
+    }
+  } catch (err) {
+    await unlink(path).catch(() => {});
+    throw err;
+  }
+  return Object.freeze({ path, json });
+}
+
+async function writePinEntry(payload: unknown): Promise<PinEntry> {
+  const file = await writeContextTempFile(payload);
+  return { path: file.path, json: file.json, refs: 1 };
+}
+
+/** One holder's view of a pin. `shared` entries are evicted from the map at refcount zero. */
+function handleFor(key: unknown, entry: PinEntry, shared: boolean): PinnedContext {
+  let released = false;
+  return Object.freeze({
+    path: entry.path,
+    json: entry.json,
+    release: async (): Promise<void> => {
+      if (released) return; // idempotent per handle, so a `finally` cannot double-decrement
+      released = true;
+      entry.refs -= 1;
+      if (entry.refs > 0) return;
+      if (shared) pins.delete(key);
+      await unlink(entry.path).catch(() => {});
+    },
+  });
+}
+
+/**
+ * Acquire a shared context file for `payload`. Holders must `release()` exactly once; the file
+ * is unlinked when the last one does.
+ */
+export async function pinContext(payload: unknown): Promise<PinnedContext> {
+  // Only arrays are shared. Their identity is a meaningful version key; a string's is not
+  // (two equal strings may or may not be the same reference), and the string payloads here are
+  // one-off child prompts with nothing to share anyway.
+  if (!Array.isArray(payload)) {
+    return handleFor(payload, await writePinEntry(payload), false);
+  }
+
+  const existing = pins.get(payload);
+  if (existing !== undefined) {
+    const entry = await existing;
+    entry.refs += 1;
+    return handleFor(payload, entry, true);
+  }
+
+  // Insert synchronously, BEFORE any await, so a concurrent caller sees this write in flight.
+  const pending = writePinEntry(payload);
+  pins.set(payload, pending);
+  try {
+    return handleFor(payload, await pending, true);
+  } catch (err) {
+    // Evict the rejected promise so a later caller retries instead of awaiting a poisoned pin.
+    pins.delete(payload);
+    throw err;
+  }
+}
+
+/** Live pin count. Exported for tests asserting the sharing and the unlink-once behaviour. */
+export function pinnedCount(): number {
+  return pins.size;
+}

package/src/sandbox/interrupts.ts

@@ -0,0 +1,145 @@
+/**
+ * The interrupt surface: what the worker can ask the host for mid-exec, and how each request is
+ * turned into a reply frame.
+ *
+ * Split from sandbox.ts, which owns the subprocess and the JSONL pump. Adding a sandbox function
+ * touches this file and worker.py; the transport underneath does not change.
+ */
+
+import type { WorkerInterrupt } from "./protocol.ts";
+import { writeContextTempFile } from "./context-file.ts";
+import { errorMessage, formatError } from "../util/errors.ts";
+
+/** Result of a host-side library pack requested by `load_library`. */
+export interface LibraryLoadResult {
+  readonly payload: unknown;     // always ContextFile[] under lib/<id>/
+  readonly files?: number;
+  readonly chars: number;
+  readonly sourceId: string;
+  readonly pathPrefix: string;
+  /** Host already has this library — no pack, empty payload. */
+  readonly alreadyLoaded?: boolean;
+}
+
+/**
+ * Per-interrupt routing context for the sub-LLM handlers.
+ *
+ * Only the four sub-call kinds can be spawned, so only they carry it; load_library is
+ * always synchronous within one exec.
+ */
+export interface SubcallOpts {
+  /** Started via `spawn()` — route to session-scoped state, not the current invocation. */
+  readonly detached: boolean;
+  /**
+   * `rlm_query(paths=[…])` — path prefixes narrowing the child's inherited context.
+   * Absent on every other path; `llm_query` never carries it.
+   */
+  readonly paths?: readonly string[];
+}
+
+/** Handlers the bridge installs to service sub-LLM interrupts. Return the reply payload. */
+export interface SubLlmHandlers {
+  llmQuery(prompt: string, model: string | null, depth: number, opts: SubcallOpts): Promise<string>;
+  llmQueryBatched(prompts: readonly string[], model: string | null, depth: number, opts: SubcallOpts): Promise<string[]>;
+  rlmQuery(prompt: string, model: string | null, depth: number, opts: SubcallOpts): Promise<string>;
+  rlmQueryBatched(prompts: readonly string[], model: string | null, depth: number, opts: SubcallOpts): Promise<string[]>;
+  loadLibrary(source: string, depth: number): Promise<LibraryLoadResult>;
+}
+
+/** Narrow an unknown JSON value to a frozen string array. Non-strings and blanks are dropped. */
+function toStringArray(value: unknown): readonly string[] | undefined {
+  if (!Array.isArray(value)) return undefined;
+  const out = new Array<string>(value.length);
+  let n = 0;
+  for (let i = 0; i < value.length; i++) {
+    const item: unknown = value[i];
+    if (typeof item === "string" && item.trim() !== "") out[n++] = item;
+  }
+  out.length = n;
+  return n > 0 ? Object.freeze(out) : undefined;
+}
+
+/** Default handlers — every sandbox function refuses until a bridge installs a real one. */
+export const REJECT: SubLlmHandlers = {
+  llmQuery: async () => formatError("sub-LLM bridge not configured"),
+  llmQueryBatched: async (p) => p.map(() => formatError("sub-LLM bridge not configured")),
+  rlmQuery: async () => formatError("sub-LLM bridge not configured"),
+  rlmQueryBatched: async (p) => p.map(() => formatError("sub-LLM bridge not configured")),
+  loadLibrary: async () => { throw new Error("load_library not configured"); },
+};
+
+/** Body of a reply frame — the union of every handler's payload shape. */
+export interface ReplyBody {
+  response?: string;
+  responses?: string[];
+  path?: string;
+  json?: boolean;
+  files?: number;
+  chars?: number;
+  source_id?: string;
+  path_prefix?: string;
+  already_loaded?: boolean;
+  error?: string;
+}
+
+/**
+ * Service one interrupt and hand the reply body to `reply`.
+ *
+ * Errors are replied, never thrown: the caller invokes this from the stdio pump, where a
+ * rejection would surface as an unhandled promise and leave the worker parked forever.
+ */
+export async function serviceInterrupt(
+  msg: WorkerInterrupt,
+  h: SubLlmHandlers,
+  reply: (rid: string, body: ReplyBody) => void,
+): Promise<void> {
+  const d = msg.depth;
+  const opts: SubcallOpts = Object.freeze({
+    detached: msg.detached === true,
+    // Only the recursive kinds carry a context slice; the value crossed JSON, so guard it.
+    paths: msg.type === "rlm_query" || msg.type === "rlm_query_batched"
+      ? toStringArray(msg.paths)
+      : undefined,
+  });
+  try {
+    if (msg.type === "llm_query") {
+      const response = await h.llmQuery(msg.prompt ?? "", msg.model ?? null, d, opts);
+      reply(msg.rid, { response });
+    } else if (msg.type === "rlm_query") {
+      const response = await h.rlmQuery(msg.prompt ?? "", msg.model ?? null, d, opts);
+      reply(msg.rid, { response });
+    } else if (msg.type === "llm_query_batched") {
+      const responses = await h.llmQueryBatched(msg.prompts ?? [], msg.model ?? null, d, opts);
+      reply(msg.rid, { responses });
+    } else if (msg.type === "rlm_query_batched") {
+      const responses = await h.rlmQueryBatched(msg.prompts ?? [], msg.model ?? null, d, opts);
+      reply(msg.rid, { responses });
+    } else if (msg.type === "load_library") {
+      const lib = await h.loadLibrary(msg.source ?? "", d);
+      if (lib.alreadyLoaded) {
+        // No temp file — worker short-circuits on already_loaded.
+        reply(msg.rid, {
+          already_loaded: true,
+          files: 0,
+          chars: lib.chars,
+          source_id: lib.sourceId,
+          path_prefix: lib.pathPrefix,
+        });
+      } else {
+        const { path, json: isJson } = await writeContextTempFile(lib.payload);
+        // Worker reads then unlinks (worker._load_library). Host must not unlink here —
+        // if the worker is SIGKILLed before os.remove, the temp file leaks in tmpdir (acceptable).
+        reply(msg.rid, {
+          path,
+          json: isJson,
+          files: lib.files,
+          chars: lib.chars,
+          source_id: lib.sourceId,
+          path_prefix: lib.pathPrefix,
+        });
+      }
+    }
+  } catch (err) {
+    reply(msg.rid, { error: errorMessage(err) });
+  }
+}

package/src/sandbox/protocol.ts

@@ -10,8 +10,6 @@
 export type WorkerRequest =
   | { readonly id: string; readonly type: "exec"; readonly code: string }
   | { readonly id: string; readonly type: "load_context"; readonly path: string; readonly index?: number; readonly json: boolean }
-  | { readonly id: string; readonly type: "snapshot"; readonly path: string; readonly nonce: string }
-  | { readonly id: string; readonly type: "restore"; readonly path: string; readonly nonce: string }
   | { readonly id: string; readonly type: "shutdown" };
 
 /** Reply the parent sends to satisfy a sub-LLM interrupt. */
@@ -20,11 +18,9 @@ export interface LlmReply {
   readonly rid: string;
   readonly response?: string;
   readonly responses?: readonly string[];
-  readonly answers?: readonly AskAnswer[];
-  /** load_library reply: temp file with the packed payload (+ resume index / namespace). */
+  /** load_library reply: temp file with the packed payload (+ namespace metadata). */
   readonly path?: string;
   readonly json?: boolean;
-  readonly index?: number;
   readonly files?: number;
   readonly chars?: number;
   readonly source_id?: string;
@@ -60,31 +56,8 @@ export type InterruptKind =
   | "llm_query_batched"
   | "rlm_query"
   | "rlm_query_batched"
-  | "advance_phase"
-  | "save_artifact"
-  | "ask_user_question"
-  | "todo"
   | "load_library";
 
-export interface AskOption {
-  readonly label: string;
-  readonly description?: string;
-  readonly preview?: string;
-}
-
-export interface AskQuestion {
-  readonly question: string;
-  readonly header: string;
-  readonly multiSelect?: boolean;
-  readonly options: readonly AskOption[];
-}
-
-export interface AskAnswer {
-  readonly question: string;
-  readonly selected: readonly string[];
-  readonly custom?: string;
-}
-
 interface InterruptBase {
   readonly rid: string;
   readonly depth: number;
@@ -100,45 +73,19 @@ interface PromptInterrupt extends InterruptBase {
   readonly type: "llm_query" | "rlm_query";
   readonly prompt?: string;
   readonly model?: string | null;
+  /**
+   * `rlm_query` only — path prefixes narrowing the child's inherited context. Never sent for
+   * `llm_query`, whose frame stays byte-identical to before.
+   */
+  readonly paths?: readonly string[];
 }
 
 interface BatchedPromptInterrupt extends InterruptBase {
   readonly type: "llm_query_batched" | "rlm_query_batched";
   readonly prompts?: readonly string[];
   readonly model?: string | null;
-}
-
-interface AdvancePhaseInterrupt extends InterruptBase {
-  readonly type: "advance_phase";
-  readonly phase?: string;
-  readonly summary?: string;
-}
-
-interface SaveArtifactInterrupt extends InterruptBase {
-  readonly type: "save_artifact";
-  readonly artifactKind?: string;
-  readonly content?: string;
-}
-
-export interface AskUserQuestionInterrupt extends InterruptBase {
-  readonly type: "ask_user_question";
-  readonly questions: readonly AskQuestion[];
-}
-
-export interface TodoInterrupt extends InterruptBase {
-  readonly type: "todo";
-  readonly action: "create" | "update" | "list" | "get" | "delete" | "clear";
-  readonly id?: number;
-  readonly subject?: string;
-  readonly description?: string;
-  readonly status?: "pending" | "in_progress" | "completed" | "deleted";
-  readonly activeForm?: string;
-  readonly blockedBy?: readonly number[];
-  readonly addBlockedBy?: readonly number[];
-  readonly removeBlockedBy?: readonly number[];
-  readonly owner?: string;
-  readonly filterStatus?: string;
-  readonly includeDeleted?: boolean;
+  /** `rlm_query_batched` only — one prefix set shared by every prompt in the batch. */
+  readonly paths?: readonly string[];
 }
 
 export interface LoadLibraryInterrupt extends InterruptBase {
@@ -150,10 +97,6 @@ export interface LoadLibraryInterrupt extends InterruptBase {
 export type WorkerInterrupt =
   | PromptInterrupt
   | BatchedPromptInterrupt
-  | AdvancePhaseInterrupt
-  | SaveArtifactInterrupt
-  | AskUserQuestionInterrupt
-  | TodoInterrupt
   | LoadLibraryInterrupt;
 
 export type WorkerMessage = WorkerResponse | WorkerInterrupt;
@@ -163,10 +106,6 @@ export const INTERRUPT_KINDS = Object.freeze(new Set<InterruptKind>([
   "llm_query_batched",
   "rlm_query",
   "rlm_query_batched",
-  "advance_phase",
-  "save_artifact",
-  "ask_user_question",
-  "todo",
   "load_library",
 ]));
 

package/src/sandbox/py/guards.py

@@ -0,0 +1,150 @@
+"""Execution guardrails for the RLM sandbox worker.
+
+The restricted builtin table, the reserved-name set that keeps scaffold functions out of
+SHOW_VARS, the protocol writer that must always reach the REAL stdout (user prints are
+captured into a buffer), and the per-exec stall alarm.
+
+This is steering, not a security boundary: `__import__` and `open` are deliberately available,
+so model code can still reach the network and the filesystem. What it does buy is that the
+scaffold cannot be clobbered silently and that a blocked builtin explains itself instead of
+failing as "'NoneType' object is not callable".
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import signal
+import sys
+from contextlib import contextmanager
+from typing import Any
+
+# Capture the REAL stdio before exec() redirects sys.stdout/sys.stderr into buffers.
+# All protocol writes must go to the real stdout even while user code's prints are captured.
+REAL_STDOUT = sys.stdout
+REAL_STDIN = sys.stdin
+REAL_STDERR = sys.stderr
+
+
+def _builtin(name: str):
+    return __builtins__[name] if isinstance(__builtins__, dict) else getattr(__builtins__, name, None)
+
+
+# Restricted builtins: enough for real data work, minus the dangerous reflection escapes.
+_SAFE_BUILTINS = {
+    name: _builtin(name)
+    for name in (
+        "abs", "all", "any", "ascii", "bin", "bool", "bytearray", "bytes", "callable",
+        "chr", "classmethod", "complex", "dict", "dir", "divmod", "enumerate", "filter",
+        "float", "format", "frozenset", "getattr", "hasattr", "hash", "hex", "id", "int",
+        "isinstance", "issubclass", "iter", "len", "list", "map", "max", "min", "next",
+        "object", "oct", "ord", "pow", "print", "property", "range", "repr", "reversed",
+        "round", "set", "setattr", "slice", "sorted", "staticmethod", "str", "sum", "super",
+        "tuple", "type", "vars", "zip", "delattr", "memoryview", "__import__", "__build_class__",
+        "Exception", "BaseException", "ValueError", "TypeError", "KeyError", "IndexError",
+        "AttributeError", "FileNotFoundError", "OSError", "IOError", "RuntimeError",
+        "NameError", "ImportError", "StopIteration", "AssertionError", "NotImplementedError",
+        "ArithmeticError", "ZeroDivisionError", "LookupError", "Warning", "True", "False", "None",
+    )
+}
+# `open` is allowed for data work; eval/exec/compile/input/globals/locals are not.
+# _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] = _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",
+        "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.
+# Only the single name `context` is the packed world. Legacy context_N names are filtered out.
+_CONTEXT_NAME = re.compile(r"context(_\d+)?\Z")
+
+def _send(obj: dict[str, Any]) -> None:
+    REAL_STDOUT.write(json.dumps(obj, ensure_ascii=False) + "\n")
+    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}"