fix(run-store): three more index-less waitpoint losses on the read path

The previous fix stored the complete id set but left three places still deriving
it from the ordered list, and the ordered list holds only batch-indexed ids.

A carry-forward decided on the order alone. Two DIFFERENT single waits both
present an empty order, so they compared equal, the second inherited the first's
cycle, and a read returned the wrong waitpoint entirely. The comparison now
requires the id set to match as well.

The dequeue site built its Redis refs from the ordered list while the delegate
connects the complete set in Postgres, so an index-less waitpoint reached
Postgres and never reached Redis. Refs are now built from the complete set, with
the index taken from the ordered list where the id appears in it.

The entry decode derived the set from the order too, which meant getLatest and
getById returned an incomplete set. That is the hot read: findLatestExecutionSnapshot
hydrates the waitpoint rows from it, so a resume would have fetched no row at all
for a single wait. The read scripts now return the stored set alongside the order.

Four tests, each verified against its own defect: two consecutive single waits
keep separate cycles, a repeated one still carries forward, the dequeue snapshot
keeps an index-less id, and the hot read hydrates its row.

Author: d-cs

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

@@ -457,11 +457,12 @@ export class RedisSnapshotStore {
       }
 
       const headOrder = reply[1] ?? "";
+      const headDistinct = reply[2] ?? "";
       const rows: SnapshotRead[] = [];
       // Tracks whether the Lua-chosen head row (always the first, i === 2) itself survives the
       // env filter below -- headOrder must never be attributed to a different, surviving row.
       let headSurvived = false;
