fix(run-store): bound connectedRuns relation hydrator on the dedicated schema

The dedicated-schema hydrator for a waitpoint's connectedRuns relation fetched an unbounded list of connections and full run rows. Bound it per parent with a window function so a waitpoint with many connections cannot pull a large result set for a display-only field.

Author: d-cs

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

@@ -0,0 +1,127 @@
+// RED→GREEN for bounding the DISPLAY `connectedRuns` RELATION hydration (findWaitpoint /
+// findManyWaitpoints with `include`/`select: { connectedRuns }`), NOT the id-list helper
+// `findWaitpointConnectedRunIds` (covered separately in connectedRunsBounded.test.ts).
+//
+// The dedicated-schema hydrator `hydrateConnectedRuns` goes through the generic
+// `batchHydrateJoinRelation`, which fetched EVERY WaitpointRunConnection link with no per-parent
+// cap — so a heavily-fanned-in waitpoint hydrated an unbounded connectedRuns list, bypassing the
+// CONNECTED_RUNS_LIMIT the presenter/`findWaitpointConnectedRunIds` already enforce.
+//
+// The fix bounds it per parent via a window function + existence-JOIN to TaskRun, mirroring the
+// bounded helper: a dangling connection row never occupies a LIMIT slot, and the returned relation
+// is capped at CONNECTED_RUNS_LIMIT.
+
+import { heteroRunOpsPostgresTest } from "@internal/testcontainers";
+import { describe, expect } from "vitest";
+import { CONNECTED_RUNS_LIMIT, PostgresRunStore } from "./PostgresRunStore.js";
+
+function seedEnvironmentDedicated(suffix: string) {
+  return {
+    organization: { id: `org_${suffix}` },
+    project: { id: `proj_${suffix}` },
+    environment: { id: `env_${suffix}` },
+  };
+}
+
+function taskRunData(opts: {
+  id: string;
+  friendlyId: string;
+  organizationId: string;
+  projectId: string;
+  runtimeEnvironmentId: string;
+}) {
+  return {
+    id: opts.id,
+    engine: "V2" as const,
+    status: "PENDING" as const,
+    friendlyId: opts.friendlyId,
+    runtimeEnvironmentId: opts.runtimeEnvironmentId,
+    environmentType: "DEVELOPMENT" as const,
+    organizationId: opts.organizationId,
+    projectId: opts.projectId,
+    taskIdentifier: "my-task",
+    payload: "{}",
+    payloadType: "application/json",
+    traceContext: {},
+    traceId: `trace_${opts.id}`,
+    spanId: `span_${opts.id}`,
+    queue: "task/my-task",
+    isTest: false,
+    taskEventStore: "taskEvent",
+    depth: 0,
+  };
+}
+
+describe("PostgresRunStore.findWaitpoint — connectedRuns relation is bounded", () => {
+  heteroRunOpsPostgresTest(
+    "dedicated schema: include connectedRuns caps at CONNECTED_RUNS_LIMIT and skips dangling connections",
+    async ({ prisma17 }) => {
+      const store = new PostgresRunStore({
+        prisma: prisma17 as never,
+        readOnlyPrisma: prisma17 as never,
+        schemaVariant: "dedicated",
+      });
+
+      const env = seedEnvironmentDedicated("cr_incl_new");
+      const waitpointId = "waitpoint_cr_incl_new";
+      await prisma17.waitpoint.create({
+        data: {
+          id: waitpointId,
+          friendlyId: "wp_cr_incl_new",
+          type: "MANUAL",
+          status: "PENDING",
+          idempotencyKey: `idem_${waitpointId}`,
+          userProvidedIdempotencyKey: false,
+          projectId: env.project.id,
+          environmentId: env.environment.id,
+        },
+      });
+
+      // Dangling connections (taskRunId points at runs that were never created). Legal on the
+      // dedicated schema — WaitpointRunConnection.taskRunId is a scalar with no FK.
+      const danglingRunIds = Array.from({ length: 20 }, (_, i) => `run_cr_dangling_${i}`);
+      for (const runId of danglingRunIds) {
+        await store.blockRunWithWaitpointEdges({
+          runId,
+          waitpointIds: [waitpointId],
+          projectId: env.project.id,
+        });
+      }
+
+      // More REAL connected runs than the limit.
+      const realRunIds = Array.from({ length: CONNECTED_RUNS_LIMIT + 3 }, (_, i) => `run_cr_real_${i}`);
+      for (const [i, id] of realRunIds.entries()) {
+        await prisma17.taskRun.create({
+          data: taskRunData({
+            id,
+            friendlyId: `run_cr_incl_new_${i}`,
+            organizationId: env.organization.id,
+            projectId: env.project.id,
+            runtimeEnvironmentId: env.environment.id,
+          }),
+        });
+        await store.blockRunWithWaitpointEdges({
+          runId: id,
+          waitpointIds: [waitpointId],
+          projectId: env.project.id,
+        });
+      }
+
+      const waitpoint = (await store.findWaitpoint({
+        where: { id: waitpointId },
+        include: { connectedRuns: true },
+      })) as unknown as { connectedRuns: { id: string }[] } | null;
+
+      expect(waitpoint).not.toBeNull();
+      const connectedRuns = waitpoint!.connectedRuns;
+
+      // Bounded.
+      expect(connectedRuns.length).toBeLessThanOrEqual(CONNECTED_RUNS_LIMIT);
+      // Never a dangling id — every returned run must be a REAL run.
+      for (const run of connectedRuns) {
+        expect(realRunIds).toContain(run.id);
+        expect(danglingRunIds).not.toContain(run.id);
+      }
+    }
+  );
+});

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

