explain insert at cursor, better diff streaming

3 files changed, 72 insertions, 11 deletions

diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index b0d9d719..d4d067ba 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -265,6 +265,23 @@ class DefaultModelEditCodeStep(Step):
         return file_prefix, rif.contents, file_suffix, model_to_use, max_tokens
 
     def compile_prompt(self, file_prefix: str, contents: str, file_suffix: str, sdk: ContinueSDK) -> str:
+        if contents.strip() == "":
+            # Seperate prompt for insertion at the cursor, the other tends to cause it to repeat whole file
+            prompt = dedent(f"""\
+<file_prefix>
+{file_prefix}
+</file_prefix>
+<insertion_code_here>
+<file_suffix>
+{file_suffix}
+</file_suffix>
+<user_request>
+{self.user_input}
+</user_request>
+
+Please output the code to be inserted at the cursor in order to fulfill the user_request. Do NOT preface your answer or write anything other than code. You should not write any tags, just the code. Make sure to correctly indent the code:""")
+            return prompt
+
         prompt = self._prompt
         if file_prefix.strip() != "":
             prompt += dedent(f"""
@@ -306,15 +323,32 @@ 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
+        lines_to_display = []
+
+        async def sendDiffUpdate(lines: List[str], sdk: ContinueSDK, final: bool = False):
+            nonlocal full_file_contents_lines, rif, lines_to_display
 
             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:]
+
+            # Don't do this at the very end, just show the inserted code
+            if final:
+                lines_to_display = []
+            # Only recalculate at every new-line, because this is sort of expensive
+            elif completion.endswith("\n"):
+                contents_lines = rif.contents.split("\n")
+                rewritten_lines = 0
+                for line in lines:
+                    for i in range(rewritten_lines, len(contents_lines)):
+                        if difflib.SequenceMatcher(None, line, contents_lines[i]).ratio() > 0.7 and contents_lines[i].strip() != "":
+                            rewritten_lines = i + 1
+                            break
+                lines_to_display = contents_lines[rewritten_lines:]
+
             new_file_contents = "\n".join(
-                full_prefix_lines) + "\n" + completion + "\n" + "\n".join(full_suffix_lines)
+                full_prefix_lines) + "\n" + completion + "\n" + "\n".join(lines_to_display) + "\n" + "\n".join(full_suffix_lines)
 
             step_index = sdk.history.current_index
 
@@ -495,7 +529,7 @@ class DefaultModelEditCodeStep(Step):
                 completion_lines_covered += 1
                 current_line_in_file += 1
 
-            await sendDiffUpdate(lines + [common_whitespace + unfinished_line], sdk)
+            await sendDiffUpdate(lines + [common_whitespace if unfinished_line.startswith("<") else (common_whitespace + unfinished_line)], sdk)
 
         # Add the unfinished line
         if unfinished_line != "" and not self.line_to_be_ignored(unfinished_line, completion_lines_covered == 0) and not self.is_end_line(unfinished_line):
@@ -505,7 +539,7 @@ class DefaultModelEditCodeStep(Step):
             completion_lines_covered += 1
             current_line_in_file += 1
 
-        await sendDiffUpdate(lines, sdk)
+        await sendDiffUpdate(lines, sdk, final=True)
 
         if False:
             # If the current block isn't empty, add that suggestion
diff --git a/extension/react-app/src/components/ComboBox.tsx b/extension/react-app/src/components/ComboBox.tsx
index 801c3a03..585a0584 100644
--- a/extension/react-app/src/components/ComboBox.tsx
+++ b/extension/react-app/src/components/ComboBox.tsx
@@ -180,6 +180,26 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
 
   useImperativeHandle(ref, () => downshiftProps, [downshiftProps]);
 
+  const [metaKeyPressed, setMetaKeyPressed] = useState(false);
+  useEffect(() => {
+    const handleKeyDown = (e: KeyboardEvent) => {
+      if (e.key === "Meta") {
+        setMetaKeyPressed(true);
+      }
+    };
+    const handleKeyUp = (e: KeyboardEvent) => {
+      if (e.key === "Meta") {
+        setMetaKeyPressed(false);
+      }
+    };
+    window.addEventListener("keydown", handleKeyDown);
+    window.addEventListener("keyup", handleKeyUp);
+    return () => {
+      window.removeEventListener("keydown", handleKeyDown);
+      window.removeEventListener("keyup", handleKeyUp);
+    };
+  });
+
   useEffect(() => {
     if (!inputRef.current) {
       return;
@@ -272,7 +292,7 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
       <div className="flex px-2" ref={divRef} hidden={!downshiftProps.isOpen}>
         <MainTextInput
           disabled={props.disabled}
-          placeholder="Ask a question, give instructions, or type '/' to see slash commands"
+          placeholder="Ask a question, give instructions, or type '/' to see slash commands. ⌘⏎ to edit."
           {...getInputProps({
             onChange: (e) => {
               const target = e.target as HTMLTextAreaElement;
@@ -357,10 +377,12 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
             ))}
         </Ul>
       </div>
-      {/* <span className="text-trueGray-400 ml-auto m-auto text-xs text-right">
-        Highlight code to include as context. Currently open file included by
-        default. {highlightedCodeSections.length === 0 && ""}
-      </span> */}
+      {highlightedCodeSections.length === 0 &&
+        (downshiftProps.inputValue?.startsWith("/edit") || metaKeyPressed) && (
+          <div className="text-trueGray-400 pr-4 text-xs text-right">
+            Inserting at cursor
+          </div>
+        )}
       <ContextDropdown
         onMouseEnter={() => {
           setHoveringContextDropdown(true);
diff --git a/extension/src/diffs.ts b/extension/src/diffs.ts
index 1dc292e1..3ea6b4f8 100644
--- a/extension/src/diffs.ts
+++ b/extension/src/diffs.ts
@@ -164,7 +164,12 @@ class DiffManager {
     // Stop the step at step_index in case it is still streaming
     ideProtocolClient.deleteAtIndex(diffInfo.step_index);
 
-    this.cleanUpDiff(diffInfo);
+    vscode.workspace.textDocuments
+      .find((doc) => doc.uri.fsPath === newFilepath)
+      ?.save()
+      .then(() => {
+        this.cleanUpDiff(diffInfo);
+      });
   }
 }