Suffix server command ids on the client and add the Marketplace icon

package.json

@@ -1,7 +1,7 @@
 {
     "name": "aws-iac-vscode",
     "displayName": "AWS CloudFormation",
-    "version": "1.0.0",
+    "version": "1.1.0",
     "description": "AWS CloudFormation language support: validation, completion, hover documentation, and cfn-lint diagnostics for JSON and YAML templates.",
     "publisher": "amazonwebservices",
     "author": "Amazon Web Services (AWS)",
@@ -16,6 +16,7 @@
     },
     "qna": "https://github.com/aws-cloudformation/cloudformation-vscode/issues",
     "pricing": "Free",
+    "icon": "resources/marketplace/aws-cloudformation-icon.png",
     "galleryBanner": {
         "color": "#232F3E",
         "theme": "dark"

src/extension.ts

@@ -4,6 +4,7 @@ import { AwsCredentialsService } from './auth/AwsCredentials';
 import { restartCommand, RestartServerCommand, updateRegion } from './commands/Commands';
 import { CfnInlineCompletionProvider } from './inlineCompletion/InlineCompletionProvider';
 import { CfnDevLspServerProvider } from './lsp-server/CfnDevLspServerProvider';
+import { CfnLanguageClient } from './lsp-server/CfnLanguageClient';
 import { CfnRemoteLspServerProvider } from './lsp-server/CfnRemoteLspServerProvider';
 import { CfnDocumentSelector, cfnInitializationOptions, cfnServerOptions } from './lsp-server/LspClientConfig';
 import { LspLauncher, LspResolver } from './lsp-server/LspLauncher';
@@ -129,7 +130,7 @@ async function initialize(context: ExtensionContext) {
                 errorHandler,
             };
 
-            return new LanguageClient(ExtensionId, ExtensionName, cfnServerOptions(serverFile), clientOptions);
+            return new CfnLanguageClient(ExtensionId, ExtensionName, cfnServerOptions(serverFile), clientOptions);
         };
 
         const launcher = new LspLauncher(resolver, clientFactory, {

src/lsp-server/CfnLanguageClient.ts

@@ -0,0 +1,30 @@
+import { ExecuteCommandRequest } from 'vscode-languageclient';
+import {
+    DynamicFeature,
+    LanguageClient,
+    LanguageClientOptions,
+    ServerOptions,
+    StaticFeature,
+} from 'vscode-languageclient/node';
+import { advertisedCommands } from './LspCommands';
+import { suffixCodeActionCommands, SuffixedExecuteCommandFeature } from './SuffixedExecuteCommandFeature';
+
+/** A `LanguageClient` that registers the server's command ids with VS Code under this client's suffix. */
+export class CfnLanguageClient extends LanguageClient {
+    constructor(id: string, name: string, serverOptions: ServerOptions, clientOptions: LanguageClientOptions) {
+        super(id, name, serverOptions, clientOptions);
+        this.middleware.provideCodeActions = async (document, range, context, token, next) => {
+            const actions = await next(document, range, context, token);
+            if (actions) {
+                suffixCodeActionCommands(actions, advertisedCommands(this));
+            }
+            return actions;
+        };
+    }
+
+    override registerFeature(feature: StaticFeature | DynamicFeature<unknown>): void {
+        const replaced =
+            'registrationType' in feature && feature.registrationType.method === ExecuteCommandRequest.method;
+        super.registerFeature(replaced ? new SuffixedExecuteCommandFeature(this) : feature);
+    }
+}

src/lsp-server/LspCommands.ts

@@ -3,16 +3,28 @@ import { CfnLspClientName } from '../utils/ExtensionConfig';
 
 export const UpdateRegionCommand = '/command/region/update';
 
+const CommandSuffix = `.${CfnLspClientName}`;
+
+/** The ids the running server advertised through `executeCommandProvider`. */
+export function advertisedCommands(client: LanguageClient): string[] {
+    return client.initializeResult?.capabilities.executeCommandProvider?.commands ?? [];
+}
+
 /**
  * Returns the id under which the running server registered `baseCommand`: suffixed when the server honours
  * `aws.commandSuffix`, unsuffixed when it predates the option.
  */
 export function advertisedCommand(client: LanguageClient, baseCommand: string): string {
-    const advertised = client.initializeResult?.capabilities.executeCommandProvider?.commands ?? [];
-    const suffixed = `${baseCommand}.${CfnLspClientName}`;
+    const advertised = advertisedCommands(client);
+    const suffixed = suffixedCommand(baseCommand);
     const command = [suffixed, baseCommand].find((candidate) => advertised.includes(candidate));
     if (!command) {
         throw new Error(`Language server does not advertise ${suffixed} (advertised: ${advertised.join(', ')})`);
     }
     return command;
 }
+
+/** The id this client registers `command` with VS Code under, whether or not the server already suffixed it. */
+export function suffixedCommand(command: string): string {
+    return command.endsWith(CommandSuffix) ? command : `${command}${CommandSuffix}`;
+}

src/lsp-server/SuffixedExecuteCommandFeature.ts

@@ -0,0 +1,97 @@
+import { randomUUID } from 'crypto';
+import { CodeAction, Command, commands, Disposable } from 'vscode';
+import { ExecuteCommandRequest } from 'vscode-languageclient';
+import {
+    ClientCapabilities,
+    DynamicFeature,
+    ExecuteCommandRegistrationOptions,
+    FeatureState,
+    LanguageClient,
+    RegistrationData,
+    ServerCapabilities,
+} from 'vscode-languageclient/node';
+import { suffixedCommand } from './LspCommands';
+
+type CommandClient = Pick<LanguageClient, 'sendRequest' | 'handleFailedRequest'>;
+
+/**
+ * Replaces the client library's execute-command feature. The server's command ids are registered with VS Code under
+ * this client's suffix even when the server predates `aws.commandSuffix` and advertises unsuffixed ids, so another
+ * extension hosting the same server in the window cannot collide with ours. Executions are forwarded to the server
+ * under the id it advertised.
+ */
+export class SuffixedExecuteCommandFeature implements DynamicFeature<ExecuteCommandRegistrationOptions> {
+    private readonly registrations = new Map<string, Disposable[]>();
+
+    constructor(private readonly client: CommandClient) {}
+
+    get registrationType() {
+        return ExecuteCommandRequest.type;
+    }
+
+    getState(): FeatureState {
+        return { kind: 'workspace', id: ExecuteCommandRequest.method, registrations: this.registrations.size > 0 };
+    }
+
+    fillClientCapabilities(capabilities: ClientCapabilities): void {
+        (capabilities.workspace ??= {}).executeCommand = { dynamicRegistration: true };
+    }
+
+    initialize(capabilities: ServerCapabilities): void {
+        if (capabilities.executeCommandProvider) {
+            this.register({ id: randomUUID(), registerOptions: capabilities.executeCommandProvider });
+        }
+    }
+
+    register(data: RegistrationData<ExecuteCommandRegistrationOptions>): void {
+        this.registrations.set(
+            data.id,
+            data.registerOptions.commands.map((command) =>
+                commands.registerCommand(
+                    suffixedCommand(command),
+                    async (...args: unknown[]) => await this.execute(command, args),
+                ),
+            ),
+        );
+    }
+
+    unregister(id: string): void {
+        for (const registration of this.registrations.get(id) ?? []) {
+            registration.dispose();
+        }
+        this.registrations.delete(id);
+    }
+
+    clear(): void {
+        for (const id of this.registrations.keys()) {
+            this.unregister(id);
+        }
+    }
+
+    private async execute(command: string, args: unknown[]): Promise<unknown> {
+        try {
+            return await this.client.sendRequest<unknown>(ExecuteCommandRequest.method, { command, arguments: args });
+        } catch (error) {
+            // `defaultValue` is a required parameter: the client returns it for errors it swallows and rethrows the rest.
+            // eslint-disable-next-line unicorn/no-useless-undefined
+            return this.client.handleFailedRequest<unknown>(ExecuteCommandRequest.type, undefined, error, undefined);
+        }
+    }
+}
+
+/**
+ * Points the commands of server-provided code actions at the ids {@link SuffixedExecuteCommandFeature} registered.
+ * Commands the server did not advertise, such as ones contributed by the extension itself, are left alone.
+ */
+export function suffixCodeActionCommands(actions: (Command | CodeAction)[], advertised: readonly string[]): void {
+    for (const action of actions) {
+        const command = isCommand(action) ? action : action.command;
+        if (command && advertised.includes(command.command)) {
+            command.command = suffixedCommand(command.command);
+        }
+    }
+}
+
+function isCommand(action: Command | CodeAction): action is Command {
+    return typeof action.command === 'string';
+}

tst/lsp-server/LspCommands.test.ts

@@ -1,6 +1,11 @@
 import { describe, it, expect } from 'vitest';
 import { LanguageClient } from 'vscode-languageclient/node';
-import { advertisedCommand, UpdateRegionCommand } from '../../src/lsp-server/LspCommands';
+import {
+    advertisedCommand,
+    advertisedCommands,
+    suffixedCommand,
+    UpdateRegionCommand,
+} from '../../src/lsp-server/LspCommands';
 
 function clientAdvertising(commands: string[] | undefined): LanguageClient {
     return {
@@ -46,3 +51,28 @@ describe('advertisedCommand', () => {
         expect(() => advertisedCommand(client, UpdateRegionCommand)).toThrow('does not advertise');
     });
 });
+
+describe('suffixedCommand', () => {
+    it('appends the client suffix to an unsuffixed id', () => {
+        expect(suffixedCommand(UpdateRegionCommand)).toBe('/command/region/update.iac-vscode');
+    });
+
+    it('keeps an id the server already suffixed for this client', () => {
+        expect(suffixedCommand('/command/region/update.iac-vscode')).toBe('/command/region/update.iac-vscode');
+    });
+
+    it('suffixes an id carrying another client suffix rather than replacing it', () => {
+        expect(suffixedCommand('/command/region/update.other')).toBe('/command/region/update.other.iac-vscode');
+    });
+});
+
+describe('advertisedCommands', () => {
+    it('returns the ids the server advertised', () => {
+        expect(advertisedCommands(clientAdvertising(['/command/region/update']))).toEqual(['/command/region/update']);
+    });
+
+    it('returns no ids before initialization or when the server provides no commands', () => {
+        expect(advertisedCommands({} as unknown as LanguageClient)).toEqual([]);
+        expect(advertisedCommands(clientAdvertising(undefined))).toEqual([]);
+    });
+});