pi-rewind-hook 1.0.0 → 1.1.1

package/CHANGELOG.md

@@ -0,0 +1,36 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+## [1.1.1] - 2024-12-27
+
+### Fixed
+- Use `before_branch` event instead of `branch` for proper hook timing (thanks @badlogic)
+- Cancel branch when user dismisses restore options menu
+
+## [1.1.0] - 2024-12-27
+
+### Added
+- "Undo last file rewind" option - restore files to state before last rewind
+- Checkpoints now capture uncommitted and untracked files (not just HEAD)
+- Git repo detection - hook gracefully skips in non-git directories
+
+### Changed
+- Checkpoints use `git write-tree` with temp index to capture working directory state
+- Pruning excludes before-restore ref and current session's resume checkpoint
+
+### Fixed
+- Code-only restore options now properly skip conversation restore
+
+## [1.0.0] - 2024-12-19
+
+### Added
+- Initial release
+- Automatic checkpoints at session start and each turn
+- `/branch` integration with restore options:
+  - Restore all (files + conversation)
+  - Conversation only (keep current files)
+  - Code only (restore files, keep conversation)
+- Resume checkpoint for pre-session messages
+- Automatic pruning (keeps last 100 checkpoints)
+- Cross-platform installation via `npx pi-rewind-hook`

package/README.md

@@ -2,6 +2,12 @@
 
 A Pi agent hook that enables rewinding file changes during coding sessions. Creates automatic checkpoints using git refs, allowing you to restore files to previous states while optionally preserving conversation history.
 
+## Screenshots
+
+![Selecting a message to branch from](rewind1.png)
+
+![Choosing a restore option](rewind2.png)
+
 ## Requirements
 
 - Pi agent v0.18.0+

package/index.ts

@@ -1,28 +1,139 @@
 import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
+import { exec as execCb } from "child_process";
+import { mkdtemp, rm } from "fs/promises";
+import { tmpdir } from "os";
+import { join } from "path";
+import { promisify } from "util";
+
+const execAsync = promisify(execCb);
 
 const REF_PREFIX = "refs/pi-checkpoints/";
+const BEFORE_RESTORE_PREFIX = "before-restore-";
 const MAX_CHECKPOINTS = 100;
 
