@hicaru/pi-rlm 0.1.0

package/src/bridge/model.ts

@@ -0,0 +1,97 @@
+/**
+ * modelComplete — a single, serverless, in-process LLM completion.
+ *
+ * This is the one place that talks to a provider. It resolves the API key from pi's
+ * ModelRegistry (keys live here, never in the sandbox) and calls pi-ai's `completeSimple`.
+ * Used both for `llm_query` (one user prompt) and for the headless RLM root (full history).
+ */
+
+import { type Api, completeSimple, type Message, type Model, type ThinkingLevel, type Usage } from "@earendil-works/pi-ai";
+import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
+
+export type Role = "system" | "user" | "assistant";
+export interface ChatMsg {
+  readonly role: Role;
+  readonly content: string;
+}
+
+export interface CompleteOptions {
+  readonly model: Model<Api>;
+  readonly registry: ModelRegistry;
+  readonly system?: string;
+  readonly maxTokens?: number;
+  readonly temperature?: number;
+  readonly reasoning?: ThinkingLevel;
+  readonly signal?: AbortSignal;
+}
+
+export interface CompleteResult {
+  readonly text: string;
+  readonly usage: Usage;
+}
+
+/** Build a synthetic AssistantMessage (pi-ai requires the full shape for history replay). */
+function assistantMessage(text: string, model: Model<Api>): Message {
+  return {
+    role: "assistant",
+    content: [{ type: "text", text }],
+    api: model.api,
+    provider: model.provider,
+    model: model.id,
+    usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+    stopReason: "stop",
+    timestamp: Date.now(),
+  };
+}
+
+function toPiMessages(messages: readonly ChatMsg[], model: Model<Api>): { readonly systemPrompt?: string; readonly messages: Message[] } {
+  let systemPrompt: string | undefined;
+  const out: Message[] = [];
+  for (const m of messages) {
+    if (m.role === "system") {
+      systemPrompt = systemPrompt ? `${systemPrompt}\n\n${m.content}` : m.content;
+    } else if (m.role === "user") {
+      out.push({ role: "user", content: m.content, timestamp: Date.now() });
+    } else {
+      out.push(assistantMessage(m.content, model));
+    }
+  }
+  return { systemPrompt, messages: out };
+}
+
+/** Extract the assistant's plain text from a completion. */
+function extractText(content: readonly { readonly type: string; readonly text?: string }[]): string {
+  return content
+    .filter((c) => c.type === "text" && typeof c.text === "string")
+    .map((c) => c.text)
+    .join("");
+}
+
+export async function modelComplete(messages: readonly ChatMsg[], opts: CompleteOptions): Promise<CompleteResult> {
+  const auth = await opts.registry.getApiKeyAndHeaders(opts.model);
+  if (!auth.ok) throw new Error(`auth for ${opts.model.provider}/${opts.model.id}: ${auth.error}`);
+
+  const built = toPiMessages(messages, opts.model);
+  const systemPrompt = opts.system
+    ? built.systemPrompt
+      ? `${opts.system}\n\n${built.systemPrompt}`
+      : opts.system
+    : built.systemPrompt;
+
+  const msg = await completeSimple(
+    opts.model,
+    { systemPrompt, messages: built.messages },
+    {
+      apiKey: auth.apiKey,
+      headers: auth.headers,
+      maxTokens: opts.maxTokens,
+      temperature: opts.temperature,
+      reasoning: opts.reasoning,
+      signal: opts.signal,
+    },
+  );
+  if (msg.stopReason === "error" || msg.stopReason === "aborted") {
+    throw new Error(msg.errorMessage ?? msg.stopReason);
+  }
+  return { text: extractText(msg.content), usage: msg.usage };
+}

package/src/bridge/pi-interactive.ts

