Wire AES-GCM encryption into serialization layer

.changeset/e2e-encryption.md

@@ -0,0 +1,5 @@
+---
+"@workflow/core": patch
+---
+
+Wire AES-GCM encryption into serialization layer with stream support

packages/cli/src/lib/inspect/output.ts

@@ -1,4 +1,6 @@
+import { importKey } from '@workflow/core/encryption';
 import {
+  type EncryptionKeyParam,
   getDeserializeStream,
   getExternalRevivers,
 } from '@workflow/core/serialization';
@@ -780,16 +782,32 @@ export const showStream = async (
   streamId: string,
   opts: InspectCLIOptions = {}
 ) => {
-  if (opts.runId || opts.stepId) {
+  if (opts.stepId) {
     logger.warn(
-      'Filtering by run-id or step-id is not supported when showing a stream, ignoring filter.'
+      'Filtering by step-id is not supported when showing a stream, ignoring filter.'
     );
   }
   const rawStream = await world.readFromStream(streamId);
 
+  // Resolve the encryption key if a runId is provided (needed for encrypted streams).
+  // The key is passed as a promise so stream construction is synchronous —
+  // it will be resolved lazily on the first encrypted frame.
+  let encryptionKey: EncryptionKeyParam;
+  if (opts.runId) {
+    encryptionKey = (async () => {
+      const rawKey = await world.getEncryptionKeyForRun?.(opts.runId!);
+      return rawKey ? await importKey(rawKey) : undefined;
+    })();
+  }
+
   // Deserialize the stream to get JavaScript objects
-  const revivers = getExternalRevivers(globalThis, [], '');
-  const transform = getDeserializeStream(revivers);
+  const revivers = getExternalRevivers(
+    globalThis,
+    [],
+    opts.runId ?? '',
+    encryptionKey
+  );
+  const transform = getDeserializeStream(revivers, encryptionKey);
   const stream = rawStream.pipeThrough(transform);
 
   logger.info('Streaming to stdout, press CTRL+C to abort.');

packages/core/src/async-deserialization-ordering.test.ts

@@ -34,6 +34,7 @@ function setupWorkflowContext(events: Event[]): WorkflowOrchestratorContext {
     globalThis: context.globalThis,
     eventsConsumer: new EventsConsumer(events, {
       onUnconsumedEvent: () => {},
+      getPromiseQueue: () => Promise.resolve(),
     }),
     invocationsQueue: new Map(),
     generateUlid: () => ulid(workflowStartedAt),

packages/core/src/events-consumer.test.ts

@@ -16,16 +16,19 @@ function createMockEvent(overrides: Partial<Event> = {}): Event {
 }
 
 // Default options for tests that don't care about onUnconsumedEvent
-const defaultOptions = { onUnconsumedEvent: vi.fn() };
+const defaultOptions = {
+  onUnconsumedEvent: vi.fn(),
+  getPromiseQueue: () => Promise.resolve(),
+};
 
 // Helper function to wait for next tick
 function waitForNextTick(): Promise<void> {
   return new Promise((resolve) => process.nextTick(resolve));
 }
 
-// Helper to wait for setTimeout(0) macrotask (used by deferred unconsumed event check)
-function waitForMacrotask(): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, 0));
+// Helper to wait for the unconsumed event check (promiseQueue .then() + setTimeout(100ms))
+function waitForUnconsumedCheck(): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, 150));
 }
 
 describe('EventsConsumer', () => {
@@ -159,7 +162,10 @@ describe('EventsConsumer', () => {
     it('should process all callbacks when none return true and call onUnconsumedEvent', async () => {
       const event = createMockEvent();
       const onUnconsumedEvent = vi.fn();
-      const consumer = new EventsConsumer([event], { onUnconsumedEvent });
+      const consumer = new EventsConsumer([event], {
+        onUnconsumedEvent,
+        getPromiseQueue: () => Promise.resolve(),
+      });
       const callback1 = vi
         .fn()
         .mockReturnValue(EventConsumerResult.NotConsumed);
@@ -181,8 +187,9 @@ describe('EventsConsumer', () => {
       expect(consumer.eventIndex).toBe(0);
       expect(consumer.callbacks).toEqual([callback1, callback2, callback3]);
 
-      // onUnconsumedEvent is deferred via setTimeout to allow new subscribes
-      await waitForMacrotask();
+      // onUnconsumedEvent is deferred via promise queue .then() + setTimeout(0)
+      await waitForNextTick();
+      await waitForUnconsumedCheck();
       expect(onUnconsumedEvent).toHaveBeenCalledWith(event);
     });
 
@@ -352,24 +359,33 @@ describe('EventsConsumer', () => {
     it('should call onUnconsumedEvent when a non-null event is not consumed by any callback', async () => {
       const event = createMockEvent();
       const onUnconsumedEvent = vi.fn();
-      const consumer = new EventsConsumer([event], { onUnconsumedEvent });
+      const consumer = new EventsConsumer([event], {
+        onUnconsumedEvent,
+        getPromiseQueue: () => Promise.resolve(),
+      });
       const callback = vi.fn().mockReturnValue(EventConsumerResult.NotConsumed);
 
       consumer.subscribe(callback);
+      // Wait for: nextTick(consume) → .then() microtask → setTimeout(0) macrotask
+      await waitForNextTick();
       await waitForNextTick();
-      await waitForMacrotask();
+      await waitForUnconsumedCheck();
 
       expect(onUnconsumedEvent).toHaveBeenCalledWith(event);
     });
 
     it('should NOT call onUnconsumedEvent for null event (end-of-events)', async () => {
       const onUnconsumedEvent = vi.fn();
-      const consumer = new EventsConsumer([], { onUnconsumedEvent });
+      const consumer = new EventsConsumer([], {
+        onUnconsumedEvent,
+        getPromiseQueue: () => Promise.resolve(),
+      });
       const callback = vi.fn().mockReturnValue(EventConsumerResult.NotConsumed);
 
       consumer.subscribe(callback);
       await waitForNextTick();
-      await waitForMacrotask();
+      await waitForNextTick();
+      await waitForUnconsumedCheck();
 
       expect(callback).toHaveBeenCalledWith(null);
       expect(onUnconsumedEvent).not.toHaveBeenCalled();
@@ -378,7 +394,10 @@ describe('EventsConsumer', () => {
     it('should cancel pending unconsumed check when a new callback subscribes', async () => {
       const event = createMockEvent();
       const onUnconsumedEvent = vi.fn();
-      const consumer = new EventsConsumer([event], { onUnconsumedEvent });
+      const consumer = new EventsConsumer([event], {
+        onUnconsumedEvent,
+        getPromiseQueue: () => Promise.resolve(),
+      });
       const callback1 = vi
         .fn()
         .mockReturnValue(EventConsumerResult.NotConsumed);
@@ -390,7 +409,8 @@ describe('EventsConsumer', () => {
       const callback2 = vi.fn().mockReturnValue(EventConsumerResult.Finished);
       consumer.subscribe(callback2);
       await waitForNextTick();
-      await waitForMacrotask();
+      await waitForNextTick();
+      await waitForUnconsumedCheck();
 
       // The new callback consumed the event, so onUnconsumedEvent should NOT be called
       expect(onUnconsumedEvent).not.toHaveBeenCalled();

packages/core/src/events-consumer.ts

@@ -21,24 +21,36 @@ type EventConsumerCallback = (event: Event | null) => EventConsumerResult;
 export interface EventsConsumerOptions {
   /**
    * Callback invoked when a non-null event cannot be consumed by any registered
-   * callback, indicating an orphaned or invalid event in the event log. Called
-   * asynchronously after a macrotask delay to allow pending callback
-   * subscriptions to settle first.
+   * callback, indicating an orphaned or invalid event in the event log. The
+   * check is deferred until after the promise queue has drained, ensuring that
+   * any pending async work (e.g., deserialization/decryption) completes and
+   * downstream subscribe() calls have a chance to cancel the check first.
    */
   onUnconsumedEvent: (event: Event) => void;
+  /**
+   * Returns the current promise queue. The unconsumed event check is chained
+   * onto this queue so it only fires after all pending async work (e.g.,
+   * deserialization) has completed. This prevents false positives when async
+   * deserialization delays the resolve() that triggers the next subscribe().
+   */
+  getPromiseQueue: () => Promise<void>;
 }
 
 export class EventsConsumer {
   eventIndex: number;
   readonly events: Event[] = [];
   readonly callbacks: EventConsumerCallback[] = [];
   private onUnconsumedEvent: (event: Event) => void;
-  private pendingUnconsumedCheck: ReturnType<typeof setTimeout> | null = null;
+  private getPromiseQueue: () => Promise<void>;
+  private pendingUnconsumedCheck: Promise<void> | null = null;
+  private pendingUnconsumedTimeout: ReturnType<typeof setTimeout> | null = null;
+  private unconsumedCheckVersion = 0;
 
   constructor(events: Event[], options: EventsConsumerOptions) {
     this.events = events;
     this.eventIndex = 0;
     this.onUnconsumedEvent = options.onUnconsumedEvent;
+    this.getPromiseQueue = options.getPromiseQueue;
   }
 
   /**
@@ -52,10 +64,16 @@ export class EventsConsumer {
    */
   subscribe(fn: EventConsumerCallback) {
     this.callbacks.push(fn);
-    // Cancel any pending unconsumed check since a new callback may consume the event
+    // Cancel any pending unconsumed check since a new callback may consume the event.
+    // Incrementing the version causes any in-flight promise chain check to no-op.
+    // Also clear the pending setTimeout if it hasn't fired yet.
     if (this.pendingUnconsumedCheck !== null) {
-      clearTimeout(this.pendingUnconsumedCheck);
+      this.unconsumedCheckVersion++;
       this.pendingUnconsumedCheck = null;
+      if (this.pendingUnconsumedTimeout !== null) {
+        clearTimeout(this.pendingUnconsumedTimeout);
+        this.pendingUnconsumedTimeout = null;
+      }
     }
     process.nextTick(this.consume);
   }
@@ -90,18 +108,29 @@ export class EventsConsumer {
 
     // If we reach here, all callbacks returned NotConsumed.
     // If the current event is non-null (a real event, not end-of-events),
-    // schedule a deferred check. We use setTimeout (macrotask) so that any
-    // pending process.nextTick microtasks (e.g., new subscribes from the
-    // workflow code) can complete first. If the event is still unconsumed
-    // when the timeout fires, it's truly orphaned.
+    // schedule a deferred check. We chain onto the promiseQueue so that any
+    // pending async work (e.g., deserialization/decryption that triggers
+    // resolve() → user code → subscribe()) completes first. If the event
+    // is still unconsumed after the queue drains, it's truly orphaned.
     if (currentEvent !== null) {
-      const unconsumedIndex = this.eventIndex;
-      this.pendingUnconsumedCheck = setTimeout(() => {
-        this.pendingUnconsumedCheck = null;
-        if (this.eventIndex === unconsumedIndex) {
-          this.onUnconsumedEvent(currentEvent);
-        }
-      }, 0);
+      const checkVersion = ++this.unconsumedCheckVersion;
+      this.pendingUnconsumedCheck = this.getPromiseQueue().then(() => {
+        // Use a delayed setTimeout after the queue drains. The delay must be
+        // long enough for promise chains to propagate across the VM boundary
+        // (from resolve() in the host context through to the workflow code
+        // calling subscribe() in the VM context). Node.js does not guarantee
+        // that setTimeout(0) fires after all cross-context microtasks settle,
+        // so we use a small but non-zero delay. Any subscribe() call that
+        // arrives during this window will cancel the check via version
+        // invalidation + clearTimeout.
+        this.pendingUnconsumedTimeout = setTimeout(() => {
+          this.pendingUnconsumedTimeout = null;
+          if (this.unconsumedCheckVersion === checkVersion) {
+            this.pendingUnconsumedCheck = null;
+            this.onUnconsumedEvent(currentEvent);
+          }
+        }, 100);
+      });
     }
   };
 }

packages/core/src/runtime/run.ts

@@ -8,7 +8,7 @@ import {
   type WorkflowRunStatus,
   type World,
 } from '@workflow/world';
-import { importKey } from '../encryption.js';
+import { type CryptoKey, importKey } from '../encryption.js';
 import {
   getExternalRevivers,
   hydrateWorkflowReturnValue,
@@ -63,11 +63,35 @@ export class Run<TResult> {
    */
   private world: World;
 
+  /**
+   * Cached encryption key resolution. Resolved once on first use and
+   * reused for returnValue, getReadable(), etc.
+   * @internal
+   */
+  private encryptionKeyPromise: Promise<CryptoKey | undefined> | null = null;
+
   constructor(runId: string) {
     this.runId = runId;
     this.world = getWorld();
   }
 
+  /**
+   * Resolves and caches the encryption key for this run.
+   * The key is the same for the lifetime of a run, so it only needs
+   * to be resolved once.
+   * @internal
+   */
+  private getEncryptionKey(): Promise<CryptoKey | undefined> {
+    if (!this.encryptionKeyPromise) {
+      this.encryptionKeyPromise = (async () => {
+        const run = await this.world.runs.get(this.runId);
+        const rawKey = await this.world.getEncryptionKeyForRun?.(run);
+        return rawKey ? await importKey(rawKey) : undefined;
+      })();
+    }
+    return this.encryptionKeyPromise;
+  }
+
   /**
    * Interrupts pending `sleep()` calls, resuming the workflow early.
    *
@@ -153,7 +177,15 @@ export class Run<TResult> {
   ): ReadableStream<R> {
     const { ops = [], global = globalThis, startIndex, namespace } = options;
     const name = getWorkflowRunStreamId(this.runId, namespace);
-    return getExternalRevivers(global, ops, this.runId).ReadableStream({
+    // Pass the key as a promise — it will be resolved lazily inside
+    // the first async transform() call of the deserialize stream.
+    const encryptionKey = this.getEncryptionKey();
+    return getExternalRevivers(
+      global,
+      ops,
+      this.runId,
+      encryptionKey
+    ).ReadableStream({
       name,
       startIndex,
     }) as ReadableStream<R>;
@@ -170,8 +202,7 @@ export class Run<TResult> {
         const run = await this.world.runs.get(this.runId);
 
         if (run.status === 'completed') {
-          const rawKey = await this.world.getEncryptionKeyForRun?.(run);
-          const encryptionKey = rawKey ? await importKey(rawKey) : undefined;
+          const encryptionKey = await this.getEncryptionKey();
           return await hydrateWorkflowReturnValue(
             run.output,
             this.runId,

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

@@ -356,6 +356,7 @@ const stepHandler = getWorldHandlers().createQueueHandler(
                   },
                   ops,
                   closureVars: hydratedInput.closureVars,
+                  encryptionKey,
                 },
                 () => stepFn.apply(thisVal, args)
               );