fix(custom-blocks): drain child log finalization when the parent is cancelled

Author: TheodoreSpeaks

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

@@ -18,6 +18,7 @@ const {
   mockGetCustomBlockAuthority,
   mockGetUserEmailById,
   mockAdmitCustomBlockChildExecution,
+  mockTrackChildFinalization,
   mockBuildTraceSpans,
   mockSafeStart,
   mockSafeComplete,
@@ -33,6 +34,7 @@ const {
   mockGetCustomBlockAuthority: vi.fn(),
   mockGetUserEmailById: vi.fn(),
   mockAdmitCustomBlockChildExecution: vi.fn(),
+  mockTrackChildFinalization: vi.fn(),
   mockBuildTraceSpans: vi.fn(),
   mockSafeStart: vi.fn(),
   mockSafeComplete: vi.fn(),
@@ -63,6 +65,7 @@ vi.mock('@/lib/logs/execution/trace-spans/trace-spans', () => ({
 
 vi.mock('@/lib/workflows/custom-blocks/child-execution', () => ({
   admitCustomBlockChildExecution: mockAdmitCustomBlockChildExecution,
+  trackChildFinalization: mockTrackChildFinalization,
   buildCustomBlockCorrelation: (params: Record<string, any>) =>
     params.invokerExecutionId
       ? { source: 'custom_block', executionId: params.invokerExecutionId }
@@ -1247,6 +1250,24 @@ describe('WorkflowBlockHandler', () => {
       expect(error.cause).toBeUndefined()
     })
 
+    it('registers the finalization so a cancelled parent can still drain it', async () => {
+      await handler.execute(customBlockContext(), customBlock(), {})
+
+      expect(mockTrackChildFinalization).toHaveBeenCalledTimes(1)
+      const [invokerId, promise] = mockTrackChildFinalization.mock.calls[0]
+      expect(invokerId).toBe('parent-execution-id')
+      expect(promise).toBeInstanceOf(Promise)
+    })
+
+    it('registers the finalization 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')
+    })
+
     it('never leaks the source workflow name when the child returns success: false', async () => {
       mockExecutorExecute.mockResolvedValue({
         success: false,

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

@@ -14,6 +14,7 @@ import {
   admitCustomBlockChildExecution,
   buildCustomBlockCorrelation,
   createChildCancellationSignal,
+  trackChildFinalization,
 } from '@/lib/workflows/custom-blocks/child-execution'
 import { getCustomBlockAuthority } from '@/lib/workflows/custom-blocks/operations'
 import { extractInputFieldsFromBlocks } from '@/lib/workflows/input-format'
@@ -601,7 +602,10 @@ export class WorkflowBlockHandler implements BlockHandler {
       const duration = performance.now() - startTime
 
       if (childSession && childSessionStarted) {
-        await this.finalizeChildSession(childSession, executionResult, duration, childWorkflowInput)
+        await this.trackFinalization(
+          ctx,
+          this.finalizeChildSession(childSession, executionResult, duration, childWorkflowInput)
+        )
         childSessionFinalized = true
       }
 
@@ -652,7 +656,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.failChildSession(childSession, error)
+        await this.trackFinalization(ctx, this.failChildSession(childSession, error))
       }
 
       // A custom block is checked FIRST and unconditionally: errors this invocation
@@ -713,6 +717,22 @@ export class WorkflowBlockHandler implements BlockHandler {
     }
   }
 
+  /**
+   * 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,6 +32,8 @@ import {
   buildCustomBlockCorrelation,
   CustomBlockAdmissionError,
   createChildCancellationSignal,
+  trackChildFinalization,
+  waitForChildFinalizations,
 } from '@/lib/workflows/custom-blocks/child-execution'
 import { isBoundarySafeError } from '@/executor/errors/boundary'
 
@@ -242,3 +244,48 @@ describe('createChildCancellationSignal durable backstop', () => {
     expect(signal.aborted).toBe(false)
   })
 })
+
+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
+    })
+    let finished = false
+    const finalization = durableWrite.then(() => {
+      finished = true
+    })
+
+    trackChildFinalization('invoker-1', finalization)
+
+    // The engine returned without draining; the write is still in flight.
+    expect(finished).toBe(false)
+
+    const drained = waitForChildFinalizations('invoker-1')
+    resolveWrite()
+    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')))
+
+    await expect(waitForChildFinalizations('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()
+  })
+
+  it('releases its entry so a run that never drains cannot leak', async () => {
+    const finalization = Promise.resolve()
+    trackChildFinalization('invoker-3', finalization)
+    await finalization
+    await Promise.resolve()
+
+    // Draining after self-cleanup still resolves rather than hanging.
+    await expect(waitForChildFinalizations('invoker-3')).resolves.toBeUndefined()
+  })
+})

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

@@ -142,3 +142,60 @@ export async function createChildCancellationSignal(params: {
     },
   }
 }
+
+/**
+ * Child-session finalizations 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.
+ *
+ * Only the immediate invoker is tracked. A finalization orphaned *below* another
+ * abandoned child still falls back to the reaper.
+ */
+const pendingChildFinalizations = new Map<string, Set<Promise<unknown>>>()
+
+/** Registers a child-session finalization against its invoking run. */
+export function trackChildFinalization(
+  invokerExecutionId: string | undefined,
+  promise: Promise<unknown>
+): void {
+  if (!invokerExecutionId) return
+
+  let pending = pendingChildFinalizations.get(invokerExecutionId)
+  if (!pending) {
+    pending = new Set()
+    pendingChildFinalizations.set(invokerExecutionId, pending)
+  }
+  pending.add(promise)
+
+  // 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(() => {})
+    .finally(() => {
+      pending.delete(promise)
+      if (pending.size === 0) pendingChildFinalizations.delete(invokerExecutionId)
+    })
+}
+
+/** Awaits every child-session finalization registered for this run. */
+export async function waitForChildFinalizations(
+  invokerExecutionId: string | undefined
+): Promise<void> {
+  if (!invokerExecutionId) return
+
+  const pending = pendingChildFinalizations.get(invokerExecutionId)
+  if (!pending) 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])
+  }
+  pendingChildFinalizations.delete(invokerExecutionId)
+}

apps/sim/lib/workflows/executor/execution-core.ts

@@ -22,6 +22,7 @@ import { redactLargeValueRefsInValue } from '@/lib/logs/execution/pii-large-valu
 import { redactObjectStrings } from '@/lib/logs/execution/pii-redaction'
 import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
 import { getUserEmailById } from '@/lib/users/queries'
+import { waitForChildFinalizations } from '@/lib/workflows/custom-blocks/child-execution'
 import { getCustomBlockRowsForWorkspace } from '@/lib/workflows/custom-blocks/operations'
 import {
   loadDeployedWorkflowState,
@@ -375,6 +376,10 @@ async function executeWorkflowCoreImpl(
     while (pendingLifecycleCallbacks.size > 0) {
       await Promise.allSettled([...pendingLifecycleCallbacks])
     }
+    // A custom block's child owns a separate log row whose terminal write the
+    // engine does not drain on cancel/timeout — await it here so the row is not
+    // left `running` for the reaper when this run finishes or the worker exits.
+    await waitForChildFinalizations(executionId)
   }
 
   try {