1  /**
   2   * Rewind Extension - session-ledger based exact file restoration for pi branching
   3   *
   4   * Rewind v2 stores exact rewind metadata in hidden session custom entries and keeps
   5   * snapshot commits reachable through a single repo-local store ref.
   6   */
   7  
   8  import { getAgentDir, type ExtensionAPI, type ExtensionContext } from "@mariozechner/pi-coding-agent";
   9  import { exec as execCb } from "child_process";
  10  import { existsSync, readFileSync, realpathSync } from "fs";
  11  import { mkdtemp, readdir, readFile, rm, stat } from "fs/promises";
  12  import { tmpdir } from "os";
  13  import { dirname, isAbsolute, join, relative, resolve } from "path";
  14  import { promisify } from "util";
  15  
  16  const execAsync = promisify(execCb);
  17  
  18  const LEGACY_REF_PREFIX = "refs/pi-checkpoints/";
  19  const STORE_REF = "refs/pi-rewind/store";
  20  const STATUS_KEY = "rewind";
  21  const FORK_PREFERENCE_SOURCE_ALLOWLIST = new Set(["fork-from-first"]);
  22  const LEGACY_ZERO_SHA = "0000000000000000000000000000000000000000";
  23  const RETENTION_SWEEP_THRESHOLD = 50;
  24  const RETENTION_VERSION = 2;
  25  const EMPTY_TREE_SHA = "4b825dc642cb6eb9a060e54bf8d69288fbee4904";
  26  
  27  type ExecFn = (cmd: string, args: string[]) => Promise<{ stdout: string; stderr: string; code: number }>;
  28  
  29  type GitExecResult = Awaited<ReturnType<ExecFn>>;
  30  
  31  type BindingTuple = [entryId: string, snapshotIndex: number];
  32  
  33  interface RewindSettings {
  34    rewind?: {
  35      silentCheckpoints?: boolean;
  36      retention?: {
  37        maxSnapshots?: number;
  38        maxAgeDays?: number;
  39        pinLabeledEntries?: boolean;
  40      };
  41    };
  42  }
  43  
  44  interface RewindTurnData {
  45    v: 2;
  46    snapshots: string[];
  47    bindings: BindingTuple[];
  48  }
  49  
  50  interface RewindOpData {
  51    v: 2;
  52    snapshots: string[];
  53    bindings?: BindingTuple[];
  54    current?: number;
  55    undo?: number;
  56  }
  57  
  58  interface ActivePromptCollector {
  59    snapshots: string[];
  60    bindings: BindingTuple[];
  61    promptText?: string;
  62    pendingUserCommitSha?: string;
  63  }
  64  
  65  interface ExactState {
  66    commitSha: string;
  67    treeSha: string;
  68  }
  69  
  70  interface ActiveBranchState {
  71    currentCommitSha?: string;
  72    currentTreeSha?: string;
  73    undoCommitSha?: string;
  74  }
  75  
  76  interface PendingResultingState {
  77    currentCommitSha: string;
  78    undoCommitSha?: string;
  79  }
  80  
  81  interface ParsedLedgerReference {
  82    commitSha: string;
  83    entryId?: string;
  84    timestamp: number;
  85    kind: "binding" | "current" | "undo";
  86  }
  87  
  88  interface ParsedSessionLedger {
  89    sessionFile: string;
  90    sessionId?: string;
  91    cwd?: string;
  92    parentSession?: string;
  93    entryToCommit: Map<string, string>;
  94    labeledEntryIds: Set<string>;
  95    references: ParsedLedgerReference[];
  96    latestCurrentCommitSha?: string;
  97    latestUndoCommitSha?: string;
  98  }
  99  
 100  interface LegacyRef {
 101    refName: string;
 102    commitSha: string;
 103    sessionId?: string;
 104    entryId: string;
 105    scoped: boolean;
 106  }
 107  
 108  interface SessionLikeMessageEntry {
 109    type: "message";
 110    id: string;
 111    parentId: string | null;
 112    timestamp: string;
 113    message: {
 114      role: string;
 115      timestamp?: number;
 116      content?: unknown;
 117    };
 118  }
 119  
 120  interface SessionLikeCustomEntry {
 121    type: "custom";
 122    id: string;
 123    parentId: string | null;
 124    timestamp: string;
 125    customType: string;
 126    data?: unknown;
 127  }
 128  
 129  interface SessionLikeLabelEntry {
 130    type: "label";
 131    targetId: string;
 132    label?: string;
 133  }
 134  
 135  interface SessionLikeBranchSummaryEntry {
 136    type: "branch_summary";
 137    id: string;
 138  }
 139  
 140  interface SessionLikeGenericEntry {
 141    type: string;
 142    id?: string;
 143    parentId?: string | null;
 144    timestamp?: string;
 145    message?: {
 146      role?: string;
 147      timestamp?: number;
 148      content?: unknown;
 149    };
 150    customType?: string;
 151    data?: unknown;
 152    targetId?: string;
 153    label?: string;
 154  }
 155  
 156  type SessionLikeEntry =
 157    | SessionLikeMessageEntry
 158    | SessionLikeCustomEntry
 159    | SessionLikeLabelEntry
 160    | SessionLikeBranchSummaryEntry
 161    | SessionLikeGenericEntry;
 162  
 163  let cachedSettings: RewindSettings | null = null;
 164  
 165  function getSettingsFilePath(): string {
 166    return join(getAgentDir(), "settings.json");
 167  }
 168  
 169  function getDefaultSessionsDir(): string {
 170    return join(getAgentDir(), "sessions");
 171  }
 172  
 173  function getSettings(): RewindSettings {
 174    if (cachedSettings) {
 175      return cachedSettings;
 176    }
 177  
 178    try {
 179      cachedSettings = JSON.parse(readFileSync(getSettingsFilePath(), "utf-8")) as RewindSettings;
 180    } catch {
 181      cachedSettings = {};
 182    }
 183  
 184    return cachedSettings;
 185  }
 186  
 187  function getSilentCheckpointsSetting(): boolean {
 188    return getSettings().rewind?.silentCheckpoints === true;
 189  }
 190  
 191  function getRetentionSettings(): NonNullable<NonNullable<RewindSettings["rewind"]>["retention"]> | undefined {
 192    return getSettings().rewind?.retention;
 193  }
 194  
 195  function isRewindTurnData(value: unknown): value is RewindTurnData {
 196    if (!value || typeof value !== "object") return false;
 197    const data = value as Partial<RewindTurnData>;
 198    return data.v === 2 && Array.isArray(data.snapshots) && Array.isArray(data.bindings);
 199  }
 200  
 201  function isRewindOpData(value: unknown): value is RewindOpData {
 202    if (!value || typeof value !== "object") return false;
 203    const data = value as Partial<RewindOpData>;
 204    return data.v === 2 && Array.isArray(data.snapshots);
 205  }
 206  
 207  function canonicalizePath(value: string): string {
 208    const resolvedValue = resolve(value);
 209    try {
 210      return realpathSync.native(resolvedValue);
 211    } catch {
 212      return resolvedValue;
 213    }
 214  }
 215  
 216  function isInsidePath(targetPath: string, parentPath: string): boolean {
 217    const resolvedTarget = canonicalizePath(targetPath);
 218    const resolvedParent = canonicalizePath(parentPath);
 219    const rel = relative(resolvedParent, resolvedTarget);
 220    return rel === "" || (!rel.startsWith("..") && !isAbsolute(rel));
 221  }
 222  
 223  function toTimestamp(value: string | undefined): number {
 224    if (!value) return 0;
 225    const parsed = Date.parse(value);
 226    return Number.isFinite(parsed) ? parsed : 0;
 227  }
 228  
 229  function getTextContent(content: unknown): string {
 230    if (typeof content === "string") return content;
 231    if (!Array.isArray(content)) return "";
 232    return content
 233      .filter((block): block is { type: string; text?: string } => !!block && typeof block === "object")
 234      .filter((block) => block.type === "text")
 235      .map((block) => block.text ?? "")
 236      .join("\n");
 237  }
 238  
 239  function updateLabelSet(labelIds: Set<string>, entry: SessionLikeLabelEntry) {
 240    if (!entry.targetId) return;
 241    if (entry.label && entry.label.trim()) {
 242      labelIds.add(entry.targetId);
 243      return;
 244    }
 245    labelIds.delete(entry.targetId);
 246  }
 247  
 248  function applyBindings(target: Map<string, string>, snapshots: string[], bindings?: BindingTuple[]) {
 249    if (!bindings) return;
 250    for (const [entryId, snapshotIndex] of bindings) {
 251      const commitSha = snapshots[snapshotIndex];
 252      if (entryId && commitSha) {
 253        target.set(entryId, commitSha);
 254      }
 255    }
 256  }
 257  
 258  function addReferences(target: ParsedLedgerReference[], snapshots: string[], timestamp: number, data: RewindTurnData | RewindOpData) {
 259    if ("bindings" in data && data.bindings) {
 260      for (const [entryId, snapshotIndex] of data.bindings) {
 261        const commitSha = snapshots[snapshotIndex];
 262        if (!commitSha) continue;
 263        target.push({ commitSha, entryId, timestamp, kind: "binding" });
 264      }
 265    }
 266  
 267    if ("current" in data && typeof data.current === "number") {
 268      const commitSha = snapshots[data.current];
 269      if (commitSha) {
 270        target.push({ commitSha, timestamp, kind: "current" });
 271      }
 272    }
 273  
 274    if ("undo" in data && typeof data.undo === "number") {
 275      const commitSha = snapshots[data.undo];
 276      if (commitSha) {
 277        target.push({ commitSha, timestamp, kind: "undo" });
 278      }
 279    }
 280  }
 281  
 282  function resolveBindingSnapshotIndex(snapshots: string[], commitSha: string): number {
 283    const existingIndex = snapshots.indexOf(commitSha);
 284    if (existingIndex >= 0) return existingIndex;
 285    snapshots.push(commitSha);
 286    return snapshots.length - 1;
 287  }
 288  
 289  function addBindingToCollector(collector: ActivePromptCollector, entryId: string, commitSha: string) {
 290    const snapshotIndex = resolveBindingSnapshotIndex(collector.snapshots, commitSha);
 291    collector.bindings.push([entryId, snapshotIndex]);
 292  }
 293  
 294  function getCommitFromData(data: RewindOpData, indexKey: "current" | "undo"): string | undefined {
 295    const snapshotIndex = data[indexKey];
 296    return typeof snapshotIndex === "number" ? data.snapshots[snapshotIndex] : undefined;
 297  }
 298  
 299  function isRestorableTreeEntry(entry: SessionLikeEntry | undefined): boolean {
 300    if (!entry) return false;
 301    if (entry.type === "message") {
 302      return entry.message.role === "user" || entry.message.role === "assistant";
 303    }
 304    return entry.type === "branch_summary" || entry.type === "compaction";
 305  }
 306  
 307  function isAssistantMessageEntry(entry: SessionLikeEntry): entry is SessionLikeMessageEntry {
 308    return entry.type === "message" && entry.message.role === "assistant";
 309  }
 310  
 311  function findLatestUserMessageEntry(entries: SessionLikeEntry[]): SessionLikeMessageEntry | null {
 312    for (let index = entries.length - 1; index >= 0; index -= 1) {
 313      const entry = entries[index];
 314      if (entry.type === "message" && entry.message.role === "user") {
 315        return entry;
 316      }
 317    }
 318    return null;
 319  }
 320  
 321  function findLatestMatchingUserMessageEntry(
 322    entries: SessionLikeEntry[],
 323    promptText: string | null | undefined,
 324  ): SessionLikeMessageEntry | null {
 325    if (!promptText) return null;
 326  
 327    for (let index = entries.length - 1; index >= 0; index -= 1) {
 328      const entry = entries[index];
 329      if (entry.type !== "message" || entry.message.role !== "user") continue;
 330      if (getTextContent(entry.message.content) === promptText) {
 331        return entry;
 332      }
 333    }
 334  
 335    return null;
 336  }
 337  
 338  function findAssistantEntryForTurn(entries: SessionLikeEntry[], message: { timestamp?: number; content?: unknown }): SessionLikeMessageEntry | null {
 339    const targetTimestamp = message.timestamp;
 340    const targetText = getTextContent(message.content);
 341  
 342    for (let index = entries.length - 1; index >= 0; index -= 1) {
 343      const entry = entries[index];
 344      if (!isAssistantMessageEntry(entry)) continue;
 345  
 346      if (targetTimestamp !== undefined && entry.message.timestamp === targetTimestamp) {
 347        return entry;
 348      }
 349  
 350      if (targetText && getTextContent(entry.message.content) === targetText) {
 351        return entry;
 352      }
 353    }
 354  
 355    return null;
 356  }
 357  
 358  export default function rewindExtension(pi: ExtensionAPI) {
 359    const entryToCommit = new Map<string, string>();
 360    const parsedSessionCache = new Map<string, { mtimeMs: number; ledger: ParsedSessionLedger }>();
 361  
 362    let repoRoot: string | null = null;
 363    let sessionId: string | null = null;
 364    let currentSessionFile: string | undefined;
 365    let currentParentSession: string | undefined;
 366    let currentSessionCwd: string | undefined;
 367    let isGitRepo = false;
 368    let lastExact: ExactState | null = null;
 369    let activeBranchState: ActiveBranchState = {};
 370    let promptCollector: ActivePromptCollector | null = null;
 371    let pendingForkState: PendingResultingState | null = null;
 372    let pendingTreeState: PendingResultingState | null = null;
 373    let activePromptText: string | null = null;
 374      let newSnapshotsSinceSweep = 0;
 375      let sweepRunning = false;
 376      let sweepCompletedThisSession = false;
 377    let forceConversationOnlyOnNextFork = false;
 378    let forceConversationOnlySource: string | null = null;
 379  
 380    function notify(ctx: ExtensionContext, message: string, level: "info" | "warning" | "error" = "info") {
 381      if (!ctx.hasUI) return;
 382      if (level === "info" && getSilentCheckpointsSetting()) return;
 383      ctx.ui.notify(message, level);
 384    }
 385  
 386    function updateStatus(ctx: ExtensionContext) {
 387      if (!ctx.hasUI) return;
 388      if (!isGitRepo || getSilentCheckpointsSetting()) {
 389        ctx.ui.setStatus(STATUS_KEY, undefined);
 390        return;
 391      }
 392  
 393      const uniqueSnapshots = new Set(entryToCommit.values()).size;
 394      const theme = ctx.ui.theme;
 395      ctx.ui.setStatus(
 396        STATUS_KEY,
 397        theme.fg("dim", "◆ ") + theme.fg("muted", `${entryToCommit.size} points / ${uniqueSnapshots} snapshots`),
 398      );
 399    }
 400  
 401    function resetState() {
 402      entryToCommit.clear();
 403      parsedSessionCache.clear();
 404      repoRoot = null;
 405      sessionId = null;
 406      currentSessionFile = undefined;
 407      currentParentSession = undefined;
 408      currentSessionCwd = undefined;
 409      isGitRepo = false;
 410      lastExact = null;
 411      activeBranchState = {};
 412      promptCollector = null;
 413      pendingForkState = null;
 414      pendingTreeState = null;
 415        activePromptText = null;
 416        newSnapshotsSinceSweep = 0;
 417        sweepCompletedThisSession = false;
 418      forceConversationOnlyOnNextFork = false;
 419      forceConversationOnlySource = null;
 420      cachedSettings = null;
 421    }
 422  
 423    function syncSessionIdentity(ctx: ExtensionContext) {
 424      sessionId = ctx.sessionManager.getSessionId();
 425      currentSessionFile = ctx.sessionManager.getSessionFile();
 426      currentParentSession = ctx.sessionManager.getHeader()?.parentSession;
 427      currentSessionCwd = ctx.sessionManager.getCwd();
 428    }
 429  
 430    async function execGitChecked(args: string[]): Promise<GitExecResult> {
 431      const result = await pi.exec("git", args);
 432      if (result.code !== 0) {
 433        const stderr = result.stderr.trim();
 434        throw new Error(stderr || `git ${args.join(" ")} failed with code ${result.code}`);
 435      }
 436      return result;
 437    }
 438  
 439    async function getRepoRoot(exec: ExecFn): Promise<string> {
 440      if (repoRoot) return repoRoot;
 441      const result = await exec("git", ["rev-parse", "--show-toplevel"]);
 442      if (result.code !== 0) {
 443        const stderr = result.stderr.trim();
 444        throw new Error(stderr || `git rev-parse --show-toplevel failed with code ${result.code}`);
 445      }
 446      repoRoot = result.stdout.trim();
 447      return repoRoot;
 448    }
 449  
 450    async function captureWorktreeTree(): Promise<{ treeSha: string }> {
 451      const root = await getRepoRoot(pi.exec);
 452      const tempDir = await mkdtemp(join(tmpdir(), "pi-rewind-"));
 453      const tempIndex = join(tempDir, "index");
 454  
 455      try {
 456        const env = { ...process.env, GIT_INDEX_FILE: tempIndex };
 457        await execAsync("git add -A", { cwd: root, env });
 458        const { stdout } = await execAsync("git write-tree", { cwd: root, env });
 459        return { treeSha: stdout.trim() };
 460      } finally {
 461        await rm(tempDir, { recursive: true, force: true }).catch(() => undefined);
 462      }
 463    }
 464  
 465    async function getCommitTreeSha(commitSha: string): Promise<string> {
 466      const result = await execGitChecked(["show", "-s", "--format=%T", commitSha]);
 467      return result.stdout.trim();
 468    }
 469  
 470    async function commitExists(commitSha: string): Promise<boolean> {
 471      const result = await pi.exec("git", ["cat-file", "-e", `${commitSha}^{commit}`]);
 472      return result.code === 0;
 473    }
 474  
 475    async function getStoreHead(): Promise<string | undefined> {
 476      const result = await pi.exec("git", ["rev-parse", "--verify", STORE_REF]);
 477      if (result.code !== 0) {
 478        return undefined;
 479      }
 480      const value = result.stdout.trim();
 481      return value || undefined;
 482    }
 483  
 484    async function createStoreKeepaliveCommit(snapshotCommitSha: string, previousStoreHead?: string): Promise<string> {
 485      const args = ["commit-tree", EMPTY_TREE_SHA];
 486  
 487      if (previousStoreHead) {
 488        args.push("-p", previousStoreHead);
 489      }
 490  
 491      args.push("-p", snapshotCommitSha, "-m", "pi rewind store");
 492      const result = await execGitChecked(args);
 493      return result.stdout.trim();
 494    }
 495  
 496    async function appendSnapshotToStore(commitSha: string): Promise<void> {
 497      let attempts = 0;
 498  
 499      while (attempts < 5) {
 500        attempts += 1;
 501        const oldHead = await getStoreHead();
 502        const keepaliveCommit = await createStoreKeepaliveCommit(commitSha, oldHead);
 503  
 504        try {
 505          if (oldHead) {
 506            await execGitChecked(["update-ref", STORE_REF, keepaliveCommit, oldHead]);
 507          } else {
 508            await execGitChecked(["update-ref", STORE_REF, keepaliveCommit, LEGACY_ZERO_SHA]);
 509          }
 510          return;
 511        } catch {
 512          // Retry if another process updated the store ref concurrently
 513        }
 514      }
 515  
 516      throw new Error("failed to update rewind store ref");
 517    }
 518  
 519    async function rewriteStoreToLiveSet(liveCommitShas: string[]): Promise<"rewritten" | "preserved-empty"> {
 520      const uniqueLiveCommits = [...new Set(liveCommitShas.filter(Boolean))];
 521      if (uniqueLiveCommits.length === 0) {
 522        return "preserved-empty";
 523      }
 524  
 525      let head: string | undefined;
 526      for (const commitSha of uniqueLiveCommits) {
 527        head = await createStoreKeepaliveCommit(commitSha, head);
 528      }
 529  
 530      const oldHead = await getStoreHead();
 531      if (oldHead) {
 532        await execGitChecked(["update-ref", STORE_REF, head!, oldHead]);
 533        return "rewritten";
 534      }
 535  
 536      await execGitChecked(["update-ref", STORE_REF, head!, LEGACY_ZERO_SHA]);
 537      return "rewritten";
 538    }
 539  
 540    async function ensureSnapshotForTree(treeSha: string): Promise<string> {
 541      if (lastExact && lastExact.treeSha === treeSha) {
 542        return lastExact.commitSha;
 543      }
 544  
 545      const result = await execGitChecked(["commit-tree", treeSha, "-m", "pi rewind snapshot"]);
 546      const commitSha = result.stdout.trim();
 547      await appendSnapshotToStore(commitSha);
 548      lastExact = { commitSha, treeSha };
 549      newSnapshotsSinceSweep += 1;
 550      return commitSha;
 551    }
 552  
 553    async function ensureSnapshotForCurrentWorktree(): Promise<string> {
 554      const { treeSha } = await captureWorktreeTree();
 555      return ensureSnapshotForTree(treeSha);
 556    }
 557  
 558    async function deletePathsFromWorkingTree(paths: string[]) {
 559      if (paths.length === 0) return;
 560      const root = await getRepoRoot(pi.exec);
 561  
 562      for (const repoRelativePath of paths) {
 563        const absolutePath = resolve(root, repoRelativePath);
 564        if (!isInsidePath(absolutePath, root)) {
 565          throw new Error(`refusing to delete path outside repo root: ${repoRelativePath}`);
 566        }
 567        await rm(absolutePath, { recursive: true, force: true });
 568      }
 569    }
 570  
 571    async function getDeletedPaths(currentTreeSha: string, targetTreeSha: string): Promise<string[]> {
 572      const result = await execGitChecked([
 573        "diff",
 574        "--name-only",
 575        "--diff-filter=D",
 576        "-z",
 577        currentTreeSha,
 578        targetTreeSha,
 579        "--",
 580      ]);
 581  
 582      return result.stdout.split("\0").filter(Boolean);
 583    }
 584  
 585    async function restoreCommitExactly(targetCommitSha: string): Promise<{ changed: boolean; undoCommitSha?: string; targetTreeSha: string }> {
 586      const { treeSha: currentTreeSha } = await captureWorktreeTree();
 587      const targetTreeSha = await getCommitTreeSha(targetCommitSha);
 588  
 589      if (currentTreeSha === targetTreeSha) {
 590        lastExact = { commitSha: targetCommitSha, treeSha: targetTreeSha };
 591        return { changed: false, targetTreeSha };
 592      }
 593  
 594      const undoCommitSha = await ensureSnapshotForTree(currentTreeSha);
 595      const pathsToDelete = await getDeletedPaths(currentTreeSha, targetTreeSha);
 596      await deletePathsFromWorkingTree(pathsToDelete);
 597      await execGitChecked(["restore", `--source=${targetCommitSha}`, "--worktree", "--", "."]);
 598      lastExact = { commitSha: targetCommitSha, treeSha: targetTreeSha };
 599      return { changed: true, undoCommitSha, targetTreeSha };
 600    }
 601  
 602    function bindPendingPromptUser(entries: SessionLikeEntry[], collector: ActivePromptCollector) {
 603      if (!collector.pendingUserCommitSha) return;
 604  
 605      const userEntry = findLatestMatchingUserMessageEntry(entries, collector.promptText) ?? findLatestUserMessageEntry(entries);
 606      if (!userEntry) return;
 607      if (collector.bindings.some(([entryId]) => entryId === userEntry.id)) {
 608        collector.pendingUserCommitSha = undefined;
 609        return;
 610      }
 611  
 612      addBindingToCollector(collector, userEntry.id, collector.pendingUserCommitSha);
 613      collector.pendingUserCommitSha = undefined;
 614    }
 615  
 616    function appendRewindTurn(ctx: ExtensionContext, collector: ActivePromptCollector) {
 617      if (collector.bindings.length === 0) return;
 618  
 619      const data: RewindTurnData = {
 620        v: RETENTION_VERSION,
 621        snapshots: collector.snapshots,
 622        bindings: collector.bindings,
 623      };
 624  
 625      pi.appendEntry("rewind-turn", data);
 626      applyBindings(entryToCommit, data.snapshots, data.bindings);
 627  
 628      const latestBinding = data.bindings[data.bindings.length - 1];
 629      if (latestBinding) {
 630        activeBranchState.currentCommitSha = data.snapshots[latestBinding[1]];
 631        activeBranchState.currentTreeSha = lastExact?.commitSha === activeBranchState.currentCommitSha ? lastExact.treeSha : undefined;
 632      }
 633  
 634      updateStatus(ctx);
 635    }
 636  
 637    function appendRewindOp(ctx: ExtensionContext, data: RewindOpData) {
 638      const hasBindings = Boolean(data.bindings?.length);
 639      const hasCurrent = typeof data.current === "number";
 640      const hasUndo = typeof data.undo === "number";
 641      if (!hasBindings && !hasCurrent && !hasUndo) return;
 642  
 643      pi.appendEntry("rewind-op", data);
 644      applyBindings(entryToCommit, data.snapshots, data.bindings);
 645  
 646      const currentCommitSha = getCommitFromData(data, "current");
 647      if (currentCommitSha) {
 648        activeBranchState.currentCommitSha = currentCommitSha;
 649        activeBranchState.currentTreeSha = lastExact?.commitSha === currentCommitSha ? lastExact.treeSha : undefined;
 650      }
 651  
 652      const undoCommitSha = getCommitFromData(data, "undo");
 653      if (undoCommitSha) {
 654        activeBranchState.undoCommitSha = undoCommitSha;
 655      }
 656  
 657      updateStatus(ctx);
 658    }
 659  
 660    function buildCurrentSessionLedger(ctx: ExtensionContext): ParsedSessionLedger {
 661      const ledger: ParsedSessionLedger = {
 662        sessionFile: currentSessionFile ?? "",
 663        sessionId: ctx.sessionManager.getSessionId(),
 664        cwd: ctx.sessionManager.getCwd(),
 665        parentSession: ctx.sessionManager.getHeader()?.parentSession,
 666        entryToCommit: new Map<string, string>(),
 667        labeledEntryIds: new Set<string>(),
 668        references: [],
 669      };
 670  
 671      for (const rawEntry of ctx.sessionManager.getEntries() as SessionLikeEntry[]) {
 672        if (rawEntry.type === "custom" && rawEntry.customType === "rewind-turn" && isRewindTurnData(rawEntry.data)) {
 673          applyBindings(ledger.entryToCommit, rawEntry.data.snapshots, rawEntry.data.bindings);
 674          addReferences(ledger.references, rawEntry.data.snapshots, toTimestamp(rawEntry.timestamp), rawEntry.data);
 675          continue;
 676        }
 677  
 678        if (rawEntry.type === "custom" && rawEntry.customType === "rewind-op" && isRewindOpData(rawEntry.data)) {
 679          applyBindings(ledger.entryToCommit, rawEntry.data.snapshots, rawEntry.data.bindings);
 680          addReferences(ledger.references, rawEntry.data.snapshots, toTimestamp(rawEntry.timestamp), rawEntry.data);
 681          const currentCommitSha = getCommitFromData(rawEntry.data, "current");
 682          if (currentCommitSha) ledger.latestCurrentCommitSha = currentCommitSha;
 683          const undoCommitSha = getCommitFromData(rawEntry.data, "undo");
 684          if (undoCommitSha) ledger.latestUndoCommitSha = undoCommitSha;
 685          continue;
 686        }
 687  
 688        if (rawEntry.type === "label") {
 689          updateLabelSet(ledger.labeledEntryIds, rawEntry);
 690        }
 691      }
 692  
 693      return ledger;
 694    }
 695  
 696    async function parseSessionLedgerFile(sessionFile: string): Promise<ParsedSessionLedger | null> {
 697      try {
 698        const fileStat = await stat(sessionFile);
 699        const cached = parsedSessionCache.get(sessionFile);
 700        if (cached && cached.mtimeMs === fileStat.mtimeMs) {
 701          return cached.ledger;
 702        }
 703  
 704        const content = await readFile(sessionFile, "utf-8");
 705        const ledger: ParsedSessionLedger = {
 706          sessionFile,
 707          entryToCommit: new Map<string, string>(),
 708          labeledEntryIds: new Set<string>(),
 709          references: [],
 710        };
 711  
 712        const hasRewindEntries = content.includes('"rewind-');
 713  
 714        if (!hasRewindEntries) {
 715          // Fast path: extract session header only, skip line-by-line JSON parsing
 716          let pos = 0;
 717          for (let i = 0; i < 5 && pos < content.length; i++) {
 718            const nextNewline = content.indexOf("\n", pos);
 719            const line = nextNewline >= 0 ? content.substring(pos, nextNewline) : content.substring(pos);
 720            pos = nextNewline >= 0 ? nextNewline + 1 : content.length;
 721            if (!line) continue;
 722            try {
 723              const entry = JSON.parse(line);
 724              if (entry?.type === "session") {
 725                ledger.sessionId = entry.id;
 726                ledger.cwd = entry.cwd;
 727                ledger.parentSession = entry.parentSession;
 728                break;
 729              }
 730            } catch { continue; }
 731          }
 732          parsedSessionCache.set(sessionFile, { mtimeMs: fileStat.mtimeMs, ledger });
 733          return ledger;
 734        }
 735  
 736        const lines = content.split("\n").filter(Boolean);
 737  
 738        for (const line of lines) {
 739          let entry: any;
 740          try {
 741            entry = JSON.parse(line);
 742          } catch {
 743            continue;
 744          }
 745  
 746          if (entry?.type === "session") {
 747            ledger.sessionId = entry.id;
 748            ledger.cwd = entry.cwd;
 749            ledger.parentSession = entry.parentSession;
 750            continue;
 751          }
 752  
 753          if (entry?.type === "custom" && entry?.customType === "rewind-turn" && isRewindTurnData(entry.data)) {
 754            applyBindings(ledger.entryToCommit, entry.data.snapshots, entry.data.bindings);
 755            addReferences(ledger.references, entry.data.snapshots, toTimestamp(entry.timestamp), entry.data);
 756            continue;
 757          }
 758  
 759          if (entry?.type === "custom" && entry?.customType === "rewind-op" && isRewindOpData(entry.data)) {
 760            applyBindings(ledger.entryToCommit, entry.data.snapshots, entry.data.bindings);
 761            addReferences(ledger.references, entry.data.snapshots, toTimestamp(entry.timestamp), entry.data);
 762            const currentCommitSha = getCommitFromData(entry.data, "current");
 763            if (currentCommitSha) ledger.latestCurrentCommitSha = currentCommitSha;
 764            const undoCommitSha = getCommitFromData(entry.data, "undo");
 765            if (undoCommitSha) ledger.latestUndoCommitSha = undoCommitSha;
 766            continue;
 767          }
 768  
 769          if (entry?.type === "label") {
 770            updateLabelSet(ledger.labeledEntryIds, entry);
 771          }
 772        }
 773  
 774        parsedSessionCache.set(sessionFile, { mtimeMs: fileStat.mtimeMs, ledger });
 775        return ledger;
 776      } catch {
 777        return null;
 778      }
 779    }
 780  
 781    async function listLegacyRefs(): Promise<LegacyRef[]> {
 782      try {
 783        const result = await execGitChecked([
 784          "for-each-ref",
 785          "--format=%(refname) %(objectname)",
 786          LEGACY_REF_PREFIX,
 787        ]);
 788  
 789        return result.stdout
 790          .split("\n")
 791          .map((line) => line.trim())
 792          .filter(Boolean)
 793          .map((line) => {
 794            const [refName, commitSha] = line.split(/\s+/, 2);
 795            const shortRef = refName.replace(LEGACY_REF_PREFIX, "");
 796            if (shortRef.startsWith("checkpoint-resume-") || shortRef.startsWith("before-restore-")) {
 797              return null;
 798            }
 799  
 800            const scoped = shortRef.match(/^checkpoint-([a-f0-9-]{36})-(\d+)-(.+)$/);
 801            if (scoped) {
 802              return {
 803                refName,
 804                commitSha,
 805                sessionId: scoped[1],
 806                entryId: scoped[3],
 807                scoped: true,
 808              } satisfies LegacyRef;
 809            }
 810  
 811            const unscoped = shortRef.match(/^checkpoint-(\d+)-(.+)$/);
 812            if (unscoped) {
 813              return {
 814                refName,
 815                commitSha,
 816                entryId: unscoped[2],
 817                scoped: false,
 818              } satisfies LegacyRef;
 819            }
 820  
 821            return null;
 822          })
 823          .filter((value): value is LegacyRef => value !== null);
 824      } catch {
 825        return [];
 826      }
 827    }
 828  
 829    async function importLegacyRefsIfNeeded(ctx: ExtensionContext) {
 830      if (!sessionId) return;
 831  
 832      const legacyRefs = await listLegacyRefs();
 833      if (legacyRefs.length === 0) return;
 834  
 835      const snapshots: string[] = [];
 836      const bindings: BindingTuple[] = [];
 837      const refsToDelete: string[] = [];
 838      const entries = ctx.sessionManager.getEntries() as SessionLikeEntry[];
 839  
 840      for (const entry of entries) {
 841        if (entry.type !== "message" || entry.message.role !== "user") continue;
 842        if (entryToCommit.has(entry.id)) continue;
 843  
 844        const legacyRef = legacyRefs.find((candidate) => {
 845          if (candidate.entryId !== entry.id) return false;
 846          if (candidate.scoped) return candidate.sessionId === sessionId;
 847          return true;
 848        });
 849  
 850        if (!legacyRef) continue;
 851        if (!(await commitExists(legacyRef.commitSha))) continue;
 852  
 853        await appendSnapshotToStore(legacyRef.commitSha);
 854        const snapshotIndex = resolveBindingSnapshotIndex(snapshots, legacyRef.commitSha);
 855        bindings.push([entry.id, snapshotIndex]);
 856        if (legacyRef.scoped) {
 857          refsToDelete.push(legacyRef.refName);
 858        }
 859      }
 860  
 861      if (bindings.length === 0) return;
 862  
 863      appendRewindOp(ctx, {
 864        v: RETENTION_VERSION,
 865        snapshots,
 866        bindings,
 867      });
 868  
 869      for (const refName of refsToDelete) {
 870        try {
 871          const result = await pi.exec("git", ["update-ref", "-d", refName]);
 872          if (result.code !== 0) {
 873            throw new Error(result.stderr.trim() || `git update-ref -d ${refName} failed with code ${result.code}`);
 874          }
 875        } catch {
 876          // Ignore legacy cleanup failures
 877        }
 878      }
 879    }
 880  
 881    async function resolveLegacyAncestorCommit(entryId: string, targetSessionId: string | undefined): Promise<string | undefined> {
 882      if (!targetSessionId) return undefined;
 883      const legacyRefs = await listLegacyRefs();
 884      const match = legacyRefs.find((ref) => ref.scoped && ref.sessionId === targetSessionId && ref.entryId === entryId);
 885      if (!match) return undefined;
 886      return (await commitExists(match.commitSha)) ? match.commitSha : undefined;
 887    }
 888  
 889    async function resolveEntrySnapshotWithLineage(entryId: string, sessionFile = currentSessionFile): Promise<string | undefined> {
 890      let cursor = sessionFile;
 891  
 892      while (cursor) {
 893        const ledger = cursor === currentSessionFile ? buildCurrentSessionLedgerFromMemory() : await parseSessionLedgerFile(cursor);
 894        if (!ledger) break;
 895  
 896        const commitSha = ledger.entryToCommit.get(entryId);
 897        if (commitSha && (await commitExists(commitSha))) {
 898          return commitSha;
 899        }
 900  
 901        const legacyCommitSha = await resolveLegacyAncestorCommit(entryId, ledger.sessionId);
 902        if (legacyCommitSha) {
 903          return legacyCommitSha;
 904        }
 905  
 906        cursor = ledger.parentSession;
 907      }
 908  
 909      return undefined;
 910    }
 911  
 912    function buildCurrentSessionLedgerFromMemory(): ParsedSessionLedger {
 913      return {
 914        sessionFile: currentSessionFile ?? "",
 915        sessionId: sessionId ?? undefined,
 916        cwd: currentSessionCwd,
 917        parentSession: currentParentSession,
 918        entryToCommit: new Map(entryToCommit),
 919        labeledEntryIds: new Set(),
 920        references: [],
 921        latestCurrentCommitSha: activeBranchState.currentCommitSha,
 922        latestUndoCommitSha: activeBranchState.undoCommitSha,
 923      };
 924    }
 925  
 926    async function reconstructState(ctx: ExtensionContext) {
 927      entryToCommit.clear();
 928      activeBranchState = {};
 929      lastExact = null;
 930  
 931      const currentLedger = buildCurrentSessionLedger(ctx);
 932      for (const [entryId, commitSha] of currentLedger.entryToCommit.entries()) {
 933        entryToCommit.set(entryId, commitSha);
 934      }
 935  
 936      let latestVisibleBindingCommitSha: string | undefined;
 937      for (const entry of ctx.sessionManager.getBranch() as SessionLikeEntry[]) {
 938        const boundCommitSha = entry.id ? entryToCommit.get(entry.id) : undefined;
 939        if (boundCommitSha && isRestorableTreeEntry(entry)) {
 940          latestVisibleBindingCommitSha = boundCommitSha;
 941        }
 942  
 943        if (entry.type === "custom" && entry.customType === "rewind-op" && isRewindOpData(entry.data)) {
 944          const currentCommitSha = getCommitFromData(entry.data, "current");
 945          if (currentCommitSha) {
 946            activeBranchState.currentCommitSha = currentCommitSha;
 947          }
 948          const undoCommitSha = getCommitFromData(entry.data, "undo");
 949          if (undoCommitSha) {
 950            activeBranchState.undoCommitSha = undoCommitSha;
 951          }
 952        }
 953      }
 954  
 955      if (!activeBranchState.currentCommitSha) {
 956        activeBranchState.currentCommitSha = latestVisibleBindingCommitSha;
 957      }
 958  
 959      if (activeBranchState.currentCommitSha && (await commitExists(activeBranchState.currentCommitSha))) {
 960        activeBranchState.currentTreeSha = await getCommitTreeSha(activeBranchState.currentCommitSha);
 961        const { treeSha: worktreeTreeSha } = await captureWorktreeTree();
 962  
 963        if (activeBranchState.currentTreeSha === worktreeTreeSha) {
 964          lastExact = {
 965            commitSha: activeBranchState.currentCommitSha,
 966            treeSha: activeBranchState.currentTreeSha,
 967          };
 968        }
 969      }
 970    }
 971  
 972    async function discoverSessionFiles(): Promise<string[]> {
 973      const roots = new Set<string>();
 974      const defaultSessionsDir = getDefaultSessionsDir();
 975      if (existsSync(defaultSessionsDir)) {
 976        roots.add(defaultSessionsDir);
 977      }
 978      if (currentSessionFile) {
 979        roots.add(dirname(currentSessionFile));
 980      }
 981  
 982      const discovered = new Set<string>();
 983      const stack = [...roots];
 984      while (stack.length > 0) {
 985        const dir = stack.pop();
 986        if (!dir) continue;
 987  
 988        let entries: Awaited<ReturnType<typeof readdir>>;
 989        try {
 990          entries = await readdir(dir, { withFileTypes: true });
 991        } catch {
 992          continue;
 993        }
 994  
 995        for (const entry of entries) {
 996          const fullPath = join(dir, entry.name);
 997          if (entry.isDirectory()) {
 998            stack.push(fullPath);
 999            continue;
1000          }
1001          if (entry.isFile() && entry.name.endsWith(".jsonl")) {
1002            discovered.add(fullPath);
1003          }
1004        }
1005      }
1006  
1007      let ancestorCursor = currentSessionFile;
1008      while (ancestorCursor) {
1009        discovered.add(ancestorCursor);
1010        const ledger = ancestorCursor === currentSessionFile ? buildCurrentSessionLedgerFromMemory() : await parseSessionLedgerFile(ancestorCursor);
1011        ancestorCursor = ledger?.parentSession;
1012      }
1013  
1014      return [...discovered];
1015    }
1016  
1017    async function maybeSweepRetention(ctx: ExtensionContext, reason: "startup" | "new-snapshots" | "shutdown") {
1018      const retention = getRetentionSettings();
1019      if (!retention) return;
1020      if (reason === "new-snapshots" && newSnapshotsSinceSweep < RETENTION_SWEEP_THRESHOLD) return;
1021      if (reason === "shutdown" && sweepCompletedThisSession && newSnapshotsSinceSweep < RETENTION_SWEEP_THRESHOLD) return;
1022      if (!repoRoot) return;
1023      if (sweepRunning) return;
1024      sweepRunning = true;
1025      try {
1026        await runRetentionSweep(ctx, reason);
1027      } finally {
1028        sweepRunning = false;
1029      }
1030    }
1031  
1032    async function runRetentionSweep(ctx: ExtensionContext, reason: "startup" | "new-snapshots" | "shutdown") {
1033      const retention = getRetentionSettings();
1034      if (!retention) return;
1035  
1036      const sessionFiles = await discoverSessionFiles();
1037      const ledgers: ParsedSessionLedger[] = [];
1038  
1039      for (const sessionFile of sessionFiles) {
1040        const ledger = sessionFile === currentSessionFile ? buildCurrentSessionLedger(ctx) : await parseSessionLedgerFile(sessionFile);
1041        if (!ledger?.cwd) continue;
1042        if (!isInsidePath(ledger.cwd, repoRoot)) continue;
1043        ledgers.push(ledger);
1044      }
1045  
1046      const latestReferenceByCommit = new Map<string, number>();
1047      const pinnedCommits = new Set<string>();
1048      const currentCommits = new Set<string>();
1049      const undoCommits = new Set<string>();
1050  
1051      for (const ledger of ledgers) {
1052        for (const reference of ledger.references) {
1053          const prev = latestReferenceByCommit.get(reference.commitSha) ?? 0;
1054          if (reference.timestamp > prev) {
1055            latestReferenceByCommit.set(reference.commitSha, reference.timestamp);
1056          }
1057          if (reference.kind === "binding" && retention.pinLabeledEntries && reference.entryId && ledger.labeledEntryIds.has(reference.entryId)) {
1058            pinnedCommits.add(reference.commitSha);
1059          }
1060        }
1061  
1062        if (ledger.latestCurrentCommitSha) {
1063          currentCommits.add(ledger.latestCurrentCommitSha);
1064        }
1065        if (ledger.latestUndoCommitSha) {
1066          undoCommits.add(ledger.latestUndoCommitSha);
1067        }
1068      }
1069  
1070      for (const commitSha of [...currentCommits, ...undoCommits]) {
1071        if (await commitExists(commitSha)) {
1072          pinnedCommits.add(commitSha);
1073        }
1074      }
1075  
1076      let candidates = [...latestReferenceByCommit.entries()]
1077        .filter(([commitSha]) => !pinnedCommits.has(commitSha))
1078        .sort((left, right) => right[1] - left[1]);
1079  
1080      if (typeof retention.maxAgeDays === "number" && retention.maxAgeDays >= 0) {
1081        const cutoff = Date.now() - retention.maxAgeDays * 24 * 60 * 60 * 1000;
1082        candidates = candidates.filter(([, timestamp]) => timestamp >= cutoff);
1083      }
1084  
1085      if (typeof retention.maxSnapshots === "number" && retention.maxSnapshots >= 0 && candidates.length > retention.maxSnapshots) {
1086        candidates = candidates.slice(0, retention.maxSnapshots);
1087      }
1088  
1089      const liveSet = [...new Set([...pinnedCommits, ...candidates.map(([commitSha]) => commitSha)])];
1090      const existingLiveSet: string[] = [];
1091      for (const commitSha of liveSet) {
1092        if (await commitExists(commitSha)) {
1093          existingLiveSet.push(commitSha);
1094        }
1095      }
1096  
1097      const rewriteResult = await rewriteStoreToLiveSet(existingLiveSet);
1098      if (rewriteResult === "preserved-empty") {
1099        return;
1100      }
1101  
1102      // Skip gc on background startup sweeps to avoid racing with concurrent snapshot creation
1103      if (reason !== "startup") {
1104        try {
1105          const result = await pi.exec("git", ["gc", "--auto"]);
1106          if (result.code !== 0) {
1107            throw new Error(result.stderr.trim() || `git gc --auto failed with code ${result.code}`);
1108          }
1109        } catch {
1110          // Best effort only
1111        }
1112      }
1113  
1114      newSnapshotsSinceSweep = 0;
1115      sweepCompletedThisSession = true;
1116      updateStatus(ctx);
1117    }
1118  
1119    async function initializeForSession(ctx: ExtensionContext) {
1120      resetState();
1121      syncSessionIdentity(ctx);
1122  
1123      try {
1124        const result = await pi.exec("git", ["rev-parse", "--is-inside-work-tree"]);
1125        isGitRepo = result.code === 0 && result.stdout.trim() === "true";
1126      } catch {
1127        isGitRepo = false;
1128      }
1129  
1130      if (!isGitRepo) {
1131        updateStatus(ctx);
1132        return;
1133      }
1134  
1135      await getRepoRoot(pi.exec);
1136      await reconstructState(ctx);
1137      await importLegacyRefsIfNeeded(ctx);
1138      await reconstructState(ctx);
1139      updateStatus(ctx);
1140      maybeSweepRetention(ctx, "startup").catch(() => {});
1141    }
1142  
1143    pi.events.on("rewind:fork-preference", (data: any) => {
1144      if (data?.mode !== "conversation-only") return;
1145      if (typeof data?.source !== "string") return;
1146      if (!FORK_PREFERENCE_SOURCE_ALLOWLIST.has(data.source)) return;
1147      forceConversationOnlyOnNextFork = true;
1148      forceConversationOnlySource = data.source;
1149    });
1150  
1151    pi.on("before_agent_start", async (event) => {
1152      activePromptText = event.prompt;
1153    });
1154  
1155    pi.on("session_start", async (_event, ctx) => {
1156      await initializeForSession(ctx);
1157    });
1158  
1159    pi.on("session_switch", async (_event, ctx) => {
1160      await initializeForSession(ctx);
1161    });
1162  
1163    pi.on("session_fork", async (_event, ctx) => {
1164      syncSessionIdentity(ctx);
1165      if (!isGitRepo || !pendingForkState) {
1166        await reconstructState(ctx);
1167        updateStatus(ctx);
1168        return;
1169      }
1170  
1171      const snapshots = [pendingForkState.currentCommitSha];
1172      const data: RewindOpData = { v: RETENTION_VERSION, snapshots, current: 0 };
1173      if (pendingForkState.undoCommitSha) {
1174        data.snapshots.push(pendingForkState.undoCommitSha);
1175        data.undo = 1;
1176      }
1177  
1178      appendRewindOp(ctx, data);
1179      pendingForkState = null;
1180      await reconstructState(ctx);
1181      updateStatus(ctx);
1182    });
1183  
1184    pi.on("session_tree", async (event, ctx) => {
1185      syncSessionIdentity(ctx);
1186      if (!isGitRepo || !pendingTreeState) {
1187        await reconstructState(ctx);
1188        updateStatus(ctx);
1189        return;
1190      }
1191  
1192      const snapshots = [pendingTreeState.currentCommitSha];
1193      const data: RewindOpData = { v: RETENTION_VERSION, snapshots, current: 0 };
1194      if (pendingTreeState.undoCommitSha) {
1195        data.snapshots.push(pendingTreeState.undoCommitSha);
1196        data.undo = 1;
1197      }
1198      if (event.summaryEntry?.id) {
1199        data.bindings = [[event.summaryEntry.id, 0]];
1200      }
1201  
1202      appendRewindOp(ctx, data);
1203      pendingTreeState = null;
1204      await reconstructState(ctx);
1205      updateStatus(ctx);
1206    });
1207  
1208    pi.on("session_compact", async (event, ctx) => {
1209      syncSessionIdentity(ctx);
1210      if (!isGitRepo) return;
1211  
1212      let currentCommitSha = activeBranchState.currentCommitSha;
1213      if (!currentCommitSha) {
1214        currentCommitSha = await ensureSnapshotForCurrentWorktree();
1215      }
1216  
1217      appendRewindOp(ctx, {
1218        v: RETENTION_VERSION,
1219        snapshots: [currentCommitSha],
1220        bindings: [[event.compactionEntry.id, 0]],
1221      });
1222      await reconstructState(ctx);
1223      updateStatus(ctx);
1224    });
1225  
1226    pi.on("session_shutdown", async (_event, ctx) => {
1227      syncSessionIdentity(ctx);
1228      if (!isGitRepo) return;
1229      await maybeSweepRetention(ctx, "shutdown");
1230    });
1231  
1232    pi.on("turn_start", async (event, ctx) => {
1233      if (!isGitRepo) return;
1234      if (event.turnIndex !== 0) return;
1235  
1236      const { treeSha } = await captureWorktreeTree();
1237      const commitSha = await ensureSnapshotForTree(treeSha);
1238      promptCollector = {
1239        snapshots: [],
1240        bindings: [],
1241        promptText: activePromptText ?? undefined,
1242        pendingUserCommitSha: commitSha,
1243      };
1244  
1245      bindPendingPromptUser(ctx.sessionManager.getBranch() as SessionLikeEntry[], promptCollector);
1246    });
1247  
1248    pi.on("turn_end", async (event, ctx) => {
1249      if (!isGitRepo || !promptCollector) return;
1250  
1251      const branchEntries = ctx.sessionManager.getBranch() as SessionLikeEntry[];
1252      bindPendingPromptUser(branchEntries, promptCollector);
1253  
1254      if (event.message.role !== "assistant") return;
1255  
1256      const assistantEntry = findAssistantEntryForTurn(branchEntries, event.message);
1257      if (!assistantEntry) return;
1258  
1259      const { treeSha } = await captureWorktreeTree();
1260      const commitSha = await ensureSnapshotForTree(treeSha);
1261      addBindingToCollector(promptCollector, assistantEntry.id, commitSha);
1262    });
1263  
1264    pi.on("agent_end", async (_event, ctx) => {
1265      if (!isGitRepo || !promptCollector) return;
1266      bindPendingPromptUser(ctx.sessionManager.getBranch() as SessionLikeEntry[], promptCollector);
1267      appendRewindTurn(ctx, promptCollector);
1268      promptCollector = null;
1269      activePromptText = null;
1270      await reconstructState(ctx);
1271      updateStatus(ctx);
1272      await maybeSweepRetention(ctx, "new-snapshots");
1273    });
1274  
1275    pi.on("session_before_fork", async (event, ctx) => {
1276      const shouldForceConversationOnly = forceConversationOnlyOnNextFork;
1277      const forcedBySource = forceConversationOnlySource;
1278      forceConversationOnlyOnNextFork = false;
1279      forceConversationOnlySource = null;
1280  
1281      if (!isGitRepo) return;
1282      if (!ctx.hasUI) {
1283        pendingForkState = { currentCommitSha: await ensureSnapshotForCurrentWorktree() };
1284        return;
1285      }
1286  
1287      const targetCommitSha = await resolveEntrySnapshotWithLineage(event.entryId);
1288      const hasUndo = Boolean(activeBranchState.undoCommitSha && (await commitExists(activeBranchState.undoCommitSha)));
1289  
1290      if (shouldForceConversationOnly) {
1291        pendingForkState = { currentCommitSha: await ensureSnapshotForCurrentWorktree() };
1292        notify(ctx, `Rewind: using conversation-only fork (keep current files)${forcedBySource ? ` (${forcedBySource})` : ""}`);
1293        return;
1294      }
1295  
1296      const options = ["Conversation only (keep current files)"];
1297      if (targetCommitSha) {
1298        options.push("Restore all (files + conversation)", "Code only (restore files, keep conversation)");
1299      }
1300      if (hasUndo) {
1301        options.push("Undo last file rewind");
1302      }
1303  
1304      const choice = await ctx.ui.select("Restore Options", options);
1305      if (!choice) {
1306        notify(ctx, "Rewind cancelled");
1307        return { cancel: true };
1308      }
1309  
1310      if (choice === "Undo last file rewind") {
1311        const restore = await restoreCommitExactly(activeBranchState.undoCommitSha!);
1312        pendingForkState = {
1313          currentCommitSha: activeBranchState.undoCommitSha!,
1314          undoCommitSha: restore.undoCommitSha,
1315        };
1316        notify(ctx, "Files restored to before last rewind");
1317        return;
1318      }
1319  
1320      if (choice === "Conversation only (keep current files)") {
1321        pendingForkState = { currentCommitSha: await ensureSnapshotForCurrentWorktree() };
1322        return;
1323      }
1324  
1325      if (!targetCommitSha) {
1326        notify(ctx, "No exact rewind point available for that entry", "error");
1327        return { cancel: true };
1328      }
1329  
1330      const restore = await restoreCommitExactly(targetCommitSha);
1331      pendingForkState = {
1332        currentCommitSha: targetCommitSha,
1333        undoCommitSha: restore.undoCommitSha,
1334      };
1335      notify(ctx, "Files restored from rewind point");
1336  
1337      if (choice === "Code only (restore files, keep conversation)") {
1338        return { skipConversationRestore: true };
1339      }
1340    });
1341  
1342    pi.on("session_before_tree", async (event, ctx) => {
1343      if (!isGitRepo || !ctx.hasUI) return;
1344  
1345      const targetEntry = ctx.sessionManager.getEntry(event.preparation.targetId) as SessionLikeEntry | undefined;
1346      const targetCommitSha = isRestorableTreeEntry(targetEntry)
1347        ? await resolveEntrySnapshotWithLineage(event.preparation.targetId, currentSessionFile)
1348        : undefined;
1349      const hasUndo = Boolean(activeBranchState.undoCommitSha && (await commitExists(activeBranchState.undoCommitSha)));
1350  
1351      const options = ["Keep current files"];
1352      if (targetCommitSha) {
1353        options.push("Restore files to that point");
1354      }
1355      if (hasUndo) {
1356        options.push("Undo last file rewind");
1357      }
1358      options.push("Cancel navigation");
1359  
1360      const choice = await ctx.ui.select("Restore Options", options);
1361      if (!choice || choice === "Cancel navigation") {
1362        notify(ctx, "Navigation cancelled");
1363        return { cancel: true };
1364      }
1365  
1366      if (choice === "Undo last file rewind") {
1367        const restore = await restoreCommitExactly(activeBranchState.undoCommitSha!);
1368        const snapshots = [activeBranchState.undoCommitSha!];
1369        const data: RewindOpData = { v: RETENTION_VERSION, snapshots, current: 0 };
1370        if (restore.undoCommitSha) {
1371          data.snapshots.push(restore.undoCommitSha);
1372          data.undo = 1;
1373        }
1374        appendRewindOp(ctx, data);
1375        notify(ctx, "Files restored to before last rewind");
1376        await reconstructState(ctx);
1377        return { cancel: true };
1378      }
1379  
1380      if (choice === "Keep current files") {
1381        pendingTreeState = { currentCommitSha: await ensureSnapshotForCurrentWorktree() };
1382        return;
1383      }
1384  
1385      if (!targetCommitSha) {
1386        notify(ctx, "Exact file rewind is only available for user, assistant, compaction, and summary nodes", "error");
1387        return { cancel: true };
1388      }
1389  
1390      const restore = await restoreCommitExactly(targetCommitSha);
1391      pendingTreeState = {
1392        currentCommitSha: targetCommitSha,
1393        undoCommitSha: restore.undoCommitSha,
1394      };
1395      notify(ctx, "Files restored to rewind point");
1396    });
1397  }