From 1d7163f2222ef6dea01c347ce9a432a7c8a13de3 Mon Sep 17 00:00:00 2001
From: Hope Hadfield <hhadfiel@redhat.com>
Date: Tue, 3 Sep 2024 11:50:19 -0400
Subject: [PATCH] Add support for SnippetTextEdit in code actions (#3730)

- Add snippetEditSupport extended client capability from JDT-LS
- Handle possible code action conversion (LSP to VS Code) in middleware.resolveCodeAction
---
 src/extension.ts | 62 ++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 57 insertions(+), 5 deletions(-)

diff --git a/src/extension.ts b/src/extension.ts
index 3e10278a2..85f08bba0 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -6,8 +6,8 @@ import * as fse from 'fs-extra';
 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, IndentAction, InputBoxOptions, languages, MarkdownString, QuickPickItemKind, RelativePattern, TextDocument, TextEditorRevealType, UIKind, Uri, ViewColumn, window, workspace, WorkspaceConfiguration } from 'vscode';
-import { CancellationToken, CodeActionParams, CodeActionRequest, Command, CompletionRequest, DidChangeConfigurationNotification, ExecuteCommandParams, ExecuteCommandRequest, LanguageClientOptions, RevealOutputChannelOn } from 'vscode-languageclient';
+import { CodeActionContext, commands, CompletionItem, ConfigurationTarget, Diagnostic, env, EventEmitter, ExtensionContext, extensions, IndentAction, InputBoxOptions, languages, MarkdownString, QuickPickItemKind, Range, RelativePattern, SnippetString, SnippetTextEdit, TextDocument, TextEditorRevealType, UIKind, Uri, ViewColumn, window, workspace, WorkspaceConfiguration, WorkspaceEdit } from 'vscode';
+import { CancellationToken, CodeActionParams, CodeActionRequest, CodeActionResolveRequest, Command, CompletionRequest, DidChangeConfigurationNotification, ExecuteCommandParams, ExecuteCommandRequest, LanguageClientOptions, RevealOutputChannelOn } from 'vscode-languageclient';
 import { LanguageClient } from 'vscode-languageclient/node';
 import { apiManager } from './apiManager';
 import { ClientErrorHandler } from './clientErrorHandler';
@@ -235,6 +235,7 @@ export async function activate(context: ExtensionContext): Promise<ExtensionAPI>
 						extractInterfaceSupport: true,
 						advancedUpgradeGradleSupport: true,
 						executeClientCommandSupport: true,
+						snippetEditSupport: true,
 					},
 					triggerFiles,
 				},
@@ -303,7 +304,60 @@ export async function activate(context: ExtensionContext): Promise<ExtensionAPI>
 						}, (error) => {
 							return client.handleFailedRequest(CodeActionRequest.type, token, error, []);
 						});
-					}
+					},
+
+					resolveCodeAction: async (item, token, next) => {
+						const client: LanguageClient = standardClient.getClient();
+						const documentUris = [];
+						const snippetEdits = [];
+						const resolveCodeAction = async (item, token) => {
+							return client.sendRequest(CodeActionResolveRequest.type, client.code2ProtocolConverter.asCodeActionSync(item), token).then(async (result) => {
+								if (token.isCancellationRequested) {
+									return item;
+								}
+								const docChanges = result.edit !== undefined ? result.edit.documentChanges : undefined;
+								if (docChanges !== undefined) {
+									for (const docChange of docChanges) {
+										if ("textDocument" in docChange) {
+											for (const edit of docChange.edits) {
+												if ("snippet" in edit) {
+													documentUris.push(docChange.textDocument.uri);
+													snippetEdits.push(new SnippetTextEdit(client.protocol2CodeConverter.asRange((edit as any).range), new SnippetString((edit as any).snippet.value)));
+												}
+											}
+										}
+									}
+									const codeAction = await client.protocol2CodeConverter.asCodeAction(result, token);
+									const docEdits = codeAction.edit !== undefined? codeAction.edit.entries() : [];
+									for (const docEdit of docEdits) {
+										const uri = docEdit[0];
+										if (documentUris.includes(uri.toString())) {
+											const editList = [];
+											for (const edit of docEdit[1]) {
+												let isSnippet = false;
+												snippetEdits.forEach((snippet, index) => {
+													if (edit.range.isEqual(snippet.range) && documentUris[index] === uri.toString()) {
+														editList.push(snippet);
+														isSnippet = true;
+													}
+												});
+												if (!isSnippet) {
+													editList.push(edit);
+												}
+											}
+											codeAction.edit.set(uri, null);
+											codeAction.edit.set(uri, editList);
+										}
+									}
+									return codeAction;
+								}
+								return await client.protocol2CodeConverter.asCodeAction(result, token);
+							}, (error) => {
+								return client.handleFailedRequest(CodeActionResolveRequest.type, token, error, item);
+							});
+						 };
+						return resolveCodeAction(item, token);
+					},
 				},
 				revealOutputChannelOn: RevealOutputChannelOn.Never,
 				errorHandler: new ClientErrorHandler(extensionName),
@@ -1189,5 +1243,3 @@ function registerRestartJavaLanguageServerCommand(context: ExtensionContext) {
 		}
 	}));
 }
-
-