@hicaru/pi-rlm 0.2.0 → 0.2.2

package/src/sandbox/sandbox-manager.ts

@@ -6,6 +6,7 @@
 
 import { PythonSandbox, type SubLlmHandlers } from "./sandbox.ts";
 import type { ReplResult } from "./protocol.ts";
+import { mergeLibraryIntoContext } from "../context/library-context.ts";
 
 /** Static configuration for sandbox creation — set once, reused across getOrCreate calls. */
 export interface SandboxManagerConfig {
@@ -14,6 +15,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;
 }
@@ -32,6 +35,18 @@ export class SandboxManager {
 
   constructor(private readonly config: SandboxManagerConfig) {}
 
+  /**
+   * Append a library payload to the context this manager replays on death-recreate.
+   *
+   * The live worker has ALREADY appended it in-process (worker.py `_append_library`), so this
+   * deliberately does not reload — it only keeps the host's replay copy truthful. Without it a
+   * recreate silently rolls the sandbox back to a repo-only context, and any child inheriting
+   * this payload would never see the library. Dedups by `lib/<id>/` prefix.
+   */
+  appendLibrary(payload: unknown): void {
+    this.contextPayload = mergeLibraryIntoContext(this.contextPayload, payload);
+  }
+
   /**
    * Lazy get-or-create the sandbox. On first call, spawns PythonSandbox with the
    * given handlers. Subsequent calls return the existing sandbox immediately.
@@ -61,6 +76,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 +101,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 +127,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 +144,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

@@ -8,48 +8,23 @@
  */
 
 import { type ChildProcessWithoutNullStreams, spawn } from "node:child_process";
-import { writeFile, unlink } from "node:fs/promises";
-import { tmpdir } from "node:os";
+import { once } from "node:events";
 import { dirname, join } from "node:path";
 import { fileURLToPath } from "node:url";
 import {
   isInterrupt,
   isWorkerMessage,
-  type AskAnswer,
-  type AskQuestion,
   type ParentMessage,
   type ReplResult,
-  type WorkerInterrupt,
   type WorkerMessage,
   type WorkerRequest,
   type WorkerResponse,
 } from "./protocol.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 index: number;        // resume-sidecar index (not a REPL var name); -1 if alreadyLoaded
-  readonly files?: number;
-  readonly chars: number;
-  readonly sourceId: string;
-  readonly pathPrefix: string;
-  /** Host already has this library — no pack, no sidecar, empty payload. */
-  readonly alreadyLoaded?: boolean;
-}
+import { pinContext, type PinnedContext } from "./context-file.ts";
+import { REJECT, serviceInterrupt, type ReplyBody, type SubLlmHandlers } from "./interrupts.ts";
+import { trace, traceEnabled } from "../util/trace.ts";
 
-/** 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[]>;
-  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[]>;
-  todo(action: string, params: Record<string, unknown>, depth: number): Promise<string>;
-  loadLibrary(source: string, depth: number): Promise<LibraryLoadResult>;
-}
+export type { LibraryLoadResult, SubcallOpts, SubLlmHandlers } from "./interrupts.ts";
 
 export interface SandboxOptions {
   /** Sandbox recursion depth label (passed to the worker, used in interrupt routing). */
@@ -69,15 +44,16 @@ export interface SandboxOptions {
   /** Sub-LLM prompt cap (chars) — sizes llm_query_chunked chunks inside the worker. */
   readonly maxPromptChars?: number;
   /**
-   * When true, the worker rejects open() write modes (pipeline read-only runs).
-   * Native repl() data work leaves this false so scratch-file writes still work.
+   * 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 readOnly?: boolean;
+  readonly awaitTimeoutS?: number;
 }
 
-const WORKER_PATH = join(dirname(fileURLToPath(import.meta.url)), "worker.py");
+const WORKER_PATH = join(dirname(fileURLToPath(import.meta.url)), "py", "worker.py");
 const STDERR_TAIL_CHARS = 8_192;
-const TODO_PROTO_KEYS = new Set(["type", "rid", "depth", "action"]);
+/** How long dispose() waits for a clean worker exit before escalating to SIGKILL. */
+const SHUTDOWN_GRACE_MS = 50;
 
 // The sandbox runs untrusted model-authored code; it must never inherit provider secrets.
 const SENSITIVE_ENV = /API[_-]?KEY|ACCESS[_-]?KEY|SECRET|TOKEN|PASSWORD|CREDENTIAL|ANTHROPIC|OPENAI|_KEY$/i;
@@ -90,22 +66,6 @@ function sanitizedEnv(): NodeJS.ProcessEnv {
   return env;
 }
 
-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")),
-  advancePhase: async () => formatError("phase advancement not available"),
-  saveArtifact: async () => formatError("save_artifact not available"),
-  askUserQuestion: async (questions) => questions.map((q) => ({
-    question: q.question,
-    selected: [],
-    custom: formatError("ask_user_question not configured"),
-  })),
-  todo: async () => formatError("todo not configured"),
-  loadLibrary: async () => { throw new Error("load_library not configured"); },
-};
-
 /** Distributive omit so each union member keeps its own fields (plain Omit collapses to shared keys). */
 type RequestBody = WorkerRequest extends infer T ? (T extends { id: string } ? Omit<T, "id"> : never) : never;
 
