pi-rlm 0.1.0

package/src/engine.ts

@@ -0,0 +1,462 @@
+import { promises as fs } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+import { completeWithBackend } from "./backends";
+import { plannerPrompt, solverPrompt, synthesisPrompt } from "./prompts";
+import { PlannerDecision, RlmNode, RlmRunResult, RunArtifacts, StartRunInput } from "./types";
+import { extractFirstJsonObject, normalizeTask, shortTask, toErrorMessage } from "./utils";
+
+interface EngineInput extends StartRunInput {
+  runId: string;
+}
+
+interface EngineState {
+  nodeCounter: number;
+  nodesVisited: number;
+  maxDepthSeen: number;
+}
+
+type ProgressFn = (line: string) => void;
+
+export async function runRlmEngine(
+  input: EngineInput,
+  ctx: ExtensionContext,
+  signal?: AbortSignal,
+  progress?: ProgressFn
+): Promise<RlmRunResult> {
+  const startedAt = Date.now();
+  const artifacts = await createArtifacts(input.runId);
+  const log = createEventLogger(artifacts.eventsPath);
+  const state: EngineState = {
+    nodeCounter: 0,
+    nodesVisited: 0,
+    maxDepthSeen: 0
+  };
+  const activeSignal = signal ?? new AbortController().signal;
+
+  progress?.(`RLM run ${input.runId} started (${input.backend}, mode=${input.mode})`);
+  log("run_start", {
+    runId: input.runId,
+    backend: input.backend,
+    mode: input.mode,
+    maxDepth: input.maxDepth,
+    maxNodes: input.maxNodes,
+    maxBranching: input.maxBranching,
+    concurrency: input.concurrency
+  });
+
+  try {
+    const root = await runNode({ task: input.task, depth: 0, lineage: [] });
+
+    const finalOutput = root.result ?? "(no final output)";
+
+    await fs.writeFile(artifacts.treePath, JSON.stringify(root, null, 2), "utf8");
+    await fs.writeFile(artifacts.outputPath, finalOutput, "utf8");
+
+    const durationMs = Date.now() - startedAt;
+    const result: RlmRunResult = {
+      runId: input.runId,
+      backend: input.backend,
+      final: finalOutput,
+      root,
+      artifacts,
+      stats: {
+        nodesVisited: state.nodesVisited,
+        maxDepthSeen: state.maxDepthSeen,
+        durationMs
+      }
+    };
+
+    log("run_end", {
+      runId: input.runId,
+      durationMs,
+      nodesVisited: state.nodesVisited,
+      maxDepthSeen: state.maxDepthSeen,
+      finalChars: finalOutput.length
+    });
+
+    if (root.status === "failed") {
+      throw new Error(root.error ?? "RLM root node failed");
+    }
+
+    progress?.(`RLM run ${input.runId} completed in ${durationMs}ms`);
+    return result;
+  } finally {
+    await log.flush();
+  }
+
+  async function runNode(params: {
+    task: string;
+    depth: number;
+    lineage: string[];
+  }): Promise<RlmNode> {
+    const nodeId = `n${++state.nodeCounter}`;
+    state.nodesVisited += 1;
+    state.maxDepthSeen = Math.max(state.maxDepthSeen, params.depth);
+
+    const node: RlmNode = {
+      id: nodeId,
+      depth: params.depth,
+      task: params.task,
+      status: "running",
+      startedAt: Date.now(),
+      children: []
+    };
+
+    progress?.(`[${node.id}] depth=${params.depth} ${shortTask(params.task, 72)}`);
+    log("node_start", {
+      nodeId: node.id,
+      depth: params.depth,
+      task: params.task,
+      nodesVisited: state.nodesVisited
+    });
+
+    if (activeSignal.aborted) {
+      node.status = "cancelled";
+      node.error = "Run cancelled";
+      node.finishedAt = Date.now();
+      log("node_cancelled", { nodeId: node.id, depth: node.depth });
+      throw new Error("RLM run cancelled");
+    }
+
+    const normalized = normalizeTask(params.task);
+    const remainingNodeBudget = Math.max(0, input.maxNodes - state.nodesVisited);
+
+    try {
+      const forcedReason = getForcedSolveReason({
+        depth: params.depth,
+        normalizedTask: normalized,
+        lineage: params.lineage
+      });
+
+      if (forcedReason || input.mode === "solve") {
+        const reason = forcedReason ?? "mode=solve";
+        node.decision = { action: "solve", reason };
+        node.result = await solveNode(node, reason);
+        node.status = "completed";
+        node.finishedAt = Date.now();
+        log("node_end", {
+          nodeId: node.id,
+          action: "solve",
+          reason,
+          chars: node.result.length,
+          durationMs: node.finishedAt - node.startedAt
+        });
+        return node;
+      }
+
+      const decision = await planNode({
+        task: params.task,
+        depth: params.depth,
+        maxDepth: input.maxDepth,
+        maxBranching: input.maxBranching,
+        remainingNodeBudget
+      });
+
+      node.decision = {
+        action: decision.action,
+        reason: decision.reason
+      };
+
+      if (decision.action === "solve") {
+        node.result = await solveNode(node, decision.reason);
+        node.status = "completed";
+        node.finishedAt = Date.now();
+        log("node_end", {
+          nodeId: node.id,
+          action: "solve",
+          reason: decision.reason,
+          chars: node.result.length,
+          durationMs: node.finishedAt - node.startedAt
+        });
+        return node;
+      }
+
+      const subtasks = sanitizeSubtasks(decision.subtasks ?? [], params.task).slice(
+        0,
+        input.maxBranching
+      );
+
+      if (subtasks.length < 2) {
+        node.decision = {
+          action: "solve",
+          reason: "planner returned insufficient valid subtasks"
+        };
+        node.result = await solveNode(node, node.decision.reason);
+        node.status = "completed";
+        node.finishedAt = Date.now();
+        log("node_end", {
+          nodeId: node.id,
+          action: "solve",
+          reason: node.decision.reason,
+          chars: node.result.length,
+          durationMs: node.finishedAt - node.startedAt
+        });
+        return node;
+      }
+
+      progress?.(`[${node.id}] decomposing into ${subtasks.length} subtasks`);
+      log("node_decompose", {
+        nodeId: node.id,
+        subtasks,
+        reason: decision.reason
+      });
+
+      node.children = await mapConcurrent(subtasks, input.concurrency, async (subtask) => {
+        return runNode({
+          task: subtask,
+          depth: params.depth + 1,
+          lineage: [...params.lineage, normalized]
+        });
+      });
+
+      node.result = await synthesizeNode(node);
+      node.status = "completed";
+      node.finishedAt = Date.now();
+      log("node_end", {
+        nodeId: node.id,
+        action: "decompose",
+        chars: node.result.length,
+        children: node.children.length,
+        durationMs: node.finishedAt - node.startedAt
+      });
+      return node;
+    } catch (error) {
+      const message = toErrorMessage(error);
+      if (activeSignal.aborted || message.toLowerCase().includes("cancel")) {
+        node.status = "cancelled";
+        node.error = message;
+        node.finishedAt = Date.now();
+        log("node_cancelled", {
+          nodeId: node.id,
+          error: message,
+          durationMs: node.finishedAt - node.startedAt
+        });
+        throw error;
+      }
+
+      node.status = "failed";
+      node.error = message;
+      node.finishedAt = Date.now();
+      log("node_error", {
+        nodeId: node.id,
+        error: message,
+        durationMs: node.finishedAt - node.startedAt
+      });
+
+      node.result = `Node failed: ${message}`;
+      return node;
+    }
+  }
+
+  async function planNode(args: {
+    task: string;
+    depth: number;
+    maxDepth: number;
+    maxBranching: number;
+    remainingNodeBudget: number;
+  }): Promise<PlannerDecision> {
+    if (input.mode === "decompose") {
+      const forced = await callModel("planner", plannerPrompt(args));
+      const parsedForced = parsePlannerDecision(forced);
+      if (parsedForced.action === "decompose") {
+        return parsedForced;
+      }
+      return {
+        action: "decompose",
+        reason: "mode=decompose requested, but planner output was invalid",
+        subtasks: []
+      };
+    }
+
+    const raw = await callModel("planner", plannerPrompt(args));
+    return parsePlannerDecision(raw);
+  }
+
+  async function solveNode(node: RlmNode, forceReason: string): Promise<string> {
+    const prompt = solverPrompt({
+      task: node.task,
+      depth: node.depth,
+      maxDepth: input.maxDepth,
+      forceReason
+    });
+    return callModel("solver", prompt, node.id);
+  }
+
+  async function synthesizeNode(node: RlmNode): Promise<string> {
+    const prompt = synthesisPrompt({
+      task: node.task,
+      depth: node.depth,
+      children: node.children
+    });
+    return callModel("synthesizer", prompt, node.id);
+  }
+
+  async function callModel(stage: string, promptText: string, nodeId?: string): Promise<string> {
+    if (activeSignal.aborted) {
+      throw new Error("RLM run cancelled");
+    }
+
+    log("backend_call", {
+      nodeId,
+      stage,
+      backend: input.backend,
+      model: input.model,
+      promptChars: promptText.length
+    });
+
+    const output = await completeWithBackend(
+      {
+        backend: input.backend,
+        prompt: promptText,
+        cwd: input.cwd,
+        model: input.model,
+        toolsProfile: input.toolsProfile,
+        timeoutMs: input.timeoutMs,
+        signal: activeSignal
+      },
+      ctx
+    );
+
+    log("backend_result", {
+      nodeId,
+      stage,
+      outputChars: output.length
+    });
+
+    return output;
+  }
+
+  function getForcedSolveReason(args: {
+    depth: number;
+    normalizedTask: string;
+    lineage: string[];
+  }): string | undefined {
+    if (args.depth >= input.maxDepth) {
+      return "maxDepth reached";
+    }
+
+    if (state.nodesVisited >= input.maxNodes) {
+      return "maxNodes reached";
+    }
+
+    if (args.lineage.includes(args.normalizedTask)) {
+      return "cycle detected in task lineage";
+    }
+
+    return undefined;
+  }
+
+}
+
+function parsePlannerDecision(raw: string): PlannerDecision {
+  const jsonCandidate = extractFirstJsonObject(raw);
+  if (!jsonCandidate) {
+    return {
+      action: "solve",
+      reason: "planner JSON parse failed"
+    };
+  }
+
+  try {
+    const parsed = JSON.parse(jsonCandidate) as {
+      action?: unknown;
+      reason?: unknown;
+      subtasks?: unknown;
+    };
+
+    const action = parsed.action === "decompose" ? "decompose" : "solve";
+    const reason = typeof parsed.reason === "string" ? parsed.reason : "planner did not provide reason";
+    const subtasks = Array.isArray(parsed.subtasks)
+      ? parsed.subtasks.filter((item): item is string => typeof item === "string")
+      : undefined;
+
+    return { action, reason, subtasks };
+  } catch {
+    return {
+      action: "solve",
+      reason: "planner JSON was invalid"
+    };
+  }
+}
+
+function sanitizeSubtasks(subtasks: string[], parentTask: string): string[] {
+  const parentNormalized = normalizeTask(parentTask);
+  const deduped = new Set<string>();
+  const cleaned: string[] = [];
+
+  for (const subtask of subtasks) {
+    const value = subtask.trim();
+    if (!value) continue;
+
+    const normalized = normalizeTask(value);
+    if (!normalized || normalized === parentNormalized) continue;
+    if (deduped.has(normalized)) continue;
+
+    deduped.add(normalized);
+    cleaned.push(value);
+  }
+
+  return cleaned;
+}
+
+async function mapConcurrent<T, R>(
+  items: T[],
+  concurrency: number,
+  worker: (item: T, index: number) => Promise<R>
+): Promise<R[]> {
+  if (items.length === 0) return [];
+
+  const limit = Math.max(1, concurrency);
+  const results = new Array<R>(items.length);
+  let nextIndex = 0;
+
+  await Promise.all(
+    Array.from({ length: Math.min(limit, items.length) }, async () => {
+      while (nextIndex < items.length) {
+        const current = nextIndex;
+        nextIndex += 1;
+        results[current] = await worker(items[current], current);
+      }
+    })
+  );
+
+  return results;
+}
+
+async function createArtifacts(runId: string): Promise<RunArtifacts> {
+  const dir = join(tmpdir(), "pi-rlm-runs", runId);
+  await fs.mkdir(dir, { recursive: true });
+  return {
+    dir,
+    eventsPath: join(dir, "events.jsonl"),
+    treePath: join(dir, "tree.json"),
+    outputPath: join(dir, "output.md")
+  };
+}
+
+function createEventLogger(path: string): {
+  (type: string, payload: Record<string, unknown>): void;
+  flush: () => Promise<void>;
+} {
+  let tail = Promise.resolve();
+
+  const write = (type: string, payload: Record<string, unknown>): void => {
+    const line = `${JSON.stringify({
+      ts: new Date().toISOString(),
+      type,
+      ...payload
+    })}\n`;
+
+    tail = tail
+      .then(() => fs.appendFile(path, line, "utf8"))
+      .catch(() => undefined);
+  };
+
+  write.flush = async (): Promise<void> => {
+    await tail;
+  };
+
+  return write;
+}

