@luxexchange/websocket 1.0.0

package/src/client/__tests__/testUtils.ts

@@ -0,0 +1,113 @@
+import { MockWebSocket } from '@luxexchange/websocket/src/client/__tests__/MockWebSocket'
+import { createWebSocketClient } from '@luxexchange/websocket/src/client/createWebSocketClient'
+import { createZustandConnectionStore } from '@luxexchange/websocket/src/store/createZustandConnectionStore'
+import type { ConnectionStore } from '@luxexchange/websocket/src/store/types'
+import type { CreateWebSocketClientOptions, WebSocketClient } from '@luxexchange/websocket/src/types'
+import { vi } from 'vitest'
+
+export interface TestParams {
+  channel: string
+  id: string
+}
+
+export interface TestMessage {
+  data: string
+  price?: number
+}
+
+interface TestHandler {
+  subscribe: ReturnType<typeof vi.fn>
+  unsubscribe: ReturnType<typeof vi.fn>
+  subscribeBatch: ReturnType<typeof vi.fn>
+  unsubscribeBatch: ReturnType<typeof vi.fn>
+  refreshSession: ReturnType<typeof vi.fn>
+}
+
+interface TestClientResult {
+  client: WebSocketClient<TestParams, TestMessage>
+  mockSocket: MockWebSocket
+  handler: TestHandler
+  connectionStore: ConnectionStore
+}
+
+/**
+ * Creates a test WebSocket client with mocked dependencies.
+ * Returns the client, mock socket, mock handler, and connection store for assertions.
+ */
+export function createTestClient(
+  overrides?: Partial<CreateWebSocketClientOptions<TestParams, TestMessage>>,
+): TestClientResult {
+  const mockSocket = new MockWebSocket()
+  const connectionStore = createZustandConnectionStore()
+
+  const handler: TestHandler = {
+    subscribe: vi.fn().mockResolvedValue(undefined),
+    unsubscribe: vi.fn().mockResolvedValue(undefined),
+    subscribeBatch: vi.fn().mockResolvedValue(undefined),
+    unsubscribeBatch: vi.fn().mockResolvedValue(undefined),
+    refreshSession: vi.fn().mockResolvedValue(undefined),
+  }
+
+  const client = createWebSocketClient<TestParams, TestMessage>({
+    config: { url: 'wss://test.example.com' },
+    connectionStore,
+    subscriptionHandler: overrides?.subscriptionHandler ?? handler,
+    parseMessage: (raw) => {
+      const msg = raw as { channel?: string; key?: string; data?: TestMessage }
+      if (msg.channel && msg.key && msg.data) {
+        return { channel: msg.channel, key: msg.key, data: msg.data }
+      }
+      return null
+    },
+    parseConnectionMessage: (raw) => {
+      const msg = raw as { type?: string; connectionId?: string }
+      if (msg.type === 'connected' && msg.connectionId) {
+        return { connectionId: msg.connectionId }
+      }
+      return null
+    },
+    createSubscriptionKey: (channel, params) => `${channel}:${params.id}`,
+    socketFactory: () => mockSocket,
+    ...overrides,
+  })
+
+  return { client, mockSocket, handler, connectionStore }
+}
+
+interface ConnectViaSubscribeOptions {
+  client: WebSocketClient<TestParams, TestMessage>
+  mockSocket: MockWebSocket
+  connectionId?: string
+}
+
+/**
+ * Triggers a connection by subscribing, then simulates socket open and connection message.
+ * Returns the unsubscribe function from the subscription that triggered the connection.
+ */
+export function connectViaSubscribe(options: ConnectViaSubscribeOptions): () => void {
+  const { client, mockSocket, connectionId = 'conn-123' } = options
+  const unsub = client.subscribe({
+    channel: '__connect',
+    params: { channel: '__connect', id: '__connect' },
+    onMessage: vi.fn(),
+  })
+  mockSocket.simulateOpen()
+  mockSocket.simulateMessage({ type: 'connected', connectionId })
+  return unsub
+}
+
+/**
+ * Flush all pending microtasks and promises.
+ */
+export async function flushMicrotasks(): Promise<void> {
+  // queueMicrotask runs before setTimeout, but we need to also flush
+  // any promises that are scheduled from within those microtasks
+  await new Promise((resolve) => setTimeout(resolve, 0))
+}
+
+/**
+ * Flush all pending promises in the microtask queue.
+ */
+export function flushPromises(): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, 0))
+}