@@ -19,11 +19,12 @@ type CallCounts = { findMany: number; findFirst: number };
 function countingClient(
   real: RunOpsPrismaClient,
   delegateNames: string[]
-): { client: RunOpsPrismaClient; counts: Record<string, CallCounts> } {
+): { client: RunOpsPrismaClient; counts: Record<string, CallCounts>; queryRaw: { count: number } } {
   const counts: Record<string, CallCounts> = {};
   for (const name of delegateNames) {
     counts[name] = { findMany: 0, findFirst: 0 };
   }
+  const queryRaw = { count: 0 };
   const wrapped = new Map(
     delegateNames.map((name) => [
       name,
@@ -42,10 +43,17 @@ function countingClient(
       if (typeof prop === "string" && wrapped.has(prop)) {
         return wrapped.get(prop);
       }
+      // Tally the grouped raw query the bounded connectedRuns hydrator issues, delegating unchanged.
+      if (prop === "$queryRaw") {
+        return (...args: unknown[]) => {
+          queryRaw.count++;
+          return (target as any).$queryRaw(...args);
+        };
+      }
       return (target as any)[prop];
     },
   }) as RunOpsPrismaClient;
-  return { client, counts };
+  return { client, counts, queryRaw };
 }
 
 function seedEnvironmentDedicated(suffix: string) {
@@ -219,8 +227,11 @@ describe("PostgresRunStore dedicated relation hydrators — grouped batch reads"
         expect(connected[0].id).toBe(runIds[i]);
       }
 
-      // GROUPED: one join query + one target query for the WHOLE batch, never one pair per parent.
-      expect(counting.counts.waitpointRunConnection.findMany).toBe(1);
+      // GROUPED: one bounded join query (raw, ROW_NUMBER-per-parent) + one target query for the
+      // WHOLE batch, never one pair per parent. The bounded hydrator replaces the delegate
+      // `waitpointRunConnection.findMany` with a single `$queryRaw`.
+      expect(counting.queryRaw.count).toBe(1);
+      expect(counting.counts.waitpointRunConnection.findMany).toBe(0);
       expect(counting.counts.taskRun.findMany).toBe(1);
       expect(counting.counts.taskRun.findFirst).toBe(0);
     }

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

@@ -336,16 +336,47 @@ const hydrateBlockingTaskRuns: DedicatedRelationHydrator = async (client, parent
   return byParent;
 };
 
-// Display connections for a waitpoint: WaitpointRunConnection → TaskRun rows.
-const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parents, projection) =>
-  batchHydrateJoinRelation(
-    client.waitpointRunConnection,
-    client.taskRun,
-    parents.map((p) => p.id as string),
-    "waitpointId",
-    "taskRunId",
-    projection
-  );
+// Display connections for a waitpoint: WaitpointRunConnection → TaskRun rows. Bounded per parent to
+// CONNECTED_RUNS_LIMIT via a window function + existence-JOIN to TaskRun, mirroring the id-list
+// helper findWaitpointConnectedRunIds: a dangling (run-less) connection row never occupies a LIMIT
+// slot, and a heavily-fanned-in waitpoint never hydrates an unbounded connectedRuns list. This is
+// the DISPLAY relation only — functional blocking reads (hydrateBlockingTaskRuns) stay uncapped.
+const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parents, projection) => {
+  const parentIds = parents.map((p) => p.id as string);
+  const byParent = new Map<string, unknown[]>(parentIds.map((id) => [id, []]));
+  if (parentIds.length === 0) {
+    return byParent;
+  }
+  // One grouped query for the bounded edges across every parent: ROW_NUMBER partitioned per
+  // waitpoint keeps at most CONNECTED_RUNS_LIMIT rows per parent (uses @@index([waitpointId])).
+  const links = (await client.$queryRaw`
+    SELECT ranked."waitpointId" AS "waitpointId", ranked."taskRunId" AS "taskRunId"
+    FROM (
+      SELECT c."waitpointId", c."taskRunId",
+        ROW_NUMBER() OVER (PARTITION BY c."waitpointId" ORDER BY c."id") AS rn
+      FROM "WaitpointRunConnection" c
+      JOIN "TaskRun" t ON t."id" = c."taskRunId"
+      WHERE c."waitpointId" = ANY(${parentIds}::text[])
+    ) ranked
+    WHERE ranked.rn <= ${CONNECTED_RUNS_LIMIT}
+  `) as { waitpointId: string; taskRunId: string }[];
+  if (links.length === 0) {
+    return byParent;
+  }
+  const targetIds = [...new Set(links.map((l) => l.taskRunId))];
+  const rows = (await client.taskRun.findMany(
+    targetFindManyArgs({ id: { in: targetIds } }, projection, ["id"])
+  )) as Record<string, unknown>[];
+  const byTargetId = new Map(rows.map((r) => [r.id as string, r]));
+  for (const link of links) {
+    const target = byTargetId.get(link.taskRunId);
+    const bucket = byParent.get(link.waitpointId);
+    if (target && bucket) {
+      bucket.push(applyProjection(target, projection));
+    }
+  }
+  return byParent;
+};
 
 // Snapshots that completed a waitpoint: CompletedWaitpoint join → TaskRunExecutionSnapshot rows.
 const hydrateCompletedExecutionSnapshots: DedicatedRelationHydrator = async (