7 files changed, 290 insertions, 190 deletions

diff --git a/extension/react-app/src/hooks/ContinueNotebookClientProtocol.ts b/extension/react-app/src/hooks/ContinueNotebookClientProtocol.ts
new file mode 100644
index 00000000..75fd7373
--- /dev/null
+++ b/extension/react-app/src/hooks/ContinueNotebookClientProtocol.ts
@@ -0,0 +1,13 @@
+abstract class AbstractContinueNotebookClientProtocol {
+  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 AbstractContinueNotebookClientProtocol;
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/useContinueNotebookProtocol.ts b/extension/react-app/src/hooks/useContinueNotebookProtocol.ts
new file mode 100644
index 00000000..d5ffbf09
--- /dev/null
+++ b/extension/react-app/src/hooks/useContinueNotebookProtocol.ts
@@ -0,0 +1,49 @@
+import AbstractContinueNotebookClientProtocol from "./ContinueNotebookClientProtocol";
+// import { Messenger, WebsocketMessenger } from "../../../src/util/messenger";
+import { Messenger, WebsocketMessenger } from "./messenger";
+import { VscodeMessenger } from "./vscodeMessenger";
+
+class ContinueNotebookClientProtocol extends AbstractContinueNotebookClientProtocol {
+  messenger: Messenger;
+  // Server URL must contain the session ID param
+  serverUrlWithSessionId: string;
+
+  constructor(
+    serverUrlWithSessionId: string,
+    useVscodeMessagePassing: boolean = false
+  ) {
+    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 ContinueNotebookClientProtocol;
diff --git a/extension/react-app/src/hooks/useWebsocket.ts b/extension/react-app/src/hooks/useWebsocket.ts
index 6e8e68fa..b98be577 100644
--- a/extension/react-app/src/hooks/useWebsocket.ts
+++ b/extension/react-app/src/hooks/useWebsocket.ts
@@ -1,158 +1,39 @@
 import React, { useEffect, useState } from "react";
 import { RootStore } from "../redux/store";
 import { useSelector } from "react-redux";
+import ContinueNotebookClientProtocol from "./useContinueNotebookProtocol";
 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,
-  useVscodeMessagePassing: boolean = true
-) {
+function useContinueNotebookProtocol(useVscodeMessagePassing: boolean = false) {
   const sessionId = useSelector((state: RootStore) => state.config.sessionId);
-  const [websocket, setWebsocket] = useState<Messenger | undefined>(undefined);
+  const serverHttpUrl = useSelector((state: RootStore) => state.config.apiUrl);
+  const [client, setClient] = useState<
+    ContinueNotebookClientProtocol | 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);
-    console.log("Using vscode message passing", useVscodeMessagePassing);
-
-    const onClose = (messenger: Messenger) => {
-      console.log("Websocket closed");
-      setWebsocket(undefined);
-    };
-
-    const onOpen = (messenger: Messenger) => {
-      console.log("Websocket opened");
-      messenger.send(JSON.stringify({ sessionId }));
-    };
-
-    const url =
-      serverUrl.replace("http", "ws") +
+    const serverUrlWithSessionId =
+      serverHttpUrl.replace("http", "ws") +
       "/notebook/ws?session_id=" +
       encodeURIComponent(sessionId);
 
-    const messenger: Messenger = useVscodeMessagePassing
-      ? new VscodeMessenger(url, onMessage, onOpen, onClose)
-      : await WebsocketMessenger.connect(
-          url,
-          sessionId,
-          onMessage,
-          onOpen,
-          onClose
-        );
-
-    setWebsocket(messenger);
-
-    return messenger;
-  }
-
-  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 };
+    console.log("Creating websocket", serverUrlWithSessionId);
+    console.log("Using vscode message passing", useVscodeMessagePassing);
+    const newClient = new ContinueNotebookClientProtocol(
+      serverUrlWithSessionId,
+      useVscodeMessagePassing
+    );
+    setClient(newClient);
+  }, [sessionId, serverHttpUrl]);
+
+  return client;
 }
