@hicaru/pi-rlm 0.2.0 → 0.2.1

package/src/sandbox/sandbox-manager.ts

@@ -14,6 +14,8 @@ export interface SandboxManagerConfig {
   readonly python: string;
   readonly sandboxInitTimeoutMs: number;
   readonly maxPromptChars: number;
+  /** Max seconds the worker waits for a host reply while parked in rlm_await. */
+  readonly awaitTimeoutS: number;
   readonly signal?: AbortSignal;
   readonly onSandboxDiscarded?: () => void;
 }
@@ -61,6 +63,7 @@ export class SandboxManager {
       signal: this.config.signal,
       initTimeoutMs: this.config.sandboxInitTimeoutMs,
       maxPromptChars: this.config.maxPromptChars,
+      awaitTimeoutS: this.config.awaitTimeoutS,
       handlers,
     }).then(async (s) => {
       // Load context on first creation if available.
@@ -85,19 +88,19 @@ export class SandboxManager {
    * a promise queue (second call waits for the first, no interleaving). On failure the sandbox
    * is nullified so the next call recreates it (death-recreate).
    */
-  async exec(code: string): Promise<ReplResult> {
-    return this.execQueued(code);
+  async exec(code: string, signal?: AbortSignal): Promise<ReplResult> {
+    return this.execQueued(code, undefined, signal);
   }
 
   /**
    * Execute code after running `setup` inside the serialized execution slot, so per-invocation
    * handler state (emitter, limits, depth) always matches the active REPL run.
    */
-  async execWithSetup(code: string, setup: () => void): Promise<ReplResult> {
-    return this.execQueued(code, setup);
+  async execWithSetup(code: string, setup: () => void, signal?: AbortSignal): Promise<ReplResult> {
+    return this.execQueued(code, setup, signal);
   }
 
-  private async execQueued(code: string, setup?: () => void): Promise<ReplResult> {
+  private async execQueued(code: string, setup?: () => void, signal?: AbortSignal): Promise<ReplResult> {
     if (!this.sandbox) throw new Error("Sandbox not initialized — call getOrCreate first");
 
     // Serialize: queue behind any in-flight execution
@@ -111,7 +114,7 @@ export class SandboxManager {
       const sandbox = this.sandbox;
       if (!sandbox) throw new Error("Sandbox not initialized — previous execution disposed it");
       setup?.();
-      return await sandbox.exec(code);
+      return await sandbox.exec(code, signal);
     } catch (err) {
       // Death-recreate: worker died — nullify so next repl() recreates
       if (this.sandbox) {
@@ -128,6 +131,17 @@ export class SandboxManager {
     }
   }
 
+  /**
+   * Keep the live sandbox's request watchdog from firing while it is legitimately idle.
+   *
+   * The watchdog only refreshes on frames arriving at THIS sandbox, but a detached
+   * rlm_query child does its work in its own sandbox — so without a heartbeat a healthy
+   * long-running child would trip death-recreate and destroy the REPL namespace.
+   */
+  refreshWatchdog(): void {
+    this.sandbox?.refreshWatchdog();
+  }
+
   /** True if the sandbox is alive and not disposed. */
   get isAlive(): boolean {
     return this.sandbox !== null && !this.disposed;

package/src/sandbox/sandbox.ts

@@ -25,6 +25,7 @@ import {
   type WorkerResponse,
 } from "./protocol.ts";
 import { errorMessage, formatError } from "../util/errors.ts";
+import { trace, traceEnabled } from "../util/trace.ts";
 
 /** Result of a host-side library pack requested by `load_library`. */
 export interface LibraryLoadResult {
@@ -38,12 +39,23 @@ export interface LibraryLoadResult {
   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; the interactive and
+ * pipeline handlers are always synchronous within one exec.
+ */
+export interface SubcallOpts {
+  /** Started via `spawn()` — route to session-scoped state, not the current invocation. */
+  readonly detached: boolean;
+}
+
 /** Handlers the bridge installs to service sub-LLM interrupts. Return the reply payload. */
 export interface SubLlmHandlers {
-  llmQuery(prompt: string, model: string | null, depth: number): Promise<string>;
-  llmQueryBatched(prompts: readonly string[], model: string | null, depth: number): Promise<string[]>;
-  rlmQuery(prompt: string, model: string | null, depth: number): Promise<string>;
-  rlmQueryBatched(prompts: readonly string[], model: string | null, depth: number): Promise<string[]>;
+  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[]>;
   advancePhase(phase: string, summary: string | undefined, depth: number): Promise<string>;
   saveArtifact(kind: string, content: string, depth: number): Promise<string>;
   askUserQuestion(questions: readonly AskQuestion[], depth: number): Promise<AskAnswer[]>;
@@ -73,6 +85,11 @@ export interface SandboxOptions {
    * Native repl() data work leaves this false so scratch-file writes still work.
    */
   readonly readOnly?: boolean;
+  /**
+   * Max seconds the worker will wait for a host reply while parked in `_drain_until`
+   * (rlm_await / sync sub-call). Defaults to the worker's own RLM_AWAIT_TIMEOUT_S (600).
+   */
+  readonly awaitTimeoutS?: number;
 }
 
 const WORKER_PATH = join(dirname(fileURLToPath(import.meta.url)), "worker.py");
@@ -161,6 +178,9 @@ export class PythonSandbox {
     if (opts.readOnly) {
       workerArgs.push("--read-only");
     }
+    if (opts.awaitTimeoutS !== undefined) {
+      workerArgs.push("--await-timeout", String(opts.awaitTimeoutS));
+    }
     this.proc = spawn(
       python,
       workerArgs,
@@ -231,8 +251,8 @@ export class PythonSandbox {
     return file;
   }
 
-  async exec(code: string): Promise<ReplResult> {
-    const res = await this.request({ type: "exec", code });
+  async exec(code: string, signal?: AbortSignal): Promise<ReplResult> {
+    const res = await this.request({ type: "exec", code }, signal);
     if (!res.ok) throw new Error(res.error ?? "exec failed");
     return {
       stdout: res.stdout ?? "",
@@ -297,12 +317,27 @@ export class PythonSandbox {
     });
   }
 
-  private request(payload: RequestBody): Promise<WorkerResponse> {
+  private request(payload: RequestBody, signal?: AbortSignal): Promise<WorkerResponse> {
     if (this.disposed) return Promise.reject(new Error("sandbox disposed"));
+    if (signal?.aborted) return Promise.reject(new Error("repl execution aborted"));
     const id = `r${++this.seq}`;
     return new Promise<WorkerResponse>((resolve, reject) => {
       const timer = this.createWatchdog(id, payload.type, reject);
-      this.pending.set(id, { resolve, reject, timer, requestType: payload.type });
+      // Cancel == kill. The worker may be parked inside `_drain_until` with no other way out;
+      // `proc.on("exit") -> failAll` settles this request and SandboxManager's catch recreates
+      // the sandbox. REPL variables are lost — the documented price of interrupting.
+      const onAbort = (): void => {
+        this.pending.delete(id);
+        clearTimeout(timer);
+        try { this.proc.kill("SIGKILL"); } catch { /* already dead */ }
+        reject(new Error("repl execution aborted — REPL variables were reset"));
+      };
+      signal?.addEventListener("abort", onAbort, { once: true });
+      const once = <T>(settle: (value: T) => void) => (value: T): void => {
+        signal?.removeEventListener("abort", onAbort);
+        settle(value);
+      };
+      this.pending.set(id, { resolve: once(resolve), reject: once(reject), timer, requestType: payload.type });
       this.send({ id, ...payload } as ParentMessage);
     });
   }
@@ -323,7 +358,23 @@ export class PythonSandbox {
     }
   }
 
+  /**
+   * Refresh the parent-side request watchdog for every pending request.
+   * Used during long mid-exec work that does not
+   * produce additional worker interrupts on this sandbox.
+   */
+  refreshWatchdog(): void {
+    this.touchPending();
+  }
+
   private send(msg: ParentMessage): void {
+    if (traceEnabled) {
+      trace("frame.out", {
+        frame: msg.type,
+        id: "id" in msg ? msg.id : undefined,
+        rid: "rid" in msg ? msg.rid : undefined,
+      });
+    }
     this.proc.stdin.write(`${JSON.stringify(msg)}\n`);
   }
 
@@ -353,6 +404,19 @@ export class PythonSandbox {
   }
 
   private dispatch(msg: WorkerMessage): void {
+    if (traceEnabled) {
+      if (isInterrupt(msg)) {
+        trace("frame.in", {
+          frame: msg.type,
+          rid: msg.rid,
+          depth: msg.depth,
+          detached: msg.detached === true,
+          prompts: "prompts" in msg ? msg.prompts?.length : 1,
+        });
+      } else {
+        trace("frame.in", { frame: "response", id: msg.id, ok: msg.ok });
+      }
+    }
     if (isInterrupt(msg)) {
       this.touchPending();
       void this.serviceInterrupt(msg);
@@ -368,18 +432,19 @@ export class PythonSandbox {
   private async serviceInterrupt(msg: WorkerInterrupt): Promise<void> {
     const h = this.handlers;
     const d = msg.depth;
+    const opts: SubcallOpts = { detached: msg.detached === true };
     try {
       if (msg.type === "llm_query") {
-        const response = await h.llmQuery(msg.prompt ?? "", msg.model ?? null, d);
+        const response = await h.llmQuery(msg.prompt ?? "", msg.model ?? null, d, opts);
         this.reply(msg.rid, { response });
       } else if (msg.type === "rlm_query") {
-        const response = await h.rlmQuery(msg.prompt ?? "", msg.model ?? null, d);
+        const response = await h.rlmQuery(msg.prompt ?? "", msg.model ?? null, d, opts);
         this.reply(msg.rid, { response });
       } else if (msg.type === "llm_query_batched") {
-        const responses = await h.llmQueryBatched(msg.prompts ?? [], msg.model ?? null, d);
+        const responses = await h.llmQueryBatched(msg.prompts ?? [], msg.model ?? null, d, opts);
         this.reply(msg.rid, { responses });
       } else if (msg.type === "rlm_query_batched") {
-        const responses = await h.rlmQueryBatched(msg.prompts ?? [], msg.model ?? null, d);
+        const responses = await h.rlmQueryBatched(msg.prompts ?? [], msg.model ?? null, d, opts);
         this.reply(msg.rid, { responses });
       } else if (msg.type === "advance_phase") {
         const response = await h.advancePhase(msg.phase ?? "", msg.summary, d);