1  import assert from "node:assert/strict";
  2  import { execFile } from "node:child_process";
  3  import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
  4  import { mkdtemp, rm, writeFile } from "node:fs/promises";
  5  import test from "node:test";
  6  import os from "node:os";
  7  import path from "node:path";
  8  import { promisify } from "node:util";
  9  
 10  import rewindExtension from "./index.ts";
 11  
 12  const execFileAsync = promisify(execFile);
 13  const STORE_REF = "refs/pi-rewind/store";
 14  
 15  type RewindEntry = Record<string, unknown>;
 16  type EventHandler = (event: any, ctx: any) => Promise<any> | any;
 17  
 18  class SessionManagerStub {
 19    private readonly header: { type: "session"; version: number; id: string; timestamp: string; cwd: string; parentSession?: string };
 20    private entries: RewindEntry[];
 21    private readonly sessionFile: string;
 22  
 23    constructor(options: {
 24      sessionFile: string;
 25      id: string;
 26      cwd: string;
 27      parentSession?: string;
 28      entries?: RewindEntry[];
 29    }) {
 30      this.sessionFile = options.sessionFile;
 31      this.header = {
 32        type: "session",
 33        version: 3,
 34        id: options.id,
 35        timestamp: new Date().toISOString(),
 36        cwd: options.cwd,
 37        parentSession: options.parentSession,
 38      };
 39      this.entries = options.entries ?? [];
 40      this.flush();
 41    }
 42  
 43    flush(): void {
 44      mkdirSync(path.dirname(this.sessionFile), { recursive: true });
 45      const lines = [this.header, ...this.entries].map((entry) => JSON.stringify(entry)).join("\n") + "\n";
 46      writeFileSync(this.sessionFile, lines);
 47    }
 48  
 49    replaceEntries(entries: RewindEntry[]): void {
 50      this.entries = entries;
 51      this.flush();
 52    }
 53  
 54    appendCustom(customType: string, data: unknown): void {
 55      const parentId = (this.entries.at(-1)?.id as string | undefined) ?? null;
 56      this.entries.push({
 57        type: "custom",
 58        customType,
 59        data,
 60        id: `${customType}-${this.entries.length + 1}`,
 61        parentId,
 62        timestamp: new Date().toISOString(),
 63      });
 64      this.flush();
 65    }
 66  
 67    getSessionId(): string {
 68      return this.header.id;
 69    }
 70  
 71    getSessionFile(): string {
 72      return this.sessionFile;
 73    }
 74  
 75    getHeader(): { parentSession?: string } {
 76      return { parentSession: this.header.parentSession };
 77    }
 78  
 79    getCwd(): string {
 80      return this.header.cwd;
 81    }
 82  
 83    getEntries(): RewindEntry[] {
 84      return this.entries;
 85    }
 86  
 87    getBranch(): RewindEntry[] {
 88      return this.entries;
 89    }
 90  
 91    getEntry(entryId: string): RewindEntry | undefined {
 92      return this.entries.find((entry) => entry.id === entryId);
 93    }
 94  }
 95  
 96  async function runGit(repoRoot: string, args: string[]): Promise<{ stdout: string; stderr: string; code: number }> {
 97    try {
 98      const { stdout, stderr } = await execFileAsync("git", args, { cwd: repoRoot });
 99      return { stdout, stderr, code: 0 };
100    } catch (error: any) {
101      return {
102        stdout: error.stdout ?? "",
103        stderr: error.stderr ?? error.message ?? "",
104        code: error.code ?? 1,
105      };
106    }
107  }
108  
109  async function runGitChecked(repoRoot: string, args: string[]): Promise<{ stdout: string; stderr: string; code: number }> {
110    const result = await runGit(repoRoot, args);
111    if (result.code !== 0) {
112      throw new Error(`git ${args.join(" ")} failed: ${result.stderr || `exit ${result.code}`}`);
113    }
114    return result;
115  }
116  
117  async function gitStdout(repoRoot: string, args: string[]): Promise<string> {
118    return (await runGitChecked(repoRoot, args)).stdout.trim();
119  }
120  
121  async function revParseOptional(repoRoot: string, ref: string): Promise<string | undefined> {
122    try {
123      return await gitStdout(repoRoot, ["rev-parse", ref]);
124    } catch {
125      return undefined;
126    }
127  }
128  
129  async function isAncestor(repoRoot: string, ancestor: string, descendant: string): Promise<boolean> {
130    try {
131      await runGitChecked(repoRoot, ["merge-base", "--is-ancestor", ancestor, descendant]);
132      return true;
133    } catch {
134      return false;
135    }
136  }
137  
138  async function captureSnapshot(repoRoot: string): Promise<string> {
139    await runGitChecked(repoRoot, ["add", "-A"]);
140    const treeSha = await gitStdout(repoRoot, ["write-tree"]);
141    return await gitStdout(repoRoot, ["commit-tree", treeSha, "-m", "rewind snapshot test"]);
142  }
143  
144  async function createHarness(options: {
145    settings?: Record<string, unknown>;
146  } = {}) {
147    const root = await mkdtemp(path.join(os.tmpdir(), "rewind-ext-test-"));
148    const repoRoot = path.join(root, "repo");
149    const agentDir = path.join(root, "agent");
150    const sessionsDir = path.join(agentDir, "sessions", "--repo--");
151    mkdirSync(repoRoot, { recursive: true });
152    mkdirSync(sessionsDir, { recursive: true });
153    writeFileSync(path.join(agentDir, "settings.json"), JSON.stringify(options.settings ?? {}, null, 2) + "\n");
154  
155    const originalAgentDir = process.env.PI_CODING_AGENT_DIR;
156    process.env.PI_CODING_AGENT_DIR = agentDir;
157  
158    await runGitChecked(repoRoot, ["init"]);
159    await runGitChecked(repoRoot, ["config", "user.name", "Rewind Test"]);
160    await runGitChecked(repoRoot, ["config", "user.email", "rewind@example.com"]);
161  
162    const handlers = new Map<string, EventHandler>();
163    const eventHandlers = new Map<string, (data: any) => void>();
164    const execCalls: string[][] = [];
165    const notifications: Array<{ message: string; level: string }> = [];
166    const statusUpdates: Array<{ key: string; value: string | undefined }> = [];
167    const selectCalls: Array<{ title: string; options: string[] }> = [];
168    const pendingSelections: string[] = [];
169  
170    const currentSession = new SessionManagerStub({
171      sessionFile: path.join(sessionsDir, "session-1.jsonl"),
172      id: "session-1",
173      cwd: repoRoot,
174    });
175    let activeSession = currentSession;
176  
177    const api = {
178      exec: async (cmd: string, args: string[]) => {
179        execCalls.push([cmd, ...args]);
180        if (cmd !== "git") {
181          throw new Error(`Unsupported command in test harness: ${cmd}`);
182        }
183        return runGit(repoRoot, args);
184      },
185      appendEntry: (customType: string, data: unknown) => {
186        activeSession.appendCustom(customType, data);
187      },
188      on: (eventName: string, handler: EventHandler) => {
189        handlers.set(eventName, handler);
190      },
191      events: {
192        on: (eventName: string, handler: (data: any) => void) => {
193          eventHandlers.set(eventName, handler);
194        },
195      },
196    } as any;
197  
198    rewindExtension(api);
199  
200    function createContext(sessionManager: SessionManagerStub, hasUI = true): any {
201      return {
202        cwd: repoRoot,
203        hasUI,
204        sessionManager,
205        ui: {
206          notify: (message: string, level: string) => {
207            notifications.push({ message, level });
208          },
209          setStatus: (key: string, value: string | undefined) => {
210            statusUpdates.push({ key, value });
211          },
212          select: async (title: string, choices: string[]) => {
213            selectCalls.push({ title, options: choices });
214            return pendingSelections.shift();
215          },
216          theme: {
217            fg: (_color: string, text: string) => text,
218          },
219        },
220      };
221    }
222  
223    return {
224      repoRoot,
225      agentDir,
226      currentSession,
227      execCalls,
228      notifications,
229      selectCalls,
230      statusUpdates,
231      enqueueSelection(choice: string) {
232        pendingSelections.push(choice);
233      },
234      async writeRepoFile(relativePath: string, content: string) {
235        const filePath = path.join(repoRoot, relativePath);
236        mkdirSync(path.dirname(filePath), { recursive: true });
237        await writeFile(filePath, content);
238      },
239      readRepoFile(relativePath: string) {
240        return readFileSync(path.join(repoRoot, relativePath), "utf-8");
241      },
242      createSession(options: { id: string; parentSession?: string; entries?: RewindEntry[] }) {
243        return new SessionManagerStub({
244          sessionFile: path.join(sessionsDir, `${options.id}.jsonl`),
245          id: options.id,
246          cwd: repoRoot,
247          parentSession: options.parentSession,
248          entries: options.entries,
249        });
250      },
251      async invoke(eventName: string, event: any, sessionManager = activeSession, hasUI = true) {
252        const handler = handlers.get(eventName);
253        assert.ok(handler, `missing handler for ${eventName}`);
254        activeSession = sessionManager;
255        return handler(event, createContext(sessionManager, hasUI));
256      },
257      async captureSnapshot() {
258        return captureSnapshot(repoRoot);
259      },
260      async revParseStore() {
261        return revParseOptional(repoRoot, STORE_REF);
262      },
263      async updateStoreRef(commitSha: string) {
264        await runGitChecked(repoRoot, ["update-ref", STORE_REF, commitSha]);
265      },
266      async isAncestor(ancestor: string, descendant: string) {
267        return isAncestor(repoRoot, ancestor, descendant);
268      },
269      eventHandlers,
270      async cleanup() {
271        if (originalAgentDir === undefined) {
272          delete process.env.PI_CODING_AGENT_DIR;
273        } else {
274          process.env.PI_CODING_AGENT_DIR = originalAgentDir;
275        }
276        await rm(root, { recursive: true, force: true });
277      },
278    };
279  }
280  
281  test("/fork undo restores files into a child session instead of cancelling the fork", async () => {
282    const harness = await createHarness({ settings: { rewind: { silentCheckpoints: true } } });
283  
284    try {
285      await harness.writeRepoFile("notes.txt", "current state\n");
286      const currentCommit = await harness.captureSnapshot();
287      await harness.writeRepoFile("notes.txt", "undo target\n");
288      const undoCommit = await harness.captureSnapshot();
289      await harness.writeRepoFile("notes.txt", "current state\n");
290  
291      harness.currentSession.replaceEntries([
292        {
293          type: "message",
294          id: "user-1",
295          parentId: null,
296          timestamp: new Date().toISOString(),
297          message: { role: "user", content: [{ type: "text", text: "Fork from here" }] },
298        },
299        {
300          type: "custom",
301          id: "rewind-op-1",
302          parentId: "user-1",
303          timestamp: new Date().toISOString(),
304          customType: "rewind-op",
305          data: { v: 2, snapshots: [currentCommit, undoCommit], current: 0, undo: 1 },
306        },
307      ]);
308  
309      await harness.invoke("session_start", {});
310      harness.enqueueSelection("Undo last file rewind");
311  
312      const result = await harness.invoke("session_before_fork", { entryId: "user-1" });
313      assert.equal(result, undefined);
314      assert.equal(harness.readRepoFile("notes.txt"), "undo target\n");
315  
316      const currentSessionRewindOps = harness.currentSession.getEntries().filter((entry) => entry.type === "custom" && entry.customType === "rewind-op");
317      assert.equal(currentSessionRewindOps.length, 1);
318  
319      const childSession = harness.createSession({
320        id: "session-2",
321        parentSession: harness.currentSession.getSessionFile(),
322      });
323      await harness.invoke("session_fork", {}, childSession);
324  
325      const childRewindOps = childSession.getEntries().filter((entry) => entry.type === "custom" && entry.customType === "rewind-op");
326      assert.equal(childRewindOps.length, 1);
327      assert.deepEqual(childRewindOps[0]?.data, {
328        v: 2,
329        snapshots: [undoCommit, currentCommit],
330        current: 0,
331        undo: 1,
332      });
333    } finally {
334      await harness.cleanup();
335    }
336  });
337  
338  test("first mutating turn creates a reachable store ref even when retention is omitted", async () => {
339    const harness = await createHarness({
340      settings: { rewind: { silentCheckpoints: true } },
341    });
342  
343    try {
344      const assistantTimestamp = Date.now();
345      harness.currentSession.replaceEntries([
346        {
347          type: "message",
348          id: "user-1",
349          parentId: null,
350          timestamp: new Date(assistantTimestamp - 1000).toISOString(),
351          message: { role: "user", content: [{ type: "text", text: "Please create the file" }] },
352        },
353        {
354          type: "message",
355          id: "assistant-1",
356          parentId: "user-1",
357          timestamp: new Date(assistantTimestamp).toISOString(),
358          message: {
359            role: "assistant",
360            timestamp: assistantTimestamp,
361            content: [{ type: "text", text: "Created the file" }],
362          },
363        },
364      ]);
365  
366      await harness.invoke("session_start", {});
367      await harness.invoke("before_agent_start", { prompt: "Please create the file" });
368      await harness.invoke("turn_start", { turnIndex: 0 });
369      await harness.writeRepoFile("tests/rewind-smoke/a.txt", "smoke test\n");
370      await harness.invoke("turn_end", {
371        message: {
372          role: "assistant",
373          timestamp: assistantTimestamp,
374          content: [{ type: "text", text: "Created the file" }],
375        },
376      });
377      await harness.invoke("agent_end", {});
378  
379      const rewindTurnEntries = harness.currentSession.getEntries().filter((entry) => entry.type === "custom" && entry.customType === "rewind-turn");
380      assert.equal(rewindTurnEntries.length, 1);
381      const snapshots = (rewindTurnEntries[0]?.data as { snapshots: string[] }).snapshots;
382      assert.equal(snapshots.length, 2);
383  
384      const storeHead = await harness.revParseStore();
385      assert.ok(storeHead);
386      assert.equal(await harness.isAncestor(snapshots[0], storeHead), true);
387      assert.equal(await harness.isAncestor(snapshots[1], storeHead), true);
388  
389    } finally {
390      await harness.cleanup();
391    }
392  });
393  
394  test("startup does not touch the keepalive ref when rewind.retention is omitted", async () => {
395    const harness = await createHarness({ settings: { rewind: { silentCheckpoints: true } } });
396  
397    try {
398      await harness.writeRepoFile("tracked.txt", "keepalive\n");
399      const snapshotCommit = await harness.captureSnapshot();
400      await harness.updateStoreRef(snapshotCommit);
401  
402      await harness.invoke("session_start", {});
403  
404      assert.equal(await harness.revParseStore(), snapshotCommit);
405      assert.equal(harness.execCalls.some((call) => call[0] === "git" && call[1] === "gc"), false);
406      assert.equal(harness.execCalls.some((call) => call[0] === "git" && call[1] === "update-ref" && call.includes(STORE_REF)), false);
407    } finally {
408      await harness.cleanup();
409    }
410  });
411  
412  test("retention preserves the keepalive ref when discovery yields an empty live set", async () => {
413    const harness = await createHarness({ settings: { rewind: { retention: { maxSnapshots: 10 } } } });
414  
415    try {
416      await harness.writeRepoFile("tracked.txt", "keepalive\n");
417      const snapshotCommit = await harness.captureSnapshot();
418      await harness.updateStoreRef(snapshotCommit);
419  
420      await harness.invoke("session_start", {});
421  
422      assert.equal(await harness.revParseStore(), snapshotCommit);
423      assert.equal(harness.execCalls.some((call) => call[0] === "git" && call[1] === "gc"), false);
424    } finally {
425      await harness.cleanup();
426    }
427  });
428  
429  test("retention rewrites the keepalive ref when a live snapshot exists", async () => {
430    const harness = await createHarness({
431      settings: { rewind: { retention: { maxSnapshots: 10 } } },
432    });
433  
434    try {
435      await harness.writeRepoFile("tracked.txt", "stale state\n");
436      const staleCommit = await harness.captureSnapshot();
437      await harness.writeRepoFile("tracked.txt", "current live state\n");
438      const liveCommit = await harness.captureSnapshot();
439      assert.notEqual(liveCommit, staleCommit);
440      await harness.updateStoreRef(staleCommit);
441  
442      harness.currentSession.replaceEntries([
443        {
444          type: "custom",
445          id: "rewind-op-1",
446          parentId: null,
447          timestamp: new Date().toISOString(),
448          customType: "rewind-op",
449          data: { v: 2, snapshots: [liveCommit], current: 0 },
450        },
451      ]);
452  
453      await harness.invoke("session_start", {});
454  
455      // Retention sweep runs in the background on startup; poll for completion
456      const deadline = Date.now() + 3000;
457      let storeHead: string | undefined;
458      while (Date.now() < deadline) {
459        storeHead = await harness.revParseStore();
460        if (storeHead && await harness.isAncestor(liveCommit, storeHead)) break;
461        await new Promise(r => setTimeout(r, 50));
462      }
463      assert.ok(storeHead);
464      assert.equal(await harness.isAncestor(liveCommit, storeHead!), true);
465    } finally {
466      await harness.cleanup();
467    }
468  });
469