archy/neode-ui/src/api/websocket.ts

// WebSocket handler for real-time updates

import type { Update, PatchOperation } from '../types/api'
import { applyPatch } from 'fast-json-patch'

type WebSocketCallback = (update: Update) => void

export class WebSocketClient {
  private ws: WebSocket | null = null
  private callbacks: Set<WebSocketCallback> = new Set()
  private reconnectAttempts = 0
  private maxReconnectAttempts = 10
  private reconnectDelay = 1000
  private shouldReconnect = true
  private url: string
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null
  private isConnecting = false

  constructor(url: string = '/ws/db') {
    this.url = url
  }

  connect(): Promise<void> {
    return new Promise((resolve, reject) => {
      // If already connected, resolve immediately
      if (this.ws && this.ws.readyState === WebSocket.OPEN) {
        console.log('[WebSocket] Already connected, skipping')
        resolve()
        return
      }
      
      // If connecting, wait for it
      if (this.ws && this.ws.readyState === WebSocket.CONNECTING) {
        console.log('[WebSocket] Already connecting, waiting...')
        const checkInterval = setInterval(() => {
          if (this.ws) {
            if (this.ws.readyState === WebSocket.OPEN) {
              clearInterval(checkInterval)
              resolve()
            } else if (this.ws.readyState === WebSocket.CLOSED) {
              clearInterval(checkInterval)
              // Connection failed, will be handled by onclose
              reject(new Error('Connection closed during connect'))
            }
          } else {
            clearInterval(checkInterval)
            reject(new Error('WebSocket was cleared'))
          }
        }, 100)
        
        // Timeout after 5 seconds
        setTimeout(() => {
          clearInterval(checkInterval)
          if (this.ws && this.ws.readyState !== WebSocket.OPEN) {
            reject(new Error('Connection timeout'))
          }
        }, 5000)
        return
      }
      
      // Close existing connection if any (but don't prevent reconnection)
      if (this.ws) {
        const oldWs = this.ws
        this.ws = null
        // Temporarily disable reconnection to prevent loop
        const wasReconnecting = this.shouldReconnect
        this.shouldReconnect = false
        oldWs.onclose = null // Remove close handler
        oldWs.close()
        // Restore reconnection flag after a moment
        setTimeout(() => {
          this.shouldReconnect = wasReconnecting
        }, 100)
      }
      
      // Reset shouldReconnect flag when explicitly connecting
      this.shouldReconnect = true
      // Reset reconnect attempts only if we're explicitly connecting (not auto-reconnecting)
      // This allows reconnection attempts to continue
      
      // In development, Vite proxies /ws to the backend
      // In production, use the same host as the page
      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
      const host = window.location.host
      const wsUrl = `${protocol}//${host}${this.url}`
      
      console.log('[WebSocket] Connecting to:', wsUrl)

      this.ws = new WebSocket(wsUrl)
      
      // Timeout handler in case connection hangs
      const connectionTimeout = setTimeout(() => {
        if (this.ws && this.ws.readyState === WebSocket.CONNECTING) {
          console.warn('WebSocket connection timeout, retrying...')
          this.ws.close()
          reject(new Error('Connection timeout'))
        }
      }, 3000) // 3 second timeout

      this.ws.onopen = () => {
        clearTimeout(connectionTimeout)
        this.isConnecting = false
        this.reconnectAttempts = 0
        console.log('[WebSocket] Connected successfully')
        resolve()
      }

      this.ws.onerror = (error) => {
        clearTimeout(connectionTimeout)
        this.isConnecting = false
        console.error('[WebSocket] Connection error:', error)
        // Don't reject immediately - let onclose handle reconnection
        // This prevents errors from blocking reconnection
      }

      this.ws.onmessage = (event) => {
        try {
          const update: Update = JSON.parse(event.data)
          this.callbacks.forEach((callback) => callback(update))
        } catch (error) {
          console.error('Failed to parse WebSocket message:', error)
        }
      }

      this.ws.onclose = (event) => {
        clearTimeout(connectionTimeout)
        this.isConnecting = false
        console.log('[WebSocket] Closed', { code: event.code, reason: event.reason, wasClean: event.wasClean })
        
        // Clear the WebSocket reference
        this.ws = null
        
        // Don't reconnect if we explicitly disconnected
        if (!this.shouldReconnect) {
          console.log('[WebSocket] Reconnection disabled')
          return
        }
        
        // Always try to reconnect unless we've exceeded max attempts
        // Code 1001 (Going Away) happens on HMR reloads - reconnect IMMEDIATELY
        if (this.reconnectAttempts < this.maxReconnectAttempts) {
          // Immediate reconnection for HMR (code 1001) - no delay
          const isHMR = event.code === 1001 || event.code === 1006
          const delay = isHMR ? 0 : (this.reconnectAttempts === 0 ? 100 : Math.min(this.reconnectDelay * Math.pow(2, this.reconnectAttempts), 5000))
          console.log(`[WebSocket] Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts + 1}/${this.maxReconnectAttempts}, code: ${event.code}, HMR: ${isHMR})`)
          
          // Clear any existing reconnect timer
          if (this.reconnectTimer) {
            clearTimeout(this.reconnectTimer)
            this.reconnectTimer = null
          }
          
          const doReconnect = () => {
            // Check again if we should reconnect (might have been disabled)
            if (!this.shouldReconnect) {
              return
            }
            
            // Don't increment attempts for HMR disconnects - they're expected
            if (!isHMR) {
              this.reconnectAttempts++
            }
            
            console.log('[WebSocket] Attempting reconnection...')
            this.connect().catch((err) => {
              console.error('[WebSocket] Reconnection failed:', err)
              // onclose will be called again and will retry
            })
          }
          
          if (delay === 0) {
            // Immediate reconnection for HMR
            doReconnect()
          } else {
            this.reconnectTimer = setTimeout(() => {
              this.reconnectTimer = null
              doReconnect()
            }, delay)
          }
        } else {
          console.warn('[WebSocket] Max reconnection attempts reached')
          this.shouldReconnect = false
        }
      }
    })
  }

