diff --git a/src/core/streaming.ts b/src/core/streaming.ts index e9d844b6e7..af9813f267 100644 --- a/src/core/streaming.ts +++ b/src/core/streaming.ts @@ -11,20 +11,28 @@ type Bytes = string | ArrayBuffer | Uint8Array | null | undefined; type StreamTeeQueue = { readonly length: number; + readonly canceled: boolean; enqueue: (value: Promise>) => void; dequeue: () => Promise> | undefined; + cancel: () => void; }; function createStreamTeeQueue(): StreamTeeQueue { let entries: (Promise> | 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(): StreamTeeQueue { return value; }, + cancel() { + canceled = true; + entries.length = 0; + head = 0; + }, }; } @@ -69,6 +82,7 @@ export class Stream implements AsyncIterable { /** Abort controller for the underlying request and all branches created with `tee()`. */ controller: AbortController; #client: OpenAI | undefined; + #isTeeBranch = false; private iterator: () => AsyncIterator; /** Wraps an asynchronous event iterator and the controller that owns its request. */ @@ -278,14 +292,24 @@ export class Stream implements AsyncIterable { /** * 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, Stream] { + const { controller } = this; const left = createStreamTeeQueue(); const right = createStreamTeeQueue(); const iterator = this.iterator(); const teeIterator = (queue: StreamTeeQueue): AsyncIterator => ({ next: () => { + if (queue.canceled) { + return Promise.resolve({ value: undefined, done: true }); + } if (queue.length === 0) { const result = iterator.next(); left.enqueue(result); @@ -293,20 +317,34 @@ export class Stream implements AsyncIterable { } 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) => { + 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; return makeReadableStream({ @@ -327,11 +365,19 @@ export class Stream implements AsyncIterable { ctrl.error(err); } }, - async cancel() { - await iter.return?.(); - }, + cancel: () => this.#cancelIterator(iter, controller), }); } + + async #cancelIterator(iterator: AsyncIterator, controller: AbortController): Promise { + const returnMethod = iterator.return; + if (returnMethod) { + if (!this.#isTeeBranch) { + controller.abort(); + } + await Reflect.apply(returnMethod, iterator, []); + } + } } function createAbortableSSESource(body: NonNullable, signal: AbortSignal) { diff --git a/tests/lib/stream-tee-queue-performance.test.ts b/tests/lib/stream-tee-queue-performance.test.ts index 51069277ec..555a2d798f 100644 --- a/tests/lib/stream-tee-queue-performance.test.ts +++ b/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) { diff --git a/tests/lib/streaming-core.test.ts b/tests/lib/streaming-core.test.ts index c4fb89400f..e8137a7372 100644 --- a/tests/lib/streaming-core.test.ts +++ b/tests/lib/streaming-core.test.ts @@ -4,6 +4,7 @@ import { APIError, APIUserAbortError, OpenAIError } from 'openai/core/error'; import { Stream, _iterSSEMessages } from 'openai/core/streaming'; import * as lineDecoders from 'openai/internal/decoders/line'; import { ReadableStreamFrom } from 'openai/internal/shims'; +import type { ResponseTextDeltaEvent } from 'openai/resources/responses/responses'; const encoder = new TextEncoder(); @@ -344,6 +345,16 @@ describe('Stream.tee', () => { }); describe('Stream.toReadableStream', () => { + const event: ResponseTextDeltaEvent = { + type: 'response.output_text.delta', + sequence_number: 0, + item_id: 'msg_test', + output_index: 0, + content_index: 0, + delta: 'hello', + logprobs: [], + }; + test('round-trips structured values as newline-delimited JSON', async () => { const original = new Stream( () => @@ -361,8 +372,12 @@ describe('Stream.toReadableStream', () => { await expect(collect(roundTripped)).resolves.toEqual([{ id: 1 }, { id: 2 }]); }); - test('closes the source iterator when the readable stream is canceled', async () => { + test.each([false, true])('closes the source iterator (return rejects: %s)', async (rejects) => { + const failure = new OpenAIError('source cleanup failed'); const returned = vi.fn().mockResolvedValue({ done: true, value: undefined }); + if (rejects) { + returned.mockRejectedValue(failure); + } const source = new Stream( () => ({ next: vi.fn().mockResolvedValue({ done: false, value: { id: 1 } }), @@ -373,9 +388,230 @@ describe('Stream.toReadableStream', () => { const reader = source.toReadableStream().getReader(); await reader.read(); - await reader.cancel(); + await (rejects ? expect(reader.cancel()).rejects.toBe(failure) : reader.cancel()); expect(returned).toHaveBeenCalledTimes(1); + reader.releaseLock(); + }); + + test('reads an iterator return hook once and preserves its receiver', async () => { + const returned = vi.fn().mockResolvedValue({ done: true, value: undefined }); + let returnReads = 0; + const iterator = { + next: vi.fn().mockResolvedValue({ done: false, value: { id: 1 } }), + get return() { + returnReads += 1; + if (returnReads > 1) { + throw new OpenAIError('return hook was read more than once'); + } + return returned; + }, + }; + const source = new Stream(() => iterator, new AbortController()); + const reader = source.toReadableStream().getReader(); + + try { + await reader.read(); + await expect(reader.cancel()).resolves.toBeUndefined(); + expect(returnReads).toBe(1); + expect(returned).toHaveBeenCalledTimes(1); + expect(returned.mock.contexts[0]).toBe(iterator); + } finally { + reader.releaseLock(); + } + }); + + test.each([ + ['SSE', false], + ['SSE', true], + ['NDJSON', false], + ['NDJSON', true], + ] as const)('cancels a pending %s read (after one event: %s)', async (format, afterEvent) => { + const cancel = vi.fn(); + const pull = vi.fn((controller: ReadableStreamDefaultController) => { + if (afterEvent && pull.mock.calls.length === 1) { + const json = JSON.stringify(event); + controller.enqueue(encoder.encode(format === 'SSE' ? `data: ${json}\n\n` : `${json}\n`)); + } + }); + const body = new ReadableStream({ pull, cancel }, { highWaterMark: 0 }); + const fetch = vi.fn(async () => new Response(body, { headers: { 'content-type': 'text/event-stream' } })); + const client = new OpenAI({ apiKey: 'test-key', maxRetries: 0, fetch }); + const stream = + format === 'SSE' + ? await client.responses.create({ model: 'gpt-4o', input: 'hello', stream: true }) + : Stream.fromReadableStream(body, new AbortController()); + const reader = stream.toReadableStream().getReader(); + let pending: ReturnType | undefined; + let cancellation: Promise | undefined; + + try { + if (afterEvent) { + const first = await reader.read(); + expect(new TextDecoder().decode(first.value)).toBe(`${JSON.stringify(event)}\n`); + } + pending = reader.read(); + await vi.waitFor(() => expect(pull).toHaveBeenCalledTimes(afterEvent ? 2 : 1)); + cancellation = reader.cancel(); + + await expect(settlesSoon(cancellation)).resolves.toBeUndefined(); + await expect(pending).resolves.toEqual({ done: true, value: undefined }); + expect(stream.controller.signal.aborted).toBe(true); + expect(cancel).toHaveBeenCalledTimes(1); + expect(body.locked).toBe(false); + expect(fetch).toHaveBeenCalledTimes(format === 'SSE' ? 1 : 0); + } finally { + stream.controller.abort(); + await (cancellation ?? reader.cancel()); + await pending; + reader.releaseLock(); + } + }); + + test('canceling a pending tee branch leaves its sibling usable', async () => { + const cancel = vi.fn(); + const pull = vi.fn(); + let sendEvent = () => {}; + const body = new ReadableStream( + { + start(controller) { + sendEvent = () => { + controller.enqueue(encoder.encode(`data: ${JSON.stringify(event)}\n\n`)); + controller.close(); + }; + }, + pull, + cancel, + }, + { highWaterMark: 0 }, + ); + const client = new OpenAI({ + apiKey: 'test-key', + maxRetries: 0, + fetch: async () => new Response(body, { headers: { 'content-type': 'text/event-stream' } }), + }); + const stream = await client.responses.create({ model: 'gpt-4o', input: 'hello', stream: true }); + const [left, right] = stream.tee(); + const reader = left.toReadableStream().getReader(); + const sibling = right[Symbol.asyncIterator](); + const pending = reader.read(); + + try { + await vi.waitFor(() => expect(pull).toHaveBeenCalledTimes(1)); + await expect(settlesSoon(reader.cancel())).resolves.toBeUndefined(); + expect(stream.controller.signal.aborted).toBe(false); + expect(cancel).not.toHaveBeenCalled(); + + sendEvent(); + await expect(settlesSoon(sibling.next())).resolves.toEqual({ done: false, value: event }); + await expect(settlesSoon(sibling.next())).resolves.toEqual({ done: true, value: undefined }); + expect(stream.controller.signal.aborted).toBe(false); + expect(body.locked).toBe(false); + } finally { + stream.controller.abort(); + await reader.cancel(); + await pending; + reader.releaseLock(); + } + }); + + test.each([false, true])( + 'keeps a long-running sibling usable after tee cancellation (nested: %s)', + async (nested) => { + const cancel = vi.fn(); + let sequence = 0; + const body = new ReadableStream( + { + pull(controller) { + controller.enqueue( + encoder.encode(`data: ${JSON.stringify({ ...event, sequence_number: sequence })}\n\n`), + ); + sequence += 1; + }, + cancel, + }, + { highWaterMark: 0 }, + ); + const client = new OpenAI({ + apiKey: 'test-key', + maxRetries: 0, + fetch: async () => new Response(body, { headers: { 'content-type': 'text/event-stream' } }), + }); + const stream = await client.responses.create({ model: 'gpt-4o', input: 'hello', stream: true }); + const [left, right] = stream.tee(); + const branches = nested ? left.tee() : [left]; + const readers = branches.map((branch) => branch.toReadableStream().getReader()); + const sibling = right[Symbol.asyncIterator](); + + try { + await Promise.all(readers.map((reader) => reader.read())); + await Promise.all(readers.map((reader) => reader.cancel())); + + const results = await Promise.all(Array.from({ length: 2048 }, () => sibling.next())); + expect(results.every((result) => !result.done)).toBe(true); + expect(results.map((result) => result.value.sequence_number)).toEqual( + Array.from({ length: 2048 }, (_, index) => index), + ); + await expect(Promise.all(readers.map((reader) => reader.read()))).resolves.toEqual( + readers.map(() => ({ done: true, value: undefined })), + ); + expect(stream.controller.signal.aborted).toBe(false); + expect(cancel).not.toHaveBeenCalled(); + await expect(sibling.next()).resolves.toMatchObject({ + done: false, + value: { sequence_number: 2048 }, + }); + } finally { + stream.controller.abort(); + await sibling.next(); + await Promise.all(readers.map((reader) => reader.cancel())); + for (const reader of readers) { + reader.releaseLock(); + } + } + }, + ); + + test('canceling both pending tee branches closes their response once', async () => { + const cancel = vi.fn(); + const pull = vi.fn(); + const body = new ReadableStream({ pull, cancel }, { highWaterMark: 0 }); + const client = new OpenAI({ + apiKey: 'test-key', + maxRetries: 0, + fetch: async () => new Response(body, { headers: { 'content-type': 'text/event-stream' } }), + }); + const stream = await client.responses.create({ model: 'gpt-4o', input: 'hello', stream: true }); + const [left, right] = stream.tee(); + const first = left.toReadableStream().getReader(); + const second = right.toReadableStream().getReader(); + const readers = [first, second]; + const pending = readers.map((reader) => reader.read()); + + try { + await vi.waitFor(() => expect(pull).toHaveBeenCalledTimes(1)); + await expect(settlesSoon(first.cancel())).resolves.toBeUndefined(); + expect(stream.controller.signal.aborted).toBe(false); + expect(cancel).not.toHaveBeenCalled(); + await expect(settlesSoon(second.cancel())).resolves.toBeUndefined(); + + expect(stream.controller.signal.aborted).toBe(true); + expect(cancel).toHaveBeenCalledTimes(1); + expect(body.locked).toBe(false); + await expect(Promise.all(pending)).resolves.toEqual([ + { done: true, value: undefined }, + { done: true, value: undefined }, + ]); + await Promise.all(readers.map((reader) => reader.cancel())); + expect(cancel).toHaveBeenCalledTimes(1); + } finally { + stream.controller.abort(); + await Promise.all(readers.map((reader) => reader.cancel())); + await Promise.all(pending); + for (const reader of readers) { + reader.releaseLock(); + } + } }); test('propagates iterator failures through the readable stream', async () => {