+type ExecFn = (cmd: string, args: string[]) => Promise<{ stdout: string; stderr: string; code: number }>;
+
 export default function (pi: HookAPI) {
   const checkpoints = new Map<number, string>();
   let resumeCheckpoint: string | null = null;
+  let repoRoot: string | null = null;
+  let isGitRepo = false;
 
   console.error(`[rewind] Hook loaded`);
 
-  pi.on("session_start", async (event, ctx) => {
-    if (!ctx.hasUI) return;
+  async function findBeforeRestoreRef(exec: ExecFn): Promise<{ refName: string; commitSha: string } | null> {
+    try {
+      const result = await exec("git", [
+        "for-each-ref",
+        "--sort=-creatordate",
+        "--count=1",
+        "--format=%(refname) %(objectname)",
+        `${REF_PREFIX}${BEFORE_RESTORE_PREFIX}*`,
+      ]);
 
-    const checkpointId = `checkpoint-resume-${Date.now()}`;
+      const line = result.stdout.trim();
+      if (!line) return null;
+
+      const [refName, commitSha] = line.split(" ");
+      return { refName, commitSha };
+    } catch {
+      return null;
+    }
+  }
+
+  async function getRepoRoot(exec: ExecFn): Promise<string> {
+    if (repoRoot) return repoRoot;
+    const result = await exec("git", ["rev-parse", "--show-toplevel"]);
+    repoRoot = result.stdout.trim();
+    return repoRoot;
+  }
+
+  async function captureWorktree(exec: ExecFn): Promise<string> {
+    const root = await getRepoRoot(exec);
+    const tmpDir = await mkdtemp(join(tmpdir(), "pi-rewind-"));
+    const tmpIndex = join(tmpDir, "index");
+
+    try {
+      const env = { ...process.env, GIT_INDEX_FILE: tmpIndex };
+      await execAsync("git add -A", { cwd: root, env });
+      const { stdout: treeSha } = await execAsync("git write-tree", { cwd: root, env });
+
+      const { stdout: commitSha } = await execAsync(
+        `git commit-tree ${treeSha.trim()} -m "rewind backup"`,
+        { cwd: root }
+      );
+      return commitSha.trim();
+    } finally {
+      await rm(tmpDir, { recursive: true, force: true }).catch(() => {});
+    }
+  }
 
+  async function restoreWithBackup(
+    exec: ExecFn,
+    targetRef: string,
+    notify: (msg: string, level: "success" | "error" | "info") => void
+  ): Promise<boolean> {
     try {
-      await ctx.exec("git", [
+      const existingBackup = await findBeforeRestoreRef(exec);
+
+      const backupCommit = await captureWorktree(exec);
+      const newBackupId = `${BEFORE_RESTORE_PREFIX}${Date.now()}`;
+      await exec("git", [
+        "update-ref",
+        `${REF_PREFIX}${newBackupId}`,
+        backupCommit,
+      ]);
+      console.error(`[rewind] Created backup: ${newBackupId}`);
+
+      if (existingBackup) {
+        await exec("git", ["update-ref", "-d", existingBackup.refName]);
+        console.error(`[rewind] Deleted old backup: ${existingBackup.refName}`);
+      }
+
+      await exec("git", ["checkout", targetRef, "--", "."]);
+      return true;
+    } catch (err) {
+      console.error(`[rewind] Failed to restore: ${err}`);
+      notify(`Failed to restore files: ${err}`, "error");
+      return false;
+    }
+  }
+
+  async function createCheckpointFromWorktree(exec: ExecFn, checkpointId: string): Promise<boolean> {
+    try {
+      const commitSha = await captureWorktree(exec);
+      await exec("git", [
         "update-ref",
         `${REF_PREFIX}${checkpointId}`,
-        "HEAD",
+        commitSha,
       ]);
+      return true;
+    } catch {
+      return false;
+    }
+  }
+
+  pi.on("session", async (event, ctx) => {
+    if (event.reason !== "start") return;
+    if (!ctx.hasUI) return;
+
+    try {
+      const result = await ctx.exec("git", ["rev-parse", "--is-inside-work-tree"]);
+      isGitRepo = result.stdout.trim() === "true";
+    } catch {
+      isGitRepo = false;
+    }
+
+    if (!isGitRepo) return;
 
-      resumeCheckpoint = checkpointId;
-      console.error(`[rewind] Created resume checkpoint: ${checkpointId}`);
+    const checkpointId = `checkpoint-resume-${Date.now()}`;
+
+    try {
+      const success = await createCheckpointFromWorktree(ctx.exec, checkpointId);
+      if (success) {
+        resumeCheckpoint = checkpointId;
+        console.error(`[rewind] Created resume checkpoint: ${checkpointId}`);
+      }
     } catch (err) {
       console.error(`[rewind] Failed to create resume checkpoint: ${err}`);
     }
@@ -30,29 +141,28 @@ export default function (pi: HookAPI) {
 
   pi.on("turn_start", async (event, ctx) => {
     if (!ctx.hasUI) return;
+    if (!isGitRepo) return;
 
     const checkpointId = `checkpoint-${event.timestamp}`;
 
     try {
-      await ctx.exec("git", [
-        "update-ref",
-        `${REF_PREFIX}${checkpointId}`,
-        "HEAD",
-      ]);
-
-      checkpoints.set(event.turnIndex, checkpointId);
-      console.error(
-        `[rewind] Created checkpoint ${checkpointId} for turn ${event.turnIndex}`
-      );
-
-      await pruneCheckpoints(ctx);
+      const success = await createCheckpointFromWorktree(ctx.exec, checkpointId);
+      if (success) {
+        checkpoints.set(event.turnIndex, checkpointId);
+        console.error(
+          `[rewind] Created checkpoint ${checkpointId} for turn ${event.turnIndex}`
+        );
+        await pruneCheckpoints(ctx.exec);
+      }
     } catch (err) {
       console.error(`[rewind] Failed to create checkpoint: ${err}`);
     }
   });
 
-  pi.on("branch", async (event, ctx) => {
+  pi.on("session", async (event, ctx) => {
+    if (event.reason !== "before_branch") return;
     if (!ctx.hasUI) return;
+    if (!isGitRepo) return;
 
     let checkpointId = checkpoints.get(event.targetTurnIndex);
     let usingResumeCheckpoint = false;
@@ -62,63 +172,84 @@ export default function (pi: HookAPI) {
       usingResumeCheckpoint = true;
     }
 
-    if (!checkpointId) {
+    const beforeRestoreRef = await findBeforeRestoreRef(ctx.exec);
+    const hasUndo = !!beforeRestoreRef;
+
+    if (!checkpointId && !hasUndo) {
       ctx.ui.notify(
         "No checkpoint available for this message",
         "info"
       );
-      return undefined;
+      return;
     }
 
-    const options = usingResumeCheckpoint
-      ? [
-          "Restore to session start (files + conversation)",
-          "Conversation only (keep current files)",
-          "Restore to session start (files only, keep conversation)",
-        ]
-      : [
-          "Restore all (files + conversation)",
-          "Conversation only (keep current files)",
-          "Code only (restore files, keep conversation)",
-        ];
+    const options: string[] = [];
+    
+    if (checkpointId) {
+      if (usingResumeCheckpoint) {
+        options.push("Restore to session start (files + conversation)");
+        options.push("Conversation only (keep current files)");
+        options.push("Restore to session start (files only, keep conversation)");
+      } else {
+        options.push("Restore all (files + conversation)");
+        options.push("Conversation only (keep current files)");
+        options.push("Code only (restore files, keep conversation)");
+      }
+    }
+
+    if (hasUndo) {
+      options.push("Undo last file rewind");
+    }
 
     const choice = await ctx.ui.select("Restore Options", options);
 
     if (!choice) {
       ctx.ui.notify("Rewind cancelled", "info");
-      return { skipConversationRestore: true };
+      return { cancel: true };
     }
 
     if (choice.startsWith("Conversation only")) {
-      return undefined;
+      return;
     }
 
-    try {
-      const ref = `${REF_PREFIX}${checkpointId}`;
-      await ctx.exec("git", ["checkout", ref, "--", "."]);
-      console.error(`[rewind] Restored files from ${checkpointId}`);
+    const isCodeOnly = choice === "Code only (restore files, keep conversation)" ||
+      choice === "Restore to session start (files only, keep conversation)";
+
+    if (choice === "Undo last file rewind") {
+      const success = await restoreWithBackup(
+        ctx.exec,
+        beforeRestoreRef!.commitSha,
+        ctx.ui.notify.bind(ctx.ui)
+      );
+      if (success) {
+        ctx.ui.notify("Files restored to before last rewind", "success");
+      }
+      return { skipConversationRestore: true };
+    }
+
+    const ref = `${REF_PREFIX}${checkpointId}`;
+    const success = await restoreWithBackup(
+      ctx.exec,
+      ref,
+      ctx.ui.notify.bind(ctx.ui)
+    );
+    if (success) {
       ctx.ui.notify(
         usingResumeCheckpoint 
           ? "Files restored to session start" 
           : "Files restored from checkpoint",
         "success"
       );
-    } catch (err) {
-      console.error(`[rewind] Failed to restore: ${err}`);
-      ctx.ui.notify(`Failed to restore files: ${err}`, "error");
-      return { skipConversationRestore: true };
     }
 
-    if (choice.includes("files only") || choice.startsWith("Code only")) {
+    if (isCodeOnly) {
       return { skipConversationRestore: true };
     }
-
-    return undefined;
   });
 
-  async function pruneCheckpoints(ctx: { exec: (cmd: string, args: string[]) => Promise<{ stdout: string; stderr: string; code: number }> }) {
+  async function pruneCheckpoints(exec: ExecFn) {
     try {
-      const result = await ctx.exec("git", [
+      const result = await exec("git", [
         "for-each-ref",
         "--sort=creatordate",
         "--format=%(refname)",
@@ -126,11 +257,15 @@ export default function (pi: HookAPI) {
       ]);
 
       const refs = result.stdout.trim().split("\n").filter(Boolean);
+      const currentResumeRef = resumeCheckpoint ? `${REF_PREFIX}${resumeCheckpoint}` : null;
+      const checkpointRefs = refs.filter(r => 
+        !r.includes(BEFORE_RESTORE_PREFIX) && r !== currentResumeRef
+      );
 
-      if (refs.length > MAX_CHECKPOINTS) {
-        const toDelete = refs.slice(0, refs.length - MAX_CHECKPOINTS);
+      if (checkpointRefs.length > MAX_CHECKPOINTS) {
+        const toDelete = checkpointRefs.slice(0, checkpointRefs.length - MAX_CHECKPOINTS);
         for (const ref of toDelete) {
-          await ctx.exec("git", ["update-ref", "-d", ref]);
+          await exec("git", ["update-ref", "-d", ref]);
           console.error(`[rewind] Pruned old checkpoint: ${ref}`);
         }
       }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-rewind-hook",
-  "version": "1.0.0",
+  "version": "1.1.1",
   "description": "Rewind hook for Pi agent - automatic git checkpoints with file/conversation restore",
   "bin": {
     "pi-rewind-hook": "./install.js"