package/src/prompts.ts

@@ -0,0 +1,83 @@
+import { RlmNode } from "./types";
+
+export function plannerPrompt(input: {
+  task: string;
+  depth: number;
+  maxDepth: number;
+  maxBranching: number;
+  remainingNodeBudget: number;
+}): string {
+  return [
+    "You are a recursion controller for a recursive language model run.",
+    "Decide whether the task should be solved directly, or decomposed into subtasks.",
+    "",
+    "Return ONLY a JSON object with this schema:",
+    '{"action":"solve"|"decompose","reason":"...","subtasks":["..."]}',
+    "",
+    "Rules:",
+    "- Use action=solve if the task is atomic enough for one model pass.",
+    "- Use action=decompose only when decomposition is clearly beneficial.",
+    `- If action=decompose, return 2 to ${input.maxBranching} subtasks (never more).`,
+    "- Subtasks must be clear, non-empty strings.",
+    "- Do not include markdown or prose outside JSON.",
+    "",
+    `Current depth: ${input.depth} / ${input.maxDepth}`,
+    `Remaining node budget: ${input.remainingNodeBudget}`,
+    "",
+    "Task:",
+    input.task
+  ].join("\n");
+}
+
+export function solverPrompt(input: {
+  task: string;
+  depth: number;
+  maxDepth: number;
+  forceReason?: string;
+}): string {
+  return [
+    "You are a worker node in a recursive language model run.",
+    "Solve the task directly and return a concrete answer.",
+    "",
+    `Depth: ${input.depth} / ${input.maxDepth}`,
+    input.forceReason ? `Note: forced direct solve because ${input.forceReason}` : "",
+    "",
+    "Task:",
+    input.task
+  ]
+    .filter(Boolean)
+    .join("\n");
+}
+
+export function synthesisPrompt(input: {
+  task: string;
+  depth: number;
+  children: RlmNode[];
+}): string {
+  const childBlocks = input.children
+    .map((child, idx) => {
+      const status = child.status.toUpperCase();
+      const result = child.result ?? child.error ?? "(no output)";
+      return [
+        `### Child ${idx + 1} (${status})`,
+        `Subtask: ${child.task}`,
+        "Output:",
+        result
+      ].join("\n");
+    })
+    .join("\n\n");
+
+  return [
+    "You are the synthesizer node in a recursive language model run.",
+    "Combine child results into one final response to the parent task.",
+    "Be explicit about uncertainties if child outputs conflict.",
+    "",
+    `Depth: ${input.depth}`,
+    "",
+    "Parent task:",
+    input.task,
+    "",
+    "Child outputs:",
+    childBlocks
+  ].join("\n");
+}

