@luxexchange/websocket 1.0.1 → 1.0.3

package/src/subscriptions/SubscriptionManager.ts

@@ -0,0 +1,283 @@
+import type {
+  SubscribeInput,
+  SubscriptionEntry,
+  SubscriptionManagerOptions,
+} from '@luxexchange/websocket/src/subscriptions/types'
+
+/**
+ * Manages subscription lifecycle with reference counting and microtask batching.
+ *
+ * Key features:
+ * - Reference counting: Only calls REST subscribe on first subscriber, unsubscribe on last
+ * - Microtask batching: Coalesces subscribe/unsubscribe calls within the same microtask
+ * - Eager cancellation: subscribe cancels a pending unsubscribe for the same key (and vice versa),
+ *   so navigations that unmount and remount the same subscription produce zero API calls
+ * - Auto-resubscribe: On reconnect, resubscribes all active subscriptions with new connectionId
+ * - Deduplication: Multiple subscribers to same params share one subscription
+ * - Message routing: Routes incoming messages to appropriate callbacks
+ */
+export class SubscriptionManager<TParams, TMessage> {
+  private subscriptions = new Map<string, SubscriptionEntry<TParams, TMessage>>()
+  private readonly handler: SubscriptionManagerOptions<TParams>['handler']
+  private readonly createKey: SubscriptionManagerOptions<TParams>['createKey']
+  private readonly onError?: SubscriptionManagerOptions<TParams>['onError']
+  private readonly onSubscriptionCountChange?: SubscriptionManagerOptions<TParams>['onSubscriptionCountChange']
+  private connectionId: string | null = null
+
+  private pendingSubscribes = new Map<string, TParams>()
+  private pendingUnsubscribes = new Map<string, TParams>()
+  private subscribeFlushScheduled = false
+  private unsubscribeFlushScheduled = false
+
+  constructor(options: SubscriptionManagerOptions<TParams>) {
+    this.handler = options.handler
+    this.createKey = options.createKey
+    this.onError = options.onError
+    this.onSubscriptionCountChange = options.onSubscriptionCountChange
+  }
+
+  /**
+   * Set the current connection ID. Called when connection is established.
+   */
+  setConnectionId(connectionId: string | null): void {
+    this.connectionId = connectionId
+  }
+
+  /**
+   * Get the current connection ID.
+   */
+  getConnectionId(): string | null {
+    return this.connectionId
+  }
+
+  /**
+   * Subscribe to a channel with given params.
+   * Synchronous — returns an unsubscribe function immediately.
+   * The actual REST subscribe call is batched via queueMicrotask.
+   */
+  subscribe(input: SubscribeInput<TParams, TMessage>): () => void {
+    const { channel, params, callback } = input
+    const key = this.createKey(channel, params)
+    let entry = this.subscriptions.get(key)
+    let isNewEntry = false
+
+    if (entry) {
+      entry.subscriberCount++
+      if (callback) {
+        entry.callbacks.add(callback)
+      }
+    } else {
+      isNewEntry = true
+      entry = {
+        channel,
+        params,
+        callbacks: new Set(callback ? [callback] : []),
+        subscriberCount: 1,
+      }
+      this.subscriptions.set(key, entry)
+
+      if (this.pendingUnsubscribes.has(key)) {
+        // Re-subscribing to a key that was just unsubscribed in this microtask window.
+        // Cancel the pending unsubscribe — the server still has this subscription.
+        this.pendingUnsubscribes.delete(key)
+      } else {
+        // Genuinely new subscription — queue the REST subscribe call
+        this.pendingSubscribes.set(key, params)
+        this.scheduleSubscribeFlush()
+      }
+    }
+
+    if (isNewEntry) {
+      this.onSubscriptionCountChange?.(this.subscriptions.size)
+    }
+
+    let unsubscribed = false
+    return () => {
+      if (unsubscribed) {
+        return
+      }
+      unsubscribed = true
+      this.handleUnsubscribe(key, callback)
+    }
+  }
+
+  private handleUnsubscribe(key: string, callback: ((message: TMessage) => void) | undefined): void {
+    const entry = this.subscriptions.get(key)
+    if (!entry) {
+      return
+    }
+
+    if (callback) {
+      entry.callbacks.delete(callback)
+    }
+
+    entry.subscriberCount--
+
+    // If no more subscribers, remove subscription entirely
+    if (entry.subscriberCount === 0) {
+      this.subscriptions.delete(key)
+
+      if (this.pendingSubscribes.has(key)) {
+        // Unsubscribing a key that was just subscribed in this microtask window.
+        // Cancel the pending subscribe — no need to tell the server about either.
+        this.pendingSubscribes.delete(key)
+      } else {
+        // Queue the REST unsubscribe call
+        this.pendingUnsubscribes.set(key, entry.params)
+        this.scheduleUnsubscribeFlush()
+      }
+
+      this.onSubscriptionCountChange?.(this.subscriptions.size)
+    }
+  }
+
+  private scheduleSubscribeFlush(): void {
+    if (!this.subscribeFlushScheduled) {
+      this.subscribeFlushScheduled = true
+      queueMicrotask(() => {
+        const params = [...this.pendingSubscribes.values()]
+        this.pendingSubscribes.clear()
+        this.subscribeFlushScheduled = false
+        if (params.length > 0 && this.connectionId) {
+          this.executeBatchSubscribe(params)
+        }
+      })
+    }
+  }
+
+  private scheduleUnsubscribeFlush(): void {
+    if (!this.unsubscribeFlushScheduled) {
+      this.unsubscribeFlushScheduled = true
+      queueMicrotask(() => {
+        const params = [...this.pendingUnsubscribes.values()]
+        this.pendingUnsubscribes.clear()
+        this.unsubscribeFlushScheduled = false
+        if (params.length > 0 && this.connectionId) {
+          this.executeBatchUnsubscribe(params)
+        }
+      })
+    }
+  }
+
+  private executeBatchSubscribe(params: TParams[]): void {
+    const { connectionId } = this
+    if (!connectionId) {
+      return
+    }
+    if (this.handler.subscribeBatch) {
+      this.handler.subscribeBatch(connectionId, params).catch((error) => {
+        this.onError?.(error, 'subscribe')
+      })
+    } else {
+      Promise.all(params.map((p) => this.handler.subscribe(connectionId, p))).catch((error) => {
+        this.onError?.(error, 'subscribe')
+      })
+    }
+  }
+
+  private executeBatchUnsubscribe(params: TParams[]): void {
+    const { connectionId } = this
+    if (!connectionId) {
+      return
+    }
+    if (this.handler.unsubscribeBatch) {
+      this.handler.unsubscribeBatch(connectionId, params).catch((error) => {
+        this.onError?.(error, 'unsubscribe')
+      })
+    } else {
+      Promise.all(params.map((p) => this.handler.unsubscribe(connectionId, p))).catch((error) => {
+        this.onError?.(error, 'unsubscribe')
+      })
+    }
+  }
+
+  /**
+   * Resubscribe all active subscriptions with a new connection ID.
+   * Called after reconnection. Uses subscribeBatch when available.
+   */
+  async resubscribeAll(connectionId: string): Promise<void> {
+    this.connectionId = connectionId
+
+    const allParams = Array.from(this.subscriptions.values()).map((entry) => entry.params)
+    if (allParams.length === 0) {
+      return
+    }
+
+    if (this.handler.subscribeBatch) {
+      try {
+        await this.handler.subscribeBatch(connectionId, allParams)
+      } catch (error) {
+        this.onError?.(error, 'resubscribe')
+      }
+    } else {
+      const subscribePromises = allParams.map((params) =>
+        this.handler.subscribe(connectionId, params).catch((error) => {
+          this.onError?.(error, 'resubscribe')
+        }),
+      )
+      await Promise.all(subscribePromises)
+    }
+  }
+
+  /**
+   * Dispatch an incoming message to appropriate callbacks.
+   */
+  dispatch(key: string, message: TMessage): void {
+    const entry = this.subscriptions.get(key)
+    if (!entry) {
+      return
+    }
+
+    for (const callback of entry.callbacks) {
+      try {
+        callback(message)
+      } catch (error) {
+        this.onError?.(error, 'dispatch')
+      }
+    }
+  }
+
+  /**
+   * Get all active subscriptions.
+   */
+  getActiveSubscriptions(): Array<{ channel: string; params: TParams; subscriberCount: number }> {
+    return Array.from(this.subscriptions.entries()).map(([, entry]) => ({
+      channel: entry.channel,
+      params: entry.params,
+      subscriberCount: entry.subscriberCount,
+    }))
+  }
+
+  /**
+   * Check if there are any active subscriptions.
+   */
+  hasActiveSubscriptions(): boolean {
+    return this.subscriptions.size > 0
+  }
+
+  /**
+   * Clear all subscriptions without calling unsubscribe API.
+   * Used on disconnect.
+   */
+  clear(): void {
+    this.subscriptions.clear()
+    this.pendingSubscribes.clear()
+    this.pendingUnsubscribes.clear()
+    this.connectionId = null
+  }
+
+  /**
+   * Refresh the session if the handler supports it.
+   */
+  async refreshSession(): Promise<void> {
+    if (!this.connectionId || !this.handler.refreshSession) {
+      return
+    }
+
+    try {
+      await this.handler.refreshSession(this.connectionId)
+    } catch (error) {
+      this.onError?.(error, 'refreshSession')
+    }
+  }
+}

