vers-rlm pi extension — Recursive Language Model with branchable Firecracker VMs as sandbox (snapshot, branch, restore)

index.ts

/**
 * Vers RLM Extension
 *
 * Combines Recursive Language Models with Vers VMs to give the sub-agent
 * a full Linux environment it can explore iteratively. The sub-agent writes
 * JavaScript that calls VM tools (exec, readFile, writeFile) and LLM tools
 * (llmQuery). It can also snapshot and branch the VM for speculative execution.
 *
 * The key idea: RLM provides iterative reasoning through code. Vers VMs
 * provide branchable, snapshotable compute. Together you get branchable
 * reasoning — fork a line of investigation, try multiple approaches, keep
 * the best result.
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";
import { execFile } from "node:child_process";
import { writeFile, mkdir } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
  RLM,
  type LanguageModel,
  type Message,
  type ToolFunction,
} from "ts-rlm";

// =============================================================================
// Vers API Client (minimal, self-contained)
// =============================================================================

const DEFAULT_BASE_URL = "https://api.vers.sh/api/v1";

function loadVersKey(): string {
  try {
    const homedir = process.env.HOME || process.env.USERPROFILE || "";
    const keysPath = join(homedir, ".vers", "keys.json");
    const data = require("fs").readFileSync(keysPath, "utf-8");
    return JSON.parse(data)?.keys?.VERS_API_KEY || "";
  } catch {
    return "";
  }
}

class VersAPI {
  private apiKey: string;
  private baseURL: string;
  private keyPathCache = new Map<string, string>();

  constructor() {
    this.apiKey =
      process.env.VERS_API_KEY || loadVersKey() || "";
    this.baseURL = (
      process.env.VERS_BASE_URL || DEFAULT_BASE_URL
    ).replace(/\/$/, "");
  }

  private async request<T>(
    method: string,
    path: string,
    body?: unknown
  ): Promise<T> {
    const res = await fetch(`${this.baseURL}${path}`, {
      method,
      headers: {
        "Content-Type": "application/json",
        ...(this.apiKey ? { Authorization: `Bearer ${this.apiKey}` } : {}),
      },
      body: body !== undefined ? JSON.stringify(body) : undefined,
    });
    if (!res.ok) {
      const text = await res.text().catch(() => "");
      throw new Error(`Vers API ${method} ${path} (${res.status}): ${text}`);
    }
    const ct = res.headers.get("content-type") || "";
    if (ct.includes("application/json")) return res.json() as Promise<T>;
    return undefined as T;
  }

  async createVM(config?: {
    vcpu_count?: number;
    mem_size_mib?: number;
    fs_size_mib?: number;
  }): Promise<string> {
    const resp = await this.request<{ vm_id: string }>(
      "POST",
      "/vm/new_root?wait_boot=true",
      { vm_config: config || {} }
    );
    return resp.vm_id;
  }

  async deleteVM(vmId: string): Promise<void> {
    await this.request("DELETE", `/vm/${encodeURIComponent(vmId)}`);
  }

  async commitVM(vmId: string): Promise<string> {
    const resp = await this.request<{ commit_id: string }>(
      "POST",
      `/vm/${encodeURIComponent(vmId)}/commit`
    );
    return resp.commit_id;
  }

  async branchVM(vmId: string): Promise<string> {
    const resp = await this.request<{ vm_id: string }>(
      "POST",
      `/vm/${encodeURIComponent(vmId)}/branch`
    );
    return resp.vm_id;
  }

  async restoreVM(commitId: string): Promise<string> {
    const resp = await this.request<{ vm_id: string }>(
      "POST",
      "/vm/from_commit",
      { commit_id: commitId }
    );
    return resp.vm_id;
  }

  async getSSHKey(
    vmId: string
  ): Promise<{ ssh_port: number; ssh_private_key: string }> {
    return this.request("GET", `/vm/${encodeURIComponent(vmId)}/ssh_key`);
  }

  async ensureKeyFile(vmId: string): Promise<string> {
    const existing = this.keyPathCache.get(vmId);
    if (existing) return existing;

    const keyInfo = await this.getSSHKey(vmId);
    const keyDir = join(tmpdir(), "vers-rlm-keys");
    await mkdir(keyDir, { recursive: true });
    const keyPath = join(keyDir, `vers-${vmId.slice(0, 12)}.pem`);
    await writeFile(keyPath, keyInfo.ssh_private_key, { mode: 0o600 });
    this.keyPathCache.set(vmId, keyPath);
    return keyPath;
  }

  async sshArgs(vmId: string): Promise<string[]> {
    const keyPath = await this.ensureKeyFile(vmId);
    const hostname = `${vmId}.vm.vers.sh`;
    return [
      "-i", keyPath,
      "-o", "StrictHostKeyChecking=no",
      "-o", "UserKnownHostsFile=/dev/null",
      "-o", "LogLevel=ERROR",
      "-o", "ConnectTimeout=30",
      "-o",
      `ProxyCommand=openssl s_client -connect %h:443 -servername %h -quiet 2>/dev/null`,
      `root@${hostname}`,
    ];
  }

  async exec(
    vmId: string,
    command: string,
    timeoutMs = 120000
  ): Promise<{ stdout: string; stderr: string; exitCode: number }> {
    const args = await this.sshArgs(vmId);
    return new Promise((resolve, reject) => {
      execFile(
        "ssh",
        [...args, command],
        { maxBuffer: 10 * 1024 * 1024, timeout: timeoutMs },
        (err, stdout, stderr) => {
          if (
            err &&
            typeof (err as any).code === "string" &&
            stdout === "" &&
            stderr === ""
          ) {
            reject(new Error(`SSH failed: ${err.message}`));
            return;
          }
          const exitCode = (err as any)?.status ?? (err ? 1 : 0);
          resolve({
            stdout: stdout?.toString() ?? "",
            stderr: stderr?.toString() ?? "",
            exitCode,
          });
        }
      );
    });
  }
}

// =============================================================================
// PiLanguageModel adapter (same as in rlm extension)
// =============================================================================

class PiLanguageModel implements LanguageModel {
  constructor(private ctx: { modelRegistry: any; model: any }) {}

  async complete(messages: Message[]): Promise<string> {
    const model = this.ctx.model;
    if (!model) throw new Error("No model configured in pi.");

    const apiKey = await this.ctx.modelRegistry.getApiKey(model);
    if (!apiKey) throw new Error(`No API key for provider "${model.provider}".`);

    const api = model.api;
    const baseUrl = model.baseUrl;

    if (api === "anthropic-messages") {
      return this.completeAnthropic(messages, model, apiKey, baseUrl);
    } else {
      return this.completeOpenAI(messages, model, apiKey, baseUrl);
    }
  }

  private async completeAnthropic(
    messages: Message[], model: any, apiKey: string, baseUrl?: string
  ): Promise<string> {
    const base = baseUrl || "https://api.anthropic.com";
    const url = `${base}${base.endsWith("/v1") ? "" : "/v1"}/messages`;
    const systemMsg = messages.find((m) => m.role === "system");
    const nonSystem = messages.filter((m) => m.role !== "system");

    const res = await fetch(url, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "x-api-key": apiKey,
        "anthropic-version": "2023-06-01",
      },
      body: JSON.stringify({
        model: model.id,
        max_tokens: 8192,
        system: systemMsg?.content,
        messages: nonSystem.map((m) => ({ role: m.role, content: m.content })),
      }),
    });

    if (!res.ok) {
      const text = await res.text();
      throw new Error(`Anthropic API error (${res.status}): ${text}`);
    }

    const data = (await res.json()) as any;
    return data.content?.find((c: any) => c.type === "text")?.text ?? "";
  }

  private async completeOpenAI(
    messages: Message[], model: any, apiKey: string, baseUrl?: string
  ): Promise<string> {
    const base = baseUrl || "https://api.openai.com";
    const url = `${base}${base.includes("/v1") ? "" : "/v1"}/chat/completions`;

    const res = await fetch(url, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${apiKey}`,
      },
      body: JSON.stringify({ model: model.id, messages }),
    });

    if (!res.ok) {
      const text = await res.text();
      throw new Error(`OpenAI API error (${res.status}): ${text}`);
    }

    const data = (await res.json()) as any;
    return data.choices?.[0]?.message?.content ?? "";
  }
}

// =============================================================================
// VM Tool Factories — these create the tools the RLM sub-agent can call
// =============================================================================

function createVMTools(
  vers: VersAPI,
  vmIdRef: { current: string },
  onUpdate?: (text: string) => void
): Record<string, ToolFunction> {
  const log = (msg: string) => onUpdate?.(msg);

  return {
    /**
     * Execute a bash command on the VM. Returns stdout+stderr.
     * Usage: await exec("ls -la /tmp")
     * Usage: await exec("python3 -c 'print(1+1)'")
     * Usage: await exec("apt-get install -y jq && cat data.json | jq '.items[]'")
     */
    async exec(command: string): Promise<string> {
      log(`🖥️  exec: ${command.slice(0, 120)}${command.length > 120 ? "…" : ""}`);
      const result = await vers.exec(vmIdRef.current, command);
      const output = [result.stdout, result.stderr].filter(Boolean).join("\n");
      const exitInfo =
        result.exitCode !== 0 ? `\n[exit code: ${result.exitCode}]` : "";
      return (output + exitInfo).trim() || "(no output)";
    },

    /**
     * Read a file from the VM. Returns the file contents as a string.
     * Usage: await readFile("/etc/hostname")
     */
    async readFile(path: string): Promise<string> {
      log(`📄 readFile: ${path}`);
      const result = await vers.exec(vmIdRef.current, `cat ${JSON.stringify(path)}`);
      if (result.exitCode !== 0) {
        return `[Error reading ${path}]: ${result.stderr}`;
      }
      return result.stdout;
    },

    /**
     * Write content to a file on the VM.
     * Usage: await writeFile("/tmp/script.py", "print('hello')")
     */
    async writeFile(path: string, content: string): Promise<string> {
      log(`✏️  writeFile: ${path}`);
      // Use heredoc to handle special characters
      const marker = "VERS_RLM_EOF_" + Math.random().toString(36).slice(2, 8);
      const cmd = `cat > ${JSON.stringify(path)} << '${marker}'\n${content}\n${marker}`;
      const result = await vers.exec(vmIdRef.current, cmd);
      if (result.exitCode !== 0) {
        return `[Error writing ${path}]: ${result.stderr}`;
      }
      return `Wrote ${content.length} bytes to ${path}`;
    },

    /**
     * Snapshot the current VM state. Returns a commit ID you can restore from.
     * Use this before risky operations so you can rollback.
     * Usage: const commitId = await vmSnapshot()
     */
    async vmSnapshot(): Promise<string> {
      log("📸 vmSnapshot: committing VM state...");
      const commitId = await vers.commitVM(vmIdRef.current);
      log(`📸 vmSnapshot: committed → ${commitId}`);
      return commitId;
    },

    /**
     * Branch the VM — create a clone with identical state.
     * Returns the new VM ID. The original VM continues unchanged.
     * Use for speculative execution: branch, try something risky on the
     * branch, check the result, then decide which VM to keep.
     * Usage: const branchVmId = await vmBranch()
     */
    async vmBranch(): Promise<string> {
      log("🔀 vmBranch: branching VM...");
      const newVmId = await vers.branchVM(vmIdRef.current);
      log(`🔀 vmBranch: created branch → ${newVmId}`);
      return newVmId;
    },

    /**
     * Switch execution to a different VM. Use after vmBranch() to
     * run commands on the branched VM.
     * Usage: vmSwitch(branchVmId)
     */
    async vmSwitch(newVmId: string): Promise<string> {
      const oldVmId = vmIdRef.current;
      vmIdRef.current = newVmId;
      log(`🔄 vmSwitch: ${oldVmId.slice(0, 8)}… → ${newVmId.slice(0, 8)}…`);
      return `Switched from ${oldVmId} to ${newVmId}`;
    },

    /**
     * Restore a VM from a previously created snapshot (commit ID).
     * Creates a new VM from that state and switches to it.
     * Usage: const restoredId = await vmRestore(commitId)
     */
    async vmRestore(commitId: string): Promise<string> {
      log(`⏪ vmRestore: restoring from ${commitId}...`);
      const newVmId = await vers.restoreVM(commitId);
      vmIdRef.current = newVmId;
      log(`⏪ vmRestore: restored → ${newVmId}`);
      return newVmId;
    },

    /**
     * Get the current VM ID.
     * Usage: const vmId = vmId()
     */
    async currentVmId(): Promise<string> {
      return vmIdRef.current;
    },
  };
}

