handle and explain tracebacks

4 files changed, 32 insertions, 10 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index b8f2695d..3ccce89a 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -13,6 +13,7 @@ from ..steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
 from ..libs.util.telemetry import capture_event
 from .sdk import ContinueSDK
 import asyncio
+from ..libs.util.step_name_to_steps import get_step_from_name
 
 
 class Autopilot(ContinueBaseModel):
@@ -88,9 +89,15 @@ class Autopilot(ContinueBaseModel):
             self._manual_edits_buffer.append(edit)
             # TODO: You're storing a lot of unecessary data here. Can compress into EditDiffs on the spot, and merge.
             # self._manual_edits_buffer = merge_file_edit(self._manual_edits_buffer, edit)
-
-    def handle_traceback(self, traceback: str):
-        raise NotImplementedError
+            # Note that this is being overriden to do nothing in DemoAgent
+
+    async def handle_command_output(self, output: str):
+        is_traceback = False
+        if is_traceback:
+            for tb_step in self.continue_sdk.config.on_traceback:
+                step = get_step_from_name(tb_step.step_name)(
+                    output=output, **tb_step.params)
+                await self._run_singular_step(step)
 
     _step_depth: int = 0
 
diff --git a/continuedev/src/continuedev/core/config.py b/continuedev/src/continuedev/core/config.py
index 23be8133..d8b29f5b 100644
--- a/continuedev/src/continuedev/core/config.py
+++ b/continuedev/src/continuedev/core/config.py
@@ -12,6 +12,11 @@ class SlashCommand(BaseModel):
     params: Optional[Dict] = {}
 
 
+class OnTracebackSteps(BaseModel):
+    step_name: str
+    params: Optional[Dict] = {}
+
+
 class ContinueConfig(BaseModel):
     """
     A pydantic class for the continue config file.
@@ -48,6 +53,8 @@ class ContinueConfig(BaseModel):
             step_name="FeedbackStep",
         )
     ]
+    on_traceback: Optional[List[OnTracebackSteps]] = [
+        OnTracebackSteps(step_name="DefaultOnTracebackStep")]
 
 
 def load_config(config_file: str) -> ContinueConfig:
diff --git a/continuedev/src/continuedev/libs/util/step_name_to_steps.py b/continuedev/src/continuedev/libs/util/step_name_to_steps.py
index b2bb838a..2c4474af 100644
--- a/continuedev/src/continuedev/libs/util/step_name_to_steps.py
+++ b/continuedev/src/continuedev/libs/util/step_name_to_steps.py
@@ -10,6 +10,7 @@ from ...recipes.AddTransformRecipe.main import AddTransformRecipe
 from ...recipes.CreatePipelineRecipe.main import CreatePipelineRecipe
 from ...recipes.DDtoBQRecipe.main import DDtoBQRecipe
 from ...recipes.DeployPipelineAirflowRecipe.main import DeployPipelineAirflowRecipe
+from ...steps.on_traceback import DefaultOnTracebackStep
 
 # This mapping is used to convert from string in ContinueConfig json to corresponding Step class.
 # Used for example in slash_commands and steps_on_startup
@@ -23,6 +24,7 @@ step_name_to_step_class = {
     "CreatePipelineRecipe": CreatePipelineRecipe,
     "DDtoBQRecipe": DDtoBQRecipe,
     "DeployPipelineAirflowRecipe": DeployPipelineAirflowRecipe,
+    "DefaultOnTracebackStep": DefaultOnTracebackStep,
 }
 
 
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index c53149d8..c66cc142 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -1,5 +1,4 @@
 # This is a separate server from server/main.py
-import asyncio
 from functools import cached_property
 import json
 import os
@@ -10,11 +9,13 @@ from uvicorn.main import Server
 
 from ..libs.util.queue import AsyncSubscriptionQueue
 from ..models.filesystem import FileSystem, RangeInFile, EditDiff, RealFileSystem
-from ..models.main import Traceback
 from ..models.filesystem_edit import AddDirectory, AddFile, DeleteDirectory, DeleteFile, FileSystemEdit, FileEdit, FileEditWithFullContents, RenameDirectory, RenameFile, SequentialFileSystemEdit
 from pydantic import BaseModel
 from .gui import SessionManager, session_manager
 from .ide_protocol import AbstractIdeProtocolServer
+import asyncio
+import nest_asyncio
+nest_asyncio.apply()
 
 
 router = APIRouter(prefix="/ide", tags=["ide"])
@@ -135,6 +136,9 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             fileEdits = list(
                 map(lambda d: FileEditWithFullContents.parse_obj(d), data["fileEdits"]))
             self.onFileEdits(fileEdits)
+        elif message_type == "commandOutput":
+            output = data["output"]
+            self.onCommandOutput(output)
         elif message_type in ["highlightedCode", "openFiles", "readFile", "editFile", "workspaceDirectory", "getUserSecret", "runCommand", "uniqueId"]:
             self.sub_queue.post(message_type, data)
         else:
@@ -189,11 +193,6 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
     def onAcceptRejectSuggestion(self, suggestionId: str, accepted: bool):
         pass
 
-    def onTraceback(self, traceback: Traceback):
-        # Same as below, maybe not every autopilot?
-        for _, session in self.session_manager.sessions.items():
-            session.autopilot.handle_traceback(traceback)
-
     def onFileSystemUpdate(self, update: FileSystemEdit):
         # Access to Autopilot (so SessionManager)
         pass
@@ -211,6 +210,13 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         for _, session in self.session_manager.sessions.items():
             session.autopilot.handle_manual_edits(edits)
 
+    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))
+
     # Request information. Session doesn't matter.
     async def getOpenFiles(self) -> List[str]:
         resp = await self._send_and_receive_json({}, OpenFilesResponse, "openFiles")