-      for (let i = 2; i + 3 < reply.length; i += 4) {
+      for (let i = 3; i + 3 < reply.length; i += 4) {
         // orderKnown is false here: headOrder covers only the head row, resolved separately below.
         const decoded = this.#decode(
           [reply[i], reply[i + 1], reply[i + 2], reply[i + 3], ""],
@@ -471,13 +472,17 @@ export class RedisSnapshotStore {
         );
         if (decoded) {
           rows.push(decoded);
-          if (i === 2) headSurvived = true;
+          if (i === 3) headSurvived = true;
         }
       }
 
       rows.reverse();
       const head = headSurvived ? rows[rows.length - 1] : undefined;
-      const headWaitpointIds = decodeWaitpointIds(head !== undefined, head ? headOrder : "");
+      const headWaitpointIds = decodeWaitpointIds(
+        head !== undefined,
+        head ? headOrder : "",
+        head ? headDistinct : ""
+      );
       if (head) {
         head.completedWaitpointIds = headWaitpointIds;
         if (head.cycle) {
@@ -516,11 +521,12 @@ export class RedisSnapshotStore {
       if (reply === null) return { kind: "miss" };
 
       const headOrder = reply[1] ?? "";
+      const headDistinct = reply[2] ?? "";
       const rows: SnapshotRead[] = [];
       // Tracks whether the Lua-chosen head row (always the first, i === 2) survives the env filter,
       // so headOrder is never attributed to a different, surviving row.
       let headSurvived = false;
-      for (let i = 2; i + 3 < reply.length; i += 4) {
+      for (let i = 3; i + 3 < reply.length; i += 4) {
         const decoded = this.#decode(
           [reply[i], reply[i + 1], reply[i + 2], reply[i + 3], ""],
           opts?.environmentId,
@@ -529,13 +535,17 @@ export class RedisSnapshotStore {
         );
         if (decoded) {
           rows.push(decoded);
-          if (i === 2) headSurvived = true;
+          if (i === 3) headSurvived = true;
         }
       }
 
       rows.reverse();
       const head = headSurvived ? rows[rows.length - 1] : undefined;
-      const headWaitpointIds = decodeWaitpointIds(head !== undefined, head ? headOrder : "");
+      const headWaitpointIds = decodeWaitpointIds(
+        head !== undefined,
+        head ? headOrder : "",
+        head ? headDistinct : ""
+      );
       if (head) {
         head.completedWaitpointIds = headWaitpointIds;
         if (head.cycle) {
@@ -568,7 +578,7 @@ export class RedisSnapshotStore {
     orderKnown: boolean
   ): SnapshotRead | null {
     if (!reply || reply.length === 0) return null;
-    const [id, raw, seqStr, pointer, orderJson] = reply;
+    const [id, raw, seqStr, pointer, orderJson, distinctJson] = reply;
     const entry = JSON.parse(raw) as Record<string, unknown>;
     if (environmentId !== undefined && entry.environmentId !== environmentId) return null;
     const read: SnapshotRead = {
@@ -582,7 +592,7 @@ export class RedisSnapshotStore {
       const [cs, count] = pointer.split(":");
       read.cycle = { cycleSeq: Number(cs), count: Number(count) };
       if (orderKnown) {
-        const ids = decodeWaitpointIds(true, orderJson);
+        const ids = decodeWaitpointIds(true, orderJson, distinctJson ?? "");
         read.completedWaitpointIds = ids;
         this.#checkCycleMismatch(runId, Number(count), ids.order.length);
       }
@@ -737,7 +747,7 @@ export class RedisSnapshotStore {
         local vals = redis.call('HMGET', eKey, id, id .. '#s', id .. '#c')
         if not vals[1] then return nil end
         -- Coerce every element: a Lua false TRUNCATES the returned array at that position.
-        return { id, vals[1], vals[2] or '', vals[3] or '', orderFor(vals[3]) }
+        return { id, vals[1], vals[2] or '', vals[3] or '', orderFor(vals[3]), distinctFor(vals[3]) }
       `,
     });
 
@@ -749,7 +759,7 @@ export class RedisSnapshotStore {
         if not cur then return nil end
         local vals = redis.call('HMGET', eKey, cur, cur .. '#s', cur .. '#c')
         if not vals[1] then return nil end
-        return { cur, vals[1], vals[2] or '', vals[3] or '', orderFor(vals[3]) }
+        return { cur, vals[1], vals[2] or '', vals[3] or '', orderFor(vals[3]), distinctFor(vals[3]) }
       `,
     });
 
@@ -785,7 +795,7 @@ export class RedisSnapshotStore {
         -- compare is a chronological compare. Walking newest-first lets the scan stop at the first
         -- entry at or before the cursor, which makes its length the length of the ANSWER rather
         -- than the length of the run's history.
-        local out = { '', '' }
+        local out = { '', '', '' }
         local headId = nil
         local offset = 0
         local page = limit
@@ -820,7 +830,9 @@ export class RedisSnapshotStore {
         end
 
         if headId then
-          out[2] = orderFor(redis.call('HGET', eKey, headId .. '#c'))
+          local headPointer = redis.call('HGET', eKey, headId .. '#c')
+          out[2] = orderFor(headPointer)
+          out[3] = distinctFor(headPointer)
         end
         return out
       `,
@@ -852,7 +864,7 @@ export class RedisSnapshotStore {
         -- The head is the newest SURVIVING entry, and it is the only one whose cycle key is read.
         -- Deriving the order after the loop keeps it paired with the row it is attached to: a row
         -- dropped for a missing body must not donate its cycle data to the next one.
-        local out = { sinceRaw, '' }
+        local out = { sinceRaw, '', '' }
         local headId = nil
         for i = 1, #ids do
           local id = ids[i]
@@ -866,7 +878,9 @@ export class RedisSnapshotStore {
           end
         end
         if headId then
-          out[2] = orderFor(redis.call('HGET', eKey, headId .. '#c'))
+          local headPointer = redis.call('HGET', eKey, headId .. '#c')
+          out[2] = orderFor(headPointer)
+          out[3] = distinctFor(headPointer)
         end
         return out
       `,

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

@@ -20,7 +20,7 @@ import type {
   SnapshotEntryInput,
   SnapshotRead,
 } from "./redisSnapshotStore.js";
-import { deriveOrder } from "./redisSnapshotStore.js";
+import { deriveDistinctIds, deriveOrder } from "./redisSnapshotStore.js";
 import {
   entryFromCompletion,
   entryFromCreateExecutionSnapshot,
@@ -384,9 +384,13 @@ export class TaskRunExecutionSnapshotStore extends DelegatingRunStore {
       "lockRunToWorker",
       entryFromLock(ctx, withStamp.snapshot),
       withStamp.snapshot.previousSnapshotId,
-      // The lock site already carries the resolved order, so the refs are rebuilt from it rather
-      // than re-derived: its index IS the position in that list.
-      withStamp.snapshot.completedWaitpointOrder.map((id, index) => ({ id, index }))
+      // Built from the COMPLETE id set, which is what the delegate connects in Postgres, with the
+      // index taken from the ordered list where the id appears in it. Building from the ordered list
+      // instead would drop every id with no batch index, exactly the ids Postgres still records.
+      lockCycleRefs(
+        withStamp.snapshot.completedWaitpointIds,
+        withStamp.snapshot.completedWaitpointOrder
+      )
     );
     return result;
   }
@@ -579,12 +583,21 @@ export class TaskRunExecutionSnapshotStore extends DelegatingRunStore {
     }
 
     const order = deriveOrder(completedWaitpoints);
+    const distinct = deriveDistinctIds(completedWaitpoints);
 
     try {
       const head = await this.redis.getLatest(runId);
-      const previous = head?.completedWaitpointIds?.order;
-
-      if (head?.cycle && previous && sameOrder(previous, order)) {
+      const previousIds = head?.completedWaitpointIds;
+
+      // Both halves must match. Comparing the order alone is not enough: it holds only indexed ids,
+      // so two DIFFERENT single waits both present an empty order and would compare equal, and the
+      // second would inherit the first's waitpoint set instead of minting its own.
+      if (
+        head?.cycle &&
+        previousIds &&
+        sameOrder(previousIds.order, order) &&
+        sameSet(previousIds.distinctIds, distinct)
+      ) {
         return { kind: "carryForward", cycleSeq: head.cycle.cycleSeq };
       }
     } catch (error) {
@@ -899,3 +912,31 @@ export class TaskRunExecutionSnapshotStore extends DelegatingRunStore {
 function sameOrder(a: string[], b: string[]): boolean {
   return a.length === b.length && a.every((id, index) => id === b[index]);
 }
+
+/**
+ * Turns the lock site's two lists into cycle refs. `completedWaitpointIds` is the complete set the
+ * delegate connects; `completedWaitpointOrder` gives a position only to the ids that have one, and a
+ * repeated id keeps each of its positions.
+ */
+function lockCycleRefs(ids: string[], order: string[]): { id: string; index?: number }[] {
+  const refs: { id: string; index?: number }[] = [];
+  const indexed = new Set<string>();
+
+  order.forEach((id, index) => {
+    refs.push({ id, index });
+    indexed.add(id);
+  });
+
+  for (const id of ids) {
+    if (!indexed.has(id)) refs.push({ id });
+  }
+
+  return refs;
+}
+
+/** Membership only, for the id set, which has no meaningful order. */
+function sameSet(a: string[], b: string[]): boolean {
+  if (a.length !== b.length) return false;
+  const seen = new Set(a);
+  return b.every((id) => seen.has(id));
+}

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

@@ -305,6 +305,124 @@ describe("completed-waitpoint cycles", () => {
     }
   );
 
+  containerTest(
+    "two consecutive index-less waits do not share a cycle",
+    async ({ prisma, redisOptions }) => {
+      const { decorated, redis } = build(prisma as never, redisOptions as never);
+      const probe = createRedisClient(redisOptions, { onError: () => {} });
+      try {
+        const env = await seedSnapshotEnvironment(prisma);
+        const runId = await seedRun(decorated, redis, env);
+        const [wpA, wpB] = await seedSnapshotWaitpoints(prisma, env, 2);
+
+        // Neither wait has a batch index, so both present an EMPTY order. Deciding carry-forward on
+        // the order alone makes them compare equal, and the second silently inherits the first's
+        // waitpoint set: its own result is never stored and a read returns the wrong id.
+        const first = await decorated.createExecutionSnapshot(
+          resumeInput(runId, env, [{ id: wpA }], "first single wait")
+        );
+        const second = await decorated.createExecutionSnapshot(
+          resumeInput(runId, env, [{ id: wpB }], "second single wait")
+        );
+
+        expect((await redis.getSnapshotWaitpointIds(runId, first.id)).distinctIds).toEqual([wpA]);
+        expect((await redis.getSnapshotWaitpointIds(runId, second.id)).distinctIds).toEqual([wpB]);
+        expect((await probe.keys(`snap:{${runId}}:wp:*`)).length).toBe(2);
+      } finally {
+        await Promise.all([redis.quit(), probe.quit().catch(() => {})]);
+      }
+    }
+  );
+
+  containerTest(
+    "the same index-less wait repeated does still carry forward",
+    async ({ prisma, redisOptions }) => {
+      const { decorated, redis } = build(prisma as never, redisOptions as never);
+      const probe = createRedisClient(redisOptions, { onError: () => {} });
+      try {
+        const env = await seedSnapshotEnvironment(prisma);
+        const runId = await seedRun(decorated, redis, env);
+        const [wpA] = await seedSnapshotWaitpoints(prisma, env, 1);
+
+        // The copy-forward case must survive the stricter comparison: the same id set, still one key.
+        await decorated.createExecutionSnapshot(resumeInput(runId, env, [{ id: wpA }], "wait"));
+        const carried = await decorated.createExecutionSnapshot(
+          resumeInput(runId, env, [{ id: wpA }], "carry")
+        );
+
+        expect((await probe.keys(`snap:{${runId}}:wp:*`)).length).toBe(1);
+        expect((await redis.getSnapshotWaitpointIds(runId, carried.id)).distinctIds).toEqual([wpA]);
+      } finally {
+        await Promise.all([redis.quit(), probe.quit().catch(() => {})]);
+      }
+    }
+  );
+
+  containerTest(
+    "the dequeue snapshot keeps an index-less waitpoint",
+    async ({ prisma, redisOptions }) => {
+      const { decorated, redis } = build(prisma as never, redisOptions as never);
+      try {
+        const env = await seedSnapshotEnvironment(prisma);
+        const runId = await seedRun(decorated, redis, env);
+        const [wpA, wpB] = await seedSnapshotWaitpoints(prisma, env, 2);
+        const head = await redis.getLatest(runId);
+        const snapshotId = generateInternalId();
+
+        // Postgres connects completedWaitpointIds, the COMPLETE set. Building the Redis refs from
+        // completedWaitpointOrder instead drops every id that has no position in it.
+        await decorated.lockRunToWorker(runId, {
+          lockedAt: new Date(),
+          lockedById: undefined,
+          lockedToVersionId: undefined,
+          lockedQueueId: undefined,
+          startedAt: new Date(),
+          baseCostInCents: 0,
+          machinePreset: "small-1x",
+          taskVersion: "1.0.0",
+          snapshot: {
+            id: snapshotId,
+            previousSnapshotId: head!.id,
+            attemptNumber: 1,
+            environmentId: env.id,
+            environmentType: env.type,
+            projectId: env.projectId,
+            organizationId: env.organizationId,
+            completedWaitpointIds: [wpA, wpB],
+            completedWaitpointOrder: [wpA],
+          },
+        } as never);
+
+        const ids = await redis.getSnapshotWaitpointIds(runId, snapshotId);
+        expect([...ids.distinctIds].sort()).toEqual([wpA, wpB].sort());
+        expect(ids.order).toEqual([wpA]);
+      } finally {
+        await redis.quit();
+      }
+    }
+  );
+
+  containerTest(
+    "findLatestExecutionSnapshot hydrates an index-less waitpoint row",
+    async ({ prisma, redisOptions }) => {
+      const { decorated, redis } = build(prisma as never, redisOptions as never);
+      try {
+        const env = await seedSnapshotEnvironment(prisma);
+        const runId = await seedRun(decorated, redis, env);
+        const [wpA] = await seedSnapshotWaitpoints(prisma, env, 1);
+
+        await decorated.createExecutionSnapshot(resumeInput(runId, env, [{ id: wpA }], "single"));
+
+        // The hot read hydrates the rows from the id set, so an incomplete set means the resume
+        // gets no waitpoint at all.
+        const latest = await decorated.findLatestExecutionSnapshot(runId);
+        expect(latest!.completedWaitpoints.map((w) => w.id)).toEqual([wpA]);
+      } finally {
+        await redis.quit();
+      }
+    }
+  );
+
   containerTest(
     "findLatestExecutionSnapshot returns the index oracle",
     async ({ prisma, redisOptions }) => {