5 files changed, 268 insertions, 55 deletions
diff --git a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
new file mode 100644
index 00000000..18a91de7
--- /dev/null
+++ b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
@@ -0,0 +1,13 @@
+abstract class AbstractContinueGUIClientProtocol {
+  abstract sendMainInput(input: string): void;
+
+  abstract reverseToIndex(index: number): void;
+
+  abstract sendRefinementInput(input: string, index: number): void;
+
+  abstract sendStepUserInput(input: string, index: number): void;
+
+  abstract onStateUpdate(state: any): void;
+}
+
+export default AbstractContinueGUIClientProtocol;
diff --git a/extension/react-app/src/hooks/messenger.ts b/extension/react-app/src/hooks/messenger.ts
new file mode 100644
index 00000000..e2a0bab8
--- /dev/null
+++ b/extension/react-app/src/hooks/messenger.ts
@@ -0,0 +1,108 @@
+// console.log("Websocket import");
+// const WebSocket = require("ws");
+
+export abstract class Messenger {
+  abstract send(messageType: string, data: object): void;
+
+  abstract onMessageType(
+    messageType: string,
+    callback: (data: object) => void
+  ): void;
+
+  abstract onMessage(callback: (messageType: string, data: any) => void): void;
+
+  abstract onOpen(callback: () => void): void;
+
+  abstract onClose(callback: () => void): void;
+
+  abstract sendAndReceive(messageType: string, data: any): Promise<any>;
+}
+
+export class WebsocketMessenger extends Messenger {
+  websocket: WebSocket;
+  private onMessageListeners: {
+    [messageType: string]: ((data: object) => void)[];
+  } = {};
+  private onOpenListeners: (() => void)[] = [];
+  private onCloseListeners: (() => void)[] = [];
+  private serverUrl: string;
+
+  _newWebsocket(): WebSocket {
+    // // Dynamic import, because WebSocket is builtin with browser, but not with node. And can't use require in browser.
+    // if (typeof process === "object") {
+    //   console.log("Using node");
+    //   // process is only available in Node
+    //   var WebSocket = require("ws");
+    // }
+
+    const newWebsocket = new WebSocket(this.serverUrl);
+    for (const listener of this.onOpenListeners) {
+      this.onOpen(listener);
+    }
+    for (const listener of this.onCloseListeners) {
+      this.onClose(listener);
+    }
+    for (const messageType in this.onMessageListeners) {
+      for (const listener of this.onMessageListeners[messageType]) {
+        this.onMessageType(messageType, listener);
+      }
+    }
+    return newWebsocket;
+  }
+
+  constructor(serverUrl: string) {
+    super();
+    this.serverUrl = serverUrl;
+    this.websocket = this._newWebsocket();
+  }
+
+  send(messageType: string, data: object) {
+    const payload = JSON.stringify({ messageType, data });
+    if (this.websocket.readyState === this.websocket.OPEN) {
+      this.websocket.send(payload);
+    } else {
+      if (this.websocket.readyState !== this.websocket.CONNECTING) {
+        this.websocket = this._newWebsocket();
+      }
+      this.websocket.addEventListener("open", () => {
+        this.websocket.send(payload);
+      });
+    }
+  }
+
+  sendAndReceive(messageType: string, data: any): Promise<any> {
+    return new Promise((resolve, reject) => {
+      const eventListener = (data: any) => {
+        // THIS ISN"T GETTING CALLED
+        resolve(data);
+        this.websocket.removeEventListener("message", eventListener);
+      };
+      this.onMessageType(messageType, eventListener);
+      this.send(messageType, data);
+    });
+  }
+
+  onMessageType(messageType: string, callback: (data: any) => void): void {
+    this.websocket.addEventListener("message", (event: any) => {
+      const msg = JSON.parse(event.data);
+      if (msg.messageType === messageType) {
+        callback(msg.data);
+      }
+    });
+  }
+
+  onMessage(callback: (messageType: string, data: any) => void): void {
+    this.websocket.addEventListener("message", (event) => {
+      const msg = JSON.parse(event.data);
+      callback(msg.messageType, msg.data);
+    });
+  }
+
+  onOpen(callback: () => void): void {
+    this.websocket.addEventListener("open", callback);
+  }
+
+  onClose(callback: () => void): void {
+    this.websocket.addEventListener("close", callback);
+  }
+}
diff --git a/extension/react-app/src/hooks/useContinueGUIProtocol.ts b/extension/react-app/src/hooks/useContinueGUIProtocol.ts
new file mode 100644
index 00000000..a3a1d0c9
--- /dev/null
+++ b/extension/react-app/src/hooks/useContinueGUIProtocol.ts
@@ -0,0 +1,49 @@
+import AbstractContinueGUIClientProtocol from "./ContinueGUIClientProtocol";
+// import { Messenger, WebsocketMessenger } from "../../../src/util/messenger";
+import { Messenger, WebsocketMessenger } from "./messenger";
+import { VscodeMessenger } from "./vscodeMessenger";
+
+class ContinueGUIClientProtocol extends AbstractContinueGUIClientProtocol {
+  messenger: Messenger;
+  // Server URL must contain the session ID param
+  serverUrlWithSessionId: string;
+
+  constructor(
+    serverUrlWithSessionId: string,
+    useVscodeMessagePassing: boolean
+  ) {
+    super();
+    this.serverUrlWithSessionId = serverUrlWithSessionId;
+    if (useVscodeMessagePassing) {
+      this.messenger = new VscodeMessenger(serverUrlWithSessionId);
+    } else {
+      this.messenger = new WebsocketMessenger(serverUrlWithSessionId);
+    }
+  }
+
+  sendMainInput(input: string) {
+    this.messenger.send("main_input", { input });
+  }
+
+  reverseToIndex(index: number) {
+    this.messenger.send("reverse_to_index", { index });
+  }
+
+  sendRefinementInput(input: string, index: number) {
+    this.messenger.send("refinement_input", { input, index });
+  }
+
+  sendStepUserInput(input: string, index: number) {
+    this.messenger.send("step_user_input", { input, index });
+  }
+
+  onStateUpdate(callback: (state: any) => void) {
+    this.messenger.onMessageType("state_update", (data: any) => {
+      if (data.state) {
+        callback(data.state);
+      }
+    });
+  }
+}
+
+export default ContinueGUIClientProtocol;
diff --git a/extension/react-app/src/hooks/useWebsocket.ts b/extension/react-app/src/hooks/useWebsocket.ts
index 147172bd..e762666f 100644
--- a/extension/react-app/src/hooks/useWebsocket.ts
+++ b/extension/react-app/src/hooks/useWebsocket.ts
@@ -1,67 +1,39 @@
 import React, { useEffect, useState } from "react";
 import { RootStore } from "../redux/store";
 import { useSelector } from "react-redux";
