From ca0c499be8e76b5ac574f6e247544639fc389422 Mon Sep 17 00:00:00 2001 From: Alex Langenfeld Date: Wed, 2 Sep 2026 10:55:28 -0500 Subject: [PATCH] fix(core): drain step stream writes before completion Signed-off-by: Alex Langenfeld --- .changeset/step-stream-drain-barrier.md | 5 + .../docs/v5/configuration/runtime-tuning.mdx | 5 + packages/core/src/flushable-stream.test.ts | 126 +++++++ packages/core/src/flushable-stream.ts | 201 +++++++++++ .../core/src/runtime/step-executor.test.ts | 333 +++++++++++++++++- packages/core/src/runtime/step-executor.ts | 99 ++++-- packages/core/src/serialization.test.ts | 10 +- packages/core/src/serialization.ts | 63 ++-- packages/core/src/step/context-storage.ts | 2 + packages/core/src/step/writable-stream.ts | 31 +- 10 files changed, 805 insertions(+), 70 deletions(-) create mode 100644 .changeset/step-stream-drain-barrier.md diff --git a/.changeset/step-stream-drain-barrier.md b/.changeset/step-stream-drain-barrier.md new file mode 100644 index 0000000000..b09a73f5ff --- /dev/null +++ b/.changeset/step-stream-drain-barrier.md @@ -0,0 +1,5 @@ +--- +'@workflow/core': patch +--- + +Wait for writes queued by released step stream writers to reach durable storage before recording step completion. diff --git a/docs/content/docs/v5/configuration/runtime-tuning.mdx b/docs/content/docs/v5/configuration/runtime-tuning.mdx index fc4061de12..346631a98c 100644 --- a/docs/content/docs/v5/configuration/runtime-tuning.mdx +++ b/docs/content/docs/v5/configuration/runtime-tuning.mdx @@ -314,6 +314,11 @@ These variables are primarily for tests, debugging, or unusual deployments. - Group-commit window for the *leading* chunk of an idle stream. `0` sends it at once; a positive value holds it up to that many milliseconds to collect a group. This opt-in setting trades first-chunk latency for larger batches and can benefit slow-but-steady producers. Chunks arriving while a request is already in flight always coalesce into the next group regardless of this setting. - Also available as `streamFlushIntervalMs` on Worlds that expose it (the env var, when set, takes precedence over the World option). +### `WORKFLOW_STEP_STREAM_DRAIN_TIMEOUT_MS` + +- Default: `30000` (30 seconds) +- Maximum time a step waits for writes queued by a released workflow stream writer to be acknowledged by the Workflow server before recording `step_completed`. A timeout fails the step instead of exposing a completion while released-writer data remains client-side. A writer intentionally kept locked does not trigger this durability wait. + ### `WORKFLOW_STREAM_MAX_INFLIGHT_CHUNKS` - Default: `1000` diff --git a/packages/core/src/flushable-stream.test.ts b/packages/core/src/flushable-stream.test.ts index 2595a4b0ce..0ae01ef86f 100644 --- a/packages/core/src/flushable-stream.test.ts +++ b/packages/core/src/flushable-stream.test.ts @@ -1,10 +1,12 @@ import { afterEach, describe, expect, it } from 'vitest'; import { createFlushableState, + drainFlushableSnapshot, flushablePipe, LOCK_POLL_INTERVAL_MS, pollReadableLock, pollWritableLock, + trackFlushableWritable, } from './flushable-stream.js'; import { STREAM_DRAIN_SYMBOL } from './symbols.js'; @@ -446,6 +448,130 @@ describe('flushablePipe drain barrier (group-commit sinks)', () => { delete process.env.WORKFLOW_STREAM_MAX_BYTES_PER_BATCH; }); + it('waits for a produced frame to reach the sink before draining', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + let drained = false; + const sink = new WritableStream({ + async write() { + await writeGate; + }, + }); + Object.defineProperty(sink, STREAM_DRAIN_SYMBOL, { + value: async () => { + drained = true; + }, + }); + const transform = new TransformStream(); + const state = createFlushableState(); + const pipe = flushablePipe(transform.readable, sink, state).catch(() => {}); + const writable = trackFlushableWritable(transform.writable, state); + const writer = writable.getWriter(); + + await writer.write(new Uint8Array([1])); + const snapshot = drainFlushableSnapshot(state); + await tick(); + expect(drained).toBe(false); + + releaseWrite(); + await snapshot; + expect(drained).toBe(true); + + await writer.close(); + await pipe; + }); + + it('propagates an idle downstream failure through the tracked writable', async () => { + const downstreamError = new Error('downstream failed while producer idle'); + const transform = new TransformStream(); + const state = createFlushableState(); + const sink = new WritableStream({ + write() { + throw downstreamError; + }, + }); + const pipe = flushablePipe(transform.readable, sink, state).catch(() => {}); + const writable = trackFlushableWritable(transform.writable, state); + let sourceCancelled = false; + const source = new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array([1])); + }, + cancel() { + sourceCancelled = true; + }, + }); + + await expect(source.pipeTo(writable)).rejects.toThrow( + 'downstream failed while producer idle' + ); + expect(sourceCancelled).toBe(true); + await pipe; + }); + + it('drains an accepted prefix before rejecting its snapshot', async () => { + let releaseDrain!: () => void; + const drainGate = new Promise((resolve) => { + releaseDrain = resolve; + }); + const { sink } = makeDrainSink(() => drainGate); + const { source, controller } = makeControlledSource(); + const state = createFlushableState(); + state.producedFrames = 2; + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + controller().enqueue(new Uint8Array([1])); + await tick(); + controller().error(new Error('second frame failed')); + await tick(); + + let settled = false; + const snapshot = drainFlushableSnapshot(state).finally(() => { + settled = true; + }); + await tick(); + expect(settled).toBe(false); + + releaseDrain(); + await expect(snapshot).rejects.toThrow('second frame failed'); + await pipe; + }); + + it('rejects snapshot waiters with the upstream pipe error', async () => { + const { sink } = makeDrainSink(async () => {}); + const source = new ReadableStream({ + start(controller) { + controller.error(new Error('producer failed before acceptance')); + }, + }); + const state = createFlushableState(); + state.producedFrames = 1; + const pipe = flushablePipe(source, sink, state).catch(() => {}); + + const snapshot = drainFlushableSnapshot(state); + + await expect(snapshot).rejects.toThrow('producer failed before acceptance'); + await pipe; + }); + + it('reports a pipe error to a snapshot registered after failure', async () => { + const { sink } = makeDrainSink(async () => {}); + const source = new ReadableStream({ + start(controller) { + controller.error(new Error('producer already failed')); + }, + }); + const state = createFlushableState(); + state.producedFrames = 1; + await flushablePipe(source, sink, state).catch(() => {}); + + await expect(drainFlushableSnapshot(state)).rejects.toThrow( + 'producer already failed' + ); + }); + it('adopts the sink drain barrier onto the flushable state', async () => { const { sink } = makeDrainSink(async () => {}); const { source, controller } = makeControlledSource(); diff --git a/packages/core/src/flushable-stream.ts b/packages/core/src/flushable-stream.ts index 3842e2fb88..79b7cabef5 100644 --- a/packages/core/src/flushable-stream.ts +++ b/packages/core/src/flushable-stream.ts @@ -119,6 +119,30 @@ const getLockPollIntervalMs = (): number => export interface FlushableStreamState extends PromiseWithResolvers { /** Number of write operations currently in flight to the server */ pendingOps: number; + /** Frames emitted by this pipe's producer. */ + producedFrames: number; + /** Produced frames accepted by this pipe's sink. */ + acceptedFrames: number; + /** Terminal pipe error, retained for snapshots registered after failure. */ + pipeError?: unknown; + /** + * If the user-facing writable is unlocked, enqueue an ordered checkpoint and + * durably drain every write ahead of it. Returns false while a writer remains + * locked, so lock-held streams never block step completion. + */ + settleReleasedWrites?: () => Promise; + /** Whether release settlement waits for explicit step-end arming. */ + deferReleaseSettlement?: boolean; + /** Whether step-end processing has armed released-writer settlement. */ + releaseSettlementArmed?: boolean; + /** Whether the user-facing writable has begun a normal close. */ + userWritableClosing?: boolean; + /** Step-end snapshot waiters blocked until their target reaches the sink. */ + frameWaiters: Array<{ + target: number; + resolve: () => void; + reject: (error: unknown) => void; + }>; /** Whether the `done` promise has been resolved */ doneResolved: boolean; /** Whether the underlying stream has actually closed/errored */ @@ -140,6 +164,9 @@ export function createFlushableState(): FlushableStreamState { const state: FlushableStreamState = { ...withResolvers(), pendingOps: 0, + producedFrames: 0, + acceptedFrames: 0, + frameWaiters: [], doneResolved: false, streamEnded: false, }; @@ -229,6 +256,154 @@ function resolveAfterDrain(state: FlushableStreamState): void { ); } +/** Record a frame synchronously when its producer emits it into this pipe. */ +function markFlushableFrameProduced(state: FlushableStreamState): void { + state.producedFrames++; +} + +/** + * Capture a step-end producer watermark, wait until the pipe has handed every + * frame through that watermark to its sink, then await the sink's durability + * barrier. Unlike {@link FlushableStreamState.promise}, this never waits for a + * user writer lock to be released. + */ +export async function drainFlushableSnapshot( + state: FlushableStreamState +): Promise { + const drainBeforeThrow = async (error: unknown): Promise => { + // A later frame can fail after an earlier frame was early-acked by the + // group-commit sink. Keep the earlier accepted prefix durable before + // surfacing the producer error, matching flushablePipe's failure path. + await state.drainBarrier?.().catch(() => {}); + throw error; + }; + + const target = state.producedFrames; + if (state.acceptedFrames < target) { + if (state.pipeError !== undefined) { + return drainBeforeThrow(state.pipeError); + } + try { + await new Promise((resolve, reject) => { + state.frameWaiters.push({ target, resolve, reject }); + }); + } catch (error) { + return drainBeforeThrow(error); + } + } + if (state.pipeError !== undefined) { + return drainBeforeThrow(state.pipeError); + } + await state.drainBarrier?.(); +} + +/** + * Mark byte-stream chunks at the producer side of a flushable pipe. Serialized + * streams should instead use `getSerializeStream`'s synchronous output hook. + */ +/** + * Wrap the user-facing producer boundary of a flushable writable. A completed + * write through this handle has a sequence number before step-end snapshots, + * while the original writable remains the input to the serialization pipe. + */ +export function trackFlushableWritable( + writable: WritableStream, + state: FlushableStreamState, + WritableStreamConstructor: typeof WritableStream = WritableStream +): WritableStream { + const targetWriter = writable.getWriter(); + const checkpoint = Symbol('workflow-stream-release-checkpoint'); + let trackedController: WritableStreamDefaultController; + const tracked = new WritableStreamConstructor({ + start(controller) { + trackedController = controller; + }, + async write(chunk) { + if (chunk === checkpoint) return; + markFlushableFrameProduced(state); + await targetWriter.write(chunk); + }, + async close() { + state.userWritableClosing = true; + try { + await targetWriter.close(); + } finally { + targetWriter.releaseLock(); + } + }, + async abort(reason) { + try { + await targetWriter.abort(reason); + } finally { + targetWriter.releaseLock(); + } + }, + }) as WritableStream; + + // A downstream sink can fail after its early-acknowledged write has already + // resolved. Forward that terminal error into the public writable even while + // its producer is idle, so native pipeTo() rejects and cancels its source. + targetWriter.closed.catch((error) => { + trackedController.error(error); + }); + + let settlement: Promise | undefined; + state.settleReleasedWrites = (): Promise => { + state.releaseSettlementArmed = true; + if (settlement) return settlement; + if (tracked.locked) return Promise.resolve(false); + if (state.userWritableClosing) { + settlement = state.promise.then(() => true); + return settlement; + } + + let checkpointWriter: WritableStreamDefaultWriter; + try { + checkpointWriter = tracked.getWriter(); + } catch { + // Closed/errored streams settle through the pipe's normal completion. + return Promise.resolve(false); + } + + settlement = (async () => { + try { + // Native writable ordering puts this behind writes queued by the + // released writer, including write promises it did not await. + try { + await checkpointWriter.write(checkpoint as T); + } catch (checkpointError) { + // A close can move from requested to terminal between getWriter() + // and write(). Its normal pipe completion already drains the sink; + // an errored pipe rejects here with its actual failure instead. + try { + await state.promise; + return true; + } catch (pipeError) { + throw pipeError ?? checkpointError; + } + } + await drainFlushableSnapshot(state); + if (!state.doneResolved) { + state.doneResolved = true; + state.resolve(); + } + return true; + } catch (error) { + if (!state.doneResolved) { + state.doneResolved = true; + state.reject(error); + } + throw error; + } finally { + checkpointWriter.releaseLock(); + } + })(); + return settlement; + }; + + return tracked; +} + /** * Polls a WritableStream to check if the user has released their lock. * Resolves the done promise when lock is released and no pending ops remain. @@ -256,6 +431,20 @@ export function pollWritableLock( return; } + if (state.settleReleasedWrites) { + if ( + (!state.deferReleaseSettlement || state.releaseSettlementArmed) && + !writable.locked + ) { + clearInterval(intervalId); + state.writablePollingInterval = undefined; + void state.settleReleasedWrites().catch(() => { + // Failure is surfaced through state.promise. + }); + } + return; + } + // Check if lock is released (not closed) and no pending ops if (isWritableUnlockedNotClosed(writable) && state.pendingOps === 0) { state.doneResolved = true; @@ -393,13 +582,25 @@ async function flushablePipePerChunk( state.pendingOps++; try { await writer.write(readResult.value); + state.acceptedFrames++; + const ready = state.frameWaiters.filter( + (waiter) => waiter.target <= state.acceptedFrames + ); + state.frameWaiters = state.frameWaiters.filter( + (waiter) => waiter.target > state.acceptedFrames + ); + for (const waiter of ready) waiter.resolve(); } finally { state.pendingOps--; } } } catch (err) { state.streamEnded = true; + state.pipeError = err; cancelReason = err; + const frameWaiters = state.frameWaiters; + state.frameWaiters = []; + for (const waiter of frameWaiters) waiter.reject(err); // Against an early-ack sink, chunks can still be buffered or in flight // when the pipe fails (pendingOps only counts un-acked writes). Deliver // that accepted prefix before settling the failure: once the state diff --git a/packages/core/src/runtime/step-executor.test.ts b/packages/core/src/runtime/step-executor.test.ts index d62633b1db..04fa881194 100644 --- a/packages/core/src/runtime/step-executor.test.ts +++ b/packages/core/src/runtime/step-executor.test.ts @@ -5,14 +5,18 @@ import type { Event, World } from '@workflow/world'; import { SPEC_VERSION_CURRENT } from '@workflow/world'; import { createWorld } from '@workflow/world-local'; import { afterEach, describe, expect, it, vi } from 'vitest'; +import { LOCK_POLL_INTERVAL_MS } from '../flushable-stream.js'; import { registerStepFunction } from '../private.js'; import { dehydrateStepArguments, hydrateStepError } from '../serialization.js'; +import { getWritable } from '../step/writable-stream.js'; +import { STREAM_NAME_SYMBOL, STREAM_SERVER_RUN_ID_SYMBOL } from '../symbols.js'; import { COMPUTE_INSTANCE_ID } from './compute-instance.js'; import { executeStep } from './step-executor.js'; import { UNSERIALIZABLE_STEP_INPUT_MARKER, unserializableStepInputPlaceholder, } from './unserializable-step.js'; +import { setWorld } from './world.js'; // The retry ceiling (`authoritativeAttempt`) is what bounds a step that keeps // timing out: a timeout hard-kills the body without writing any error, so the @@ -34,8 +38,16 @@ async function setupRunningStep(opts: { onBody: () => void; register?: boolean; createStep?: boolean; + stepArgs?: unknown[]; }): Promise<{ runId: string; stepId: string }> { - const { world, stepName, onBody, register = true, createStep = true } = opts; + const { + world, + stepName, + onBody, + register = true, + createStep = true, + stepArgs = [], + } = opts; const runInput = await dehydrateStepArguments([], 'run', undefined); const created = await world.events.create(null, { eventType: 'run_created', @@ -55,7 +67,11 @@ async function setupRunningStep(opts: { const stepId = 'step_timeout_1'; if (createStep) { - const stepInput = await dehydrateStepArguments([], runId, undefined); + const stepInput = await dehydrateStepArguments( + { args: stepArgs, closureVars: undefined, thisVal: undefined }, + runId, + undefined + ); await world.events.create(runId, { eventType: 'step_created', specVersion: SPEC_VERSION_CURRENT, @@ -83,6 +99,64 @@ function makeWorld(): World { return createWorld({ dataDir, tag: `t${counter}` }); } +async function runWritableStep(options: { + releaseLock: boolean; + awaitWrite?: boolean; + delayBeforeWriterMs?: number; + closeAfterRelease?: boolean; + writeImpl?: () => Promise; +}): Promise<{ + execution: Promise>>; + world: World; + runId: string; + stepId: string; +}> { + const world = makeWorld(); + setWorld(world); + if (options.writeImpl) { + world.streams.write = vi.fn( + options.writeImpl + ) as typeof world.streams.write; + } + + const stepName = uniqueStepName(); + const { runId, stepId } = await setupRunningStep({ + world, + stepName, + onBody: () => {}, + register: false, + }); + registerStepFunction(stepName, async () => { + const writable = getWritable(); + if (options.delayBeforeWriterMs) { + await new Promise((resolve) => + setTimeout(resolve, options.delayBeforeWriterMs) + ); + } + const writer = writable.getWriter(); + const write = writer.write('snapshot'); + if (options.awaitWrite !== false) await write; + if (options.releaseLock) writer.releaseLock(); + if (options.closeAfterRelease) await writable.close(); + return 'ok'; + }); + + return { + execution: executeStep({ + world, + workflowRunId: runId, + workflowName: 'wf', + workflowStartedAt: Date.now(), + stepId, + stepName, + authoritativeAttempt: 1, + }), + world, + runId, + stepId, + }; +} + async function eventsFor( world: World, runId: string, @@ -95,6 +169,261 @@ async function eventsFor( ); } +describe('executeStep — stream durability barrier', () => { + afterEach(() => { + setWorld(undefined); + delete process.env.WORKFLOW_STEP_STREAM_DRAIN_TIMEOUT_MS; + counter += 1; + }); + + it('writes step_completed only after a released writer drains', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + const { execution, world, runId, stepId } = await runWritableStep({ + releaseLock: true, + writeImpl: () => writeGate, + }); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(0); + + releaseWrite(); + await expect(execution).resolves.toMatchObject({ + type: 'completed', + hasPendingOps: false, + }); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(1); + }); + + it('does not durably block a step that keeps its writer lock', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + const { execution } = await runWritableStep({ + releaseLock: false, + awaitWrite: false, + writeImpl: () => writeGate, + }); + + await expect(execution).resolves.toMatchObject({ + type: 'completed', + hasPendingOps: true, + }); + releaseWrite(); + }); + + it('does not settle before the step acquires and releases its writer', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + const { execution, world, runId, stepId } = await runWritableStep({ + releaseLock: true, + delayBeforeWriterMs: LOCK_POLL_INTERVAL_MS * 3, + writeImpl: () => writeGate, + }); + + await new Promise((resolve) => + setTimeout(resolve, LOCK_POLL_INTERVAL_MS * 5) + ); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(0); + + releaseWrite(); + await expect(execution).resolves.toMatchObject({ type: 'completed' }); + }); + + it('orders unsettled writes before the release checkpoint', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + const { execution, world, runId, stepId } = await runWritableStep({ + releaseLock: true, + awaitWrite: false, + writeImpl: () => writeGate, + }); + + await new Promise((resolve) => setTimeout(resolve, 520)); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(0); + + releaseWrite(); + await expect(execution).resolves.toMatchObject({ type: 'completed' }); + }); + + it('drains a revived forwarded writable argument before completion', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + const world = makeWorld(); + setWorld(world); + world.streams.write = vi.fn(() => writeGate) as typeof world.streams.write; + + const forwarded = new WritableStream(); + Object.defineProperty(forwarded, STREAM_NAME_SYMBOL, { + value: 'strm_forwarded', + }); + Object.defineProperty(forwarded, STREAM_SERVER_RUN_ID_SYMBOL, { + value: 'wrun_forwarded_owner', + }); + const stepName = uniqueStepName(); + const { runId, stepId } = await setupRunningStep({ + world, + stepName, + onBody: () => {}, + register: false, + stepArgs: [forwarded], + }); + registerStepFunction(stepName, async (writable: WritableStream) => { + const writer = writable.getWriter(); + await writer.write('forwarded snapshot'); + writer.releaseLock(); + return 'ok'; + }); + + const execution = executeStep({ + world, + workflowRunId: runId, + workflowName: 'wf', + workflowStartedAt: Date.now(), + stepId, + stepName, + authoritativeAttempt: 1, + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(0); + + releaseWrite(); + await expect(execution).resolves.toMatchObject({ + type: 'completed', + hasPendingOps: false, + }); + }); + + it('allows a released writable to close normally before step end', async () => { + const { execution, world, runId, stepId } = await runWritableStep({ + releaseLock: true, + closeAfterRelease: true, + }); + + await expect(execution).resolves.toMatchObject({ + type: 'completed', + hasPendingOps: false, + }); + expect(await eventsFor(world, runId, stepId, 'step_retrying')).toHaveLength( + 0 + ); + expect(await eventsFor(world, runId, stepId, 'step_failed')).toHaveLength( + 0 + ); + }); + + it('does not complete successfully when the drain times out', async () => { + process.env.WORKFLOW_STEP_STREAM_DRAIN_TIMEOUT_MS = '10'; + const { execution, world, runId, stepId } = await runWritableStep({ + releaseLock: true, + writeImpl: () => new Promise(() => {}), + }); + + await expect(execution).resolves.toMatchObject({ type: 'retry' }); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(0); + }); + + it('does not complete successfully when the drain fails', async () => { + const { execution, world, runId, stepId } = await runWritableStep({ + releaseLock: true, + writeImpl: async () => { + throw new Error('stream write failed'); + }, + }); + + await expect(execution).resolves.toMatchObject({ type: 'retry' }); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(0); + }); + + it('an aborted stream does not bypass another stream drain', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + const world = makeWorld(); + setWorld(world); + world.streams.write = vi.fn(async (_runId, name) => { + if (name.endsWith('_aborted')) { + throw Object.assign(new Error('client disconnected'), { + name: 'AbortError', + }); + } + await writeGate; + }) as typeof world.streams.write; + + const stepName = uniqueStepName(); + const { runId, stepId } = await setupRunningStep({ + world, + stepName, + onBody: () => {}, + register: false, + }); + registerStepFunction(stepName, async () => { + const aborted = getWritable({ namespace: 'aborted' }).getWriter(); + const durable = getWritable({ namespace: 'durable' }).getWriter(); + await aborted.write('a'); + await durable.write('b'); + aborted.releaseLock(); + durable.releaseLock(); + return 'ok'; + }); + + const execution = executeStep({ + world, + workflowRunId: runId, + workflowName: 'wf', + workflowStartedAt: Date.now(), + stepId, + stepName, + authoritativeAttempt: 1, + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect( + await eventsFor(world, runId, stepId, 'step_completed') + ).toHaveLength(0); + + releaseWrite(); + await expect(execution).resolves.toMatchObject({ type: 'completed' }); + }); + + it.each([ + 'AbortError', + 'ResponseAborted', + ])('tolerates a client disconnect named %s during drain', async (name) => { + const { execution } = await runWritableStep({ + releaseLock: true, + writeImpl: async () => { + throw Object.assign(new Error('client disconnected'), { name }); + }, + }); + + await expect(execution).resolves.toMatchObject({ type: 'completed' }); + }); +}); + describe('executeStep — retry ceiling (authoritativeAttempt)', () => { afterEach(() => { counter += 1; diff --git a/packages/core/src/runtime/step-executor.ts b/packages/core/src/runtime/step-executor.ts index a4136268b8..7ba2e7b15f 100644 --- a/packages/core/src/runtime/step-executor.ts +++ b/packages/core/src/runtime/step-executor.ts @@ -29,6 +29,8 @@ import { SPEC_VERSION_CURRENT, SPEC_VERSION_SUPPORTS_COMPRESSION, } from '@workflow/world'; +import { envNumber } from '@workflow/world/env-config'; +import type { FlushableStreamState } from '../flushable-stream.js'; import { runtimeLogger, stepLogger } from '../logger.js'; import { getStepFunction } from '../private.js'; import type { PayloadKey } from '../serialization/encryption.js'; @@ -74,6 +76,61 @@ import { isUnserializableStepInputPlaceholder } from './unserializable-step.js'; import { safeWaitUntil } from './wait-until.js'; export const DEFAULT_STEP_MAX_RETRIES = 3; +export const STEP_STREAM_DRAIN_TIMEOUT_MS = 30_000; + +export function getStepStreamDrainTimeoutMs(): number { + return envNumber( + 'WORKFLOW_STEP_STREAM_DRAIN_TIMEOUT_MS', + STEP_STREAM_DRAIN_TIMEOUT_MS, + { integer: true, min: 1 } + ); +} + +function isClientDisconnectError(error: unknown): boolean { + const name = (error as { name?: unknown })?.name; + return name === 'AbortError' || name === 'ResponseAborted'; +} + +async function settleReleasedStepStreams( + states: FlushableStreamState[] +): Promise { + if (states.length === 0) return; + + const timeoutMs = getStepStreamDrainTimeoutMs(); + let timeout: ReturnType | undefined; + try { + await Promise.race([ + Promise.all( + states.map((state) => + (state.settleReleasedWrites?.() ?? Promise.resolve(false)).catch( + (error) => { + // A disconnected client may abandon one response stream, but it + // must not let that rejection bypass durability waits for other + // streams written by the same step. + if (!isClientDisconnectError(error)) throw error; + return false; + } + ) + ) + ), + new Promise((_, reject) => { + timeout = setTimeout( + () => + reject( + new WorkflowRuntimeError( + `Timed out draining step stream writes after ${timeoutMs}ms` + ) + ), + timeoutMs + ); + }), + ]); + } catch (error) { + if (!isClientDisconnectError(error)) throw error; + } finally { + if (timeout !== undefined) clearTimeout(timeout); + } +} /** * Extract the inline delta from a step-terminal `events.create` result, @@ -937,6 +994,7 @@ export async function executeStep( // outside the try so the failure path below can also drain them. const preCompletionOps: Promise[] = []; const ops: Promise[] = []; + const streamStates: FlushableStreamState[] = []; let opsSettled = true; // Latency telemetry to attach to this step's terminal event. Computed @@ -992,7 +1050,8 @@ export async function executeStep( ops, globalThis, {}, - params.workflowDeploymentId + params.workflowDeploymentId, + streamStates ); const durationMs = Date.now() - startTime; hydrateSpan?.setAttributes({ @@ -1129,6 +1188,7 @@ export async function executeStep( rootRunId: params.rootRunId, ops, preCompletionOps, + streamStates, closureVars: hydratedInput.closureVars, encryptionKey, // Turbo optimistic start runs this body before `run_started` is @@ -1202,21 +1262,12 @@ export async function executeStep( return dehydrated; }); - // Flush pending ops (stream writes, etc.) with a short inline wait. - // WorkflowServerWritableStream acks writes on buffer entry - // (group-commit batching); durability is enforced by its drain - // barrier, which the flushable state's completion awaits after - // lock release. Most ops settle within ~200ms (lock-release - // polling + one batched HTTP flush). - // If ops don't settle in 500ms (e.g., WritableStream kept open - // across steps), waitUntil handles the rest. + // Arm the background flush before the durability wait so lock-held + // streams and late close/writes keep their existing lifecycle even if a + // drain fails. The drain snapshots only frames produced by this step and + // does not depend on writer-lock release. if (ops.length > 0) { const opsPromise = Promise.all(ops); - // The race below surfaces failures inline when ops settle quickly; - // if the 500ms timeout wins, the failure is only observed here. The - // promise handed to waitUntil must never reject (an unconsumed - // waitUntil rejection crashes the process as unhandledRejection), - // so unexpected failures are logged instead. safeWaitUntil(opsPromise, (err) => { runtimeLogger.warn('Background flush of step stream ops failed', { workflowRunId, @@ -1224,20 +1275,28 @@ export async function executeStep( error: err instanceof Error ? err.message : String(err), }); }); - opsSettled = await Promise.race([ + + // Start the V2 inline-loop heuristic concurrently with durability so + // a held lock costs max(500ms, PUT RTT), not 500ms plus the PUT RTT. + const opsSettledPromise = Promise.race([ opsPromise.then( () => true as const, (err) => { - // Ignore expected client disconnect errors (e.g., browser - // refresh during streaming) - const isAbortError = - err?.name === 'AbortError' || err?.name === 'ResponseAborted'; - if (isAbortError) return true as const; + if (isClientDisconnectError(err)) return true as const; throw err; } ), new Promise((r) => setTimeout(() => r(false), 500)), ]); + // The durability wait can outlive an immediate op rejection. Observe + // this branch now; the awaited copy below still surfaces the error. + opsSettledPromise.catch(() => {}); + + await settleReleasedStepStreams(streamStates); + + // This outcome is only the inline-loop heuristic. Durability was + // established independently above. + opsSettled = await opsSettledPromise; } // Optimistic start: the body ran before `step_started` was confirmed. diff --git a/packages/core/src/serialization.test.ts b/packages/core/src/serialization.test.ts index fab0d3f68b..af2d03cb06 100644 --- a/packages/core/src/serialization.test.ts +++ b/packages/core/src/serialization.test.ts @@ -12,6 +12,10 @@ import { WORKFLOW_DESERIALIZE, WORKFLOW_SERIALIZE } from '@workflow/serde'; import { beforeAll, describe, expect, it, vi } from 'vitest'; import { registerSerializationClass } from './class-serialization.js'; import { decrypt, encrypt, importKey } from './encryption.js'; +import { + drainFlushableSnapshot, + type FlushableStreamState, +} from './flushable-stream.js'; import { getStepFunction, registerStepFunction } from './private.js'; import { bytesToBase64, deriveRunKeyPair } from './sealed-box.js'; import { @@ -592,6 +596,7 @@ describe('workflow arguments', () => { noEncryptionKey ); const ops: Promise[] = []; + const streamStates: FlushableStreamState[] = []; const hydrated = (await hydrateStepArguments( serialized, 'wrun_child', @@ -599,11 +604,14 @@ describe('workflow arguments', () => { ops, globalThis, {}, - 'dpl_child' + 'dpl_child', + streamStates )) as WritableStream; const writer = hydrated.getWriter(); await writer.write('cross-deployment'); + expect(streamStates).toHaveLength(1); + await Promise.all(streamStates.map(drainFlushableSnapshot)); await writer.close(); await Promise.all(ops); diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index 0fc3b2b1a0..9c794f52cd 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -13,6 +13,7 @@ import { monotonicFactory } from 'ulid'; import { importKey } from './encryption.js'; import { createFlushableState, + type FlushableStreamState, flushablePipe, getMaxBufferedBytes, getMaxBytesPerBatch, @@ -20,6 +21,7 @@ import { getMaxInflightChunks, pollReadableLock, pollWritableLock, + trackFlushableWritable, } from './flushable-stream.js'; import { getStepFunction } from './private.js'; // V2: use getWorldLazy in step-side code paths so Turbopack can statically @@ -3287,7 +3289,8 @@ function getStepRevivers( ops: Promise[], runId: string, cryptoKey: EncryptionKeyParam, - deploymentId?: string + deploymentId?: string, + streamStates?: FlushableStreamState[] ): Partial { return { ...getCommonRevivers(global), @@ -3509,6 +3512,8 @@ function getStepRevivers( // Create flushable state for this stream const state = createFlushableState(); + state.deferReleaseSettlement = streamStates !== undefined; + streamStates?.push(state); ops.push(state.promise); // Start the flushable pipe in the background @@ -3516,8 +3521,13 @@ function getStepRevivers( // Errors are handled via state.reject }); - // Start polling to detect when user releases lock - pollWritableLock(serialize.writable, state); + // Track completed user writes independently of lock release. + const writable = trackFlushableWritable( + serialize.writable, + state, + global.WritableStream + ); + pollWritableLock(writable, state); // Record the underlying `(runId, name)` so downstream reducers can // recognize that this writable is already backed by a workflow @@ -3525,23 +3535,19 @@ function getStepRevivers( // the child passes this writable on to a grandchild), the // external reducer needs both to emit the original `runId` in // the descriptor. - Object.defineProperty(serialize.writable, STREAM_NAME_SYMBOL, { + Object.defineProperty(writable, STREAM_NAME_SYMBOL, { value: value.name, writable: false, }); - Object.defineProperty(serialize.writable, STREAM_SERVER_RUN_ID_SYMBOL, { + Object.defineProperty(writable, STREAM_SERVER_RUN_ID_SYMBOL, { value: targetRunId, writable: false, }); if (targetDeploymentId) { - Object.defineProperty( - serialize.writable, - STREAM_SERVER_DEPLOYMENT_ID_SYMBOL, - { - value: targetDeploymentId, - writable: false, - } - ); + Object.defineProperty(writable, STREAM_SERVER_DEPLOYMENT_ID_SYMBOL, { + value: targetDeploymentId, + writable: false, + }); } // Keep the owner's public key on the handle so a further forward stays on // the zero-lookup sealed path. @@ -3565,27 +3571,19 @@ function getStepRevivers( targetRunId === runId && isRunPayloadKeys(cryptoKey) ) { - Object.defineProperty( - serialize.writable, - STREAM_SERVER_PUBLIC_KEY_SYMBOL, - { - value: bytesToBase64(cryptoKey.keyPair.publicKey), - writable: false, - } - ); + Object.defineProperty(writable, STREAM_SERVER_PUBLIC_KEY_SYMBOL, { + value: bytesToBase64(cryptoKey.keyPair.publicKey), + writable: false, + }); } if (typeof value.encryptionPublicKey === 'string') { - Object.defineProperty( - serialize.writable, - STREAM_SERVER_PUBLIC_KEY_SYMBOL, - { - value: value.encryptionPublicKey, - writable: false, - } - ); + Object.defineProperty(writable, STREAM_SERVER_PUBLIC_KEY_SYMBOL, { + value: value.encryptionPublicKey, + writable: false, + }); } - return serialize.writable; + return writable; }, AbortController: (value) => reviveAbortController(value, ops, runId), @@ -3941,14 +3939,15 @@ export async function hydrateStepArguments( ops: Promise[] = [], global: Record = globalThis, extraRevivers: Record any> = {}, - deploymentId?: string + deploymentId?: string, + streamStates?: FlushableStreamState[] ): Promise { const compressionStats: CompressionStats = {}; const result = await stepModule.deserialize(value, key, { global, extraRevivers: { ...getStreamAndRequestRevivers( - getStepRevivers(global, ops, runId, key, deploymentId) + getStepRevivers(global, ops, runId, key, deploymentId, streamStates) ), ...extraRevivers, }, diff --git a/packages/core/src/step/context-storage.ts b/packages/core/src/step/context-storage.ts index c6a33c28d5..572baf7338 100644 --- a/packages/core/src/step/context-storage.ts +++ b/packages/core/src/step/context-storage.ts @@ -59,6 +59,8 @@ export type StepContext = { closureVars?: Record; encryptionKey?: PayloadKey; writables?: Map; + /** Server-bound stream pipes whose current frames must drain at step end. */ + streamStates?: FlushableStreamState[]; /** * Turbo mode only: a promise that resolves once the backgrounded * `run_started` has landed (the run exists). Set when the step body runs diff --git a/packages/core/src/step/writable-stream.ts b/packages/core/src/step/writable-stream.ts index b3bfd264ea..5adf82b679 100644 --- a/packages/core/src/step/writable-stream.ts +++ b/packages/core/src/step/writable-stream.ts @@ -3,6 +3,7 @@ import { createFlushableState, flushablePipe, pollWritableLock, + trackFlushableWritable, } from '../flushable-stream.js'; import { bytesToBase64 } from '../sealed-box.js'; import { @@ -81,6 +82,8 @@ export function getWritable( // The target run is the workflow run that owns this step, which (per // version skew protection) is on this same SDK version, so byte-stream // framing is always safe here. + const state = createFlushableState(); + state.deferReleaseSettlement = ctx.streamStates !== undefined; const serialize = getSerializeStream( // In turbo optimistic start the body runs before `run_started` is durable. // Thread the run-ready barrier so that a nested ReadableStream written into @@ -109,14 +112,16 @@ export function getWritable( name, ctx.runReadyBarrier ); - const state = createFlushableState(); + ctx.streamStates ??= []; + ctx.streamStates.push(state); ctx.ops.push(state.promise); flushablePipe(serialize.readable, serverWritable, state).catch(() => { // Errors are handled via state.reject }); - pollWritableLock(serialize.writable, state); + const writable = trackFlushableWritable(serialize.writable, state); + pollWritableLock(writable, state); // Tag the writable with its underlying `(runId, name)` so downstream // reducers can recognize that it's already backed by a workflow @@ -125,36 +130,32 @@ export function getWritable( // dehydrated descriptor, so the child's reviver can open the // writable against the original `(runId, name)` directly, with no // in-process bridge tied to this step's lifetime. - Object.defineProperty(serialize.writable, STREAM_NAME_SYMBOL, { + Object.defineProperty(writable, STREAM_NAME_SYMBOL, { value: name, writable: false, }); - Object.defineProperty(serialize.writable, STREAM_SERVER_RUN_ID_SYMBOL, { + Object.defineProperty(writable, STREAM_SERVER_RUN_ID_SYMBOL, { value: runId, writable: false, }); if (ctx.workflowDeploymentId) { - Object.defineProperty( - serialize.writable, - STREAM_SERVER_DEPLOYMENT_ID_SYMBOL, - { - value: ctx.workflowDeploymentId, - writable: false, - } - ); + Object.defineProperty(writable, STREAM_SERVER_DEPLOYMENT_ID_SYMBOL, { + value: ctx.workflowDeploymentId, + writable: false, + }); } // Publish this run's X25519 public key on the handle so that a run this // writable is forwarded to can seal frames without looking anything up. // The key is already resolved on the step context, so this costs nothing // here and saves the receiver a round trip. if (isRunPayloadKeys(ctx.encryptionKey)) { - Object.defineProperty(serialize.writable, STREAM_SERVER_PUBLIC_KEY_SYMBOL, { + Object.defineProperty(writable, STREAM_SERVER_PUBLIC_KEY_SYMBOL, { value: bytesToBase64(ctx.encryptionKey.keyPair.publicKey), writable: false, }); } - cache.set(name, { writable: serialize.writable, state }); + cache.set(name, { writable, state }); - return serialize.writable as WritableStream; + return writable as WritableStream; }