3 files changed, 22 insertions, 14 deletions
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index f6cc129e..bbf93f65 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -10,7 +10,7 @@ import * as vscode from "vscode";
 import fkill from "fkill";
 import { sendTelemetryEvent, TelemetryEvent } from "../telemetry";
 
-const MAX_RETRIES = 0;
+const MAX_RETRIES = 3;
 async function retryThenFail(
   fn: () => Promise<any>,
   retries: number = MAX_RETRIES
@@ -197,7 +197,7 @@ async function setupPythonEnv() {
       } else if (stderr) {
         if (stderr.includes("running scripts is disabled on this system")) {
           vscode.window.showErrorMessage(
-             "A Python virtual enviroment cannot be activated because running scripts is disabled for this user. Please enable signed scripts to run with this command in PowerShell: `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser`, reload VS Code, and then try again."
+            "A Python virtual enviroment cannot be activated because running scripts is disabled for this user. Please enable signed scripts to run with this command in PowerShell: `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser`, reload VS Code, and then try again."
           );
         }
         throw new Error(stderr);
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);
   }
 }