fix(custom-blocks): track the whole child run, not just its finalization

Author: TheodoreSpeaks

apps/sim/executor/handlers/workflow/workflow-handler.test.ts

@@ -18,7 +18,7 @@ const {
   mockGetCustomBlockAuthority,
   mockGetUserEmailById,
   mockAdmitCustomBlockChildExecution,
-  mockTrackChildFinalization,
+  mockTrackChildRun,
   mockBuildTraceSpans,
   mockSafeStart,
   mockSafeComplete,
@@ -34,7 +34,7 @@ const {
   mockGetCustomBlockAuthority: vi.fn(),
   mockGetUserEmailById: vi.fn(),
   mockAdmitCustomBlockChildExecution: vi.fn(),
-  mockTrackChildFinalization: vi.fn(),
+  mockTrackChildRun: vi.fn(),
   mockBuildTraceSpans: vi.fn(),
   mockSafeStart: vi.fn(),
   mockSafeComplete: vi.fn(),
@@ -65,7 +65,7 @@ vi.mock('@/lib/logs/execution/trace-spans/trace-spans', () => ({
 
 vi.mock('@/lib/workflows/custom-blocks/child-execution', () => ({
   admitCustomBlockChildExecution: mockAdmitCustomBlockChildExecution,
-  trackChildFinalization: mockTrackChildFinalization,
+  trackChildRun: mockTrackChildRun,
   buildCustomBlockCorrelation: (params: Record<string, any>) =>
     params.invokerExecutionId
       ? { source: 'custom_block', executionId: params.invokerExecutionId }
@@ -1250,22 +1250,32 @@ describe('WorkflowBlockHandler', () => {
       expect(error.cause).toBeUndefined()
     })
 
-    it('registers the finalization so a cancelled parent can still drain it', async () => {
+    it('registers the child run BEFORE executing it, and settles it when done', async () => {
+      // Ordering is the whole point: a cancelled parent drains while the child is
+      // still inside `execute`, so registering after it would find nothing.
+      let registeredBeforeExecute = false
+      mockExecutorExecute.mockImplementation(async () => {
+        registeredBeforeExecute = mockTrackChildRun.mock.calls.length === 1
+        return { success: true, output: { data: 'ok' } }
+      })
+
       await handler.execute(customBlockContext(), customBlock(), {})
 
-      expect(mockTrackChildFinalization).toHaveBeenCalledTimes(1)
-      const [invokerId, promise] = mockTrackChildFinalization.mock.calls[0]
+      expect(registeredBeforeExecute).toBe(true)
+      const [invokerId, childRun] = mockTrackChildRun.mock.calls[0]
       expect(invokerId).toBe('parent-execution-id')
-      expect(promise).toBeInstanceOf(Promise)
+      // Settled in `finally`, so the invoking run's drain can complete.
+      await expect(childRun).resolves.toBeUndefined()
     })
 
-    it('registers the finalization on the failure path too', async () => {
+    it('settles the registered child run on the failure path too', async () => {
       mockExecutorExecute.mockRejectedValue(new Error('boom'))
 
       await handler.execute(customBlockContext(), customBlock(), {}).catch(() => {})
 
-      expect(mockTrackChildFinalization).toHaveBeenCalledTimes(1)
-      expect(mockTrackChildFinalization.mock.calls[0][0]).toBe('parent-execution-id')
+      expect(mockTrackChildRun).toHaveBeenCalledTimes(1)
+      expect(mockTrackChildRun.mock.calls[0][0]).toBe('parent-execution-id')
+      await expect(mockTrackChildRun.mock.calls[0][1]).resolves.toBeUndefined()
     })
 
     it('never leaks the source workflow name when the child returns success: false', async () => {

apps/sim/executor/handlers/workflow/workflow-handler.ts

@@ -14,7 +14,7 @@ import {
   admitCustomBlockChildExecution,
   buildCustomBlockCorrelation,
   createChildCancellationSignal,
-  trackChildFinalization,
+  trackChildRun,
 } from '@/lib/workflows/custom-blocks/child-execution'
 import { getCustomBlockAuthority } from '@/lib/workflows/custom-blocks/operations'
 import { extractInputFieldsFromBlocks } from '@/lib/workflows/input-format'
@@ -275,6 +275,8 @@ export class WorkflowBlockHandler implements BlockHandler {
     /** Large-value id list shared with the child (and any nested custom blocks). */
     let sharedLargeValueIds: string[] | undefined
     let childCancellation: { signal: AbortSignal; dispose: () => void } | undefined
+    /** Settled in `finally` once the child is fully done — see `trackChildRun`. */
+    let settleChildRun: (() => void) | undefined
     try {
       // A custom block runs the source's latest deployment; if the source has been
       // undeployed there's nothing to run. `BoundarySafeError` marks the message as
@@ -485,6 +487,17 @@ export class WorkflowBlockHandler implements BlockHandler {
           parentSignal: ctx.abortSignal,
           parentExecutionId: ctx.executionId,
         })
+        // Registered BEFORE the child starts and settled in `finally`, so the
+        // tracked promise spans execution AND the terminal log write. A cancelled
+        // parent drains at a moment when the child is still inside `execute`, so
+        // registering only the finalization step would find nothing to await.
+        trackChildRun(
+          ctx.executionId,
+          new Promise<void>((resolve) => {
+            settleChildRun = resolve
+          })
+        )
+
         // Large values are scoped by execution id, so the parent must be able to
         // read a large exposed output the child produced. ONE array is shared down
         // the whole chain rather than copied per hop: a nested custom block pushes
@@ -618,10 +631,7 @@ export class WorkflowBlockHandler implements BlockHandler {
       const duration = performance.now() - startTime
 
       if (childSession && childSessionStarted) {
-        await this.trackFinalization(
-          ctx,
-          this.finalizeChildSession(childSession, executionResult, duration, childWorkflowInput)
-        )
+        await this.finalizeChildSession(childSession, executionResult, duration, childWorkflowInput)
         childSessionFinalized = true
       }
 
@@ -672,7 +682,7 @@ export class WorkflowBlockHandler implements BlockHandler {
       // The child's own log row records the real failure in the source workspace,
       // so the publisher sees what the consumer deliberately cannot.
       if (childSession && childSessionStarted && !childSessionFinalized) {
-        await this.trackFinalization(ctx, this.failChildSession(childSession, error))
+        await this.failChildSession(childSession, error)
       }
 
       // A custom block is checked FIRST and unconditionally: errors this invocation
@@ -730,25 +740,12 @@ export class WorkflowBlockHandler implements BlockHandler {
       // A custom block inside a loop would otherwise leak one abort listener and
       // one cancellation subscription per iteration.
       childCancellation?.dispose()
+      // Releases the invoking run's drain: reached on every exit path, including
+      // when this handler's promise was abandoned by a cancelled parent engine.
+      settleChildRun?.()
     }
   }
 
-  /**
-   * Awaits a child-session finalization while ALSO registering it against the
-   * invoking run. A cancelled or timed-out parent engine returns without
-   * draining its in-flight node promises, which would abandon this `await`
-   * mid-write; the registration lets the invoking run's completion path finish
-   * the durable write instead of leaving the child's row `running` for the
-   * stale-execution reaper.
-   */
-  private async trackFinalization(
-    ctx: ExecutionContext,
-    finalization: Promise<void>
-  ): Promise<void> {
-    trackChildFinalization(ctx.executionId, finalization)
-    await finalization
-  }
-
   /**
    * Completes a custom-block child's own logging session, so the publisher gets a
    * full run record — trace waterfall, duration, and ledger rows — in the source

apps/sim/lib/workflows/custom-blocks/child-execution.test.ts

@@ -32,8 +32,8 @@ import {
   buildCustomBlockCorrelation,
   CustomBlockAdmissionError,
   createChildCancellationSignal,
-  trackChildFinalization,
-  waitForChildFinalizations,
+  trackChildRun,
+  waitForChildRuns,
 } from '@/lib/workflows/custom-blocks/child-execution'
 import { isBoundarySafeError } from '@/executor/errors/boundary'
 
@@ -245,47 +245,60 @@ describe('createChildCancellationSignal durable backstop', () => {
   })
 })
 
-describe('child finalization tracking', () => {
-  it('lets an invoking run await a finalization the engine abandoned', async () => {
-    let resolveWrite: () => void = () => {}
-    const durableWrite = new Promise<void>((resolve) => {
-      resolveWrite = resolve
-    })
+describe('child run tracking', () => {
+  it('waits for a child the engine abandoned mid-execution', async () => {
+    // The scenario this exists for: the parent cancelled, so at drain time the
+    // child has NOT finished. Registration must already have happened.
+    let settle: () => void = () => {}
     let finished = false
-    const finalization = durableWrite.then(() => {
+    const childRun = new Promise<void>((resolve) => {
+      settle = resolve
+    }).then(() => {
       finished = true
     })
 
-    trackChildFinalization('invoker-1', finalization)
+    trackChildRun('invoker-1', childRun)
 
-    // The engine returned without draining; the write is still in flight.
+    const drained = waitForChildRuns('invoker-1')
     expect(finished).toBe(false)
-
-    const drained = waitForChildFinalizations('invoker-1')
-    resolveWrite()
+    settle()
     await drained
 
     expect(finished).toBe(true)
   })
 
-  it('does not let a rejected finalization break the drain', async () => {
-    trackChildFinalization('invoker-2', Promise.reject(new Error('db down')))
+  it('gives up on a hung child rather than wedging the invoking run', async () => {
+    vi.useFakeTimers()
+    try {
+      // Never settles — an unbounded drain would hang the parent forever.
+      trackChildRun('invoker-hung', new Promise<void>(() => {}))
+
+      const drained = waitForChildRuns('invoker-hung')
+      await vi.advanceTimersByTimeAsync(10_000)
+
+      await expect(drained).resolves.toBeUndefined()
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
+  it('does not let a rejected child break the drain', async () => {
+    trackChildRun('invoker-2', Promise.reject(new Error('db down')))
 
-    await expect(waitForChildFinalizations('invoker-2')).resolves.toBeUndefined()
+    await expect(waitForChildRuns('invoker-2')).resolves.toBeUndefined()
   })
 
   it('is a no-op for a run that registered nothing', async () => {
-    await expect(waitForChildFinalizations('invoker-never')).resolves.toBeUndefined()
-    await expect(waitForChildFinalizations(undefined)).resolves.toBeUndefined()
+    await expect(waitForChildRuns('invoker-never')).resolves.toBeUndefined()
+    await expect(waitForChildRuns(undefined)).resolves.toBeUndefined()
   })
 
   it('releases its entry so a run that never drains cannot leak', async () => {
-    const finalization = Promise.resolve()
-    trackChildFinalization('invoker-3', finalization)
-    await finalization
+    const childRun = Promise.resolve()
+    trackChildRun('invoker-3', childRun)
+    await childRun
     await Promise.resolve()
 
-    // Draining after self-cleanup still resolves rather than hanging.
-    await expect(waitForChildFinalizations('invoker-3')).resolves.toBeUndefined()
+    await expect(waitForChildRuns('invoker-3')).resolves.toBeUndefined()
   })
 })

apps/sim/lib/workflows/custom-blocks/child-execution.ts

@@ -1,3 +1,4 @@
+import { createLogger } from '@sim/logger'
 import {
   type BillingAttributionSnapshot,
   checkAttributedUsageLimits,
@@ -10,6 +11,8 @@ import {
 } from '@/lib/execution/cancellation'
 import { BoundarySafeError } from '@/executor/errors/boundary'
 
+const logger = createLogger('CustomBlockChildExecution')
+
 /**
  * The payer has no headroom for this custom-block child run.
  *
@@ -144,58 +147,91 @@ export async function createChildCancellationSignal(params: {
 }
 
 /**
- * Child-session finalizations still in flight, keyed by the INVOKING run's
- * execution id.
+ * Custom-block child runs still in flight, keyed by the INVOKING run's execution id.
  *
  * A cancelled or timed-out parent engine returns without draining its in-flight
- * node promises, so the custom-block handler's finalization would otherwise be
- * abandoned mid-write: the invoking run finishes, the worker tears down, and the
- * child's log row is left `running` until the stale-execution reaper sweeps it
- * minutes later. Registering here lets the run's own completion path await the
- * durable write instead of racing it.
+ * node promises, so the child's execution AND its terminal log write are both
+ * abandoned mid-flight: the invoking run finishes, the worker tears down, and the
+ * child's row is left `running` until the stale-execution reaper sweeps it.
+ *
+ * The tracked promise therefore covers the WHOLE child lifecycle and is registered
+ * BEFORE the child starts. Registering only the finalization step would be useless
+ * on exactly the path this exists for: at drain time the child is still inside
+ * `execute`, so nothing would be registered yet and the drain would no-op.
  *
- * Only the immediate invoker is tracked. A finalization orphaned *below* another
- * abandoned child still falls back to the reaper.
+ * Only the immediate invoker is tracked. A child orphaned *below* another abandoned
+ * child still falls back to the reaper.
  */
-const pendingChildFinalizations = new Map<string, Set<Promise<unknown>>>()
+const pendingChildRuns = new Map<string, Set<Promise<unknown>>>()
+
+/** How long a run will wait on abandoned children before leaving them to the reaper. */
+const CHILD_RUN_DRAIN_TIMEOUT_MS = 10_000
 
-/** Registers a child-session finalization against its invoking run. */
-export function trackChildFinalization(
+/**
+ * Registers a child run against its invoking run. The promise must settle when the
+ * child is fully done — execution finished AND its session written.
+ */
+export function trackChildRun(
   invokerExecutionId: string | undefined,
-  promise: Promise<unknown>
+  childRun: Promise<unknown>
 ): void {
   if (!invokerExecutionId) return
 
-  let pending = pendingChildFinalizations.get(invokerExecutionId)
+  let pending = pendingChildRuns.get(invokerExecutionId)
   if (!pending) {
     pending = new Set()
-    pendingChildFinalizations.set(invokerExecutionId, pending)
+    pendingChildRuns.set(invokerExecutionId, pending)
   }
-  pending.add(promise)
+  pending.add(childRun)
 
   // Self-cleaning so an invoker that never drains cannot leak the entry. The
-  // `catch` also marks the rejection handled — callers use `allSettled`.
-  void promise
+  // `catch` also marks the rejection handled — the drain uses `allSettled`.
+  void childRun
     .catch(() => {})
     .finally(() => {
-      pending.delete(promise)
-      if (pending.size === 0) pendingChildFinalizations.delete(invokerExecutionId)
+      pending.delete(childRun)
+      if (pending.size === 0) pendingChildRuns.delete(invokerExecutionId)
     })
 }
 
-/** Awaits every child-session finalization registered for this run. */
-export async function waitForChildFinalizations(
-  invokerExecutionId: string | undefined
-): Promise<void> {
+/**
+ * Awaits every child run registered for this invoking run, bounded.
+ *
+ * The bound matters: on the normal path the handler already awaited its child, so
+ * this returns immediately. It only blocks when a child was abandoned, where the
+ * bridged abort should end it in milliseconds. Waiting unbounded would trade a
+ * stale `running` row — which the reaper fixes — for a wedged parent run, which
+ * nothing fixes.
+ */
+export async function waitForChildRuns(invokerExecutionId: string | undefined): Promise<void> {
   if (!invokerExecutionId) return
 
-  const pending = pendingChildFinalizations.get(invokerExecutionId)
-  if (!pending) return
+  const pending = pendingChildRuns.get(invokerExecutionId)
+  if (!pending || pending.size === 0) return
 
-  // Re-check: a settling finalization can register nothing further, but a
-  // nested child can still be added while we await.
-  while (pending.size > 0) {
-    await Promise.allSettled([...pending])
+  let timer: ReturnType<typeof setTimeout> | undefined
+  const timedOut = new Promise<'timeout'>((resolve) => {
+    timer = setTimeout(() => resolve('timeout'), CHILD_RUN_DRAIN_TIMEOUT_MS)
+  })
+
+  const drained = (async () => {
+    // Re-check: a nested child can still register while we await.
+    while (pending.size > 0) {
+      await Promise.allSettled([...pending])
+    }
+    return 'drained' as const
+  })()
+
+  try {
+    if ((await Promise.race([drained, timedOut])) === 'timeout') {
+      logger.warn('Abandoned custom block child did not finish; leaving it to the reaper', {
+        invokerExecutionId,
+        outstanding: pending.size,
+      })
+      return
+    }
+    pendingChildRuns.delete(invokerExecutionId)
+  } finally {
+    if (timer) clearTimeout(timer)
   }
-  pendingChildFinalizations.delete(invokerExecutionId)
 }