From d63263ee3d080132795a55648b31e4894832459f Mon Sep 17 00:00:00 2001 From: Alex Langenfeld Date: Thu, 27 Aug 2026 17:24:13 -0500 Subject: [PATCH] [core] Prefetch stream read encryption keys Signed-off-by: Alex Langenfeld --- .changeset/stream-read-key-prefetch.md | 5 + .../src/reconnecting-framed-stream.test.ts | 158 ++++++++++++++++++ packages/core/src/runtime/run.ts | 23 +-- packages/core/src/runtime/runs.test.ts | 156 +++++++++++++++++ packages/core/src/serialization.ts | 148 ++++++++++++++-- 5 files changed, 470 insertions(+), 20 deletions(-) create mode 100644 .changeset/stream-read-key-prefetch.md diff --git a/.changeset/stream-read-key-prefetch.md b/.changeset/stream-read-key-prefetch.md new file mode 100644 index 0000000000..66b20fa841 --- /dev/null +++ b/.changeset/stream-read-key-prefetch.md @@ -0,0 +1,5 @@ +--- +'@workflow/core': patch +--- + +Prefetch run encryption keys when reading workflow streams. diff --git a/packages/core/src/reconnecting-framed-stream.test.ts b/packages/core/src/reconnecting-framed-stream.test.ts index ce22485f7b..6b2805a77e 100644 --- a/packages/core/src/reconnecting-framed-stream.test.ts +++ b/packages/core/src/reconnecting-framed-stream.test.ts @@ -461,6 +461,164 @@ describe('createReconnectingFramedStream', () => { expect(chunks).toEqual([payloadFrame(1), payloadFrame(2), payloadFrame(3)]); }); + it('starts key resolution concurrently with the first stream GET', async () => { + let resolveStream: (stream: ReadableStream) => void; + const streamPromise = new Promise>((resolve) => { + resolveStream = resolve; + }); + let resolveKey: () => void; + const keyPromise = new Promise((resolve) => { + resolveKey = resolve; + }); + const get = vi.fn().mockReturnValue(streamPromise); + const prefetchKey = vi.fn().mockReturnValue(keyPromise); + setWorld({ + specVersion: SPEC_VERSION_CURRENT, + streams: { get }, + } as unknown as World); + + const read = readAll( + createReconnectingFramedStream(RUN_ID, 's', 0, prefetchKey) + ); + await vi.waitFor(() => { + expect(get).toHaveBeenCalledOnce(); + expect(prefetchKey).toHaveBeenCalledOnce(); + }); + + resolveKey?.(); + resolveStream?.( + scriptedStream([ + { kind: 'value', value: payloadFrame(7) }, + { kind: 'close' }, + ]) + ); + await expect(read).resolves.toEqual([payloadFrame(7)]); + }); + + it('finishes key resolution before the first raw frame', async () => { + let releaseFrame: () => void; + const frameReady = new Promise((resolve) => { + releaseFrame = resolve; + }); + const prefetchKey = vi.fn().mockResolvedValue(undefined); + const { world } = makeWorldWithScriptedStreams({ + 0: () => + new ReadableStream({ + async pull(controller) { + await frameReady; + controller.enqueue(payloadFrame(7)); + controller.close(); + }, + }), + }); + setWorld(world); + + const read = readAll( + createReconnectingFramedStream(RUN_ID, 's', 0, prefetchKey) + ); + await vi.waitFor(() => expect(prefetchKey).toHaveBeenCalledOnce()); + // The resolver has already settled by the time the raw frame is released. + await Promise.resolve(); + releaseFrame?.(); + await expect(read).resolves.toEqual([payloadFrame(7)]); + expect(prefetchKey).toHaveBeenCalledOnce(); + }); + + it('prefetches one key promise across reconnects', async () => { + const prefetchKey = vi.fn().mockResolvedValue(undefined); + const { world, calls } = makeWorldWithScriptedStreams({ + 0: () => + scriptedStream([ + { kind: 'value', value: payloadFrame(1) }, + { kind: 'error', err: new Error('connection reset') }, + ]), + 1: () => + scriptedStream([ + { kind: 'value', value: payloadFrame(2) }, + { kind: 'close' }, + ]), + }); + setWorld(world); + + await expect( + readAll(createReconnectingFramedStream(RUN_ID, 's', 0, prefetchKey)) + ).resolves.toEqual([payloadFrame(1), payloadFrame(2)]); + expect(calls).toEqual([0, 1]); + expect(prefetchKey).toHaveBeenCalledOnce(); + }); + + it('keeps a stream GET failure primary when its speculative key lookup also fails', async () => { + const streamError = new Error('stream connection failed'); + const keyError = new Error('key lookup failed'); + const unhandled = vi.fn(); + process.once('unhandledRejection', unhandled); + setWorld({ + specVersion: SPEC_VERSION_CURRENT, + streams: { get: vi.fn().mockRejectedValue(streamError) }, + } as unknown as World); + + await expect( + readAll( + createReconnectingFramedStream(RUN_ID, 's', -1, () => + Promise.reject(keyError) + ) + ) + ).rejects.toThrow('stream connection failed'); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(unhandled).not.toHaveBeenCalled(); + }); + + it('observes a rejected speculative key lookup after cancellation', async () => { + const keyError = new Error('key lookup failed'); + const unhandled = vi.fn(); + process.once('unhandledRejection', unhandled); + const stream = createReconnectingFramedStream(RUN_ID, 's', 0, () => + Promise.reject(keyError) + ); + const reader = stream.getReader(); + const pending = reader.read(); + await reader.cancel(); + await expect(pending).resolves.toMatchObject({ done: true }); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(unhandled).not.toHaveBeenCalled(); + }); + + it('cancels an acquired underlying reader while a prefetched key is pending', async () => { + let cancelCount = 0; + let resolveKey: () => void; + const prefetchKey = vi.fn().mockReturnValue( + new Promise((resolve) => { + resolveKey = resolve; + }) + ); + const source = new ReadableStream({ + pull() { + // Keep the first raw read pending until the consumer cancels. + }, + cancel() { + cancelCount++; + }, + }); + const get = vi.fn().mockResolvedValue(source); + setWorld({ + specVersion: SPEC_VERSION_CURRENT, + streams: { get }, + } as unknown as World); + + const reader = createReconnectingFramedStream( + RUN_ID, + 's', + 0, + prefetchKey + ).getReader(); + const pendingRead = reader.read(); + await vi.waitFor(() => expect(get).toHaveBeenCalledOnce()); + await reader.cancel(); + await expect(pendingRead).resolves.toMatchObject({ done: true }); + expect(cancelCount).toBe(1); + resolveKey?.(); + }); + it('threads runId through to streams.get', async () => { const getSpy = vi.fn( async (_runId: string, _name: string, _startIndex?: number) => diff --git a/packages/core/src/runtime/run.ts b/packages/core/src/runtime/run.ts index 08afacf859..2f116612fe 100644 --- a/packages/core/src/runtime/run.ts +++ b/packages/core/src/runtime/run.ts @@ -13,7 +13,7 @@ import { type PayloadKey, } from '../serialization/encryption.js'; import { - getExternalRevivers, + getRunReadableStream, hydrateRunError, hydrateWorkflowReturnValue, } from '../serialization.js'; @@ -224,9 +224,9 @@ export class Run { } /** - * Defer fetching the run and its encryption key until serialized stream data - * is actually read. An empty or metadata-only stream must not start an - * unobserved run lookup. + * Defers fetching the run and its encryption key until a readable is first + * consumed. The first pull prefetches it so an encrypted first frame can join + * the lookup; this also applies to an empty consumed stream. * @internal */ #getEncryptionKeyLazily(): () => Promise { @@ -370,15 +370,18 @@ export class Run { 'use step'; const { ops = [], global = globalThis, startIndex, namespace } = options; const name = getWorkflowRunStreamId(this.runId, namespace); - // The resolver starts only when the deserialize stream sees its first - // chunk, so creating or probing an empty stream cannot reject in the - // background. + // The resolver starts only on the readable's first pull, so construction + // is inert. A consumed empty stream still performs the speculative lookup + // to keep the first encrypted-frame path concurrent. const encryptionKey = this.#getEncryptionKeyLazily(); - const stream = getExternalRevivers(global, ops, this.runId, encryptionKey) - .ReadableStream!({ + const stream = getRunReadableStream( + global, + ops, + this.runId, name, startIndex, - }) as ReadableStream; + encryptionKey + ); const worldPromise = this.#lazyWorldPromise; const runId = this.runId; diff --git a/packages/core/src/runtime/runs.test.ts b/packages/core/src/runtime/runs.test.ts index 10b540f100..b57d13464f 100644 --- a/packages/core/src/runtime/runs.test.ts +++ b/packages/core/src/runtime/runs.test.ts @@ -25,10 +25,12 @@ vi.mock('../serialization.js', async (importActual) => { }); import { registerSerializationClass } from '../class-serialization.js'; +import { deriveRunPayloadKeys } from '../serialization/encryption.js'; import { dehydrateRunError, dehydrateStepReturnValue, dehydrateWorkflowReturnValue, + getSerializeStream, hydrateStepReturnValue, } from '../serialization.js'; import { getReturnValuePollIntervalMs, Run } from './run.js'; @@ -267,6 +269,21 @@ describe('Run.getReadable', () => { setWorld(undefined as unknown as World); }); + async function encryptedFrames(value: unknown, material: Uint8Array) { + const serialize = getSerializeStream( + {}, + await deriveRunPayloadKeys(material) + ); + const reader = serialize.readable.getReader(); + const read = reader.read(); + const writer = serialize.writable.getWriter(); + await writer.write(value); + await writer.close(); + const first = await read; + if (!first.value) throw new Error('Expected serialized frame'); + return first.value; + } + it('does not fetch the run encryption key for an empty stream', async () => { const world = createMockWorld(); world.getEncryptionKeyForRun = vi.fn().mockResolvedValue(undefined); @@ -284,9 +301,148 @@ describe('Run.getReadable', () => { new Run('wrun_123').getReadable(); await new Promise((resolve) => setTimeout(resolve, 0)); + expect(world.streams.get).not.toHaveBeenCalled(); expect(world.runs.get).not.toHaveBeenCalled(); expect(world.getEncryptionKeyForRun).not.toHaveBeenCalled(); }); + + it('prefetches the run key when a consumed stream is empty', async () => { + const world = createMockWorld(); + world.getEncryptionKeyForRun = vi.fn().mockResolvedValue(undefined); + world.streams = { + get: vi.fn().mockResolvedValue( + new ReadableStream({ + start(controller) { + controller.close(); + }, + }) + ), + } as unknown as World['streams']; + setWorld(world); + + await expect( + new Run('wrun_123').getReadable().getReader().read() + ).resolves.toMatchObject({ done: true }); + expect(world.streams.get).toHaveBeenCalledOnce(); + expect(world.runs.get).toHaveBeenCalledOnce(); + expect(world.getEncryptionKeyForRun).toHaveBeenCalledOnce(); + }); + + it('resolves supplied ops when the caller releases an open readable lock', async () => { + const ops: Promise[] = []; + const material = new Uint8Array(32).fill(6); + const frame = await encryptedFrames({ open: true }, material); + const world = createMockWorld(); + world.getEncryptionKeyForRun = vi.fn().mockResolvedValue(material); + world.streams = { + get: vi.fn().mockResolvedValue( + new ReadableStream({ + start(controller) { + controller.enqueue(frame); + // Deliberately remain open: releaseLock(), not EOF, is the signal. + }, + }) + ), + } as unknown as World['streams']; + setWorld(world); + + const reader = new Run('wrun_123').getReadable({ ops }).getReader(); + await reader.read(); + reader.releaseLock(); + await expect(Promise.all(ops)).resolves.toEqual([undefined]); + }); + + it('starts stream GET and the cached run-key lookup on first read', async () => { + const material = new Uint8Array(32).fill(7); + const frame = await encryptedFrames({ first: true }, material); + let resolveRun: (run: any) => void; + const runPromise = new Promise((resolve) => { + resolveRun = resolve; + }); + const world = createMockWorld(); + world.runs.get = vi.fn().mockReturnValue(runPromise); + world.getEncryptionKeyForRun = vi.fn().mockResolvedValue(material); + world.streams = { + get: vi.fn().mockResolvedValue( + new ReadableStream({ + start(controller) { + controller.enqueue(frame); + controller.close(); + }, + }) + ), + } as unknown as World['streams']; + setWorld(world); + + const read = new Run('wrun_123').getReadable().getReader().read(); + await vi.waitFor(() => { + expect(world.streams.get).toHaveBeenCalledOnce(); + expect(world.runs.get).toHaveBeenCalledOnce(); + }); + resolveRun?.({ + runId: 'wrun_123', + deploymentId: 'test-deployment', + }); + + await expect(read).resolves.toMatchObject({ value: { first: true } }); + expect(world.getEncryptionKeyForRun).toHaveBeenCalledOnce(); + }); + + it('reuses one run-key promise across readable sessions', async () => { + const material = new Uint8Array(32).fill(8); + const frame = await encryptedFrames({ reusable: true }, material); + const world = createMockWorld(); + world.getEncryptionKeyForRun = vi.fn().mockResolvedValue(material); + world.streams = { + get: vi.fn().mockImplementation( + async () => + new ReadableStream({ + start(controller) { + controller.enqueue(frame); + controller.close(); + }, + }) + ), + } as unknown as World['streams']; + setWorld(world); + + const run = new Run('wrun_123'); + await expect(run.getReadable().getReader().read()).resolves.toMatchObject({ + value: { reusable: true }, + }); + await expect(run.getReadable().getReader().read()).resolves.toMatchObject({ + value: { reusable: true }, + }); + + expect(world.runs.get).toHaveBeenCalledOnce(); + expect(world.getEncryptionKeyForRun).toHaveBeenCalledOnce(); + expect(world.streams.get).toHaveBeenCalledTimes(2); + }); + + it('surfaces a prefetched key failure when an encrypted frame is consumed', async () => { + const material = new Uint8Array(32).fill(9); + const frame = await encryptedFrames({ secret: true }, material); + const keyError = new Error('key lookup failed'); + const world = createMockWorld(); + world.getEncryptionKeyForRun = vi.fn().mockRejectedValue(keyError); + world.streams = { + get: vi.fn().mockResolvedValue( + new ReadableStream({ + start(controller) { + controller.enqueue(frame); + controller.close(); + }, + }) + ), + } as unknown as World['streams']; + setWorld(world); + + await expect( + new Run('wrun_123').getReadable().getReader().read() + ).rejects.toThrow('key lookup failed'); + expect(world.runs.get).toHaveBeenCalledOnce(); + expect(world.getEncryptionKeyForRun).toHaveBeenCalledOnce(); + }); }); describe('Run.wakeUp', () => { diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index 51a8e667c6..be043f918c 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -698,6 +698,29 @@ function recordReadTimeToFirstChunk( })(); } +/** + * Record speculative run-key resolution independently from raw stream TTFC. + * This phase never carries key material and is intentionally separate from + * `workflow.stream.read`, whose endpoint remains the first raw frame. + */ +function recordStreamReadKeyResolution( + startEpochMs: number, + runId: string, + name: string, + succeeded: boolean +): void { + void (async () => { + await recordElapsedSpan('workflow.stream.read.resolve_key', startEpochMs, { + kind: await getSpanKind('CLIENT'), + attributes: { + 'workflow.run.id': runId, + 'workflow.stream.name': name, + 'workflow.stream.read.key_succeeded': succeeded, + }, + }); + })(); +} + /** * Emit the client-observed read-completion span when a stream read drains: * back-dated to the read dispatch, so its duration is the total read, with @@ -891,7 +914,9 @@ const getFramedStreamMaxTotalReconnects = (): number => export function createReconnectingFramedStream( runId: string, name: string, - startIndex?: number + startIndex?: number, + prefetchEncryptionKey: () => Promise = async () => + undefined ): ReadableStream { const reconnectSupported = startIndex === undefined || startIndex >= 0; let currentStartIndex = startIndex ?? 0; @@ -909,6 +934,26 @@ export function createReconnectingFramedStream( let firstChunkReported = false; let chunksDelivered = 0; let bytesDelivered = 0; + let keyPrefetched = false; + + function prefetchKey(): void { + if (keyPrefetched) return; + keyPrefetched = true; + const keyStart = Date.now(); + // The raw stream and key lookup deliberately race. Observe a speculative + // failure here: if the stream is cancelled or fails before an encrypted + // frame reaches the deserialize transform, this promise otherwise has no + // consumer and would become an unhandled rejection. The transform still + // awaits the same promise and surfaces the original error on consumption. + void prefetchEncryptionKey().then( + (key) => { + // Unencrypted runs do not need a key-resolution span. Their resolver + // still runs speculatively so an encrypted frame can join it. + if (key) recordStreamReadKeyResolution(keyStart, runId, name, true); + }, + () => recordStreamReadKeyResolution(keyStart, runId, name, false) + ); + } async function connect(): Promise { if (canceled) return false; @@ -1000,6 +1045,11 @@ export function createReconnectingFramedStream( pull: async (controller) => { if (canceled) return; if (readStart === undefined) readStart = Date.now(); + // Begin resolving the key before awaiting streams.get()/the first raw + // frame. This is intentionally not awaited: buffered raw frames remain + // bounded by the Web Streams backpressure chain while the deserialize + // transform joins this promise only if an encrypted frame needs it. + prefetchKey(); // Loop until we emit something, hit EOF, or fatally error. Reads that // only extend the in-flight-frame buffer don't enqueue anything; we // keep reading rather than returning empty-handed. @@ -2736,6 +2786,69 @@ async function getForwardedWritableEncryptionKey( return rawKey ? await importKey(rawKey, ['encrypt']) : undefined; } +/** + * Create a run's object readable without dispatching its stream GET or + * encryption-key lookup until the caller reads it. The external reviver starts + * its background pipe immediately, so this boundary belongs here—next to the + * source/transform assembly—rather than in `Run`. + * + * @internal + */ +export function getRunReadableStream( + global: Record, + ops: Promise[], + runId: string, + name: string, + startIndex: number | undefined, + cryptoKey: EncryptionKeyParam +): ReadableStream { + let reader: ReadableStreamDefaultReader | undefined; + let lockState: ReturnType | undefined; + let lockPollingStarted = false; + let userReadable: ReadableStream; + + userReadable = new ReadableStream( + { + async pull(controller) { + try { + if (!reader) { + const stream = getExternalRevivers(global, ops, runId, cryptoKey, { + onReadableState: (state) => { + lockState = state; + }, + }).ReadableStream!({ name, startIndex }) as ReadableStream; + reader = stream.getReader(); + if (lockState && !lockPollingStarted) { + lockPollingStarted = true; + // The caller owns this wrapper's reader, so polling it preserves + // the documented releaseLock() completion signal. + pollReadableLock(userReadable, lockState); + } + } + const result = await reader.read(); + if (result.done) controller.close(); + else controller.enqueue(result.value); + } catch (error) { + controller.error(error); + } + }, + async cancel(reason) { + await reader?.cancel(reason).catch(() => {}); + }, + }, + // A positive default high-water mark would run pull at construction to + // fill the queue, turning an unread Run.getReadable() into I/O. + { highWaterMark: 0 } + ); + return userReadable; +} + +/** Options for externally revived object streams. @internal */ +type ExternalReviverOptions = { + /** Receives completion state when a wrapper owns the public readable. */ + onReadableState?: (state: ReturnType) => void; +}; + /** * Revivers for deserialization boundary from the client side, * receiving the return value from the workflow handler. @@ -2748,7 +2861,8 @@ export function getExternalRevivers( global: Record = globalThis, ops: Promise[], runId: string, - cryptoKey: EncryptionKeyParam + cryptoKey: EncryptionKeyParam, + options?: ExternalReviverOptions ): Partial { return { ...getCommonRevivers(global), @@ -2840,33 +2954,47 @@ export function getExternalRevivers( // Errors are handled via state.reject }); - // Start polling to detect when user releases lock - pollReadableLock(userReadable, state); + // Direct reviver callers hold this readable. A future public wrapper + // can provide the state and poll the readable it hands to the caller. + if (options?.onReadableState) options.onReadableState(state); + else pollReadableLock(userReadable, state); return userReadable; } else { // Non-byte streams carry length-prefixed frames, so we can count // completed frames and transparently reconnect when the server // stream connection times out mid-run. + // Memoize this resolver per readable session. The first raw pull + // starts it concurrently with the stream GET; getDeserializeStream + // joins the same promise when an encrypted frame arrives. This also + // avoids invoking arbitrary EncryptionKeyParam callbacks twice. + let keyPromise: Promise | undefined; + const resolveKey = (): Promise => { + keyPromise ??= resolveEncryptionKey(cryptoKey); + return keyPromise; + }; const readable = createReconnectingFramedStream( runId, value.name, - value.startIndex + value.startIndex, + resolveKey ); const transform = getDeserializeStream( - getExternalRevivers(global, ops, runId, cryptoKey), - cryptoKey + getExternalRevivers(global, ops, runId, resolveKey), + resolveKey ); const state = createFlushableState(); ops.push(state.promise); - // Start the flushable pipe in the background + // Start the flushable pipe in the background. flushablePipe(readable, transform.writable, state).catch(() => { // Errors are handled via state.reject }); - // Start polling to detect when user releases lock - pollReadableLock(transform.readable, state); + // Direct reviver callers hold this readable. The public Run factory + // wraps it for first-pull laziness and polls that wrapper instead. + if (options?.onReadableState) options.onReadableState(state); + else pollReadableLock(transform.readable, state); return transform.readable; }