fix: unnecessary restore prompt, wrong checkpoint on tree nav, and leaf position persistence

src/commands.ts

@@ -211,18 +211,17 @@ export async function handleForkRestore(
   const entry = ctx.sessionManager.getEntry(event.entryId);
   const targetTs = entry?.timestamp ? new Date(entry.timestamp).getTime() : Date.now();
 
-  // Find best checkpoint
+  // Find best checkpoint and determine if files have changed since that point
   const sorted = [...state.checkpoints.values()].sort((a, b) => b.timestamp - a.timestamp);
-  const target = sorted.find((cp) => cp.timestamp <= targetTs) ?? sorted[sorted.length - 1];
+  const targetCp = sorted.find((cp) => cp.timestamp <= targetTs) ?? (sorted.length > 0 ? sorted[sorted.length - 1] : undefined);
+  const latestCp = sorted[0];
 
-  if (!target && state.resumeCheckpoint) {
-    // Use resume checkpoint as fallback
-  }
-
-  const cp = target || state.resumeCheckpoint;
+  // Offer file restore only when target differs from latest (changes happened after target)
+  const hasFileChanges = latestCp && targetCp && latestCp.id !== targetCp.id;
+  const cp = hasFileChanges ? targetCp : state.resumeCheckpoint;
 
   const options: string[] = ["Conversation only (keep files)"];
-  if (cp) {
+  if (hasFileChanges) {
     options.push("Restore all (files + conversation)");
     options.push("Code only (restore files, keep conversation)");
   }
@@ -260,20 +259,28 @@ export async function handleForkRestore(
 
 export async function handleTreeRestore(
   state: RewindState,
-  event: { preparation: { targetId: string } },
+  event: { preparation: { targetId: string; oldLeafId: string | null } },
   ctx: any,
 ): Promise<{ cancel: true } | undefined> {
   if (!state.gitAvailable || !state.repoRoot || !state.sessionId) return undefined;
   if (!ctx.hasUI) return undefined;
 
-  const entry = ctx.sessionManager.getEntry(event.preparation.targetId);
-  const targetTs = entry?.timestamp ? new Date(entry.timestamp).getTime() : Date.now();
+  const targetEntry = ctx.sessionManager.getEntry(event.preparation.targetId);
+  const targetTs = targetEntry?.timestamp ? new Date(targetEntry.timestamp).getTime() : Date.now();
 
   const sorted = [...state.checkpoints.values()].sort((a, b) => b.timestamp - a.timestamp);
-  const cp = sorted.find((c) => c.timestamp <= targetTs) ?? state.resumeCheckpoint;
+  const targetCp = sorted.find((c) => c.timestamp <= targetTs) ?? state.resumeCheckpoint;
+
+  // Only show restore when navigation crosses a checkpoint boundary (old leaf vs target)
+  const oldLeafEntry = event.preparation.oldLeafId
+    ? ctx.sessionManager.getEntry(event.preparation.oldLeafId)
+    : undefined;
+  const oldLeafTs = oldLeafEntry?.timestamp ? new Date(oldLeafEntry.timestamp).getTime() : Date.now();
+  const oldLeafCp = sorted.find((c) => c.timestamp <= oldLeafTs) ?? state.resumeCheckpoint;
+  const hasFileChanges = oldLeafCp && targetCp && oldLeafCp.id !== targetCp.id;
 
   const options: string[] = ["Keep current files"];
-  if (cp) options.push("Restore files to that point");
+  if (hasFileChanges) options.push("Restore files to that point");
   if (state.redoStack.length > 0) options.push("↩ Undo last rewind");
   options.push("Cancel navigation");
 
@@ -289,8 +296,8 @@ export async function handleTreeRestore(
     return { cancel: true };
   }
 
-  if (cp) {
-    await performRestore(state, ctx, cp, "files");
+  if (targetCp) {
+    await performRestore(state, ctx, targetCp, "files");
     ctx.ui.notify("Files restored to checkpoint", "info");
   }
 

src/index.ts

@@ -30,6 +30,38 @@ import {
 import { createInitialState, resetState } from "./state.js";
 import { updateStatus, clearStatus } from "./ui.js";
 import { registerCommands, handleForkRestore, handleTreeRestore } from "./commands.js";
+import { existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
+import * as os from "os";
+import * as path from "path";
+
+function repoKey(repoRoot) {
+  return repoRoot.replace(/[^a-zA-Z0-9_-]/g, "-");
+}
+
+function leafFilePath(repoRoot) {
+  return path.join(os.homedir(), ".pi", "agent", "pi-rewind", `leaf-${repoKey(repoRoot)}.json`);
+}
+
+function saveLeafId(repoRoot, leafId) {
+  const file = leafFilePath(repoRoot);
+  const dir = path.dirname(file);
+  if (!existsSync(dir)) mkdirSync(dir, { recursive: true });
+  writeFileSync(file, JSON.stringify({ leafId }));
+}
+
+function loadLeafId(repoRoot) {
+  const file = leafFilePath(repoRoot);
+  if (!existsSync(file)) return null;
+  try {
+    const data = JSON.parse(readFileSync(file, "utf-8"));
+    return data?.leafId || null;
+  } catch { return null; }
+}
+
+function clearLeafId(repoRoot) {
+  const file = leafFilePath(repoRoot);
+  try { writeFileSync(file, JSON.stringify({ leafId: null })); } catch {}
+}
 
 /** Truncate a string to maxLen, adding ellipsis if needed */
 function truncate(s: string, maxLen: number): string {
@@ -104,6 +136,14 @@ export default function (pi: ExtensionAPI) {
 
     if (ctx.hasUI) updateStatus(state, ctx);
 
+    // Restore leaf position from saved navigation point
+    if (state.repoRoot) {
+      const savedLeaf = loadLeafId(state.repoRoot);
+      if (savedLeaf && ctx.sessionManager.getEntry(savedLeaf)) {
+        ctx.sessionManager.branch(savedLeaf);
+      }
+    }
+
     // Prune old sessions in background (non-blocking)
     if (state.repoRoot && state.sessionId) {
       const root = state.repoRoot;