@@ -0,0 +1,86 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { InteractiveDeps } from "../core/types.ts";
+import type { AskAnswer, AskQuestion } from "../sandbox/protocol.ts";
+import { formatError } from "../util/errors.ts";
+import { createTodoFallback } from "./fallback-todo.ts";
+
+interface ToolInvoker {
+  readonly callTool?: (name: string, params: unknown) => Promise<unknown>;
+}
+
+function hasAnswers(value: unknown): value is { readonly answers: readonly unknown[] } {
+  return typeof value === "object" && value !== null && Array.isArray((value as { readonly answers?: unknown }).answers);
+}
+
+function isAskAnswer(value: unknown): value is AskAnswer {
+  if (typeof value !== "object" || value === null) return false;
+  const candidate = value as { readonly question?: unknown; readonly selected?: unknown; readonly custom?: unknown };
+  return typeof candidate.question === "string"
+    && Array.isArray(candidate.selected)
+    && candidate.selected.every((item) => typeof item === "string")
+    && (candidate.custom === undefined || typeof candidate.custom === "string");
+}
+
+function normalizeAnswers(result: unknown): AskAnswer[] | undefined {
+  if (!hasAnswers(result)) return undefined;
+  const answers = result.answers;
+  return answers.every(isAskAnswer) ? Array.from(answers) : undefined;
+}
+
+async function askViaUi(ctx: ExtensionContext, questions: readonly AskQuestion[]): Promise<AskAnswer[]> {
+  if (!ctx.hasUI) throw new Error("ask_user_question requires UI");
+  const answers = new Array<AskAnswer>(questions.length);
+  for (let i = 0; i < questions.length; i++) {
+    const q = questions[i];
+    if (!q) {
+      answers[i] = { question: "", selected: [], custom: formatError("malformed question") };
+      continue;
+    }
+    if (q.multiSelect) {
+      const selected: string[] = [];
+      while (true) {
+        const pick = await ctx.ui.select(`${q.header}: ${q.question}`, [...q.options.map((o) => o.label), "Done"]);
+        if (!pick || pick === "Done") break;
+        if (!selected.includes(pick)) selected.push(pick);
+      }
+      answers[i] = { question: q.question, selected };
+      continue;
+    }
+    const pick = await ctx.ui.select(`${q.header}: ${q.question}`, [...q.options.map((o) => o.label), "Type something."]);
+    if (!pick) answers[i] = { question: q.question, selected: [], custom: formatError("user cancelled") };
+    else if (pick === "Type something.") answers[i] = { question: q.question, selected: [], custom: await ctx.ui.input(q.question) ?? "" };
+    else answers[i] = { question: q.question, selected: [pick] };
+  }
+  return answers;
+}
+
+export function createPiInteractiveDeps(ctx: ExtensionContext): InteractiveDeps {
+  const fallbackTodo = createTodoFallback();
+  return Object.freeze({
+    onAskUserQuestion: async (questions: readonly AskQuestion[]): Promise<AskAnswer[]> => {
+      const callTool = (ctx as unknown as ToolInvoker).callTool;
+      if (typeof callTool === "function") {
+        try {
+          const result = await callTool.call(ctx, "ask_user_question", { questions });
+          const answers = normalizeAnswers(result);
+          if (answers) return answers;
+        } catch {
+          // Fall through to native UI fallback when the extension tool is not registered or fails.
+        }
+      }
+      return askViaUi(ctx, questions);
+    },
+    onTodo: async (action: string, params: Record<string, unknown>): Promise<string> => {
+      const callTool = (ctx as unknown as ToolInvoker).callTool;
+      if (typeof callTool === "function") {
+        try {
+          const result = await callTool.call(ctx, "todo", { action, ...params });
+          return typeof result === "string" ? result : JSON.stringify(result);
+        } catch {
+          // Fall through to in-process task store when the extension tool is not registered or fails.
+        }
+      }
+      return fallbackTodo(action, params);
+    },
+  });
+}

package/src/bridge/rlm-query.ts