// =============================================================================
// Tool instructions for the RLM sub-agent
// =============================================================================

const VM_TOOL_INSTRUCTIONS = `
You have access to a full Linux VM (Ubuntu) that you can control via these tools:

## VM Execution Tools
- \`exec(command)\` — Run any bash command on the VM. Returns stdout+stderr. You can install packages, compile code, run scripts, curl APIs, etc.
- \`readFile(path)\` — Read a file from the VM filesystem.
- \`writeFile(path, content)\` — Write a file to the VM filesystem.

## VM State Management Tools
- \`vmSnapshot()\` — Snapshot the VM state. Returns a commit ID. Use before risky operations.
- \`vmBranch()\` — Clone the VM into a new identical VM. Returns the new VM ID. The original continues unchanged.
- \`vmSwitch(vmId)\` — Switch execution to a different VM (e.g., a branch).
- \`vmRestore(commitId)\` — Restore from a snapshot. Creates a new VM and switches to it.
- \`currentVmId()\` — Get the current VM ID.

## Strategy
- Use \`exec()\` freely — you have a full Linux environment with root access
- Install tools as needed: \`exec("apt-get update && apt-get install -y jq ripgrep")\`
- For risky operations, snapshot first: \`const checkpoint = await vmSnapshot()\`
- For exploring alternatives, branch: \`const altVm = await vmBranch()\`, then \`vmSwitch(altVm)\`
- Use \`llmQuery(prompt)\` for semantic analysis of command output or file contents
- Call \`FINAL({...})\` when you have the answer

## Important
- Commands run as root on the VM
- The VM has network access (can curl, wget, git clone, etc.)
- State persists between exec() calls (installed packages, created files, etc.)
- Each branch is a full copy-on-write clone — instant and cheap
`.trim();