package/src/subscriptions/types.ts

@@ -0,0 +1,21 @@
+import type { SubscriptionHandler } from '@luxexchange/websocket/src/types'
+
+export interface SubscriptionEntry<TParams, TMessage> {
+  channel: string
+  params: TParams
+  callbacks: Set<(message: TMessage) => void>
+  subscriberCount: number
+}
+
+export interface SubscriptionManagerOptions<TParams> {
+  handler: SubscriptionHandler<TParams>
+  createKey: (channel: string, params: TParams) => string
+  onError?: (error: unknown, operation: string) => void
+  onSubscriptionCountChange?: (count: number) => void
+}
+
+export interface SubscribeInput<TParams, TMessage> {
+  channel: string
+  params: TParams
+  callback?: (message: TMessage) => void
+}

package/src/types.ts

@@ -0,0 +1,88 @@
+/**
+ * Core types for generic WebSocket client functionality
+ */
+
+import type { ConnectionStore } from '@luxexchange/websocket/src/store/types'
+
+export type ConnectionStatus = 'disconnected' | 'connecting' | 'connected' | 'reconnecting'
+
+/**
+ * Minimal interface for WebSocket-like objects.
+ * Allows injection of mock WebSocket implementations for testing.
+ */
+export interface WebSocketLike {
+  readyState: number
+  addEventListener(event: string, handler: (event: unknown) => void): void
+  close(): void
+}
+
+/**
+ * Options for creating a WebSocket connection.
+ * These match the options expected by PartySocket.
+ */
+export interface SocketFactoryOptions {
+  maxReconnectionDelay: number
+  minReconnectionDelay: number
+  reconnectionDelayGrowFactor: number
+  connectionTimeout: number
+  maxRetries: number
+  debug: boolean
+}
+
+export interface ConnectionConfig {
+  url: string
+  maxReconnectionDelay?: number
+  minReconnectionDelay?: number
+  connectionTimeout?: number
+  maxRetries?: number
+  debug?: boolean
+}
+
+/**
+ * Subscription handler - consumers provide REST API implementation
+ * for managing server-side subscription lifecycle
+ */
+export interface SubscriptionHandler<TParams = unknown> {
+  subscribe: (connectionId: string, params: TParams) => Promise<void>
+  unsubscribe: (connectionId: string, params: TParams) => Promise<void>
+  subscribeBatch?: (connectionId: string, params: TParams[]) => Promise<void>
+  unsubscribeBatch?: (connectionId: string, params: TParams[]) => Promise<void>
+  refreshSession?: (connectionId: string) => Promise<void>
+}
+
+export interface SubscriptionOptions<TParams, TMessage> {
+  channel: string
+  params: TParams
+  onMessage?: (message: TMessage) => void
+}
+
+export interface WebSocketClient<TParams = unknown, TMessage = unknown> {
+  isConnected: () => boolean
+  getConnectionStatus: () => ConnectionStatus
+  getConnectionId: () => string | null
+  subscribe: (options: SubscriptionOptions<TParams, TMessage>) => () => void
+  onStatusChange: (callback: (status: ConnectionStatus) => void) => () => void
+  onConnectionEstablished: (callback: (connectionId: string) => void) => () => void
+}
+
+export interface CreateWebSocketClientOptions<TParams, TMessage> {
+  config: ConnectionConfig
+  connectionStore: ConnectionStore
+  subscriptionHandler: SubscriptionHandler<TParams>
+  parseMessage: (raw: unknown) => { channel: string; key: string; data: TMessage } | null
+  parseConnectionMessage: (raw: unknown) => { connectionId: string } | null
+  createSubscriptionKey: (channel: string, params: TParams) => string
+  onError?: (error: unknown) => void
+  onRawMessage?: (message: unknown) => void
+  /**
+   * Interval (ms) for automatically refreshing the server session.
+   * When set, the client starts a timer on connect and stops it on disconnect.
+   * Requires the subscriptionHandler to implement refreshSession.
+   */
+  sessionRefreshIntervalMs?: number
+  /**
+   * Optional factory for creating WebSocket instances.
+   * Defaults to PartySocket. Primarily used for testing with mock WebSockets.
+   */
+  socketFactory?: (url: string, options: SocketFactoryOptions) => WebSocketLike
+}