@@ -0,0 +1,78 @@
+/**
+ * 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`.
+ */
+
+import type { RunRlm } from "../core/types.ts";
+import type { LlmBridge } from "./llm-query.ts";
+import type { RlmEmitter } from "../tool/rlm-events.ts";
+import { checkResourceLimits } from "../core/resource-limits.ts";
+import { formatError } from "../util/errors.ts";
+import { mapPool } from "../util/concurrency.ts";
+
+export interface RlmHandlers {
+  rlmQuery(prompt: string, model: string | null, depth: number): Promise<string>;
+  rlmQueryBatched(prompts: string[], model: string | null, depth: number): Promise<string[]>;
+}
+
+export interface RlmBridgeOptions {
+  readonly run: RunRlm;
+  readonly llm: LlmBridge;
+  /** Live RlmDetails reporting via onUpdate. Required — replaces SubcallObserver for recursive subcalls. */
+  readonly emitter: RlmEmitter;
+  readonly maxDepth: number;
+  readonly maxConcurrent: number;
+  /** Parent subcall ID that this run is attached under. */
+  readonly parentNodeId?: string;
+  /** Returns the parent's remaining budget/timeout for seeding child runs. */
+  readonly remainingBudget?: () => { readonly budgetUsd?: number; readonly timeoutMs?: number };
+  /** 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 {
+  async function child(prompt: string, model: string | null, depth: number): Promise<string> {
+    const childDepth = depth + 1;
+    // At the cap, a child RLM would just be an LM — short-circuit to a one-shot llm_query.
+    if (childDepth >= opts.maxDepth) return opts.llm.llmQuery(prompt, model, depth);
+    let subId: string | undefined;
+    try {
+      const rem = opts.remainingBudget?.() ?? {};
+      // 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) return limitError;
+      subId = opts.emitter.emitSubcallCreated({
+        kind: "rlm", parentId: opts.parentNodeId, label: "rlm_query",
+        model: model ?? undefined, detail: prompt.slice(0, 60),
+        depth: childDepth,
+      });
+      const res = 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);
+      opts.emitter.emitSubcallUpdated({ id: subId,
+        status: "done", resultPreview: res.answer.slice(0, 200),
+      });
+      return res.answer;
+    } catch (err) {
+      const msg = err instanceof Error ? err.message : String(err);
+      if (subId) opts.emitter.emitSubcallUpdated({ id: subId, status: "error", detail: msg });
+      return formatError(`child RLM failed - ${msg}`);
+    }
+  }
+
+  return {
+    rlmQuery: (prompt, model, depth) => child(prompt, model, depth),
+    rlmQueryBatched: (prompts, model, depth) => mapPool(prompts, opts.maxConcurrent, (p) => child(p, model, depth)),
+  };
+}

package/src/commands/rlm-config.ts

@@ -0,0 +1,42 @@
+/** `/rlm-config` — choose worker model, reasoning level, and run settings (smart is always pi's active model). */
+
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { modelRef } from "../config/settings.ts";
+import { cheapestModel, type RlmController } from "../mode/rlm-mode.ts";
+import { setRlmModeStatus } from "../ui/status.ts";
+import { showConfigPanel } from "../ui/config-panel.ts";
+import { selectModel } from "../ui/model-picker.ts";
+
+export async function runRlmConfig(controller: RlmController, ctx: ExtensionContext): Promise<boolean> {
+  const models = ctx.modelRegistry.getAvailable();
+
+  const worker = await selectModel(ctx, "Worker model (sub-LLM / llm_query)", models, controller.workerModel, controller.config.subSampling.reasoning);
+  if (worker) {
+    controller.workerModel = worker.model;
+    controller.config.subSampling.reasoning = worker.thinkingLevel;
+  }
+
+  await showConfigPanel(ctx, controller.config);
+
+  const effectiveWorker = controller.workerModel ?? cheapestModel(ctx.modelRegistry);
+  controller.savedWorkerRef = modelRef(controller.workerModel) ?? modelRef(effectiveWorker);
+  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);
+
+  const w = controller.workerModel;
+  ctx.ui.notify(
+    `RLM: worker=${w ? `${w.provider}/${w.id}` : "(cheapest)"}${controller.config.subSampling.reasoning ? `/${controller.config.subSampling.reasoning}` : ""}`,
+    "info",
+  );
+  return worker !== undefined;
+}
+
+export function registerRlmConfigCommand(pi: ExtensionAPI, controller: RlmController): void {
+  pi.registerCommand("rlm-config", {
+    description: "Configure RLM worker model and run settings.",
+    handler: async (_args, ctx) => {
+      await runRlmConfig(controller, ctx);
+    },
+  });
+}

package/src/commands/rlm.ts

