fix(streaming): interrupt pending reads on cancellation

src/core/streaming.ts

@@ -11,20 +11,28 @@ type Bytes = string | ArrayBuffer | Uint8Array | null | undefined;
 
 type StreamTeeQueue<Item> = {
   readonly length: number;
+  readonly canceled: boolean;
   enqueue: (value: Promise<IteratorResult<Item>>) => void;
   dequeue: () => Promise<IteratorResult<Item>> | undefined;
+  cancel: () => void;
 };
 
 function createStreamTeeQueue<Item>(): StreamTeeQueue<Item> {
   let entries: (Promise<IteratorResult<Item>> | undefined)[] = [];
   let head = 0;
+  let canceled = false;
 
   return {
     get length() {
       return entries.length - head;
     },
+    get canceled() {
+      return canceled;
+    },
     enqueue(value) {
-      entries.push(value);
+      if (!canceled) {
+        entries.push(value);
+      }
     },
     dequeue() {
       if (head === entries.length) {
@@ -45,6 +53,11 @@ function createStreamTeeQueue<Item>(): StreamTeeQueue<Item> {
 
       return value;
     },
+    cancel() {
+      canceled = true;
+      entries.length = 0;
+      head = 0;
+    },
   };
 }
 
@@ -69,6 +82,7 @@ export class Stream<Item> implements AsyncIterable<Item> {
   /** Abort controller for the underlying request and all branches created with `tee()`. */
   controller: AbortController;
   #client: OpenAI | undefined;
+  #isTeeBranch = false;
   private iterator: () => AsyncIterator<Item>;
 
   /** Wraps an asynchronous event iterator and the controller that owns its request. */
@@ -278,35 +292,59 @@ export class Stream<Item> implements AsyncIterable<Item> {
   /**
    * Splits the stream into two streams which can be
    * independently read from at different speeds.
+   * Closing a branch discards its buffered events without stopping its sibling.
+   * Future reads on that branch finish immediately; previously issued `next()`
+   * promises remain shared with its sibling and may still resolve with events.
+   * Closing both branches invokes the source iterator's `return()` when available.
+   * For {@link Stream.fromReadableStream}, closing both branches before iteration
+   * starts does not cancel the supplied readable; cancel that readable directly.
    */
   tee(): [Stream<Item>, Stream<Item>] {
+    const { controller } = this;
     const left = createStreamTeeQueue<Item>();
     const right = createStreamTeeQueue<Item>();
     const iterator = this.iterator();
 
     const teeIterator = (queue: StreamTeeQueue<Item>): AsyncIterator<Item> => ({
       next: () => {
+        if (queue.canceled) {
+          return Promise.resolve({ value: undefined, done: true });
+        }
         if (queue.length === 0) {
           const result = iterator.next();
           left.enqueue(result);
           right.enqueue(result);
         }
         return queue.dequeue()!;
       },
+      return: async () => {
+        if (!queue.canceled) {
+          queue.cancel();
+          if (left.canceled && right.canceled) {
+            await this.#cancelIterator(iterator, controller);
+          }
+        }
+        return { value: undefined, done: true };
+      },
     });
 
-    return [
-      new Stream(() => teeIterator(left), this.controller, this.#client),
-      new Stream(() => teeIterator(right), this.controller, this.#client),
-    ];
+    const branch = (queue: StreamTeeQueue<Item>) => {
+      const stream = new Stream(() => teeIterator(queue), controller, this.#client);
+      stream.#isTeeBranch = true;
+      return stream;
+    };
+    return [branch(left), branch(right)];
   }
 
   /**
    * Converts this stream to a newline-separated ReadableStream of
    * JSON stringified values in the stream
    * which can be turned back into a Stream with `Stream.fromReadableStream()`.
+   * Canceling a response-backed readable aborts its request. Canceling a tee
+   * branch discards its buffered events and leaves sibling consumers running.
    */
   toReadableStream(): ReadableStream {
+    const { controller } = this;
     let iter: AsyncIterator<Item>;
 
     return makeReadableStream({
@@ -327,11 +365,19 @@ export class Stream<Item> implements AsyncIterable<Item> {
           ctrl.error(err);
         }
       },
-      async cancel() {
-        await iter.return?.();
-      },
+      cancel: () => this.#cancelIterator(iter, controller),
     });
   }
+
+  async #cancelIterator(iterator: AsyncIterator<Item>, controller: AbortController): Promise<void> {
+    const returnMethod = iterator.return;
+    if (returnMethod) {
+      if (!this.#isTeeBranch) {
+        controller.abort();
+      }
+      await Reflect.apply(returnMethod, iterator, []);
+    }
+  }
 }
 
 function createAbortableSSESource(body: NonNullable<Response['body']>, signal: AbortSignal) {

tests/lib/stream-tee-queue-performance.test.ts

@@ -136,6 +136,44 @@ describe('Stream.tee queue performance', () => {
     expect(next).toHaveBeenCalledTimes(QUEUE_SIZE);
   });
 
+  test.each([false, true])(
+    'releases canceled queues while a sibling keeps reading (nested: %s)',
+    async (nested) => {
+      const { stream, controller } = createNumberedStream(QUEUE_SIZE + 4);
+      const [left, right] = stream.tee();
+      const sibling = right[Symbol.asyncIterator]();
+      const { result: first, queues } = capturePromiseQueues(() => sibling.next());
+      await first;
+      const buffered = [...queues].find((queue) => queue[0] === first);
+      if (!buffered) {
+        throw new Error('Expected to capture the lagging reader queue');
+      }
+      await Promise.all(Array.from({ length: 3 }, () => sibling.next()));
+      expect(buffered).toHaveLength(4);
+      const branches = nested ? left.tee() : [left];
+      const readers = branches.map((branch) => branch.toReadableStream().getReader());
+
+      try {
+        await Promise.all(readers.map((reader) => reader.read()));
+        await Promise.all(readers.map((reader) => reader.cancel()));
+        expect(buffered).toHaveLength(0);
+
+        const results = await Promise.all(Array.from({ length: QUEUE_SIZE }, () => sibling.next()));
+        expect(results.map(({ value }) => value)).toEqual(
+          Array.from({ length: QUEUE_SIZE }, (_, index) => index + 4),
+        );
+        expect(buffered).toHaveLength(0);
+        expect(controller.signal.aborted).toBe(false);
+      } finally {
+        controller.abort();
+        await Promise.all(readers.map((reader) => reader.cancel()));
+        for (const reader of readers) {
+          reader.releaseLock();
+        }
+      }
+    },
+  );
+
   test('replays values independently for interleaved readers without extra source pulls', async () => {
     const { stream, next } = createNumberedStream(6);
     const [left, right] = stream.tee();
@@ -200,7 +238,7 @@ describe('Stream.tee queue performance', () => {
 
     expect(left.controller).toBe(controller);
     expect(right.controller).toBe(controller);
-    expect(left[Symbol.asyncIterator]().return).toBeUndefined();
+    expect(typeof left[Symbol.asyncIterator]().return).toBe('function');
 
     for await (const value of left) {
       if (value === 1) {