Merge pull request #232 from continuedev/asyncio-error-catching

Asyncio error catching

6 files changed, 78 insertions, 39 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 5c3baafd..615e7657 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -20,6 +20,7 @@ import asyncio
 from ..libs.util.step_name_to_steps import get_step_from_name
 from ..libs.util.traceback_parsers import get_python_traceback, get_javascript_traceback
 from openai import error as openai_errors
+from ..libs.util.create_async_task import create_async_task
 
 
 def get_error_title(e: Exception) -> str:
@@ -341,7 +342,8 @@ class Autopilot(ContinueBaseModel):
             # Update subscribers with new description
             await self.update_subscribers()
 
-        asyncio.create_task(update_description())
+        create_async_task(update_description(),
+                          self.continue_sdk.ide.unique_id)
 
         return observation
 
diff --git a/continuedev/src/continuedev/libs/util/create_async_task.py b/continuedev/src/continuedev/libs/util/create_async_task.py
new file mode 100644
index 00000000..608d4977
--- /dev/null
+++ b/continuedev/src/continuedev/libs/util/create_async_task.py
@@ -0,0 +1,23 @@
+from typing import Coroutine, Union
+import traceback
+from .telemetry import capture_event
+import asyncio
+import nest_asyncio
+nest_asyncio.apply()
+
+
+def create_async_task(coro: Coroutine, unique_id: Union[str, None] = None):
+    """asyncio.create_task and log errors by adding a callback"""
+    task = asyncio.create_task(coro)
+
+    def callback(future: asyncio.Future):
+        try:
+            future.result()
+        except Exception as e:
+            print("Exception caught from async task: ", e)
+            capture_event(unique_id or "None", "async_task_error", {
+                "error_title": e.__str__() or e.__repr__(), "error_message": traceback.format_tb(e.__traceback__)
+            })
+
+    task.add_done_callback(callback)
+    return task
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 8e9b1fb9..ae53be00 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -2,14 +2,14 @@ import json
 from fastapi import Depends, Header, WebSocket, APIRouter
 from typing import Any, List, Type, TypeVar, Union
 from pydantic import BaseModel
+import traceback
 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 +102,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 +188,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": traceback.format_tb(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..93996edd 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Type, TypeVar, Union
 import uuid
 from fastapi import WebSocket, Body, APIRouter
 from uvicorn.main import Server
+import traceback
 
 from ..libs.util.telemetry import capture_event
 from ..libs.util.queue import AsyncSubscriptionQueue
@@ -15,8 +16,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 +250,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 +413,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": traceback.format_tb(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):
diff --git a/continuedev/src/continuedev/steps/search_directory.py b/continuedev/src/continuedev/steps/search_directory.py
index 2eecc99c..bfb97630 100644
--- a/continuedev/src/continuedev/steps/search_directory.py
+++ b/continuedev/src/continuedev/steps/search_directory.py
@@ -6,6 +6,7 @@ from ..models.filesystem import RangeInFile
 from ..models.main import Range
 from ..core.main import Step
 from ..core.sdk import ContinueSDK
+from ..libs.util.create_async_task import create_async_task
 import os
 import re
 
@@ -60,9 +61,9 @@ class EditAllMatchesStep(Step):
         # Search all files for a given string
         range_in_files = find_all_matches_in_dir(self.pattern, self.directory or await sdk.ide.getWorkspaceDirectory())
 
-        tasks = [asyncio.create_task(sdk.edit_file(
+        tasks = [create_async_task(sdk.edit_file(
             range=range_in_file.range,
             filename=range_in_file.filepath,
             prompt=self.user_request
-        )) for range_in_file in range_in_files]
+        ), sdk.ide.unique_id) for range_in_file in range_in_files]
         await asyncio.gather(*tasks)