refactor(run-store): hold RoutingRunStore's stores in a keyed shard map

RoutingRunStore held two named store fields and took its routing policy
from the order the statements happened to run in. It now holds a
Map<ShardKey, RunStore>, and the three policies are readable data:

- #probeOrder — sequential probe for a lookup with no routable id; the
  first non-null result wins and the last entry owns the not-found throw
- #precedence — ascending authority for a merge, so the highest-authority
  shard wins a duplicate id
- #idlessRouteShard / #idlessWaitpointShard — the two id-less defaults,
  which differ by role and were previously two unrelated literals

The two orders are the reverse of one another, which is why they are
separate fields: some sites observe the result-array order, others only
the winner of a duplicate.

Four helpers absorb the twenty-six hand-written fan-outs (#probeFirst,
#fanOut, #fanOutPartitioned, #shardsExcept), and #shardKeyOf replaces the
inline residency-to-store ternaries. The constructor keeps its exact
options type, so every existing call site compiles unchanged and the
@ts-expect-error onLegacyRead lock keeps firing.

One behaviour changes. findManyTaskRunWaitpoints merged its edge rows
NEW-first into a last-wins dedupe, so a duplicate edge id resolved to the
LEGACY row — the opposite of the rule every sibling merge follows, and of
what dedupeEdgesById's own comment claimed. It now resolves NEW-wins, with
a test pinning it. The case needs one edge id present on both stores at
once, which is the drain-mirror window; the drain seam is removed, so it
is unreachable in practice.

Also refreshes the sixteen stale runOpsStore.ts line references in
runOpsStore.mixedResidency.test.ts.

Author: d-cs

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

@@ -195,10 +195,10 @@ async function seedSharedEnv(prisma14: PrismaClient, suffix: string) {
 }
 
 describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id #new coexisting)", () => {
-  // ── Case 1: findRuns by a MIXED bounded id-set (#findRunsByIdSet, runOpsStore.ts:294) ──
+  // ── Case 1: findRuns by a MIXED bounded id-set (#findRunsByIdSet, runOpsStore.ts:449) ──
   // A list-hydrate id set spans cuid (legacy) + run-ops id (new) ids plus a run-ops id absent from legacy.
   // Both resident runs returned; take/skip applied GLOBALLY post-merge; orderBy honored; the absent
-  // run-ops id short-circuits (never probed on LEGACY, :309).
+  // run-ops id short-circuits (never probed on LEGACY, #fanOutPartitioned, :165).
   heteroRunOpsPostgresTest(
     "case 1: findRuns by a mixed id-set returns both DBs' runs, ordered, take/skip global",
     async ({ prisma14, prisma17 }) => {
@@ -306,7 +306,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 2: findRuns by an OPEN predicate (#findRunsOpen, runOpsStore.ts:319) ──
+  // ── Case 2: findRuns by an OPEN predicate (#findRunsOpen, runOpsStore.ts:464) ──
   // No id set → query BOTH stores, union, dedup by id (NEW wins). Filter by a shared scalar
   // (runtimeEnvironmentId + status) that matches rows on both DBs.
   heteroRunOpsPostgresTest(
@@ -355,7 +355,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 3: expireRunsBatch with a MIXED id list (runOpsStore.ts:474) ──
+  // ── Case 3: expireRunsBatch with a MIXED id list (runOpsStore.ts:723) ──
   // Partitions run-ops id→NEW / cuid→LEGACY; each leg called only when non-empty; counts summed; each row
   // updated on its OWN DB only.
   heteroRunOpsPostgresTest(
@@ -388,7 +388,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 4: clearIdempotencyKey fan-out arm (byFriendlyIds, runOpsStore.ts:358) ──
+  // ── Case 4: clearIdempotencyKey fan-out arm (byFriendlyIds, runOpsStore.ts:590) ──
   // byFriendlyIds spans mixed residency → fan out to both, sum the count, each row cleared on its home.
   heteroRunOpsPostgresTest(
     "case 4: clearIdempotencyKey byFriendlyIds clears across both DBs and sums the count",
@@ -429,7 +429,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 5: countPendingWaitpoints scattered across both DBs (runOpsStore.ts:731) ──
+  // ── Case 5: countPendingWaitpoints scattered across both DBs (runOpsStore.ts:1146) ──
   // A run's pending waitpoints can be split across both stores mid-drain → count on each and sum.
   heteroRunOpsPostgresTest(
     "case 5: countPendingWaitpoints sums PENDING waitpoints scattered across both DBs",
@@ -465,7 +465,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 6: findManyWaitpoints { id: { in: [...mixed...] } } (runOpsStore.ts:793) ──
+  // ── Case 6: findManyWaitpoints { id: { in: [...mixed...] } } (runOpsStore.ts:1326) ──
   // Merge waitpoints from both DBs for a mixed id set.
   heteroRunOpsPostgresTest(
     "case 6: findManyWaitpoints merges a mixed id set from both DBs",
@@ -495,7 +495,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
 
   // ── Case 8: findExecutionSnapshot / findManyExecutionSnapshots OPEN (no runId) where ──
   // A by-snapshot-id-only lookup (snapshot ids are non-classifiable cuids) must fan out NEW→LEGACY
-  // (findExecutionSnapshot, :675) / merge both (findManyExecutionSnapshots, :688). Seed a snapshot on
+  // (findExecutionSnapshot, :1008) / merge both (findManyExecutionSnapshots, :1023). Seed a snapshot on
   // EACH DB (one run-ops run on #new, one cuid run on #legacy) and read with a no-runId where.
   heteroRunOpsPostgresTest(
     "case 8: findExecutionSnapshot/findManyExecutionSnapshots with an open where reach both DBs",
@@ -544,7 +544,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 9a: findRun with an UNCLASSIFIABLE where (spanId) on a #legacy run (#findRunUnrouted, :213) ──
+  // ── Case 9a: findRun with an UNCLASSIFIABLE where (spanId) on a #legacy run (#findRunUnrouted, :365) ──
   // A run-ops run on #new and a cuid run on #legacy each carry a distinct spanId. A spanId where can't
   // be id-classified → fan out NEW-first then LEGACY. The legacy-resident run must be found.
   heteroRunOpsPostgresTest(
@@ -586,7 +586,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 9b: findRunOrThrow with an UNCLASSIFIABLE where (spanId) on a #legacy run (:593) ──
+  // ── Case 9b: findRunOrThrow with an UNCLASSIFIABLE where (spanId) on a #legacy run (#findRunOrThrowUnrouted, :864) ──
   // The throwing twin must match findRun's fan-out: an unclassifiable where whose only matching run
   // lives on #legacy must NOT throw. A NEW-only fallback would miss the legacy run and throw.
   heteroRunOpsPostgresTest(
@@ -628,7 +628,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 7: findManyTaskRunWaitpoints with edges whose relations STRADDLE DBs (runOpsStore.ts:876) ──
+  // ── Case 7: findManyTaskRunWaitpoints with edges whose relations STRADDLE DBs (runOpsStore.ts:1573) ──
   // An edge co-locates with its RUN, but its `waitpoint`/`taskRun` relations can live on the OTHER DB
   // (a cuid token blocking a run-ops run, and vice versa). The per-leg scalar query is stripped of the
   // relation keys; the router re-hydrates `waitpoint`/`taskRun` across BOTH DBs. Exercises BOTH
@@ -721,7 +721,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 7b: the "blocking waitpoint not found on either DB" HARD ERROR (runOpsStore.ts:917) ──
+  // ── Case 7b: the "blocking waitpoint not found on either DB" HARD ERROR (#hydrateEdgeWaitpointsCrossDb, runOpsStore.ts:1637) ──
   // An edge whose `waitpointId` resolves on NEITHER DB must throw rather than leave a null status that
   // would strand (hang) or wrongly unblock the run.
   heteroRunOpsPostgresTest(
@@ -748,7 +748,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 10: findBatchTaskRunById / findBatchTaskRunByFriendlyId NEW-then-LEGACY probe (:1124,:1137) ──
+  // ── Case 10: findBatchTaskRunById / findBatchTaskRunByFriendlyId NEW-then-LEGACY probe (:1755, :1768) ──
   // A batch resident on #legacy AND a run-ops-id batch landed on #new (the control-plane window mints
   // cuid ids, but a run-ops batch resides on #new) are BOTH found via the probe, regardless of id-shape.
   heteroRunOpsPostgresTest(
@@ -792,7 +792,7 @@ describe("RoutingRunStore — mixed-residency matrix (cuid #legacy + run-ops id
     }
   );
 
-  // ── Case 11a: updateManyWaitpoints with a NO-ID (batch) where fans out to both and sums (:822) ──
+  // ── Case 11a: updateManyWaitpoints with a NO-ID (batch) where fans out to both and sums (:1518) ──
   // A batch where (no single routable id, e.g. completedByTaskRunId IS NULL + status PENDING) must
   // apply on BOTH DBs and sum the count.
   heteroRunOpsPostgresTest(

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

@@ -0,0 +1,180 @@
+import { describe, expect, it } from "vitest";
+import { RoutingRunStore } from "./runOpsStore.js";
+import type { ReadClient, RunStore } from "./types.js";
+
+// Pins the routing ALGEBRA: probe order, merge precedence, and the two id-less fallbacks that
+// differ by role. No DB — each slot is a fake RunStore recording into ONE shared ordered log, so a
+// sequential probe's leg order and a merge's winner are both observable. Real two-DB topology stays
+// with the heteroRunOpsPostgresTest suites.
+//
+// MUST NOT assert invocation order for a PARALLEL fan-out: both legs are issued before either
+// resolves, so the order they are created in is not a behaviour.
+
+type Slot = "new" | "legacy";
+
+type Call = { slot: Slot; method: string };
+
+type FakeConfig = {
+  // Rows this store returns from findRun / findRuns / findRunOrThrow, regardless of filter.
+  runs?: Array<Record<string, unknown>>;
+  // Edge rows this store returns from findManyTaskRunWaitpoints, regardless of filter.
+  edges?: Array<Record<string, unknown>>;
+  // Waitpoint rows this store returns from findWaitpoint, regardless of filter.
+  waitpoint?: Record<string, unknown> | null;
+};
+
+type FakeStore = RunStore & {
+  slot: Slot;
+  primaryReadClient: { __primary: Slot };
+};
+
+function fakeStore(slot: Slot, log: Call[], config: FakeConfig = {}): FakeStore {
+  const record = (method: string) => log.push({ slot, method });
+  const runs = config.runs ?? [];
+
+  const store: Partial<FakeStore> = {
+    slot,
+    primaryReadClient: { __primary: slot },
+
+    findRun: ((_where: unknown, _args?: unknown) => {
+      record("findRun");
+      return Promise.resolve((runs[0] ?? null) as never);
+    }) as FakeStore["findRun"],
+
+    findRunOnPrimary: ((_where: unknown, _args?: unknown) => {
+      record("findRunOnPrimary");
+      return Promise.resolve((runs[0] ?? null) as never);
+    }) as FakeStore["findRunOnPrimary"],
+
+    findRunOrThrow: ((_where: unknown, _args?: unknown) => {
+      record("findRunOrThrow");
+      if (runs[0] === undefined) {
+        return Promise.reject(new Error(`no run on ${slot}`)) as never;
+      }
+      return Promise.resolve(runs[0] as never);
+    }) as FakeStore["findRunOrThrow"],
+
+    findRuns: ((_args: unknown, _client?: ReadClient) => {
+      record("findRuns");
+      return Promise.resolve(runs as never);
+    }) as FakeStore["findRuns"],
+
+    createRun: ((_params: unknown) => {
+      record("createRun");
+      return Promise.resolve({ slot } as never);
+    }) as FakeStore["createRun"],
+
+    createTaskRunCheckpoint: ((_args: unknown) => {
+      record("createTaskRunCheckpoint");
+      return Promise.resolve({ slot } as never);
+    }) as FakeStore["createTaskRunCheckpoint"],
+
+    findWaitpoint: ((_args: unknown, _client?: ReadClient) => {
+      record("findWaitpoint");
+      return Promise.resolve((config.waitpoint ?? null) as never);
+    }) as FakeStore["findWaitpoint"],
+
+    updateWaitpoint: ((_args: unknown) => {
+      record("updateWaitpoint");
+      return Promise.resolve({ slot } as never);
+    }) as FakeStore["updateWaitpoint"],
+
+    findManyTaskRunWaitpoints: ((_args: unknown, _client?: ReadClient) => {
+      record("findManyTaskRunWaitpoints");
+      return Promise.resolve((config.edges ?? []) as never);
+    }) as FakeStore["findManyTaskRunWaitpoints"],
+  };
+
+  return store as unknown as FakeStore;
+}
+
+// Deterministic residency by id prefix via the classify seam (no dependence on id-shape rules).
+function buildRouter(newConfig: FakeConfig = {}, legacyConfig: FakeConfig = {}) {
+  const log: Call[] = [];
+  const newStore = fakeStore("new", log, newConfig);
+  const legacyStore = fakeStore("legacy", log, legacyConfig);
+  const router = new RoutingRunStore({
+    new: newStore,
+    legacy: legacyStore,
+    classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"),
+  });
+  return { router, newStore, legacyStore, log };
+}
+
+const trace = (log: Call[]) => log.map((c) => `${c.slot}:${c.method}`);
+
+describe("RoutingRunStore #probeOrder — new then legacy, sequential", () => {
+  it("probes new BEFORE legacy for an unrouted findRun", async () => {
+    const { router, log } = buildRouter();
+    await router.findRun({ spanId: "span_x" });
+    expect(trace(log)).toEqual(["new:findRun", "legacy:findRun"]);
+  });
+
+  it("stops at the first non-null leg and never consults legacy", async () => {
+    const { router, log } = buildRouter({ runs: [{ id: "r1" }] });
+    await router.findRun({ spanId: "span_x" });
+    expect(trace(log)).toEqual(["new:findRun"]);
+  });
+
+  it("gives the LAST probe leg the canonical not-found throw", async () => {
+    const { router, log } = buildRouter();
+    await expect(router.findRunOrThrow({ spanId: "span_x" })).rejects.toThrow("no run on legacy");
+    // new is probed with the nullable read; only legacy is asked to throw.
+    expect(trace(log)).toEqual(["new:findRun", "legacy:findRunOrThrow"]);
+  });
+
+  it("probes each store's own primary for a read-your-writes unrouted findRun", async () => {
+    const { router, log } = buildRouter();
+    await router.findRunOnPrimary({ spanId: "span_x" });
+    expect(trace(log)).toEqual(["new:findRunOnPrimary", "legacy:findRunOnPrimary"]);
+  });
+});
+
+describe("RoutingRunStore #precedence — NEW wins a merge", () => {
+  it("keeps the NEW row for a duplicate run id on an open predicate", async () => {
+    const { router } = buildRouter(
+      { runs: [{ id: "dup", from: "new" }] },
+      { runs: [{ id: "dup", from: "legacy" }] }
+    );
+    const rows = (await router.findRuns({
+      where: { runtimeEnvironmentId: "env_1" },
+      select: { id: true, from: true },
+    })) as Array<{ id: string; from: string }>;
+    expect(rows).toHaveLength(1);
+    expect(rows[0]?.from).toBe("new");
+  });
+
+  // Every merge in the router MUST resolve a duplicate id NEW-wins, edges included.
+  it("keeps the NEW row for a duplicate edge id on a waitpoint-keyed edge read", async () => {
+    const { router } = buildRouter(
+      { edges: [{ id: "edge_dup", taskRunId: "new_run" }] },
+      { edges: [{ id: "edge_dup", taskRunId: "legacy_run" }] }
+    );
+    const edges = (await router.findManyTaskRunWaitpoints({
+      where: { waitpointId: "waitpoint_x" },
+      select: { id: true, taskRunId: true },
+    })) as Array<{ id: string; taskRunId: string }>;
+    expect(edges).toHaveLength(1);
+    expect(edges[0]?.taskRunId).toBe("new_run");
+  });
+});
+
+describe("RoutingRunStore id-less fallbacks — the two defaults differ by role", () => {
+  it("routes an id-less create to new (#idlessRouteShard)", async () => {
+    const { router, log } = buildRouter();
+    await router.createRun({ data: {} } as never);
+    expect(trace(log)).toEqual(["new:createRun"]);
+  });
+
+  it("routes an id-less checkpoint create to new (#idlessRouteShard)", async () => {
+    const { router, log } = buildRouter();
+    await router.createTaskRunCheckpoint({ data: {} } as never);
+    expect(trace(log)).toEqual(["new:createTaskRunCheckpoint"]);
+  });
+
+  it("routes an id-less waitpoint update to legacy (#idlessWaitpointShard)", async () => {
+    const { router, log } = buildRouter();
+    await router.updateWaitpoint({ where: { idempotencyKey: "k" }, data: {} } as never);
+    expect(trace(log)).toEqual(["legacy:updateWaitpoint"]);
+  });
+});