Agent-JAE/default-extensions/checkpoints.ts

import type { ExtensionAPI } from "@jaeswift/jae-coding-agent";

export default function (pi: ExtensionAPI) {
  let checkpointCount = 0;

  async function createCheckpoint(ctx: any, label: string): Promise<boolean> {
    try {
      // Check if in git repo
      const statusResult = await pi.exec("git", ["status", "--porcelain"], { cwd: ctx.cwd });
      if (statusResult.code !== 0) return false;

      // Only checkpoint if there are changes
      if (!statusResult.stdout.trim()) return false;

      // Stage all and create stash
      await pi.exec("git", ["add", "-A"], { cwd: ctx.cwd });
      checkpointCount++;
      const tag = `jae-checkpoint-${checkpointCount}`;
      const result = await pi.exec(
        "git",
        ["stash", "push", "-m", `${tag}: ${label}`, "--include-untracked"],
        { cwd: ctx.cwd },
      );

      if (result.code === 0 && !result.stdout.includes("No local changes")) {
        // Immediately pop to restore working state - the stash entry remains in reflog
        await pi.exec("git", ["stash", "pop"], { cwd: ctx.cwd });
        return true;
      }
      return false;
    } catch {
      return false;
    }
  }

  // Auto-checkpoint on file mutations
  pi.on("tool_execution_start", async (event, ctx) => {
    const mutatingTools = ["write", "edit"];
    if (mutatingTools.includes(event.toolName)) {
      await createCheckpoint(ctx, `before ${event.toolName}`);
    }
  });

  // Also checkpoint on bash commands that might modify files
  pi.on("tool_call", async (event, ctx) => {
    if (event.toolName === "bash") {
      const cmd = (event.input as any)?.command || "";
      const dangerousPatterns = ["rm ", "mv ", "cp ", "sed ", "chmod ", "chown ", "truncate", "dd "];
      if (dangerousPatterns.some((p) => cmd.includes(p))) {
        await createCheckpoint(ctx, `before bash: ${cmd.substring(0, 50)}`);
      }
    }
  });

  pi.registerCommand("undo", {
    description: "Undo last file change by restoring from git stash checkpoint",
    handler: async (_args, ctx) => {
      try {
        // List stashes to find jae checkpoints
        const listResult = await pi.exec("git", ["stash", "list"], { cwd: ctx.cwd });
        if (listResult.code !== 0) {
          ctx.ui.notify("Not in a git repository.", "error");
          return;
        }

        const stashes = listResult.stdout.trim().split("\n").filter(Boolean);
        const jaeStashes = stashes.filter((s) => s.includes("jae-checkpoint-"));

        if (jaeStashes.length === 0) {
          ctx.ui.notify("No JAE checkpoints available to undo.", "warning");
          return;
        }

        // Show what we're undoing
        const latest = jaeStashes[0];
        const confirmed = await ctx.ui.confirm(
          "Undo Checkpoint",
          `Restore to: ${latest}\n\nThis will discard current changes and restore the checkpoint.`,
        );

        if (!confirmed) {
          ctx.ui.notify("Undo cancelled.", "info");
          return;
        }

        // Hard reset and apply stash
        const stashRef = latest.split(":")[0]; // e.g., "stash@{0}"
        await pi.exec("git", ["checkout", "."], { cwd: ctx.cwd });
        await pi.exec("git", ["clean", "-fd"], { cwd: ctx.cwd });
        const applyResult = await pi.exec("git", ["stash", "apply", stashRef], { cwd: ctx.cwd });

        if (applyResult.code === 0) {
          ctx.ui.notify(`\u23EA Restored checkpoint: ${latest}`, "info");
        } else {
          ctx.ui.notify(`Failed to apply stash: ${applyResult.stderr}`, "error");
        }
      } catch (err) {
        ctx.ui.notify(`Undo error: ${err}`, "error");
      }
    },
  });

  pi.registerCommand("checkpoints", {
    description: "List available JAE checkpoints",
    handler: async (_args, ctx) => {
      const listResult = await pi.exec("git", ["stash", "list"], { cwd: ctx.cwd });
      if (listResult.code !== 0) {
        ctx.ui.notify("Not in a git repository.", "error");
        return;
      }

      const stashes = listResult.stdout.trim().split("\n").filter(Boolean);
      const jaeStashes = stashes.filter((s) => s.includes("jae-checkpoint-"));

      if (jaeStashes.length === 0) {
        ctx.ui.notify("No JAE checkpoints found.", "info");
        return;
      }

      ctx.ui.notify(
        `\u{1F4CC} JAE Checkpoints (${jaeStashes.length}):\n${jaeStashes.map((s) => `  ${s}`).join("\n")}`,
        "info",
      );
    },
  });
}