Reveal the edited file after applying a code action

src/commands.ts

@@ -74,6 +74,11 @@ export namespace Commands {
      */
     export const APPLY_WORKSPACE_EDIT = 'java.apply.workspaceEdit';
 
+    /**
+     * Reveal the change made by a code action if it is in a document other than the active one
+     */
+    export const REVEAL_EDIT = 'java.reveal.edit';
+
     /**
      * Execute Workspace Command
      */

src/extension.ts

@@ -8,7 +8,7 @@ import * as os from 'os';
 import * as path from 'path';
 import * as semver from 'semver';
 import {
-	CodeActionContext, commands, CompletionItem, ConfigurationTarget, Diagnostic, env, EventEmitter, ExtensionContext, extensions,
+	CodeAction, CodeActionContext, commands, CompletionItem, ConfigurationTarget, Diagnostic, env, EventEmitter, ExtensionContext, extensions,
 	InputBoxOptions, Location, MarkdownString, RelativePattern,
 	SnippetString, SnippetTextEdit, TextDocument, TextEditorRevealType, UIKind, Uri, version, ViewColumn, window, workspace, WorkspaceConfiguration
 } from 'vscode';
@@ -325,7 +325,7 @@ export async function activate(context: ExtensionContext): Promise<ExtensionAPI>
 							result.push(client.protocol2CodeConverter.asCommand(item));
 						}
 						else {
-							result.push(await client.protocol2CodeConverter.asCodeAction(item));
+							result.push(addRevealCommand(client, await client.protocol2CodeConverter.asCodeAction(item), item.edit));
 						}
 					}
 					return result;
@@ -384,9 +384,9 @@ export async function activate(context: ExtensionContext): Promise<ExtensionAPI>
 								codeAction.edit.set(uri, editList);
 							}
 						}
-						return codeAction;
+						return addRevealCommand(client, codeAction, result.edit);
 					}
-					return await client.protocol2CodeConverter.asCodeAction(result, token);
+					return addRevealCommand(client, await client.protocol2CodeConverter.asCodeAction(result, token), result.edit);
 				}, (error) => {
 					return client.handleFailedRequest(CodeActionResolveRequest.type, token, error, item);
 				});
@@ -1376,3 +1376,38 @@ export function escapeSnippetLiterals(value: string): string {
         .replace(/\\(?!,)/g, '\\\\')       // Escape backslashes, but preserve escaped commas
         .replace(/\$(?!\{)/g, '\\$'); // Escape $ only if NOT followed by {
 }
+
+/**
+ * VS Code applies the edit of a code action itself and doesn't show the edited file.
+ * Attach a command, which VS Code runs after applying the edit, to reveal the change
+ * when it is made in a document other than the active one.
+ * Any command the action already has is passed along and run first.
+ */
+function addRevealCommand(client: LanguageClient, codeAction: CodeAction, lspEdit: { documentChanges?: any[], changes?: { [uri: string]: any[] } } | undefined): CodeAction {
+	if (lspEdit && codeAction.command?.command !== Commands.REVEAL_EDIT) {
+		// Taken from the raw edit because WorkspaceEdit.entries() doesn't include snippet edits
+		let target: { uri: string, range: any } | undefined;
+		for (const docChange of lspEdit.documentChanges ?? []) {
+			if ("textDocument" in docChange && docChange.edits.length > 0) {
+				target = { uri: docChange.textDocument.uri, range: docChange.edits[0].range };
+				break;
+			}
+		}
+		if (!target) {
+			for (const [uri, edits] of Object.entries(lspEdit.changes ?? {})) {
+				if (edits.length > 0) {
+					target = { uri, range: edits[0].range };
+					break;
+				}
+			}
+		}
+		if (target) {
+			codeAction.command = {
+				title: 'Reveal edit',
+				command: Commands.REVEAL_EDIT,
+				arguments: [Uri.parse(target.uri), client.protocol2CodeConverter.asRange(target.range), codeAction.command],
+			};
+		}
+	}
+	return codeAction;
+}

src/sourceAction.ts

@@ -1,7 +1,7 @@
 'use strict';
 
-import { Disposable, ExtensionContext, TextEditorRevealType, Uri, ViewColumn, commands, window, workspace } from 'vscode';
-import { CodeActionParams, WorkspaceEdit } from 'vscode-languageclient';
+import { Disposable, ExtensionContext, Uri, ViewColumn, commands, window } from 'vscode';
+import { CodeActionParams } from 'vscode-languageclient';
 import { LanguageClient } from 'vscode-languageclient/node';
 import { Commands } from './commands';
 import { getActiveLanguageClient } from './extension';
@@ -25,7 +25,7 @@ import {
     OrganizeImportsRequest,
     VariableBinding
 } from './protocol';
-import { applyWorkspaceEdit } from './standardLanguageClient';
+import { applyWorkspaceEdit, revealWorkspaceEdit } from './standardLanguageClient';
 
 export function registerCommands(languageClient: LanguageClient, context: ExtensionContext) {
     registerOverrideMethodsCommand(languageClient, context);
@@ -412,18 +412,3 @@ function registerGenerateDelegateMethodsCommand(languageClient: LanguageClient,
         await revealWorkspaceEdit(workspaceEdit, languageClient);
     }));
 }