@@ -0,0 +1,165 @@
+/** `/rlm` — toggle persistent Recursive Language Model mode. */
+
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { createPiInteractiveDeps } from "../bridge/pi-interactive.ts";
+import type { RlmController, RunHandle } from "../mode/rlm-mode.ts";
+import { postRlmGuide } from "../ui/intro.ts";
+import { clearRlmStatus, setRlmModeStatus } from "../ui/status.ts";
+import { listRunIds, readContextSidecar, readHeader, resolveRunId } from "../state/index.ts";
+import { DEFAULT_RUN_DIR } from "../config/defaults.ts";
+import { reconstructRlmState } from "../state/resume.ts";
+import type { ReconstructResult } from "../state/resume.ts";
+import type { RunHeader } from "../state/rows.ts";
+import { buildRlmSystemPrompt } from "../prompts/system.ts";
+import { RlmEmitter } from "../tool/rlm-events.ts";
+import { RlmEventAggregator } from "../tool/rlm-aggregator.ts";
+import { createTelemetrySink } from "../telemetry/index.ts";
+import { applyEdits } from "../patch/index.ts";
+import { tryExtractDiff } from "../core/answer.ts";
+
+export function registerRlmCommand(pi: ExtensionAPI, controller: RlmController): void {
+  pi.registerCommand("rlm", {
+    description: "Toggle persistent RLM mode (route plain prompts through the RLM engine).",
+    handler: async (_args, ctx) => {
+      const enabled = controller.toggle();
+      setRlmModeStatus(ctx.ui, controller);
+      ctx.ui.notify(`RLM mode ${enabled ? "ON" : "OFF"}`, "info");
+    },
+  });
+
+  pi.registerCommand("rlm-stop", {
+    description: "Abort the in-progress RLM run.",
+    handler: async (_args, ctx) => {
+      if (!controller.isBusy()) {
+        ctx.ui.notify("No RLM run in progress.", "info");
+        return;
+      }
+      controller.abort();
+      ctx.ui.notify("RLM run aborted.", "info");
+    },
+  });
+
+  pi.registerCommand("rlm-help", {
+    description: "Show the RLM startup guide and command cheatsheet.",
+    handler: async () => {
+      postRlmGuide(pi, controller);
+    },
+  });
+
+  pi.registerCommand("rlm-resume", {
+    description: "Resume an interrupted RLM run (default @latest).",
+    handler: async (args, ctx) => {
+      if (controller.isBusy()) {
+        ctx.ui.notify("RLM is busy (use /rlm-stop to cancel).", "warning");
+        return;
+      }
+      const ref = args.trim() || "@latest";
+      const dir = controller.config.runLog?.dir ?? DEFAULT_RUN_DIR;
+      const cwd = ctx.cwd ?? process.cwd();
+      const runId = await resolveRunId(cwd, dir, ref);
+      if (!runId) { ctx.ui.notify(`No resumable RLM run for '${ref}'.`, "error"); return; }
+      const header = await readHeader(cwd, dir, runId);
+      if (!header) { ctx.ui.notify(`Run ${runId} has no header.`, "error"); return; }
+      const systemPrompt = buildRlmSystemPrompt(
+        { contextType: header.context.type, contextChars: header.context.chars, rootPrompt: header.rootPrompt },
+        {
+          orchestrator: header.meta.orchestrator,
+          recursion: 1 < header.meta.maxDepth,
+          askUserQuestion: controller.config.askUserQuestion,
+          todo: controller.config.todo,
+        },
+      );
+      let recon: ReconstructResult;
+      try { recon = await reconstructRlmState(cwd, dir, runId, systemPrompt); }
+      catch (e) {
+        ctx.ui.notify(`RLM resume failed: corrupt run state — ${e instanceof Error ? e.message : String(e)}`, "error");
+        return;
+      }
+      if (!recon.ok) { ctx.ui.notify(`Cannot resume ${runId}: ${recon.reason}.`, "error"); return; }
+      if (recon.terminated) { ctx.ui.notify(`Run ${runId} already finished.`, "info"); return; }
+      const context = await readContextSidecar(cwd, dir, runId, header.context.json);
+      if (context === undefined) // R-C2: warn instead of silently resuming on empty context
+        ctx.ui.notify(`Warning: context sidecar missing for ${runId} — resuming without original context.`, "warning");
+      await executeRlmRunWithResume(pi, controller, ctx, recon, header, context ?? "");
+    },
+  });
+
+  pi.registerCommand("rlm-runs", {
+    description: "List recent RLM runs.",
+    handler: async (_args, ctx) => {
+      const dir = controller.config.runLog?.dir ?? DEFAULT_RUN_DIR;
+      const ids = (await listRunIds(ctx.cwd ?? process.cwd(), dir)).slice(0, 20);
+      ctx.ui.notify(ids.length ? ids.join("\n") : "No RLM runs recorded.", "info");
+    },
+  });
+
+  pi.registerShortcut?.("ctrl+shift+r", {
+    description: "Toggle RLM mode (off also stops a running query)",
+    handler: async (ctx) => {
+      const enabled = controller.toggle();
+      setRlmModeStatus(ctx.ui, controller);
+      ctx.ui.notify(`RLM mode ${enabled ? "ON" : "OFF"}`, "info");
+    },
+  });
+}
+
+async function executeRlmRunWithResume(
+  pi: ExtensionAPI,
+  controller: RlmController,
+  ctx: ExtensionContext,
+  recon: ReconstructResult & { ok: true },
+  header: RunHeader,
+  context: unknown,
+): Promise<void> {
+  let handle: RunHandle | undefined;
+  let sink: Awaited<ReturnType<typeof createTelemetrySink>>;
+  let emitter: RlmEmitter | undefined;
+  let detachSink: (() => void) | undefined;
+  let aggregator: RlmEventAggregator | undefined;
+  try {
+    sink = await createTelemetrySink(controller.config.telemetry);
+    emitter = new RlmEmitter();
+    if (sink) detachSink = emitter.attachSink(sink);
+    aggregator = new RlmEventAggregator(emitter, (partial) => {
+      const d = partial.details;
+      if (!d) return;
+      const turn = d.turns.max > 0 ? ` · turn ${d.turns.current}/${d.turns.max}` : "";
+      const cost = d.totals.costUsd > 0 ? ` · $${d.totals.costUsd.toFixed(4)}` : "";
+      const glyph = d.status === "running" ? "⏳" : d.status === "done" ? "✓" : "✗";
+      ctx.ui.setWidget?.("rlm-status", [`${glyph} RLM resume${turn}${cost}`], { placement: "aboveEditor" });
+    });
+    emitter.emitRootPrompt(header.rootPrompt);
+    const interactive = createPiInteractiveDeps(ctx);
+    if (controller.config.todo) {
+      for (const row of recon.todoRows) await interactive.onTodo?.(row.action, row.params);
+    }
+    handle = controller.start(ctx, { kind: "resume", resume: recon, context }, emitter, {
+      onAskUserQuestion: controller.config.askUserQuestion ? interactive.onAskUserQuestion : undefined,
+      onTodo: controller.config.todo ? interactive.onTodo : undefined,
+    });
+  } catch (e) {
+    ctx.ui.notify(`RLM resume failed: ${e instanceof Error ? e.message : String(e)}`, "error");
+    return;
+  }
+  pi.sendMessage({ customType: "rlm-question", content: `[resume] ${header.rootPrompt}`, display: true });
+  const { done } = handle;
+  try {
+    const result = await done;
+    pi.sendMessage({ customType: "rlm-answer", content: result.answer, display: true });
+    const proposedDiffs = result.diffs?.length ? result.diffs : tryExtractDiff(result.answer);
+    await applyEdits(
+      result.edits ?? [],
+      proposedDiffs,
+      ctx,
+    );
+  } catch (e) {
+    ctx.ui.notify(`RLM resume failed: ${e instanceof Error ? e.message : String(e)}`, "error");
+  } finally {
+    clearRlmStatus(ctx.ui);
+    ctx.ui.setWidget?.("rlm-status", undefined);
+    detachSink?.();
+    aggregator?.dispose();
+    emitter?.shutdown();
+    try { await sink?.shutdown(); } catch { /* best-effort */ }
+  }
+}