package/src/runs.ts

@@ -0,0 +1,128 @@
+import { RlmRunResult, RunRecord, RunStatus, StartRunInput } from "./types";
+import { createRunId, toErrorMessage } from "./utils";
+
+const maxRecords = 200;
+
+export class RunStore {
+  private readonly records = new Map<string, RunRecord>();
+
+  start(
+    input: StartRunInput,
+    executor: (runId: string, signal: AbortSignal) => Promise<RlmRunResult>,
+    externalSignal?: AbortSignal
+  ): RunRecord {
+    const id = createRunId();
+    const controller = new AbortController();
+
+    if (externalSignal) {
+      if (externalSignal.aborted) {
+        controller.abort();
+      } else {
+        externalSignal.addEventListener(
+          "abort",
+          () => {
+            controller.abort();
+          },
+          { once: true }
+        );
+      }
+    }
+
+    const record: RunRecord = {
+      id,
+      input,
+      status: "running",
+      createdAt: Date.now(),
+      startedAt: Date.now(),
+      controller,
+      promise: Promise.resolve(null as unknown as RlmRunResult)
+    };
+
+    record.promise = (async () => {
+      try {
+        const result = await executor(id, controller.signal);
+        record.status = "completed";
+        record.finishedAt = Date.now();
+        record.result = result;
+        return result;
+      } catch (error) {
+        const message = toErrorMessage(error);
+        record.finishedAt = Date.now();
+        if (controller.signal.aborted || message.toLowerCase().includes("cancel")) {
+          record.status = "cancelled";
+        } else {
+          record.status = "failed";
+        }
+        record.error = message;
+        throw error;
+      } finally {
+        this.prune();
+      }
+    })();
+
+    this.records.set(id, record);
+    this.prune();
+    return record;
+  }
+
+  get(id: string): RunRecord | undefined {
+    return this.records.get(id);
+  }
+
+  list(): RunRecord[] {
+    return Array.from(this.records.values()).sort((a, b) => b.createdAt - a.createdAt);
+  }
+
+  async wait(id: string, timeoutMs: number): Promise<{
+    status: RunStatus;
+    record: RunRecord;
+    done: boolean;
+  }> {
+    const record = this.records.get(id);
+    if (!record) {
+      throw new Error(`Unknown run id: ${id}`);
+    }
+
+    if (record.status !== "running") {
+      return { status: record.status, record, done: true };
+    }
+
+    let timeoutHandle: NodeJS.Timeout | undefined;
+    try {
+      await Promise.race([
+        record.promise.then(() => undefined).catch(() => undefined),
+        new Promise((resolve) => {
+          timeoutHandle = setTimeout(resolve, timeoutMs);
+        })
+      ]);
+    } finally {
+      if (timeoutHandle) clearTimeout(timeoutHandle);
+    }
+
+    const done = record.status !== "running";
+    return { status: record.status, record, done };
+  }
+
+  cancel(id: string): RunRecord {
+    const record = this.records.get(id);
+    if (!record) {
+      throw new Error(`Unknown run id: ${id}`);
+    }
+
+    if (record.status === "running") {
+      record.controller.abort();
+    }
+
+    return record;
+  }
+
+  private prune(): void {
+    const records = this.list();
+    if (records.length <= maxRecords) return;
+
+    for (const record of records.slice(maxRecords)) {
+      if (record.status === "running") continue;
+      this.records.delete(record.id);
+    }
+  }
+}

