diff --git a/.changeset/forwarded-writable-key-lazy.md b/.changeset/forwarded-writable-key-lazy.md new file mode 100644 index 0000000000..13011dc9a7 --- /dev/null +++ b/.changeset/forwarded-writable-key-lazy.md @@ -0,0 +1,5 @@ +--- +'@workflow/core': patch +--- + +Fix an unhandled rejection that could exit the process when looking up the encryption key for a forwarded writable stream failed (for example a run metadata request that timed out) before anything was written to the stream. The lookup now starts on the first write and its failure rejects that stream instead. diff --git a/packages/core/src/serialization.test.ts b/packages/core/src/serialization.test.ts index 1dc9198a1e..52c9f37e9b 100644 --- a/packages/core/src/serialization.test.ts +++ b/packages/core/src/serialization.test.ts @@ -1,5 +1,8 @@ import { runInContext } from 'node:vm'; -import type { WorkflowRuntimeError } from '@workflow/errors'; +import { + WorkflowWorldError, + type WorkflowRuntimeError, +} from '@workflow/errors'; import { FatalError, HookConflictError, @@ -552,6 +555,80 @@ describe('workflow arguments', () => { expect(text).toContain('dpl_parent'); }); + it('does not reject unhandled when the forwarded key lookup fails before the first write', async () => { + const { getWorldLazy } = await import('./runtime/get-world-lazy.js'); + const timeout = new WorkflowWorldError( + 'GET /v2/runs/wrun_parent timed out', + { + code: 'TIMEOUT', + } + ); + const runsGet = vi.fn().mockRejectedValue(timeout); + const getEncryptionKeyForRun = vi.fn(); + const world = { + ...makeMockWorld(), + runs: { get: runsGet }, + getEncryptionKeyForRun, + } as any; + vi.mocked(getWorldLazy).mockReturnValue(world); + + const unhandledRejections: unknown[] = []; + const onUnhandledRejection = (reason: unknown) => { + unhandledRejections.push(reason); + }; + process.on('unhandledRejection', onUnhandledRejection); + + try { + // No deployment id and no public key on the forwarded writable, so the + // key lookup has to read the parent run's metadata. + const parentWritable = new WritableStream(); + Object.defineProperty(parentWritable, STREAM_NAME_SYMBOL, { + value: 'strm_parentstreamname', + writable: false, + }); + Object.defineProperty(parentWritable, STREAM_SERVER_RUN_ID_SYMBOL, { + value: 'wrun_parent', + writable: false, + }); + + const serialized = await dehydrateStepArguments( + parentWritable, + 'wrun_child', + noEncryptionKey + ); + const ops: Promise[] = []; + const hydrated = (await hydrateStepArguments( + serialized, + 'wrun_child', + noEncryptionKey, + ops, + globalThis, + {} + )) as WritableStream; + + // Nothing is written yet, so nothing should have been requested and + // nothing can reject. + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(runsGet).not.toHaveBeenCalled(); + expect(unhandledRejections).toEqual([]); + + // The failure belongs to the stream that needed the key: the serialize + // transform errors the stream, so it shows up on the writer. + const writer = hydrated.getWriter(); + await writer.write('hello').catch(() => {}); + await expect(writer.closed).rejects.toMatchObject({ + cause: { code: 'TIMEOUT' }, + }); + expect(runsGet).toHaveBeenCalledTimes(1); + + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(unhandledRejections).toEqual([]); + } finally { + process.off('unhandledRejection', onUnhandledRejection); + vi.mocked(getWorldLazy).mockImplementation(() => makeMockWorld() as any); + } + }); + it('uses the forwarded stream deployment to resolve its encryption key', async () => { const { getWorldLazy } = await import('./runtime/get-world-lazy.js'); const getEncryptionKeyForRun = vi diff --git a/packages/core/src/serialization.ts b/packages/core/src/serialization.ts index be043f918c..227f57b0fb 100644 --- a/packages/core/src/serialization.ts +++ b/packages/core/src/serialization.ts @@ -2786,6 +2786,30 @@ async function getForwardedWritableEncryptionKey( return rawKey ? await importKey(rawKey, ['encrypt']) : undefined; } +/** + * Defer the forwarded-writable key lookup until something is written. + * + * The thunk runs the lookup once and memoizes the result, so a failure (for + * example a run metadata request that timed out) surfaces on the write that + * needed the key instead of as an unhandled rejection on a writable nobody + * has written to. + */ +function lazyForwardedWritableEncryptionKey( + runId: string, + deploymentId: string | undefined, + encryptionPublicKey: string | undefined +): () => Promise { + let keyPromise: Promise | undefined; + return () => { + keyPromise ??= getForwardedWritableEncryptionKey( + runId, + deploymentId, + encryptionPublicKey + ); + return keyPromise; + }; +} + /** * Create a run's object readable without dispatching its stream GET or * encryption-key lookup until the caller reads it. The external reviver starts @@ -3007,7 +3031,7 @@ export function getExternalRevivers( const targetKey: EncryptionKeyParam = targetRunId === runId ? cryptoKey - : getForwardedWritableEncryptionKey( + : lazyForwardedWritableEncryptionKey( targetRunId, value.deploymentId, value.encryptionPublicKey @@ -3420,10 +3444,11 @@ function getStepRevivers( // Cross-run case (parent → child via `start()`): the descriptor // carries the original `runId` and `name`. Open a server writable // against the original `(runId, name)` and resolve THAT run's key - // for encryption. The resolution is async but doesn't need to - // block reviver return: `getSerializeStream` accepts the - // `Promise` directly and awaits it lazily - // on the first chunk written. The key is imported encrypt-only + // for encryption. The lookup does not start until the first chunk + // is written: `getSerializeStream` calls the key thunk lazily, so a + // failed lookup errors the stream instead of becoming an unhandled + // rejection on a writable nobody wrote to. The key is imported + // encrypt-only // so the receiving run can never decrypt anything else on the // owning run's stream; it can only contribute new writes. const targetRunId = typeof value.runId === 'string' ? value.runId : runId; @@ -3436,7 +3461,7 @@ function getStepRevivers( const targetKey: EncryptionKeyParam = targetRunId === runId ? cryptoKey - : getForwardedWritableEncryptionKey( + : lazyForwardedWritableEncryptionKey( targetRunId, targetDeploymentId, value.encryptionPublicKey