refactor: :loud_sound: show errors to IDE before GUI connected

4 files changed, 33 insertions, 19 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index cf18c56b..57ecceb2 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -190,7 +190,7 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we
         posthog_logger.capture_event("gui_error", {
             "error_title": e.__str__() or e.__repr__(), "error_message": err_msg})
 
-        await protocol.session.autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
+        await session.autopilot.ide.showMessage(err_msg)
 
         raise e
     finally:
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 6124f3bd..b026f4f2 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -1,7 +1,7 @@
 # This is a separate server from server/main.py
 import json
 import os
-from typing import Any, List, Type, TypeVar, Union
+from typing import Any, Coroutine, List, Type, TypeVar, Union
 import uuid
 from fastapi import WebSocket, APIRouter
 from starlette.websockets import WebSocketState, WebSocketDisconnect
@@ -211,8 +211,6 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         else:
             raise ValueError("Unknown message type", message_type)
 
-    # ------------------------------- #
-    # Request actions in IDE, doesn't matter which Session
     async def showSuggestion(self, file_edit: FileEdit):
         await self._send_json("showSuggestion", {
             "edit": file_edit.dict()
@@ -232,6 +230,11 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             "open": open
         })
 
+    async def showMessage(self, message: str):
+        await self._send_json("showMessage", {
+            "message": message
+        })
+
     async def showVirtualFile(self, name: str, contents: str):
         await self._send_json("showVirtualFile", {
             "name": name,
@@ -275,13 +278,12 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         # Just need connect the suggestionId to the IDE (and the gui)
         return any([r.accepted for r in responses])
 
-    # ------------------------------- #
-    # Here needs to pass message onto the Autopilot OR Autopilot just subscribes.
-    # 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 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}")
 
     def onAcceptRejectSuggestion(self, accepted: bool):
         posthog_logger.capture_event("accept_reject_suggestion", {
@@ -442,6 +444,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
     try:
+        # Accept the websocket connection
         await websocket.accept()
         logger.debug(f"Accepted websocket connection from {websocket.client}")
         await websocket.send_json({"messageType": "connected", "data": {}})
@@ -453,6 +456,7 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
             logger.debug("Failed to start MeiliSearch")
             logger.debug(e)
 
+        # Message handler
         def handle_msg(msg):
             message = json.loads(msg)
 
@@ -465,6 +469,7 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
             create_async_task(
                 ideProtocolServer.handle_json(message_type, data), ideProtocolServer.on_error)
 
+        # Initialize the IDE Protocol Server
         ideProtocolServer = IdeProtocolServer(session_manager, websocket)
         if session_id is not None:
             session_manager.registered_ides[session_id] = ideProtocolServer
@@ -475,20 +480,23 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
         for other_msg in other_msgs:
             handle_msg(other_msg)
 
+        # Handle messages
         while AppStatus.should_exit is False:
             message = await websocket.receive_text()
             handle_msg(message)
 
-        logger.debug("Closing ide websocket")
     except WebSocketDisconnect as e:
-        logger.debug("IDE wbsocket disconnected")
+        logger.debug("IDE websocket 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": err_msg})
 
-        await session_manager.sessions[session_id].autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
+        if session_id is not None and session_id in session_manager.sessions:
+            await session_manager.sessions[session_id].autopilot.continue_sdk.run_step(DisplayErrorStep(e=e))
+        elif ideProtocolServer is not None:
+            await ideProtocolServer.showMessage(f"Error in Continue server: {err_msg}")
 
         raise e
     finally:
diff --git a/continuedev/src/continuedev/server/ide_protocol.py b/continuedev/src/continuedev/server/ide_protocol.py
index 0ae7e7fa..72b410d4 100644
--- a/continuedev/src/continuedev/server/ide_protocol.py
+++ b/continuedev/src/continuedev/server/ide_protocol.py
@@ -24,6 +24,10 @@ class AbstractIdeProtocolServer(ABC):
         """Set whether a file is open"""
 
     @abstractmethod
+    async def showMessage(self, message: str):
+        """Show a message to the user"""
+
+    @abstractmethod
     async def showVirtualFile(self, name: str, contents: str):
         """Show a virtual file"""
 
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index b5580fe8..56c92307 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -1,15 +1,14 @@
 import os
 import traceback
 from fastapi import WebSocket
-from typing import Any, Dict, List, Union
+from typing import Any, Coroutine, Dict, Union
 from uuid import uuid4
 import json
 
 from fastapi.websockets import WebSocketState
 
-from ..plugins.steps.core.core import DisplayErrorStep, MessageStep
+from ..plugins.steps.core.core import MessageStep
 from ..libs.util.paths import getSessionFilePath, getSessionsFolderPath
-from ..models.filesystem_edit import FileEditWithFullContents
 from ..core.main import FullState, HistoryNode
 from ..core.autopilot import Autopilot
 from .ide_protocol import AbstractIdeProtocolServer
@@ -90,8 +89,11 @@ class SessionManager:
             ))
             logger.warning(f"Error loading context manager: {e}")
 
-        create_async_task(autopilot.run_policy(), lambda e: autopilot.continue_sdk.run_step(
-            DisplayErrorStep(e=e)))
+        def on_error(e: Exception) -> Coroutine:
+            err_msg = '\n'.join(traceback.format_exception(e))
+            return ide.showMessage(f"Error in Continue server: {err_msg}")
+
+        create_async_task(autopilot.run_policy(), on_error)
         return session
 
     async def remove_session(self, session_id: str):