feat: :loud_sound: display any server errors to the GUI

3 files changed, 67 insertions, 57 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 2adb680e..98a5aea0 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -8,6 +8,7 @@ import traceback
 from uvicorn.main import Server
 
 from .session_manager import session_manager, Session
+from ..plugins.steps.core.core import DisplayErrorStep, MessageStep
 from .gui_protocol import AbstractGUIProtocolServer
 from ..libs.util.queue import AsyncSubscriptionQueue
 from ..libs.util.telemetry import posthog_logger
@@ -70,94 +71,88 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
         resp = await self._receive_json(message_type)
         return resp_model.parse_obj(resp)
 
+    def on_error(self, e: Exception):
+        return self.session.autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
+
     def handle_json(self, message_type: str, data: Any):
-        try:
-            if message_type == "main_input":
-                self.on_main_input(data["input"])
-            elif message_type == "step_user_input":
-                self.on_step_user_input(data["input"], data["index"])
-            elif message_type == "refinement_input":
-                self.on_refinement_input(data["input"], data["index"])
-            elif message_type == "reverse_to_index":
-                self.on_reverse_to_index(data["index"])
-            elif message_type == "retry_at_index":
-                self.on_retry_at_index(data["index"])
-            elif message_type == "clear_history":
-                self.on_clear_history()
-            elif message_type == "delete_at_index":
-                self.on_delete_at_index(data["index"])
-            elif message_type == "delete_context_with_ids":
-                self.on_delete_context_with_ids(data["ids"])
-            elif message_type == "toggle_adding_highlighted_code":
-                self.on_toggle_adding_highlighted_code()
-            elif message_type == "set_editing_at_indices":
-                self.on_set_editing_at_indices(data["indices"])
-            elif message_type == "show_logs_at_index":
-                self.on_show_logs_at_index(data["index"])
-            elif message_type == "select_context_item":
-                self.select_context_item(data["id"], data["query"])
-        except Exception as e:
-            logger.debug(e)
+        if message_type == "main_input":
+            self.on_main_input(data["input"])
+        elif message_type == "step_user_input":
+            self.on_step_user_input(data["input"], data["index"])
+        elif message_type == "refinement_input":
+            self.on_refinement_input(data["input"], data["index"])
+        elif message_type == "reverse_to_index":
+            self.on_reverse_to_index(data["index"])
+        elif message_type == "retry_at_index":
+            self.on_retry_at_index(data["index"])
+        elif message_type == "clear_history":
+            self.on_clear_history()
+        elif message_type == "delete_at_index":
+            self.on_delete_at_index(data["index"])
+        elif message_type == "delete_context_with_ids":
+            self.on_delete_context_with_ids(data["ids"])
+        elif message_type == "toggle_adding_highlighted_code":
+            self.on_toggle_adding_highlighted_code()
+        elif message_type == "set_editing_at_indices":
+            self.on_set_editing_at_indices(data["indices"])
+        elif message_type == "show_logs_at_index":
+            self.on_show_logs_at_index(data["index"])
+        elif message_type == "select_context_item":
+            self.select_context_item(data["id"], data["query"])
 
     def on_main_input(self, input: str):
         # Do something with user input
-        create_async_task(self.session.autopilot.accept_user_input(
-            input), self.session.autopilot.continue_sdk.ide.unique_id)
+        create_async_task(
+            self.session.autopilot.accept_user_input(input), self.on_error)
 
     def on_reverse_to_index(self, index: int):
         # Reverse the history to the given index
-        create_async_task(self.session.autopilot.reverse_to_index(
-            index), self.session.autopilot.continue_sdk.ide.unique_id)
+        create_async_task(
+            self.session.autopilot.reverse_to_index(index), self.on_error)
 
     def on_step_user_input(self, input: str, index: int):
         create_async_task(
-            self.session.autopilot.give_user_input(input, index), self.session.autopilot.continue_sdk.ide.unique_id)
+            self.session.autopilot.give_user_input(input, index), self.on_error)
 
     def on_refinement_input(self, input: str, index: int):
         create_async_task(
-            self.session.autopilot.accept_refinement_input(input, index), self.session.autopilot.continue_sdk.ide.unique_id)
+            self.session.autopilot.accept_refinement_input(input, index), self.on_error)
 
     def on_retry_at_index(self, index: int):
         create_async_task(
-            self.session.autopilot.retry_at_index(index), self.session.autopilot.continue_sdk.ide.unique_id)
+            self.session.autopilot.retry_at_index(index), self.on_error)
 
     def on_clear_history(self):
-        create_async_task(self.session.autopilot.clear_history(
-        ), self.session.autopilot.continue_sdk.ide.unique_id)
+        create_async_task(
+            self.session.autopilot.clear_history(), self.on_error)
 
     def on_delete_at_index(self, index: int):
