From e7d6c40452b206d883679ad88922244dc8fc560d Mon Sep 17 00:00:00 2001 From: Haileyesus <118998054+blackmammoth@users.noreply.github.com> Date: Tue, 3 Feb 2026 12:05:15 +0300 Subject: [PATCH] Refactor WebSocket context + centralize platform flag (#363) * fix: remove unnecessary websocket.js file and replace its usage directly in `WebSocketContext` * fix: connect() doesn't need to be async * fix: update WebSocket context import to use useWebSocket hook * fix: use `useRef` for WebSocketContext The main issue with using states was, previously the websocket never closed properly on unmount, so multiple connections could be opened. This was because the useEffect cleanup function was closing an old websocket (that was initialized to null) instead of the current one. We could have fixed this by adding `ws` to the useEffect dependency array, but this was unnecessary since `ws` doesn't affect rendering so we shouldn't use a state. * fix: replace `WebSocketContext` default value with null and add type definitions * fix: add type definition for WebSocket URL and remove redundant protocol declaration * fix: Prevent WebSocket reconnection attempts after unmount Right now, when the WebSocketContext component unmounts, there is still a pending reconnection attempt that tries to reconnect the WebSocket after 3 seconds. * refactor: Extract WebSocket URL construction into a separate function * refactor: Centralize platform mode detection using IS_PLATFORM constant; use `token` from Auth context in WebSocket connection * refactor: Use IS_PLATFORM constant for platform detection in authenticatedFetch function (backend) * refactor: move IS_PLATFORM to config file for both frontend and backend The reason we couldn't place it in shared/modelConstants.js is that the frontend uses Vite which requires import.meta.env for environment variables, while the backend uses process.env. Therefore, we created separate config files for the frontend (src/constants/config.ts) and backend (server/constants/config.js). * refactor: update import path for IS_PLATFORM constant to use config file * refactor: replace `messages` with `latestMessage` in WebSocket context and related components Why? Because, messages was only being used to access the latest message in the components it's used in. * refactor: optimize WebSocket connection handling with useCallback and useMemo * refactor: comment out debug log for render count in AppContent component * refactor(backend): update environment variable handling and replace VITE_IS_PLATFORM with IS_PLATFORM constant * refactor: update WebSocket connection effect to depend on token changes for reconnection --------- --- server/constants/config.js | 5 ++ server/index.js | 22 +---- server/load-env.js | 24 ++++++ server/middleware/auth.js | 5 +- server/routes/agent.js | 5 +- shared/modelConstants.js | 2 +- src/App.jsx | 16 ++-- src/components/ChatInterface.jsx | 11 +-- src/components/GitPanel.jsx | 1 + src/components/LoginModal.jsx | 5 +- src/components/MainContent.jsx | 6 +- src/components/Onboarding.jsx | 4 +- src/components/ProtectedRoute.jsx | 3 +- src/components/Shell.jsx | 4 +- src/components/Sidebar.jsx | 5 +- src/constants/config.ts | 5 ++ src/contexts/AuthContext.jsx | 3 +- src/contexts/TaskMasterContext.jsx | 9 +-- src/contexts/WebSocketContext.jsx | 29 ------- src/contexts/WebSocketContext.tsx | 125 +++++++++++++++++++++++++++++ src/utils/api.js | 5 +- src/utils/websocket.js | 94 ---------------------- 22 files changed, 209 insertions(+), 179 deletions(-) create mode 100644 server/constants/config.js create mode 100644 server/load-env.js create mode 100644 src/constants/config.ts delete mode 100644 src/contexts/WebSocketContext.jsx create mode 100644 src/contexts/WebSocketContext.tsx delete mode 100755 src/utils/websocket.js diff --git a/server/constants/config.js b/server/constants/config.js new file mode 100644 index 0000000..580a985 --- /dev/null +++ b/server/constants/config.js @@ -0,0 +1,5 @@ +/** + * Environment Flag: Is Platform + * Indicates if the app is running in Platform mode (hosted) or OSS mode (self-hosted) + */ +export const IS_PLATFORM = process.env.VITE_IS_PLATFORM === 'true'; \ No newline at end of file diff --git a/server/index.js b/server/index.js index 1c42d30..1db726d 100755 --- a/server/index.js +++ b/server/index.js @@ -1,5 +1,6 @@ #!/usr/bin/env node -// Load environment variables from .env file +// Load environment variables before other imports execute +import './load-env.js'; import fs from 'fs'; import path from 'path'; import { fileURLToPath } from 'url'; @@ -28,22 +29,6 @@ const c = { dim: (text) => `${colors.dim}${text}${colors.reset}`, }; -try { - const envPath = path.join(__dirname, '../.env'); - const envFile = fs.readFileSync(envPath, 'utf8'); - envFile.split('\n').forEach(line => { - const trimmedLine = line.trim(); - if (trimmedLine && !trimmedLine.startsWith('#')) { - const [key, ...valueParts] = trimmedLine.split('='); - if (key && valueParts.length > 0 && !process.env[key]) { - process.env[key] = valueParts.join('=').trim(); - } - } - }); -} catch (e) { - console.log('No .env file found or error reading it:', e.message); -} - console.log('PORT from env:', process.env.PORT); import express from 'express'; @@ -76,6 +61,7 @@ import userRoutes from './routes/user.js'; import codexRoutes from './routes/codex.js'; import { initializeDatabase } from './database/db.js'; import { validateApiKey, authenticateToken, authenticateWebSocket } from './middleware/auth.js'; +import { IS_PLATFORM } from './constants/config.js'; // File system watcher for projects folder let projectsWatcher = null; @@ -200,7 +186,7 @@ const wss = new WebSocketServer({ console.log('WebSocket connection attempt to:', info.req.url); // Platform mode: always allow connection - if (process.env.VITE_IS_PLATFORM === 'true') { + if (IS_PLATFORM) { const user = authenticateWebSocket(null); // Will return first user if (!user) { console.log('[WARN] Platform mode: No user found in database'); diff --git a/server/load-env.js b/server/load-env.js new file mode 100644 index 0000000..21280a4 --- /dev/null +++ b/server/load-env.js @@ -0,0 +1,24 @@ +// Load environment variables from .env before other imports execute. +import fs from 'fs'; +import path from 'path'; +import { fileURLToPath } from 'url'; +import { dirname } from 'path'; + +const __filename = fileURLToPath(import.meta.url); +const __dirname = dirname(__filename); + +try { + const envPath = path.join(__dirname, '../.env'); + const envFile = fs.readFileSync(envPath, 'utf8'); + envFile.split('\n').forEach(line => { + const trimmedLine = line.trim(); + if (trimmedLine && !trimmedLine.startsWith('#')) { + const [key, ...valueParts] = trimmedLine.split('='); + if (key && valueParts.length > 0 && !process.env[key]) { + process.env[key] = valueParts.join('=').trim(); + } + } + }); +} catch (e) { + console.log('No .env file found or error reading it:', e.message); +} diff --git a/server/middleware/auth.js b/server/middleware/auth.js index 9231e4e..ab12e0c 100644 --- a/server/middleware/auth.js +++ b/server/middleware/auth.js @@ -1,5 +1,6 @@ import jwt from 'jsonwebtoken'; import { userDb } from '../database/db.js'; +import { IS_PLATFORM } from '../constants/config.js'; // Get JWT secret from environment or use default (for development) const JWT_SECRET = process.env.JWT_SECRET || 'claude-ui-dev-secret-change-in-production'; @@ -21,7 +22,7 @@ const validateApiKey = (req, res, next) => { // JWT authentication middleware const authenticateToken = async (req, res, next) => { // Platform mode: use single database user - if (process.env.VITE_IS_PLATFORM === 'true') { + if (IS_PLATFORM) { try { const user = userDb.getFirstUser(); if (!user) { @@ -80,7 +81,7 @@ const generateToken = (user) => { // WebSocket authentication function const authenticateWebSocket = (token) => { // Platform mode: bypass token validation, return first user - if (process.env.VITE_IS_PLATFORM === 'true') { + if (IS_PLATFORM) { try { const user = userDb.getFirstUser(); if (user) { diff --git a/server/routes/agent.js b/server/routes/agent.js index f633034..3ef2620 100644 --- a/server/routes/agent.js +++ b/server/routes/agent.js @@ -11,6 +11,7 @@ import { spawnCursor } from '../cursor-cli.js'; import { queryCodex } from '../openai-codex.js'; import { Octokit } from '@octokit/rest'; import { CLAUDE_MODELS, CURSOR_MODELS, CODEX_MODELS } from '../../shared/modelConstants.js'; +import { IS_PLATFORM } from '../constants/config.js'; const router = express.Router(); @@ -18,7 +19,7 @@ const router = express.Router(); * Middleware to authenticate agent API requests. * * Supports two authentication modes: - * 1. Platform mode (VITE_IS_PLATFORM=true): For managed/hosted deployments where + * 1. Platform mode (IS_PLATFORM=true): For managed/hosted deployments where * authentication is handled by an external proxy. Requests are trusted and * the default user context is used. * @@ -28,7 +29,7 @@ const router = express.Router(); const validateExternalApiKey = (req, res, next) => { // Platform mode: Authentication is handled externally (e.g., by a proxy layer). // Trust the request and use the default user context. - if (process.env.VITE_IS_PLATFORM === 'true') { + if (IS_PLATFORM) { try { const user = userDb.getFirstUser(); if (!user) { diff --git a/shared/modelConstants.js b/shared/modelConstants.js index 7d4347f..a327c9f 100644 --- a/shared/modelConstants.js +++ b/shared/modelConstants.js @@ -62,4 +62,4 @@ export const CODEX_MODELS = { ], DEFAULT: 'gpt-5.2' -}; +}; \ No newline at end of file diff --git a/src/App.jsx b/src/App.jsx index b65f65b..3ded1e6 100644 --- a/src/App.jsx +++ b/src/App.jsx @@ -31,7 +31,7 @@ import { ThemeProvider } from './contexts/ThemeContext'; import { AuthProvider } from './contexts/AuthContext'; import { TaskMasterProvider } from './contexts/TaskMasterContext'; import { TasksSettingsProvider } from './contexts/TasksSettingsContext'; -import { WebSocketProvider, useWebSocketContext } from './contexts/WebSocketContext'; +import { WebSocketProvider, useWebSocket } from './contexts/WebSocketContext'; import ProtectedRoute from './components/ProtectedRoute'; import { useVersionCheck } from './hooks/useVersionCheck'; import useLocalStorage from './hooks/useLocalStorage'; @@ -40,11 +40,15 @@ import { I18nextProvider, useTranslation } from 'react-i18next'; import i18n from './i18n/config.js'; +// ! Move to a separate file called AppContent.ts // Main App component with routing function AppContent() { const navigate = useNavigate(); const { sessionId } = useParams(); const { t } = useTranslation('common'); + // * This is a tracker for avoiding excessive re-renders during development + const renderCountRef = useRef(0); + // console.log(`AppContent render count: ${renderCountRef.current++}`); const { updateAvailable, latestVersion, currentVersion, releaseInfo } = useVersionCheck('siteboon', 'claudecodeui'); const [showVersionModal, setShowVersionModal] = useState(false); @@ -81,7 +85,7 @@ function AppContent() { // Triggers ChatInterface to reload messages without switching sessions const [externalMessageUpdate, setExternalMessageUpdate] = useState(0); - const { ws, sendMessage, messages } = useWebSocketContext(); + const { ws, sendMessage, latestMessage } = useWebSocket(); // Ref to track loading progress timeout for cleanup const loadingProgressTimeoutRef = useRef(null); @@ -175,9 +179,7 @@ function AppContent() { // Handle WebSocket messages for real-time project updates useEffect(() => { - if (messages.length > 0) { - const latestMessage = messages[messages.length - 1]; - + if (latestMessage) { // Handle loading progress updates if (latestMessage.type === 'loading_progress') { if (loadingProgressTimeoutRef.current) { @@ -277,7 +279,7 @@ function AppContent() { loadingProgressTimeoutRef.current = null; } }; - }, [messages, selectedProject, selectedSession, activeSessions]); + }, [latestMessage, selectedProject, selectedSession, activeSessions]); const fetchProjects = async () => { try { @@ -916,7 +918,7 @@ function AppContent() { setActiveTab={setActiveTab} ws={ws} sendMessage={sendMessage} - messages={messages} + latestMessage={latestMessage} isMobile={isMobile} isPWA={isPWA} onMenuClick={() => setSidebarOpen(true)} diff --git a/src/components/ChatInterface.jsx b/src/components/ChatInterface.jsx index 3ee5a84..d3d7bea 100644 --- a/src/components/ChatInterface.jsx +++ b/src/components/ChatInterface.jsx @@ -43,6 +43,8 @@ import { CLAUDE_MODELS, CURSOR_MODELS, CODEX_MODELS } from '../../shared/modelCo import { safeJsonParse } from '../lib/utils.js'; +// ! Move all utility functions to utils/chatUtils.ts + // Helper function to decode HTML entities in text function decodeHtmlEntities(text) { if (!text) return text; @@ -1860,7 +1862,7 @@ const ImageAttachment = ({ file, onRemove, uploadProgress, error }) => { // - onReplaceTemporarySession: Called to replace temporary session ID with real WebSocket session ID // // This ensures uninterrupted chat experience by pausing sidebar refreshes during conversations. -function ChatInterface({ selectedProject, selectedSession, ws, sendMessage, messages, onFileOpen, onInputFocusChange, onSessionActive, onSessionInactive, onSessionProcessing, onSessionNotProcessing, processingSessions, onReplaceTemporarySession, onNavigateToSession, onShowSettings, autoExpandTools, showRawParameters, showThinking, autoScrollToBottom, sendByCtrlEnter, externalMessageUpdate, onTaskClick, onShowAllTasks }) { +function ChatInterface({ selectedProject, selectedSession, ws, sendMessage, latestMessage, onFileOpen, onInputFocusChange, onSessionActive, onSessionInactive, onSessionProcessing, onSessionNotProcessing, processingSessions, onReplaceTemporarySession, onNavigateToSession, onShowSettings, autoExpandTools, showRawParameters, showThinking, autoScrollToBottom, sendByCtrlEnter, externalMessageUpdate, onTaskClick, onShowAllTasks }) { const { tasksEnabled, isTaskMasterInstalled } = useTasksSettings(); const { t } = useTranslation('chat'); const [input, setInput] = useState(() => { @@ -3242,8 +3244,7 @@ function ChatInterface({ selectedProject, selectedSession, ws, sendMessage, mess useEffect(() => { // Handle WebSocket messages - if (messages.length > 0) { - const latestMessage = messages[messages.length - 1]; + if (latestMessage) { const messageData = latestMessage.data?.message || latestMessage.data; // Filter messages by session ID to prevent cross-session interference @@ -4068,7 +4069,7 @@ function ChatInterface({ selectedProject, selectedSession, ws, sendMessage, mess } } - }, [messages]); + }, [latestMessage]); // Load file list when project changes useEffect(() => { @@ -4879,7 +4880,7 @@ function ChatInterface({ selectedProject, selectedSession, ws, sendMessage, mess }; - +// ! Unused const handleNewSession = () => { setChatMessages([]); setInput(''); diff --git a/src/components/GitPanel.jsx b/src/components/GitPanel.jsx index 3abd4c5..b077b05 100644 --- a/src/components/GitPanel.jsx +++ b/src/components/GitPanel.jsx @@ -960,6 +960,7 @@ function GitPanel({ selectedProject, isMobile, onFileOpen }) { {gitStatus.details && (

{gitStatus.details}

)} + {/* // ! This can be a custom component that can be reused for " Tip: Create a new project..." as well */}

