fix(run-store): hydrate a snapshot's checkpoint by id at redis-only

At redis-only the TaskRunExecutionSnapshot row is not written to Postgres, but
the checkpoint hydration read it back through that row, so a suspended run
resumed with a null checkpoint and restarted with no state to restore from.

The TaskRunCheckpoint row itself is never suppressed and the Redis entry
carries the checkpointId, so read the checkpoint directly by id instead. Adds
a residency-aware findTaskRunCheckpointById to the run store (routed to the
run's co-located store) and points the decorator's hydration at it. Covered by
a new redis-only checkpoint test that fails on the old path and passes now.

Author: d-cs

internal-packages/run-store/src/PostgresRunStore.ts

@@ -57,7 +57,7 @@ export interface RunOpsCapableClient {
   taskRunAttempt: RunOpsDelegate<"create" | "findFirst" | "findMany" | "update">;
   taskRunExecutionSnapshot: RunOpsDelegate<"create" | "findFirst" | "findMany">;
   taskRunWaitpoint: RunOpsDelegate<"deleteMany" | "findMany">;
-  taskRunCheckpoint: RunOpsDelegate<"create">;
+  taskRunCheckpoint: RunOpsDelegate<"create" | "findFirst">;
   checkpoint: RunOpsDelegate<"create" | "findFirst">;
   checkpointRestoreEvent: RunOpsDelegate<"create" | "findFirst">;
   taskRunDependency: RunOpsDelegate<"create" | "findFirst" | "findMany">;
@@ -2616,6 +2616,19 @@ export class PostgresRunStore implements RunStore {
     return prisma.taskRunCheckpoint.create(args) as Promise<Prisma.TaskRunCheckpointGetPayload<T>>;
   }
 
+  async findTaskRunCheckpointById(
+    checkpointId: string,
+    // `ownerRunId` selects residency at the router; a single store has one client and ignores it.
+    _ownerRunId: string,
+    client?: ReadClient
+  ): Promise<Prisma.TaskRunCheckpointGetPayload<{}> | null> {
+    const prisma = client ?? this.readOnlyPrisma;
+
+    return prisma.taskRunCheckpoint.findFirst({
+      where: { id: checkpointId },
+    }) as Promise<Prisma.TaskRunCheckpointGetPayload<{}> | null>;
+  }
+
   // --- BatchTaskRun (run-ops) ---
 
   async createBatchTaskRun(

internal-packages/run-store/src/delegatingRunStore.ts

@@ -623,6 +623,14 @@ export class DelegatingRunStore implements RunStore {
     return this.delegate.createTaskRunCheckpoint(args, ownerRunId, tx);
   }
 
+  findTaskRunCheckpointById(
+    checkpointId: string,
+    ownerRunId: string,
+    client?: ReadClient
+  ): Promise<Prisma.TaskRunCheckpointGetPayload<{}> | null> {
+    return this.delegate.findTaskRunCheckpointById(checkpointId, ownerRunId, client);
+  }
+
   createBatchTaskRun(
     data: CreateBatchTaskRunData,
     tx?: PrismaClientOrTransaction

internal-packages/run-store/src/runOpsStore.ts

@@ -2083,6 +2083,20 @@ export class RoutingRunStore implements RunStore {
     return store.createTaskRunCheckpoint(args, ownerRunId, undefined);
   }
 
+  async findTaskRunCheckpointById(
+    checkpointId: string,
+    ownerRunId: string,
+    client?: ReadClient
+  ): Promise<Prisma.TaskRunCheckpointGetPayload<{}> | null> {
+    // Co-located with its owner run, so route by ownerRunId and read that store's own primary.
+    const store = this.#route(ownerRunId);
+    return store.findTaskRunCheckpointById(
+      checkpointId,
+      ownerRunId,
+      RoutingRunStore.#ownPrimary(store, client)
+    );
+  }
+
   // ---------------------------------------------------------------------------
   // BatchTaskRun (run-ops). Route by id-shape: run-ops id→NEW, cuid→LEGACY.
   // ---------------------------------------------------------------------------

internal-packages/run-store/src/runStoreMethodNames.ts

@@ -68,6 +68,7 @@ export const RUN_STORE_METHOD_NAMES = [
   "deleteManyTaskRunWaitpoints",
   "findTaskRunAttempt",
   "createTaskRunCheckpoint",
+  "findTaskRunCheckpointById",
   "createBatchTaskRun",
   "updateBatchTaskRun",
   "findBatchTaskRunById",

internal-packages/run-store/src/taskRunExecutionSnapshotStore.checkpoint.test.ts

@@ -107,6 +107,76 @@ describe("a checkpoint on a snapshot served from Redis", () => {
     }
   );
 
+  // The shipping redis-only PAIR: decorator at redis-only OVER a store with snapshotWrites:false, so
+  // the TaskRunExecutionSnapshot row is NOT written to Postgres. The checkpoint ROW still lives in
+  // Postgres (only snapshot rows are suppressed), and the Redis entry carries checkpointId. A resume
+  // MUST still re-attach the checkpoint row. If hydration reads it through the (suppressed) snapshot
+  // row, the resumed run gets a null checkpoint and restarts with no state to restore from.
+  containerTest(
+    "at redis-only, a suspended run's checkpoint still hydrates (snapshot row suppressed)",
+    async ({ prisma, redisOptions }) => {
+      const redis = new RedisSnapshotStore({ redisOptions, completedTtlMs: COMPLETED_TTL_MS });
+      // snapshotWrites:false is redis-only: run mutations land, snapshot rows do not.
+      const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma, snapshotWrites: false });
+      const decorated = new TaskRunExecutionSnapshotStore(store as unknown as RunStore, {
+        store: redis,
+        mode: "redis-only",
+        readPercent: 100,
+      });
+      try {
+        const env = await seedSnapshotEnvironment(prisma);
+        const runId = generateInternalId();
+
+        await decorated.createRun({
+          data: buildCreateRunData(runId, env),
+          snapshot: birthSnapshot(env),
+        });
+
+        const checkpoint = await prisma.taskRunCheckpoint.create({
+          data: {
+            friendlyId: `checkpoint_${generateInternalId()}`,
+            type: "DOCKER",
+            location: "s3://bucket/redis-only-checkpoint.tar",
+            imageRef: "registry/image@sha256:def",
+            reason: "suspend at redis-only",
+            projectId: env.projectId,
+            runtimeEnvironmentId: env.id,
+          },
+        });
+
+        // Suspend transition naming the checkpoint. At redis-only the snapshot row is suppressed in PG.
+        await decorated.createExecutionSnapshot({
+          run: { id: runId, status: "WAITING_TO_RESUME" },
+          snapshot: { executionStatus: "SUSPENDED", description: "Run was suspended" },
+          checkpointId: checkpoint.id,
+          environmentId: env.id,
+          environmentType: env.type,
+          projectId: env.projectId,
+          organizationId: env.organizationId,
+        } as never);
+
+        // Confirm the redis-only premise: NO snapshot row for this run exists in Postgres.
+        const pgSnapshotCount = await prisma.taskRunExecutionSnapshot.count({ where: { runId } });
+        expect(pgSnapshotCount).toBe(0);
+
+        // The resume read is served from Redis. It MUST still carry the checkpoint id AND the row.
+        const fromRedis = await decorated.findLatestExecutionSnapshot(runId);
+        expect(fromRedis).not.toBeNull();
+        expect(fromRedis!.executionStatus).toBe("SUSPENDED");
+        expect(fromRedis!.checkpointId).toBe(checkpoint.id);
+
+        // The load-bearing assertion: the checkpoint ROW is re-attached, hydrated directly from the
+        // TaskRunCheckpoint table (not via the suppressed snapshot row), so the run can restore.
+        expect(fromRedis!.checkpoint).not.toBeNull();
+        expect(fromRedis!.checkpoint!.id).toBe(checkpoint.id);
+        expect(fromRedis!.checkpoint!.location).toBe("s3://bucket/redis-only-checkpoint.tar");
+        expect(fromRedis!.checkpoint!.imageRef).toBe("registry/image@sha256:def");
+      } finally {
+        await redis.quit();
+      }
+    }
+  );
+
   containerTest(
     "costs no Postgres read when the snapshot has no checkpoint",
     async ({ prisma, redisOptions }) => {

internal-packages/run-store/src/taskRunExecutionSnapshotStore.ts

@@ -1463,7 +1463,7 @@ export class TaskRunExecutionSnapshotStore extends DelegatingRunStore {
     const entry = read.entry as Record<string, unknown>;
 
     const checkpoint = entry.checkpointId
-      ? await this.#hydrateCheckpoint(runId, read.id, client)
+      ? await this.#hydrateCheckpoint(runId, entry.checkpointId as string, client)
       : null;
 
     // `completedWaitpointOrder` is a scalar column, NOT the join. The engine reads it off the head
@@ -1555,14 +1555,13 @@ export class TaskRunExecutionSnapshotStore extends DelegatingRunStore {
    */
   async #hydrateCheckpoint(
     runId: string,
-    snapshotId: string,
+    checkpointId: string,
     client?: ReadClient
   ): Promise<unknown> {
-    const row = await this.delegate.findExecutionSnapshot(
-      { where: { id: snapshotId, runId }, include: { checkpoint: true } },
-      client
-    );
-    return (row as { checkpoint?: unknown } | null)?.checkpoint ?? null;
+    // Read the checkpoint ROW directly by id, NOT via the snapshot row: at redis-only the snapshot
+    // row is suppressed, so hydrating through it returns null and a resumed run loses its checkpoint.
+    // The TaskRunCheckpoint row is never suppressed. runId routes to the run's co-located store.
+    return this.delegate.findTaskRunCheckpointById(checkpointId, runId, client);
   }
 
   async #enqueueRepair(entry: SnapshotEntryInput): Promise<void> {

internal-packages/run-store/src/types.ts

@@ -882,6 +882,15 @@ export interface RunStore {
     tx?: PrismaClientOrTransaction
   ): Promise<Prisma.TaskRunCheckpointGetPayload<T>>;
 
+  // Residency-aware direct checkpoint read by id. A snapshot served from Redis carries only the
+  // checkpointId; the checkpoint hydrates by reading TaskRunCheckpoint directly, NOT via the snapshot
+  // row, which at redis-only is suppressed. `ownerRunId` routes to the run's co-located store.
+  findTaskRunCheckpointById(
+    checkpointId: string,
+    ownerRunId: string,
+    client?: ReadClient
+  ): Promise<Prisma.TaskRunCheckpointGetPayload<{}> | null>;
+
   // --- BatchTaskRun (run-ops) ---
   // Batch row is born on the run-ops store at create. `findBatchTaskRunById`
   // reads the primary by default (worker reads the just-written row; replica lag).