2 files changed, 38 insertions, 55 deletions

diff --git a/extension/react-app/src/tabs/chat/MessageDiv.tsx b/extension/react-app/src/tabs/chat/MessageDiv.tsx
index d7c79721..9bdd8638 100644
--- a/extension/react-app/src/tabs/chat/MessageDiv.tsx
+++ b/extension/react-app/src/tabs/chat/MessageDiv.tsx
@@ -20,7 +20,8 @@ const Container = styled.div`
   margin: 3px;
   width: fit-content;
   max-width: 75%;
-  overflow: scroll;
+  overflow-y: scroll;
+  scrollbar-gutter: stable both-edges;
   word-wrap: break-word;
   -ms-word-wrap: break-word;
   height: fit-content;
diff --git a/extension/react-app/src/tabs/notebook.tsx b/extension/react-app/src/tabs/gui.tsx
index a9c69c5b..42ad4ed5 100644
--- a/extension/react-app/src/tabs/notebook.tsx
+++ b/extension/react-app/src/tabs/gui.tsx
@@ -14,10 +14,12 @@ import StepContainer from "../components/StepContainer";
 import { useSelector } from "react-redux";
 import { RootStore } from "../redux/store";
 import useContinueWebsocket from "../hooks/useWebsocket";
+import useContinueGUIProtocol from "../hooks/useWebsocket";
 
-let TopNotebookDiv = styled.div`
+let TopGUIDiv = styled.div`
   display: grid;
   grid-template-columns: 1fr;
+  overflow: scroll;
 `;
 
 let UserInputQueueItem = styled.div`
@@ -28,17 +30,15 @@ let UserInputQueueItem = styled.div`
   text-align: center;
 `;
 
-interface NotebookProps {
+interface GUIProps {
   firstObservation?: any;
 }
 
-function Notebook(props: NotebookProps) {
-  const serverUrl = useSelector((state: RootStore) => state.config.apiUrl);
-
+function GUI(props: GUIProps) {
   const [waitingForSteps, setWaitingForSteps] = useState(false);
   const [userInputQueue, setUserInputQueue] = useState<string[]>([]);
   const [history, setHistory] = useState<History | undefined>();
-  //   {
+  // {
   //   timeline: [
   //     {
   //       step: {
@@ -154,33 +154,19 @@ function Notebook(props: NotebookProps) {
   //     },
   //   ],
   //   current_index: 0,
-  // } as any
-  // );
+  // } 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 = useContinueGUIProtocol();
 
-  // 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,22 @@ 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>
+    <TopGUIDiv>
+      {typeof client === "undefined" && (
+        <>
+          <Loader></Loader>
+          <p style={{ textAlign: "center" }}>
+            Trying to reconnect with server...
+          </p>
+        </>
+      )}
       {history?.timeline.map((node: HistoryNode, index: number) => {
         return (
           <StepContainer
@@ -237,17 +226,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);
             }}
           />
         );
@@ -278,8 +260,8 @@ function Notebook(props: NotebookProps) {
         }}
       ></MainTextInput>
       <ContinueButton onClick={onMainTextInput}></ContinueButton>
-    </TopNotebookDiv>
+    </TopGUIDiv>
   );
 }
 
-export default Notebook;
+export default GUI;