catch errors during asyncio tasks

3 files changed, 47 insertions, 36 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 8e9b1fb9..8daad73a 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -7,9 +7,8 @@ from uvicorn.main import Server
 from .session_manager import SessionManager, session_manager, Session
 from .gui_protocol import AbstractGUIProtocolServer
 from ..libs.util.queue import AsyncSubscriptionQueue
-import asyncio
-import nest_asyncio
-nest_asyncio.apply()
+from ..libs.util.telemetry import capture_event
+from ..libs.util.create_async_task import create_async_task
 
 router = APIRouter(prefix="/gui", tags=["gui"])
 
@@ -102,51 +101,60 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
 
     def on_main_input(self, input: str):
         # Do something with user input
-        asyncio.create_task(self.session.autopilot.accept_user_input(input))
+        create_async_task(self.session.autopilot.accept_user_input(
+            input), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_reverse_to_index(self, index: int):
         # Reverse the history to the given index
-        asyncio.create_task(self.session.autopilot.reverse_to_index(index))
+        create_async_task(self.session.autopilot.reverse_to_index(
+            index), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_step_user_input(self, input: str, index: int):
-        asyncio.create_task(
-            self.session.autopilot.give_user_input(input, index))
+        create_async_task(
+            self.session.autopilot.give_user_input(input, index), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_refinement_input(self, input: str, index: int):
-        asyncio.create_task(
-            self.session.autopilot.accept_refinement_input(input, index))
+        create_async_task(
+            self.session.autopilot.accept_refinement_input(input, index), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_retry_at_index(self, index: int):
-        asyncio.create_task(
-            self.session.autopilot.retry_at_index(index))
+        create_async_task(
+            self.session.autopilot.retry_at_index(index), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_change_default_model(self, model: str):
-        asyncio.create_task(self.session.autopilot.change_default_model(model))
+        create_async_task(self.session.autopilot.change_default_model(
+            model), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_clear_history(self):
-        asyncio.create_task(self.session.autopilot.clear_history())
+        create_async_task(self.session.autopilot.clear_history(
+        ), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_delete_at_index(self, index: int):
-        asyncio.create_task(self.session.autopilot.delete_at_index(index))
+        create_async_task(self.session.autopilot.delete_at_index(
+            index), self.session.autopilot.continue_sdk.ide.unique_id)
 
     def on_delete_context_at_indices(self, indices: List[int]):
-        asyncio.create_task(
-            self.session.autopilot.delete_context_at_indices(indices)
+        create_async_task(
+            self.session.autopilot.delete_context_at_indices(
+                indices), self.session.autopilot.continue_sdk.ide.unique_id
         )
 
     def on_toggle_adding_highlighted_code(self):
-        asyncio.create_task(
-            self.session.autopilot.toggle_adding_highlighted_code()
+        create_async_task(
+            self.session.autopilot.toggle_adding_highlighted_code(
+            ), self.session.autopilot.continue_sdk.ide.unique_id
         )
 
     def on_set_editing_at_indices(self, indices: List[int]):
-        asyncio.create_task(
-            self.session.autopilot.set_editing_at_indices(indices)
+        create_async_task(
+            self.session.autopilot.set_editing_at_indices(
+                indices), self.session.autopilot.continue_sdk.ide.unique_id
         )
 
     def on_set_pinned_at_indices(self, indices: List[int]):
-        asyncio.create_task(
-            self.session.autopilot.set_pinned_at_indices(indices)
+        create_async_task(
+            self.session.autopilot.set_pinned_at_indices(
+                indices), self.session.autopilot.continue_sdk.ide.unique_id
         )
 
 
@@ -179,6 +187,8 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we
 
     except Exception as e:
         print("ERROR in gui websocket: ", e)
+        capture_event(session.autopilot.continue_sdk.ide.unique_id, "gui_error", {
+                      "error_title": e.__str__() or e.__repr__(), "error_message": e.__traceback__})
         raise e
     finally:
         print("Closing gui websocket")
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index e4a6266a..f84f3de2 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -15,8 +15,7 @@ from pydantic import BaseModel
 from .gui import SessionManager, session_manager
 from .ide_protocol import AbstractIdeProtocolServer
 import asyncio
-import nest_asyncio
-nest_asyncio.apply()
+from ..libs.util.create_async_task import create_async_task
 
 
 router = APIRouter(prefix="/ide", tags=["ide"])
@@ -250,24 +249,25 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
     def onDeleteAtIndex(self, index: int):
         for _, session in self.session_manager.sessions.items():
-            asyncio.create_task(session.autopilot.delete_at_index(index))
+            create_async_task(
+                session.autopilot.delete_at_index(index), self.unique_id)
 
     def onCommandOutput(self, output: str):
         # Send the output to ALL autopilots.
         # Maybe not ideal behavior
         for _, session in self.session_manager.sessions.items():
-            asyncio.create_task(
-                session.autopilot.handle_command_output(output))
+            create_async_task(
+                session.autopilot.handle_command_output(output), self.unique_id)
 
     def onHighlightedCodeUpdate(self, range_in_files: List[RangeInFileWithContents]):
         for _, session in self.session_manager.sessions.items():
-            asyncio.create_task(
-                session.autopilot.handle_highlighted_code(range_in_files))
+            create_async_task(
+                session.autopilot.handle_highlighted_code(range_in_files), self.unique_id)
 
     def onMainUserInput(self, input: str):
         for _, session in self.session_manager.sessions.items():
-            asyncio.create_task(
-                session.autopilot.accept_user_input(input))
+            create_async_task(
+                session.autopilot.accept_user_input(input), self.unique_id)
 
     # Request information. Session doesn't matter.
     async def getOpenFiles(self) -> List[str]:
@@ -412,5 +412,7 @@ async def websocket_endpoint(websocket: WebSocket):
         await websocket.close()
     except Exception as e:
         print("Error in ide websocket: ", e)
+        capture_event(ideProtocolServer.unique_id, "gui_error", {
+                      "error_title": e.__str__() or e.__repr__(), "error_message": e.__traceback__})
         await websocket.close()
         raise e
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index 99a38146..873a379e 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -1,3 +1,4 @@
+from asyncio import BaseEventLoop
 from fastapi import WebSocket
 from typing import Any, Dict, List, Union
 from uuid import uuid4
@@ -7,9 +8,7 @@ from ..core.policy import DemoPolicy
 from ..core.main import FullState
 from ..core.autopilot import Autopilot
 from .ide_protocol import AbstractIdeProtocolServer
-import asyncio
-import nest_asyncio
-nest_asyncio.apply()
+from ..libs.util.create_async_task import create_async_task
 
 
 class Session:
@@ -38,7 +37,7 @@ class DemoAutopilot(Autopilot):
 
 class SessionManager:
     sessions: Dict[str, Session] = {}
-    _event_loop: Union[asyncio.BaseEventLoop, None] = None
+    _event_loop: Union[BaseEventLoop, None] = None
 
     def get_session(self, session_id: str) -> Session:
         if session_id not in self.sessions:
@@ -57,7 +56,7 @@ class SessionManager:
             })
 
         autopilot.on_update(on_update)
-        asyncio.create_task(autopilot.run_policy())
+        create_async_task(autopilot.run_policy())
         return session_id
 
     def remove_session(self, session_id: str):