Skip to content

refactor: share existing WebSocket send and open handling #2947

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
markstuart-oai merged 1 commit into from
Oct 15, 2026
Merged

refactor: share existing WebSocket send and open handling #2947

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 116 additions & 0 deletions src/internal/ws-connection.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
import { OpenAIError } from '../core/error';
import { ReadyState } from './ws-adapter';
import type { WebSocketLike } from './ws-adapter';
import { flattenRawData } from './ws';
import type { RawWebSocketData, SendQueue } from './ws';

// Read the public socket on each access; callers may replace it through an accessor.
interface SocketOwner {
socket: WebSocketLike;
}
type OnError = (message: string, cause?: unknown) => void;

/** Sends or queues an application event using the existing connection state. */
export function sendWebSocketEvent<T>(
owner: SocketOwner,
event: T,
queue: SendQueue<T>,
isReconnecting: () => boolean,
onError: OnError,
): void {
if (!owner.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

if (isReconnecting() || owner.socket.readyState === ReadyState.CONNECTING) {
if (!queue.enqueue(event)) {
onError('send queue is full, message discarded');
}
return;
}
if (owner.socket.readyState !== ReadyState.OPEN) {
onError('cannot send on a closed WebSocket');
return;
}
try {
owner.socket.send(JSON.stringify(event));
} catch (error) {
onError('could not send data', error);
}
}

/** Sends or snapshots a raw frame using the same admission rules as application events. */
export function sendRawWebSocketData<T>(
owner: SocketOwner,
data: RawWebSocketData,
queue: SendQueue<T>,
isReconnecting: () => boolean,
onError: OnError,
): void {
if (!owner.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

if (isReconnecting() || owner.socket.readyState === ReadyState.CONNECTING) {
if (!queue.enqueueRaw(data)) {
onError('send queue is full, message discarded');
}
return;
}
if (owner.socket.readyState !== ReadyState.OPEN) {
onError('cannot send on a closed WebSocket');
return;
}
try {
owner.socket.send(flattenRawData(data));
} catch (error) {
onError('could not send data', error);
}
}

/** Flushes queued frames with the caller's existing failed-send policy. */
export function flushWebSocketSendQueue<T>(
owner: SocketOwner,
queue: SendQueue<T>,
onError: OnError,
options?: { requeueFailed?: boolean },
): void {
if (!owner.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

try {
queue.flush((data) => owner.socket.send(flattenRawData(data)), options);
} catch (error) {
onError('could not send queued data', error);
}
}

/** Resolves when the socket opens, or rejects if it errors or closes first. */
export function waitForWebSocketOpen(socket: WebSocketLike): Promise<void> {
// oxlint-disable-next-line promise/avoid-new -- Convert the existing socket event callbacks into the connection attempt's promise.
return new Promise<void>((resolve, reject) => {
/* oxlint-disable eslint/no-use-before-define -- All mutually dependent callbacks are initialized before socket listeners are registered. */
const cleanup = () => {
socket.off('open', onOpen);
socket.off('error', onError);
socket.off('close', onFail);
};
/* oxlint-enable eslint/no-use-before-define */
const onOpen = () => {
cleanup();
resolve();
};
const onError = (error: Error) => {
cleanup();
reject(error);
};
const onFail = () => {
cleanup();
reject(new Error('socket closed before open'));
};
socket.once('open', onOpen);
socket.once('error', onError);
socket.once('close', onFail);
});
}
99 changes: 22 additions & 77 deletions src/resources/beta/responses/ws-base.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
// File generated from our OpenAPI spec by Castiron. See CONTRIBUTING.md for details.

import { ResponsesEmitter, ResponsesStreamMessage, WebSocketError, buildURL } from './internal-base';
import {
sendWebSocketEvent,
sendRawWebSocketData,
flushWebSocketSendQueue,
waitForWebSocketOpen,
} from '../../../internal/ws-connection';
import { InternalEventEmitter } from '../../../core/EventEmitter';
import { sleep } from '../../../internal/utils/sleep';
import { type WebSocketLike, ReadyState } from '../../../internal/ws-adapter';
Expand All @@ -10,7 +16,6 @@ import {
recordWebSocketError,
rawByteLength,
type WebSocketStreamOptions,
flattenRawData,
isRecoverableClose,
type RawWebSocketData,
type ReconnectingEvent,
Expand Down Expand Up @@ -128,47 +133,23 @@ export abstract class ResponsesWSBase<TSocket extends WebSocketLike> extends Res
private _cancelRefresh: (() => void) | undefined;

send(event: ResponsesAPI.BetaResponsesClientEvent) {
if (!this.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

if (this._isReconnecting || this.socket.readyState === ReadyState.CONNECTING) {
if (!this._sendQueue.enqueue(event)) {
this._onError(null, 'send queue is full, message discarded', undefined);
}
return;
}
if (this.socket.readyState !== ReadyState.OPEN) {
this._onError(null, 'cannot send on a closed WebSocket', undefined);
return;
}
try {
this.socket.send(JSON.stringify(event));
} catch (err) {
this._onError(null, 'could not send data', err);
}
sendWebSocketEvent(
this,
event,
this._sendQueue,
() => this._isReconnecting,
(message, cause) => this._onError(null, message, cause),
);
}

sendRaw(data: RawWebSocketData) {
if (!this.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

if (this._isReconnecting || this.socket.readyState === ReadyState.CONNECTING) {
if (!this._sendQueue.enqueueRaw(data)) {
this._onError(null, 'send queue is full, message discarded', undefined);
}
return;
}
if (this.socket.readyState !== ReadyState.OPEN) {
this._onError(null, 'cannot send on a closed WebSocket', undefined);
return;
}
try {
this.socket.send(flattenRawData(data));
} catch (err) {
this._onError(null, 'could not send data', err);
}
sendRawWebSocketData(
this,
data,
this._sendQueue,
() => this._isReconnecting,
(message, cause) => this._onError(null, message, cause),
);
}

close(props?: { code: number; reason: string }) {
Expand Down Expand Up @@ -661,7 +642,7 @@ export abstract class ResponsesWSBase<TSocket extends WebSocketLike> extends Res
this.socket.once('close', resolve);
});

await this._awaitOpen(this.socket);
await waitForWebSocketOpen(this.socket);

this._internalEvents._emit('socketSwap', oldSocket, this.socket);
this._isReconnecting = false;
Expand Down Expand Up @@ -698,44 +679,8 @@ export abstract class ResponsesWSBase<TSocket extends WebSocketLike> extends Res
this._emitPermanentClose(closeCode, `reconnect failed after ${maxRetries} attempts`);
}

/**
* Resolves once the socket is open, rejects if it errors or closes first
*/
private _awaitOpen(socket: TSocket): Promise<void> {
return new Promise<void>((resolve, reject) => {
const cleanup = () => {
socket.off('open', onOpen);
socket.off('error', onError);
socket.off('close', onFail);
};
const onOpen = () => {
cleanup();
resolve();
};
const onError = (err: Error) => {
cleanup();
reject(err);
};
const onFail = () => {
cleanup();
reject(new Error('socket closed before open'));
};
socket.once('open', onOpen);
socket.once('error', onError);
socket.once('close', onFail);
});
}

private _flushSendQueue(): void {
if (!this.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

try {
this._sendQueue.flush((data) => this.socket.send(flattenRawData(data)));
} catch (err) {
this._onError(null, 'could not send queued data', err);
}
flushWebSocketSendQueue(this, this._sendQueue, (message, cause) => this._onError(null, message, cause));
}

/**
Expand Down
101 changes: 24 additions & 77 deletions src/resources/live/forks/ws-base.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,19 @@
// File generated from our OpenAPI spec by Castiron. See CONTRIBUTING.md for details.

import { ForksEmitter, ForksStreamMessage, WebSocketError, buildURL } from './internal-base';
import {
sendWebSocketEvent,
sendRawWebSocketData,
flushWebSocketSendQueue,
waitForWebSocketOpen,
} from '../../../internal/ws-connection';
import { InternalEventEmitter } from '../../../core/EventEmitter';
import { sleep } from '../../../internal/utils/sleep';
import { type WebSocketLike, ReadyState } from '../../../internal/ws-adapter';
import {
SendQueue,
getMaxBufferedEvents,
type WebSocketStreamOptions,
flattenRawData,
isRecoverableClose,
type RawWebSocketData,
type ReconnectingEvent,
Expand Down Expand Up @@ -107,47 +112,23 @@ export abstract class ForksWSBase<TSocket extends WebSocketLike> extends ForksEm
protected abstract _createSocket(url: URL, authHeaders: Record<string, string>): TSocket;

send(event: ForksAPI.ForkClientEvent) {
if (!this.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

if (this._isReconnecting || this.socket.readyState === ReadyState.CONNECTING) {
if (!this._sendQueue.enqueue(event)) {
this._onError(null, 'send queue is full, message discarded', undefined);
}
return;
}
if (this.socket.readyState !== ReadyState.OPEN) {
this._onError(null, 'cannot send on a closed WebSocket', undefined);
return;
}
try {
this.socket.send(JSON.stringify(event));
} catch (err) {
this._onError(null, 'could not send data', err);
}
sendWebSocketEvent(
this,
event,
this._sendQueue,
() => this._isReconnecting,
(message, cause) => this._onError(null, message, cause),
);
}

sendRaw(data: RawWebSocketData) {
if (!this.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

if (this._isReconnecting || this.socket.readyState === ReadyState.CONNECTING) {
if (!this._sendQueue.enqueueRaw(data)) {
this._onError(null, 'send queue is full, message discarded', undefined);
}
return;
}
if (this.socket.readyState !== ReadyState.OPEN) {
this._onError(null, 'cannot send on a closed WebSocket', undefined);
return;
}
try {
this.socket.send(flattenRawData(data));
} catch (err) {
this._onError(null, 'could not send data', err);
}
sendRawWebSocketData(
this,
data,
this._sendQueue,
() => this._isReconnecting,
(message, cause) => this._onError(null, message, cause),
);
}

close(props?: { code: number; reason: string }) {
Expand Down Expand Up @@ -609,7 +590,7 @@ export abstract class ForksWSBase<TSocket extends WebSocketLike> extends ForksEm
this.socket.once('close', resolve);
});

await this._awaitOpen(this.socket);
await waitForWebSocketOpen(this.socket);

this._internalEvents._emit('socketSwap', oldSocket, this.socket);
this._isReconnecting = false;
Expand Down Expand Up @@ -637,44 +618,10 @@ export abstract class ForksWSBase<TSocket extends WebSocketLike> extends ForksEm
this._emitPermanentClose(closeCode, `reconnect failed after ${maxRetries} attempts`);
}

/**
* Resolves once the socket is open, rejects if it errors or closes first
*/
private _awaitOpen(socket: TSocket): Promise<void> {
return new Promise<void>((resolve, reject) => {
const cleanup = () => {
socket.off('open', onOpen);
socket.off('error', onError);
socket.off('close', onFail);
};
const onOpen = () => {
cleanup();
resolve();
};
const onError = (err: Error) => {
cleanup();
reject(err);
};
const onFail = () => {
cleanup();
reject(new Error('socket closed before open'));
};
socket.once('open', onOpen);
socket.once('error', onError);
socket.once('close', onFail);
});
}

private _flushSendQueue(): void {
if (!this.socket) {
throw new OpenAIError('Internal error: failed to initialize socket. Please report this issue.');
}

try {
this._sendQueue.flush((data) => this.socket.send(flattenRawData(data)), { requeueFailed: false });
} catch (err) {
this._onError(null, 'could not send queued data', err);
}
flushWebSocketSendQueue(this, this._sendQueue, (message, cause) => this._onError(null, message, cause), {
requeueFailed: false,
});
}

/**
Expand Down
Loading
Loading