Newer
Older
cortex-hub / ui / client-app / src / services / websocket.js
// src/services/websocket.js

import { createSession } from "./apiService";

/**
 * Connects to the WebSocket server, establishing a new session first.
 * @param {function(Object): void} onMessageCallback - Callback for incoming messages.
 * @param {function(): void} onOpenCallback - Callback when the connection is opened.
 * @param {function(): void} onCloseCallback - Callback when the connection is closed.
 * @param {function(Error): void} onErrorCallback - Callback when an error occurs.
 * @returns {Promise<{ws: WebSocket, sessionId: string}>} The WebSocket instance and the session ID.
 */
export const connectToWebSocket = async (
  onMessageCallback,
  onOpenCallback,
  onCloseCallback,
  onErrorCallback
) => {
  try {
    let sessionId = localStorage.getItem("sessionId");

    if (!sessionId) {
      // No existing session, so create one via API
      const session = await createSession();
      sessionId = session.id;
    
      // Store it in localStorage for reuse
      localStorage.setItem("sessionId", sessionId);
    }
    
    // You now have a valid sessionId, either reused or newly created
    console.log("Using session ID:", sessionId);

    // Then, use the session ID to open the WebSocket connection
    const websocketUrl = `ws://localhost:8001/ws/workspace/${sessionId}`;
    const ws = new WebSocket(websocketUrl);

    ws.onopen = () => {
      console.log("WebSocket connected");
      if (onOpenCallback) onOpenCallback();
    };

    ws.onmessage = (event) => {
      const message = JSON.parse(event.data);
      if (onMessageCallback) onMessageCallback(message);
    };

    ws.onclose = () => {
      console.log("WebSocket disconnected");
      if (onCloseCallback) onCloseCallback();
    };

    ws.onerror = (error) => {
      console.error("WebSocket error:", error);
      if (onErrorCallback) onErrorCallback(error);
    };

    return { ws, sessionId };

  } catch (error) {
    console.error("Failed to create session or connect to WebSocket:", error);
    if (onErrorCallback) onErrorCallback(error);
    throw error;
  }
};