fix(checkpoints): a revert never applies one checkout's snapshots to another

apps/server/src/checkpointing/Layers/CheckpointRevert.test.ts

@@ -8,6 +8,7 @@ import {
   ThreadCheckoutMovedActivityKind,
   type ThreadCheckoutMovedPayload,
   ThreadId,
+  TurnId,
 } from "@threadlines/contracts";
 import * as Effect from "effect/Effect";
 import * as Layer from "effect/Layer";
@@ -25,7 +26,8 @@ import {
 import { CheckpointRevert } from "../Services/CheckpointRevert.ts";
 import { CheckpointStore } from "../Services/CheckpointStore.ts";
 import { makeCheckpointStoreStub } from "../testing/CheckpointStoreStub.ts";
-import { checkpointRefForThreadTurn } from "../Utils.ts";
+import type { VcsCheckpointCheckout } from "../../vcs/VcsDriver.ts";
+import { checkpointPreTurnRefForThreadTurnCount, checkpointRefForThreadTurn } from "../Utils.ts";
 import { CheckpointRevertLive } from "./CheckpointRevert.ts";
 
 const PROJECT = ProjectId.make("project-1");
@@ -45,14 +47,36 @@ const thread = (overrides: Partial<OrchestrationThread> = {}) =>
     worktreePath: worktree,
     activities: [],
     checkpoints: [1, 2, 3].map((turnCount) => ({
+      turnId: TurnId.make(`turn-${turnCount}`),
       checkpointTurnCount: turnCount,
       checkpointRef: checkpointRefForThreadTurn(THREAD, turnCount),
     })),
     ...overrides,
   }) as unknown as OrchestrationThread;
 