  subscribe(callback: WebSocketCallback): () => void {
    this.callbacks.add(callback)
    return () => {
      this.callbacks.delete(callback)
    }
  }

  disconnect(): void {
    this.shouldReconnect = false
    this.reconnectAttempts = 0
    this.isConnecting = false
    
    // Clear reconnect timer
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer)
      this.reconnectTimer = null
    }
    
    if (this.ws) {
      // Remove handlers to prevent reconnection
      this.ws.onclose = null
      this.ws.onerror = null
      try {
        this.ws.close()
      } catch (e) {
        // Ignore errors
      }
      this.ws = null
    }
  }
  
  reset(): void {
    this.disconnect()
    this.callbacks.clear()
  }

  isConnected(): boolean {
    return this.ws?.readyState === WebSocket.OPEN
  }
}

// Create singleton that persists across HMR
let wsClientInstance: WebSocketClient | null = null

function getWebSocketClient(): WebSocketClient {
  if (typeof window === 'undefined') {
    // SSR - create new instance
    if (!wsClientInstance) {
      wsClientInstance = new WebSocketClient()
    }
    return wsClientInstance
  }
  
  // Check if we have a persisted instance from HMR
  const existing = (window as any).__archipelago_ws_client
  if (existing && existing instanceof WebSocketClient) {
    // Check if the WebSocket is still valid
    if (existing.ws && existing.ws.readyState === WebSocket.OPEN) {
      console.log('[WebSocket] Using existing connected client from HMR')
      wsClientInstance = existing
      return existing
    }
  }
  
  // Create new instance
  if (!wsClientInstance) {
    wsClientInstance = new WebSocketClient()
    if (typeof window !== 'undefined') {
      (window as any).__archipelago_ws_client = wsClientInstance
    }
    console.log('[WebSocket] Created new client instance')
  }
  
  return wsClientInstance
}

// Lazy initialization - only create when accessed
let _wsClient: WebSocketClient | null = null

export const wsClient: WebSocketClient = (() => {
  if (_wsClient) {
    return _wsClient
  }
  try {
    _wsClient = getWebSocketClient()
    return _wsClient
  } catch (error) {
    console.error('[WebSocket] Error initializing client:', error)
    // Fallback to new instance
    _wsClient = new WebSocketClient()
    return _wsClient
  }
})()

// Helper to apply patches to data
export function applyDataPatch<T>(data: T, patch: PatchOperation[]): T {
  // Validate patch is an array before applying
  if (!Array.isArray(patch) || patch.length === 0) {
    console.warn('Invalid or empty patch received, returning original data')
    return data
  }
  
  try {
  const result = applyPatch(data, patch as any, false, false)
  return result.newDocument as T
  } catch (error) {
    console.error('Failed to apply patch:', error, 'Patch:', patch)
    return data // Return original data on error
  }
}