4 files changed, 23 insertions, 15 deletions
diff --git a/extension/package-lock.json b/extension/package-lock.json
index 7b1ad703..71525620 100644
--- a/extension/package-lock.json
+++ b/extension/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "continue",
-  "version": "0.0.118",
+  "version": "0.0.120",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "continue",
-      "version": "0.0.118",
+      "version": "0.0.120",
       "license": "Apache-2.0",
       "dependencies": {
         "@electron/rebuild": "^3.2.10",
diff --git a/extension/package.json b/extension/package.json
index 548cdcd6..7475a56d 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -14,7 +14,7 @@
   "displayName": "Continue",
   "pricing": "Free",
   "description": "The open-source coding autopilot",
-  "version": "0.0.118",
+  "version": "0.0.120",
   "publisher": "Continue",
   "engines": {
     "vscode": "^1.67.0"
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 9b16a7a2..679d94ba 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -164,7 +164,7 @@ class IdeProtocolClient {
         this.showSuggestion(data.edit);
         break;
       case "showDiff":
-        this.showDiff(data.filepath, data.replacement);
+        this.showDiff(data.filepath, data.replacement, data.step_index);
         break;
       case "openGUI":
       case "connected":
@@ -243,8 +243,8 @@ class IdeProtocolClient {
     );
   }
 
-  showDiff(filepath: string, replacement: string) {
-    diffManager.writeDiff(filepath, replacement);
+  showDiff(filepath: string, replacement: string, step_index: number) {
+    diffManager.writeDiff(filepath, replacement, step_index);
   }
 
   openFile(filepath: string) {
@@ -431,6 +431,10 @@ class IdeProtocolClient {
   sendMainUserInput(input: string) {
     this.messenger?.send("mainUserInput", { input });
   }
+
+  deleteAtIndex(index: number) {
+    this.messenger?.send("deleteAtIndex", { index });
+  }
 }
 
 export default IdeProtocolClient;
diff --git a/extension/src/diffs.ts b/extension/src/diffs.ts
index 52a54046..dbfd8f59 100644
--- a/extension/src/diffs.ts
+++ b/extension/src/diffs.ts
@@ -8,6 +8,7 @@ interface DiffInfo {
   originalFilepath: string;
   newFilepath: string;
   editor?: vscode.TextEditor;
+  step_index: number;
 }
 
 export const DIFF_DIRECTORY = path.join(os.homedir(), ".continue", "diffs");
@@ -36,8 +37,7 @@ class DiffManager {
 
   private openDiffEditor(
     originalFilepath: string,
-    newFilepath: string,
-    newContent: string
+    newFilepath: string
   ): vscode.TextEditor | undefined {
     // If the file doesn't yet exist, don't open the diff editor
     if (!fs.existsSync(newFilepath)) {
@@ -62,7 +62,11 @@ class DiffManager {
     return editor;
   }
 
-  writeDiff(originalFilepath: string, newContent: string): string {
+  writeDiff(
+    originalFilepath: string,
+    newContent: string,
+    step_index: number
+  ): string {
     this.setupDirectory();
 
     // Create or update existing diff
@@ -77,6 +81,7 @@ class DiffManager {
       const diffInfo: DiffInfo = {
         originalFilepath,
         newFilepath,
+        step_index,
       };
       this.diffs.set(newFilepath, diffInfo);
     }
@@ -84,11 +89,7 @@ class DiffManager {
     // Open the editor if it hasn't been opened yet
     const diffInfo = this.diffs.get(newFilepath);
     if (diffInfo && !diffInfo?.editor) {
-      diffInfo.editor = this.openDiffEditor(
-        originalFilepath,
-        newFilepath,
-        newContent
-      );
+      diffInfo.editor = this.openDiffEditor(originalFilepath, newFilepath);
       this.diffs.set(newFilepath, diffInfo);
     }
 
@@ -101,7 +102,7 @@ class DiffManager {
       vscode.window.showTextDocument(diffInfo.editor.document);
       vscode.commands.executeCommand("workbench.action.closeActiveEditor");
     }
-    // this.diffs.delete(diffInfo.newFilepath);
+    this.diffs.delete(diffInfo.newFilepath);
     fs.unlinkSync(diffInfo.newFilepath);
   }
 
@@ -138,6 +139,9 @@ class DiffManager {
       return;
     }
 
+    // Stop the step at step_index in case it is still streaming
+    ideProtocolClient.deleteAtIndex(diffInfo.step_index);
+
     this.cleanUpDiff(diffInfo);
   }
 }