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

7 files changed, 131 insertions, 83 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 313ceded..b1c4f471 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,16 @@ 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
+
+        # Filter out rifs from ~/.continue/diffs folder
+        range_in_files = [
+            rif for rif in range_in_files if not os.path.dirname(rif.filepath) == os.path.expanduser("~/.continue/diffs")]
+
         workspace_path = self.continue_sdk.ide.workspace_directory
         for rif in range_in_files:
             rif.filepath = os.path.basename(rif.filepath)
@@ -186,6 +195,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/config.py b/continuedev/src/continuedev/core/config.py
index 9e8541dc..ff7b8cb0 100644
--- a/continuedev/src/continuedev/core/config.py
+++ b/continuedev/src/continuedev/core/config.py
@@ -33,11 +33,11 @@ DEFAULT_SLASH_COMMANDS = [
         description="Edit code in the current file or the highlighted code",
         step_name="EditHighlightedCodeStep",
     ),
-    SlashCommand(
-        name="explain",
-        description="Reply to instructions or a question with previous steps and the highlighted code or current file as context",
-        step_name="SimpleChatStep",
-    ),
+    # SlashCommand(
+    #     name="explain",
+    #     description="Reply to instructions or a question with previous steps and the highlighted code or current file as context",
+    #     step_name="SimpleChatStep",
+    # ),
     SlashCommand(
         name="config",
         description="Open the config file to create new and edit existing slash commands",
@@ -71,7 +71,10 @@ class ContinueConfig(BaseModel):
     allow_anonymous_telemetry: Optional[bool] = True
     default_model: Literal["gpt-3.5-turbo", "gpt-3.5-turbo-16k",
                            "gpt-4"] = 'gpt-4'
-    custom_commands: Optional[List[CustomCommand]] = []
+    custom_commands: Optional[List[CustomCommand]] = [CustomCommand(
+        name="test",
+        prompt="Write a comprehensive set of unit tests for the selected code. It should setup, run tests that check for correctness including important edge cases, and teardown. Ensure that the tests are complete and sophisticated. Give the tests just as chat output, don't edit any file.",
+    )]
     slash_commands: Optional[List[SlashCommand]] = DEFAULT_SLASH_COMMANDS
     on_traceback: Optional[List[OnTracebackSteps]] = [
         OnTracebackSteps(step_name="DefaultOnTracebackStep")]
@@ -126,7 +129,7 @@ def load_global_config() -> ContinueConfig:
         config_path = os.path.join(global_dir, 'config.json')
         if not os.path.exists(config_path):
             with open(config_path, 'w') as f:
-                json.dump(dict(ContinueConfig()), f)
+                json.dump(ContinueConfig().dict(), f)
         with open(config_path, 'r') as f:
             try:
                 config_dict = json.load(f)
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/core/policy.py b/continuedev/src/continuedev/core/policy.py
index b0853380..fc9266ab 100644
--- a/continuedev/src/continuedev/core/policy.py
+++ b/continuedev/src/continuedev/core/policy.py
@@ -1,5 +1,5 @@
 from textwrap import dedent
-from typing import List, Tuple, Type
+from typing import List, Tuple, Type, Union
 
 from ..steps.welcome import WelcomeStep
 from .config import ContinueConfig
@@ -22,6 +22,34 @@ from ..libs.util.step_name_to_steps import get_step_from_name
 from ..steps.custom_command import CustomCommandStep
 
 
+def parse_slash_command(inp: str, config: ContinueConfig) -> Union[None, Step]:
+    """
+    Parses a slash command, returning the command name and the rest of the input.
+    """
+    if inp.startswith("/"):
+        command_name = inp.split(" ")[0]
+        after_command = " ".join(inp.split(" ")[1:])
+
+        for slash_command in config.slash_commands:
+            if slash_command.name == command_name[1:]:
+                params = slash_command.params
+                params["user_input"] = after_command
+                return get_step_from_name(slash_command.step_name, params)
+    return None
+
+
+def parse_custom_command(inp: str, config: ContinueConfig) -> Union[None, Step]:
+    command_name = inp.split(" ")[0]
+    after_command = " ".join(inp.split(" ")[1:])
+    for custom_cmd in config.custom_commands:
+        if custom_cmd.name == command_name[1:]:
+            slash_command = parse_slash_command(custom_cmd.prompt, config)
+            if slash_command is not None:
+                return slash_command
+            return CustomCommandStep(name=custom_cmd.name, prompt=custom_cmd.prompt, user_input=after_command)
+    return None
+
+
 class DemoPolicy(Policy):
     ran_code_last: bool = False
 
@@ -46,34 +74,14 @@ class DemoPolicy(Policy):
             # This could be defined with ObservationTypePolicy. Ergonomics not right though.
             user_input = observation.user_input
 
-            if user_input.startswith("/"):
-                command_name = user_input.split(" ")[0]
-                after_command = " ".join(user_input.split(" ")[1:])
-                for slash_command in config.slash_commands:
-                    if slash_command.name == command_name[1:]:
-                        params = slash_command.params
-                        params["user_input"] = after_command
-                        return get_step_from_name(slash_command.step_name, params)
-
-                for custom_cmd in config.custom_commands:
-                    if custom_cmd.name == command_name[1:]:
-                        return CustomCommandStep(name=custom_cmd.name, prompt=custom_cmd.prompt, user_input=after_command)
+            slash_command = parse_slash_command(user_input, config)
+            if slash_command is not None:
+                return slash_command
 
-            # return EditHighlightedCodeStep(user_input=user_input)
-            return ChatWithFunctions(user_input=user_input)
-            return NLDecisionStep(user_input=user_input, steps=[
-                (EditHighlightedCodeStep(user_input=user_input),
-                 "Edit the highlighted code"),
-                # AnswerQuestionChroma(question=user_input),
-                # EditFileChroma(request=user_input),
-                (SimpleChatStep(user_input=user_input),
-                 "Respond to the user with a chat message. Can answer questions about code or anything else."),
-            ], default_step=EditHighlightedCodeStep(user_input=user_input))
+            custom_command = parse_custom_command(user_input, config)
+            if custom_command is not None:
+                return custom_command
 
-        state = history.get_current()
+            return SimpleChatStep(user_input=user_input)
 
-        if observation is not None and isinstance(observation, TracebackObservation):
-            self.ran_code_last = False
-            return SolveTracebackStep(traceback=observation.traceback)
-        else:
-            return None
+        return None
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/chat.py b/continuedev/src/continuedev/steps/chat.py
index db3f9d7f..8f88244c 100644
--- a/continuedev/src/continuedev/steps/chat.py
+++ b/continuedev/src/continuedev/steps/chat.py
@@ -107,12 +107,8 @@ class RunTerminalCommandStep(Step):
     description: str = "Run a terminal command."
     command: str
 
-    async def describe(self, models: Models) -> Coroutine[Any, Any, Coroutine[str, None, None]]:
-        return f"Ran the terminal command `{self.command}`."
-
     async def run(self, sdk: ContinueSDK):
-        await sdk.wait_for_user_confirmation(f"Run the following terminal command?\n\n```bash\n{self.command}\n```")
-        await sdk.run(self.command)
+        self.description = f"Copy this command and run in your terminal:\n\n```bash\n{self.command}\n```"
 
 
 class ViewDirectoryTreeStep(Step):
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index b9f0da35..3a7c8876 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(
@@ -295,6 +296,17 @@ class DefaultModelEditCodeStep(Step):
         prompt = self.compile_prompt(file_prefix, contents, file_suffix, sdk)
         full_file_contents_lines = full_file_contents.split("\n")
 
+        async def sendDiffUpdate(lines: List[str], sdk: ContinueSDK):
+            nonlocal full_file_contents_lines, rif
+
+            completion = "\n".join(lines)
+
+            full_prefix_lines = full_file_contents_lines[:rif.range.start.line]
+            full_suffix_lines = full_file_contents_lines[rif.range.end.line + 1:]
+            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)
+
         # Important state variables
         # -------------------------
         original_lines = [] if rif.contents == "" else rif.contents.split("\n")
@@ -320,8 +332,9 @@ class DefaultModelEditCodeStep(Step):
 
             # Highlight the line to show progress
             line_to_highlight = current_line_in_file - len(current_block_lines)
-            await sdk.ide.highlightCode(RangeInFile(filepath=rif.filepath, range=Range.from_shorthand(
-                line_to_highlight, 0, line_to_highlight, 0)), "#FFFFFF22" if len(current_block_lines) == 0 else "#00FF0022")
+            if False:
+                await sdk.ide.highlightCode(RangeInFile(filepath=rif.filepath, range=Range.from_shorthand(
+                    line_to_highlight, 0, line_to_highlight, 0)), "#FFFFFF22" if len(current_block_lines) == 0 else "#00FF0022")
 
             if len(current_block_lines) == 0:
                 # Set this as the start of the next block
@@ -370,12 +383,14 @@ class DefaultModelEditCodeStep(Step):
                 replacement = "\n".join(current_block_lines)
                 start_line = current_block_start
                 end_line = current_block_start + index_of_last_line_in_block
-                await sdk.ide.showSuggestion(FileEdit(
-                    filepath=rif.filepath,
-                    range=Range.from_shorthand(
-                        start_line, 0, end_line, 0),
-                    replacement=replacement
-                ))
+
+                if False:
+                    await sdk.ide.showSuggestion(FileEdit(
+                        filepath=rif.filepath,
+                        range=Range.from_shorthand(
+                            start_line, 0, end_line, 0),
+                        replacement=replacement
+                    ))
 
                 # Reset current block / update variables
                 current_line_in_file += 1
@@ -435,16 +450,16 @@ class DefaultModelEditCodeStep(Step):
                 chunk_lines.pop()  # because this will be an empty string
             else:
                 unfinished_line = chunk_lines.pop()
-            lines.extend(chunk_lines)
+            lines.extend(map(lambda l: common_whitespace + l, chunk_lines))
+
+            if True:
+                await sendDiffUpdate(lines, sdk)
 
             # Deal with newly accumulated lines
             for line in chunk_lines:
                 # Trailing whitespace doesn't matter
                 line = line.rstrip()
 
-                # Add the common whitespace that was removed before prompting
-                line = common_whitespace + line
-
                 # Lines that should signify the end of generation
                 if self.is_end_line(line):
                     break
@@ -463,7 +478,9 @@ class DefaultModelEditCodeStep(Step):
                     break
 
                 # If none of the above, insert the line!
-                await handle_generated_line(line)
+                if False:
+                    await handle_generated_line(line)
+
                 completion_lines_covered += 1
                 current_line_in_file += 1
 
@@ -475,34 +492,37 @@ class DefaultModelEditCodeStep(Step):
             completion_lines_covered += 1
             current_line_in_file += 1
 
-        # If the current block isn't empty, add that suggestion
-        if len(current_block_lines) > 0:
-            # We have a chance to back-track here for blank lines that are repeats of the end of the original
-            # Don't want to have the same ending in both the original and the generated, can just leave it there
-            num_to_remove = 0
-            for i in range(-1, -len(current_block_lines) - 1, -1):
-                if len(original_lines_below_previous_blocks) == 0:
-                    break
-                if current_block_lines[i] == original_lines_below_previous_blocks[-1]:
-                    num_to_remove += 1
-                    original_lines_below_previous_blocks.pop()
-                else:
-                    break
-            current_block_lines = current_block_lines[:-
-                                                      num_to_remove] if num_to_remove > 0 else current_block_lines
-
-            # It's also possible that some lines match at the beginning of the block
-            # while len(current_block_lines) > 0 and len(original_lines_below_previous_blocks) > 0 and current_block_lines[0] == original_lines_below_previous_blocks[0]:
-            #     current_block_lines.pop(0)
-            #     original_lines_below_previous_blocks.pop(0)
-            #     current_block_start += 1
-
-            await sdk.ide.showSuggestion(FileEdit(
-                filepath=rif.filepath,
-                range=Range.from_shorthand(
-                    current_block_start, 0, current_block_start + len(original_lines_below_previous_blocks), 0),
-                replacement="\n".join(current_block_lines)
-            ))
+        await sendDiffUpdate(lines, sdk)
+
+        if False:
+            # If the current block isn't empty, add that suggestion
+            if len(current_block_lines) > 0:
+                # We have a chance to back-track here for blank lines that are repeats of the end of the original
+                # Don't want to have the same ending in both the original and the generated, can just leave it there
+                num_to_remove = 0
+                for i in range(-1, -len(current_block_lines) - 1, -1):
+                    if len(original_lines_below_previous_blocks) == 0:
+                        break
+                    if current_block_lines[i] == original_lines_below_previous_blocks[-1]:
+                        num_to_remove += 1
+                        original_lines_below_previous_blocks.pop()
+                    else:
+                        break
+                current_block_lines = current_block_lines[:-
+                                                          num_to_remove] if num_to_remove > 0 else current_block_lines
+
+                # It's also possible that some lines match at the beginning of the block
+                # while len(current_block_lines) > 0 and len(original_lines_below_previous_blocks) > 0 and current_block_lines[0] == original_lines_below_previous_blocks[0]:
+                #     current_block_lines.pop(0)
+                #     original_lines_below_previous_blocks.pop(0)
+                #     current_block_start += 1
+
+                await sdk.ide.showSuggestion(FileEdit(
+                    filepath=rif.filepath,
+                    range=Range.from_shorthand(
+                        current_block_start, 0, current_block_start + len(original_lines_below_previous_blocks), 0),
+                    replacement="\n".join(current_block_lines)
+                ))
 
         # Record the completion
         completion = "\n".join(lines)