5s timeout on websocket connections

5 files changed, 18 insertions, 10 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 9a411fbe..4201353e 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -1,3 +1,4 @@
+import asyncio
 import json
 from fastapi import Depends, Header, WebSocket, APIRouter
 from starlette.websockets import WebSocketState, WebSocketDisconnect
@@ -60,8 +61,12 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
             "data": data
         })
 
-    async def _receive_json(self, message_type: str) -> Any:
-        return await self.sub_queue.get(message_type)
+    async def _receive_json(self, message_type: str, timeout: int = 5) -> Any:
+        try:
+            return await asyncio.wait_for(self.sub_queue.get(message_type), timeout=timeout)
+        except asyncio.TimeoutError:
+            raise Exception(
+                "GUI Protocol _receive_json timed out after 5 seconds")
 
     async def _send_and_receive_json(self, data: Any, resp_model: Type[T], message_type: str) -> T:
         await self._send_json(message_type, data)
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 77b13483..e5e8de02 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -156,8 +156,12 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             "data": data
         })
 
-    async def _receive_json(self, message_type: str) -> Any:
-        return await self.sub_queue.get(message_type)
+    async def _receive_json(self, message_type: str, timeout: int = 5) -> Any:
+        try:
+            return await asyncio.wait_for(self.sub_queue.get(message_type), timeout=timeout)
+        except asyncio.TimeoutError:
+            raise Exception(
+                "IDE Protocol _receive_json timed out after 5 seconds")
 
     async def _send_and_receive_json(self, data: Any, resp_model: Type[T], message_type: str) -> T:
         await self._send_json(message_type, data)
diff --git a/extension/package-lock.json b/extension/package-lock.json
index f793abae..b86cb10e 100644
--- a/extension/package-lock.json
+++ b/extension/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "continue",
-  "version": "0.0.171",
+  "version": "0.0.172",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "continue",
-      "version": "0.0.171",
+      "version": "0.0.172",
       "license": "Apache-2.0",
       "dependencies": {
         "@electron/rebuild": "^3.2.10",
diff --git a/extension/package.json b/extension/package.json
index 38dc4542..6b719723 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -14,7 +14,7 @@
   "displayName": "Continue",
   "pricing": "Free",
   "description": "The open-source coding autopilot",
-  "version": "0.0.171",
+  "version": "0.0.172",
   "publisher": "Continue",
   "engines": {
     "vscode": "^1.67.0"
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index be1c220c..7a0d24d4 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -28,8 +28,6 @@ async function retryThenFail(
         );
         console.log("Execution policy stdout: ", stdout);
         console.log("Execution policy stderr: ", stderr);
-        // Then reload the window for this to take effect
-        await vscode.commands.executeCommand("workbench.action.reloadWindow");
       }
     }
 
@@ -447,7 +445,8 @@ export async function startContinuePythonServer() {
           console.log(`stdout: ${data}`);
           if (
             data.includes("Uvicorn running on") || // Successfully started the server
-            data.includes("address already in use") // The server is already running (probably a simultaneously opened VS Code window)
+            data.includes("only one usage of each socket address") || // [windows] The server is already running (probably a simultaneously opened VS Code window)
+            data.includes("address already in use") // [mac/linux] The server is already running (probably a simultaneously opened VS Code window)
           ) {
             console.log("Successfully started Continue python server");
             resolve(null);