VS Code API active text editor is undefined after clicking on a webview

Question

I'm developing a Visual Studio Code extension project that involves interacting with a webview to provide code refactoring suggestions generated by an LLM. On the webview the code suggestion is displayed and the user can click 'Accept' button to replace the selected code with the new suggestion. However, I've encountered an issue where the webview becomes active after clicking anywhere on the webview, including sometimes when clicking on the 'Accept' button. Then when 'Accept' command is executed the selected code is not replaced because vscode.window.activeTextEditor is undefined.

The issue is intermittent, doesn't happen every time!

Here's the relevant code from my extension:

commands.ts:

import * as vscode from "vscode";
import { getLoadingWebviewContent, getWebviewContent } from "./utils/utils";
import { llmProxyConnector } from "./connector/LLMProxyConnector";

class Commands {
  private panel: vscode.WebviewPanel | undefined;
  private suggestion: string | undefined;

  public acceptSuggestionCommand() {
    const editor = vscode.window.activeTextEditor;

    if (this.panel?.webview && editor && this.suggestion) {
      const { selection } = editor;

      if (selection) {
        editor.edit((editBuilder) => {
          editBuilder.replace(selection, this.suggestion as string);
        });
        this.panel.dispose();
        vscode.window.showInformationMessage("Accepted suggestion");
      } else {
        vscode.window.showInformationMessage(
          "No selection in the active editor."
        );
      }
    } else {
      vscode.window.showInformationMessage(
        "No active editor or suggestion available."
      );
    }
  }

  public async showSuggestionsCommand() {
    const editor = vscode.window.activeTextEditor;

    if (editor) {
      const selectedText = editor.document.getText(editor.selection);

      if (selectedText) {
        vscode.window.showInformationMessage(
          "Generating suggestion for the code you have selected."
        );
        if (!this.panel) {
          this.panel = vscode.window.createWebviewPanel(
            "suggestionPanel",
            "Code Suggestion",
            { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true },
            {
              enableScripts: true,
              retainContextWhenHidden: true,
            }
          );
        }
        this.panel.webview.html = getLoadingWebviewContent();

        const response = await llmProxyConnector.fetchLLMSuggestions(
          selectedText
        );

        const { codeSnippet, explanation } =
          this.parseCodeAndExplanation(response.choices[0].text);
        this.suggestion = codeSnippet;
        this.panel.webview.html = getWebviewContent(
          selectedText,
          this.suggestion,
          explanation
        );

        this.panel.webview.onDidReceiveMessage(
          this.onDidReceiveMessageCallback.bind(this)
        );
        this.panel.onDidDispose(() => {
          this.panel = undefined;
        });
      } else {
        vscode.window.showInformationMessage("Please select a text first!");
      }
    }
  }

  // Other command methods...

  private parseCodeAndExplanation(input: string) {
    const delimiter = "<===>";
    const parts = input.split(delimiter);

    const codeSnippet = parts[0].trim();
    const explanation = parts[1].trim();
    return { codeSnippet, explanation };
  }

  private onDidReceiveMessageCallback(message: any) {
    switch (message.command) {
      case "acceptSuggestion":
        vscode.commands.executeCommand("ai-code-simplifier.acceptSuggestion");
        break;
      case "declineSuggestion":
        vscode.commands.executeCommand("ai-code-simplifier.declineSuggestion");
        break;
      case "webviewClose":
        this.panel?.dispose();
        break;
      case "focusOnEditor":
        vscode.commands.executeCommand(
          "workbench.action.focusActiveEditorGroup"
        );
        break;
    }
  }
}

export const commands = new Commands();

webviewScript.ts:

export const webviewScript = `
  const vscode = acquireVsCodeApi();

  function acceptSuggestion() {
      vscode.postMessage({ command: 'acceptSuggestion' });
  }

  function declineSuggestion() {
      vscode.postMessage({ command: 'declineSuggestion' });
  }

  window.addEventListener('message', (event) => {
    const message = event.data;

    switch (message.command) {
      case 'acceptSuggestion':
        acceptSuggestion();
        break;
      case 'declineSuggestion':
        declineSuggestion();
        break;
      case 'webviewClose':
        vscode.postMessage({ command: 'webviewClose' });
        break;
    }
  });

  document.body.addEventListener('click', (event) => {
    event.stopPropagation();
    setTimeout(() => {
      vscode.postMessage({ command: 'focusOnEditor' });
    }, 0);
  });
`;

In commands.ts, the showSuggestionsCommand method creates a webview panel when the user selects text in the editor. The onDidReceiveMessageCallback method handles messages received from the webview, including commands to accept or decline suggestions.

Steps to Reproduce:

• Start the extension
• Open a text file.
• Select some text in the editor.
• Trigger the command to show suggestions (e.g., Ctrl+Shift+P -> Show Code Suggestions).

Expected Results:

The webview panel should appear with code suggestions based on the selected text. Clicking the 'Accept' button within the webview panel should replace the selected text in the editor with the suggested code snippet. After accepting or declining a suggestion, the focus should remain on the editor, and activeTextEditor should remain defined.

Actual Results:

Sometimes, after clicking on the webview panel or even on 'Accept' button, the activeTextEditor becomes undefined, webview tab becomes active instead of the editor and the selected text is not replaced. The message "No active editor or suggestion available." is displayed when attempting to execute the 'Accept' command.

Additional Context:

To mitigate the issue, I've attempted to focus back on the editor by executing the 'workbench.action.focusActiveEditorGroup' command when clicking on the webview, but this doesn't always prevent 'activeTextEditor' from becoming undefined.

Environment:

VS Code version: 1.84.2 NodeJS version: 16.14.2 Operating System: macOS Sonoma 14.1

I'm seeking assistance in identifying the root cause of this issue and finding a reliable solution to ensure that 'activeTextEditor' remains defined after interacting with the webview.