diff --git a/src/lib/AssistantStream.ts b/src/lib/AssistantStream.ts index 9a4ecaa343..f6be42dfad 100644 --- a/src/lib/AssistantStream.ts +++ b/src/lib/AssistantStream.ts @@ -117,66 +117,15 @@ export class AssistantStream /** Iterates over cloned raw assistant events; stopping early aborts the underlying request. */ [Symbol.asyncIterator](): AsyncIterator { - const pushQueue: AssistantStreamEvent[] = []; - const readQueue: { - resolve: (chunk: AssistantStreamEvent | undefined) => void; - reject: (err: unknown) => void; - }[] = []; - let done = false; - - //Catch all for passing along all events - this.on('event', (event) => { - const eventCopy = structuredClone(event); - const reader = readQueue.shift(); - if (reader) { - reader.resolve(eventCopy); - } else { - pushQueue.push(eventCopy); - } - }); - - this.on('end', () => { - done = true; - for (const reader of readQueue) { - reader.resolve(undefined); - } - readQueue.length = 0; - }); - - this.on('abort', (err) => { - done = true; - for (const reader of readQueue) { - reader.reject(err); - } - readQueue.length = 0; - }); - - this.on('error', (err) => { - done = true; - for (const reader of readQueue) { - reader.reject(err); - } - readQueue.length = 0; - }); - - return { - next: async (): Promise> => { - if (!pushQueue.length) { - if (done) { - return { value: undefined, done: true }; - } - return new Promise((resolve, reject) => - readQueue.push({ resolve, reject }), - ).then((chunk) => (chunk ? { value: chunk, done: false } : { value: undefined, done: true })); - } - const chunk = pushQueue.shift()!; - return { value: chunk, done: false }; + return this._createIterator( + (push) => { + //Catch all for passing along all events + const onEvent = (event: AssistantStreamEvent) => push(structuredClone(event)); + this.on('event', onEvent); + return () => this.off('event', onEvent); }, - return: async () => { - this.abort(); - return { value: undefined, done: true }; - }, - }; + { onReturn: () => this.abort() }, + ); } /** Restores an assistant stream from events serialized by `toReadableStream()`. */ diff --git a/src/lib/ChatCompletionStream.ts b/src/lib/ChatCompletionStream.ts index 0662a84287..2066aead2a 100644 --- a/src/lib/ChatCompletionStream.ts +++ b/src/lib/ChatCompletionStream.ts @@ -774,64 +774,14 @@ export class ChatCompletionStream /** Iterates over raw API chunks; stopping iteration early aborts the underlying request. */ [Symbol.asyncIterator](this: ChatCompletionStream): AsyncIterator { - const pushQueue: ChatCompletionChunk[] = []; - const readQueue: { - resolve: (chunk: ChatCompletionChunk | undefined) => void; - reject: (err: unknown) => void; - }[] = []; - let done = false; - - this.on('chunk', (chunk) => { - const reader = readQueue.shift(); - if (reader) { - reader.resolve(chunk); - } else { - pushQueue.push(chunk); - } - }); - - this.on('end', () => { - done = true; - for (const reader of readQueue) { - reader.resolve(undefined); - } - readQueue.length = 0; - }); - - this.on('abort', (err) => { - done = true; - for (const reader of readQueue) { - reader.reject(err); - } - readQueue.length = 0; - }); - - this.on('error', (err) => { - done = true; - for (const reader of readQueue) { - reader.reject(err); - } - readQueue.length = 0; - }); - - return { - next: async (): Promise> => { - if (!pushQueue.length) { - if (done) { - return { value: undefined, done: true }; - } - return new Promise((resolve, reject) => - readQueue.push({ resolve, reject }), - ).then((chunk) => (chunk ? { value: chunk, done: false } : { value: undefined, done: true })); - } - const chunk = pushQueue.shift()!; - return { value: chunk, done: false }; + return this._createIterator( + (push) => { + const onChunk = (chunk: ChatCompletionChunk) => push(chunk); + this.on('chunk', onChunk); + return () => this.off('chunk', onChunk); }, - return: async () => { - this.abort(); - return { value: undefined, done: true }; - }, - }; + { onReturn: () => this.abort() }, + ); } /** Serializes raw completion chunks into a readable stream for transfer to another runtime. */ diff --git a/src/lib/ChatCompletionStreamingRunner.ts b/src/lib/ChatCompletionStreamingRunner.ts index 575a7936f3..6682024b14 100644 --- a/src/lib/ChatCompletionStreamingRunner.ts +++ b/src/lib/ChatCompletionStreamingRunner.ts @@ -76,89 +76,39 @@ export class ChatCompletionStreamingRunner /** Serializes completion chunks and tool-result messages for replay in another runtime. */ override toReadableStream(): ReadableStream { - const pushQueue: ChatCompletionReadableStreamItem[] = []; - const readQueue: { - resolve: (event: ChatCompletionReadableStreamItem | undefined) => void; - reject: (err: unknown) => void; - }[] = []; - let done = false; let lastChunk: ChatCompletionChunk | undefined; let toolCallIds: string[] | undefined; - const pushEvent = (event: ChatCompletionReadableStreamItem) => { - const reader = readQueue.shift(); - if (reader) { - reader.resolve(event); - } else { - pushQueue.push(event); - } - }; - - this.on('chunk', (chunk) => { - lastChunk = chunk; - pushEvent(chunk); - }); - this.on('message', (message: ChatCompletionMessageParam) => { - if (isAssistantMessage(message)) { - toolCallIds = message.tool_calls?.map((toolCall) => toolCall.id); - return; - } - - if (isToolMessage(message)) { - if (!lastChunk) { - throw new OpenAIError('cannot serialize a tool message before receiving any chunks'); - } - pushEvent(makeChatCompletionReadableStreamMessageChunk(lastChunk, message, toolCallIds)); - } - }); - - this.on('end', () => { - done = true; - for (const reader of readQueue) { - reader.resolve(undefined); - } - readQueue.length = 0; - }); - - this.on('abort', (err) => { - done = true; - for (const reader of readQueue) { - reader.reject(err); - } - readQueue.length = 0; - }); - - this.on('error', (err) => { - done = true; - for (const reader of readQueue) { - reader.reject(err); - } - readQueue.length = 0; - }); + const iterator = this._createIterator( + (push) => { + const onChunk = (chunk: ChatCompletionChunk) => { + lastChunk = chunk; + push(chunk); + }; + const onMessage = (message: ChatCompletionMessageParam) => { + if (isAssistantMessage(message)) { + toolCallIds = message.tool_calls?.map((toolCall) => toolCall.id); + return; + } - const iterator = (): AsyncIterator => ({ - next: async (): Promise> => { - if (!pushQueue.length) { - if (done) { - return { value: undefined, done: true }; + if (isToolMessage(message)) { + if (!lastChunk) { + throw new OpenAIError('cannot serialize a tool message before receiving any chunks'); + } + push(makeChatCompletionReadableStreamMessageChunk(lastChunk, message, toolCallIds)); } - return new Promise((resolve, reject) => - readQueue.push({ resolve, reject }), - ).then((event) => (event ? { value: event, done: false } : { value: undefined, done: true })); - } - const event = pushQueue.shift(); - if (!event) { - return { value: undefined, done: true }; - } - return { value: event, done: false }; + }; + this.on('chunk', onChunk); + this.on('message', onMessage); + return () => { + this.off('chunk', onChunk); + this.off('message', onMessage); + }; }, - return: async () => { - this.abort(); - return { value: undefined, done: true }; - }, - }); + { onReturn: () => this.abort() }, + ); - const stream = new Stream(iterator, this.controller); + const stream = new Stream(() => iterator, this.controller); return stream.toReadableStream(); } diff --git a/src/lib/EventStream.ts b/src/lib/EventStream.ts index dae896495f..d324f2f77f 100644 --- a/src/lib/EventStream.ts +++ b/src/lib/EventStream.ts @@ -225,17 +225,50 @@ export class EventStream { event: Event, ): AsyncIterableIterator> { type Parameters = EventParameters; - type Result = IteratorResult; + return this._createIterator( + (push) => { + const onEvent = (...args: Parameters) => push(args); + this.on(event, onEvent as EventListener); + return () => this.off(event, onEvent as EventListener); + }, + { + // When iterating the 'error' or 'abort' event itself, yield it as a + // value instead of rejecting the iterator. + rejectOnError: event !== 'error', + rejectOnAbort: event !== 'abort', + }, + ); + } + + /** + * Shared buffered async-iterator adapter over this stream's events. + * + * `attach` registers the producer listener(s) with the given `push` and + * returns a cleanup function that removes them. Termination is handled + * here: the iterator ends when the stream ends, listeners are removed on + * end/return, and a terminal error is retained until buffered values have + * drained so it is surfaced even when no reader was waiting when it fired. + */ + protected _createIterator( + attach: (push: (value: T) => void) => () => void, + { + rejectOnError = true, + rejectOnAbort = true, + onReturn, + }: { rejectOnError?: boolean; rejectOnAbort?: boolean; onReturn?: () => void } = {}, + ): AsyncIterableIterator { + type Result = IteratorResult; type Reader = { resolve: (result: Result) => void; reject: (error: OpenAIError) => void; }; - const pushQueue: Parameters[] = []; + const pushQueue: T[] = []; const readQueue: Reader[] = []; let ended = this.ended; let failure: OpenAIError | undefined; let failureDelivered = false; + let detach: () => void = () => undefined; const doneResult = (): Result => ({ value: undefined as never, done: true }); const finishReaders = () => { @@ -251,24 +284,24 @@ export class EventStream { readQueue.shift()!.reject(failure); }; const cleanup = () => { - this.off(event, onEvent as EventListener); + detach(); this.off('end', onEnd); - if (event !== 'error') { + if (rejectOnError) { this.off('error', onFailure); } - if (event !== 'abort') { + if (rejectOnAbort) { this.off('abort', onFailure); } }; - const onEvent = (...args: Parameters) => { + const push = (value: T) => { if (ended) { return; } const reader = readQueue.shift(); if (reader) { - reader.resolve({ value: args, done: false }); + reader.resolve({ value, done: false }); } else { - pushQueue.push(args); + pushQueue.push(value); } }; const onFailure = (error: OpenAIError) => { @@ -287,21 +320,20 @@ export class EventStream { }; if (!ended) { - this.on(event, onEvent as EventListener); + detach = attach(push); this.on('end', onEnd); - if (event !== 'error') { + if (rejectOnError) { this.on('error', onFailure); } - if (event !== 'abort') { + if (rejectOnAbort) { this.on('abort', onFailure); } } return { - next: () => { - const value = pushQueue.shift(); - if (value) { - return Promise.resolve({ value, done: false }); + next: (): Promise => { + if (pushQueue.length) { + return Promise.resolve({ value: pushQueue.shift()!, done: false }); } if (failure && !failureDelivered) { @@ -322,6 +354,14 @@ export class EventStream { pushQueue.length = 0; cleanup(); finishReaders(); + if (onReturn) { + // The consumer explicitly ended iteration, so any failure the + // onReturn callback triggers (e.g. aborting the stream) is + // self-inflicted; mark the stream's terminal promise as handled so + // it does not surface as an unhandled rejection. + void this.done().catch(() => undefined); + onReturn(); + } return Promise.resolve(doneResult()); }, [Symbol.asyncIterator]() { diff --git a/src/lib/responses/ResponseStream.ts b/src/lib/responses/ResponseStream.ts index 0834aede59..bd7336656f 100644 --- a/src/lib/responses/ResponseStream.ts +++ b/src/lib/responses/ResponseStream.ts @@ -265,17 +265,14 @@ export class ResponseStream /** Iterates over response events; stopping iteration early aborts the underlying request. */ [Symbol.asyncIterator](this: ResponseStream): AsyncIterator { - const iterator = this.events('event'); - return { - next: async () => { - const result = await iterator.next(); - return result.done ? { value: undefined, done: true } : { value: result.value[0], done: false }; + return this._createIterator( + (push) => { + const onEvent = (event: ResponseStreamEvent) => push(event); + this.on('event', onEvent); + return () => this.off('event', onEvent); }, - return: async () => { - this.abort(); - return { value: undefined, done: true }; - }, - }; + { onReturn: () => this.abort() }, + ); } /** diff --git a/tests/lib/AssistantStream.test.ts b/tests/lib/AssistantStream.test.ts index 8393ca4a34..1058be4c7d 100644 --- a/tests/lib/AssistantStream.test.ts +++ b/tests/lib/AssistantStream.test.ts @@ -659,6 +659,38 @@ describe('AssistantStream factories and async iteration', () => { await expect(pendingAbort).rejects.toBe(abortError); }); + test('drains cloned queued events before rejecting a terminal stream error', async () => { + const runner = new AssistantStream(); + const iterator = runner[Symbol.asyncIterator](); + const event = completedRun('run_original'); + const error = new OpenAIError('stream failed after an event'); + + runner._emit('event', event as AssistantStreamEvent); + event.data.id = 'run_mutated_after_emit'; + runner._emit('error', error); + + await expect(iterator.next()).resolves.toEqual({ + value: completedRun('run_original'), + done: false, + }); + await expect(iterator.next()).rejects.toBe(error); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); + + test('drains queued events before rejecting a terminal stream abort', async () => { + const runner = new AssistantStream(); + const iterator = runner[Symbol.asyncIterator](); + const event = completedRun(); + const error = new APIUserAbortError(); + + runner._emit('event', event as AssistantStreamEvent); + runner._emit('abort', error); + + await expect(iterator.next()).resolves.toEqual({ value: event, done: false }); + await expect(iterator.next()).rejects.toBe(error); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); + test('closes pending event reads when an otherwise idle stream ends', async () => { const runner = new AssistantStream(); const pending = runner[Symbol.asyncIterator]().next(); diff --git a/tests/lib/ChatCompletionStream.test.ts b/tests/lib/ChatCompletionStream.test.ts index 98a9caeea4..06e4e39b94 100644 --- a/tests/lib/ChatCompletionStream.test.ts +++ b/tests/lib/ChatCompletionStream.test.ts @@ -1,5 +1,6 @@ import { vi } from 'vitest'; import type OpenAI from 'openai'; +import { OpenAIError } from 'openai/error'; import { zodResponseFormat } from 'openai/helpers/zod'; import { ChatCompletionStream } from 'openai/lib/ChatCompletionStream'; import type { ChatCompletionSnapshot } from 'openai/lib/ChatCompletionStream'; @@ -841,6 +842,198 @@ describe('.stream()', () => { expect(capturedLogProbs?.length).toEqual(choice?.logprobs?.refusal?.length); }); + it('surfaces a mid-stream error when chunks are buffered before consumption', async () => { + const chunks = [ + { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { role: 'assistant', content: 'hel' }, finish_reason: null }], + }, + { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { content: 'lo' }, finish_reason: null }], + }, + ] as unknown as OpenAI.Chat.ChatCompletionChunk[]; + // Yield valid chunks, then throw to error the stream after they have been + // delivered (mimics a connection drop mid-response). + const readable = new Stream(async function* failingChunks() { + for (const chunk of chunks) { + yield chunk; + } + throw new Error('network boom'); + }, new AbortController()).toReadableStream(); + + const stream = ChatCompletionStream.fromReadableStream(readable); + // Grab the iterator (registering its listeners) but do not consume yet, so + // the valid chunks and the error land while no reader is waiting: they + // buffer in the iterator's internal queue instead of rejecting a pending + // reader. + const iterator = stream[Symbol.asyncIterator](); + // Wait for the stream's terminal signal so the chunks and the error have + // definitely been emitted before we start reading. + await stream.done().catch(() => {}); + + const collected: OpenAI.Chat.ChatCompletionChunk[] = []; + let caught: unknown = null; + try { + for await (const chunk of { [Symbol.asyncIterator]: () => iterator }) { + collected.push(chunk); + } + } catch (error) { + caught = error; + } + + expect(collected).toHaveLength(chunks.length); + expect(caught).toBeInstanceOf(OpenAIError); + expect((caught as OpenAIError).message).toBe('network boom'); + }); + + it('rejects a pending read exactly once when the stream errors while a reader is waiting', async () => { + const chunks = [ + { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { role: 'assistant', content: 'hel' }, finish_reason: null }], + }, + { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { content: 'lo' }, finish_reason: null }], + }, + ] as unknown as OpenAI.Chat.ChatCompletionChunk[]; + const readable = new Stream(async function* failingChunks() { + for (const chunk of chunks) { + yield chunk; + } + throw new Error('network boom'); + }, new AbortController()).toReadableStream(); + + const stream = ChatCompletionStream.fromReadableStream(readable); + // Consume eagerly so each read is awaiting when its chunk (and finally the + // error) arrives, exercising the pending-reader path rather than the + // buffered path. + const iterator = stream[Symbol.asyncIterator](); + + await expect(iterator.next()).resolves.toMatchObject({ done: false }); + await expect(iterator.next()).resolves.toMatchObject({ done: false }); + const caught = await iterator.next().then( + () => null, + (error) => error, + ); + expect(caught).toBeInstanceOf(OpenAIError); + expect((caught as OpenAIError).message).toBe('network boom'); + // The failure is delivered exactly once; iteration then ends cleanly. + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); + + it('aborts the stream when the consumer breaks out of iteration', async () => { + const readable = new Stream( + async function* unfinishedChunks(): AsyncGenerator { + yield { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { role: 'assistant', content: 'hel' }, finish_reason: null }], + } as unknown as OpenAI.Chat.ChatCompletionChunk; + // Hang so the only way the consumer stops is by breaking out. + await Promise.race([]); + }, + new AbortController(), + ).toReadableStream(); + + const stream = ChatCompletionStream.fromReadableStream(readable); + for await (const chunk of stream) { + if (chunk.choices[0]?.delta.content === 'hel') { + break; + } + } + + expect(stream.controller.signal.aborted).toBe(true); + }); + + it('returns done immediately when iterating after the stream has ended', async () => { + const chunks = [ + { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { role: 'assistant', content: 'hello' }, finish_reason: 'stop' }], + }, + ] as unknown as OpenAI.Chat.ChatCompletionChunk[]; + const readable = new Stream(async function* completeChunks() { + for (const chunk of chunks) { + yield chunk; + } + }, new AbortController()).toReadableStream(); + + const stream = ChatCompletionStream.fromReadableStream(readable); + await stream.done(); + + const iterator = stream[Symbol.asyncIterator](); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); + + it('toReadableStream surfaces a mid-stream error when items are buffered before consumption', async () => { + const chunks = [ + { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { role: 'assistant', content: 'hel' }, finish_reason: null }], + }, + { + id: 'chatcmpl-test', + object: 'chat.completion.chunk', + created: 1, + model: 'gpt-4', + choices: [{ index: 0, delta: { content: 'lo' }, finish_reason: null }], + }, + ] as unknown as OpenAI.Chat.ChatCompletionChunk[]; + const readable = new Stream(async function* failingChunks() { + for (const chunk of chunks) { + yield chunk; + } + throw new Error('network boom'); + }, new AbortController()).toReadableStream(); + + const runner = ChatCompletionStreamingRunner.fromReadableStream(readable); + // Bridge to a ReadableStream immediately (registering its listeners) but + // do not read from it until the runner has already errored, so the chunks + // and the error land while nothing is pulling: they buffer in the + // adapter's internal queue instead of rejecting a pending reader. + const proxied = Stream.fromReadableStream( + runner.toReadableStream(), + new AbortController(), + ); + await runner.done().catch(() => {}); + + const collected: OpenAI.Chat.ChatCompletionChunk[] = []; + let caught: unknown = null; + try { + for await (const chunk of proxied) { + collected.push(chunk); + } + } catch (error) { + caught = error; + } + + expect(collected).toHaveLength(chunks.length); + expect(caught).toBeInstanceOf(OpenAIError); + expect((caught as OpenAIError).message).toBe('network boom'); + }); + it('preserves the existing streamed function-call detail type', () => { const legacyFunction: ChatCompletionSnapshot.Choice.Message.ToolCall.Function = { name: 'get_weather', diff --git a/tests/lib/EventStream.test.ts b/tests/lib/EventStream.test.ts index 9254063e14..8febd20533 100644 --- a/tests/lib/EventStream.test.ts +++ b/tests/lib/EventStream.test.ts @@ -1,5 +1,5 @@ import { vi } from 'vitest'; -import { OpenAIError } from 'openai/error'; +import { APIUserAbortError, OpenAIError } from 'openai/error'; import { EventStream } from 'openai/lib/EventStream'; import type { BaseEvents } from 'openai/lib/EventStream'; @@ -16,6 +16,10 @@ class TestStream extends EventStream { this._emit('error', error); } + emitAbort(error: APIUserAbortError) { + this._emit('abort', error); + } + end() { this._emit('end'); } @@ -41,6 +45,29 @@ describe('EventStream.emitted', () => { await expect(pending).rejects.toBe(failure); }); + test('removes the error listener after the requested event arrives', async () => { + const stream = new TestStream(); + const removeListener = vi.spyOn(stream, 'off'); + const pending = stream.emitted('foo'); + + stream.emitFoo('received', 4); + + await expect(pending).resolves.toEqual(['received', 4]); + expect(removeListener).toHaveBeenCalledWith('error', expect.any(Function)); + }); + + test('removes the requested-event listener when an error arrives first', async () => { + const stream = new TestStream(); + const removeListener = vi.spyOn(stream, 'off'); + const pending = stream.emitted('foo'); + const failure = new OpenAIError('stream failed'); + + stream.emitError(failure); + + await expect(pending).rejects.toBe(failure); + expect(removeListener).toHaveBeenCalledWith('foo', expect.any(Function)); + }); + test('resolves rather than rejects when waiting for the error event itself', async () => { const stream = new TestStream(); const pending = stream.emitted('error'); @@ -99,6 +126,60 @@ describe('EventStream.events', () => { await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); }); + test('drains queued events before rejecting on abort', async () => { + const stream = new TestStream(); + const iterator = stream.events('foo'); + const error = new APIUserAbortError(); + + stream.emitFoo('first', 1); + stream.emitAbort(error); + + await expect(iterator.next()).resolves.toEqual({ value: ['first', 1], done: false }); + await expect(iterator.next()).rejects.toBe(error); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); + + test("yields the 'error' event as a value instead of rejecting when iterating it", async () => { + const stream = new TestStream(); + const iterator = stream.events('error'); + const error = new OpenAIError('oops'); + + stream.emitError(error); + + await expect(iterator.next()).resolves.toEqual({ value: [error], done: false }); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); + + test("yields the 'abort' event as a value instead of rejecting when iterating it", async () => { + const stream = new TestStream(); + const iterator = stream.events('abort'); + const error = new APIUserAbortError(); + + stream.emitAbort(error); + + await expect(iterator.next()).resolves.toEqual({ value: [error], done: false }); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); + + test.each(['end', 'return'] as const)( + 'removes producer and lifecycle listeners on %s', + async (termination) => { + const stream = new TestStream(); + const removeListener = vi.spyOn(stream, 'off'); + const iterator = stream.events('foo'); + + if (termination === 'end') { + stream.end(); + } else { + await iterator.return?.(); + } + + for (const event of ['foo', 'end', 'error', 'abort'] as const) { + expect(removeListener).toHaveBeenCalledWith(event, expect.any(Function)); + } + }, + ); + test('does not suppress errors after iterator cleanup', async () => { const stream = new TestStream(); const iterator = stream.events('foo'); diff --git a/tests/lib/ResponseStream.test.ts b/tests/lib/ResponseStream.test.ts index e6a57dddb7..a5400c70d8 100644 --- a/tests/lib/ResponseStream.test.ts +++ b/tests/lib/ResponseStream.test.ts @@ -638,6 +638,56 @@ describe('.stream()', () => { } expect(final.output_text).toBe('The answer is 42'); }); + + it('surfaces a mid-stream error when events are buffered before consumption', async () => { + // Two valid events, then a malformed delta that references a missing output + // index so accumulation throws mid-stream (the stream itself closes cleanly, + // so the two earlier events are delivered). + const validEvents: ResponseStreamEvent[] = [ + { type: 'response.created', sequence_number: 0, response: makeResponse() }, + { + type: 'response.output_item.added', + sequence_number: 1, + output_index: 0, + item: { id: 'msg_1', type: 'message', role: 'assistant', status: 'in_progress', content: [] }, + }, + ]; + const malformedEvent = { + type: 'response.output_text.delta', + sequence_number: 2, + item_id: 'msg_1', + output_index: 99, + content_index: 0, + delta: 'boom', + logprobs: [], + } as unknown as ResponseStreamEvent; + + const stream = ResponseStream.fromReadableStream( + readableStreamFromEvents([...validEvents, malformedEvent]), + ); + // Grab the iterator (registering its listeners) but do not consume yet, so + // the valid events and the error land while no reader is waiting: they + // buffer in the iterator's internal queue instead of rejecting a pending + // reader. + const iterator = stream[Symbol.asyncIterator](); + // Wait for the stream's terminal signal so the events and the error have + // definitely been emitted before we start reading. + await expect(stream.done()).rejects.toThrow('missing output at index 99'); + + await expect(iterator.next()).resolves.toEqual({ value: validEvents[0], done: false }); + await expect(iterator.next()).resolves.toEqual({ value: validEvents[1], done: false }); + + const failure = await iterator.next().then( + () => { + throw new Error('Expected the response iterator to reject'); + }, + (error: unknown) => error, + ); + + expect(failure).toBeInstanceOf(OpenAIError); + expect((failure as OpenAIError).message).toBe('missing output at index 99'); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); }); function readableStreamFromEvents(events: ResponseStreamEvent[]) { diff --git a/tests/streaming/assistants/assistant.test.ts b/tests/streaming/assistants/assistant.test.ts index 8d03ba64b0..752e10d839 100644 --- a/tests/streaming/assistants/assistant.test.ts +++ b/tests/streaming/assistants/assistant.test.ts @@ -1,4 +1,4 @@ -import OpenAI from 'openai'; +import OpenAI, { OpenAIError } from 'openai'; import { ReadableStreamFrom } from 'openai/internal/shims'; import { AssistantStream } from 'openai/lib/AssistantStream'; import { Stream } from 'openai/streaming'; @@ -96,4 +96,58 @@ describe('assistant tests', () => { expect(deltas).toEqual(['E', 'ddy']); }); + + test('surfaces a mid-stream error when events are buffered before consumption', async () => { + const encoder = new TextEncoder(); + const events = [ + { + event: 'thread.message.created', + data: { + id: 'msg_1', + content: [], + }, + }, + { + event: 'thread.message.delta', + data: { + id: 'msg_1', + delta: { + content: [{ index: 0, type: 'text', text: { value: 'hi', annotations: [] } }], + }, + }, + }, + ]; + // Yield valid events, then throw to error the stream after they have been + // delivered (mimics a connection drop mid-run). + async function* eventsWithFailure() { + for (const event of events) { + yield encoder.encode(JSON.stringify(event) + '\n'); + } + throw new Error('assistant boom'); + } + + const input = ReadableStreamFrom(eventsWithFailure()); + const assistantStream = AssistantStream.fromReadableStream(input); + // Grab the iterator (registering its listeners) but do not consume yet, so + // the valid events and the error land while no reader is waiting: they + // buffer in the iterator's internal queue instead of rejecting a pending + // reader. + const iterator = assistantStream[Symbol.asyncIterator](); + // Wait for the stream's terminal signal so the events and the error have + // definitely been emitted before we start reading. + const failure: unknown = await assistantStream.done().catch((error: unknown) => error); + + await expect(iterator.next()).resolves.toMatchObject({ + done: false, + value: { event: 'thread.message.created' }, + }); + await expect(iterator.next()).resolves.toMatchObject({ + done: false, + value: { event: 'thread.message.delta' }, + }); + expect(failure).toBeInstanceOf(OpenAIError); + expect((failure as OpenAIError).message).toBe('assistant boom'); + await expect(iterator.next()).rejects.toBe(failure); + await expect(iterator.next()).resolves.toEqual({ value: undefined, done: true }); + }); });