11"use client" ;
22
3- import React , { createContext , useContext , useEffect , useRef , useState , type ReactNode } from "react" ;
4- import type { LogEntry , WebSocketLogMessage } from "../lib/types/logs" ;
3+ import React , { createContext , useContext , useEffect , useRef , useState , useCallback , type ReactNode } from "react" ;
54import { getWebSocketUrl } from "@/lib/utils/port" ;
65
6+ type MessageHandler = ( data : any ) => void ;
7+
78interface WebSocketContextType {
89 isConnected : boolean ;
910 ws : React . RefObject < WebSocket | null > ;
10- setMessageHandler : ( handler : ( log : LogEntry , operation : "create" | "update" ) => void ) => void ;
11+ subscribe : ( channel : string , handler : MessageHandler ) => ( ) => void ;
12+ send : ( data : any ) => void ;
1113}
1214
1315const WebSocketContext = createContext < WebSocketContextType | null > ( null ) ;
1416
1517interface WebSocketProviderProps {
1618 children : ReactNode ;
19+ path ?: string ;
1720}
1821
1922// Global reference to maintain state across component remounts
2023let globalWsRef : WebSocket | null = null ;
21- let globalMessageHandler : ( ( log : LogEntry , operation : "create" | "update" ) => void ) | null = null ;
24+ const messageHandlers = new Map < string , Set < MessageHandler > > ( ) ;
2225
23- export function WebSocketProvider ( { children } : WebSocketProviderProps ) {
26+ export function WebSocketProvider ( { children, path = "/ws" } : WebSocketProviderProps ) {
2427 const wsRef = useRef < WebSocket | null > ( globalWsRef ) ;
2528 const reconnectTimeoutRef = useRef < ReturnType < typeof setTimeout > | null > ( null ) ;
29+ const pingTimerRef = useRef < ReturnType < typeof setInterval > | null > ( null ) ;
30+ const retryCountRef = useRef ( 0 ) ;
2631 const [ isConnected , setIsConnected ] = useState ( false ) ;
2732
28- const setMessageHandler = ( handler : ( log : LogEntry , operation : "create" | "update" ) => void ) => {
29- globalMessageHandler = handler ;
33+ const subscribe = useCallback < ( channel : string , handler : MessageHandler ) => ( ) => void > ( ( channel , handler ) => {
34+ if ( ! messageHandlers . has ( channel ) ) {
35+ messageHandlers . set ( channel , new Set ( ) ) ;
36+ }
37+ messageHandlers . get ( channel ) ! . add ( handler ) ;
38+
39+ // Return unsubscribe function
40+ return ( ) => {
41+ const handlers = messageHandlers . get ( channel ) ;
42+ if ( handlers ) {
43+ handlers . delete ( handler ) ;
44+ if ( handlers . size === 0 ) {
45+ messageHandlers . delete ( channel ) ;
46+ }
47+ }
48+ } ;
49+ } , [ ] ) ;
50+
51+ const send = ( data : any ) => {
52+ if ( wsRef . current ?. readyState === WebSocket . OPEN ) {
53+ try {
54+ wsRef . current . send ( typeof data === "string" ? data : JSON . stringify ( data ) ) ;
55+ } catch ( error ) {
56+ console . error ( "Failed to send WebSocket message:" , error ) ;
57+ }
58+ }
3059 } ;
3160
3261 useEffect ( ( ) => {
@@ -35,7 +64,7 @@ export function WebSocketProvider({ children }: WebSocketProviderProps) {
3564 return ;
3665 }
3766
38- const wsUrl = getWebSocketUrl ( "/ws/logs" ) ;
67+ const wsUrl = getWebSocketUrl ( path ) ;
3968
4069 const ws = new WebSocket ( wsUrl ) ;
4170 wsRef . current = ws ;
@@ -44,18 +73,44 @@ export function WebSocketProvider({ children }: WebSocketProviderProps) {
4473 ws . onopen = ( ) => {
4574 console . log ( "WebSocket connected" ) ;
4675 setIsConnected ( true ) ;
76+ retryCountRef . current = 0 ; // Reset retry count on successful connection
77+
4778 // Clear any pending reconnection attempts
4879 if ( reconnectTimeoutRef . current ) {
4980 clearTimeout ( reconnectTimeoutRef . current ) ;
5081 reconnectTimeoutRef . current = null ;
5182 }
83+
84+ // Start heartbeat/ping to keep connection alive
85+ if ( pingTimerRef . current ) {
86+ clearInterval ( pingTimerRef . current ) ;
87+ }
88+ pingTimerRef . current = setInterval ( ( ) => {
89+ if ( ws . readyState === WebSocket . OPEN ) {
90+ try {
91+ ws . send ( "ping" ) ;
92+ } catch ( error ) {
93+ console . error ( "Ping failed:" , error ) ;
94+ }
95+ }
96+ } , 25000 ) ; // Ping every 25 seconds
5297 } ;
5398
5499 ws . onmessage = ( event ) => {
55100 try {
56- const data : WebSocketLogMessage = JSON . parse ( event . data ) ;
57- if ( data . type === "log" && globalMessageHandler ) {
58- globalMessageHandler ( data . payload , data . operation ) ;
101+ const data = JSON . parse ( event . data ) ;
102+ const messageType = data . type || "default" ;
103+
104+ // Notify all subscribers for this message type
105+ const handlers = messageHandlers . get ( messageType ) ;
106+ if ( handlers ) {
107+ handlers . forEach ( ( handler ) => handler ( data ) ) ;
108+ }
109+
110+ // Also notify wildcard subscribers
111+ const wildcardHandlers = messageHandlers . get ( "*" ) ;
112+ if ( wildcardHandlers ) {
113+ wildcardHandlers . forEach ( ( handler ) => handler ( data ) ) ;
59114 }
60115 } catch ( error ) {
61116 console . error ( "Failed to parse WebSocket message:" , error ) ;
@@ -65,11 +120,23 @@ export function WebSocketProvider({ children }: WebSocketProviderProps) {
65120 ws . onclose = ( ) => {
66121 console . log ( "WebSocket disconnected, attempting to reconnect..." ) ;
67122 setIsConnected ( false ) ;
68- // Attempt to reconnect after 5 seconds
69- reconnectTimeoutRef . current = setTimeout ( connect , 5000 ) ;
123+
124+ // Clear ping timer
125+ if ( pingTimerRef . current ) {
126+ clearInterval ( pingTimerRef . current ) ;
127+ pingTimerRef . current = null ;
128+ }
129+
130+ // Exponential backoff: 0.5s, 1s, 2s, 4s, 8s, 16s, 32s (max)
131+ retryCountRef . current = Math . min ( retryCountRef . current + 1 , 6 ) ;
132+ const delay = Math . pow ( 2 , retryCountRef . current ) * 500 ;
133+ console . log ( `Reconnecting in ${ delay } ) ;
134+
135+ reconnectTimeoutRef . current = setTimeout ( connect , delay ) ;
70136 } ;
71137
72138 ws . onerror = ( error ) => {
139+ console . error ( "WebSocket error:" , error ) ;
73140 setIsConnected ( false ) ;
74141 ws . close ( ) ;
75142 } ;
@@ -84,10 +151,14 @@ export function WebSocketProvider({ children }: WebSocketProviderProps) {
84151 clearTimeout ( reconnectTimeoutRef . current ) ;
85152 reconnectTimeoutRef . current = null ;
86153 }
154+ if ( pingTimerRef . current ) {
155+ clearInterval ( pingTimerRef . current ) ;
156+ pingTimerRef . current = null ;
157+ }
87158 } ;
88- } , [ ] ) ;
159+ } , [ path ] ) ;
89160
90- return < WebSocketContext . Provider value = { { isConnected, ws : wsRef , setMessageHandler } } > { children } </ WebSocketContext . Provider > ;
161+ return < WebSocketContext . Provider value = { { isConnected, ws : wsRef , subscribe , send } } > { children } </ WebSocketContext . Provider > ;
91162}
92163
93164export function useWebSocket ( ) {
0 commit comments