highlighted code improvements

4 files changed, 19 insertions, 1 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 313ceded..29be3b79 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -69,7 +69,8 @@ class Autopilot(ContinueBaseModel):
             user_input_queue=self._main_user_input_queue,
             default_model=self.continue_sdk.config.default_model,
             highlighted_ranges=self._highlighted_ranges,
-            slash_commands=self.get_available_slash_commands()
+            slash_commands=self.get_available_slash_commands(),
+            adding_highlighted_code=self._adding_highlighted_code,
         )
 
     def get_available_slash_commands(self) -> List[Dict]:
@@ -140,8 +141,12 @@ class Autopilot(ContinueBaseModel):
                     await self._run_singular_step(step)
 
     _highlighted_ranges: List[RangeInFileWithContents] = []
+    _adding_highlighted_code: bool = False
 
     async def handle_highlighted_code(self, range_in_files: List[RangeInFileWithContents]):
+        if not self._adding_highlighted_code:
+            return
+
         workspace_path = self.continue_sdk.ide.workspace_directory
         for rif in range_in_files:
             rif.filepath = os.path.basename(rif.filepath)
@@ -186,6 +191,10 @@ class Autopilot(ContinueBaseModel):
         self._highlighted_ranges = kept_ranges
         await self.update_subscribers()
 
+    async def toggle_adding_highlighted_code(self):
+        self._adding_highlighted_code = not self._adding_highlighted_code
+        await self.update_subscribers()
+
     async def _run_singular_step(self, step: "Step", is_future_step: bool = False) -> Coroutine[Observation, None, None]:
         # Allow config to set disallowed steps
         if step.__class__.__name__ in self.continue_sdk.config.disallowed_steps:
diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
index 8bad09d1..28fd964e 100644
--- a/continuedev/src/continuedev/core/main.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -207,6 +207,7 @@ class FullState(ContinueBaseModel):
     default_model: str
     highlighted_ranges: List[RangeInFileWithContents]
     slash_commands: List[SlashCommandDescription]
+    adding_highlighted_code: bool
 
 
 class ContinueSDK:
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 4e960f7c..fa573b37 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -85,6 +85,8 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
                 self.on_delete_at_index(data["index"])
             elif message_type == "delete_context_at_indices":
                 self.on_delete_context_at_indices(data["indices"])
+            elif message_type == "toggle_adding_highlighted_code":
+                self.on_toggle_adding_highlighted_code()
         except Exception as e:
             print(e)
 
@@ -128,6 +130,11 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
             self.session.autopilot.delete_context_at_indices(indices)
         )
 
+    def on_toggle_adding_highlighted_code(self):
+        asyncio.create_task(
+            self.session.autopilot.toggle_adding_highlighted_code()
+        )
+
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(websocket_session)):
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index b215b317..c74412ba 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -286,6 +286,7 @@ class DefaultModelEditCodeStep(Step):
         return "```" in line or "<modified_code_to_edit>" in line or "<file_prefix>" in line or "</file_prefix>" in line or "<file_suffix>" in line or "</file_suffix>" in line or "<user_request>" in line or "</user_request>" in line or "<code_to_edit>" in line
 
     async def stream_rif(self, rif: RangeInFileWithContents, sdk: ContinueSDK):
+        await sdk.ide.saveFile(rif.filepath)
         full_file_contents = await sdk.ide.readFile(rif.filepath)
 
         file_prefix, contents, file_suffix, model_to_use = await self.get_prompt_parts(