package/src/utils/backoff.test.ts

@@ -0,0 +1,48 @@
+import { addJitter, getDefaultJitteredDelay } from '@luxexchange/websocket/src/utils/backoff'
+import { describe, expect, it, vi } from 'vitest'
+
+describe('backoff utilities', () => {
+  describe('addJitter', () => {
+    it('returns value in range [minDelay, minDelay + jitterRange]', () => {
+      const minDelay = 1000
+      const jitterRange = 4000
+
+      // Run multiple times to test randomness
+      for (let i = 0; i < 100; i++) {
+        const result = addJitter(minDelay, jitterRange)
+        expect(result).toBeGreaterThanOrEqual(minDelay)
+        expect(result).toBeLessThan(minDelay + jitterRange)
+      }
+    })
+
+    it('returns exact minDelay when Math.random returns 0', () => {
+      vi.spyOn(Math, 'random').mockReturnValue(0)
+
+      const result = addJitter(1000, 4000)
+
+      expect(result).toBe(1000)
+
+      vi.restoreAllMocks()
+    })
+
+    it('returns near max when Math.random returns near 1', () => {
+      vi.spyOn(Math, 'random').mockReturnValue(0.999)
+
+      const result = addJitter(1000, 4000)
+
+      expect(result).toBeCloseTo(4996, 0)
+
+      vi.restoreAllMocks()
+    })
+  })
+
+  describe('getDefaultJitteredDelay', () => {
+    it('returns value between 1000 and 5000', () => {
+      for (let i = 0; i < 100; i++) {
+        const result = getDefaultJitteredDelay()
+        expect(result).toBeGreaterThanOrEqual(1000)
+        expect(result).toBeLessThan(5000)
+      }
+    })
+  })
+})

