@@ -113,7 +113,8 @@ class CountingPostgresRunStore extends PostgresRunStore {
113113async function driveBatchToAllChildrenComplete (
114114 engine : RunEngine ,
115115 prisma : PrismaClient ,
116- friendlyPrefix : string
116+ friendlyPrefix : string ,
117+ completeChildren = true
117118) {
118119 const environment = await setupAuthenticatedEnvironment ( prisma , "PRODUCTION" ) ;
119120 const parentTask = "parent-task" ;
@@ -183,7 +184,7 @@ async function driveBatchToAllChildrenComplete(
183184 queue : `task/${ childTask } ,
184185 isTest : false ,
185186 tags : [ ] ,
186- resumeParentOnCompletion : true ,
187+ resumeParentOnCompletion : completeChildren ,
187188 parentTaskRunId : parentRun . id ,
188189 batch : { id : batch . id , index : 0 } ,
189190 } ,
@@ -206,39 +207,41 @@ async function driveBatchToAllChildrenComplete(
206207 queue : `task/${ childTask } ,
207208 isTest : false ,
208209 tags : [ ] ,
209- resumeParentOnCompletion : true ,
210+ resumeParentOnCompletion : completeChildren ,
210211 parentTaskRunId : parentRun . id ,
211212 batch : { id : batch . id , index : 1 } ,
212213 } ,
213214 prisma
214215 ) ;
215216
216- for ( const child of [ child1 , child2 ] ) {
217+ if ( completeChildren ) {
218+ for ( const child of [ child1 , child2 ] ) {
219+ await setTimeout ( 500 ) ;
220+ const dequeued = await engine . dequeueFromWorkerQueue ( {
221+ consumerId : "test_consumer" ,
222+ workerQueue : "main" ,
223+ } ) ;
224+ const match = dequeued . find ( ( d ) => d . run . id === child . id ) ?? dequeued [ 0 ] ;
225+ assertNonNullable ( match ) ;
226+ const attempt = await engine . startRunAttempt ( {
227+ runId : match . run . id ,
228+ snapshotId : match . snapshot . id ,
229+ } ) ;
230+ await engine . completeRunAttempt ( {
231+ runId : attempt . run . id ,
232+ snapshotId : attempt . snapshot . id ,
233+ completion : {
234+ id : attempt . run . id ,
235+ ok : true ,
236+ output : '{"foo":"bar"}' ,
237+ outputType : "application/json" ,
238+ } ,
239+ } ) ;
240+ }
241+
217242 await setTimeout ( 500 ) ;
218- const dequeued = await engine . dequeueFromWorkerQueue ( {
219- consumerId : "test_consumer" ,
220- workerQueue : "main" ,
221- } ) ;
222- const match = dequeued . find ( ( d ) => d . run . id === child . id ) ?? dequeued [ 0 ] ;
223- assertNonNullable ( match ) ;
224- const attempt = await engine . startRunAttempt ( {
225- runId : match . run . id ,
226- snapshotId : match . snapshot . id ,
227- } ) ;
228- await engine . completeRunAttempt ( {
229- runId : attempt . run . id ,
230- snapshotId : attempt . snapshot . id ,
231- completion : {
232- id : attempt . run . id ,
233- ok : true ,
234- output : '{"foo":"bar"}' ,
235- outputType : "application/json" ,
236- } ,
237- } ) ;
238243 }
239244
240- await setTimeout ( 500 ) ;
241-
242245 return { environment, batch, parentRun, child1, child2 } ;
243246}
244247
@@ -308,6 +311,57 @@ describe("RunEngine #tryCompleteBatch store routing", () => {
308311 }
309312 ) ;
310313
314+ containerTest (
315+ "retries waitpoint completion for a completed batch whose parent was not resumed" ,
316+ async ( { prisma, redisOptions } ) => {
317+ const engine = new RunEngine ( createEngineOptions ( redisOptions , prisma ) ) ;
318+
319+ try {
320+ const { batch, parentRun } = await driveBatchToAllChildrenComplete (
321+ engine ,
322+ prisma ,
323+ "run_batch_completion_recovery" ,
324+ false
325+ ) ;
326+
327+ const waitpoint = await prisma . waitpoint . findFirstOrThrow ( {
328+ where : { completedByBatchId : batch . id } ,
329+ } ) ;
330+ expect ( waitpoint . status ) . toBe ( "PENDING" ) ;
331+ expect (
332+ await prisma . taskRunWaitpoint . count ( {
333+ where : { taskRunId : parentRun . id , waitpointId : waitpoint . id } ,
334+ } )
335+ ) . toBe ( 1 ) ;
336+
337+ await prisma . batchTaskRun . update ( {
338+ where : { id : batch . id } ,
339+ data : { status : "COMPLETED" , resumedAt : null } ,
340+ } ) ;
341+
342+ await engine . batchSystem . performCompleteBatch ( { batchId : batch . id } ) ;
343+
344+ const recoveredBatch = await prisma . batchTaskRun . findFirstOrThrow ( {
345+ where : { id : batch . id } ,
346+ } ) ;
347+ expect ( recoveredBatch . resumedAt ) . not . toBeNull ( ) ;
348+
349+ const recoveredWaitpoint = await prisma . waitpoint . findFirstOrThrow ( {
350+ where : { id : waitpoint . id } ,
351+ } ) ;
352+ expect ( recoveredWaitpoint . status ) . toBe ( "COMPLETED" ) ;
353+
354+ await setTimeout ( 1_000 ) ;
355+ expect ( await prisma . taskRunWaitpoint . count ( { where : { taskRunId : parentRun . id } } ) ) . toBe ( 0 ) ;
356+ const parentExecution = await engine . getRunExecutionData ( { runId : parentRun . id } ) ;
357+ assertNonNullable ( parentExecution ) ;
358+ expect ( parentExecution . snapshot . executionStatus ) . not . toBe ( "EXECUTING_WITH_WAITPOINTS" ) ;
359+ } finally {
360+ await engine . quit ( ) ;
361+ }
362+ }
363+ ) ;
364+
311365 // The member-run read is driven by batchId only and does not rely on the
312366 // BatchTaskRun.runtimeEnvironmentId FK. A second batch (distinct batchId) must not leak members
313367 // into the first batch's batchId-scoped read.
0 commit comments