-export default useContinueWebsocket;
+export default useContinueNotebookProtocol;
diff --git a/extension/react-app/src/hooks/vscodeMessenger.ts b/extension/react-app/src/hooks/vscodeMessenger.ts
new file mode 100644
index 00000000..746c4302
--- /dev/null
+++ b/extension/react-app/src/hooks/vscodeMessenger.ts
@@ -0,0 +1,68 @@
+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") {
+        if (event.data.message.messageType === messageType) {
+          callback(event.data.message.data);
+        }
+      }
+    });
+  }
+
+  onMessage(callback: (messageType: string, data: any) => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingMessage") {
+        callback(event.data.message.messageType, event.data.message.data);
+      }
+    });
+  }
+
+  sendAndReceive(messageType: string, data: any): Promise<any> {
+    return new Promise((resolve) => {
+      const handler = (event: any) => {
+        if (event.data.type === "websocketForwardingMessage") {
+          if (event.data.message.messageType === messageType) {
+            window.removeEventListener("message", handler);
+            resolve(event.data.message.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();
+      }
+    });
+  }
+}
diff --git a/extension/react-app/src/tabs/notebook.tsx b/extension/react-app/src/tabs/notebook.tsx
index a9c69c5b..02c9ff31 100644
--- a/extension/react-app/src/tabs/notebook.tsx
+++ b/extension/react-app/src/tabs/notebook.tsx
@@ -14,6 +14,7 @@ import StepContainer from "../components/StepContainer";
 import { useSelector } from "react-redux";
 import { RootStore } from "../redux/store";
 import useContinueWebsocket from "../hooks/useWebsocket";
+import useContinueNotebookProtocol from "../hooks/useWebsocket";
 
 let TopNotebookDiv = styled.div`
   display: grid;
@@ -33,8 +34,6 @@ interface NotebookProps {
 }
 
 function Notebook(props: NotebookProps) {
-  const serverUrl = useSelector((state: RootStore) => state.config.apiUrl);
-
   const [waitingForSteps, setWaitingForSteps] = useState(false);
   const [userInputQueue, setUserInputQueue] = useState<string[]>([]);
   const [history, setHistory] = useState<History | undefined>();
@@ -157,30 +156,17 @@ function Notebook(props: NotebookProps) {
   // } as any
   // );
 
-  const { send: websocketSend } = useContinueWebsocket(serverUrl, (msg) => {
-    let data = JSON.parse(msg.data);
-    if (data.messageType === "state") {
-      setWaitingForSteps(data.state.active);
-      setHistory(data.state.history);
-      setUserInputQueue(data.state.user_input_queue);
-    }
-  });
+  const client = useContinueNotebookProtocol();
 
-  // useEffect(() => {
-  //   (async () => {
-  //     if (sessionId && props.firstObservation) {
-  //       let resp = await fetch(serverUrl + "/observation", {
-  //         method: "POST",
-  //         headers: new Headers({
-  //           "x-continue-session-id": sessionId,
-  //         }),
-  //         body: JSON.stringify({
-  //           observation: props.firstObservation,
-  //         }),
-  //       });
-  //     }
-  //   })();
-  // }, [props.firstObservation]);
+  useEffect(() => {
+    console.log("CLIENT ON STATE UPDATE: ", client, client?.onStateUpdate);
+    client?.onStateUpdate((state) => {
+      console.log("Received state update: ", state);
+      setWaitingForSteps(state.active);
+      setHistory(state.history);
+      setUserInputQueue(state.user_input_queue);
+    });
+  }, [client]);
 
   const mainTextInputRef = useRef<HTMLTextAreaElement>(null);
 
@@ -201,14 +187,12 @@ function Notebook(props: NotebookProps) {
 
   const onMainTextInput = () => {
     if (mainTextInputRef.current) {
-      let value = mainTextInputRef.current.value;
+      if (!client) return;
+      let input = mainTextInputRef.current.value;
       setWaitingForSteps(true);
-      websocketSend({
-        messageType: "main_input",
-        value: value,
-      });
+      client.sendMainInput(input);
       setUserInputQueue((queue) => {
-        return [...queue, value];
+        return [...queue, input];
       });
       mainTextInputRef.current.value = "";
       mainTextInputRef.current.style.height = "";
@@ -216,17 +200,20 @@ function Notebook(props: NotebookProps) {
   };
 
   const onStepUserInput = (input: string, index: number) => {
+    if (!client) return;
     console.log("Sending step user input", input, index);
-    websocketSend({
-      messageType: "step_user_input",
-      value: input,
-      index,
-    });
+    client.sendStepUserInput(input, index);
   };
 
   // const iterations = useSelector(selectIterations);
   return (
     <TopNotebookDiv>
+      {typeof client === "undefined" && (
+        <>
+          <Loader></Loader>
+          <p>Server disconnected</p>
+        </>
+      )}
       {history?.timeline.map((node: HistoryNode, index: number) => {
         return (
           <StepContainer
@@ -237,17 +224,10 @@ function Notebook(props: NotebookProps) {
             inFuture={index > history?.current_index}
             historyNode={node}
             onRefinement={(input: string) => {
-              websocketSend({
-                messageType: "refinement_input",
-                value: input,
-                index,
-              });
+              client?.sendRefinementInput(input, index);
             }}
             onReverse={() => {
-              websocketSend({
-                messageType: "reverse",
-                index,
-              });
+              client?.reverseToIndex(index);
             }}
           />
         );
diff --git a/extension/react-app/src/vscode/index.ts b/extension/react-app/src/vscode/index.ts
index 7e373cd9..0785aa4d 100644
--- a/extension/react-app/src/vscode/index.ts
+++ b/extension/react-app/src/vscode/index.ts
@@ -5,6 +5,7 @@ declare const vscode: any;
 
 export function postVscMessage(type: string, data: any) {
   if (typeof vscode === "undefined") {
+    console.log("Unable to send message: vscode is undefined");
     return;
   }
   vscode.postMessage({