Merge branch 'main' of github.com:continuedev/continue

6 files changed, 67 insertions, 11 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 49fbacc5..e59c6ed4 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -12,7 +12,7 @@ from .observation import Observation, InternalErrorObservation
 from ..server.ide_protocol import AbstractIdeProtocolServer
 from ..libs.util.queue import AsyncSubscriptionQueue
 from ..models.main import ContinueBaseModel
-from .main import Context, ContinueCustomException, Policy, History, FullState, Step, HistoryNode
+from .main import Context, ContinueCustomException, HighlightedRangeContext, Policy, History, FullState, Step, HistoryNode
 from ..steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
 from ..libs.util.telemetry import capture_event
 from .sdk import ContinueSDK
@@ -140,11 +140,24 @@ class Autopilot(ContinueBaseModel):
                         tb_step.step_name, {"output": output, **tb_step.params})
                     await self._run_singular_step(step)
 
-    _highlighted_ranges: List[RangeInFileWithContents] = []
+    _highlighted_ranges: List[HighlightedRangeContext] = []
     _adding_highlighted_code: bool = False
 
+    def _make_sure_is_editing_range(self):
+        """If none of the highlighted ranges are currently being edited, the first should be selected"""
+        if len(self._highlighted_ranges) == 0:
+            return
+        if not any(map(lambda x: x.editing, self._highlighted_ranges)):
+            self._highlighted_ranges[0].editing = True
+
     async def handle_highlighted_code(self, range_in_files: List[RangeInFileWithContents]):
-        if not self._adding_highlighted_code:
+        if not self._adding_highlighted_code and len(self._highlighted_ranges) > 0:
+            return
+
+        # If un-highlighting, then remove the range
+        if len(self._highlighted_ranges) == 1 and len(range_in_files) == 1 and range_in_files[0].range.start == range_in_files[0].range.end:
+            self._highlighted_ranges = []
+            await self.update_subscribers()
             return
 
         # Filter out rifs from ~/.continue/diffs folder
@@ -160,20 +173,25 @@ class Autopilot(ContinueBaseModel):
         for i, rif in enumerate(self._highlighted_ranges):
             found_overlap = False
             for new_rif in range_in_files:
-                if rif.filepath == new_rif.filepath and rif.range.overlaps_with(new_rif.range):
+                if rif.range.filepath == new_rif.filepath and rif.range.range.overlaps_with(new_rif.range):
                     found_overlap = True
                     break
 
                 # Also don't allow multiple ranges in same file with same content. This is useless to the model, and avoids
                 # the bug where cmd+f causes repeated highlights
-                if rif.filepath == new_rif.filepath and rif.contents == new_rif.contents:
+                if rif.range.filepath == new_rif.filepath and rif.range.contents == new_rif.contents:
                     found_overlap = True
                     break
 
             if not found_overlap:
                 new_ranges.append(rif)
 
-        self._highlighted_ranges = new_ranges + range_in_files
+        self._highlighted_ranges = new_ranges + [HighlightedRangeContext(
+            range=rif, editing=False, pinned=False
+        ) for rif in range_in_files]
+
+        self._make_sure_is_editing_range()
+
         await self.update_subscribers()
 
     _step_depth: int = 0
@@ -193,12 +211,25 @@ class Autopilot(ContinueBaseModel):
             if i not in indices:
                 kept_ranges.append(rif)
         self._highlighted_ranges = kept_ranges
+
+        self._make_sure_is_editing_range()
+
         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 set_editing_at_indices(self, indices: List[int]):
+        for i in range(len(self._highlighted_ranges)):
+            self._highlighted_ranges[i].editing = i in indices
+        await self.update_subscribers()
+
+    async def set_pinned_at_indices(self, indices: List[int]):
+        for i in range(len(self._highlighted_ranges)):
+            self._highlighted_ranges[i].pinned = i in indices
+        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:
@@ -359,6 +390,10 @@ class Autopilot(ContinueBaseModel):
         if len(self._main_user_input_queue) > 1:
             return
 
