@@ -126,10 +126,13 @@ function enhanceExecutionSnapshotWithWaitpoints(
126126async function getSnapshotWaitpointIds (
127127 prisma : PrismaClientOrTransaction ,
128128 snapshotId : string ,
129- runStore ?: RunStore
129+ runStore ?: RunStore ,
130+ // The owning run id, so the router can route to the run's store (the completed-waitpoint join
131+ // co-locates with the snapshot/run) instead of fanning out to both run-ops DBs.
132+ runId ?: string
130133) : Promise < string [ ] > {
131134 if ( runStore ) {
132- return runStore . findSnapshotCompletedWaitpointIds ( snapshotId , prisma ) ;
135+ return runStore . findSnapshotCompletedWaitpointIds ( snapshotId , prisma , runId ) ;
133136 }
134137
135138 const result = await prisma . $queryRaw < { B : string } [ ] > `
@@ -144,10 +147,12 @@ async function getSnapshotWaitpointIds(
144147async function getSnapshotWaitpointIdsWithPresence (
145148 prisma : PrismaClientOrTransaction ,
146149 snapshotId : string ,
147- runStore ?: RunStore
150+ runStore ?: RunStore ,
151+ // The owning run id, so the router can route to the run's store instead of fanning out.
152+ runId ?: string
148153) : Promise < { present : boolean ; ids : string [ ] } > {
149154 if ( runStore ) {
150- return runStore . findSnapshotCompletedWaitpointIdsWithPresence ( snapshotId , prisma ) ;
155+ return runStore . findSnapshotCompletedWaitpointIdsWithPresence ( snapshotId , prisma , runId ) ;
151156 }
152157
153158 const rows = await prisma . $queryRaw < { id : string ; B : string | null } [ ] > `
@@ -170,15 +175,18 @@ async function getSnapshotWaitpointIdsWithPresence(
170175async function fetchWaitpointsInChunks (
171176 prisma : PrismaClientOrTransaction ,
172177 waitpointIds : string [ ] ,
173- runStore ?: RunStore
178+ runStore ?: RunStore ,
179+ // The owning run id, so the router routes to the run's store and only falls back to the other DB
180+ // for the rare cross-tree token, instead of fanning every chunk out to both run-ops DBs.
181+ runId ?: string
174182) : Promise < Waitpoint [ ] > {
175183 if ( waitpointIds . length === 0 ) return [ ] ;
176184
177185 const allWaitpoints : Waitpoint [ ] = [ ] ;
178186 for ( let i = 0 ; i < waitpointIds . length ; i += WAITPOINT_CHUNK_SIZE ) {
179187 const chunk = waitpointIds . slice ( i , i + WAITPOINT_CHUNK_SIZE ) ;
180188 const waitpoints = runStore
181- ? await runStore . findManyWaitpoints ( { where : { id : { in : chunk } } } , prisma )
189+ ? await runStore . findManyWaitpoints ( { where : { id : { in : chunk } } } , prisma , runId )
182190 : await prisma . waitpoint . findMany ( {
183191 where : { id : { in : chunk } } ,
184192 } ) ;
@@ -347,7 +355,8 @@ export async function getExecutionSnapshotsSince(
347355 const { present, ids } = await getSnapshotWaitpointIdsWithPresence (
348356 prisma ,
349357 latestSnapshot . id ,
350- runStore
358+ runStore ,
359+ runId
351360 ) ;
352361 let waitpointIds = ids ;
353362
@@ -356,7 +365,7 @@ export async function getExecutionSnapshotsSince(
356365 // authoritative - re-read from the primary so the runner is not handed a waitpoint-less continue
357366 // (which it silently drops, hanging the run). Single-reader replicas never hit this (present stays true).
358367 if ( repairClient && repairClient !== prisma && ! present ) {
359- waitpointIds = await getSnapshotWaitpointIds ( repairClient , latestSnapshot . id , runStore ) ;
368+ waitpointIds = await getSnapshotWaitpointIds ( repairClient , latestSnapshot . id , runStore , runId ) ;
360369 readClient = repairClient ;
361370 }
362371
@@ -369,15 +378,15 @@ export async function getExecutionSnapshotsSince(
369378 // poll even against a caught-up replica.
370379 const expectedCount = new Set ( latestSnapshot . completedWaitpointOrder ?? [ ] ) . size ;
371380 if ( repairClient && repairClient !== prisma && waitpointIds . length < expectedCount ) {
372- const repaired = await getSnapshotWaitpointIds ( repairClient , latestSnapshot . id , runStore ) ;
381+ const repaired = await getSnapshotWaitpointIds ( repairClient , latestSnapshot . id , runStore , runId ) ;
373382 if ( repaired . length > waitpointIds . length ) {
374383 waitpointIds = repaired ;
375384 readClient = repairClient ;
376385 }
377386 }
378387
379388 // Step 4: Fetch waitpoints in chunks to avoid NAPI string conversion limits
380- const waitpoints = await fetchWaitpointsInChunks ( readClient , waitpointIds , runStore ) ;
389+ const waitpoints = await fetchWaitpointsInChunks ( readClient , waitpointIds , runStore , runId ) ;
381390
382391 // Step 5: Build enhanced snapshots - only latest gets waitpoints, others get empty arrays
383392 // The runner only uses completedWaitpoints from the latest snapshot anyway
0 commit comments