-
-async function revealWorkspaceEdit(workspaceEdit: WorkspaceEdit, languageClient: LanguageClient): Promise<void> {
-    const codeWorkspaceEdit = await languageClient.protocol2CodeConverter.asWorkspaceEdit(workspaceEdit);
-    if (!codeWorkspaceEdit) {
-        return;
-    }
-    for (const entry of codeWorkspaceEdit.entries()) {
-        await workspace.openTextDocument(entry[0]);
-        if (entry[1].length > 0) {
-            // reveal first available change of the workspace edit
-            window.activeTextEditor.revealRange(entry[1][0].range, TextEditorRevealType.InCenter);
-            break;
-        }
-    }
-}

src/standardLanguageClient.ts

@@ -2,7 +2,7 @@
 
 import * as net from 'net';
 import * as path from 'path';
-import { CancellationToken, CodeActionKind, commands, ConfigurationTarget, DocumentSelector, EventEmitter, ExtensionContext, extensions, languages, Location, ProgressLocation, TextEditor, Uri, ViewColumn, window, workspace, WorkspaceConfiguration } from "vscode";
+import { CancellationToken, CodeActionKind, Command, commands, ConfigurationTarget, DocumentSelector, EventEmitter, ExtensionContext, extensions, languages, Location, ProgressLocation, Range, TextEditor, TextEditorRevealType, Uri, ViewColumn, window, workspace, WorkspaceConfiguration } from "vscode";
 import { ConfigurationParams, ConfigurationRequest, LanguageClientOptions, Location as LSLocation, MessageType, Position as LSPosition, TextDocumentPositionParams, WorkspaceEdit, StaticFeature, ClientCapabilities, FeatureState, TelemetryEventNotification } from "vscode-languageclient";
 import { LanguageClient, StreamInfo } from "vscode-languageclient/node";
 import { apiManager } from "./apiManager";
@@ -438,8 +438,30 @@ export class StandardLanguageClient {
 
 		context.subscriptions.push(commands.registerCommand(Commands.NULL_ANALYSIS_SET_MODE, (status) => setNullAnalysisStatus(status)));
 
-		context.subscriptions.push(commands.registerCommand(Commands.APPLY_WORKSPACE_EDIT, (obj) => {
-			applyWorkspaceEdit(obj, this.languageClient);
+		context.subscriptions.push(commands.registerCommand(Commands.APPLY_WORKSPACE_EDIT, async (obj) => {
+			const activeUri = window.activeTextEditor?.document.uri.toString();
+			await applyWorkspaceEdit(obj, this.languageClient);
+			await revealWorkspaceEdit(obj, this.languageClient, activeUri);
+		}));
+
+		context.subscriptions.push(commands.registerCommand(Commands.REVEAL_EDIT, async (uri: Uri, range: Range, originalCommand?: Command) => {
+			if (originalCommand) {
+				await commands.executeCommand(originalCommand.command, ...(originalCommand.arguments ?? []));
+			}
+			if (window.activeTextEditor?.document.uri.toString() === uri.toString()) {
+				return;
+			}
+			// The edit may still be being applied (e.g. a snippet edit opens the file in the background),
+			// so open the document and retry until its editor is the active one.
+			const document = await workspace.openTextDocument(uri);
+			for (let i = 0; i < 10; i++) {
+				const editor = await window.showTextDocument(document, { preview: false });
+				if (window.activeTextEditor?.document.uri.toString() === uri.toString()) {
+					editor.revealRange(range, TextEditorRevealType.InCenter);
+					return;
+				}
+				await new Promise(resolve => setTimeout(resolve, 100));
+			}
 		}));
 
 		context.subscriptions.push(commands.registerCommand(Commands.NAVIGATE_TO_SUPER_IMPLEMENTATION_COMMAND, async (location: LinkLocation | Uri) => {
@@ -927,6 +949,27 @@ export async function applyWorkspaceEdit(workspaceEdit: WorkspaceEdit, languageC
 	}
 }
 
+/**
+ * Reveals the first change of the workspace edit in an editor.
+ * If `onlyIfNotUri` is given, nothing happens when the change is in that document.
+ */
+export async function revealWorkspaceEdit(workspaceEdit: WorkspaceEdit, languageClient: LanguageClient, onlyIfNotUri?: string): Promise<void> {
+	const codeWorkspaceEdit = await languageClient.protocol2CodeConverter.asWorkspaceEdit(workspaceEdit);
+	if (!codeWorkspaceEdit) {
+		return;
+	}
+	for (const [uri, edits] of codeWorkspaceEdit.entries()) {
+		if (edits.length > 0) {
+			if (uri.toString() !== onlyIfNotUri) {
+				// reveal first available change of the workspace edit
+				const editor = await window.showTextDocument(uri);
+				editor.revealRange(edits[0].range, TextEditorRevealType.InCenter);
+			}
+			break;
+		}
+	}
+}
+
 /**
  * 'workspace/willRenameFiles' already handled so we need to disable it.
  * @see fileEventHandler.registerFileEventHandlers