pi-extensions/extensions/git-checkpoint.ts

/**
 * Git Checkpoint Extension
 *
 * Creates a git stash checkpoint at the start of each turn, keyed to the
 * session entry ID. If you /fork from a past entry, you're offered the option
 * to restore the code state from that point.
 *
 * Silently skips turns where the working directory isn't inside a git repo or
 * where there are no changes to stash.
 *
 * Status bar shows the number of checkpoints saved in the current session.
 */

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  // entryId → stash ref (e.g. "refs/stash" or a full sha)
  const checkpoints = new Map<string, string>();

  async function inGitRepo(): Promise<boolean> {
    try {
      const result = await pi.exec("git", ["rev-parse", "--git-dir"]);
      return result.exitCode === 0;
    } catch {
      return false;
    }
  }

  function updateStatus(ctx: { ui: { setStatus: (id: string, text: string) => void } }): void {
    const n = checkpoints.size;
    if (n > 0) {
      ctx.ui.setStatus("git-checkpoint", `⎇ ${n} checkpoint${n === 1 ? "" : "s"}`);
    } else {
      ctx.ui.setStatus("git-checkpoint", "");
    }
  }

  pi.on("turn_start", async (_event, ctx) => {
    if (!(await inGitRepo())) return;

    // Capture the current entry ID at turn start — this is the entry the
    // user will fork from if they want to restore code to this point.
    const entryId = ctx.sessionManager.getLeafId();
    if (!entryId) return;

    // git stash create makes a stash object without touching the working tree.
    // Returns the stash ref on stdout, or empty string if nothing to stash.
    const result = await pi.exec("git", ["stash", "create"]);
    const ref = result.stdout.trim();

    if (ref) {
      checkpoints.set(entryId, ref);
      updateStatus(ctx);
    }
  });

  pi.on("session_before_fork", async (event, ctx) => {
    const ref = checkpoints.get(event.entryId);
    if (!ref) return;
    if (!ctx.hasUI) return;

    const choice = await ctx.ui.select(
      "Restore code to this checkpoint?",
      ["Yes, restore code state", "No, keep current code"],
    );

    if (choice?.startsWith("Yes")) {
      await pi.exec("git", ["stash", "apply", ref]);
      ctx.ui.notify("Code restored to checkpoint", "success");
    }
  });

  pi.on("session_shutdown", () => {
    checkpoints.clear();
  });
}