1 files changed, 112 insertions, 21 deletions
diff --git a/extension/react-app/src/hooks/useWebsocket.ts b/extension/react-app/src/hooks/useWebsocket.ts
index 147172bd..6e8e68fa 100644
--- a/extension/react-app/src/hooks/useWebsocket.ts
+++ b/extension/react-app/src/hooks/useWebsocket.ts
@@ -1,13 +1,100 @@
 import React, { useEffect, useState } from "react";
 import { RootStore } from "../redux/store";
 import { useSelector } from "react-redux";
+import { postVscMessage } from "../vscode";
+
+abstract class Messenger {
+  abstract send(data: string): void;
+}
+
+class VscodeMessenger extends Messenger {
+  url: string;
+
+  constructor(
+    url: string,
+    onMessage: (message: { data: any }) => void,
+    onOpen: (messenger: Messenger) => void,
+    onClose: (messenger: Messenger) => void
+  ) {
+    super();
+    this.url = url;
+    window.addEventListener("message", (event: any) => {
+      switch (event.data.type) {
+        case "websocketForwardingMessage":
+          onMessage(event.data);
+          break;
+        case "websocketForwardingOpen":
+          onOpen(this);
+          break;
+        case "websocketForwardingClose":
+          onClose(this);
+          break;
+      }
+    });
+
+    postVscMessage("websocketForwardingOpen", { url: this.url });
+  }
+
+  send(data: string) {
+    postVscMessage("websocketForwardingMessage", {
+      message: data,
+      url: this.url,
+    });
+  }
+}
+
+class WebsocketMessenger extends Messenger {
+  websocket: WebSocket;
+  constructor(
+    websocket: WebSocket,
+    onMessage: (message: { data: any }) => void,
+    onOpen: (messenger: Messenger) => void,
+    onClose: (messenger: Messenger) => void
+  ) {
+    super();
+    this.websocket = websocket;
+
+    websocket.addEventListener("close", () => {
+      onClose(this);
+    });
+
+    websocket.addEventListener("open", () => {
+      onOpen(this);
+    });
+
+    websocket.addEventListener("message", (event) => {
+      onMessage(event.data);
+    });
+  }
+
+  static async connect(
+    url: string,
+    sessionId: string,
+    onMessage: (message: { data: any }) => void,
+    onOpen: (messenger: Messenger) => void,
+    onClose: (messenger: Messenger) => void
+  ): Promise<WebsocketMessenger> {
+    const ws = new WebSocket(url);
+
+    return new Promise((resolve, reject) => {
+      ws.addEventListener("open", () => {
+        resolve(new WebsocketMessenger(ws, onMessage, onOpen, onClose));
+      });
+    });
+  }
+
+  send(data: string) {
+    this.websocket.send(JSON.stringify(data));
+  }
+}
 
 function useContinueWebsocket(
   serverUrl: string,
-  onMessage: (message: { data: any }) => void
+  onMessage: (message: { data: any }) => void,
+  useVscodeMessagePassing: boolean = true
 ) {
   const sessionId = useSelector((state: RootStore) => state.config.sessionId);
-  const [websocket, setWebsocket] = useState<WebSocket | undefined>(undefined);
+  const [websocket, setWebsocket] = useState<Messenger | undefined>(undefined);
 
   async function connect() {
     while (!sessionId) {
@@ -15,32 +102,36 @@ function useContinueWebsocket(
     }
 
     console.log("Creating websocket", sessionId);
+    console.log("Using vscode message passing", useVscodeMessagePassing);
 
-    const wsUrl =
-      serverUrl.replace("http", "ws") +
-      "/notebook/ws?session_id=" +
-      encodeURIComponent(sessionId);
-
-    const ws = new WebSocket(wsUrl);
-    setWebsocket(ws);
+    const onClose = (messenger: Messenger) => {
+      console.log("Websocket closed");
+      setWebsocket(undefined);
+    };
 
-    // Set up callbacks
-    ws.onopen = () => {
+    const onOpen = (messenger: Messenger) => {
       console.log("Websocket opened");
-      ws.send(JSON.stringify({ sessionId }));
+      messenger.send(JSON.stringify({ sessionId }));
     };
 
-    ws.onmessage = (msg) => {
-      onMessage(msg);
-      console.log("Got message", msg);
-    };
+    const url =
+      serverUrl.replace("http", "ws") +
+      "/notebook/ws?session_id=" +
+      encodeURIComponent(sessionId);
 
-    ws.onclose = (msg) => {
-      console.log("Websocket closed");
-      setWebsocket(undefined);
-    };
+    const messenger: Messenger = useVscodeMessagePassing
+      ? new VscodeMessenger(url, onMessage, onOpen, onClose)
+      : await WebsocketMessenger.connect(
+          url,
+          sessionId,
+          onMessage,
+          onOpen,
+          onClose
+        );
+
+    setWebsocket(messenger);
 
-    return ws;
+    return messenger;
   }
 
   async function getConnection() {