package/src/client/createWebSocketClient.ts

@@ -0,0 +1,262 @@
+import { SubscriptionManager } from '@luxexchange/websocket/src/subscriptions/SubscriptionManager'
+import type {
+  ConnectionStatus,
+  CreateWebSocketClientOptions,
+  SocketFactoryOptions,
+  SubscriptionOptions,
+  WebSocketClient,
+  WebSocketLike,
+} from '@luxexchange/websocket/src/types'
+import { getDefaultJitteredDelay } from '@luxexchange/websocket/src/utils/backoff'
+import { WebSocket as PartySocket } from 'partysocket'
+
+/** Default socket factory using PartySocket */
+function defaultSocketFactory(url: string, options: SocketFactoryOptions): WebSocketLike {
+  return new PartySocket(url, [], options)
+}
+
+/** Default configuration for WebSocket connection behavior */
+const DEFAULT_CONFIG = {
+  // Maximum delay (ms) between reconnection attempts
+  maxReconnectionDelay: 10000,
+  // Minimum delay (ms) between reconnection attempts (before jitter is applied)
+  minReconnectionDelay: 1000,
+  // Time (ms) to wait for a connection to establish before timing out
+  connectionTimeout: 4000,
+  // Maximum number of reconnection attempts before giving up
+  maxRetries: 5,
+  // Enable debug logging for connection events
+  debug: false,
+}
+
+/**
+ * Creates a generic WebSocket client with subscription management.
+ *
+ * The client handles:
+ * - Lazy connection lifecycle (connects on first subscribe, disconnects on last unsubscribe)
+ * - Subscription management via SubscriptionManager (reference counting, microtask batching)
+ * - Message parsing and routing to appropriate subscribers
+ *
+ * Consumers provide:
+ * - connectionStore: Manages connection state (status, connectionId, errors)
+ * - subscriptionHandler: REST API for subscribe/unsubscribe calls
+ * - parseMessage: Convert raw WS messages to typed messages
+ * - parseConnectionMessage: Extract connectionId from initial message
+ * - createSubscriptionKey: Create unique keys for subscriptions
+ */
+export function createWebSocketClient<TParams, TMessage>(
+  options: CreateWebSocketClientOptions<TParams, TMessage>,
+): WebSocketClient<TParams, TMessage> {
+  const {
+    config,
+    connectionStore,
+    subscriptionHandler,
+    parseMessage,
+    parseConnectionMessage,
+    createSubscriptionKey,
+    onError,
+    onRawMessage,
+    sessionRefreshIntervalMs,
+    socketFactory = defaultSocketFactory,
+  } = options
+
+  const {
+    url,
+    maxReconnectionDelay = DEFAULT_CONFIG.maxReconnectionDelay,
+    connectionTimeout = DEFAULT_CONFIG.connectionTimeout,
+    maxRetries = DEFAULT_CONFIG.maxRetries,
+    debug = DEFAULT_CONFIG.debug,
+  } = config
+
+  // Internal state
+  let socket: WebSocketLike | null = null
+  const connectionCallbacks = new Set<(connectionId: string) => void>()
+  let wasConnected = false
+  let sessionRefreshTimer: ReturnType<typeof setInterval> | null = null
+
+  function startSessionRefreshTimer(): void {
+    stopSessionRefreshTimer()
+    if (sessionRefreshIntervalMs && subscriptionHandler.refreshSession) {
+      sessionRefreshTimer = setInterval(() => {
+        subscriptionManager.refreshSession().catch((error) => onError?.(error))
+      }, sessionRefreshIntervalMs)
+    }
+  }
+
+  function stopSessionRefreshTimer(): void {
+    if (sessionRefreshTimer !== null) {
+      clearInterval(sessionRefreshTimer)
+      sessionRefreshTimer = null
+    }
+  }
+
+  // Subscription manager handles all subscription logic
+  const subscriptionManager = new SubscriptionManager<TParams, TMessage>({
+    handler: subscriptionHandler,
+    createKey: createSubscriptionKey,
+    onError: (error) => onError?.(error),
+    onSubscriptionCountChange: (count): void => {
+      if (count > 0 && !socket) {
+        connect()
+      } else if (count === 0 && socket) {
+        disconnect()
+      }
+    },
+  })
+
+  function notifyStatusChange(status: ConnectionStatus): void {
+    connectionStore.setStatus(status)
+  }
+
+  function notifyConnectionEstablished(connectionId: string): void {
+    connectionStore.setConnectionId(connectionId)
+    for (const callback of connectionCallbacks) {
+      callback(connectionId)
+    }
+  }
+
+  function connect(): void {
+    if (socket) {
+      return
+    }
+
+    notifyStatusChange('connecting')
+
+    // Add jitter to prevent thundering herd on reconnect
+    const jitteredMinDelay = getDefaultJitteredDelay()
+
+    socket = socketFactory(url, {
+      maxReconnectionDelay,
+      minReconnectionDelay: jitteredMinDelay,
+      reconnectionDelayGrowFactor: 1.3,
+      connectionTimeout,
+      maxRetries,
+      debug,
+    })
+
+    // Capture a reference to this socket so event handlers can detect
+    // stale events from a previous socket (e.g. during React Strict Mode
+    // cleanup/remount cycles where disconnect + reconnect race).
+    const thisSocket = socket
+
+    socket.addEventListener('open', () => {
+      if (socket !== thisSocket) {
+        return
+      }
+      wasConnected = true
+      startSessionRefreshTimer()
+      notifyStatusChange('connected')
+    })
+
+    socket.addEventListener('close', () => {
+      if (socket !== thisSocket) {
+        return
+      }
+      stopSessionRefreshTimer()
+      // Ignore close events after intentional disconnect (socket already nulled)
+      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+      if (!socket) {
+        return
+      }
+      if (wasConnected) {
+        notifyStatusChange('reconnecting')
+      } else {
+        notifyStatusChange('disconnected')
+      }
+    })
+
+    socket.addEventListener('error', () => {
+      if (socket !== thisSocket) {
+        return
+      }
+      onError?.(new Error('WebSocket error - check Network tab for details'))
+      connectionStore.setError(new Error('WebSocket error'))
+    })
+
+    socket.addEventListener('message', (event) => {
+      if (socket !== thisSocket) {
+        return
+      }
+      try {
+        const message: unknown = JSON.parse((event as { data: string }).data)
+        onRawMessage?.(message)
+
+        // Check for connection established message
+        const connectionInfo = parseConnectionMessage(message)
+        if (connectionInfo) {
+          const { connectionId } = connectionInfo
+          subscriptionManager.setConnectionId(connectionId)
+          notifyConnectionEstablished(connectionId)
+
+          // Always resubscribe if there are active subscriptions
+          // Covers both reconnect and initial connect with queued subs
+          if (subscriptionManager.hasActiveSubscriptions()) {
+            subscriptionManager.resubscribeAll(connectionId).catch((error) => onError?.(error))
+          }
+          return
+        }
+
+        // Try to parse as a subscription message
+        const parsed = parseMessage(message)
+        if (parsed) {
+          subscriptionManager.dispatch(parsed.key, parsed.data)
+        }
+      } catch (error) {
+        onError?.(error)
+      }
+    })
+  }
+
+  function disconnect(): void {
+    if (socket) {
+      stopSessionRefreshTimer()
+      const s = socket
+      wasConnected = false
+      socket = null
+      s.close()
+      subscriptionManager.clear()
+      connectionStore.reset()
+      notifyStatusChange('disconnected')
+    }
+  }
+
+  function isConnected(): boolean {
+    return socket?.readyState === WebSocket.OPEN
+  }
+
+  function getConnectionStatus(): ConnectionStatus {
+    return connectionStore.getStatus()
+  }
+
+  function getConnectionId(): string | null {
+    return connectionStore.getConnectionId()
+  }
+
+  function subscribe(opts: SubscriptionOptions<TParams, TMessage>): () => void {
+    return subscriptionManager.subscribe({
+      channel: opts.channel,
+      params: opts.params,
+      callback: opts.onMessage,
+    })
+  }
+
+  function onStatusChange(callback: (status: ConnectionStatus) => void): () => void {
+    return connectionStore.onStatusChange(callback)
+  }
+
+  function onConnectionEstablished(callback: (connectionId: string) => void): () => void {
+    connectionCallbacks.add(callback)
+    return () => {
+      connectionCallbacks.delete(callback)
+    }
+  }
+
+  return {
+    isConnected,
+    getConnectionStatus,
+    getConnectionId,
+    subscribe,
+    onStatusChange,
+    onConnectionEstablished,
+  }
+}

