Con-991-a: Diff Manager Improvements

extensions/vscode/src/commands.ts

@@ -257,7 +257,7 @@ const commandsMap: (
         "getDefaultModelTitle",
         undefined,
       );
-      await verticalDiffManager.streamEdit(text, modelName);
+      await verticalDiffManager.streamEdit(text, modelName, previousInput);
     } else {
       // Pick context first
       const selectedProviders = await vscode.window.showQuickPick(
@@ -302,7 +302,7 @@ const commandsMap: (
           "\n\n---\n\n" +
           text;
 
-        await verticalDiffManager.streamEdit(text, defaultModelTitle);
+        await verticalDiffManager.streamEdit(text, defaultModelTitle, previousInput);
       }
     }
   },

extensions/vscode/src/diff/verticalPerLine/handler.ts

@@ -208,7 +208,21 @@ export class VerticalPerLineDiffHandler {
     this.editor.setDecorations(indexDecorationType, []);
   }
 
-  clear(accept: boolean) {
+  public getLineDeltaBeforeLine(line: number) {
+    //Returns the number of lines removed from a file when the diff currently active is closed 
+    let totalLineDelta = 0
+    for (const range of this.greenDecorationManager.getRanges().sort((a, b) => a.start.line - b.start.line)) {
+      if (range.start.line > line){
+        break
+      }
+
+      totalLineDelta -= range.end.line - range.start.line + 1
+    }
+
+    return totalLineDelta
+  }
+
+  async clear(accept: boolean) {
     vscode.commands.executeCommand(
       "setContext",
       "continue.streamingDiff",
@@ -224,7 +238,7 @@ export class VerticalPerLineDiffHandler {
 
     this.editorToVerticalDiffCodeLens.delete(this.filepath);
 
-    this.editor.edit(
+    await this.editor.edit(
       (editBuilder) => {
         for (const range of rangesToDelete) {
           editBuilder.delete(
@@ -361,6 +375,11 @@ export class VerticalPerLineDiffHandler {
     this.redDecorationManager.shiftDownAfterLine(startLine, offset);
     this.greenDecorationManager.shiftDownAfterLine(startLine, offset);
 
+    // Shift the codelens objects
+    this.shiftCodeLensObjects(startLine, offset)
+  }
+
+  private shiftCodeLensObjects(startLine: number, offset: number){  
     // Shift the codelens objects
     const blocks =
       this.editorToVerticalDiffCodeLens
@@ -376,4 +395,19 @@ export class VerticalPerLineDiffHandler {
 
     this.refreshCodeLens();
   }
+
+  public updateLineDelta(filepath: string, startLine: number, lineDelta: number) {
+    // Retrieve the diff blocks for the given file
+    const blocks = this.editorToVerticalDiffCodeLens.get(filepath);
+    if (!blocks) {
+      return;
+    }
+
+    //update decorations
+    this.redDecorationManager.shiftDownAfterLine(startLine, lineDelta);
+    this.greenDecorationManager.shiftDownAfterLine(startLine, lineDelta);
+
+    //update code lens
+    this.shiftCodeLensObjects(startLine, lineDelta)
+  }
 }

extensions/vscode/src/diff/verticalPerLine/manager.ts

@@ -19,7 +19,11 @@ export class VerticalPerLineDiffManager {
 
   filepathToCodeLens: Map<string, VerticalDiffCodeLens[]> = new Map();
 
-  constructor(private readonly configHandler: ConfigHandler) {}
+  private userChangeListener: vscode.Disposable | undefined;
+
+  constructor(private readonly configHandler: ConfigHandler) {
+    this.userChangeListener = undefined
+  }
 
   createVerticalPerLineDiffHandler(
     filepath: string,
@@ -49,34 +53,49 @@ export class VerticalPerLineDiffManager {
     }
   }
 
-  getOrCreateVerticalPerLineDiffHandler(
-    filepath: string,
-    startLine: number,
-    endLine: number,
-  ) {
-    if (this.filepathToHandler.has(filepath)) {
-      return this.filepathToHandler.get(filepath)!;
-    } else {
-      const editor = vscode.window.activeTextEditor; // TODO
-      if (editor && editor.document.uri.fsPath === filepath) {
-        const handler = new VerticalPerLineDiffHandler(
-          startLine,
-          endLine,
-          editor,
-          this.filepathToCodeLens,
-          this.clearForFilepath.bind(this),
-          this.refreshCodeLens,
-        );
-        this.filepathToHandler.set(filepath, handler);
-        return handler;
-      } else {
-        return undefined;
+  getHandlerForFile(filepath: string) {
+    return this.filepathToHandler.get(filepath);
+  }
+
+  // Creates a listener for document changes by user.
+  private enableDocumentChangeListener(): vscode.Disposable | undefined {
+    if (this.userChangeListener) { //Only create one listener per file
+      return;
+    }
+
+    this.userChangeListener = vscode.workspace.onDidChangeTextDocument((event) => {
+      // Check if there is an active handler for the affected file
+      const filepath = event.document.uri.fsPath;
+      const handler = this.getHandlerForFile(filepath);
+      if (handler ) {
+        // If there is an active diff for that file, handle the document change
+        this.handleDocumentChange(event, handler);
       }
+    }); 
+  }
+
+  // Listener for user doc changes is disabled during updates to the text document by continue
+  public disableDocumentChangeListener() {
+    if (this.userChangeListener) {
+        this.userChangeListener.dispose();
+        this.userChangeListener = undefined;
     }
   }
 
-  getHandlerForFile(filepath: string) {
-    return this.filepathToHandler.get(filepath);
+  private handleDocumentChange(
+    event: vscode.TextDocumentChangeEvent,
+    handler: VerticalPerLineDiffHandler
+  ) {
+    // Loop through each change in the event
+    event.contentChanges.forEach((change) => {
+      // Calculate the number of lines added or removed
+      const linesAdded = change.text.split('\n').length - 1;
+      const linesDeleted = change.range.end.line - change.range.start.line;
+      const lineDelta = linesAdded - linesDeleted;
+
+      // Update the diff handler with the new line delta
+      handler.updateLineDelta(event.document.uri.fsPath, change.range.start.line, lineDelta);
+    });
   }
 
   clearForFilepath(filepath: string | undefined, accept: boolean) {
@@ -94,13 +113,15 @@ export class VerticalPerLineDiffManager {
       this.filepathToHandler.delete(filepath);
     }
 
+    this.disableDocumentChangeListener()
+
     vscode.commands.executeCommand("setContext", "continue.diffVisible", false);
   }
 
-  acceptRejectVerticalDiffBlock(
+  async acceptRejectVerticalDiffBlock(
     accept: boolean,
     filepath?: string,
-    index?: number,
+    index?: number
   ) {
     if (!filepath) {
       const activeEditor = vscode.window.activeTextEditor;
@@ -125,8 +146,11 @@ export class VerticalPerLineDiffManager {
       return;
     }
 
+    // Disable listening to file changes while continue makes changes
+    this.disableDocumentChangeListener()
+
     // CodeLens object removed from editorToVerticalDiffCodeLens here
-    handler.acceptRejectBlock(
+    await handler.acceptRejectBlock(
       accept,
       block.start,
       block.numGreen,
@@ -135,40 +159,75 @@ export class VerticalPerLineDiffManager {
 
     if (blocks.length === 1) {
       this.clearForFilepath(filepath, true);
+    } else {
+      // Re-enable listener for user changes to file
+      this.enableDocumentChangeListener()
     }
   }
 
-  async streamEdit(input: string, modelTitle: string | undefined) {
+  async streamEdit(input: string, modelTitle: string | undefined, quickEdit?: string) {
     vscode.commands.executeCommand("setContext", "continue.diffVisible", true);
 
-    const editor = vscode.window.activeTextEditor;
+    let editor = vscode.window.activeTextEditor;
+
     if (!editor) {
       return;
     }
 
     const filepath = editor.document.uri.fsPath;
-    const startLine = editor.selection.start.line;
-    const endLine = editor.selection.end.line;
+
+    // Initialize start/end at editor selection
+    let startLine = editor.selection.start.line
+    let endLine = editor.selection.end.line
 
+    // Check for existing handlers in the same file the new one will be created in
     const existingHandler = this.getHandlerForFile(filepath);
-    existingHandler?.clear(false);
+
+    if (existingHandler) {
+      let existingStartLine = existingHandler?.range.start.line
+      if (quickEdit) { // Previous diff was a quickEdit
+        // Check if user has highlighted a range   
+        let rangeBool =  (startLine != endLine) ||(editor.selection.start.character != editor.selection.end.character)
+
+        // Check if the range is different from the previous range
+        let newRangeBool = (startLine != existingHandler.range.start.line) || (endLine != existingHandler.range.end.line) 
+
+        if (!rangeBool || !newRangeBool) { 
+          // User did not highlight a new range -> use start/end from the previous quickEdit
+          startLine = existingHandler.range.start.line 
+          endLine = existingHandler.range.end.line
+        }
+      }
+
+      // Clear the previous handler 
+      let effectiveLineDelta = existingHandler.getLineDeltaBeforeLine(startLine)
+      existingHandler.clear(false)
+
+      startLine += effectiveLineDelta 
+      endLine += effectiveLineDelta
+    }
+
     await new Promise((resolve) => {
       setTimeout(resolve, 200);
     });
+
+    // Create new handler with determined start/end
     const diffHandler = this.createVerticalPerLineDiffHandler(
       filepath,
-      existingHandler?.range.start.line ?? startLine,
-      existingHandler?.range.end.line ?? endLine,
+      startLine,
+      endLine,
       input,
     );
+
     if (!diffHandler) {
+      console.warn("Issue occured while creating new vertical diff handler")
       return;
     }
 
-    let selectedRange = existingHandler?.range ?? editor.selection;
+    let selectedRange = diffHandler.range 
 
     // Only if the selection is empty, use exact prefix/suffix instead of by line
-    if (!selectedRange.isEmpty) {
+    if (selectedRange.isEmpty) {
       selectedRange = new vscode.Range(
         editor.selection.start.with(undefined, 0),
         editor.selection.end.with(undefined, Number.MAX_SAFE_INTEGER),
@@ -218,8 +277,11 @@ export class VerticalPerLineDiffManager {
           getMarkdownLanguageTagForFile(filepath),
         ),
       );
+
+      // enable a listener for user edits to file while diff is open
+      this.enableDocumentChangeListener()
     } catch (e) {
-      console.error("Error streaming diff:", e);
+      this.disableDocumentChangeListener()
       vscode.window.showErrorMessage(`Error streaming diff: ${e}`);
     } finally {
       vscode.commands.executeCommand(

extensions/vscode/src/lang-server/codeLens.ts

@@ -45,12 +45,12 @@ class VerticalPerLineCodeLensProvider implements vscode.CodeLensProvider {
         codeLenses.push(
           new vscode.CodeLens(range, {
             title: `Accept All (${getMetaKeyLabel()}⇧↩)`,
-            command: "continue.acceptVerticalDiffBlock",
+            command: "continue.acceptDiff",
             arguments: [filepath, i],
           }),
           new vscode.CodeLens(range, {
             title: `Reject All (${getMetaKeyLabel()}⇧⌫)`,
-            command: "continue.rejectVerticalDiffBlock",
+            command: "continue.rejectDiff",
             arguments: [filepath, i],
           }),
         );