@hicaru/pi-rlm 0.1.9 → 0.2.1

package/src/sandbox/sandbox.ts

@@ -24,7 +24,8 @@ import {
   type WorkerRequest,
   type WorkerResponse,
 } from "./protocol.ts";
-import { formatError } from "../util/errors.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,9 +85,15 @@ 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");
+const STDERR_TAIL_CHARS = 8_192;
 const TODO_PROTO_KEYS = new Set(["type", "rid", "depth", "action"]);
 
 // The sandbox runs untrusted model-authored code; it must never inherit provider secrets.
@@ -124,7 +142,23 @@ export class PythonSandbox {
   private readonly handlers: SubLlmHandlers;
   private readonly requestTimeoutMs: number;
   private readonly initTimeoutMs: number;
-  private stderr = "";
+  /** Bounded stderr tail (chunks, newest last) — avoids rebuilding the buffer per chunk. */
+  private readonly stderrTail: string[] = [];
+  private stderrLen = 0;
+
+  /** Bounded tail of everything written to stderr, oldest chunks already dropped. */
+  private get stderr(): string {
+    return this.stderrTail.join("");
+  }
+
+  /** Record a diagnostic on the same bounded tail as real worker stderr. */
+  private appendStderr(text: string): void {
+    this.stderrTail.push(text);
+    this.stderrLen += text.length;
+    while (this.stderrLen > STDERR_TAIL_CHARS && this.stderrTail.length > 1) {
+      this.stderrLen -= (this.stderrTail.shift() ?? "").length;
+    }
+  }
   private disposed = false;
   private ready: Promise<void>;
 
@@ -144,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,
@@ -153,9 +190,7 @@ export class PythonSandbox {
     this.proc.stdout.setEncoding("utf8");
     this.proc.stdout.on("data", (chunk: string) => this.onData(chunk));
     this.proc.stderr.setEncoding("utf8");
-    this.proc.stderr.on("data", (chunk: string) => {
-      this.stderr = (this.stderr + chunk).slice(-8192);
-    });
+    this.proc.stderr.on("data", (chunk: string) => this.appendStderr(chunk));
     this.proc.on("error", (err: NodeJS.ErrnoException) => {
       const hint = err.code === "ENOENT" ? ` ('${python}' not found — is Python installed and on PATH?)` : "";
       this.failAll(new Error(`failed to start sandbox${hint}: ${err.message}`));
@@ -189,12 +224,12 @@ export class PythonSandbox {
     return sandbox;
   }
 
-  async loadContext(payload: unknown, index?: number): Promise<number> {
+  async loadContext(payload: unknown): Promise<number> {
     const isJson = typeof payload !== "string";
     let path: string | undefined;
     try {
       path = await this.writeContextFile(payload, isJson);
-      const res = await this.request({ type: "load_context", path, index, json: isJson });
+      const res = await this.request({ type: "load_context", path, json: isJson });
       if (!res.ok) throw new Error(res.error ?? "load_context failed");
       return res.index ?? 0;
     } finally {
@@ -216,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 ?? "",
@@ -282,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);
     });
   }
@@ -318,6 +368,13 @@ export class PythonSandbox {
   }
 
   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`);
   }
 
@@ -331,11 +388,11 @@ export class PythonSandbox {
         try {
           const message = JSON.parse(line) as unknown;
           if (isWorkerMessage(message)) this.dispatch(message);
-          else this.stderr = `${this.stderr}\n[protocol] skipped invalid stdout message: ${line.slice(0, 200)}`.slice(-8192);
+          else this.appendStderr(`\n[protocol] skipped invalid stdout message: ${line.slice(0, 200)}`);
         } catch {
           // Non-JSON line on the protocol stream — likely a subprocess writing to fd 1.
           // Skip it so a rogue write doesn't kill the pump, but retain a breadcrumb for watchdog errors.
-          this.stderr = `${this.stderr}\n[protocol] skipped non-JSON stdout line: ${line.slice(0, 200)}`.slice(-8192);
+          this.appendStderr(`\n[protocol] skipped non-JSON stdout line: ${line.slice(0, 200)}`);
         }
       }
     }
@@ -347,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);
@@ -362,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);
@@ -419,7 +490,7 @@ export class PythonSandbox {
         }
       }
     } catch (err) {
-      this.reply(msg.rid, { error: err instanceof Error ? err.message : String(err) });
+      this.reply(msg.rid, { error: errorMessage(err) });
     }
   }