Fix the issue that cannot apply changes after switching editors

efewijum · Mar 22, 2023 · 7d3d7d3 · 7d3d7d3
1 parent 2526f59
commit 7d3d7d3
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 16 deletions.
diff --git a/src/extension.ts b/src/extension.ts
@@ -16,7 +16,6 @@ async function handleGenerateCodeCommand() {
     if (!editor) {
         return;
     }
-    const selection = editor.selection;
 
     // End the active session first.
     const globalState = getGlobalState();
@@ -25,7 +24,7 @@ async function handleGenerateCodeCommand() {
         activeSession.dispose();
     }
 
-    const session = new GenerateSession(input, selection, editor);
+    const session = new GenerateSession(input, editor);
     session.start();
     session.showResult();
     globalState.activeSession = session;

diff --git a/src/generate/core.ts b/src/generate/core.ts
@@ -26,11 +26,11 @@ class SelectionRange implements ISelectionRange {
 
 export async function generateCode(
     prompt: string,
-    editor: vscode.TextEditor,
+    document: vscode.TextDocument,
+    selection: vscode.Selection,
     cancellationToken: vscode.CancellationToken,
     resultStream: ResultStream<String>
 ): Promise<void> {
-    const { document, selection } = editor;
     const filePath = document.uri.fsPath;
     const workspaceDirectory =
         vscode.workspace.getWorkspaceFolder(document.uri)?.uri.fsPath ?? null;

diff --git a/src/generate/generateSession.ts b/src/generate/generateSession.ts
@@ -6,21 +6,18 @@ import { generateCode } from "./core";
 export class GenerateSession {
     #prompt: string;
     #selection: vscode.Selection;
-    #editor: vscode.TextEditor;
+    #document: vscode.TextDocument;
     #scratchpad: Scratchpad | null;
     #errorOccurred = false;
     #statusBarItem: vscode.StatusBarItem | null = null;
 
-    constructor(
-        prompt: string,
-        selection: vscode.Selection,
-        editor: vscode.TextEditor
-    ) {
-        const selectionText = editor.document.getText(selection);
+    constructor(prompt: string, editor: vscode.TextEditor) {
+        const { document, selection } = editor;
+        const selectionText = document.getText(selection);
 
         this.#prompt = prompt;
         this.#selection = selection;
-        this.#editor = editor;
+        this.#document = document;
         this.#scratchpad = new Scratchpad(selectionText);
     }
 
@@ -50,7 +47,8 @@ export class GenerateSession {
                 try {
                     await generateCode(
                         this.#prompt,
-                        this.#editor,
+                        this.#document,
+                        this.#selection,
                         token,
                         scratchpad
                     );
@@ -119,8 +117,8 @@ export class GenerateSession {
 
         const thisUriString = scratchpad.uri.toString();
         const tabGroups = vscode.window.tabGroups;
-        for (let tabGroup of tabGroups.all) {
-            for (let tab of tabGroup.tabs) {
+        for (const tabGroup of tabGroups.all) {
+            for (const tab of tabGroup.tabs) {
                 const tabInput = tab.input;
                 if (!(tabInput instanceof vscode.TabInputTextDiff)) {
                     continue;
@@ -134,6 +132,16 @@ export class GenerateSession {
         return null;
     }
 
+    #getEditor(): vscode.TextEditor | null {
+        const documentUriString = this.#document.uri.toString();
+        for (const editor of vscode.window.visibleTextEditors) {
+            if (editor.document.uri.toString() === documentUriString) {
+                return editor;
+            }
+        }
+        return null;
+    }
+
     async #showGenerationDecisionMessage() {
         if (!this.#statusBarItem) {
             const statusBarItem = vscode.window.createStatusBarItem(
@@ -196,8 +204,16 @@ export class GenerateSession {
             return;
         }
 
+        const editor = this.#getEditor();
+        if (!editor) {
+            vscode.window.showWarningMessage(
+                "Need to active the original text buffer before applying changes."
+            );
+            return;
+        }
+
         // TODO: reconcile with the modified document.
-        this.#editor
+        editor
             .edit((editBuilder) => {
                 editBuilder.replace(this.#selection, scratchpad.contents);
             })