package/src/index.ts

@@ -0,0 +1,22 @@
+// Main factory function
+export { createWebSocketClient } from './client/createWebSocketClient'
+export { createZustandConnectionStore } from './store/createZustandConnectionStore'
+// Store types (for consumers who need to access connection state)
+export type { ConnectionStore, CreateZustandConnectionStoreOptions } from './store/types'
+// Additional exports for advanced use cases
+export { SubscriptionManager } from './subscriptions/SubscriptionManager'
+// Subscription types (for consumers building custom subscription logic)
+export type { SubscribeInput, SubscriptionEntry, SubscriptionManagerOptions } from './subscriptions/types'
+export type {
+  ConnectionConfig,
+  ConnectionStatus,
+  CreateWebSocketClientOptions,
+  SocketFactoryOptions,
+  SubscriptionHandler,
+  SubscriptionOptions,
+  WebSocketClient,
+  // For testing/mocking
+  WebSocketLike,
+} from './types'
+// Utility functions
+export { addJitter, getDefaultJitteredDelay } from './utils/backoff'

package/src/store/createZustandConnectionStore.test.ts

@@ -0,0 +1,162 @@
+import { createZustandConnectionStore } from '@luxexchange/websocket/src/store/createZustandConnectionStore'
+import { describe, expect, it, vi } from 'vitest'
+
+describe('createZustandConnectionStore', () => {
+  describe('initial state', () => {
+    it('starts with disconnected status', () => {
+      const store = createZustandConnectionStore()
+      expect(store.getStatus()).toBe('disconnected')
+    })
+
+    it('starts with null connectionId', () => {
+      const store = createZustandConnectionStore()
+      expect(store.getConnectionId()).toBe(null)
+    })
+
+    it('starts with null error', () => {
+      const store = createZustandConnectionStore()
+      expect(store.getError()).toBe(null)
+    })
+  })
+
+  describe('setStatus', () => {
+    it('updates status to connecting', () => {
+      const store = createZustandConnectionStore()
+      store.setStatus('connecting')
+      expect(store.getStatus()).toBe('connecting')
+    })
+
+    it('updates status to connected', () => {
+      const store = createZustandConnectionStore()
+      store.setStatus('connected')
+      expect(store.getStatus()).toBe('connected')
+    })
+
+    it('updates status to reconnecting', () => {
+      const store = createZustandConnectionStore()
+      store.setStatus('reconnecting')
+      expect(store.getStatus()).toBe('reconnecting')
+    })
+
+    it('updates status to disconnected', () => {
+      const store = createZustandConnectionStore()
+      store.setStatus('connected')
+      store.setStatus('disconnected')
+      expect(store.getStatus()).toBe('disconnected')
+    })
+  })
+
+  describe('setConnectionId', () => {
+    it('sets connectionId', () => {
+      const store = createZustandConnectionStore()
+      store.setConnectionId('test-connection-123')
+      expect(store.getConnectionId()).toBe('test-connection-123')
+    })
+
+    it('clears connectionId with null', () => {
+      const store = createZustandConnectionStore()
+      store.setConnectionId('test-connection-123')
+      store.setConnectionId(null)
+      expect(store.getConnectionId()).toBe(null)
+    })
+  })
+
+  describe('setError', () => {
+    it('sets error', () => {
+      const store = createZustandConnectionStore()
+      const error = new Error('Connection failed')
+      store.setError(error)
+      expect(store.getError()).toBe(error)
+    })
+
+    it('clears error with null', () => {
+      const store = createZustandConnectionStore()
+      store.setError(new Error('Connection failed'))
+      store.setError(null)
+      expect(store.getError()).toBe(null)
+    })
+  })
+
+  describe('reset', () => {
+    it('resets all state to initial values', () => {
+      const store = createZustandConnectionStore()
+
+      // Modify all state
+      store.setStatus('connected')
+      store.setConnectionId('conn-123')
+      store.setError(new Error('test'))
+
+      // Reset
+      store.reset()
+
+      // Verify all reset
+      expect(store.getStatus()).toBe('disconnected')
+      expect(store.getConnectionId()).toBe(null)
+      expect(store.getError()).toBe(null)
+    })
+  })
+
+  describe('onStatusChange', () => {
+    it('calls callback when status changes', () => {
+      const store = createZustandConnectionStore()
+      const callback = vi.fn()
+
+      store.onStatusChange(callback)
+      store.setStatus('connecting')
+
+      expect(callback).toHaveBeenCalledWith('connecting')
+    })
+
+    it('does not call callback when other state changes', () => {
+      const store = createZustandConnectionStore()
+      const callback = vi.fn()
+
+      store.onStatusChange(callback)
+      store.setConnectionId('conn-123')
+      store.setError(new Error('test'))
+
+      expect(callback).not.toHaveBeenCalled()
+    })
+
+    it('returns unsubscribe function', () => {
+      const store = createZustandConnectionStore()
+      const callback = vi.fn()
+
+      const unsubscribe = store.onStatusChange(callback)
+      store.setStatus('connecting')
+      expect(callback).toHaveBeenCalledTimes(1)
+
+      unsubscribe()
+      store.setStatus('connected')
+      expect(callback).toHaveBeenCalledTimes(1)
+    })
+
+    it('tracks multiple status transitions', () => {
+      const store = createZustandConnectionStore()
+      const statuses: string[] = []
+
+      store.onStatusChange((status) => statuses.push(status))
+      store.setStatus('connecting')
+      store.setStatus('connected')
+      store.setStatus('reconnecting')
+      store.setStatus('disconnected')
+
+      expect(statuses).toEqual(['connecting', 'connected', 'reconnecting', 'disconnected'])
+    })
+  })
+
+  describe('options', () => {
+    it('creates store without devtools by default', () => {
+      const store = createZustandConnectionStore()
+      expect(store).toBeDefined()
+    })
+
+    it('creates store with custom devtools name', () => {
+      const store = createZustandConnectionStore({
+        enableDevtools: true,
+        devtoolsName: 'customStoreName',
+      })
+      expect(store).toBeDefined()
+    })
+  })
+})

