[core] Prefetch stream read encryption keys

.changeset/stream-read-key-prefetch.md

@@ -0,0 +1,5 @@
+---
+'@workflow/core': patch
+---
+
+Prefetch run encryption keys when reading workflow streams.

packages/core/src/readable-stream-telemetry.test.ts

@@ -153,7 +153,12 @@ describe('createReconnectingFramedStream read telemetry', () => {
   });
 
   it('emits read and read.complete spans with connect, ttfc, totals, and reconnects', async () => {
-    const stream = createReconnectingFramedStream('run-123', 'test-stream');
+    const stream = createReconnectingFramedStream(
+      'run-123',
+      'test-stream',
+      undefined,
+      async () => undefined
+    );
     const reader = stream.getReader();
     for (let i = 0; i < 10; i++) {
       const { done } = await reader.read();
@@ -169,6 +174,15 @@ describe('createReconnectingFramedStream read telemetry', () => {
       'number'
     );
 
+    const [keySpan] = await waitForSpans('workflow.stream.read.resolve_key', 1);
+    expect(keySpan.attributes['workflow.stream.read.key_prefetched']).toBe(
+      true
+    );
+    expect(keySpan.attributes['workflow.stream.read.key_succeeded']).toBe(true);
+    expect(
+      typeof keySpan.attributes['workflow.stream.read.key_resolve_ms']
+    ).toBe('number');
+
     const [doneSpan] = await waitForSpans('workflow.stream.read.complete', 1);
     expect(doneSpan).toBeDefined();
     expect(doneSpan.attributes['workflow.stream.read.chunks']).toBe(2);

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<Uint8Array>) => void;
+    const streamPromise = new Promise<ReadableStream<Uint8Array>>((resolve) => {
+      resolveStream = resolve;
+    });
+    let resolveKey: () => void;
+    const keyPromise = new Promise<void>((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<void>((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<void>((resolve) => {
+        resolveKey = resolve;
+      })
+    );
+    const source = new ReadableStream<Uint8Array>({
+      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) =>

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';
@@ -374,11 +374,14 @@ export class Run<TResult> {
     // chunk, so creating or probing an empty stream cannot reject in the
     // background.
     const encryptionKey = this.#getEncryptionKeyLazily();
-    const stream = getExternalRevivers(global, ops, this.runId, encryptionKey)
-      .ReadableStream!({
+    const stream = getRunReadableStream<R>(
+      global,
+      ops,
+      this.runId,
       name,
       startIndex,
-    }) as ReadableStream<R>;
+      encryptionKey
+    );
 
     const worldPromise = this.#lazyWorldPromise;
     const runId = this.runId;

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,102 @@ 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('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<any>((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', () => {