websockets through vscode messaging for codespaces

2 files changed, 207 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() {
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index 66829836..4192595c 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -16,6 +16,7 @@ import {
 import { sendTelemetryEvent, TelemetryEvent } from "./telemetry";
 import { RangeInFile, SerializedDebugContext } from "./client";
 import { addFileSystemToDebugContext } from "./util/util";
+const WebSocket = require("ws");
 
 class StreamManager {
   private _fullText: string = "";
@@ -87,6 +88,46 @@ class StreamManager {
   }
 }
 
+let websocketConnections: { [url: string]: WebsocketConnection | undefined } =
+  {};
+
+class WebsocketConnection {
+  private _ws: WebSocket;
+  private _onMessage: (message: string) => void;
+  private _onOpen: () => void;
+  private _onClose: () => void;
+
+  constructor(
+    url: string,
+    onMessage: (message: string) => void,
+    onOpen: () => void,
+    onClose: () => void
+  ) {
+    this._ws = new WebSocket(url);
+    this._onMessage = onMessage;
+    this._onOpen = onOpen;
+    this._onClose = onClose;
+
+    this._ws.onmessage = (event) => {
+      this._onMessage(event.data);
+    };
+    this._ws.onclose = () => {
+      this._onClose();
+    };
+    this._ws.onopen = () => {
+      this._onOpen();
+    };
+  }
+
+  public send(message: string) {
+    this._ws.send(message);
+  }
+
+  public close() {
+    this._ws.close();
+  }
+}
+
 let streamManager = new StreamManager();
 
 export let debugPanelWebview: vscode.Webview | undefined;
@@ -147,6 +188,39 @@ export function setupDebugPanel(
     });
   });
 
+  async function connectWebsocket(url: string) {
+    return new Promise((resolve, reject) => {
+      const onMessage = (message: any) => {
+        panel.webview.postMessage({
+          type: "websocketForwardingMessage",
+          url,
+          data: message,
+        });
+      };
+      const onOpen = () => {
+        panel.webview.postMessage({
+          type: "websocketForwardingOpen",
+          url,
+        });
+        resolve(null);
+      };
+      const onClose = () => {
+        websocketConnections[url] = undefined;
+        panel.webview.postMessage({
+          type: "websocketForwardingClose",
+          url,
+        });
+      };
+      const connection = new WebsocketConnection(
+        url,
+        onMessage,
+        onOpen,
+        onClose
+      );
+      websocketConnections[url] = connection;
+    });
+  }
+
   panel.webview.onDidReceiveMessage(async (data) => {
     switch (data.type) {
       case "onLoad": {
@@ -158,6 +232,27 @@ export function setupDebugPanel(
         });
         break;
       }
+
+      case "websocketForwardingOpen": {
+        let url = data.url;
+        if (typeof websocketConnections[url] === "undefined") {
+          await connectWebsocket(url);
+        }
+        break;
+      }
+      case "websocketForwardingMessage": {
+        let url = data.url;
+        let connection = websocketConnections[url];
+        if (typeof connection === "undefined") {
+          await connectWebsocket(url);
+        }
+        connection = websocketConnections[url];
+        if (typeof connection === "undefined") {
+          throw new Error("Failed to connect websocket in VS Code Extension");
+        }
+        connection.send(data.message);
+        break;
+      }
       case "listTenThings": {
         sendTelemetryEvent(TelemetryEvent.GenerateIdeas);
         let resp = await debugApi.listtenDebugListPost({