package/src/config/defaults.ts

@@ -0,0 +1,38 @@
+import type { RlmConfig } from "../core/types.ts";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+export const DEFAULT_RUN_DIR = join(tmpdir(), "rlm-runs");
+
+/** Frozen default sub-LLM system prompt — avoids re-allocation on every llm_query call. */
+const DEFAULT_SUB_SYSTEM_PROMPT =
+  "Answer directly and concisely. Return only the requested information. " +
+  "No preamble, no meta-commentary, no explanation of your approach. " +
+  "If listing items, use compact bullet form.";
+
+export const DEFAULT_CONFIG: Readonly<RlmConfig> = Object.freeze({
+  enabled: true,
+  maxDepth: 4,
+  maxIterations: 30,
+  execTimeoutS: 120,
+  requestTimeoutMs: 10 * 60_000,
+  maxConcurrentSubcalls: 4,
+  maxPromptChars: 400_000,
+  maxErrors: 5,
+  orchestrator: true,
+  compaction: true,
+  compactionThresholdPct: 0.65,
+  python: "python3",
+  sandboxInitTimeoutMs: 30_000,
+  askUserQuestion: true,
+  todo: true,
+  rootSampling: Object.freeze({ maxTokens: 16_384 }),
+  subSystemPrompt: DEFAULT_SUB_SYSTEM_PROMPT,
+  subSampling: Object.freeze({ maxTokens: 8192 }),
+  runLog: Object.freeze({
+    enabled: true,
+    dir: DEFAULT_RUN_DIR,
+    snapshot: true,
+    maxRuns: 50,
+  }),
+});