@hicaru/pi-rlm 0.3.0 → 0.3.2

package/src/bridge/handlers/rlm-query.ts

@@ -0,0 +1,227 @@
+/**
+ * rlm_query and rlm_batch handlers — async-by-default spawn pattern.
+ *
+ * AGENTS.md DRY #2: childRun exists once, here.
+ */
+
+import { modelRef } from "../../config/settings.ts";
+import { errorMessage, formatError } from "../../util/errors.ts";
+import { filterContextByPaths } from "../../context/merge.ts";
+import { previewText } from "../../text/preview.ts";
+import type { RlmInput, RlmResult } from "../../core/types.ts";
+import { checkResourceLimits } from "../../core/resource-limits.ts";
+import type { Invocation, SpawnResult, SubcallHandlerDeps } from "./types.ts";
+import type { SubcallOpts } from "../../sandbox/interrupts.ts";
+import { SPAWN_HINT, spawnAndRun, type SpawnDeps } from "./task-registry.ts";
+import { complete1, type Complete1Deps } from "./completion.ts";
+
+const UNWIRED = formatError("RLM bridge not wired for this invocation");
+const NO_UNMATCHED: readonly string[] = Object.freeze([]);
+
+function emptyResult(answer: string): RlmResult {
+  return {
+    answer,
+    iterations: 0,
+    costUsd: 0,
+    inputTokens: 0,
+    outputTokens: 0,
+    durationMs: 0,
+  };
+}
+
+interface ChildContext {
+  readonly context: unknown;
+  readonly unmatched: readonly string[];
+}
+
+function childContextFor(
+  deps: SubcallHandlerDeps,
+  prompt: string,
+  paths: readonly string[] | undefined,
+): ChildContext {
+  const inherited = deps.getChildContext?.();
+  if (inherited === undefined || inherited === null) {
+    return Object.freeze({ context: prompt, unmatched: NO_UNMATCHED });
+  }
+  if (paths === undefined || paths.length === 0) {
+    return Object.freeze({ context: inherited, unmatched: NO_UNMATCHED });
+  }
+  const filtered = filterContextByPaths(inherited, paths);
+  return Object.freeze({
+    context: filtered.files.length > 0 ? filtered.files : inherited,
+    unmatched: filtered.unmatched,
+  });
+}
+
+function completeDeps(deps: SubcallHandlerDeps): Complete1Deps {
+  return {
+    leafGate: deps.gates.leaf,
+    registry: deps.registry,
+    getLlmModel: deps.getLlmModel,
+    getConfig: deps.getConfig,
+    signal: deps.signal,
+    onUsage: deps.onUsage,
+  };
+}
+
+/**
+ * One child RLM run: depth cap → resource guard → depth gate → spawn engine → debit parent.
+ * Emits its own subcall node (do not wrap in emitting()).
+ */
+async function childRun(
+  deps: SubcallHandlerDeps,
+  inv: Invocation,
+  prompt: string,
+  paths: readonly string[] | undefined,
+): Promise<RlmResult> {
+  const childDepth = inv.depth + 1;
+  const run = deps.runChild;
+  const maxDepth = deps.getConfig().maxDepth;
+
+  if (run === undefined || childDepth >= maxDepth) {
+    const degrade = deps.degrade;
+    const answer =
+      degrade !== undefined
+        ? await degrade(prompt, inv.depth)
+        : await complete1(inv, prompt, () => {}, completeDeps(deps));
+    return emptyResult(answer);
+  }
+
+  const remTimeout = inv.limits.remainingTimeoutMs();
+  const limitError = checkResourceLimits({ timeoutMs: remTimeout });
+  if (limitError !== undefined) return emptyResult(limitError);
+
+  const rootModel = deps.getModel?.();
+  const modelLabel =
+    rootModel === undefined ? undefined : (modelRef(rootModel) ?? rootModel.id);
+  const subId = inv.emitter.emitSubcallCreated({
+    kind: "rlm",
+    parentId: inv.parentId,
+    label: "rlm_query",
+    model: modelLabel,
+    detail: prompt.slice(0, 60),
+    depth: childDepth,
+  });
+
+  const child = childContextFor(deps, prompt, paths);
+  const rootPrompt =
+    child.unmatched.length === 0
+      ? prompt
+      : `${prompt}\n\n[rlm] paths=${child.unmatched.join(", ")} matched no files; you received the full context.`;
+
+  const input: RlmInput = {
+    rootPrompt,
+    context: child.context,
+    depth: childDepth,
+    parentNodeId: subId,
+    remainingTimeoutMs: remTimeout,
+  };
+
+  try {
+    const res = await deps.gates.rlm.at(childDepth).run(() => run(input, inv));
+    inv.limits.addRaw(res.costUsd, res.inputTokens, res.outputTokens);
+    deps.onChildUsage?.(res.costUsd, res.inputTokens, res.outputTokens);
+    inv.emitter.emitSubcallUpdated({
+      id: subId,
+      status: "done",
+      resultPreview: res.answer.slice(0, 200),
+    });
+    return res;
+  } catch (err: unknown) {
+    const msg = errorMessage(err);
+    inv.emitter.emitSubcallUpdated({ id: subId, status: "error", detail: msg });
+    return emptyResult(formatError(`child RLM failed - ${msg}`));
+  }
+}
+
+export function createRlmQueryHandler(deps: SubcallHandlerDeps, sd: SpawnDeps) {
+  return async (
+    task: string,
+    depth: number,
+    opts: SubcallOpts,
+  ): Promise<SpawnResult> => {
+    const inv = deps.resolve(opts, depth);
+    if (inv === null) {
+      return {
+        ok: false,
+        task_id: null,
+        kind: "rlm",
+        n: 1,
+        status: "pending",
+        hint: SPAWN_HINT,
+        error: UNWIRED,
+      };
+    }
+
+    const pathArg = opts.paths;
+
+    return spawnAndRun(
+      sd,
+      "rlm",
+      1,
+      async () => {
+        const r = await childRun(deps, inv, task, pathArg);
+        return r.answer;
+      },
+      deps.trackDetached,
+      opts.detached,
+    );
+  };
+}
+
+export function createRlmBatchHandler(deps: SubcallHandlerDeps, sd: SpawnDeps) {
+  return async (
+    tasks: readonly string[],
+    depth: number,
+    opts: SubcallOpts,
+  ): Promise<SpawnResult> => {
+    const inv = deps.resolve(opts, depth);
+    if (inv === null) {
+      return {
+        ok: false,
+        task_id: null,
+        kind: "rlm_batch",
+        n: tasks.length,
+        status: "pending",
+        hint: SPAWN_HINT,
+        error: UNWIRED,
+      };
+    }
+
+    const pathArg = opts.paths;
+    const id = inv.emitter.emitSubcallCreated({
+      kind: "batch",
+      parentId: inv.parentId,
+      label: `rlm_batch ×${tasks.length}`,
+      args: previewText(tasks[0] ?? ""),
+      depth: inv.depth,
+    });
+
+    return spawnAndRun(
+      sd,
+      "rlm_batch",
+      tasks.length,
+      async () => {
+        try {
+          const results = await Promise.all(
+            tasks.map((t) => childRun(deps, inv, t, pathArg)),
+          );
+          const answers = results.map((r) => r.answer);
+          inv.emitter.emitSubcallUpdated({
+            id,
+            status: "done",
+            resultPreview: previewText(answers[0] ?? ""),
+            totalCount: answers.length,
+          });
+          return answers;
+        } catch (err: unknown) {
+          const msg = errorMessage(err);
+          inv.emitter.emitSubcallUpdated({ id, status: "error", detail: msg });
+          throw err;
+        }
+      },
+      deps.trackDetached,
+      opts.detached,
+    );
+  };
+}

