fixed websockets issue

6 files changed, 101 insertions, 86 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index cc6235e9..c0178920 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -130,29 +130,36 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(websocket_session)):
-    await websocket.accept()
-
-    print("Session started")
-    session_manager.register_websocket(session.session_id, websocket)
-    protocol = GUIProtocolServer(session)
-    protocol.websocket = websocket
-
-    # Update any history that may have happened before connection
-    await protocol.send_available_slash_commands()
-    await protocol.send_state_update()
-
-    while AppStatus.should_exit is False:
-        message = await websocket.receive_text()
-        print("Received message", message)
-        if type(message) is str:
-            message = json.loads(message)
-
-        if "messageType" not in message or "data" not in message:
-            continue
-        message_type = message["messageType"]
-        data = message["data"]
-
-        protocol.handle_json(message_type, data)
-
-    print("Closing websocket")
-    await websocket.close()
+    try:
+        print("Received websocket connection at url: ", websocket.url)
+        await websocket.accept()
+
+        print("Session started")
+        session_manager.register_websocket(session.session_id, websocket)
+        protocol = GUIProtocolServer(session)
+        protocol.websocket = websocket
+
+        # Update any history that may have happened before connection
+        await protocol.send_available_slash_commands()
+        # await protocol.send_state_update() THIS WAS CAUSING A LOT OF ISSUES. Don't uncomment.
+
+        while AppStatus.should_exit is False:
+            message = await websocket.receive_text()
+            print("Received message", message)
+            if type(message) is str:
+                message = json.loads(message)
+
+            if "messageType" not in message or "data" not in message:
+                continue
+            message_type = message["messageType"]
+            data = message["data"]
+
+            protocol.handle_json(message_type, data)
+
+    except Exception as e:
+        print("ERROR in gui websocket: ", e)
+        raise e
+    finally:
+        print("Closing gui websocket")
+        await websocket.close()
+        session_manager.remove_session(session.session_id)
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 65f3ee74..3c3555f1 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -349,22 +349,28 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    print("Accepted websocket connection from, ", websocket.client)
-    await websocket.send_json({"messageType": "connected", "data": {}})
-
-    ideProtocolServer = IdeProtocolServer(session_manager)
-    ideProtocolServer.websocket = websocket
-
-    while True:
-        message = await websocket.receive_text()
-        message = json.loads(message)
-
-        if "messageType" not in message or "data" not in message:
-            continue
-        message_type = message["messageType"]
-        data = message["data"]
-
-        await ideProtocolServer.handle_json(message_type, data)
-
-    await websocket.close()
+    try:
+        await websocket.accept()
+        print("Accepted websocket connection from, ", websocket.client)
+        await websocket.send_json({"messageType": "connected", "data": {}})
+
+        ideProtocolServer = IdeProtocolServer(session_manager)
+        ideProtocolServer.websocket = websocket
+
+        while AppStatus.should_exit is False:
+            message = await websocket.receive_text()
+            message = json.loads(message)
+
+            if "messageType" not in message or "data" not in message:
+                continue
+            message_type = message["messageType"]
+            data = message["data"]
+
+            await ideProtocolServer.handle_json(message_type, data)
+
+        print("Closing ide websocket")
+        await websocket.close()
+    except Exception as e:
+        print("Error in ide websocket: ", e)
+        await websocket.close()
+        raise e
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index ebea08a5..99a38146 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -42,7 +42,7 @@ class SessionManager:
 
     def get_session(self, session_id: str) -> Session:
         if session_id not in self.sessions:
-            raise KeyError("Session ID not recognized")
+            raise KeyError("Session ID not recognized", session_id)
         return self.sessions[session_id]
 
     def new_session(self, ide: AbstractIdeProtocolServer) -> str:
diff --git a/extension/react-app/src/tabs/gui.tsx b/extension/react-app/src/tabs/gui.tsx
index 3e0fce6d..39925fc5 100644
--- a/extension/react-app/src/tabs/gui.tsx
+++ b/extension/react-app/src/tabs/gui.tsx
@@ -58,13 +58,13 @@ function GUI(props: GUIProps) {
   const [dataSwitchChecked, setDataSwitchChecked] = useState(false);
   const dataSwitchOn = useSelector(
     (state: RootStore) => state.config.dataSwitchOn
-  )
+  );
 
   useEffect(() => {
     if (typeof dataSwitchOn !== "undefined") {
-      setDataSwitchChecked(dataSwitchOn)
+      setDataSwitchChecked(dataSwitchOn);
     }
-  }, [dataSwitchOn])
+  }, [dataSwitchOn]);
 
   const [usingFastModel, setUsingFastModel] = useState(false);
   const [waitingForSteps, setWaitingForSteps] = useState(false);
