import { useEffect, useRef, useState } from "react"; import type { Socket } from "socket.io-client"; import { io } from "socket.io-client"; import { API_BASE_URL } from "@/lib/config"; interface Task { id: string; [key: string]: unknown; } interface Event { id: string; [key: string]: unknown; } interface Project { id: string; [key: string]: unknown; } interface DeletePayload { id: string; } interface WebSocketCallbacks { onTaskCreated?: (task: Task) => void; onTaskUpdated?: (task: Task) => void; onTaskDeleted?: (payload: DeletePayload) => void; onEventCreated?: (event: Event) => void; onEventUpdated?: (event: Event) => void; onEventDeleted?: (payload: DeletePayload) => void; onProjectUpdated?: (project: Project) => void; } interface ConnectionError { message: string; type: string; description?: string; } interface UseWebSocketReturn { isConnected: boolean; socket: Socket | null; connectionError: ConnectionError | null; } /** * Check if the WebSocket URL uses secure protocol (wss://) * Logs a warning in production when using insecure ws:// */ function validateWebSocketSecurity(url: string): void { const isProduction = process.env.NODE_ENV === "production"; const isSecure = url.startsWith("https://") || url.startsWith("wss://"); if (isProduction && !isSecure) { console.warn( "[Security Warning] WebSocket connection using insecure protocol (ws://). " + "Authentication tokens may be exposed. Use wss:// in production." ); } } /** * Hook for managing WebSocket connections and real-time updates * * @param workspaceId - The workspace ID to subscribe to * @param token - Authentication token * @param callbacks - Event callbacks for real-time updates * @returns Connection status, socket instance, and connection error */ export function useWebSocket( workspaceId: string, token: string, callbacks: WebSocketCallbacks = {} ): UseWebSocketReturn { const [socket, setSocket] = useState(null); const [isConnected, setIsConnected] = useState(false); const [connectionError, setConnectionError] = useState(null); // Use refs for callbacks to prevent stale closures and unnecessary reconnects // This ensures that callback changes don't trigger useEffect re-runs const callbacksRef = useRef(callbacks); // Keep the ref up-to-date with the latest callbacks useEffect(() => { callbacksRef.current = callbacks; }, [callbacks]); useEffect(() => { // Use WebSocket URL from central config const wsUrl = API_BASE_URL; // Validate WebSocket security - warn if using insecure connection in production validateWebSocketSecurity(wsUrl); // Clear any previous connection error setConnectionError(null); // Create socket connection const newSocket = io(wsUrl, { auth: { token }, query: { workspaceId }, }); setSocket(newSocket); // Connection event handlers const handleConnect = (): void => { setIsConnected(true); }; const handleDisconnect = (): void => { setIsConnected(false); }; const handleConnectError = (error: Error): void => { setConnectionError({ message: error.message || "Connection failed", type: "connect_error", description: "Failed to establish WebSocket connection", }); setIsConnected(false); }; // Wrapper functions that always use the latest callbacks via ref // This prevents stale closure issues while avoiding reconnects on callback changes const handleTaskCreated = (task: Task): void => { callbacksRef.current.onTaskCreated?.(task); }; const handleTaskUpdated = (task: Task): void => { callbacksRef.current.onTaskUpdated?.(task); }; const handleTaskDeleted = (payload: DeletePayload): void => { callbacksRef.current.onTaskDeleted?.(payload); }; const handleEventCreated = (event: Event): void => { callbacksRef.current.onEventCreated?.(event); }; const handleEventUpdated = (event: Event): void => { callbacksRef.current.onEventUpdated?.(event); }; const handleEventDeleted = (payload: DeletePayload): void => { callbacksRef.current.onEventDeleted?.(payload); }; const handleProjectUpdated = (project: Project): void => { callbacksRef.current.onProjectUpdated?.(project); }; newSocket.on("connect", handleConnect); newSocket.on("disconnect", handleDisconnect); newSocket.on("connect_error", handleConnectError); // Register all event handlers - they'll check the ref for actual callbacks newSocket.on("task:created", handleTaskCreated); newSocket.on("task:updated", handleTaskUpdated); newSocket.on("task:deleted", handleTaskDeleted); newSocket.on("event:created", handleEventCreated); newSocket.on("event:updated", handleEventUpdated); newSocket.on("event:deleted", handleEventDeleted); newSocket.on("project:updated", handleProjectUpdated); // Cleanup on unmount or dependency change return (): void => { newSocket.off("connect", handleConnect); newSocket.off("disconnect", handleDisconnect); newSocket.off("connect_error", handleConnectError); newSocket.off("task:created", handleTaskCreated); newSocket.off("task:updated", handleTaskUpdated); newSocket.off("task:deleted", handleTaskDeleted); newSocket.off("event:created", handleEventCreated); newSocket.off("event:updated", handleEventUpdated); newSocket.off("event:deleted", handleEventDeleted); newSocket.off("project:updated", handleProjectUpdated); newSocket.disconnect(); }; // Only stable values in deps - callbacks are accessed via ref }, [workspaceId, token]); return { isConnected, socket, connectionError, }; }