+import ContinueGUIClientProtocol from "./useContinueGUIProtocol";
+import { postVscMessage } from "../vscode";
 
-function useContinueWebsocket(
-  serverUrl: string,
-  onMessage: (message: { data: any }) => void
-) {
+function useContinueGUIProtocol(useVscodeMessagePassing: boolean = true) {
   const sessionId = useSelector((state: RootStore) => state.config.sessionId);
-  const [websocket, setWebsocket] = useState<WebSocket | undefined>(undefined);
+  const serverHttpUrl = useSelector((state: RootStore) => state.config.apiUrl);
+  const [client, setClient] = useState<ContinueGUIClientProtocol | undefined>(
+    undefined
+  );
 
-  async function connect() {
-    while (!sessionId) {
-      await new Promise((resolve) => setTimeout(resolve, 300));
+  useEffect(() => {
+    if (!sessionId || !serverHttpUrl) {
+      if (useVscodeMessagePassing) {
+        postVscMessage("onLoad", {});
+      }
+      setClient(undefined);
+      return;
     }
 
-    console.log("Creating websocket", sessionId);
-
-    const wsUrl =
-      serverUrl.replace("http", "ws") +
-      "/notebook/ws?session_id=" +
+    const serverUrlWithSessionId =
+      serverHttpUrl.replace("http", "ws") +
+      "/gui/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]);
+    console.log("Creating websocket", serverUrlWithSessionId);
+    console.log("Using vscode message passing", useVscodeMessagePassing);
+    const newClient = new ContinueGUIClientProtocol(
+      serverUrlWithSessionId,
+      useVscodeMessagePassing
+    );
+    setClient(newClient);
+  }, [sessionId, serverHttpUrl]);
 
-  return { send };
+  return client;
 }
-export default useContinueWebsocket;
+export default useContinueGUIProtocol;
diff --git a/extension/react-app/src/hooks/vscodeMessenger.ts b/extension/react-app/src/hooks/vscodeMessenger.ts
new file mode 100644
index 00000000..ba19586b
--- /dev/null
+++ b/extension/react-app/src/hooks/vscodeMessenger.ts
@@ -0,0 +1,71 @@
+import { postVscMessage } from "../vscode";
+// import { Messenger } from "../../../src/util/messenger";
+import { Messenger } from "./messenger";
+
+export class VscodeMessenger extends Messenger {
+  serverUrl: string;
+
+  constructor(serverUrl: string) {
+    super();
+    this.serverUrl = serverUrl;
+    postVscMessage("websocketForwardingOpen", { url: this.serverUrl });
+  }
+
+  send(messageType: string, data: object) {
+    postVscMessage("websocketForwardingMessage", {
+      message: { messageType, data },
+      url: this.serverUrl,
+    });
+  }
+
+  onMessageType(messageType: string, callback: (data: object) => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingMessage") {
+        const data = JSON.parse(event.data.data);
+        if (data.messageType === messageType) {
+          callback(data.data);
+        }
+      }
+    });
+  }
+
+  onMessage(callback: (messageType: string, data: any) => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingMessage") {
+        const data = JSON.parse(event.data.data);
+        callback(data.messageType, data.data);
+      }
+    });
+  }
+
+  sendAndReceive(messageType: string, data: any): Promise<any> {
+    return new Promise((resolve) => {
+      const handler = (event: any) => {
+        if (event.data.type === "websocketForwardingMessage") {
+          const data = JSON.parse(event.data.data);
+          if (data.messageType === messageType) {
+            window.removeEventListener("message", handler);
+            resolve(data.data);
+          }
+        }
+      };
+      window.addEventListener("message", handler);
+      this.send(messageType, data);
+    });
+  }
+
+  onOpen(callback: () => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingOpen") {
+        callback();
+      }
+    });
+  }
+  onClose(callback: () => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingClose") {
+        callback();
+      }
+    });
+  }
+}