// =============================================================================
// Extension
// =============================================================================

export default function versRlmExtension(pi: ExtensionAPI) {
  pi.registerTool({
    name: "vers_rlm",
    label: "Vers RLM",
    description: `Run a Recursive Language Model with a full Linux VM as its sandbox. The sub-agent writes JavaScript that executes bash commands, reads/writes files, and manages VM state (snapshot, branch, restore) on a Vers Firecracker microVM.

Use this for tasks that need:
- Iterative exploration of systems, codebases, or data on a real machine
- Installing and running arbitrary tools (compilers, databases, etc.)
- Speculative execution via VM branching (try multiple approaches cheaply)
- Any task where a sandboxed JS REPL isn't enough`,
    parameters: Type.Object({
      signature: Type.String({
        description:
          'Defines inputs and outputs, e.g. "task -> findings, recommendations"',
      }),
      inputs: Type.Record(Type.String(), Type.Unknown(), {
        description: "Input data as key-value pairs matching the signature",
      }),
      instructions: Type.Optional(
        Type.String({
          description:
            "Additional instructions for the sub-agent beyond the signature",
        })
      ),
      maxIterations: Type.Optional(
        Type.Number({
          description:
            "Max REPL iterations (default: 20)",
        })
      ),
      maxLLMCalls: Type.Optional(
        Type.Number({
          description:
            "Max sub-LLM calls (default: 30)",
        })
      ),
      vmId: Type.Optional(
        Type.String({
          description:
            "Existing VM ID to use. If not provided, creates a fresh VM.",
        })
      ),
      vmConfig: Type.Optional(
        Type.Object({
          vcpu_count: Type.Optional(Type.Number()),
          mem_size_mib: Type.Optional(Type.Number()),
          fs_size_mib: Type.Optional(Type.Number()),
        }),
      ),
      keepVM: Type.Optional(
        Type.Boolean({
          description:
            "Keep the VM alive after completion (default: false). Useful for follow-up tasks.",
        })
      ),
    }),
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      const {
        signature,
        inputs,
        instructions,
        maxIterations = 20,
        maxLLMCalls = 30,
        vmId: existingVmId,
        vmConfig,
        keepVM = false,
      } = params as any;

      const vers = new VersAPI();
      const createdVMs: string[] = [];

      // Track the active VM via a ref so tools can swap it
      const vmIdRef = { current: "" };

      const update = (text: string) => {
        onUpdate?.({
          content: [{ type: "text", text }],
          details: {},
        });
      };

      try {
        // 1. Get or create a VM
        if (existingVmId) {
          vmIdRef.current = existingVmId;
          update(`Using existing VM: ${existingVmId}`);
        } else {
          update("🚀 Creating Vers VM...");
          const config = vmConfig || { vcpu_count: 2, mem_size_mib: 2048, fs_size_mib: 8192 };
          vmIdRef.current = await vers.createVM(config);
          createdVMs.push(vmIdRef.current);
          update(`🚀 VM created: ${vmIdRef.current}`);
        }

        // 2. Verify connectivity
        update("🔌 Verifying SSH connectivity...");
        const whoami = await vers.exec(vmIdRef.current, "whoami && hostname");
        update(`🔌 Connected: ${whoami.stdout.trim()}`);

        // 3. Build tools
        const vmTools = createVMTools(vers, vmIdRef, update);

        // 4. Build signature with VM instructions
        const { Signature } = await import("ts-rlm");
        const fullInstructions = [VM_TOOL_INSTRUCTIONS, instructions]
          .filter(Boolean)
          .join("\n\n");
        const sig = Signature.parse(signature, fullInstructions);

        // 5. Build LM adapter
        const lm = new PiLanguageModel(ctx);

        // 6. Create and run RLM
        const rlm = new RLM({
          signature: sig,
          maxIterations,
          maxLLMCalls,
          subLM: lm,
          tools: vmTools,
          verbose: false,
          onIteration(event) {
            const outputPreview =
              event.output.length > 500
                ? event.output.slice(0, 500) + "…"
                : event.output;
            const status = event.done ? "✓" : "⟳";
            const text = [
              `${status} Step ${event.iteration + 1}/${event.maxIterations}`,
              event.reasoning ? `  💭 ${event.reasoning}` : "",
              `  📝 ${event.code.slice(0, 300)}${event.code.length > 300 ? "…" : ""}`,
              `  📤 ${outputPreview}`,
            ]
              .filter(Boolean)
              .join("\n");

            update(text);
          },
        });

        update("🧠 Starting RLM exploration...");
        const prediction = await rlm.forward(inputs);

        // 7. Format results
        const trajectoryLines = prediction.trajectory.map(
          (step: any, i: number) => {
            const reasoning = step.reasoning
              ? `Reasoning: ${step.reasoning}\n`
              : "";
            const outputPreview =
              step.output.length > 1000
                ? step.output.slice(0, 1000) + "..."
                : step.output;
            return `=== Step ${i + 1} ===\n${reasoning}Code:\n\`\`\`js\n${step.code}\n\`\`\`\nOutput:\n${outputPreview}`;
          }
        );

        const { trajectory, finalReasoning, ...outputs } = prediction;

        const resultText = [
          `✅ Vers RLM completed in ${prediction.trajectory.length} steps.`,
          `VM: ${vmIdRef.current}${keepVM ? " (kept alive)" : " (will be deleted)"}`,
          "",
          "## Outputs",
          ...Object.entries(outputs).map(
            ([k, v]) =>
              `**${k}**: ${typeof v === "string" ? v : JSON.stringify(v, null, 2)}`
          ),
          "",
          "## Trajectory",
          ...trajectoryLines,
        ].join("\n");

        return {
          content: [{ type: "text", text: resultText }],
          details: {
            outputs,
            steps: prediction.trajectory.length,
            vmId: vmIdRef.current,
            finalReasoning: prediction.finalReasoning,
          },
        };
      } finally {
        // Cleanup: delete VMs we created (unless keepVM)
        if (!keepVM) {
          for (const id of createdVMs) {
            try {
              await vers.deleteVM(id);
              update(`🗑️ Deleted VM: ${id}`);
            } catch (e: any) {
              update(`⚠️ Failed to delete VM ${id}: ${e.message}`);
            }
          }
        }
      }
    },
  });

  pi.on("session_start", async (_event, ctx) => {
    if (ctx.hasUI) {
      ctx.ui.setStatus("vers-rlm", "vers-rlm: ready");
    }
  });
}