custom slash command expansion bug fixed

3 files changed, 13 insertions, 4 deletions

diff --git a/continuedev/src/continuedev/core/policy.py b/continuedev/src/continuedev/core/policy.py
index 7a72b596..b8363df2 100644
--- a/continuedev/src/continuedev/core/policy.py
+++ b/continuedev/src/continuedev/core/policy.py
@@ -46,7 +46,7 @@ def parse_custom_command(inp: str, config: ContinueConfig) -> Union[None, Step]:
             slash_command = parse_slash_command(custom_cmd.prompt, config)
             if slash_command is not None:
                 return slash_command
-            return CustomCommandStep(name=custom_cmd.name, description=custom_cmd.description, prompt=custom_cmd.prompt, user_input=after_command)
+            return CustomCommandStep(name=custom_cmd.name, description=custom_cmd.description, prompt=custom_cmd.prompt, user_input=after_command, slash_command=command_name)
     return None
 
 
diff --git a/continuedev/src/continuedev/steps/chat.py b/continuedev/src/continuedev/steps/chat.py
index e4fc3f72..37d1b1cc 100644
--- a/continuedev/src/continuedev/steps/chat.py
+++ b/continuedev/src/continuedev/steps/chat.py
@@ -22,10 +22,12 @@ class SimpleChatStep(Step):
     name: str = "Generating Response..."
     manage_own_chat_context: bool = True
     description: str = ""
+    messages: List[ChatMessage] = None
 
     async def run(self, sdk: ContinueSDK):
         completion = ""
-        async for chunk in sdk.models.gpt4.stream_chat(await sdk.get_chat_context()):
+        messages = self.messages or await sdk.get_chat_context()
+        async for chunk in sdk.models.gpt4.stream_chat(messages):
             if sdk.current_step_was_deleted():
                 return
 
diff --git a/continuedev/src/continuedev/steps/custom_command.py b/continuedev/src/continuedev/steps/custom_command.py
index b162679e..5a56efb0 100644
--- a/continuedev/src/continuedev/steps/custom_command.py
+++ b/continuedev/src/continuedev/steps/custom_command.py
@@ -8,6 +8,7 @@ class CustomCommandStep(Step):
     name: str
     prompt: str
     user_input: str
+    slash_command: str
     hide: bool = True
 
     async def describe(self):
@@ -15,5 +16,11 @@ class CustomCommandStep(Step):
 
     async def run(self, sdk: ContinueSDK):
         prompt_user_input = f"Task: {self.prompt}. Additional info: {self.user_input}"
-        await sdk.run_step(UserInputStep(user_input=self.user_input))
-        await sdk.run_step(SimpleChatStep())
+        messages = await sdk.get_chat_context()
+        # Find the last chat message with this slash command and replace it with the user input
+        for i in range(len(messages) - 1, -1, -1):
+            if messages[i].role == "user" and messages[i].content.startswith(self.slash_command):
+                messages[i] = messages[i].copy(
+                    update={"content": prompt_user_input})
+                break
+        await sdk.run_step(SimpleChatStep(messages=messages))