Merge branch 'main' into package-python

7 files changed, 60 insertions, 83 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index d92c51cd..a3dd854e 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -1,7 +1,7 @@
 from functools import cached_property
 import traceback
 import time
-from typing import Any, Callable, Coroutine, Dict, List, Union
+from typing import Callable, Coroutine, Dict, List, Union
 from aiohttp import ClientPayloadError
 from pydantic import root_validator
 
@@ -54,7 +54,7 @@ class Autopilot(ContinueBaseModel):
     history: History = History.from_empty()
     context: Context = Context()
     full_state: Union[FullState, None] = None
-    context_manager: Union[ContextManager, None] = None
+    context_manager: ContextManager = ContextManager()
     continue_sdk: ContinueSDK = None
 
     _on_update_callbacks: List[Callable[[FullState], None]] = []
@@ -66,19 +66,22 @@ class Autopilot(ContinueBaseModel):
     _user_input_queue = AsyncSubscriptionQueue()
     _retry_queue = AsyncSubscriptionQueue()
 
+    started: bool = False
+
     async def start(self):
         self.continue_sdk = await ContinueSDK.create(self)
         if override_policy := self.continue_sdk.config.policy_override:
             self.policy = override_policy
 
         # Load documents into the search index
-        self.context_manager = await ContextManager.create(
+        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)
+        self.started = True
 
     class Config:
         arbitrary_types_allowed = True
@@ -98,7 +101,7 @@ class Autopilot(ContinueBaseModel):
             user_input_queue=self._main_user_input_queue,
             slash_commands=self.get_available_slash_commands(),
             adding_highlighted_code=self.context_manager.context_providers[
-                "code"].adding_highlighted_code if self.context_manager is not None else False,
+                "code"].adding_highlighted_code if "code" in self.context_manager.context_providers else False,
             selected_context_items=await self.context_manager.get_selected_items() if self.context_manager is not None else [],
         )
         self.full_state = full_state
@@ -201,7 +204,7 @@ class Autopilot(ContinueBaseModel):
         await self.update_subscribers()
 
     async def set_editing_at_ids(self, ids: List[str]):
-        self.context_manager.context_providers["code"].set_editing_at_ids(ids)
+        await self.context_manager.context_providers["code"].set_editing_at_ids(ids)
         await self.update_subscribers()
 
     async def _run_singular_step(self, step: "Step", is_future_step: bool = False) -> Coroutine[Observation, None, None]:
@@ -244,7 +247,7 @@ class Autopilot(ContinueBaseModel):
         try:
             observation = await step(self.continue_sdk)
         except Exception as e:
-            if self.history.timeline[index_of_history_node].deleted:
+            if index_of_history_node >= len(self.history.timeline) or self.history.timeline[index_of_history_node].deleted:
                 # If step was deleted/cancelled, don't show error or allow retry
                 return None
 
@@ -301,7 +304,7 @@ class Autopilot(ContinueBaseModel):
         self._step_depth -= 1
 
         # Add observation to history, unless already attached error observation
-        if not caught_error:
+        if not caught_error and index_of_history_node < len(self.history.timeline):
             self.history.timeline[index_of_history_node].observation = observation
             self.history.timeline[index_of_history_node].active = False
             await self.update_subscribers()
