diff --git a/packages/exchange/src/broker/alpaca/AlpacaWebSocket.test.ts b/packages/exchange/src/broker/alpaca/AlpacaWebSocket.test.ts new file mode 100644 index 000000000..5ad65b11c --- /dev/null +++ b/packages/exchange/src/broker/alpaca/AlpacaWebSocket.test.ts @@ -0,0 +1,207 @@ +import type {EventEmitter} from 'node:events'; +import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; +import type {MinuteBarMessage} from './api/schema/StreamSchema.js'; + +interface FakeStream extends EventEmitter { + close: ReturnType; + subscribe: ReturnType; + unsubscribe: ReturnType; +} + +const fakeStreams = vi.hoisted(() => ({instances: [] as unknown[]})); + +vi.mock('./api/AlpacaStream.js', async () => { + const {EventEmitter} = await import('node:events'); + + class FakeAlpacaStream extends EventEmitter { + close = vi.fn(() => { + queueMicrotask(() => this.emit('close', this)); + }); + subscribe = vi.fn(); + unsubscribe = vi.fn(); + + constructor() { + super(); + fakeStreams.instances.push(this); + // The real stream authenticates asynchronously after the socket opens + queueMicrotask(() => this.emit('authenticated', this)); + } + } + + return {AlpacaStream: FakeAlpacaStream}; +}); + +const {alpacaWebSocket} = await import('./AlpacaWebSocket.js'); + +function getStream(index: number): FakeStream { + const stream = fakeStreams.instances[index]; + if (!stream) { + throw new Error(`No fake stream at index ${index}.`); + } + return stream as FakeStream; +} + +function createBar(symbol: string): MinuteBarMessage { + return {c: 100, h: 101, l: 99, n: 10, o: 100, S: symbol, T: 'b', t: '2025-01-15T14:30:00Z', v: 1_000, vw: 100}; +} + +/** Each test uses fresh credentials so the singleton-per-`apiKey:source` map never reuses a prior connection. */ +function createCredentials() { + return {apiKey: crypto.randomUUID(), apiSecret: 'test-secret', usePaperTrading: true}; +} + +async function flushMicrotasks(): Promise { + await new Promise(resolve => { + setTimeout(resolve, 0); + }); +} + +describe('alpacaWebSocket', {concurrent: false}, () => { + beforeEach(() => { + fakeStreams.instances.length = 0; + vi.spyOn(console, 'error').mockImplementation(() => {}); + vi.spyOn(console, 'log').mockImplementation(() => {}); + }); + + afterEach(() => { + alpacaWebSocket.removeAllListeners(); + vi.restoreAllMocks(); + }); + + describe('connect', () => { + it('reuses the existing connection for the same credentials and source', async () => { + const credentials = createCredentials(); + + const first = await alpacaWebSocket.connect(credentials, 'v2/iex'); + const second = await alpacaWebSocket.connect(credentials, 'v2/iex'); + + expect(second.connectionId).toBe(first.connectionId); + expect(fakeStreams.instances).toHaveLength(1); + }); + }); + + describe('reconnect', () => { + it('reconnects and resubscribes prior bar topics after an unexpected close', async () => { + const connection = await alpacaWebSocket.connect(createCredentials(), 'v2/iex'); + alpacaWebSocket.subscribeToBars(connection.connectionId, 'AAPL', vi.fn()); + alpacaWebSocket.subscribeToBars(connection.connectionId, 'TSLA', vi.fn()); + + const onReconnecting = vi.fn(); + const onResubscribed = vi.fn(); + alpacaWebSocket.on('reconnecting', onReconnecting); + alpacaWebSocket.on('resubscribed', onResubscribed); + + const firstStream = getStream(0); + firstStream.emit('close', firstStream); + + await vi.waitFor(() => { + expect(onResubscribed).toHaveBeenCalledTimes(1); + }); + + expect(onReconnecting).toHaveBeenCalledWith({connectionId: connection.connectionId}); + expect(onResubscribed).toHaveBeenCalledWith({connectionId: connection.connectionId, symbols: ['AAPL', 'TSLA']}); + + const secondStream = getStream(1); + expect(secondStream).not.toBe(firstStream); + expect(secondStream.subscribe).toHaveBeenCalledWith('bars', ['AAPL', 'TSLA']); + }); + + it('keeps emitting bars to existing subscribers after a reconnect', async () => { + const connection = await alpacaWebSocket.connect(createCredentials(), 'v2/iex'); + const onBar = vi.fn(); + alpacaWebSocket.subscribeToBars(connection.connectionId, 'AAPL', onBar); + + const onResubscribed = vi.fn(); + alpacaWebSocket.on('resubscribed', onResubscribed); + const firstStream = getStream(0); + firstStream.emit('close', firstStream); + await vi.waitFor(() => { + expect(onResubscribed).toHaveBeenCalledTimes(1); + }); + + getStream(1).emit('message', createBar('AAPL')); + + expect(onBar).toHaveBeenCalledTimes(1); + }); + + it('never calls process.exit when the socket drops', async () => { + const exitSpy = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('process.exit must not be called'); + }); + + const onResubscribed = vi.fn(); + alpacaWebSocket.on('resubscribed', onResubscribed); + await alpacaWebSocket.connect(createCredentials(), 'v2/iex'); + + const firstStream = getStream(0); + firstStream.emit('close', firstStream); + await vi.waitFor(() => { + expect(onResubscribed).toHaveBeenCalledTimes(1); + }); + + expect(exitSpy).not.toHaveBeenCalled(); + }); + }); + + describe('disconnect', () => { + it('does not reconnect after an intentional disconnect', async () => { + const connection = await alpacaWebSocket.connect(createCredentials(), 'v2/iex'); + const onReconnecting = vi.fn(); + alpacaWebSocket.on('reconnecting', onReconnecting); + + alpacaWebSocket.disconnect(connection.connectionId); + await flushMicrotasks(); + await flushMicrotasks(); + + expect(getStream(0).close).toHaveBeenCalledTimes(1); + expect(onReconnecting).not.toHaveBeenCalled(); + expect(fakeStreams.instances).toHaveLength(1); + }); + }); + + describe('subscribeToBars', () => { + it('emits exactly one candle per message after a subscribe/unsubscribe/subscribe cycle', async () => { + const connection = await alpacaWebSocket.connect(createCredentials(), 'v2/iex'); + const stream = getStream(0); + + const firstCallback = vi.fn(); + alpacaWebSocket.subscribeToBars(connection.connectionId, 'AAPL', firstCallback); + stream.emit('message', createBar('AAPL')); + expect(firstCallback).toHaveBeenCalledTimes(1); + + alpacaWebSocket.unsubscribeFromBars(connection.connectionId, 'AAPL'); + const secondCallback = vi.fn(); + alpacaWebSocket.subscribeToBars(connection.connectionId, 'AAPL', secondCallback); + stream.emit('message', createBar('AAPL')); + + expect(secondCallback).toHaveBeenCalledTimes(1); + expect(firstCallback).toHaveBeenCalledTimes(1); + }); + + it('only dispatches bars to callbacks of the matching symbol', async () => { + const connection = await alpacaWebSocket.connect(createCredentials(), 'v2/iex'); + const onAppleBar = vi.fn(); + alpacaWebSocket.subscribeToBars(connection.connectionId, 'AAPL', onAppleBar); + + getStream(0).emit('message', createBar('TSLA')); + + expect(onAppleBar).not.toHaveBeenCalled(); + }); + + it('survives subscription updates while the underlying socket is closed', async () => { + const connection = await alpacaWebSocket.connect(createCredentials(), 'v2/iex'); + const stream = getStream(0); + const closedSocketError = new Error('WebSocket is not open'); + stream.subscribe.mockImplementation(() => { + throw closedSocketError; + }); + stream.unsubscribe.mockImplementation(() => { + throw closedSocketError; + }); + + const callback = vi.fn(); + expect(() => alpacaWebSocket.subscribeToBars(connection.connectionId, 'AAPL', callback)).not.toThrow(); + expect(() => alpacaWebSocket.unsubscribeFromBars(connection.connectionId, 'AAPL')).not.toThrow(); + }); + }); +}); diff --git a/packages/exchange/src/broker/alpaca/AlpacaWebSocket.ts b/packages/exchange/src/broker/alpaca/AlpacaWebSocket.ts index 854158de2..fade6c1ec 100644 --- a/packages/exchange/src/broker/alpaca/AlpacaWebSocket.ts +++ b/packages/exchange/src/broker/alpaca/AlpacaWebSocket.ts @@ -1,3 +1,4 @@ +import {EventEmitter} from 'node:events'; import {ms} from 'ms'; import type {RetryConfig} from 'ts-retry-promise'; import {retry} from 'ts-retry-promise'; @@ -17,11 +18,22 @@ export interface AlpacaConnection { * Alpaca only allows 1 WebSocket connection per API key. This class manages * WebSocket connections as singletons to avoid the following error: * {"T":"error","code":406,"msg":"connection limit exceeded"} + * + * When a connection drops unexpectedly it is re-established in-process (same + * `connectionId`, same bar subscriptions) instead of killing the process — + * other trading sessions and bots share this process and must keep running. + * Hosts can observe the lifecycle via the emitted `reconnecting`, + * `resubscribed`, and `reconnect_failed` events. */ -class AlpacaWebSocket { +class AlpacaWebSocket extends EventEmitter { readonly #connections: Map = new Map(); - readonly #symbols: Map> = new Map(); + /** Bar callbacks per connection, keyed by symbol, so reconnects can restore them. */ + readonly #subscriptions: Map void>>> = new Map(); readonly #credentialToConnectionId: Map = new Map(); + /** Connection parameters retained for re-authentication after a transport drop. */ + readonly #connectionParams: Map = new Map(); + /** Connections closed via `disconnect()` — their `close` event must not trigger a reconnect. */ + readonly #intentionalCloses: Set = new Set(); /** * @see https://docs.alpaca.markets/docs/streaming-market-data#connection @@ -52,20 +64,8 @@ class AlpacaWebSocket { timeout: 'INFINITELY', } as const; - async #establishConnection(credentials: AlpacaStreamCredentials, source: string): Promise { - // Check if we already have a connection for these credentials + source - const singletonKey = `${credentials.apiKey}:${source}`; - const existingConnectionId = this.#credentialToConnectionId.get(singletonKey); - if (existingConnectionId) { - const existing = this.#connections.get(existingConnectionId); - if (existing) { - return existing; - } - } - - const connectionId = crypto.randomUUID(); - - return new Promise((resolve, reject) => { + #openStream(credentials: AlpacaStreamCredentials, source: string, connectionId: string): Promise { + return new Promise((resolve, reject) => { const stream = new AlpacaStream(credentials, source); stream.on('error', (error: unknown) => { @@ -77,30 +77,99 @@ class AlpacaWebSocket { console.log(`WebSocket streaming is subscribed with ID "${connectionId}":`, JSON.stringify(message)); }); - stream.on('authenticated', () => { + stream.once('authenticated', () => { console.log(`WebSocket streaming is authenticated with ID "${connectionId}".`); - const connection = {connectionId, stream}; - this.#connections.set(connectionId, connection); - this.#credentialToConnectionId.set(singletonKey, connectionId); - - /** - * Every close on this socket is currently a failure (transport drop). - * In this case, we do a hard-exit so the orchestrator restarts the stream. - */ - stream.once('close', () => { - console.error( - `Market-data WebSocket closed for "${connectionId}". Exiting so the orchestrator restarts the stream.` - ); - process.exit(1); - }); - - resolve(connection); + resolve(stream); }); }); } + /** + * Installs the single message handler plus the close watchdog on a freshly + * authenticated stream. Exactly one message handler exists per stream, so + * repeated subscribe/unsubscribe cycles cannot stack duplicate listeners. + */ + #wireStream(connectionId: string, stream: AlpacaStream): void { + stream.on('message', (message: StreamMessage) => { + this.#dispatchBar(connectionId, message); + }); + + stream.once('close', () => { + if (this.#intentionalCloses.delete(connectionId)) { + return; + } + console.error(`Market-data WebSocket closed unexpectedly for "${connectionId}". Reconnecting in-process.`); + this.emit('reconnecting', {connectionId}); + void this.#reconnect(connectionId); + }); + } + + #dispatchBar(connectionId: string, message: StreamMessage): void { + if (message.T !== 'b') { + return; + } + const callbacks = this.#subscriptions.get(connectionId)?.get(message.S); + if (!callbacks) { + return; + } + for (const cb of callbacks) { + cb(message); + } + } + + async #reconnect(connectionId: string): Promise { + const params = this.#connectionParams.get(connectionId); + if (!params) { + return; + } + + try { + const stream = await retry( + () => this.#openStream(params.credentials, params.source, connectionId), + this.#retryConfig + ); + + const connection = this.#connections.get(connectionId); + if (!connection) { + // Disconnected while the reconnect was in flight + stream.close(); + return; + } + + connection.stream = stream; + this.#wireStream(connectionId, stream); + + const symbols = Array.from(this.#subscriptions.get(connectionId)?.keys() ?? []); + if (symbols.length > 0) { + stream.subscribe('bars', symbols); + } + this.emit('resubscribed', {connectionId, symbols}); + } catch (error) { + console.error(`Reconnect failed permanently for "${connectionId}".`, error); + this.emit('reconnect_failed', {connectionId, error}); + } + } + async connect(credentials: AlpacaStreamCredentials, source: string): Promise { - return retry(() => this.#establishConnection(credentials, source), this.#retryConfig); + const singletonKey = `${credentials.apiKey}:${source}`; + const existingConnectionId = this.#credentialToConnectionId.get(singletonKey); + if (existingConnectionId) { + const existing = this.#connections.get(existingConnectionId); + if (existing) { + return existing; + } + } + + const connectionId = crypto.randomUUID(); + const stream = await retry(() => this.#openStream(credentials, source, connectionId), this.#retryConfig); + + const connection: AlpacaConnection = {connectionId, stream}; + this.#connections.set(connectionId, connection); + this.#credentialToConnectionId.set(singletonKey, connectionId); + this.#connectionParams.set(connectionId, {credentials, source}); + this.#wireStream(connectionId, stream); + + return connection; } /** @@ -114,23 +183,24 @@ class AlpacaWebSocket { return; } - // Track symbols per connection - let symbols = this.#symbols.get(connectionId); - if (!symbols) { - symbols = new Set(); - this.#symbols.set(connectionId, symbols); + let subscriptions = this.#subscriptions.get(connectionId); + if (!subscriptions) { + subscriptions = new Map(); + this.#subscriptions.set(connectionId, subscriptions); } - symbols.add(symbol); - connection.stream.on('message', (message: StreamMessage) => { - if (message.T === 'b' && message.S === symbol) { - cb(message); - } - }); + let callbacks = subscriptions.get(symbol); + if (!callbacks) { + callbacks = new Set(); + subscriptions.set(symbol, callbacks); + } + callbacks.add(cb); - const allSymbols = Array.from(symbols); - connection.stream.unsubscribe('bars', allSymbols); - connection.stream.subscribe('bars', allSymbols); + const allSymbols = Array.from(subscriptions.keys()); + this.#sendSubscriptionUpdate(connection, stream => { + stream.unsubscribe('bars', allSymbols); + stream.subscribe('bars', allSymbols); + }); } unsubscribeFromBars(connectionId: string, symbol: string) { @@ -139,10 +209,44 @@ class AlpacaWebSocket { return; } - const symbols = this.#symbols.get(connectionId); - symbols?.delete(symbol); + this.#subscriptions.get(connectionId)?.delete(symbol); + + this.#sendSubscriptionUpdate(connection, stream => { + stream.unsubscribe('bars', [symbol]); + }); + } + + /** + * Subscription updates go over the wire, and during a reconnect window the current + * stream may already be closed — `WebSocket#send` then throws, which would crash the + * host process. Swallowing the failure is safe: the reconnect path resubscribes every + * tracked symbol from `#subscriptions` once the fresh stream is authenticated. + */ + #sendSubscriptionUpdate(connection: AlpacaConnection, update: (stream: AlpacaStream) => void) { + try { + update(connection.stream); + } catch (error) { + console.warn(`Subscription update on "${connection.connectionId}" deferred to reconnect:`, error); + } + } - connection.stream.unsubscribe('bars', [symbol]); + disconnect(connectionId: string) { + const connection = this.#connections.get(connectionId); + if (!connection) { + return; + } + + this.#intentionalCloses.add(connectionId); + connection.stream.close(); + + this.#connections.delete(connectionId); + this.#subscriptions.delete(connectionId); + this.#connectionParams.delete(connectionId); + for (const [singletonKey, id] of this.#credentialToConnectionId) { + if (id === connectionId) { + this.#credentialToConnectionId.delete(singletonKey); + } + } } }