pause streaming on rejection

4 files changed, 20 insertions, 5 deletions

diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index 49513013..ed670799 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -192,7 +192,8 @@ class ContinueSDK(AbstractContinueSDK):
 
     async def get_chat_context(self) -> List[ChatMessage]:
         history_context = self.history.to_chat_history()
-        highlighted_code = self.__autopilot._highlighted_ranges
+        highlighted_code = [
+            hr.range for hr in self.__autopilot._highlighted_ranges]
 
         preface = "The following code is highlighted"
 
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 1d51758e..61e7ca78 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -152,6 +152,8 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             self.onAcceptRejectDiff(data["accepted"])
         elif message_type == "mainUserInput":
             self.onMainUserInput(data["input"])
+        elif message_type == "deleteAtIndex":
+            self.onDeleteAtIndex(data["index"])
         elif message_type in ["highlightedCode", "openFiles", "readFile", "editFile", "workspaceDirectory", "getUserSecret", "runCommand", "uniqueId"]:
             self.sub_queue.post(message_type, data)
         else:
@@ -164,10 +166,11 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             "edit": file_edit.dict()
         })
 
-    async def showDiff(self, filepath: str, replacement: str):
+    async def showDiff(self, filepath: str, replacement: str, step_index: int):
         await self._send_json("showDiff", {
             "filepath": filepath,
-            "replacement": replacement
+            "replacement": replacement,
+            "step_index": step_index
         })
 
     async def setFileOpen(self, filepath: str, open: bool = True):
@@ -245,6 +248,10 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         for _, session in self.session_manager.sessions.items():
             session.autopilot.handle_manual_edits(edits)
 
+    def onDeleteAtIndex(self, index: int):
+        for _, session in self.session_manager.sessions.items():
+            session.autopilot.delete_at_index(index)
+
     def onCommandOutput(self, output: str):
         # Send the output to ALL autopilots.
         # Maybe not ideal behavior
diff --git a/continuedev/src/continuedev/server/ide_protocol.py b/continuedev/src/continuedev/server/ide_protocol.py
index 2e1f78d7..dfdca504 100644
--- a/continuedev/src/continuedev/server/ide_protocol.py
+++ b/continuedev/src/continuedev/server/ide_protocol.py
@@ -96,7 +96,11 @@ class AbstractIdeProtocolServer(ABC):
         """Called when highlighted code is updated"""
 
     @abstractmethod
-    async def showDiff(self, filepath: str, replacement: str):
+    def onDeleteAtIndex(self, index: int):
+        """Called when a step is deleted at a given index"""
+
+    @abstractmethod
+    async def showDiff(self, filepath: str, replacement: str, step_index: int):
         """Show a diff"""
 
     @abstractproperty
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index f22297ae..10853828 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -305,7 +305,10 @@ class DefaultModelEditCodeStep(Step):
             full_suffix_lines = full_file_contents_lines[rif.range.end.line:]
             new_file_contents = "\n".join(
                 full_prefix_lines) + "\n" + completion + "\n" + "\n".join(full_suffix_lines)
-            await sdk.ide.showDiff(rif.filepath, new_file_contents)
+
+            step_index = sdk.history.current_index
+
+            await sdk.ide.showDiff(rif.filepath, new_file_contents, step_index)
 
         # Important state variables
         # -------------------------