Reuse resolved encryption keys and simplify payload key handling

.changeset/reuse-suspension-encryption-key.md

@@ -0,0 +1,5 @@
+---
+'@workflow/core': patch
+---
+
+Reuse resolved workflow encryption keys during suspension handling and avoid duplicate run lookups while hydrating return values.

packages/core/src/runtime.ts

@@ -1640,6 +1640,7 @@ export function workflowEntrypoint(
                             suspension: err,
                             world,
                             run: workflowRun,
+                            encryptionKey,
                             span,
                             requestId,
                             eventLog: suspensionLog,

packages/core/src/runtime/run.ts

@@ -7,6 +7,7 @@ import {
 import { WORKFLOW_DESERIALIZE, WORKFLOW_SERIALIZE } from '@workflow/serde';
 import {
   SPEC_VERSION_CURRENT,
+  type WorkflowRun,
   type WorkflowRunStatus,
   type World,
 } from '@workflow/world';
@@ -128,12 +129,12 @@ export class Run<TResult> {
    * to be resolved once.
    * @internal
    */
-  #getEncryptionKey(): Promise<PayloadKey | undefined> {
+  #getEncryptionKey(run?: WorkflowRun): Promise<PayloadKey | undefined> {
     if (!this.#encryptionKeyPromise) {
       this.#encryptionKeyPromise = (async () => {
         const world = await this.#lazyWorldPromise;
-        const run = await world.runs.get(this.runId);
-        const rawKey = await world.getEncryptionKeyForRun?.(run);
+        const resolvedRun = run ?? (await world.runs.get(this.runId));
+        const rawKey = await world.getEncryptionKeyForRun?.(resolvedRun);
         return rawKey ? await deriveRunPayloadKeys(rawKey) : undefined;
       })();
     }
@@ -335,7 +336,7 @@ export class Run<TResult> {
         const run = await world.runs.get(this.runId);
 
         if (run.status === 'completed') {
-          const encryptionKey = await this.#getEncryptionKey();
+          const encryptionKey = await this.#getEncryptionKey(run);
           return await hydrateWorkflowReturnValue(
             run.output,
             this.runId,
@@ -351,7 +352,7 @@ export class Run<TResult> {
           // Hydrate the serialized run error so the original thrown value
           // (with its type identity, cause chain, etc.) is set as the
           // `cause` on WorkflowRunFailedError.
-          const encryptionKey = await this.#getEncryptionKey();
+          const encryptionKey = await this.#getEncryptionKey(run);
           let hydratedError: unknown;
           try {
             hydratedError = await hydrateRunError(

packages/core/src/runtime/runs.test.ts

@@ -25,9 +25,11 @@ vi.mock('../serialization.js', async (importActual) => {
 });
 
 import { registerSerializationClass } from '../class-serialization.js';
+import { importKey } from '../encryption.js';
 import {
   dehydrateRunError,
   dehydrateStepReturnValue,
+  dehydrateWorkflowReturnValue,
   hydrateStepReturnValue,
 } from '../serialization.js';
 import { Run } from './run.js';
@@ -396,6 +398,41 @@ describe('Run custom serialization', () => {
   });
 });
 
+describe('Run.returnValue encryption key resolution', () => {
+  afterEach(() => {
+    setWorld(undefined as unknown as World);
+  });
+
+  it('reuses the polled run when resolving its encryption key', async () => {
+    const rawKey = new Uint8Array(32).fill(7);
+    const encryptionKey = await importKey(rawKey);
+    const world = createMockWorld({
+      run: {
+        status: 'completed',
+        output: await dehydrateWorkflowReturnValue(
+          'result',
+          'wrun_123',
+          encryptionKey
+        ),
+        completedAt: new Date(),
+      },
+    });
+    const getEncryptionKeyForRun = vi.fn().mockResolvedValue(rawKey);
+    world.getEncryptionKeyForRun = getEncryptionKeyForRun;
+    setWorld(world);
+
+    await expect(new Run('wrun_123').returnValue).resolves.toBe('result');
+
+    expect(world.runs.get).toHaveBeenCalledTimes(1);
+    expect(getEncryptionKeyForRun).toHaveBeenCalledWith(
+      expect.objectContaining({
+        runId: 'wrun_123',
+        status: 'completed',
+      })
+    );
+  });
+});
+
 describe('Run.returnValue when run.status === "failed"', () => {
   // Register the FatalError class so the run-error serialization pipeline
   // can find it during hydration (the SWC plugin does this in production).

packages/core/src/runtime/suspension-handler.test.ts

@@ -1,6 +1,12 @@
-import type { WorkflowRun, World } from '@workflow/world';
+import type { CreateEventRequest, WorkflowRun, World } from '@workflow/world';
 import { describe, expect, it, vi } from 'vitest';
 import { WorkflowSuspension } from '../global.js';
+import { deriveRunPayloadKeys } from '../serialization/encryption.js';
+import {
+  hydrateStepArguments,
+  peekFormatPrefix,
+  SerializationFormat,
+} from '../serialization.js';
 import { handleSuspension } from './suspension-handler.js';
 
 vi.mock('../version.js', () => ({ version: '0.0.0-test' }));
@@ -52,6 +58,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     expect(eventsCreate).toHaveBeenCalledWith(
@@ -67,6 +74,42 @@ describe('handleSuspension', () => {
     );
   });
 
+  it('uses the supplied encryption key without resolving it again', async () => {
+    const eventsCreate = vi.fn().mockImplementation(async (_runId, event) => ({
+      event,
+    }));
+    const world = createWorld(eventsCreate);
+    const encryptionKey = await deriveRunPayloadKeys(
+      new Uint8Array(32).fill(7)
+    );
+    const pending = new Map([
+      [
+        'hook_with_metadata',
+        {
+          type: 'hook' as const,
+          correlationId: 'hook_with_metadata',
+          token: 'order:encrypted',
+          metadata: { visibility: 'private' },
+        },
+      ],
+    ]);
+
+    await handleSuspension({
+      suspension: new WorkflowSuspension(pending, globalThis),
+      world,
+      run,
+      encryptionKey,
+    });
+
+    expect(world.getEncryptionKeyForRun).not.toHaveBeenCalled();
+    const createdEvent = eventsCreate.mock.calls[0]?.[1] as CreateEventRequest;
+    const metadata = createdEvent.eventData?.metadata as Uint8Array;
+    expect(peekFormatPrefix(metadata)).toBe(SerializationFormat.ENCRYPTED);
+    await expect(
+      hydrateStepArguments(metadata, run.runId, encryptionKey)
+    ).resolves.toEqual({ visibility: 'private' });
+  });
+
   it('marks hook.getConflict()-awaited creations without converting them into wait timeouts', async () => {
     const eventsCreate = vi.fn().mockResolvedValue({
       event: {
@@ -90,6 +133,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     expect(eventsCreate).toHaveBeenCalledWith(
@@ -136,6 +180,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     expect(result.hasAwaitedHookCreation).toBe(true);
@@ -163,6 +208,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     expect(result.lazyInlineSteps.map((s) => s.correlationId)).toEqual([
@@ -202,6 +248,7 @@ describe('handleSuspension', () => {
         suspension: new WorkflowSuspension(pending, globalThis),
         world,
         run,
+        encryptionKey: undefined,
       });
 
       // Cap of 1: only the first step is deferred; s2 and s3 are eager-created.
@@ -249,6 +296,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     // Nothing runs inline: the step keeps its eager step_created (owned) and is
@@ -279,6 +327,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     expect(result.hasAwaitedHookCreation).toBe(false);
@@ -318,6 +367,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     const hookCalls = eventsCreate.mock.calls.map(([, event]) => ({
@@ -351,6 +401,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     const hookCalls = eventsCreate.mock.calls.map(([, event]) => ({
@@ -387,6 +438,7 @@ describe('handleSuspension', () => {
       suspension: new WorkflowSuspension(pending, globalThis),
       world,
       run,
+      encryptionKey: undefined,
     });
 
     expect(result.hasHookConflict).toBe(true);

packages/core/src/runtime/suspension-handler.ts

@@ -17,7 +17,6 @@ import {
   type WorkflowRun,
   type World,
 } from '@workflow/world';
-import { importKey } from '../encryption.js';
 import type {
   AttributeInvocationQueueItem,
   HookInvocationQueueItem,
@@ -26,6 +25,7 @@ import type {
   WorkflowSuspension,
 } from '../global.js';
 import { runtimeLogger } from '../logger.js';
+import type { PayloadKey } from '../serialization/encryption.js';
 import { dehydrateStepArguments } from '../serialization.js';
 import * as Attribute from '../telemetry/semantic-conventions.js';
 import { getAbortStreamIdFromToken } from '../util.js';
@@ -36,6 +36,12 @@ export interface SuspensionHandlerParams {
   suspension: WorkflowSuspension;
   world: World;
   run: WorkflowRun;
+  /**
+   * The run's already-resolved encryption key. The workflow runtime shares
+   * this with replay payload hydration and suspension serialization so the
+   * World is not asked to derive or fetch the same key twice.
+   */
+  encryptionKey: PayloadKey | undefined;
   span?: Span;
   requestId?: string;
   /**
@@ -201,6 +207,7 @@ export async function handleSuspension({
   suspension,
   world,
   run,
+  encryptionKey,
   span,
   requestId,
   eventLog,
@@ -281,10 +288,6 @@ export async function handleSuspension({
     }
   }
 
-  // Resolve encryption key for this run
-  const rawKey = await world.getEncryptionKeyForRun?.(run);
-  const encryptionKey = rawKey ? await importKey(rawKey) : undefined;
-
   // Gate payload compression on the run's specVersion.
   const compression =
     (run.specVersion ?? 0) >= SPEC_VERSION_SUPPORTS_COMPRESSION;

packages/core/src/serialization.ts

@@ -6,7 +6,6 @@ import {
 import { envNumber } from '@workflow/world';
 import { parse, stringify, unflatten } from 'devalue';
 import { monotonicFactory } from 'ulid';
-import { bytesToBase64, decodeRunPublicKey } from './sealed-box.js';
 import { importKey } from './encryption.js';
 import {
   createFlushableState,
@@ -26,24 +25,28 @@ import { getStepFunction } from './private.js';
 // "Turbopack NFT Tracing Errors in V2 Combined Flow Route" section of
 // `docs/content/docs/changelog/eager-processing.mdx`.
 import { getWorldLazy } from './runtime/get-world-lazy.js';
-import { createOpenSession, createSealSession } from './sealed-box.js';
+import {
+  bytesToBase64,
+  createOpenSession,
+  createSealSession,
+  decodeRunPublicKey,
+} from './sealed-box.js';
 import * as clientModule from './serialization/client.js';
 import {
   type CompressionStats,
   compress,
   decompress,
 } from './serialization/compression.js';
 import {
-  aesKeyOf,
   decrypt,
   deriveRunPayloadKeys,
   type EncryptionKeyParam,
   encrypt,
   isRunPayloadKeys,
   isSealTarget,
   type PayloadKey,
-  resolveEncryptionKey,
   type RunPayloadKeys,
+  resolveEncryptionKey,
   runPayloadKeys,
   type SealTarget,
   sealTo,
@@ -124,7 +127,6 @@ export {
   deriveRunPayloadKeys,
   isSealTarget,
   isRunPayloadKeys,
-  aesKeyOf,
 };
 
 // Re-export the legacy SerializationFormatType for backwards compatibility.
@@ -382,7 +384,7 @@ export function getDeserializeStream(
         // wrong kind of key" have very different causes.
         const usable = sealed
           ? isRunPayloadKeys(keyState.key)
-          : aesKeyOf(keyState.key) !== undefined;
+          : keyState.key !== undefined && !isSealTarget(keyState.key);
         if (!usable) {
           controller.error(
             new RuntimeDecryptionError(