diff --git a/src/core/streaming.ts b/src/core/streaming.ts index 3ea5f21524..d05f8faac9 100644 --- a/src/core/streaming.ts +++ b/src/core/streaming.ts @@ -114,16 +114,44 @@ export class Stream implements AsyncIterable { async function* iterLines(): AsyncGenerator { const lineDecoder = new LineDecoder(); + const reader = readableStream.getReader(); + let closed = false; + let cancelPromise: Promise | undefined; + const cancel = () => { + cancelPromise ??= reader.cancel(); + cancelPromise.catch(() => {}); + }; - const iter = ReadableStreamToAsyncIterable(readableStream); - for await (const chunk of iter) { - for (const line of lineDecoder.decode(chunk)) { - yield line; + controller.signal.addEventListener('abort', cancel, { once: true }); + try { + if (controller.signal.aborted) { + cancel(); + return; } - } - for (const line of lineDecoder.flush()) { - yield line; + while (true) { + const { value: chunk, done } = await reader.read(); + if (done) { + closed = true; + break; + } + if (controller.signal.aborted) return; + + for (const line of lineDecoder.decode(chunk)) { + if (controller.signal.aborted) return; + yield line; + } + } + + if (controller.signal.aborted) return; + for (const line of lineDecoder.flush()) { + if (controller.signal.aborted) return; + yield line; + } + } finally { + controller.signal.removeEventListener('abort', cancel); + if (!closed) cancel(); + reader.releaseLock(); } } @@ -141,7 +169,7 @@ export class Stream implements AsyncIterable { done = true; } catch (e) { // If the user calls `stream.controller.abort()`, we should exit without throwing. - if (isAbortError(e)) return; + if (controller.signal.aborted || isAbortError(e)) return; throw e; } finally { // If the user `break`s, abort the ongoing request. diff --git a/src/lib/responses/ResponseStream.ts b/src/lib/responses/ResponseStream.ts index 1761f4bbfb..f7246076ac 100644 --- a/src/lib/responses/ResponseStream.ts +++ b/src/lib/responses/ResponseStream.ts @@ -7,6 +7,7 @@ import { type ResponseStreamEvent, } from '../../resources/responses/responses'; import { RequestOptions } from '../../internal/request-options'; +import { type ReadableStream } from '../../internal/shim-types'; import { APIUserAbortError, OpenAIError } from '../../error'; import OpenAI from '../../index'; import { type BaseEvents, EventStream } from '../EventStream'; @@ -90,6 +91,12 @@ export class ResponseStream return runner; } + static fromReadableStream(stream: ReadableStream): ResponseStream { + const runner = new ResponseStream(null); + runner._run(() => runner._fromReadableStream(stream)); + return runner; + } + #beginRequest() { if (this.ended) return; this.#currentResponseSnapshot = undefined; @@ -198,6 +205,23 @@ export class ResponseStream return this.#endRequest(); } + protected async _fromReadableStream( + readableStream: ReadableStream, + options?: RequestOptions, + ): Promise> { + this._listenForAbort(options?.signal); + this.#beginRequest(); + this._connected(); + const stream = Stream.fromReadableStream(readableStream, this.controller); + for await (const event of stream) { + this.#addEvent(event, null); + } + if (stream.controller.signal?.aborted) { + throw new APIUserAbortError(); + } + return this.#endRequest(); + } + [Symbol.asyncIterator](this: ResponseStream): AsyncIterator { const pushQueue: ResponseStreamEvent[] = []; const readQueue: { diff --git a/tests/lib/ResponseStream.test.ts b/tests/lib/ResponseStream.test.ts index d23449b34c..b54b162d18 100644 --- a/tests/lib/ResponseStream.test.ts +++ b/tests/lib/ResponseStream.test.ts @@ -1,8 +1,136 @@ +import { APIUserAbortError } from 'openai'; +import { ReadableStreamFrom } from 'openai/internal/shims'; +import { ResponseStream } from 'openai/lib/responses/ResponseStream'; +import type { Response, ResponseStreamEvent } from 'openai/resources/responses/responses'; import { makeStreamSnapshotRequest } from '../utils/mock-snapshots'; jest.setTimeout(1000 * 30); describe('.stream()', () => { + it('creates a response stream from a readable stream', async () => { + const events: ResponseStreamEvent[] = [ + { + type: 'response.created', + sequence_number: 0, + response: makeResponse(), + }, + { + type: 'response.output_item.added', + sequence_number: 1, + output_index: 0, + item: { + id: 'msg_123', + type: 'message', + role: 'assistant', + status: 'in_progress', + content: [], + }, + }, + { + type: 'response.content_part.added', + sequence_number: 2, + item_id: 'msg_123', + output_index: 0, + content_index: 0, + part: { + type: 'output_text', + annotations: [], + text: '', + }, + }, + { + type: 'response.output_text.delta', + sequence_number: 3, + item_id: 'msg_123', + output_index: 0, + content_index: 0, + delta: 'Hello world', + logprobs: [], + }, + { + type: 'response.completed', + sequence_number: 4, + response: makeResponse({ + status: 'completed', + output_text: 'Hello world', + output: [ + { + id: 'msg_123', + type: 'message', + role: 'assistant', + status: 'completed', + content: [{ type: 'output_text', annotations: [], text: 'Hello world' }], + }, + ], + }), + }, + ]; + const snapshots: string[] = []; + const stream = ResponseStream.fromReadableStream(readableStreamFromEvents(events)).on( + 'response.output_text.delta', + (event) => snapshots.push(event.snapshot), + ); + const emittedEvents: ResponseStreamEvent[] = []; + + for await (const event of stream) { + emittedEvents.push(event); + } + + const final = await stream.finalResponse(); + + expect(emittedEvents).toEqual(events); + expect(snapshots).toEqual(['Hello world']); + expect(final.output_text).toBe('Hello world'); + expect(final.output[0]).toMatchObject({ + type: 'message', + content: [{ type: 'output_text', text: 'Hello world' }], + }); + }); + + it('cancels a stalled readable stream when aborted', async () => { + const encoder = new TextEncoder(); + const created = { + type: 'response.created', + sequence_number: 0, + response: makeResponse(), + } satisfies ResponseStreamEvent; + + let resolvePullStarted!: () => void; + const pullStarted = new Promise((resolve) => { + resolvePullStarted = resolve; + }); + let resolveCancelled!: () => void; + const cancelled = new Promise((resolve) => { + resolveCancelled = resolve; + }); + const cancel = jest.fn(() => resolveCancelled()); + let pulls = 0; + const readable = new ReadableStream({ + pull(controller) { + if (pulls++ === 0) { + controller.enqueue(encoder.encode(`${JSON.stringify(created)}\n`)); + return; + } + + resolvePullStarted(); + return new Promise(() => {}); + }, + cancel, + }); + const stream = ResponseStream.fromReadableStream(readable); + + await stream.emitted('response.created'); + await pullStarted; + + const done = expect(stream.done()).rejects.toThrowError(APIUserAbortError); + stream.abort(); + + await cancelled; + await done; + expect(cancel).toHaveBeenCalledTimes(1); + expect(stream.aborted).toBe(true); + }, 5000); + it('standard text works', async () => { const deltas: string[] = []; @@ -59,3 +187,39 @@ describe('.stream()', () => { expect(final.output_text).toBe('The answer is 42'); }); }); + +function readableStreamFromEvents(events: ResponseStreamEvent[]) { + const encoder = new TextEncoder(); + return ReadableStreamFrom(events.map((event) => encoder.encode(JSON.stringify(event) + '\n'))); +} + +function makeResponse(overrides: Partial = {}): Response { + return { + id: 'resp_123', + object: 'response', + created_at: 1, + error: null, + incomplete_details: null, + instructions: null, + metadata: null, + model: 'gpt-5', + output: [], + output_text: '', + parallel_tool_calls: false, + status: 'in_progress', + temperature: null, + tool_choice: 'auto', + tools: [], + top_p: null, + max_output_tokens: null, + previous_response_id: null, + reasoning: { effort: null, summary: null }, + service_tier: null, + store: true, + text: { format: { type: 'text' }, verbosity: null }, + truncation: 'disabled', + usage: null, + user: null, + ...overrides, + } as Response; +}