@@ -548,9 +548,9 @@ function GUI(props: GUIProps) {
         }}
         hidden={!showDataSharingInfo}
       >
-        By turning on this switch, you will begin collecting accepted and 
-        rejected suggestions in .continue/suggestions.json. This data is 
-        stored locally on your machine and not sent anywhere.
+        By turning on this switch, you will begin collecting accepted and
+        rejected suggestions in .continue/suggestions.json. This data is stored
+        locally on your machine and not sent anywhere.
         <br />
         <br />
         <b>
@@ -583,15 +583,13 @@ function GUI(props: GUIProps) {
                 vscMachineId: vscMachineId,
                 dataSwitchChecked: !dataSwitchChecked,
               });
-              postVscMessage("toggleDataSwitch", {on: !dataSwitchChecked})
+              postVscMessage("toggleDataSwitch", { on: !dataSwitchChecked });
               setDataSwitchChecked((prev) => !prev);
             }}
             onColor="#12887a"
             checked={dataSwitchChecked}
           />
-          <span style={{ cursor: "help", fontSize: "14px" }}>
-            Collect Data
-          </span>
+          <span style={{ cursor: "help", fontSize: "14px" }}>Collect Data</span>
         </div>
         <HeaderButtonWithText
           onClick={() => {
@@ -616,7 +614,10 @@ function GUI(props: GUIProps) {
         >
           <Trash size="1.6em" />
         </HeaderButtonWithText>
-        <a href="https://continue.dev/docs/how-to-use-continue" className="no-underline">
+        <a
+          href="https://continue.dev/docs/how-to-use-continue"
+          className="no-underline"
+        >
           <HeaderButtonWithText text="Docs">
             <BookOpen size="1.6em" />
           </HeaderButtonWithText>
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index 2d9afaec..f1aace77 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -71,7 +71,8 @@ async function getPythonPipCommands() {
   }
 
   const version = stdout.split(" ")[1];
-  if (version < "3.7") {
+  const [major, minor] = version.split(".");
+  if (parseInt(major) !== 3 || parseInt(minor) < 7) {
     vscode.window.showErrorMessage(
       "Continue requires Python3 version 3.7 or greater. Please update your Python3 installation, reload VS Code, and try again."
     );
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 1ccc070c..b179cbf3 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -37,34 +37,34 @@ class IdeProtocolClient {
     });
 
     // Setup listeners for any file changes in open editors
-    vscode.workspace.onDidChangeTextDocument((event) => {
-      if (this._makingEdit === 0) {
-        let fileEdits: FileEditWithFullContents[] = event.contentChanges.map(
-          (change) => {
-            return {
-              fileEdit: {
-                filepath: event.document.uri.fsPath,
-                range: {
-                  start: {
-                    line: change.range.start.line,
-                    character: change.range.start.character,
-                  },
-                  end: {
-                    line: change.range.end.line,
-                    character: change.range.end.character,
-                  },
-                },
-                replacement: change.text,
-              },
-              fileContents: event.document.getText(),
-            };
-          }
-        );
-        this.messenger?.send("fileEdits", { fileEdits });
-      } else {
-        this._makingEdit--;
-      }
-    });
+    // vscode.workspace.onDidChangeTextDocument((event) => {
+    //   if (this._makingEdit === 0) {
+    //     let fileEdits: FileEditWithFullContents[] = event.contentChanges.map(
+    //       (change) => {
+    //         return {
+    //           fileEdit: {
+    //             filepath: event.document.uri.fsPath,
+    //             range: {
+    //               start: {
+    //                 line: change.range.start.line,
+    //                 character: change.range.start.character,
+    //               },
+    //               end: {
+    //                 line: change.range.end.line,
+    //                 character: change.range.end.character,
+    //               },
+    //             },
+    //             replacement: change.text,
+    //           },
+    //           fileContents: event.document.getText(),
+    //         };
+    //       }
+    //     );
+    //     this.messenger?.send("fileEdits", { fileEdits });
+    //   } else {
+    //     this._makingEdit--;
+    //   }
+    // });
   }
 
   async handleMessage(messageType: string, data: any) {