package/src/bridge/handlers/task-registry.ts

@@ -0,0 +1,202 @@
+/**
+ * Single in-memory task registry for api_v5 async-by-default spawns.
+ *
+ * DRY: this is the ONLY place that assigns task_ids, parks await waiters, and
+ * resolves/rejects entries. Handlers and session UI both consume this type —
+ * never re-implement a Map of TaskEntry elsewhere.
+ */
+
+import type { AwaitResult, SpawnResult, TaskEntry } from "./types.ts";
+
+export const SPAWN_HINT =
+  "Call await_task(task_id=...) to get the result — this is NOT the answer.";
+
+export interface SpawnDeps {
+  nextId(): number;
+  register(kind: SpawnResult["kind"], n: number, taskId: string): SpawnResult;
+  resolve(taskId: string, result: string | readonly string[]): void;
+  reject(taskId: string, error: string): void;
+}
+
+export interface AwaitDeps {
+  get(taskId: string): TaskEntry | undefined;
+  wait(taskId: string, timeoutMs?: number): Promise<TaskEntry>;
+  unawaitedIds(): readonly string[];
+}
+
+export interface TaskRegistry {
+  readonly spawnDeps: SpawnDeps;
+  readonly awaitDeps: AwaitDeps;
+  entries(): ReadonlyMap<string, TaskEntry>;
+  toAwaitResult(taskId: string): AwaitResult;
+}
+
+interface Waiter {
+  resolve: (entry: TaskEntry) => void;
+  reject: (err: Error) => void;
+  timer?: ReturnType<typeof setTimeout>;
+}
+
+function notify(waiters: Map<string, Waiter>, taskId: string, entry: TaskEntry): void {
+  const w = waiters.get(taskId);
+  if (w === undefined) return;
+  if (w.timer !== undefined) clearTimeout(w.timer);
+  w.resolve(entry);
+  waiters.delete(taskId);
+}
+
+export function createTaskRegistry(): TaskRegistry {
+  const tasks = new Map<string, TaskEntry>();
+  const waiters = new Map<string, Waiter>();
+  let counter = 0;
+
+  const spawnDeps: SpawnDeps = {
+    nextId: () => {
+      counter += 1;
+      return counter;
+    },
+    register(kind, n, taskId) {
+      const entry: TaskEntry = {
+        taskId,
+        kind,
+        n,
+        status: "pending",
+        createdAt: Date.now(),
+      };
+      tasks.set(taskId, entry);
+      return Object.freeze({
+        ok: true,
+        task_id: taskId,
+        kind,
+        n,
+        status: "pending" as const,
+        hint: SPAWN_HINT,
+      });
+    },
+    resolve(taskId, result) {
+      const entry = tasks.get(taskId);
+      if (entry === undefined) return;
+      entry.status = "done";
+      if (typeof result === "string") {
+        entry.result = result;
+      } else {
+        entry.results = Object.freeze([...result]);
+      }
+      notify(waiters, taskId, entry);
+    },
+    reject(taskId, error) {
+      const entry = tasks.get(taskId);
+      if (entry === undefined) return;
+      entry.status = "error";
+      entry.error = error;
+      const w = waiters.get(taskId);
+      if (w !== undefined) {
+        if (w.timer !== undefined) clearTimeout(w.timer);
+        w.reject(new Error(error));
+        waiters.delete(taskId);
+      }
+    },
+  };
+
+  const awaitDeps: AwaitDeps = {
+    get: (taskId) => tasks.get(taskId),
+    wait(taskId, timeoutMs) {
+      return new Promise<TaskEntry>((resolve, reject) => {
+        const entry = tasks.get(taskId);
+        if (entry !== undefined && entry.status !== "pending") {
+          resolve(entry);
+          return;
+        }
+        const timer =
+          timeoutMs !== undefined
+            ? setTimeout(() => {
+                waiters.delete(taskId);
+                const e = tasks.get(taskId);
+                if (e !== undefined && e.status === "pending") {
+                  e.status = "timeout";
+                  e.error = `Timeout after ${timeoutMs}ms`;
+                }
+                reject(new Error(`Timeout waiting for task ${taskId}`));
+              }, timeoutMs)
+            : undefined;
+        waiters.set(taskId, { resolve, reject, timer });
+      });
+    },
+    unawaitedIds: () => {
+      const ids: string[] = [];
+      for (const [id, e] of tasks) {
+        if (e.status === "pending") ids.push(id);
+      }
+      return ids;
+    },
+  };
+
+  return {
+    spawnDeps,
+    awaitDeps,
+    entries: () => tasks,
+    toAwaitResult(taskId) {
+      const entry = tasks.get(taskId);
+      if (entry === undefined) {
+        return {
+          ok: false,
+          task_id: taskId,
+          kind: "unknown",
+          status: "error",
+          error: `Task ${taskId} not found`,
+        };
+      }
+      const status =
+        entry.status === "pending" ? "error" : entry.status;
+      return {
+        ok: entry.status === "done",
+        task_id: entry.taskId,
+        kind: entry.kind,
+        status,
+        result: entry.result,
+        results: entry.results,
+        error:
+          entry.error ??
+          (entry.status === "pending" ? "Task still pending" : undefined),
+      };
+    },
+  };
+}
+
+/**
+ * Register a spawn and start background work once.
+ * Returns SpawnResult immediately; never uses non-null assertions on task_id.
+ */
+export function spawnAndRun(
+  sd: SpawnDeps,
+  kind: SpawnResult["kind"],
+  n: number,
+  work: () => Promise<string | readonly string[]>,
+  trackDetached: (<T>(run: () => Promise<T>) => Promise<T>) | undefined,
+  detached: boolean,
+): SpawnResult {
+  const id = sd.nextId();
+  const taskId = `task_${id}`;
+  const spawned = sd.register(kind, n, taskId);
+
+  const run = async (): Promise<void> => {
+    try {
+      const result = await work();
+      sd.resolve(taskId, result);
+    } catch (err: unknown) {
+      const message = err instanceof Error ? err.message : String(err);
+      sd.reject(taskId, message.startsWith("Error:") ? message : `Error: ${message}`);
+    }
+  };
+
+  if (trackDetached !== undefined && detached) {
+    void trackDetached(run).catch((err: unknown) => {
+      const message = err instanceof Error ? err.message : String(err);
+      sd.reject(taskId, message.startsWith("Error:") ? message : `Error: ${message}`);
+    });
+  } else {
+    void run();
+  }
+
+  return spawned;
+}