-        create_async_task(self.session.autopilot.delete_at_index(
-            index), self.session.autopilot.continue_sdk.ide.unique_id)
+        create_async_task(
+            self.session.autopilot.delete_at_index(index), self.on_error)
 
     def on_delete_context_with_ids(self, ids: List[str]):
         create_async_task(
-            self.session.autopilot.delete_context_with_ids(
-                ids), self.session.autopilot.continue_sdk.ide.unique_id
-        )
+            self.session.autopilot.delete_context_with_ids(ids), self.on_error)
 
     def on_toggle_adding_highlighted_code(self):
         create_async_task(
-            self.session.autopilot.toggle_adding_highlighted_code(
-            ), self.session.autopilot.continue_sdk.ide.unique_id
-        )
+            self.session.autopilot.toggle_adding_highlighted_code(), self.on_error)
 
     def on_set_editing_at_indices(self, indices: List[int]):
         create_async_task(
-            self.session.autopilot.set_editing_at_indices(
-                indices), self.session.autopilot.continue_sdk.ide.unique_id
-        )
+            self.session.autopilot.set_editing_at_indices(indices), self.on_error)
 
     def on_show_logs_at_index(self, index: int):
         name = f"continue_logs.txt"
         logs = "\n\n############################################\n\n".join(
             ["This is a log of the exact prompt/completion pairs sent/received from the LLM during this step"] + self.session.autopilot.continue_sdk.history.timeline[index].logs)
         create_async_task(
-            self.session.autopilot.ide.showVirtualFile(name, logs), self.session.autopilot.continue_sdk.ide.unique_id)
+            self.session.autopilot.ide.showVirtualFile(name, logs), self.on_error)
 
     def select_context_item(self, id: str, query: str):
         """Called when user selects an item from the dropdown"""
         create_async_task(
-            self.session.autopilot.select_context_item(id, query), self.session.autopilot.continue_sdk.ide.unique_id)
+            self.session.autopilot.select_context_item(id, query), self.on_error)
 
 
 @router.websocket("/ws")
@@ -189,9 +184,14 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we
     except WebSocketDisconnect as e:
         logger.debug("GUI websocket disconnected")
     except Exception as e:
+        # Log, send to PostHog, and send to GUI
         logger.debug(f"ERROR in gui websocket: {e}")
+        err_msg = '\n'.join(traceback.format_exception(e))
         posthog_logger.capture_event("gui_error", {
-            "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))})
+            "error_title": e.__str__() or e.__repr__(), "error_message": err_msg})
+
+        await protocol.session.autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
+
         raise e
     finally:
         logger.debug("Closing gui websocket")
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 8a269cb7..e4c07029 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -10,6 +10,7 @@ from pydantic import BaseModel
 import traceback
 import asyncio
 
+from ..plugins.steps.core.core import DisplayErrorStep, MessageStep
 from .meilisearch_server import start_meilisearch
 from ..libs.util.telemetry import posthog_logger
 from ..libs.util.queue import AsyncSubscriptionQueue
@@ -279,6 +280,9 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
     # This is where you might have triggers: plugins can subscribe to certian events
     # like file changes, tracebacks, etc...
 
+    def on_error(self, e: Exception):
+        return self.session_manager.sessions[self.session_id].autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
+
     def onAcceptRejectSuggestion(self, accepted: bool):
         posthog_logger.capture_event("accept_reject_suggestion", {
             "accepted": accepted
@@ -311,22 +315,22 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
     def onDeleteAtIndex(self, index: int):
         if autopilot := self.__get_autopilot():
-            create_async_task(autopilot.delete_at_index(index), self.unique_id)
+            create_async_task(autopilot.delete_at_index(index), self.on_error)
 
     def onCommandOutput(self, output: str):
         if autopilot := self.__get_autopilot():
             create_async_task(
-                autopilot.handle_command_output(output), self.unique_id)
+                autopilot.handle_command_output(output), self.on_error)
 
     def onHighlightedCodeUpdate(self, range_in_files: List[RangeInFileWithContents]):
         if autopilot := self.__get_autopilot():
             create_async_task(autopilot.handle_highlighted_code(
-                range_in_files), self.unique_id)
+                range_in_files), self.on_error)
 
     def onMainUserInput(self, input: str):
         if autopilot := self.__get_autopilot():
             create_async_task(
-                autopilot.accept_user_input(input), self.unique_id)
+                autopilot.accept_user_input(input), self.on_error)
 
     # Request information. Session doesn't matter.
     async def getOpenFiles(self) -> List[str]:
@@ -459,7 +463,7 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
 
             logger.debug(f"Received IDE message: {message_type}")
             create_async_task(
-                ideProtocolServer.handle_json(message_type, data))
+                ideProtocolServer.handle_json(message_type, data), ideProtocolServer.on_error)
 
         ideProtocolServer = IdeProtocolServer(session_manager, websocket)
         if session_id is not None:
@@ -480,8 +484,12 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
         logger.debug("IDE wbsocket disconnected")
     except Exception as e:
         logger.debug(f"Error in ide websocket: {e}")
+        err_msg = '\n'.join(traceback.format_exception(e))
         posthog_logger.capture_event("gui_error", {
-            "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))})
+            "error_title": e.__str__() or e.__repr__(), "error_message": err_msg})
+
+        await session_manager.sessions[session_id].autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
+
         raise e
     finally:
         logger.debug("Closing ide websocket")
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index d30411cd..cf46028f 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -6,6 +6,7 @@ import json
 
 from fastapi.websockets import WebSocketState
 
+from ..plugins.steps.core.core import DisplayErrorStep
 from ..libs.util.paths import getSessionFilePath, getSessionsFolderPath
 from ..models.filesystem_edit import FileEditWithFullContents
 from ..libs.constants.main import CONTINUE_SESSIONS_FOLDER
@@ -83,7 +84,8 @@ class SessionManager:
             })
 
         autopilot.on_update(on_update)
-        create_async_task(autopilot.run_policy())
+        create_async_task(autopilot.run_policy(
+        ), lambda e: autopilot.continue_sdk.run_step(DisplayErrorStep(e=e)))
         return session
 
     async def remove_session(self, session_id: str):