@hicaru/pi-rlm 0.2.1 → 0.3.0

package/src/sandbox/interrupts.ts

@@ -0,0 +1,160 @@
+/**
+ * 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 pack requested by `add_context`. */
+export interface AddContextResult {
+  readonly payload: unknown;     // always ContextFile[] under ctx/<id>/ (or un-prefixed for cwd)
+  readonly files?: number;
+  readonly chars: number;
+  readonly sourceId: string;
+  readonly pathPrefix: string;
+  /** Host already has this source — no pack, empty payload. */
+  readonly alreadyLoaded?: boolean;
+  /** Document-type files in the payload (fresh + cache hits). */
+  readonly documents?: number;
+  /** Documents freshly converted this call (cache hits excluded). */
+  readonly converted?: number;
+  /** Paths skipped during packing (model-facing). */
+  readonly skipped?: readonly { readonly path: string; readonly reason: string }[];
+}
+
+/**
+ * Per-interrupt routing context for the sub-LLM handlers.
+ *
+ * Only the four sub-call kinds can be spawned, so only they carry it; add_context 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[]>;
+  addContext(source: string, depth: number): Promise<AddContextResult>;
+}
+
+/** 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")),
+  addContext: async () => { throw new Error("add_context 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;
+  documents?: number;
+  converted?: number;
+  skipped?: readonly { readonly path: string; readonly reason: string }[];
+  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 === "add_context") {
+      const lib = await h.addContext(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,
+          documents: lib.documents ?? 0,
+          converted: lib.converted ?? 0,
+          skipped: lib.skipped,
+        });
+      } else {
+        const { path, json: isJson } = await writeContextTempFile(lib.payload);
+        // Worker reads then unlinks (worker._add_context). 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,
+          documents: lib.documents ?? 0,
+          converted: lib.converted ?? 0,
+          skipped: lib.skipped,
+        });
+      }
+    }
+  } 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,17 +18,21 @@ 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). */
+  /** add_context 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;
   readonly path_prefix?: string;
-  /** Host-side idempotency: library already loaded — no path payload. */
+  /** Host-side idempotency: source already loaded — no path payload. */
   readonly already_loaded?: boolean;
+  /** Document-type files in the payload (fresh + cache hits). */
+  readonly documents?: number;
+  /** Documents freshly converted this call (cache hits excluded). */
+  readonly converted?: number;
+  /** Paths skipped during packing (binary, no-converter, …). */
+  readonly skipped?: readonly { readonly path: string; readonly reason: string }[];
   readonly error?: string;
 }
 
@@ -60,30 +62,7 @@ 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;
-}
+  | "add_context";
 
 interface InterruptBase {
   readonly rid: string;
@@ -100,49 +79,23 @@ 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;
+  /** `rlm_query_batched` only — one prefix set shared by every prompt in the batch. */
+  readonly paths?: readonly string[];
 }
 
-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;
-}
-
-export interface LoadLibraryInterrupt extends InterruptBase {
-  readonly type: "load_library";
+export interface AddContextInterrupt extends InterruptBase {
+  readonly type: "add_context";
   readonly source?: string;
 }
 
@@ -150,11 +103,7 @@ export interface LoadLibraryInterrupt extends InterruptBase {
 export type WorkerInterrupt =
   | PromptInterrupt
   | BatchedPromptInterrupt
-  | AdvancePhaseInterrupt
-  | SaveArtifactInterrupt
-  | AskUserQuestionInterrupt
-  | TodoInterrupt
-  | LoadLibraryInterrupt;
+  | AddContextInterrupt;
 
 export type WorkerMessage = WorkerResponse | WorkerInterrupt;
 
@@ -163,11 +112,7 @@ 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",
+  "add_context",
 ]));
 
 function isRecord(value: unknown): value is Record<string, unknown> {

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",
+        "add_context",
+        "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}"