Tip: Run git init in your project directory to initialize git source control. diff --git a/src/components/LoginModal.jsx b/src/components/LoginModal.jsx index 6861313..811daf2 100644 --- a/src/components/LoginModal.jsx +++ b/src/components/LoginModal.jsx @@ -1,5 +1,6 @@ import { X } from 'lucide-react'; import StandaloneShell from './StandaloneShell'; +import { IS_PLATFORM } from '../constants/config'; /** * Reusable login modal component for Claude, Cursor, and Codex CLI authentication @@ -27,15 +28,13 @@ function LoginModal({ const getCommand = () => { if (customCommand) return customCommand; - const isPlatform = import.meta.env.VITE_IS_PLATFORM === 'true'; - switch (provider) { case 'claude': return isAuthenticated ? 'claude setup-token --dangerously-skip-permissions' : 'claude /exit --dangerously-skip-permissions'; case 'cursor': return 'cursor-agent login'; case 'codex': - return isPlatform ? 'codex login --device-auth' : 'codex login'; + return IS_PLATFORM ? 'codex login --device-auth' : 'codex login'; default: return isAuthenticated ? 'claude setup-token --dangerously-skip-permissions' : 'claude /exit --dangerously-skip-permissions'; } diff --git a/src/components/MainContent.jsx b/src/components/MainContent.jsx index 50949a0..c1db1b5 100644 --- a/src/components/MainContent.jsx +++ b/src/components/MainContent.jsx @@ -36,9 +36,9 @@ function MainContent({ setActiveTab, ws, sendMessage, - messages, + latestMessage, isMobile, - isPWA, + isPWA, // ! Unused onMenuClick, isLoading, onInputFocusChange, @@ -477,7 +477,7 @@ function MainContent({ selectedSession={selectedSession} ws={ws} sendMessage={sendMessage} - messages={messages} + latestMessage={latestMessage} onFileOpen={handleFileOpen} onInputFocusChange={onInputFocusChange} onSessionActive={onSessionActive} diff --git a/src/components/Onboarding.jsx b/src/components/Onboarding.jsx index 29bf762..17125ae 100644 --- a/src/components/Onboarding.jsx +++ b/src/components/Onboarding.jsx @@ -6,6 +6,7 @@ import CodexLogo from './CodexLogo'; import LoginModal from './LoginModal'; import { authenticatedFetch } from '../utils/api'; import { useAuth } from '../contexts/AuthContext'; +import { IS_PLATFORM } from '../constants/config'; const Onboarding = ({ onComplete }) => { const [currentStep, setCurrentStep] = useState(0); @@ -15,8 +16,7 @@ const Onboarding = ({ onComplete }) => { const [error, setError] = useState(''); const [activeLoginProvider, setActiveLoginProvider] = useState(null); - const isPlatform = import.meta.env.VITE_IS_PLATFORM === 'true'; - const [selectedProject] = useState({ name: 'default', fullPath: isPlatform ? '/workspace' : '' }); + const [selectedProject] = useState({ name: 'default', fullPath: IS_PLATFORM ? '/workspace' : '' }); const [claudeAuthStatus, setClaudeAuthStatus] = useState({ authenticated: false, diff --git a/src/components/ProtectedRoute.jsx b/src/components/ProtectedRoute.jsx index 56343ae..507efa8 100644 --- a/src/components/ProtectedRoute.jsx +++ b/src/components/ProtectedRoute.jsx @@ -4,6 +4,7 @@ import SetupForm from './SetupForm'; import LoginForm from './LoginForm'; import Onboarding from './Onboarding'; import { MessageSquare } from 'lucide-react'; +import { IS_PLATFORM } from '../constants/config'; const LoadingScreen = () => (

@@ -27,7 +28,7 @@ const LoadingScreen = () => ( const ProtectedRoute = ({ children }) => { const { user, isLoading, needsSetup, hasCompletedOnboarding, refreshOnboardingStatus } = useAuth(); - if (import.meta.env.VITE_IS_PLATFORM === 'true') { + if (IS_PLATFORM) { if (isLoading) { return ; } diff --git a/src/components/Shell.jsx b/src/components/Shell.jsx index 59810ad..0bf78fd 100644 --- a/src/components/Shell.jsx +++ b/src/components/Shell.jsx @@ -5,6 +5,7 @@ import { WebglAddon } from '@xterm/addon-webgl'; import { WebLinksAddon } from '@xterm/addon-web-links'; import '@xterm/xterm/css/xterm.css'; import { useTranslation } from 'react-i18next'; +import { IS_PLATFORM } from '../constants/config'; const xtermStyles = ` .xterm .xterm-screen { @@ -55,10 +56,9 @@ function Shell({ selectedProject, selectedSession, initialCommand, isPlainShell if (isConnecting || isConnected) return; try { - const isPlatform = import.meta.env.VITE_IS_PLATFORM === 'true'; let wsUrl; - if (isPlatform) { + if (IS_PLATFORM) { const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; wsUrl = `${protocol}//${window.location.host}/shell`; } else { diff --git a/src/components/Sidebar.jsx b/src/components/Sidebar.jsx index ef74465..a495cfc 100644 --- a/src/components/Sidebar.jsx +++ b/src/components/Sidebar.jsx @@ -16,6 +16,7 @@ import ProjectCreationWizard from './ProjectCreationWizard'; import { api } from '../utils/api'; import { useTaskMaster } from '../contexts/TaskMasterContext'; import { useTasksSettings } from '../contexts/TasksSettingsContext'; +import { IS_PLATFORM } from '../constants/config'; // Move formatTimeAgo outside component to avoid recreation on every render const formatTimeAgo = (dateString, currentTime, t) => { @@ -622,7 +623,7 @@ function Sidebar({
{/* Desktop Header */}
- {import.meta.env.VITE_IS_PLATFORM === 'true' ? ( + {IS_PLATFORM ? (
- {import.meta.env.VITE_IS_PLATFORM === 'true' ? ( + {IS_PLATFORM ? ( { const [error, setError] = useState(null); useEffect(() => { - if (import.meta.env.VITE_IS_PLATFORM === 'true') { + if (IS_PLATFORM) { setUser({ username: 'platform-user' }); setNeedsSetup(false); checkOnboardingStatus(); diff --git a/src/contexts/TaskMasterContext.jsx b/src/contexts/TaskMasterContext.jsx index e1e4610..77e996f 100644 --- a/src/contexts/TaskMasterContext.jsx +++ b/src/contexts/TaskMasterContext.jsx @@ -1,7 +1,7 @@ import React, { createContext, useContext, useEffect, useState, useCallback } from 'react'; import { api } from '../utils/api'; import { useAuth } from './AuthContext'; -import { useWebSocketContext } from './WebSocketContext'; +import { useWebSocket } from './WebSocketContext'; const TaskMasterContext = createContext({ // TaskMaster project state @@ -42,7 +42,7 @@ export const useTaskMaster = () => { export const TaskMasterProvider = ({ children }) => { // Get WebSocket messages from shared context to avoid duplicate connections - const { messages } = useWebSocketContext(); + const { latestMessage } = useWebSocket(); // Authentication context const { user, token, isLoading: authLoading } = useAuth(); @@ -238,9 +238,8 @@ export const TaskMasterProvider = ({ children }) => { } }, [currentProject?.name, user, token, refreshTasks]); - // Handle WebSocket messages for TaskMaster updates + // Handle WebSocket latestMessage for TaskMaster updates useEffect(() => { - const latestMessage = messages[messages.length - 1]; if (!latestMessage) return; @@ -268,7 +267,7 @@ export const TaskMasterProvider = ({ children }) => { // Ignore non-TaskMaster messages break; } - }, [messages, refreshProjects, refreshTasks, refreshMCPStatus, currentProject]); + }, [latestMessage, refreshProjects, refreshTasks, refreshMCPStatus, currentProject]); // Context value const contextValue = { diff --git a/src/contexts/WebSocketContext.jsx b/src/contexts/WebSocketContext.jsx deleted file mode 100644 index 5c9fab5..0000000 --- a/src/contexts/WebSocketContext.jsx +++ /dev/null @@ -1,29 +0,0 @@ -import React, { createContext, useContext } from 'react'; -import { useWebSocket } from '../utils/websocket'; - -const WebSocketContext = createContext({ - ws: null, - sendMessage: () => {}, - messages: [], - isConnected: false -}); - -export const useWebSocketContext = () => { - const context = useContext(WebSocketContext); - if (!context) { - throw new Error('useWebSocketContext must be used within a WebSocketProvider'); - } - return context; -}; - -export const WebSocketProvider = ({ children }) => { - const webSocketData = useWebSocket(); - - return ( - - {children} - - ); -}; - -export default WebSocketContext; \ No newline at end of file diff --git a/src/contexts/WebSocketContext.tsx b/src/contexts/WebSocketContext.tsx new file mode 100644 index 0000000..ad31f66 --- /dev/null +++ b/src/contexts/WebSocketContext.tsx @@ -0,0 +1,125 @@ +import { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react'; +import { useAuth } from './AuthContext'; +import { IS_PLATFORM } from '../constants/config'; + +type WebSocketContextType = { + ws: WebSocket | null; + sendMessage: (message: any) => void; + latestMessage: any | null; + isConnected: boolean; +}; + +const WebSocketContext = createContext(null); + +export const useWebSocket = () => { + const context = useContext(WebSocketContext); + if (!context) { + throw new Error('useWebSocket must be used within a WebSocketProvider'); + } + return context; +}; + +const buildWebSocketUrl = (token: string | null) => { + const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; + if (IS_PLATFORM) return `${protocol}//${window.location.host}/ws`; // Platform mode: Use same domain as the page (goes through proxy) + if (!token) return null; + return `${protocol}//${window.location.host}/ws?token=${encodeURIComponent(token)}`; // OSS mode: Use same host:port that served the page +}; + +const useWebSocketProviderState = (): WebSocketContextType => { + const wsRef = useRef(null); + const unmountedRef = useRef(false); // Track if component is unmounted + const [latestMessage, setLatestMessage] = useState(null); + const [isConnected, setIsConnected] = useState(false); + const reconnectTimeoutRef = useRef(null); + const { token } = useAuth(); + + useEffect(() => { + connect(); + + return () => { + unmountedRef.current = true; + if (reconnectTimeoutRef.current) { + clearTimeout(reconnectTimeoutRef.current); + } + if (wsRef.current) { + wsRef.current.close(); + } + }; + }, [token]); // everytime token changes, we reconnect + + const connect = useCallback(() => { + if (unmountedRef.current) return; // Prevent connection if unmounted + try { + // Construct WebSocket URL + const wsUrl = buildWebSocketUrl(token); + + if (!wsUrl) return console.warn('No authentication token found for WebSocket connection'); + + const websocket = new WebSocket(wsUrl); + + websocket.onopen = () => { + setIsConnected(true); + wsRef.current = websocket; + }; + + websocket.onmessage = (event) => { + try { + const data = JSON.parse(event.data); + setLatestMessage(data); + } catch (error) { + console.error('Error parsing WebSocket message:', error); + } + }; + + websocket.onclose = () => { + setIsConnected(false); + wsRef.current = null; + + // Attempt to reconnect after 3 seconds + reconnectTimeoutRef.current = setTimeout(() => { + if (unmountedRef.current) return; // Prevent reconnection if unmounted + connect(); + }, 3000); + }; + + websocket.onerror = (error) => { + console.error('WebSocket error:', error); + }; + + } catch (error) { + console.error('Error creating WebSocket connection:', error); + } + }, [token]); // everytime token changes, we reconnect + + const sendMessage = useCallback((message: any) => { + const socket = wsRef.current; + if (socket && isConnected) { + socket.send(JSON.stringify(message)); + } else { + console.warn('WebSocket not connected'); + } + }, [isConnected]); + + const value: WebSocketContextType = useMemo(() => + ({ + ws: wsRef.current, + sendMessage, + latestMessage, + isConnected + }), [sendMessage, latestMessage, isConnected]); + + return value; +}; + +export const WebSocketProvider = ({ children }: { children: React.ReactNode }) => { + const webSocketData = useWebSocketProviderState(); + + return ( + + {children} + + ); +}; + +export default WebSocketContext; diff --git a/src/utils/api.js b/src/utils/api.js index 1625546..c423875 100644 --- a/src/utils/api.js +++ b/src/utils/api.js @@ -1,6 +1,7 @@ +import { IS_PLATFORM } from "../constants/config"; + // Utility function for authenticated API calls export const authenticatedFetch = (url, options = {}) => { - const isPlatform = import.meta.env.VITE_IS_PLATFORM === 'true'; const token = localStorage.getItem('auth-token'); const defaultHeaders = {}; @@ -10,7 +11,7 @@ export const authenticatedFetch = (url, options = {}) => { defaultHeaders['Content-Type'] = 'application/json'; } - if (!isPlatform && token) { + if (!IS_PLATFORM && token) { defaultHeaders['Authorization'] = `Bearer ${token}`; } diff --git a/src/utils/websocket.js b/src/utils/websocket.js deleted file mode 100755 index 553eaf2..0000000 --- a/src/utils/websocket.js +++ /dev/null @@ -1,94 +0,0 @@ -import { useState, useEffect, useRef } from 'react'; - -export function useWebSocket() { - const [ws, setWs] = useState(null); - const [messages, setMessages] = useState([]); - const [isConnected, setIsConnected] = useState(false); - const reconnectTimeoutRef = useRef(null); - - useEffect(() => { - connect(); - - return () => { - if (reconnectTimeoutRef.current) { - clearTimeout(reconnectTimeoutRef.current); - } - if (ws) { - ws.close(); - } - }; - }, []); // Keep dependency array but add proper cleanup - - const connect = async () => { - try { - const isPlatform = import.meta.env.VITE_IS_PLATFORM === 'true'; - - // Construct WebSocket URL - let wsUrl; - - if (isPlatform) { - // Platform mode: Use same domain as the page (goes through proxy) - const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - wsUrl = `${protocol}//${window.location.host}/ws`; - } else { - // OSS mode: Connect to same host:port that served the page - const token = localStorage.getItem('auth-token'); - if (!token) { - console.warn('No authentication token found for WebSocket connection'); - return; - } - - const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - wsUrl = `${protocol}//${window.location.host}/ws?token=${encodeURIComponent(token)}`; - } - - const websocket = new WebSocket(wsUrl); - - websocket.onopen = () => { - setIsConnected(true); - setWs(websocket); - }; - - websocket.onmessage = (event) => { - try { - const data = JSON.parse(event.data); - setMessages(prev => [...prev, data]); - } catch (error) { - console.error('Error parsing WebSocket message:', error); - } - }; - - websocket.onclose = () => { - setIsConnected(false); - setWs(null); - - // Attempt to reconnect after 3 seconds - reconnectTimeoutRef.current = setTimeout(() => { - connect(); - }, 3000); - }; - - websocket.onerror = (error) => { - console.error('WebSocket error:', error); - }; - - } catch (error) { - console.error('Error creating WebSocket connection:', error); - } - }; - - const sendMessage = (message) => { - if (ws && isConnected) { - ws.send(JSON.stringify(message)); - } else { - console.warn('WebSocket not connected'); - } - }; - - return { - ws, - sendMessage, - messages, - isConnected - }; -}