From a2f41064840e9427ac46218f02e8adf5b8f46b54 Mon Sep 17 00:00:00 2001
From: Rohit <rohit.rajan031101@gmail.com>
Date: Thu, 14 Aug 2025 16:03:54 +0530
Subject: [PATCH] feat: add queue socket conn logic

---
 src/context/socket.tsx | 68 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 66 insertions(+), 2 deletions(-)

diff --git a/src/context/socket.tsx b/src/context/socket.tsx
index 4d9c95e1..05a8c989 100644
--- a/src/context/socket.tsx
+++ b/src/context/socket.tsx
@@ -1,4 +1,4 @@
-import React, { createContext, useCallback, useContext, useMemo, useState } from 'react';
+import React, { createContext, useCallback, useContext, useState, useRef, useEffect } from 'react';
 import { io, Socket } from 'socket.io-client';
 import { apiUrl } from "../apiConfig";
 
@@ -6,12 +6,16 @@ const SERVER_ENDPOINT = apiUrl;
 
 interface SocketState {
   socket: Socket | null;
+  queueSocket: Socket | null;
   id: string;
   setId: (id: string) => void;
+  connectToQueueSocket: (userId: string, onRunCompleted?: (data: any) => void) => void;
+  disconnectQueueSocket: () => void;
 };
 
 class SocketStore implements Partial<SocketState> {
-  socket = null;
+  socket: Socket | null = null;
+  queueSocket: Socket | null = null;
   id = '';
 };
 
@@ -22,7 +26,9 @@ export const useSocketStore = () => useContext(socketStoreContext);
 
 export const SocketProvider = ({ children }: { children: JSX.Element }) => {
   const [socket, setSocket] = useState<Socket | null>(socketStore.socket);
+  const [queueSocket, setQueueSocket] = useState<Socket | null>(socketStore.queueSocket);
   const [id, setActiveId] = useState<string>(socketStore.id);
+  const runCompletedCallbackRef = useRef<((data: any) => void) | null>(null);
 
   const setId = useCallback((id: string) => {
     // the socket client connection is recomputed whenever id changes -> the new browser has been initialized
@@ -39,12 +45,70 @@ export const SocketProvider = ({ children }: { children: JSX.Element }) => {
     setActiveId(id);
   }, [setSocket]);
 
+  const connectToQueueSocket = useCallback((userId: string, onRunCompleted?: (data: any) => void) => {
+    runCompletedCallbackRef.current = onRunCompleted || null;
+
+    const newQueueSocket = io(`${SERVER_ENDPOINT}/queued-run`, {
+      transports: ["websocket"],
+      rejectUnauthorized: false,
+      query: { userId }
+    });
+
+    newQueueSocket.on('connect', () => {
+      console.log('Queue socket connected for user:', userId);
+    });
+
+    newQueueSocket.on('connect_error', (error) => {
+      console.log('Queue socket connection error:', error);
+    });
+
+    newQueueSocket.on('run-completed', (completionData) => {
+      console.log('Run completed event received:', completionData);
+      if (runCompletedCallbackRef.current) {
+        runCompletedCallbackRef.current(completionData);
+      }
+    });
+
+    setQueueSocket(currentSocket => {
+      if (currentSocket) {
+        currentSocket.disconnect();
+      }
+      return newQueueSocket;
+    });
+
+    socketStore.queueSocket = newQueueSocket;
+  }, []);
+
+  const disconnectQueueSocket = useCallback(() => {
+    setQueueSocket(currentSocket => {
+      if (currentSocket) {
+        currentSocket.disconnect();
+      }
+      return null;
+    });
+
+    socketStore.queueSocket = null;
+    runCompletedCallbackRef.current = null;
+  }, []);
+
+  // Cleanup on unmount
+  useEffect(() => {
+    return () => {
+      if (queueSocket) {
+        queueSocket.disconnect();
+      }
+    };
+  }, [queueSocket]);
+
   return (
     <socketStoreContext.Provider
       value={{
         socket,
+        queueSocket,
         id,
         setId,
+        connectToQueueSocket,
+        disconnectQueueSocket,
       }}
     >
       {children}