+        # Remove context unless pinned
+        self._highlighted_ranges = [
+            hr for hr in self._highlighted_ranges if hr.pinned]
+
         # await self._request_halt()
         # Just run the step that takes user input, and
         # then up to the policy to decide how to deal with it.
diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
index 28fd964e..62cc4936 100644
--- a/continuedev/src/continuedev/core/main.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -199,13 +199,20 @@ class SlashCommandDescription(ContinueBaseModel):
     description: str
 
 
+class HighlightedRangeContext(ContinueBaseModel):
+    """Context for a highlighted range"""
+    range: RangeInFileWithContents
+    editing: bool
+    pinned: bool
+
+
 class FullState(ContinueBaseModel):
     """A full state of the program, including the history"""
     history: History
     active: bool
     user_input_queue: List[str]
     default_model: str
-    highlighted_ranges: List[RangeInFileWithContents]
+    highlighted_ranges: List[HighlightedRangeContext]
     slash_commands: List[SlashCommandDescription]
     adding_highlighted_code: bool
 
diff --git a/continuedev/src/continuedev/models/generate_json_schema.py b/continuedev/src/continuedev/models/generate_json_schema.py
index 080787a5..6cebf429 100644
--- a/continuedev/src/continuedev/models/generate_json_schema.py
+++ b/continuedev/src/continuedev/models/generate_json_schema.py
@@ -1,7 +1,7 @@
 from .main import *
 from .filesystem import RangeInFile, FileEdit
 from .filesystem_edit import FileEditWithFullContents
-from ..core.main import History, HistoryNode
+from ..core.main import History, HistoryNode, FullState
 from pydantic import schema_json_of
 import os
 
@@ -12,7 +12,7 @@ MODELS_TO_GENERATE = [
 ] + [
     FileEditWithFullContents
 ] + [
-    History, HistoryNode
+    History, HistoryNode, FullState
 ]
 
 RENAMES = {
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index fa573b37..8e9b1fb9 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -87,6 +87,10 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
                 self.on_delete_context_at_indices(data["indices"])
             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_pinned_at_indices":
+                self.on_set_pinned_at_indices(data["indices"])
         except Exception as e:
             print(e)
 
@@ -135,6 +139,16 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
             self.session.autopilot.toggle_adding_highlighted_code()
         )
 
+    def on_set_editing_at_indices(self, indices: List[int]):
+        asyncio.create_task(
+            self.session.autopilot.set_editing_at_indices(indices)
+        )
+
+    def on_set_pinned_at_indices(self, indices: List[int]):
+        asyncio.create_task(
+            self.session.autopilot.set_pinned_at_indices(indices)
+        )
+
 
 @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 3a7c8876..cbf0fe5e 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -453,7 +453,7 @@ class DefaultModelEditCodeStep(Step):
             lines.extend(map(lambda l: common_whitespace + l, chunk_lines))
 
             if True:
-                await sendDiffUpdate(lines, sdk)
+                await sendDiffUpdate(lines + [common_whitespace + unfinished_line], sdk)
 
             # Deal with newly accumulated lines
             for line in chunk_lines:
diff --git a/continuedev/src/continuedev/steps/main.py b/continuedev/src/continuedev/steps/main.py
index 3cf78c40..5ccffbfe 100644
--- a/continuedev/src/continuedev/steps/main.py
+++ b/continuedev/src/continuedev/steps/main.py
@@ -270,7 +270,7 @@ class EditHighlightedCodeStep(Step):
         if len(range_in_files) == 0:
             # Create a new file
             new_file_path = "new_file.txt"
-            await sdk.add_file(new_file_path)
+            await sdk.add_file(new_file_path, "")
             range_in_files = [RangeInFile.from_entire_file(new_file_path, "")]
 
         await sdk.run_step(DefaultModelEditCodeStep(user_input=self.user_input, range_in_files=range_in_files))