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

@@ -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';

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');

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);
+}

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) {

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) {

shared/modelConstants.js

@@ -62,4 +62,4 @@ export const CODEX_MODELS = {
   ],
 
   DEFAULT: 'gpt-5.2'
-};
+};

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)}

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('');

src/components/GitPanel.jsx

@@ -960,6 +960,7 @@ function GitPanel({ selectedProject, isMobile, onFileOpen }) {
           {gitStatus.details && (
             <p className="text-sm text-center leading-relaxed mb-6 max-w-md">{gitStatus.details}</p>
           )}
+          {/* // ! This can be a custom component that can be reused for " Tip: Create a new project..." as well */}
           <div className="p-4 bg-blue-50 dark:bg-blue-900/20 rounded-lg border border-blue-200 dark:border-blue-800 max-w-md">
             <p className="text-sm text-blue-700 dark:text-blue-300 text-center">
               <strong>Tip:</strong> Run <code className="bg-blue-100 dark:bg-blue-900 px-2 py-1 rounded font-mono text-xs">git init</code> in your project directory to initialize git source control.

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';
     }

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}

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,

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 = () => (
   <div className="min-h-screen bg-background flex items-center justify-center p-4">
@@ -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 <LoadingScreen />;
     }

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 {

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({
       <div className="md:p-4 md:border-b md:border-border">
         {/* Desktop Header */}
         <div className="hidden md:flex items-center justify-between">
-          {import.meta.env.VITE_IS_PLATFORM === 'true' ? (
+          {IS_PLATFORM ? (
             <a
               href="https://cloudcli.ai/dashboard"
               className="flex items-center gap-3 hover:opacity-80 transition-opacity group"
@@ -673,7 +674,7 @@ function Sidebar({
           style={isPWA && isMobile ? { paddingTop: '16px' } : {}}
         >
           <div className="flex items-center justify-between">
-            {import.meta.env.VITE_IS_PLATFORM === 'true' ? (
+            {IS_PLATFORM ? (
               <a
                 href="https://cloudcli.ai/dashboard"
                 className="flex items-center gap-3 active:opacity-70 transition-opacity"

src/constants/config.ts

@@ -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 = import.meta.env.VITE_IS_PLATFORM === 'true';

src/contexts/AuthContext.jsx

@@ -1,5 +1,6 @@
 import React, { createContext, useContext, useEffect, useState } from 'react';
 import { api } from '../utils/api';
+import { IS_PLATFORM } from '../constants/config';
 
 const AuthContext = createContext({
   user: null,
@@ -31,7 +32,7 @@ export const AuthProvider = ({ children }) => {
   const [error, setError] = useState(null);
 
   useEffect(() => {
-    if (import.meta.env.VITE_IS_PLATFORM === 'true') {
+    if (IS_PLATFORM) {
       setUser({ username: 'platform-user' });
       setNeedsSetup(false);
       checkOnboardingStatus();

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 = {

src/contexts/WebSocketContext.jsx

@@ -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 (
-    <WebSocketContext.Provider value={webSocketData}>
-      {children}
-    </WebSocketContext.Provider>
-  );
-};
-
-export default WebSocketContext;

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<WebSocketContextType | null>(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<WebSocket | null>(null);
+  const unmountedRef = useRef(false); // Track if component is unmounted
+  const [latestMessage, setLatestMessage] = useState<any>(null);
+  const [isConnected, setIsConnected] = useState(false);
+  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(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 (
+    <WebSocketContext.Provider value={webSocketData}>
+      {children}
+    </WebSocketContext.Provider>
+  );
+};
+
+export default WebSocketContext;

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}`;
   }
 

src/utils/websocket.js

@@ -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
-  };
-}