fix(web): keep a stream failure that lands while nobody is parked on next()

bindings/web/packages/core/src/Adapters/ProtoAdapterTypes.ts

@@ -549,6 +549,7 @@ export function streamCallback<T>(
       let callbackPtr = 0;
       let started = false;
       let finished = false;
+      let failure: unknown = null;
       let callActive = false;
       let emitsSinceYield = 0;
 
@@ -571,6 +572,10 @@ export function streamCallback<T>(
       const fail = (error: unknown): void => {
         if (finished) return;
         finished = true;
+        // Kept for the consumer that is not parked right now (a decode that
+        // throws while earlier events are still buffered, say); without it the
+        // next `next()` would read the stream as a clean end.
+        failure = error;
         while (waiters.length > 0) {
           waiters.shift()!.reject(error);
         }
@@ -700,6 +705,11 @@ export function streamCallback<T>(
           if (queue.length > 0) {
             return Promise.resolve({ value: queue.shift()!, done: false });
           }
+          if (failure !== null) {
+            const error = failure;
+            failure = null;
+            return Promise.reject(error);
+          }
           if (finished) {
             return Promise.resolve({ value: undefined as T, done: true });
           }
@@ -711,6 +721,11 @@ export function streamCallback<T>(
           try {
             onCancel?.();
           } finally {
+            // An explicit cancel outranks an error the consumer never asked
+            // about: `finish()` is a no-op once `fail()` has set `finished`,
+            // so drop the retained failure here rather than replaying it at
+            // the next `next()`.
+            failure = null;
             finish();
           }
           return Promise.resolve({ value: undefined as T, done: true });

bindings/web/packages/core/src/runtime/BackendWorkerHost.ts

@@ -73,6 +73,13 @@ interface StreamPending {
     reject(reason: unknown): void;
   }>;
   finished: boolean;
+  /**
+   * Set by whichever path failed the stream. Rejecting the parked waiters is
+   * not enough on its own: a worker that dies while the consumer is running its
+   * loop body has no waiter to reject, and `next()` would otherwise read
+   * `finished` as a clean end of stream.
+   */
+  failure?: unknown;
 }
 
 type PendingRequest = UnaryPending | StreamPending;
@@ -223,6 +230,7 @@ export class BackendWorkerHost {
     const fail = (error: unknown): void => {
       if (state.finished) return;
       state.finished = true;
+      state.failure = error;
       this.pending.delete(requestId);
       this.activeStreamIds.delete(requestId);
       while (state.waiters.length) state.waiters.shift()!.reject(error);
@@ -244,11 +252,19 @@ export class BackendWorkerHost {
           if (state.events.length) {
             return Promise.resolve({ value: state.events.shift(), done: false });
           }
+          if (state.failure !== undefined) {
+            const error = state.failure;
+            state.failure = undefined;
+            return Promise.reject(error);
+          }
           if (state.finished) return Promise.resolve({ value: undefined, done: true });
           return new Promise((resolve, reject) => state.waiters.push({ resolve, reject }));
         },
         return: (): Promise<IteratorResult<unknown>> => {
           if (started && !state.finished) this.cancel(requestId);
+          // An explicit cancel outranks an error the consumer never asked
+          // about, matching the other two iterators.
+          state.failure = undefined;
           finish();
           return Promise.resolve({ value: undefined, done: true });
         },
@@ -380,6 +396,7 @@ export class BackendWorkerHost {
       return;
     }
     pending.finished = true;
+    pending.failure = error;
     while (pending.waiters.length) pending.waiters.shift()!.reject(error);
   }
 

bindings/web/packages/core/src/runtime/OffscreenRuntimeBridge.ts

@@ -331,6 +331,11 @@ export class OffscreenRuntimeBridge {
         }> = [];
         let started = false;
         let finished = false;
+        // Rejecting the parked waiters is not enough on its own: a worker that
+        // dies while the consumer is running its loop body has no waiter to
+        // reject, and without this the next `next()` would read the stream as
+        // a clean end.
+        let failure: unknown = null;
 
         const finish = (): void => {
           if (finished) return;
@@ -344,6 +349,7 @@ export class OffscreenRuntimeBridge {
         const fail = (err: unknown): void => {
           if (finished) return;
           finished = true;
+          failure = err;
           this.pending.delete(requestId);
           while (waiters.length > 0) waiters.shift()!.reject(err);
         };
@@ -399,6 +405,11 @@ export class OffscreenRuntimeBridge {
             if (queue.length > 0) {
               return Promise.resolve({ value: queue.shift()!, done: false });
             }
+            if (failure !== null) {
+              const err = failure;
+              failure = null;
+              return Promise.reject(err);
+            }
             if (finished) {
               return Promise.resolve({ value: undefined as T, done: true });
             }
@@ -415,6 +426,11 @@ export class OffscreenRuntimeBridge {
               const cancelMsg: WorkerRequest = { type: 'cancel', requestId };
               this.worker.postMessage(cancelMsg);
             }
+            // An explicit cancel outranks an error the consumer never asked
+            // about: `finish()` is a no-op once `fail()` has set `finished`,
+            // so drop the retained failure here rather than replaying it at
+            // the next `next()`.
+            failure = null;
             finish();
             return Promise.resolve({ value: undefined as T, done: true });
           },

bindings/web/packages/core/tests/unit/runtime/StreamFailureNotParked.test.ts

@@ -0,0 +1,255 @@
+/**
+ * StreamFailureNotParked.test.ts
+ *
+ * A stream failure has to reach the consumer even when the consumer is not
+ * parked on `next()` at the instant it happens.
+ *
+ * Both hand-rolled iterators here (`OffscreenRuntimeBridge.getStreamIterator`
+ * and `streamCallback` in `ProtoAdapterTypes`) reject the waiters that happen
+ * to be parked. Neither used to keep the error anywhere, so a failure that
+ * arrived while the consumer was between `next()` calls was reported to the
+ * following `next()` as a clean end-of-stream.
+ *
+ * Run from `bindings/web/packages/core`:
+ *
+ *     npx vitest run tests/unit/runtime/StreamFailureNotParked.test.ts
+ */
+
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import {
+  OffscreenRuntimeBridge,
+  setStreamWorkerInit,
+} from '../../../src/runtime/OffscreenRuntimeBridge';
+import { setStreamWorkerFactory } from '../../../src/runtime/StreamWorkerFactoryRegistry';
+import { streamCallback, type ModalityProtoModule } from '../../../src/Adapters/ProtoAdapterTypes';
+import {
+  BackendWorkerHost,
+  type BackendWorkerLike,
+} from '../../../src/runtime/BackendWorkerHost';
+import { runBackendWorker, type BackendWorkerScope } from '../../../src/runtime/BackendWorker';
+import type {
+  BackendWorkerRequest,
+  BackendWorkerResponse,
+} from '../../../src/runtime/BackendWorkerProtocol';
+import type { ProtoCodec } from '../../../src/runtime/ProtoWasm';
+import type { WorkerRequest, WorkerResponse } from '../../../src/runtime/StreamWorker';
+
+const uint32Codec: ProtoCodec<number> = {
+  encode(_message: number) {
+    return { finish: (): Uint8Array => new Uint8Array(0) };
+  },
+  decode(input: Uint8Array): number {
+    return new DataView(input.buffer, input.byteOffset, input.byteLength).getUint32(0, true);
+  },
+};
+
+function encodeU32(value: number): Uint8Array {
+  const out = new Uint8Array(4);
+  new DataView(out.buffer).setUint32(0, value, true);
+  return out;
+}
+
+/**
+ * Posts one callback for the stream request and then stays quiet until the
+ * test calls {@link crash}. Nothing is timer-driven, so "the consumer is
+ * between `next()` calls" is a fact of the test rather than a race it hopes
+ * to win.
+ */
+class CrashingWorker {
+  onmessage: ((ev: MessageEvent<WorkerResponse>) => void) | null = null;
+  onerror: ((ev: ErrorEvent) => void) | null = null;
+
+  private requestId: string | null = null;
+
+  constructor() {
+    queueMicrotask(() => this.deliver({ type: 'ready' }));
+  }
+
+  postMessage(msg: WorkerRequest): void {
+    // Everything that is not `init` or `cancel` is a stream request, and the
+    // union narrows to the variants carrying `requestId`.
+    if (msg.type === 'init' || msg.type === 'cancel') return;
+    this.requestId = msg.requestId;
+    this.deliver({
+      type: 'callback',
+      requestId: msg.requestId,
+      payloadBytes: encodeU32(1),
+    });
+  }
+
+  /** The worker dies mid-stream, the way a WASM OOM kills it. */
+  crash(message: string): void {
+    this.deliver({
+      type: 'error',
+      requestId: this.requestId ?? undefined,
+      message,
+    });
+  }
+
+  terminate(): void {}
+
+  private deliver(data: WorkerResponse): void {
+    this.onmessage?.({ data } as MessageEvent<WorkerResponse>);
+  }
+}
+
+interface FakeModule extends ModalityProtoModule {
+  /** Emit a uint32 value through the currently-installed callback. */
+  emitValue: (callbackPtr: number, value: number) => void;
+}
+
+/** Same shape as the fake in `Adapters/StreamLiveDelivery.test.ts`. */
+function makeFakeModule(): FakeModule {
+  const heap = new Uint8Array(4 * 1024);
+  const callbacks = new Map<number, (bytesPtr: number, size: number) => unknown>();
+  let nextPtr = 1;
+  // `streamCallback` treats bytesPtr === 0 as a null sentinel, so start at 1.
+  let heapCursor = 1;
+
+  const mod: Partial<FakeModule> = {
+    HEAPU8: heap,
+    addFunction(fn, _signature) {
+      const id = nextPtr++;
+      callbacks.set(id, fn as (bytesPtr: number, size: number) => unknown);
+      return id;
+    },
+    removeFunction(ptr) {
+      callbacks.delete(ptr);
+    },
+    emitValue(callbackPtr: number, value: number): void {
+      const fn = callbacks.get(callbackPtr);
+      if (!fn) throw new Error(`emitValue: no callback at ptr ${callbackPtr}`);
+      heap.set(encodeU32(value), heapCursor);
+      const ptr = heapCursor;
+      heapCursor += 4;
+      fn(ptr, 4);
+    },
+  };
+  return mod as FakeModule;
+}
+
+describe('a stream failure that lands while the consumer is not parked', () => {
+  let worker: CrashingWorker | null = null;
+
+  beforeEach(() => {
+    OffscreenRuntimeBridge.resetForTesting();
+    setStreamWorkerFactory(null);
+    setStreamWorkerInit(null);
+  });
+
+  afterEach(() => {
+    worker = null;
+    setStreamWorkerFactory(null);
+    setStreamWorkerInit(null);
+    OffscreenRuntimeBridge.resetForTesting();
+  });
+
+  it('is raised by the worker bridge rather than read as end-of-stream', async () => {
+    setStreamWorkerInit({ wasmBytes: new ArrayBuffer(0), moduleFactoryId: 'fake-factory' });
+    setStreamWorkerFactory(() => {
+      worker = new CrashingWorker();
+      return worker as unknown as Worker;
+    });
+
+    const bridge = OffscreenRuntimeBridge.tryGet('worker');
+    expect(bridge).not.toBeNull();
+
+    const stream = bridge!.getStreamIterator(
+      { kind: 'stream.llm.generate', handle: 0, requestBytes: new Uint8Array() },
+      uint32Codec,
+    );
+    const iterator = stream[Symbol.asyncIterator]();
+
+    expect(await iterator.next()).toEqual({ value: 1, done: false });
+
+    // No waiter is parked here, which is the state a consumer is in while it
+    // renders the token it just received.
+    worker!.crash('worker died');
+
+    await expect(iterator.next()).rejects.toThrow('worker died');
+  });
+
+  it('is raised by streamCallback rather than read as end-of-stream', async () => {
+    const fake = makeFakeModule();
+
+    // The first emit wakes the parked consumer; the second is buffered, so
+    // when the call throws there is no waiter left to reject.
+    const iterator = streamCallback(fake, uint32Codec, 'crashing_test_stream', (callbackPtr) => {
+      fake.emitValue(callbackPtr, 1);
+      fake.emitValue(callbackPtr, 2);
+      throw new Error('native call failed');
+    })[Symbol.asyncIterator]();
+
+    expect(await iterator.next()).toEqual({ value: 1, done: false });
+    expect(await iterator.next()).toEqual({ value: 2, done: false });
+
+    await expect(iterator.next()).rejects.toThrow('native call failed');
+  });
+
+  it('lets an explicit cancel outrank a failure the consumer never read', async () => {
+    setStreamWorkerInit({ wasmBytes: new ArrayBuffer(0), moduleFactoryId: 'fake-factory' });
+    setStreamWorkerFactory(() => {
+      worker = new CrashingWorker();
+      return worker as unknown as Worker;
+    });
+
+    const stream = OffscreenRuntimeBridge.tryGet('worker')!.getStreamIterator(
+      { kind: 'stream.llm.generate', handle: 0, requestBytes: new Uint8Array() },
+      uint32Codec,
+    );
+    const iterator = stream[Symbol.asyncIterator]();
+
+    expect(await iterator.next()).toEqual({ value: 1, done: false });
+    worker!.crash('worker died');
+
+    // The consumer walked away instead of reading the error, so it should not
+    // be handed that error afterwards.
+    await iterator.return?.();
+    expect(await iterator.next()).toEqual({ value: undefined, done: true });
+  });
+
+  it('is raised by the backend worker host rather than read as end-of-stream', async () => {
+    // Emits one event, then keeps the stream open so the crash is what ends it.
+    class OneThenHangWorker implements BackendWorkerLike {
+      onmessage: ((event: MessageEvent<BackendWorkerResponse>) => void) | null = null;
+      onerror: ((event: ErrorEvent) => void) | null = null;
+      private readonly scope: BackendWorkerScope;
+
+      constructor() {
+        this.scope = {
+          onmessage: null,
+          postMessage: (response) =>
+            this.onmessage?.({ data: response } as MessageEvent<BackendWorkerResponse>),
+        };
+        runBackendWorker(this.scope, {
+          init: () => undefined,
+          infer: (_kind, payload) => payload,
+          stream: async function* () {
+            yield 'first';
+            await new Promise<void>(() => {});
+          },
+          cancel: () => undefined,
+        });
+      }
+
+      postMessage(request: BackendWorkerRequest): void {
+        this.scope.onmessage?.({ data: request } as MessageEvent<BackendWorkerRequest>);
+      }
+
+      terminate(): void {}
+    }
+
+    const worker = new OneThenHangWorker();
+    const host = new BackendWorkerHost(() => worker);
+    const iterator = host.stream('tts.synthesize', { text: 'hello' })[Symbol.asyncIterator]();
+
+    expect(await iterator.next()).toEqual({ value: 'first', done: false });
+
+    // The worker dies while the consumer is handling the event it just got.
+    // handleCrash fails every in-flight request, and none of them is parked.
+    worker.onerror?.({ message: 'simulated crash' } as ErrorEvent);
+
+    await expect(iterator.next()).rejects.toThrow();
+  });
+});