package/src/schema.ts

@@ -0,0 +1,29 @@
+import { StringEnum } from "@mariozechner/pi-ai";
+import { Static, Type } from "@sinclair/typebox";
+
+const opSchema = StringEnum(["start", "status", "wait", "cancel"] as const);
+const backendSchema = StringEnum(["sdk", "cli", "tmux"] as const);
+const modeSchema = StringEnum(["auto", "solve", "decompose"] as const);
+const toolsProfileSchema = StringEnum(["coding", "read-only"] as const);
+
+export const rlmToolParamsSchema = Type.Object({
+  op: Type.Optional(opSchema),
+  id: Type.Optional(Type.String({ description: "Run ID for status/wait/cancel" })),
+  task: Type.Optional(Type.String({ description: "Task to solve recursively" })),
+
+  backend: Type.Optional(backendSchema),
+  mode: Type.Optional(modeSchema),
+  async: Type.Optional(Type.Boolean({ description: "Return immediately and run in background" })),
+  model: Type.Optional(Type.String({ description: "Optional provider/model[:thinking] override" })),
+  cwd: Type.Optional(Type.String({ description: "Working directory for subcalls" })),
+  toolsProfile: Type.Optional(toolsProfileSchema),
+
+  maxDepth: Type.Optional(Type.Integer({ minimum: 0, maximum: 8 })),
+  maxNodes: Type.Optional(Type.Integer({ minimum: 1, maximum: 300 })),
+  maxBranching: Type.Optional(Type.Integer({ minimum: 1, maximum: 8 })),
+  concurrency: Type.Optional(Type.Integer({ minimum: 1, maximum: 8 })),
+  timeoutMs: Type.Optional(Type.Integer({ minimum: 1000, maximum: 3600000 })),
+  waitTimeoutMs: Type.Optional(Type.Integer({ minimum: 100, maximum: 3600000 }))
+});
+
+export type RlmToolParams = Static<typeof rlmToolParamsSchema>;