package/src/bridge/handlers/types.ts

@@ -0,0 +1,136 @@
+/**
+ * Shared types for the async-by-default subcall handler set.
+ *
+ * Every subcall (llm_query, llm_batch, rlm_query, rlm_batch) returns a SpawnResult
+ * immediately with a task_id. The model must call await(task_id) to collect the
+ * real answer. This contract is proven in rlm_test (api_v5 + batch, scores 0.89–1.0).
+ */
+
+import type { Api, Model, Usage } from "@earendil-works/pi-ai";
+import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
+import type { RlmInput, RlmResult, Sampling } from "../../core/types.ts";
+import type { SubcallGates } from "../../util/concurrency.ts";
+import type { SubcallOpts } from "../../sandbox/interrupts.ts";
+import type { RlmEmitter } from "../../tool/rlm-events.ts";
+
+// ---------------------------------------------------------------------------
+// Spawn / Await / Finish — the three shapes the model sees
+// ---------------------------------------------------------------------------
+
+/** Returned by every spawn: llm_query, llm_batch, rlm_query, rlm_batch. */
+export interface SpawnResult {
+  readonly ok: boolean;
+  readonly task_id: string | null;
+  readonly kind: "llm" | "rlm" | "llm_batch" | "rlm_batch";
+  /** Number of sub-tasks in this batch (1 for singles). */
+  readonly n: number;
+  readonly status: "pending";
+  /** Hint text reminding the model to await. */
+  readonly hint: string;
+  readonly error?: string;
+}
+
+/** Returned by await(task_id) or await(task_ids=[...]). */
+export interface AwaitResult {
+  readonly ok: boolean;
+  readonly task_id: string;
+  readonly kind: string;
+  readonly status: "done" | "error" | "timeout";
+  /** Single result (llm_query, rlm_query). */
+  readonly result?: string;
+  /** Ordered batch results (llm_batch, rlm_batch). */
+  readonly results?: readonly string[];
+  readonly error?: string;
+}
+
+/** Returned by finish(summary) — the contract boundary. */
+export interface FinishResult {
+  readonly ok: true;
+  readonly finished: true;
+}
+
+// ---------------------------------------------------------------------------
+// Invocation — per-call context (emitter, limits, depth)
+// ---------------------------------------------------------------------------
+
+export interface InvocationLimits {
+  remainingTimeoutMs(): number | undefined;
+  addUsage(usage: Usage): void;
+  addRaw(costUsd: number, inputTokens: number, outputTokens: number): void;
+}
+
+export function limitsFromRemaining(
+  remaining?: () => { readonly timeoutMs?: number },
+): InvocationLimits {
+  return Object.freeze({
+    remainingTimeoutMs: () => remaining?.().timeoutMs,
+    addUsage: (_usage: Usage) => {},
+    addRaw: (_costUsd: number, _inputTokens: number, _outputTokens: number) => {},
+  });
+}
+
+export interface Invocation {
+  readonly emitter: RlmEmitter;
+  readonly parentId: string | undefined;
+  readonly depth: number;
+  readonly limits: InvocationLimits;
+}
+
+// ---------------------------------------------------------------------------
+// Handler dependencies — supplied by the engine or repl() tool
+// ---------------------------------------------------------------------------
+
+export interface SubcallConfig {
+  readonly maxPromptChars: number;
+  readonly maxDepth: number;
+  readonly subSampling?: Sampling;
+  readonly subSystemPrompt?: string;
+}
+
+export interface SubcallHandlerDeps {
+  /** Resolve the invocation for a given interrupt (emitter + limits). */
+  readonly resolve: (opts: SubcallOpts, depth: number) => Invocation | null;
+  /** Session-wide concurrency gates. */
+  readonly gates: SubcallGates;
+  readonly registry: ModelRegistry;
+  readonly getLlmModel: () => Model<Api>;
+  readonly getConfig: () => SubcallConfig;
+  readonly signal?: AbortSignal;
+  readonly onUsage?: (usage: Usage, role: "sub") => void;
+  // Recursion
+  readonly runChild?: (input: RlmInput, inv: Invocation) => Promise<RlmResult>;
+  readonly getChildContext?: () => unknown;
+  readonly getModel?: () => Model<Api>;
+  readonly degrade?: (prompt: string, depth: number) => Promise<string>;
+  readonly onChildUsage?: (costUsd: number, inputTokens: number, outputTokens: number) => void;
+  readonly trackDetached?: <T>(run: () => Promise<T>) => Promise<T>;
+}
+
+// ---------------------------------------------------------------------------
+// Handler set returned to the sandbox (canonical names only)
+// ---------------------------------------------------------------------------
+
+/** Canonical api_v5 handler set — no legacy aliases. */
+export interface SubcallHandlers {
+  readonly llmQuery: (prompt: string, depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly llmBatch: (prompts: readonly string[], depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly rlmQuery: (task: string, depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly rlmBatch: (tasks: readonly string[], depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly awaitTask: (taskId: string | undefined, taskIds: readonly string[] | undefined, timeoutS: number | undefined, depth: number, opts: SubcallOpts) => Promise<AwaitResult>;
+  readonly finishTask: (summary: string, depth: number, opts: SubcallOpts) => Promise<FinishResult>;
+}
+
+// ---------------------------------------------------------------------------
+// Task registry entry — tracks in-flight spawns
+// ---------------------------------------------------------------------------
+
+export interface TaskEntry {
+  readonly taskId: string;
+  readonly kind: "llm" | "rlm" | "llm_batch" | "rlm_batch";
+  readonly n: number;
+  status: "pending" | "done" | "error" | "timeout";
+  result?: string;
+  results?: readonly string[];
+  error?: string;
+  readonly createdAt: number;
+}

package/src/commands/rlm-config.ts

@@ -8,7 +8,7 @@ import type { RlmController } from "../mode/rlm-mode.ts";
 import { cheapestModel } from "../mode/llm-model.ts";
 import { setRlmModeStatus } from "../ui/status.ts";
 import { showConfigPanel } from "../ui/config-panel.ts";
-import { pickableModels, selectModel } from "../ui/model-picker.ts";
+import { pickableModels, selectModel, type ModelSelection } from "../ui/model-picker.ts";
 
 /** Newer Pi hosts expose session-scoped models; 0.79 peers do not — duck-type safely. */
 function sessionScopedModels(
@@ -18,6 +18,34 @@ function sessionScopedModels(
   return Array.isArray(scoped) ? scoped as readonly { readonly model: Model<Api> }[] : undefined;
 }
 
+/**
+ * Apply a model-picker result to controller pin state.
+ *
+ * - `null` → explicit "cheapest (auto)" (clear pin; leave reasoning alone)
+ * - `ModelSelection` → pin that model (and its thinking level, which may be undefined)
+ * - `undefined` → ESC / no change
+ */
+export function applyLlmSelection(
+  controller: RlmController,
+  llm: ModelSelection | null | undefined,
+): void {
+  if (llm === undefined) return;
+  if (llm === null) {
+    controller.llmModel = undefined;
+    controller.savedLlmRef = undefined;
+    return;
+  }
+  controller.llmModel = llm.model;
+  controller.savedLlmRef = modelRef(llm.model);
+  controller.setConfig(Object.freeze({
+    ...controller.config,
+    subSampling: Object.freeze({
+      ...controller.config.subSampling,
+      reasoning: llm.thinkingLevel,
+    }),
+  }));
+}
+
 export async function runRlmConfig(controller: RlmController, ctx: ExtensionContext): Promise<boolean> {
   // Match Pi's native list: refresh so a just-added key appears, then use scoped models when
   // the session narrowed them, else every available (auth-configured) model. Never getAll().
@@ -34,23 +62,14 @@ export async function runRlmConfig(controller: RlmController, ctx: ExtensionCont
     models,
     controller.llmModel,
     controller.config.subSampling.reasoning,
+    controller.savedLlmRef,
   );
-  if (llm !== undefined) {
-    controller.llmModel = llm?.model;
-    controller.setConfig(Object.freeze({
-      ...controller.config,
-      subSampling: Object.freeze({ ...controller.config.subSampling, reasoning: llm?.thinkingLevel }),
-    }));
-  }
+  // Only an explicit choice touches the pin. ESC leaves model + reasoning alone.
+  // Choosing cheapest must NOT wipe subSampling.reasoning (null !== undefined used to).
+  applyLlmSelection(controller, llm);
 
   controller.setConfig(await showConfigPanel(ctx, controller.config));
 
-  // Only an explicit choice touches the persisted pin. ESC (`undefined`) used to fall through
-  // here and freeze whatever cheapest resolved to at that moment, which silently ended
-  // "cheapest (auto)" for every later session — including once a cheaper model appeared.
-  if (llm === null) controller.savedLlmRef = undefined;          // "⟳ cheapest (auto)"
-  else if (llm !== undefined) controller.savedLlmRef = modelRef(llm.model);
-
   const persisted = await controller.persist();
   if (!persisted) ctx.ui.notify("RLM: failed to save settings to ~/.pi/agent/rlm.json", "error");
   setRlmModeStatus(ctx.ui, controller, ctx.getContextUsage());

package/src/context/listing.ts

@@ -23,7 +23,7 @@ export function formatContextListing(context: unknown): string {
       'Use `add_context("/path/to/dir")` / `add_context("docs.pdf")` / `add_context("https://…")` for external sources.',
       "Documents (PDF, DOCX, XLSX, PPTX, CSV, …) are converted to Markdown on the way in.",
       "",
-      "Use repl({code}) and delegate semantic reading to llm_query / llm_query_batched / llm_query_chunked.",
+      "Use repl({code}) and delegate semantic reading to llm_query / llm_batch / llm_query_chunked.",
     ].join("\n");
   }
 
@@ -53,7 +53,7 @@ export function formatContextListing(context: unknown): string {
     truncated,
     "",
     "File contents are loaded in the REPL `context` variable — file-reading tools are disabled.",
-    "Use repl({code}) and delegate semantic reading to llm_query / llm_query_batched / llm_query_chunked.",
+    "Use repl({code}) and delegate semantic reading to llm_query / llm_batch / llm_query_chunked.",
   ].join("\n");
 }