-/** Where a revert to `turnCount` would work, and how, for a thread no one else is near. */
-const resolve = (subject: OrchestrationThread, turnCount: number) =>
+/** A turn's completion snapshot, and the one taken right before it started. */
+const done = (turn: number) => checkpointRefForThreadTurn(THREAD, turn);
+const start = (turn: number) => checkpointPreTurnRefForThreadTurnCount(THREAD, turn);
+
+/** Every snapshot of these turns, as taken in the folder the thread has since left. */
+const leftBehind = (...turns: ReadonlyArray<number>) =>
+  Object.fromEntries(
+    turns.flatMap((turn) => [
+      [done(turn), "elsewhere"],
+      [start(turn), "elsewhere"],
+    ]),
+  ) as Record<string, VcsCheckpointCheckout>;
+
+/**
+ * Where a revert to `turnCount` would work, and how, for a thread no one else
+ * is near. `takenIn` says where a snapshot was taken; any other one was taken
+ * here, in a thread whose snapshots all record their checkout.
+ */
+const resolve = (
+  subject: OrchestrationThread,
+  turnCount: number,
+  takenIn: Record<string, VcsCheckpointCheckout> | "unrecorded" = {},
+) =>
   Effect.gen(function* () {
     const revert = yield* CheckpointRevert;
     return yield* revert.resolveContext({ threadId: THREAD, turnCount });
@@ -61,7 +85,21 @@ const resolve = (subject: OrchestrationThread, turnCount: number) =>
       CheckpointRevertLive.pipe(
         Layer.provide(
           Layer.mergeAll(
-            Layer.succeed(CheckpointStore, makeCheckpointStoreStub()),
+            Layer.succeed(
+              CheckpointStore,
+              makeCheckpointStoreStub({
+                checkpointCheckouts: (input) =>
+                  Effect.succeed(
+                    input.checkpointRefs.map((ref) =>
+                      takenIn === "unrecorded"
+                        ? "unknown"
+                        : (takenIn[ref] ??
+                          // The older per-turn-id refs don't exist on these threads.
+                          (ref.includes("pre-turn/") ? "unknown" : "here")),
+                    ),
+                  ),
+              }),
+            ),
             Layer.succeed(ProjectionSnapshotQuery, {
               getThreadDetailById: () => Effect.succeed(Option.some(subject)),
               getProjectShellById: () =>
@@ -78,11 +116,16 @@ const resolve = (subject: OrchestrationThread, turnCount: number) =>
     Effect.runPromise,
   );
 
+const movedSince = {
+  kind: "unavailable",
+  detail: expect.stringContaining("moved to another checkout"),
+};
+
 describe("CheckpointRevert.resolveContext", () => {
   it("restores the whole folder only for a thread whose worktree is its own", async () => {
     expect(await resolve(thread(), 1)).toMatchObject({
       kind: "ready",
-      context: { mode: "workspace" },
+      context: { mode: "workspace", targetCheckpointRef: done(1) },
     });
     // A guest in another thread's worktree never does, even with that thread
     // gone and nothing else pointing at the folder.
@@ -93,7 +136,31 @@ describe("CheckpointRevert.resolveContext", () => {
     });
   });
 
-  it("never reaches back past the point an agent moved the thread into another checkout", async () => {
+  it("never undoes a turn whose snapshots were taken in another checkout", async () => {
+    // Turns one and two ran in the folder the thread has since left.
+    expect(await resolve(thread(), 1, leftBehind(1, 2))).toMatchObject(movedSince);
+    expect(await resolve(thread(), 0, leftBehind(1, 2))).toMatchObject(movedSince);
+    // A starting snapshot left over from the other folder spoils its turn too.
+    expect(await resolve(thread(), 2, { [start(3)]: "elsewhere" })).toMatchObject(movedSince);
+  });
+
+  it("undoes everything since the move from this folder's own snapshot, never the other folder's", async () => {
+    // Turn three is the only one that ran here. Reverting to turn two would
+    // restore from turn two's snapshot, which holds the other folder's files:
+    // the files come from this folder as it was right before turn three.
+    expect(await resolve(thread(), 2, leftBehind(1, 2))).toMatchObject({
+      kind: "ready",
+      context: { targetCheckpointRef: start(3) },
+    });
+    // With no snapshot of this folder from before that turn, nothing is safe to restore from.
+    expect(
+      await resolve(thread(), 2, { ...leftBehind(1, 2), [start(3)]: "unknown" }),
+    ).toMatchObject(movedSince);
+    // And reverting to the newest turn undoes nothing, so there is no such snapshot at all.
+    expect(await resolve(thread(), 3, leftBehind(3))).toMatchObject(movedSince);
+  });
+
+  it("falls back to the move on the thread's record only for snapshots that don't say where they were taken", async () => {
     const moved = thread({
       worktreeOrigin: { kind: "thread", threadId: CODER },
       activities: [
@@ -109,13 +176,10 @@ describe("CheckpointRevert.resolveContext", () => {
         },
       ] as never,
     });
-    // Turns one and two changed the folder it left.
-    const before = await resolve(moved, 1);
-    expect(before.kind).toBe("unavailable");
-    expect(before).toMatchObject({
-      detail: expect.stringContaining("moved into another checkout"),
-    });
-    // Undoing only what it did since the move stays possible.
-    expect((await resolve(moved, 2)).kind).toBe("ready");
+    // Older snapshots: turns one and two changed the folder it left.
+    expect(await resolve(moved, 1, "unrecorded")).toMatchObject(movedSince);
+    expect((await resolve(moved, 2, "unrecorded")).kind).toBe("ready");
+    // Snapshots that all say "here" outrank the record: it moved away and came back.
+    expect((await resolve(moved, 1)).kind).toBe("ready");
   });
 });

apps/server/src/checkpointing/Layers/CheckpointRevert.ts

@@ -208,15 +208,65 @@ const make = Effect.gen(function* () {
         };
       }
 
-      // Turns from before an agent moved it into another checkout changed the
-      // folder it left: undoing them here would rewrite files it never touched.
-      const movedAtTurnCount = turnCountAtLastCheckoutMove(thread);
-      if (movedAtTurnCount !== null && input.turnCount < movedAtTurnCount) {
-        return {
-          kind: "unavailable" as const,
-          detail:
-            "This thread moved into another checkout after that point. Its earlier changes are in the folder it left, so they can't be undone from here.",
-        };
+      // A thread can move to another checkout between turns (the user's own
+      // switch, an agent's move, a session that settled somewhere else), and
+      // every checkout of the repository shares these snapshots. Each one
+      // says where it was taken. A turn whose snapshots come from another
+      // checkout changed that folder: undoing it here would rewrite files it
+      // never touched, so such a revert is refused.
+      const undone = thread.checkpoints
+        .filter((checkpoint) => checkpoint.checkpointTurnCount > input.turnCount)
+        .toSorted((left, right) => left.checkpointTurnCount - right.checkpointTurnCount);
+      // Per turn: its completion snapshot, then the two refs its starting
+      // snapshot can be under (by count, or by turn id on older threads).
+      const snapshotsOf = (checkpoint: RevertCheckpointContext) =>
+        [
+          checkpoint.checkpointRef,
+          checkpointPreTurnRefForThreadTurnCount(input.threadId, checkpoint.checkpointTurnCount),
+          checkpointPreTurnRefForThreadTurn(input.threadId, checkpoint.turnId),
+        ] as const;
+      const [targetTakenIn = "unknown", ...undoneTakenIn] =
+        yield* checkpointStore.checkpointCheckouts({
+          cwd,
+          checkpointRefs: [targetCheckpointRef, ...undone.flatMap(snapshotsOf)],
+        });
+      const takenIn = (index: number, snapshot: 0 | 1 | 2) =>
+        undoneTakenIn[index * 3 + snapshot] ?? "unknown";
+      // Snapshots from before they recorded their checkout say nothing. For
+      // turns like that, the move an agent made is still on the thread's record.
+      const unrecorded = undone.some((_, index) => takenIn(index, 0) === "unknown");
+      const movedAtTurnCount = unrecorded ? turnCountAtLastCheckoutMove(thread) : null;
+      const movedSince = {
+        kind: "unavailable" as const,
+        detail:
+          "This thread moved to another checkout after that point. Its earlier changes are in the folder it left, so they can't be undone from here.",
+      };
+      if (
+        undoneTakenIn.includes("elsewhere") ||
+        (movedAtTurnCount !== null && input.turnCount < movedAtTurnCount)
+      ) {
+        return movedSince;
+      }
+
+      // The target is what files are restored from. One taken in another
+      // checkout (the turn the thread moved after) holds that folder's files;
+      // this folder as it was right before the first turn being undone is
+      // that turn's own starting snapshot, so the revert restores from it.
+      let revertTargetRef = targetCheckpointRef;
+      if (targetTakenIn === "elsewhere") {
+        const first = undone[0];
+        const start =
+          first === undefined || first.checkpointTurnCount !== input.turnCount + 1
+            ? undefined
+            : takenIn(0, 1) === "here"
+              ? snapshotsOf(first)[1]
+              : takenIn(0, 2) === "here"
+                ? snapshotsOf(first)[2]
+                : undefined;
+        if (start === undefined) {
+          return movedSince;
+        }
+        revertTargetRef = start;
       }
 
       const isolated = yield* isIsolatedWorkspaceForThread({
@@ -236,7 +286,7 @@ const make = Effect.gen(function* () {
           cwd,
           hasProviderSession: sessionCwd !== undefined,
           mode: isolated ? ("workspace" as const) : ("selective" as const),
-          targetCheckpointRef,
+          targetCheckpointRef: revertTargetRef,
           thread,
         },
       };
@@ -252,6 +302,8 @@ const make = Effect.gen(function* () {
     readonly checkpoints: ReadonlyArray<RevertCheckpointContext>;
     readonly cwd: string;
     readonly targetTurnCount: number;
+    /** What the revert restores from: where the first turn's gap starts. */
+    readonly targetCheckpointRef: CheckpointRef;
     readonly currentTurnCount: number;
   }) {
     const contested = new Set<string>();
@@ -285,8 +337,10 @@ const make = Effect.gen(function* () {
       }
 
       const previousCheckpointRef =
-        input.checkpoints.find((entry) => entry.checkpointTurnCount === turnCount - 1)
-          ?.checkpointRef ?? checkpointRefForThreadTurn(input.threadId, turnCount - 1);
+        turnCount === input.targetTurnCount + 1
+          ? input.targetCheckpointRef
+          : (input.checkpoints.find((entry) => entry.checkpointTurnCount === turnCount - 1)
+              ?.checkpointRef ?? checkpointRefForThreadTurn(input.threadId, turnCount - 1));
       const previousCommit = yield* checkpointStore.resolveCheckpointCommit({
         cwd: input.cwd,
         checkpointRef: previousCheckpointRef,
@@ -367,6 +421,7 @@ const make = Effect.gen(function* () {
       checkpoints: thread.checkpoints,
       cwd: context.cwd,
       targetTurnCount: context.turnCount,
+      targetCheckpointRef: context.targetCheckpointRef,
       currentTurnCount: context.currentTurnCount,
     });
 

apps/server/src/checkpointing/Layers/CheckpointStore.test.ts

@@ -199,6 +199,63 @@ it.layer(TestLayer)("CheckpointStoreLive", (it) => {
     );
   });
 
+  describe("checkpointCheckouts", () => {
+    it.effect("says which checkout of the repository each snapshot was taken in", () =>
+      Effect.gen(function* () {
+        const tmp = yield* makeTmpDir();
+        const worktree = path.join(yield* makeTmpDir("checkpoint-store-worktree-"), "feature");
+        yield* initRepoWithCommit(tmp);
+        yield* git(tmp, ["worktree", "add", "-b", "feature", worktree]);
+        const checkpointStore = yield* CheckpointStore;
+        const threadId = ThreadId.make("thread-two-checkouts");
+        const inProject = checkpointRefForThreadTurn(threadId, 1);
+        const inWorktree = checkpointRefForThreadTurn(threadId, 2);
+        const older = checkpointRefForThreadTurn(threadId, 3);
+        const gone = checkpointRefForThreadTurn(threadId, 4);
+
+        // A thread that worked in the project folder, then moved to the worktree.
+        yield* checkpointStore.captureCheckpoint({ cwd: tmp, checkpointRef: inProject });
+        // Taken from a folder inside it: a checkout is its root, however deep the cwd.
+        yield* git(worktree, ["mv", "README.md", "docs.md"]);
+        const fileSystem = yield* FileSystem.FileSystem;
+        yield* fileSystem.makeDirectory(path.join(worktree, "src"));
+        yield* checkpointStore.captureCheckpoint({
+          cwd: path.join(worktree, "src"),
+          checkpointRef: inWorktree,
+        });
+        // A snapshot from before they recorded their checkout.
+        const tree = yield* git(tmp, ["rev-parse", "HEAD^{tree}"]);
+        const legacy = yield* git(tmp, [
+          "commit-tree",
+          tree,
+          "-m",
+          `threadlines checkpoint ref=${older}`,
+        ]);
+        yield* git(tmp, ["update-ref", older, legacy]);
+
+        const checkpointRefs = [inProject, inWorktree, older, gone];
+        expect(yield* checkpointStore.checkpointCheckouts({ cwd: tmp, checkpointRefs })).toEqual([
+          "here",
+          "elsewhere",
+          "unknown",
+          "unknown",
+        ]);
+        // The same snapshots, asked from the worktree it moved to.
+        expect(
+          yield* checkpointStore.checkpointCheckouts({ cwd: worktree, checkpointRefs }),
+        ).toEqual(["elsewhere", "here", "unknown", "unknown"]);
+
+        // A checkout is still itself after its folder is renamed: its
+        // snapshots stay revertible there.
+        const renamed = path.join(path.dirname(worktree), "renamed");
+        yield* git(tmp, ["worktree", "move", worktree, renamed]);
+        expect(
+          yield* checkpointStore.checkpointCheckouts({ cwd: renamed, checkpointRefs }),
+        ).toEqual(["elsewhere", "here", "unknown", "unknown"]);
+      }),
+    );
+  });
+
   describe("diffCheckpoints", () => {
     it.effect("returns full oversized checkpoint diffs without truncation", () =>
       Effect.gen(function* () {

apps/server/src/checkpointing/Layers/CheckpointStore.ts

@@ -69,6 +69,13 @@ const makeCheckpointStore = Effect.gen(function* () {
     return yield* checkpoints.resolveCheckpointCommit(input);
   });
 
+  const checkpointCheckouts: CheckpointStoreShape["checkpointCheckouts"] = Effect.fn(
+    "checkpointCheckouts",
+  )(function* (input) {
+    const checkpoints = yield* resolveCheckpoints("CheckpointStore.checkpointCheckouts", input.cwd);
+    return yield* checkpoints.checkpointCheckouts(input);
+  });
+
   const diffCheckpoints: CheckpointStoreShape["diffCheckpoints"] = Effect.fn("diffCheckpoints")(
     function* (input) {
       const checkpoints = yield* resolveCheckpoints("CheckpointStore.diffCheckpoints", input.cwd);
@@ -156,6 +163,7 @@ const makeCheckpointStore = Effect.gen(function* () {
     hasCheckpointRef,
     restoreCheckpoint,
     resolveCheckpointCommit,
+    checkpointCheckouts,
     diffCheckpoints,
     listHeadMovementPaths,
     diffCheckpointAgainstHead,

apps/server/src/checkpointing/Services/CheckpointStore.ts

@@ -14,7 +14,11 @@ import * as Context from "effect/Context";
 import type * as Effect from "effect/Effect";
 
 import type { CheckpointStoreError } from "../Errors.ts";
-import type { VcsReadRevisionFileInput, VcsRevisionFile } from "../../vcs/VcsDriver.ts";
+import type {
+  VcsCheckpointCheckout,
+  VcsReadRevisionFileInput,
+  VcsRevisionFile,
+} from "../../vcs/VcsDriver.ts";
 import { CheckpointRef } from "@threadlines/contracts";
 
 export interface CaptureCheckpointInput {
@@ -54,6 +58,11 @@ export interface CheckpointRefInput {
   readonly checkpointRef: CheckpointRef;
 }
 
+export interface CheckpointCheckoutsInput {
+  readonly cwd: string;
+  readonly checkpointRefs: ReadonlyArray<CheckpointRef>;
+}
+
 /** Line counts for one repository-root-relative path. */
 export interface CheckpointFileStat {
   readonly path: string;
@@ -183,6 +192,17 @@ export interface CheckpointStoreShape {
     input: ResolveCheckpointCommitInput,
   ) => Effect.Effect<string | null, CheckpointStoreError>;
 
+  /**
+   * Where each snapshot was taken, seen from the checkout `cwd` is in, one
+   * answer per ref in order: `here`, `elsewhere` (another checkout of the
+   * repository: a thread that moved since), or `unknown` (a snapshot from
+   * before they recorded it, or a ref that is gone). A revert never applies
+   * a snapshot from elsewhere to this checkout.
+   */
+  readonly checkpointCheckouts: (
+    input: CheckpointCheckoutsInput,
+  ) => Effect.Effect<ReadonlyArray<VcsCheckpointCheckout>, CheckpointStoreError>;
+
   /**
    * Compute patch diff between two checkpoint refs.
    *