package/src/store/createZustandConnectionStore.ts

@@ -0,0 +1,74 @@
+import type { ConnectionStore, CreateZustandConnectionStoreOptions } from '@luxexchange/websocket/src/store/types'
+import type { ConnectionStatus } from '@luxexchange/websocket/src/types'
+import { devtools } from 'zustand/middleware'
+import { createStore } from 'zustand/vanilla'
+
+interface ConnectionStoreState {
+  status: ConnectionStatus
+  connectionId: string | null
+  error: Error | null
+}
+
+/**
+ * Creates a Zustand-backed ConnectionStore.
+ * This is a convenience implementation — consumers can provide any object
+ * that satisfies the ConnectionStore interface.
+ */
+export function createZustandConnectionStore(options?: CreateZustandConnectionStoreOptions): ConnectionStore {
+  const { enableDevtools = false, devtoolsName = 'websocketConnectionStore' } = options ?? {}
+
+  const store = createStore<ConnectionStoreState>()(
+    devtools(
+      () => ({
+        status: 'disconnected' as ConnectionStatus,
+        connectionId: null,
+        error: null,
+      }),
+      {
+        name: devtoolsName,
+        enabled: enableDevtools,
+      },
+    ),
+  )
+
+  return {
+    getStatus(): ConnectionStatus {
+      return store.getState().status
+    },
+
+    setStatus(status: ConnectionStatus): void {
+      store.setState({ status }, false, 'setStatus')
+    },
+
+    getConnectionId(): string | null {
+      return store.getState().connectionId
+    },
+
+    setConnectionId(id: string | null): void {
+      store.setState({ connectionId: id }, false, 'setConnectionId')
+    },
+
+    getError(): Error | null {
+      return store.getState().error
+    },
+
+    setError(error: Error | null): void {
+      store.setState({ error }, false, 'setError')
+    },
+
+    reset(): void {
+      store.setState({ status: 'disconnected', connectionId: null, error: null }, false, 'reset')
+    },
+
+    onStatusChange(callback: (status: ConnectionStatus) => void): () => void {
+      let previousStatus = store.getState().status
+      return store.subscribe(() => {
+        const currentStatus = store.getState().status
+        if (currentStatus !== previousStatus) {
+          previousStatus = currentStatus
+          callback(currentStatus)
+        }
+      })
+    },
+  }
+}

package/src/store/types.ts

@@ -0,0 +1,17 @@
+import type { ConnectionStatus } from '@luxexchange/websocket/src/types'
+
+export interface ConnectionStore {
+  getStatus(): ConnectionStatus
+  setStatus(status: ConnectionStatus): void
+  getConnectionId(): string | null
+  setConnectionId(id: string | null): void
+  getError(): Error | null
+  setError(error: Error | null): void
+  reset(): void
+  onStatusChange(callback: (status: ConnectionStatus) => void): () => void
+}
+
+export interface CreateZustandConnectionStoreOptions {
+  enableDevtools?: boolean
+  devtoolsName?: string
+}