@hicaru/pi-rlm 0.2.0 → 0.2.1

package/src/tool/subcall-store.ts

@@ -7,13 +7,19 @@
  * subcall accumulation logic.
  */
 import type { RlmEmitter, SubcallCreatedEvent, SubcallUpdatedEvent } from "./rlm-events.ts";
-import type { RlmSubcall } from "./rlm-details.ts";
+import type { RlmSubcall, SubcallStatus } from "./rlm-details.ts";
 import { EmitterListener } from "./emitter-listener.ts";
 
 type MutableSubcall = {
   -readonly [Key in keyof RlmSubcall]: RlmSubcall[Key];
 };
 
+/** Accumulated cost/tokens, shared by getTotals() and takeSettledSubtrees(). */
+export interface SubcallTotals {
+  readonly costUsd: number;
+  readonly tokens: number;
+}
+
 export class SubcallStore extends EmitterListener {
   private readonly subcalls = new Map<string, MutableSubcall>();
 
@@ -82,6 +88,56 @@ export class SubcallStore extends EmitterListener {
     return { costUsd: this.totalCostUsd, tokens: this.totalTokens };
   }
 
+  /**
+   * Remove and return every fully-settled root subtree, with its cost/tokens subtracted
+   * from the running totals so the caller can add them without double-counting.
+   *
+   * A root whose subtree still has a running node stays put. That matters because
+   * `renderCollapsedSubcallTree` walks down from `parentId === undefined`: a subcall handed
+   * over without its parent has no path from a root and is silently dropped from the tree.
+   * Handing over whole subtrees is what keeps adopted nodes renderable.
+   */
+  takeSettledSubtrees(): { readonly subcalls: readonly RlmSubcall[]; readonly totals: SubcallTotals } {
+    const children = new Map<string | undefined, MutableSubcall[]>();
+    for (const sc of this.subcalls.values()) {
+      const siblings = children.get(sc.parentId);
+      if (siblings === undefined) children.set(sc.parentId, [sc]);
+      else siblings.push(sc);
+    }
+
+    // Collect a root's subtree, or undefined when any node in it is still running.
+    const settledSubtree = (root: MutableSubcall): MutableSubcall[] | undefined => {
+      const collected: MutableSubcall[] = [];
+      const stack: MutableSubcall[] = [root];
+      while (stack.length > 0) {
+        const node = stack.pop();
+        if (node === undefined) continue;
+        if (node.status === "running") return undefined;
+        collected.push(node);
+        const kids = children.get(node.id);
+        if (kids !== undefined) stack.push(...kids);
+      }
+      return collected;
+    };
+
+    const taken: RlmSubcall[] = [];
+    let costUsd = 0;
+    let tokens = 0;
+    for (const root of children.get(undefined) ?? []) {
+      const subtree = settledSubtree(root);
+      if (subtree === undefined) continue;
+      for (const node of subtree) {
+        costUsd += node.costUsd;
+        tokens += node.tokens;
+        taken.push(Object.freeze({ ...node, status: node.status as SubcallStatus }));
+        this.subcalls.delete(node.id);
+      }
+    }
+    this.totalCostUsd -= costUsd;
+    this.totalTokens -= tokens;
+    return { subcalls: taken, totals: { costUsd, tokens } };
+  }
+
   // ── Root usage (delegated from RlmEventAggregator) ──
 
   /** Accumulate root-level usage into shared totals. Called by aggregator. */

package/src/util/concurrency.ts

@@ -1,15 +1,89 @@
-/** Fixed-size concurrency pool: run `fn` over `items` with at most `limit` in flight, preserving order. */
-export async function mapPool<T, R>(items: readonly T[], limit: number, fn: (item: T, idx: number) => Promise<R>): Promise<R[]> {
-  const out = new Array<R>(items.length);
-  let next = 0;
-  const worker = async (): Promise<void> => {
-    while (true) {
-      const index = next;
-      next += 1;
-      if (index >= items.length) return;
-      out[index] = await fn(items[index], index);
+/**
+ * Sub-call admission control.
+ *
+ * `spawn()` lets the sandbox put many requests on the wire at once, so a per-call pool no
+ * longer bounds anything: one `llm_query_chunked` over a large file posts every batch
+ * simultaneously, and each batch fans out again. The bound has to be session-wide, which is
+ * what `SubcallGates` is — constructed once at the composition root and shared by every
+ * handler.
+ */
+
+/**
+ * Counting semaphore: at most `limit` holders at once, FIFO.
+ *
+ * NOT re-entrant. A leaf completion takes exactly ONE slot, in `complete1`; wrapping a batch in
+ * a second acquisition of this same gate deadlocks as soon as the batch reaches `limit`
+ * prompts — the outer holders fill the gate and each waits for an inner slot nothing can free.
+ *
+ * FIFO means a large fan-out is not starved, but also that it is not preempted: a
+ * 500-prompt `llm_query_chunked` holds the queue until it drains, so an interactive
+ * `llm_query` issued behind it waits for the whole thing. Acceptable while sub-calls are
+ * uniform in priority; revisit if an interactive tier is ever added.
+ */
+export class Semaphore {
+  private active = 0;
+  private readonly waiters: Array<() => void> = [];
+
+  constructor(private readonly limit: number) {}
+
+  /** Hold a slot for the duration of `fn`. */
+  async run<T>(fn: () => Promise<T>): Promise<T> {
+    // `while`, not `if`: a caller arriving between the decrement below and a woken waiter's
+    // resumption would otherwise slip past the limit.
+    while (this.active >= this.limit) {
+      await new Promise<void>((resolve) => { this.waiters.push(resolve); });
     }
-  };
-  await Promise.all(Array.from({ length: Math.max(1, Math.min(limit, items.length)) }, worker));
-  return out;
+    this.active += 1;
+    try {
+      return await fn();
+    } finally {
+      this.active -= 1;
+      this.waiters.shift()?.();
+    }
+  }
+
+  /** In-flight holders. Exposed for tests asserting the bound. */
+  get inFlight(): number {
+    return this.active;
+  }
+}
+
+/**
+ * One semaphore per recursion depth.
+ *
+ * A single shared gate would deadlock: `limit` rlm_query children holding every slot, each
+ * blocked waiting for a grandchild that can never be admitted. Splitting by depth breaks the
+ * cycle — a holder at depth k only ever waits on depth k+1. Leaf LLM calls are terminal and
+ * never re-enter, so they safely share one process-wide gate.
+ */
+export class DepthGates {
+  private readonly gates = new Map<number, Semaphore>();
+
+  constructor(private readonly limit: number) {}
+
+  at(depth: number): Semaphore {
+    let gate = this.gates.get(depth);
+    if (gate === undefined) {
+      gate = new Semaphore(this.limit);
+      this.gates.set(depth, gate);
+    }
+    return gate;
+  }
+}
+
+/** Session-wide sub-call admission. Construct once; pass explicitly — never default one in. */
+export interface SubcallGates {
+  /** llm_query / llm_query_batched completions — terminal, so one shared gate. */
+  readonly leaf: Semaphore;
+  /** Recursive child engines — one gate per depth, see DepthGates. */
+  readonly rlm: DepthGates;
+}
+
+/**
+ * Worst case is `maxDepth × limit` concurrent child engines (each owning a Python
+ * subprocess) plus `limit` leaf completions, so keep `limit` modest — see
+ * DEFAULT_CONFIG.maxConcurrentSubcalls.
+ */
+export function createSubcallGates(limit: number): SubcallGates {
+  return Object.freeze({ leaf: new Semaphore(limit), rlm: new DepthGates(limit) });
 }

package/src/util/trace.ts

@@ -0,0 +1,42 @@
+/**
+ * RLM trace — one JSONL line per interesting event, for the E2E harness and for post-mortem of
+ * a hung run. Off unless RLM_TRACE_FILE is set; when off, every call costs one boolean check.
+ *
+ * Purely observational: it subscribes to the RlmEmitter that already exists and to the sandbox
+ * frames that already cross the pipe. No behaviour is duplicated here.
+ */
+import { appendFileSync } from "node:fs";
+import type { RlmEmitter } from "../tool/rlm-events.ts";
+
+const FILE = process.env.RLM_TRACE_FILE;
+
+/** Check this before building payloads on hot paths. */
+export const traceEnabled: boolean = typeof FILE === "string" && FILE.length > 0;
+
+const START = Date.now();
+
+/** Append one event. Fail-soft: tracing must never break a run (state/writes.ts convention). */
+export function trace(kind: string, data: Record<string, unknown> = {}): void {
+  if (!traceEnabled || FILE === undefined) return;
+  const now = Date.now();
+  try {
+    appendFileSync(FILE, `${JSON.stringify({ t: now, rel: now - START, pid: process.pid, kind, ...data })}\n`);
+  } catch {
+    /* ignore */
+  }
+}
+
+/** Mirror one emitter's sub-call lifecycle into the trace. Returns an unsubscribe fn. */
+export function attachTracer(emitter: RlmEmitter, scope: "turn" | "background"): () => void {
+  if (!traceEnabled) return () => {};
+  // Rename e.kind → subcallKind so it does not overwrite the outer event kind
+  // (`subcall.created` / `subcall.updated`) when `trace` spreads `data` after `kind`.
+  const offs = [
+    emitter.onSubcallCreated((e) => {
+      const { kind: subcallKind, ...rest } = e;
+      trace("subcall.created", { scope, subcallKind, ...rest });
+    }),
+    emitter.onSubcallUpdated((e) => trace("subcall.updated", { scope, ...e })),
+  ];
+  return () => { for (const off of offs) off(); };
+}

package/src/bridge/llm-query.ts

@@ -1,156 +0,0 @@
-/**
- * The `llm_query` / `llm_query_batched` bridge: turns sandbox sub-LLM interrupts into
- * real (serverless) completions on the configured *worker* model, reporting each call to the
- * RlmEmitter for progressive TUI re-rendering.
- *
- * Caps enforce the divide-and-conquer budget from the RLM method: per-prompt size and batch
- * fan-out are bounded, and batches run through a fixed-size concurrency pool.
- *
- * Every input that can change between calls (worker model, emitter, parent node, depth,
- * remaining budget) is an accessor, so a single bridge instance serves both the headless
- * engine — which binds them once per run — and the native `repl` tool, which swaps them per
- * invocation without recreating the sandbox.
- */
-
-import type { Api, Model, Usage } from "@earendil-works/pi-ai";
-import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
-import type { RlmEmitter } from "../tool/rlm-events.ts";
-import { displayModelRef, resolveModelId } from "../config/settings.ts";
-import { checkResourceLimits, type RemainingResources } from "../core/resource-limits.ts";
-import type { Sampling } from "../core/types.ts";
-import { type ChatMsg, modelComplete } from "./model.ts";
-import { previewText } from "../text/preview.ts";
-import { errorMessage, formatError, isErrorText } from "../util/errors.ts";
-import { mapPool } from "../util/concurrency.ts";
-
-/**
- * The config slice this bridge reads. Structurally satisfied by `RlmConfig`, and re-read on
- * every call so `/rlm-config` changes take effect without rebuilding the sandbox.
- */
-export interface LlmBridgeConfig {
-  readonly maxPromptChars: number;
-  readonly maxConcurrentSubcalls: number;
-  readonly subSystemPrompt?: string;
-  readonly subSampling?: Sampling;
-}
-
-export interface LlmBridgeOptions {
-  /** Resolved per call so provider/config changes between calls are picked up. */
-  readonly workerModel: () => Model<Api>;
-  readonly registry: ModelRegistry;
-  readonly config: () => LlmBridgeConfig;
-  readonly signal?: AbortSignal;
-  readonly onUsage?: (usage: Usage, model: Model<Api>) => void;
-  /** Parent run's remaining budget/timeout; checked before every sub-call. */
-  readonly remainingBudget?: () => RemainingResources | undefined;
-  /** Live RlmDetails reporting target, resolved per call. */
-  readonly emitter?: () => RlmEmitter | undefined;
-  readonly parentId?: () => string | undefined;
-  readonly depth?: () => number;
-}
-
-/** Provider failures that a smaller batch or a cheaper model might get past. */
-const RETRYABLE_HINT = /credit|402|payment|quota|rate.limit/i;
-
-export interface LlmBridge {
-  llmQuery(prompt: string, model: string | null, depth: number): Promise<string>;
-  llmQueryBatched(prompts: readonly string[], model: string | null, depth: number): Promise<string[]>;
-}
-
-/** Batch outcome summary, or undefined when every prompt succeeded. */
-function batchError(failed: number, total: number): string | undefined {
-  if (failed === 0) return undefined;
-  return failed === total
-    ? `all ${total} sub-calls failed — reduce batch size or try llm_query individually`
-    : `${failed}/${total} sub-calls failed`;
-}
-
-export function createLlmBridge(opts: LlmBridgeOptions): LlmBridge {
-  const displayModel = (model: string | null): string =>
-    displayModelRef(opts.registry, model, opts.workerModel());
-
-  // Run one completion; report cost/tokens via `track` (a per-call or per-batch accumulator).
-  async function complete1(prompt: string, model: string | null, track: (u: Usage) => void): Promise<string> {
-    const config = opts.config();
-    const rem = opts.remainingBudget?.();
-    if (rem !== undefined) {
-      const limitError = checkResourceLimits(rem);
-      if (limitError !== undefined) return limitError;
-    }
-    if (prompt.length > config.maxPromptChars) {
-      return formatError(`sub-LLM prompt exceeded the size limit (${prompt.length.toLocaleString()} chars > ${config.maxPromptChars.toLocaleString()}). Shorten or chunk the prompt before calling llm_query.`);
-    }
-    const resolved = model ? resolveModelId(opts.registry, model) : undefined;
-    if (model && !resolved) return formatError(`unknown model override '${model}'`);
-    const target = resolved ?? opts.workerModel();
-    try {
-      const messages: ChatMsg[] = [{ role: "user", content: prompt }];
-      const res = await modelComplete(messages, {
-        model: target,
-        registry: opts.registry,
-        system: config.subSystemPrompt,
-        maxTokens: config.subSampling?.maxTokens,
-        temperature: config.subSampling?.temperature,
-        reasoning: config.subSampling?.reasoning,
-        signal: opts.signal,
-      });
-      opts.onUsage?.(res.usage, target);
-      track(res.usage);
-      return res.text;
-    } catch (err) {
-      const msg = errorMessage(err);
-      const hint = RETRYABLE_HINT.test(msg) ? " — try smaller batches or individual llm_query calls" : "";
-      return formatError(`${msg}${hint}`);
-    }
-  }
-
-  return {
-    async llmQuery(prompt, model) {
-      const emitter = opts.emitter?.();
-      const id = emitter?.emitSubcallCreated({
-        kind: "llm", parentId: opts.parentId?.(), label: "llm_query",
-        model: displayModel(model), args: `prompt: ${previewText(prompt)}`,
-        depth: opts.depth?.() ?? 0,
-      });
-      let cost = 0;
-      let tokens = 0;
-      const out = await complete1(prompt, model, (u) => {
-        cost += u.cost.total;
-        tokens += u.totalTokens;
-      });
-      if (emitter && id !== undefined) emitter.emitSubcallUpdated({ id,
-        status: isErrorText(out) ? "error" : "done",
-        costUsd: cost, tokens, resultPreview: previewText(out),
-        detail: isErrorText(out) ? out : undefined,
-      });
-      return out;
-    },
-
-    async llmQueryBatched(prompts, model) {
-      const emitter = opts.emitter?.();
-      const id = emitter?.emitSubcallCreated({
-        kind: "batch", parentId: opts.parentId?.(), label: `llm_query ×${prompts.length}`,
-        model: displayModel(model), args: `prompt: ${previewText(prompts[0] ?? "")}`,
-        depth: opts.depth?.() ?? 0,
-      });
-      let cost = 0;
-      let tokens = 0;
-      const out = await mapPool(prompts, opts.config().maxConcurrentSubcalls, (p) =>
-        complete1(p, model, (u) => {
-          cost += u.cost.total;
-          tokens += u.totalTokens;
-        }),
-      );
-      const failed = out.filter(isErrorText).length;
-      const error = batchError(failed, out.length);
-      const firstPreview = previewText(out[0] ?? "");
-      const resultPreview = out.length > 1 ? `${firstPreview}  (+${out.length - 1} more)` : firstPreview;
-      if (emitter && id !== undefined) emitter.emitSubcallUpdated({ id,
-        status: error ? "error" : "done", costUsd: cost, tokens,
-        resultPreview, detail: error,
-        failedCount: failed, totalCount: out.length,
-      });
-      return out;
-    },
-  };
-}

package/src/bridge/rlm-query.ts

@@ -1,108 +0,0 @@
-/**
- * The `rlm_query` recursion bridge.
- *
- * A child RLM gets its own sandbox and iterates over the prompt as its context. At/over the
- * depth cap it degrades to a plain `llm_query` (ported from rlm/core/rlm.py `_subcall`). The
- * concurrency pool bounds parallel children for `rlm_query_batched`.
- *
- * As with the llm bridge, everything that can change between calls is an accessor, so the
- * headless engine (which binds them once per run) and the native `repl` tool (which swaps them
- * per invocation) share one implementation.
- */
-
-import type { RlmResult, RunRlm } from "../core/types.ts";
-import type { LlmBridge } from "./llm-query.ts";
-import type { RlmEmitter } from "../tool/rlm-events.ts";
-import { checkResourceLimits, type RemainingResources } from "../core/resource-limits.ts";
-import { errorMessage, formatError } from "../util/errors.ts";
-import { mapPool } from "../util/concurrency.ts";
-import { previewText } from "../text/preview.ts";
-
-/** The config slice this bridge reads; structurally satisfied by `RlmConfig`. */
-export interface RlmBridgeConfig {
-  readonly maxDepth: number;
-  readonly maxConcurrentSubcalls: number;
-}
-
-export interface RlmHandlers {
-  rlmQuery(prompt: string, model: string | null, depth: number): Promise<string>;
-  rlmQueryBatched(prompts: readonly string[], model: string | null, depth: number): Promise<string[]>;
-}
-
-export interface RlmBridgeOptions {
-  /**
-   * Spawns one child run. The engine passes its own `run` (self-recursion); the native tool
-   * passes a closure that builds a child engine bound to the current invocation's emitter.
-   */
-  readonly run: RunRlm;
-  readonly llm: LlmBridge;
-  readonly config: () => RlmBridgeConfig;
-  /** "provider/id" shown on the sub-call node; see `displayModelRef` in config/settings.ts. */
-  readonly modelLabel?: (override: string | null) => string;
-  /** Live RlmDetails reporting target, resolved per call. */
-  readonly emitter: () => RlmEmitter | undefined;
-  /** Parent subcall ID that this run is attached under, resolved per call. */
-  readonly parentNodeId?: () => string | undefined;
-  /** Returns the parent's remaining budget/timeout for seeding child runs. */
-  readonly remainingBudget?: () => RemainingResources | undefined;
-  /** Called with a child run's total cost/tokens so the parent LimitGuard debits it. */
-  readonly onChildUsage?: (costUsd: number, inputTokens: number, outputTokens: number) => void;
-}
-
-export function createRlmHandlers(opts: RlmBridgeOptions): RlmHandlers {
-  /**
-   * One child spawn. `childDepth` is the absolute depth the child will run at.
-   * Never throws: failures come back as "Error: ..." strings, matching the sandbox contract.
-   */
-  async function child(prompt: string, model: string | null, childDepth: number): Promise<string> {
-    // At the cap, a child RLM would just be an LM — short-circuit to a one-shot llm_query.
-    if (childDepth >= opts.config().maxDepth) {
-      return opts.llm.llmQuery(prompt, model, childDepth - 1);
-    }
-
-    const rem = opts.remainingBudget?.();
-    if (rem !== undefined) {
-      // Pre-spawn guard: refuse if the parent's budget or timeout is already exhausted
-      // (reference: _subcall checks remaining_budget/timeout before spawning).
-      const limitError = checkResourceLimits(rem);
-      if (limitError !== undefined) return limitError;
-    }
-
-    const emitter = opts.emitter();
-    const subId = emitter?.emitSubcallCreated({
-      kind: "rlm", parentId: opts.parentNodeId?.(), label: "rlm_query",
-      model: opts.modelLabel?.(model) ?? model ?? undefined,
-      detail: prompt.slice(0, 60),
-      depth: childDepth,
-    });
-
-    try {
-      const res: RlmResult = await opts.run({
-        rootPrompt: "",
-        context: prompt,
-        depth: childDepth,
-        parentNodeId: subId,
-        modelOverride: model ?? undefined,
-        remainingBudgetUsd: rem?.budgetUsd,
-        remainingTimeoutMs: rem?.timeoutMs,
-      });
-      opts.onChildUsage?.(res.costUsd, res.inputTokens, res.outputTokens);
-      // The child emits live cost/token deltas on the shared emitter as it runs, so the node
-      // must NOT also receive a final aggregate — that would double-count.
-      if (emitter && subId !== undefined) {
-        emitter.emitSubcallUpdated({ id: subId, status: "done", resultPreview: previewText(res.answer) });
-      }
-      return res.answer;
-    } catch (err) {
-      const msg = errorMessage(err);
-      if (emitter && subId !== undefined) emitter.emitSubcallUpdated({ id: subId, status: "error", detail: msg });
-      return formatError(`child RLM failed - ${msg}`);
-    }
-  }
-
-  return {
-    rlmQuery: (prompt, model, depth) => child(prompt, model, depth + 1),
-    rlmQueryBatched: (prompts, model, depth) =>
-      mapPool(prompts, opts.config().maxConcurrentSubcalls, (p) => child(p, model, depth + 1)),
-  };
-}