@@ -158,8 +118,8 @@ export class PythonSandbox {
     if (opts.maxPromptChars !== undefined) {
       workerArgs.push("--max-prompt-chars", String(opts.maxPromptChars));
     }
-    if (opts.readOnly) {
-      workerArgs.push("--read-only");
+    if (opts.awaitTimeoutS !== undefined) {
+      workerArgs.push("--await-timeout", String(opts.awaitTimeoutS));
     }
     this.proc = spawn(
       python,
@@ -171,11 +131,30 @@ export class PythonSandbox {
     this.proc.stdout.on("data", (chunk: string) => this.onData(chunk));
     this.proc.stderr.setEncoding("utf8");
     this.proc.stderr.on("data", (chunk: string) => this.appendStderr(chunk));
+
+    // A dead worker's pipe fails the write ASYNCHRONOUSLY: node emits 'error' on the stream, and
+    // an EventEmitter 'error' with no listener is an uncaughtException — which no try/catch
+    // around send() and no `await dispose().catch()` can intercept. Without these three
+    // listeners a worker dying mid-exec took the entire pi process down with it (EPIPE from the
+    // shutdown frame dispose() writes after the watchdog SIGKILLs).
+    // Record and swallow: the real reason is already surfaced by failAll on 'exit'.
+    const swallowPipeError = (stream: string) => (err: NodeJS.ErrnoException): void => {
+      this.appendStderr(`[rlm] worker ${stream} ${err.code ?? "error"}: ${err.message}\n`);
+    };
+    this.proc.stdin.on("error", swallowPipeError("stdin"));
+    this.proc.stdout.on("error", swallowPipeError("stdout"));
+    this.proc.stderr.on("error", swallowPipeError("stderr"));
+
     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}`));
     });
-    this.proc.on("exit", () => this.failAll(new Error(`worker exited; stderr=${this.stderr.trim()}`)));
+    // Name the cause: a SIGKILL (watchdog, abort, OOM killer) reads very differently from a
+    // Python-level crash, and this message is what reaches the user as `REPL error: …`.
+    this.proc.on("exit", (code, signal) => this.failAll(new Error(
+      `worker exited (${signal !== null ? `signal ${signal}` : `code ${code}`}); `
+      + `stderr=${this.stderr.trim()}`,
+    )));
 
     this.ready = this.waitForInit();
 
@@ -204,35 +183,31 @@ export class PythonSandbox {
     return sandbox;
   }
 
+  /**
+   * Load a payload whose pin this sandbox does not own — it acquires and releases one itself.
+   * Used by SandboxManager and tests; the engine owns its run's pin and calls the pinned form.
+   */
   async loadContext(payload: unknown): Promise<number> {
-    const isJson = typeof payload !== "string";
-    let path: string | undefined;
+    const pinned = await pinContext(payload);
     try {
-      path = await this.writeContextFile(payload, 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;
+      return await this.loadContextPinned(pinned);
     } finally {
-      if (path) await unlink(path).catch(() => {});
+      await pinned.release();
     }
   }
 
-  private async writeContextFile(payload: unknown, isJson: boolean): Promise<string> {
-    const file = join(
-      tmpdir(),
-      `rlm-ctx-${process.pid}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}.${isJson ? "json" : "txt"}`,
-    );
-    try {
-      await writeFile(file, isJson ? JSON.stringify(payload) : (payload as string));
-    } catch (e) {
-      await unlink(file).catch(() => {});
-      throw e;
-    }
-    return file;
+  /**
+   * Load from a pin the CALLER owns and will release. Keeping ownership outside means a run and
+   * every child that inherits its payload share one serialization and one file.
+   */
+  async loadContextPinned(pinned: PinnedContext): Promise<number> {
+    const res = await this.request({ type: "load_context", path: pinned.path, json: pinned.json });
+    if (!res.ok) throw new Error(res.error ?? "load_context failed");
+    return res.index ?? 0;
   }
 
-  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 ?? "",
@@ -248,38 +223,22 @@ export class PythonSandbox {
   async dispose(): Promise<void> {
     if (this.disposed) return;
     this.disposed = true;
-    try {
+    // Handshake only a live worker. The watchdog SIGKILLs before SandboxManager's catch calls
+    // dispose(), and writing the shutdown frame to a process that is already dead is exactly
+    // what produced the EPIPE that killed the host. Skipping it also drops a pointless 50ms
+    // wait from every dead-worker teardown.
+    if (this.workerAlive) {
       this.send({ id: "_shutdown", type: "shutdown" });
-    } catch {
-      /* pipe may already be gone */
+      // Wait for the worker to actually go rather than sleeping a fixed 50ms: this returns the
+      // instant it exits (the common case) and still gives up promptly if it never will, in
+      // which case the SIGKILL below finishes the job.
+      await once(this.proc, "exit", { signal: AbortSignal.timeout(SHUTDOWN_GRACE_MS) })
+        .catch(() => { /* did not exit in time — SIGKILL below */ });
     }
-    await new Promise((r) => setTimeout(r, 50));
     if (this.proc.exitCode === null) this.proc.kill("SIGKILL");
     this.failAll(new Error("sandbox disposed"));
   }
 
-  /** Pickle the worker's user namespace atomically to path (rename .tmp \u2192 final inside worker). */
-  async snapshot(path: string, nonce: string): Promise<boolean> {
-    try {
-      const res = await this.request({ type: "snapshot", path, nonce });
-      if (!res.ok) return false;
-      return res.ok;
-    } catch {
-      return false;
-    }
-  }
-
-  /** Restore user variables. Worker verifies session nonce before deserializing. */
-  async restore(path: string, nonce: string): Promise<boolean> {
-    try {
-      const res = await this.request({ type: "restore", path, nonce });
-      if (!res.ok) return false;
-      return res.ok;
-    } catch {
-      return false;
-    }
-  }
-
   // ---- internals ------------------------------------------------------------------------
 
   private waitForInit(): Promise<void> {
@@ -297,12 +256,34 @@ 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"));
+    // Reject before registering the pending entry: `send` no-ops for a dead worker, so a request
+    // queued here would otherwise sit until the watchdog fired instead of failing now.
+    if (!this.workerAlive) {
+      return Promise.reject(new Error(
+        `worker is not running (${this.exitDescription()}); request '${payload.type}' not sent`,
+      ));
+    }
     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,10 +304,64 @@ 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,
+      });
+    }
+    // Never write to a corpse. The write would fail asynchronously and, historically, take the
+    // host process with it; even with the stdin 'error' listener in place there is nothing to
+    // gain. MUST NOT throw — `reply()` calls this from serviceInterrupt's catch, where a throw
+    // would become an unhandled rejection, trading one crash for another.
+    if (!this.workerAlive) {
+      this.appendStderr(`[rlm] dropped '${msg.type}' frame: worker ${this.exitDescription()}\n`);
+      return;
+    }
     this.proc.stdin.write(`${JSON.stringify(msg)}\n`);
   }
 
+  /**
+   * Listener counts on the worker's stdio pipes. Exposed for tests.
+   *
+   * The `'error'` listener on each is load-bearing: an EventEmitter `'error'` with no listener is
+   * an uncaughtException, so a failed write to a dead worker's stdin kills the HOST process. That
+   * cannot be behaviour-tested reliably — whether the write reaches the syscall depends on
+   * whether node has reaped the child yet — so the invariant is asserted structurally.
+   */
+  get pipeErrorListenerCounts(): Readonly<Record<"stdin" | "stdout" | "stderr", number>> {
+    return Object.freeze({
+      stdin: this.proc.stdin.listenerCount("error"),
+      stdout: this.proc.stdout.listenerCount("error"),
+      stderr: this.proc.stderr.listenerCount("error"),
+    });
+  }
+
+  /** False once the worker is gone — exited, signalled, or its stdin torn down. */
+  private get workerAlive(): boolean {
+    return this.proc.exitCode === null
+      && this.proc.signalCode === null
+      && this.proc.stdin.writable;
+  }
+
+  /** How the worker went away, for diagnostics. */
+  private exitDescription(): string {
+    if (this.proc.signalCode !== null) return `killed by ${this.proc.signalCode}`;
+    if (this.proc.exitCode !== null) return `exited with code ${this.proc.exitCode}`;
+    return "stdin closed";
+  }
+
   private onData(chunk: string): void {
     this.buf += chunk;
     let nl: number;
@@ -353,9 +388,22 @@ 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);
+      void serviceInterrupt(msg, this.handlers, (rid, body) => this.reply(rid, body));
       return;
     }
     const p = this.pending.get(msg.id);
@@ -365,84 +413,7 @@ export class PythonSandbox {
     p.resolve(msg);
   }
 
-  private async serviceInterrupt(msg: WorkerInterrupt): Promise<void> {
-    const h = this.handlers;
-    const d = msg.depth;
-    try {
-      if (msg.type === "llm_query") {
-        const response = await h.llmQuery(msg.prompt ?? "", msg.model ?? null, d);
-        this.reply(msg.rid, { response });
-      } else if (msg.type === "rlm_query") {
-        const response = await h.rlmQuery(msg.prompt ?? "", msg.model ?? null, d);
-        this.reply(msg.rid, { response });
-      } else if (msg.type === "llm_query_batched") {
-        const responses = await h.llmQueryBatched(msg.prompts ?? [], msg.model ?? null, d);
-        this.reply(msg.rid, { responses });
-      } else if (msg.type === "rlm_query_batched") {
-        const responses = await h.rlmQueryBatched(msg.prompts ?? [], msg.model ?? null, d);
-        this.reply(msg.rid, { responses });
-      } else if (msg.type === "advance_phase") {
-        const response = await h.advancePhase(msg.phase ?? "", msg.summary, d);
-        this.reply(msg.rid, { response });
-      } else if (msg.type === "save_artifact") {
-        const response = await h.saveArtifact(msg.artifactKind ?? "", msg.content ?? "", d);
-        this.reply(msg.rid, { response });
-      } else if (msg.type === "ask_user_question") {
-        const answers = await h.askUserQuestion(msg.questions ?? [], d);
-        this.reply(msg.rid, { answers });
-      } else if (msg.type === "todo") {
-        const params = Object.fromEntries(
-          Object.entries(msg).filter(([key]) => !TODO_PROTO_KEYS.has(key)),
-        );
-        const response = await h.todo(msg.action ?? "list", params, d);
-        this.reply(msg.rid, { response });
-      } 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.
-          this.reply(msg.rid, {
-            already_loaded: true,
-            index: lib.index,
-            files: 0,
-            chars: lib.chars,
-            source_id: lib.sourceId,
-            path_prefix: lib.pathPrefix,
-          });
-        } else {
-          const isJson = typeof lib.payload !== "string";
-          const path = await this.writeContextFile(lib.payload, isJson);
-          // 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).
-          this.reply(msg.rid, {
-            path,
-            json: isJson,
-            index: lib.index,
-            files: lib.files,
-            chars: lib.chars,
-            source_id: lib.sourceId,
-            path_prefix: lib.pathPrefix,
-          });
-        }
-      }
-    } catch (err) {
-      this.reply(msg.rid, { error: errorMessage(err) });
-    }
-  }
-
-  private reply(rid: string, body: {
-    response?: string;
-    responses?: string[];
-    answers?: AskAnswer[];
-    path?: string;
-    json?: boolean;
-    index?: number;
-    files?: number;
-    chars?: number;
-    source_id?: string;
-    path_prefix?: string;
-    already_loaded?: boolean;
-    error?: string;
-  }): void {
+  private reply(rid: string, body: ReplyBody): void {
     if (!this.disposed) this.send({ type: "llm_reply", rid, ...body });
   }
 

package/src/text/tokens.ts

@@ -21,11 +21,37 @@ export function estimateMessageTokens(messages: { content: string }[]): number {
   return estimateTokens(chars);
 }
 
-/** Total character length of a context payload (string or list of strings). */
+/**
+ * Character length of one context entry. `ContextFile`-shaped entries report their content
+ * length; anything else falls back to its serialized form.
+ *
+ * Deliberately does NOT import `isContextFile` from context/library-context.ts: that module
+ * imports `estimateTokens` from here, so the reverse import would be a cycle. `in`-narrowing
+ * needs no type guard and no cast.
+ */
+function entryLength(entry: unknown): number {
+  if (typeof entry === "string") return entry.length;
+  if (entry !== null && typeof entry === "object" && "content" in entry) {
+    const content: unknown = entry.content;
+    if (typeof content === "string") return content.length;
+  }
+  return JSON.stringify(entry ?? "").length;
+}
+
+/**
+ * Total character length of a context payload (string, file bundle, or arbitrary value).
+ *
+ * The array branch must read each entry's `content`: `String(fileEntry)` yields
+ * "[object Object]" (15 chars), which under-reported a packed repository by ~200x. This number
+ * is what buildMetadataLine tells the model to size its batches against, and it is replayed
+ * into the system prompt, so it has to be real.
+ */
 export function contextLength(context: unknown): number {
   if (typeof context === "string") return context.length;
-  if (Array.isArray(context)) return context.reduce<number>((n, x) => n + String(x).length, 0);
-  return JSON.stringify(context ?? "").length;
+  if (!Array.isArray(context)) return JSON.stringify(context ?? "").length;
+  let total = 0; // running sum — no intermediate array, no per-entry closure
+  for (let i = 0; i < context.length; i++) total += entryLength(context[i]);
+  return total;
 }
 
 /** Human label for a context payload's type, used in the metadata prompt. */