copying from old repo

1 files changed, 67 insertions, 0 deletions

diff --git a/extension/react-app/src/hooks/useWebsocket.ts b/extension/react-app/src/hooks/useWebsocket.ts
new file mode 100644
index 00000000..147172bd
--- /dev/null
+++ b/extension/react-app/src/hooks/useWebsocket.ts
@@ -0,0 +1,67 @@
+import React, { useEffect, useState } from "react";
+import { RootStore } from "../redux/store";
+import { useSelector } from "react-redux";
+
+function useContinueWebsocket(
+  serverUrl: string,
+  onMessage: (message: { data: any }) => void
+) {
+  const sessionId = useSelector((state: RootStore) => state.config.sessionId);
+  const [websocket, setWebsocket] = useState<WebSocket | undefined>(undefined);
+
+  async function connect() {
+    while (!sessionId) {
+      await new Promise((resolve) => setTimeout(resolve, 300));
+    }
+
+    console.log("Creating websocket", sessionId);
+
+    const wsUrl =
+      serverUrl.replace("http", "ws") +
+      "/notebook/ws?session_id=" +
+      encodeURIComponent(sessionId);
+
+    const ws = new WebSocket(wsUrl);
+    setWebsocket(ws);
+
+    // Set up callbacks
+    ws.onopen = () => {
+      console.log("Websocket opened");
+      ws.send(JSON.stringify({ sessionId }));
+    };
+
+    ws.onmessage = (msg) => {
+      onMessage(msg);
+      console.log("Got message", msg);
+    };
+
+    ws.onclose = (msg) => {
+      console.log("Websocket closed");
+      setWebsocket(undefined);
+    };
+
+    return ws;
+  }
+
+  async function getConnection() {
+    if (!websocket) {
+      return await connect();
+    }
+    return websocket;
+  }
+
+  async function send(message: object) {
+    let ws = await getConnection();
+    ws.send(JSON.stringify(message));
+  }
+
+  useEffect(() => {
+    if (!sessionId) {
+      return;
+    }
+    connect();
+  }, [sessionId]);
+
+  return { send };
+}
+export default useContinueWebsocket;