diff --git a/continuedev/src/continuedev/core/context.py b/continuedev/src/continuedev/core/context.py
index e968c35c..3f5f6fd3 100644
--- a/continuedev/src/continuedev/core/context.py
+++ b/continuedev/src/continuedev/core/context.py
@@ -133,14 +133,19 @@ class ContextManager:
         """
         return sum([await provider.get_chat_messages() for provider in self.context_providers.values()], [])
 
-    def __init__(self, context_providers: List[ContextProvider]):
+    def __init__(self):
+        self.context_providers = {}
+        self.provider_titles = set()
+
+    async def start(self, context_providers: List[ContextProvider]):
+        """
+        Starts the context manager.
+        """
         self.context_providers = {
             prov.title: prov for prov in context_providers}
         self.provider_titles = {
             provider.title for provider in context_providers}
 
-    @classmethod
-    async def create(cls, context_providers: List[ContextProvider]):
         async with Client('http://localhost:7700') as search_client:
             meilisearch_running = True
             try:
@@ -154,10 +159,8 @@ class ContextManager:
             if not meilisearch_running:
                 logger.warning(
                     "MeiliSearch not running, avoiding any dependent context providers")
-                context_providers = list(
-                    filter(lambda cp: cp.title == "code", context_providers))
-
-        return cls(context_providers)
+                self.context_providers = list(
+                    filter(lambda cp: cp.title == "code", self.context_providers))
 
     async def load_index(self, workspace_dir: str):
         for _, provider in self.context_providers.items():
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index cf18c56b..7c89c5c2 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -93,8 +93,8 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
             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 == "set_editing_at_ids":
+            self.on_set_editing_at_ids(data["ids"])
         elif message_type == "show_logs_at_index":
             self.on_show_logs_at_index(data["index"])
         elif message_type == "select_context_item":
@@ -138,9 +138,9 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
         create_async_task(
             self.session.autopilot.toggle_adding_highlighted_code(), self.on_error)
 
-    def on_set_editing_at_indices(self, indices: List[int]):
+    def on_set_editing_at_ids(self, ids: List[str]):
         create_async_task(
-            self.session.autopilot.set_editing_at_indices(indices), self.on_error)
+            self.session.autopilot.set_editing_at_ids(ids), self.on_error)
 
     def on_show_logs_at_index(self, index: int):
         name = f"continue_logs.txt"
@@ -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..9797a8b7 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", {
@@ -307,7 +309,9 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
     def __get_autopilot(self):
         if self.session_id not in self.session_manager.sessions:
             return None
-        return self.session_manager.sessions[self.session_id].autopilot
+
+        autopilot = self.session_manager.sessions[self.session_id].autopilot
+        return autopilot if autopilot.started else None
 
     def onFileEdits(self, edits: List[FileEditWithFullContents]):
         if autopilot := self.__get_autopilot():
@@ -442,6 +446,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 +458,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 +471,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 +482,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/main.py b/continuedev/src/continuedev/server/main.py
index 468bc855..f8dfb009 100644
--- a/continuedev/src/continuedev/server/main.py
+++ b/continuedev/src/continuedev/server/main.py
@@ -1,5 +1,4 @@
 import asyncio
-import sys
 import time
 import psutil
 import os
@@ -8,13 +7,11 @@ from fastapi.middleware.cors import CORSMiddleware
 import atexit
 import uvicorn
 import argparse
-import logging.config
 
 
 from .ide import router as ide_router
 from .gui import router as gui_router
 from .session_manager import session_manager
-from ..libs.util.paths import getLogFilePath
 from ..libs.util.logging import logger
 
 app = FastAPI()
@@ -38,25 +35,6 @@ def health():
     return {"status": "ok"}
 
 
-class Logger(object):
-    def __init__(self, log_file: str):
-        self.terminal = sys.stdout
-        self.log = open(log_file, "a")
-
-    def write(self, message):
-        self.terminal.write(message)
-        self.log.write(message)
-
-    def flush(self):
-        # this flush method is needed for python 3 compatibility.
-        # this handles the flush command by doing nothing.
-        # you might want to specify some extra behavior here.
-        pass
-
-    def isatty(self):
-        return False
-
-
 try:
     # add cli arg for server port
     parser = argparse.ArgumentParser()
@@ -71,7 +49,6 @@ except Exception as e:
 def run_server():
     config = uvicorn.Config(app, host="127.0.0.1", port=args.port)
     server = uvicorn.Server(config)
-
     server.run()
 
 
@@ -87,32 +64,10 @@ def cleanup():
     loop.close()
 
 
-def cpu_usage_report():
-    process = psutil.Process(os.getpid())
-    # Call cpu_percent once to start measurement, but ignore the result
-    process.cpu_percent(interval=None)
-    # Wait for a short period of time
-    time.sleep(1)
-    # Call cpu_percent again to get the CPU usage over the interval
-    cpu_usage = process.cpu_percent(interval=None)
-    logger.debug(f"CPU usage: {cpu_usage}%")
-
-
 atexit.register(cleanup)
 
 if __name__ == "__main__":
     try:
-        # Uncomment to get CPU usage reports
-        # import threading
-
-        # def cpu_usage_loop():
-        #     while True:
-        #         cpu_usage_report()
-        #         time.sleep(2)
-
-        # cpu_thread = threading.Thread(target=cpu_usage_loop)
-        # cpu_thread.start()
-
         run_server()
     except Exception as e:
         logger.debug(f"Error starting Continue server: {e}")
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):