package/src/utils/backoff.ts

@@ -0,0 +1,15 @@
+/**
+ * Add jitter to a delay to prevent thundering herd on reconnect.
+ * Returns a value between minDelay and minDelay + jitterRange.
+ */
+export function addJitter(minDelay: number, jitterRange: number): number {
+  return minDelay + Math.random() * jitterRange
+}
+
+/**
+ * Default jitter configuration for WebSocket reconnection.
+ * Base delay of 1000ms with up to 4000ms additional jitter.
+ */
+export function getDefaultJitteredDelay(): number {
+  return addJitter(1000, 4000)
+}

package/tsconfig.json

@@ -1,12 +1,28 @@
 {
   "extends": "../../config/tsconfig/app.json",
-  "include": ["src/**/*.ts", "src/**/*.tsx", "src/**/*.json", "src/global.d.ts"],
-  "exclude": ["src/**/*.spec.ts", "src/**/*.spec.tsx", "src/**/*.test.ts", "src/**/*.test.tsx"],
+  "include": [
+    "src/**/*.ts",
+    "src/**/*.tsx",
+    "src/**/*.json",
+    "src/global.d.ts"
+  ],
+  "exclude": [
+    "src/**/*.spec.ts",
+    "src/**/*.spec.tsx",
+    "src/**/*.test.ts",
+    "src/**/*.test.tsx"
+  ],
   "compilerOptions": {
     "noEmit": false,
     "emitDeclarationOnly": true,
     "types": ["node", "vitest/globals"],
-    "paths": {}
+    "paths": {
+      "@luxexchange/websocket/src/*": ["./src/*"]
+    }
   },
-  "references": []
+  "references": [
+    {
+      "path": "../eslint-config"
+    }
+  ]
 }

package/tsconfig.lint.json

@@ -0,0 +1,8 @@
+{
+  "extends": "./tsconfig.json",
+  "compilerOptions": {
+    "preserveSymlinks": true
+  },
+  "include": ["**/*.ts", "**/*.tsx", "**/*.json"],
+  "exclude": ["node_modules"]
+}

package/vitest.config.ts

@@ -0,0 +1,17 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    environment: 'node',
+    coverage: {
+      exclude: [
+        '**/__generated__/**',
+        '**/node_modules/**',
+        '**/dist/**',
+        '**/*.config.*',
+        '**/scripts/**',
+        '**/.eslintrc.*',
+      ],
+    },
+  },
+})