fix: :bug: timeout on blocking processes

5 files changed, 26 insertions, 15 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index a3dd854e..f3a17d47 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -74,13 +74,16 @@ class Autopilot(ContinueBaseModel):
             self.policy = override_policy
 
         # Load documents into the search index
+        logger.debug("Starting context manager")
         await self.context_manager.start(
             self.continue_sdk.config.context_providers + [
                 HighlightedCodeContextProvider(ide=self.ide),
                 FileContextProvider(workspace_dir=self.ide.workspace_directory)
             ])
 
-        await self.context_manager.load_index(self.ide.workspace_directory)
+        logger.debug("Loading index")
+        create_async_task(self.context_manager.load_index(
+            self.ide.workspace_directory))
         self.started = True
 
     class Config:
diff --git a/continuedev/src/continuedev/core/context.py b/continuedev/src/continuedev/core/context.py
index 20725216..f83cbd34 100644
--- a/continuedev/src/continuedev/core/context.py
+++ b/continuedev/src/continuedev/core/context.py
@@ -1,5 +1,7 @@
 
 from abc import abstractmethod
+import asyncio
+import time
 from typing import Dict, List
 from meilisearch_python_async import Client
 from pydantic import BaseModel
@@ -165,6 +167,8 @@ class ContextManager:
 
     async def load_index(self, workspace_dir: str):
         for _, provider in self.context_providers.items():
+            ti = time.time()
+
             context_items = await provider.provide_context_items(workspace_dir)
             documents = [
                 {
@@ -178,10 +182,14 @@ class ContextManager:
             if len(documents) > 0:
                 try:
                     async with Client('http://localhost:7700') as search_client:
-                        await search_client.index(SEARCH_INDEX_NAME).add_documents(documents)
+                        await asyncio.wait_for(search_client.index(SEARCH_INDEX_NAME).add_documents(documents), timeout=5)
                 except Exception as e:
                     logger.debug(f"Error loading meilisearch index: {e}")
 
+            tf = time.time()
+            logger.debug(
+                f"Loaded {len(documents)} documents into meilisearch in {tf - ti} seconds for context provider {provider.title}")
+
     async def select_context_item(self, id: str, query: str):
         """
         Selects the ContextItem with the given id.
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 9797a8b7..e523d3aa 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -249,8 +249,10 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         })
 
     async def getSessionId(self):
-        session_id = (await self.session_manager.new_session(
-            self, self.session_id)).session_id
+        new_session = await asyncio.wait_for(self.session_manager.new_session(
+            self, self.session_id), timeout=5)
+        session_id = new_session.session_id
+        logger.debug(f"Sending session id: {session_id}")
         await self._send_json("getSessionId", {
             "sessionId": session_id
         })
@@ -279,11 +281,9 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         return any([r.accepted for r in responses])
 
     def on_error(self, e: Exception) -> Coroutine:
-        try:
-            return self.session_manager.sessions[self.session_id].autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
-        except:
-            err_msg = '\n'.join(traceback.format_exception(e))
-            return self.showMessage(f"Error in Continue server: {err_msg}")
+        err_msg = '\n'.join(traceback.format_exception(e))
+        e_title = e.__str__() or e.__repr__()
+        return self.showMessage(f"Error in Continue server: {e_title}\n {err_msg}")
 
     def onAcceptRejectSuggestion(self, accepted: bool):
         posthog_logger.capture_event("accept_reject_suggestion", {
diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index d07719a8..560b970c 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -56,12 +56,8 @@ export async function activateExtension(context: vscode.ExtensionContext) {
 
   // Start the server and display loader if taking > 2 seconds
   const sessionIdPromise = (async () => {
-    await new Promise((resolve) => {
-      // Start the server and set serverStarted to true when done
-      startContinuePythonServer().then(() => {
-        resolve(null);
-      });
-    });
+    // Start the server and set serverStarted to true when done
+    await startContinuePythonServer();
 
     console.log("Continue server started");
     // Initialize IDE Protocol Client
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index ff6fdd66..8818c949 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -230,6 +230,10 @@ export async function startContinuePythonServer() {
     detached: true,
     stdio: "ignore",
   });
+  child.on("error", (err: any) => {
+    console.error("Failed to start subprocess.", err);
+  });
+
   child.unref();
 
   // Write the current version of vscode extension to a file called server_version.txt