Adding support for VSCode command suffix's per extension

src/app/standalone.ts

@@ -1,7 +1,8 @@
 import './polyfills';
 import { createConnection, ProposedFeatures } from 'vscode-languageserver/node'; // eslint-disable-line no-restricted-imports
 import { InitializedParams } from 'vscode-languageserver-protocol';
-import { LspCapabilities } from '../protocol/LspCapabilities';
+import { initCommands, getCommands } from '../handlers/ExecutionHandler';
+import { buildCapabilities } from '../protocol/LspCapabilities';
 import { LspConnection } from '../protocol/LspConnection';
 import { ExtendedInitializeParams } from '../server/InitParams';
 import { ExtensionName } from '../utils/ExtensionConfig';
@@ -13,13 +14,17 @@ let server: unknown;
 async function onInitialize(params: ExtendedInitializeParams) {
     staticInitialize(params.clientInfo, params.initializationOptions?.['aws']);
 
+    const name = params.initializationOptions?.aws?.clientInfo?.extension?.name;
+    const suffix = typeof name === 'string' ? name : undefined;
+    initCommands(suffix);
+
     // Dynamically load these modules so that OTEL can instrument all the libraries first
     const { CfnInfraCore } = await import('../server/CfnInfraCore');
     const core = new CfnInfraCore(lsp.components, params);
 
     const { CfnServer } = await import('../server/CfnServer');
     server = new CfnServer(lsp.components, core);
-    return LspCapabilities;
+    return buildCapabilities(getCommands());
 }
 
 function onInitialized(params: InitializedParams) {

src/handlers/ExecutionHandler.ts

@@ -4,6 +4,34 @@ import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { TelemetryService } from '../telemetry/TelemetryService';
 import { getRegion } from '../utils/Region';
 
+export const CLEAR_DIAGNOSTIC = '/command/template/clear-diagnostic';
+export const TRACK_CODE_ACTION_ACCEPTED = '/command/codeAction/track';
+export const UPDATE_REGION = '/command/region/update';
+
+export function createCommands(suffix?: string) {
+    const s = suffix ? `.${suffix}` : '';
+    return {
+        clearDiagnostic: `${CLEAR_DIAGNOSTIC}${s}`,
+        trackCodeAction: `${TRACK_CODE_ACTION_ACCEPTED}${s}`,
+        updateRegion: `${UPDATE_REGION}${s}`,
+    };
+}
+
+export type Commands = ReturnType<typeof createCommands>;
+
+// Module-level singleton. initCommands() must be called during onInitialize
+// before any workspace/executeCommand requests arrive. This is safe because
+// LSP guarantees no requests are sent until after the initialize handshake.
+let resolvedCommands: Commands = createCommands();
+
+export function initCommands(suffix?: string) {
+    resolvedCommands = createCommands(suffix);
+}
+
+export function getCommands(): Commands {
+    return resolvedCommands;
+}
+
 export function executionHandler(
     components: ServerComponents,
 ): ServerRequestHandler<ExecuteCommandParams, unknown, never, void> {
@@ -12,7 +40,7 @@ export function executionHandler(
         TelemetryService.instance.get('ExecutionHandler').count(`count.${params.command}`, 1);
 
         switch (params.command) {
-            case CLEAR_DIAGNOSTIC: {
+            case resolvedCommands.clearDiagnostic: {
                 const args = params.arguments ?? [];
                 if (args.length >= 2) {
                     const uri = args[0] as string;
@@ -26,15 +54,15 @@ export function executionHandler(
                 }
                 break;
             }
-            case TRACK_CODE_ACTION_ACCEPTED: {
+            case resolvedCommands.trackCodeAction: {
                 const args = params.arguments ?? [];
                 if (args.length > 0) {
                     const actionType = args[0] as string;
                     TelemetryService.instance.get('CodeAction').count(`accepted.${actionType}`, 1);
                 }
                 break;
             }
-            case UPDATE_REGION: {
+            case resolvedCommands.updateRegion: {
                 const args = params.arguments ?? [];
                 if (args.length > 0) {
                     components.awsCredentials.handleIamCredentialsDelete();
@@ -49,7 +77,3 @@ export function executionHandler(
         }
     };
 }
-
-export const CLEAR_DIAGNOSTIC = '/command/template/clear-diagnostic';
-export const TRACK_CODE_ACTION_ACCEPTED = '/command/codeAction/track';
-export const UPDATE_REGION = '/command/region/update';

src/protocol/LspCapabilities.ts

@@ -1,43 +1,45 @@
 import { InitializeResult, TextDocumentSyncKind, CodeActionKind } from 'vscode-languageserver';
-import { CLEAR_DIAGNOSTIC, TRACK_CODE_ACTION_ACCEPTED, UPDATE_REGION } from '../handlers/ExecutionHandler';
+import { Commands } from '../handlers/ExecutionHandler';
 import { ExtensionName, ExtensionVersion } from '../utils/ExtensionConfig';
 
-export const LspCapabilities: InitializeResult = {
-    capabilities: {
-        textDocumentSync: {
-            openClose: true,
-            change: TextDocumentSyncKind.Incremental,
-            willSave: false,
-            willSaveWaitUntil: false,
-            save: {
-                includeText: true,
+export function buildCapabilities(commands: Commands): InitializeResult {
+    return {
+        capabilities: {
+            textDocumentSync: {
+                openClose: true,
+                change: TextDocumentSyncKind.Incremental,
+                willSave: false,
+                willSaveWaitUntil: false,
+                save: {
+                    includeText: true,
+                },
             },
-        },
-        hoverProvider: true,
-        codeActionProvider: {
-            resolveProvider: false,
-            codeActionKinds: [CodeActionKind.RefactorExtract],
-        },
-        completionProvider: {
-            triggerCharacters: ['.', '!', ':', '\n', '\t', '"'],
-            completionItem: {
-                labelDetailsSupport: true,
+            hoverProvider: true,
+            codeActionProvider: {
+                resolveProvider: false,
+                codeActionKinds: [CodeActionKind.RefactorExtract],
             },
-        },
-        definitionProvider: true,
-        documentSymbolProvider: true,
-        executeCommandProvider: {
-            commands: [CLEAR_DIAGNOSTIC, TRACK_CODE_ACTION_ACCEPTED, UPDATE_REGION],
-        },
-        workspace: {
-            workspaceFolders: {
-                supported: true,
-                changeNotifications: true,
+            completionProvider: {
+                triggerCharacters: ['.', '!', ':', '\n', '\t', '"'],
+                completionItem: {
+                    labelDetailsSupport: true,
+                },
             },
+            definitionProvider: true,
+            documentSymbolProvider: true,
+            executeCommandProvider: {
+                commands: [commands.clearDiagnostic, commands.trackCodeAction, commands.updateRegion],
+            },
+            workspace: {
+                workspaceFolders: {
+                    supported: true,
+                    changeNotifications: true,
+                },
+            },
+        },
+        serverInfo: {
+            name: ExtensionName,
+            version: ExtensionVersion,
         },
-    },
-    serverInfo: {
-        name: ExtensionName,
-        version: ExtensionVersion,
-    },
-};
+    };
+}

src/protocol/LspConnection.ts

@@ -1,8 +1,9 @@
 import { Connection, InitializeParams, InitializeResult } from 'vscode-languageserver';
 import { InitializedParams } from 'vscode-languageserver-protocol';
+import { getCommands } from '../handlers/ExecutionHandler';
 import { ExtensionName } from '../utils/ExtensionConfig';
 import { LspAuthHandlers } from './LspAuthHandlers';
-import { LspCapabilities } from './LspCapabilities';
+import { buildCapabilities } from './LspCapabilities';
 import { LspCfnEnvironmentHandlers } from './LspCfnEnvironmentHandlers';
 import { LspCommunication } from './LspCommunication';
 import { LspComponents } from './LspComponents';
@@ -44,7 +45,7 @@ export class LspConnection {
         handlers: LspConnectionHandlers = {},
     ) {
         const {
-            onInitialize = () => LspCapabilities,
+            onInitialize = () => buildCapabilities(getCommands()),
             onInitialized = () => {},
             onShutdown = () => {},
             onExit = () => {},

src/services/CodeActionService.ts

@@ -15,7 +15,7 @@ import { SyntaxTreeManager } from '../context/syntaxtree/SyntaxTreeManager';
 import { NodeSearch } from '../context/syntaxtree/utils/NodeSearch';
 import { NodeType } from '../context/syntaxtree/utils/NodeType';
 import { DocumentManager } from '../document/DocumentManager';
-import { TRACK_CODE_ACTION_ACCEPTED } from '../handlers/ExecutionHandler';
+import { getCommands } from '../handlers/ExecutionHandler';
 import { CfnInfraCore } from '../server/CfnInfraCore';
 import { CFN_VALIDATION_SOURCE } from '../stacks/actions/ValidationWorkflow';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
@@ -117,7 +117,7 @@ export class CodeActionService {
                 textEdits: [],
                 command: {
                     title: CodeActionService.REMOVE_ERROR_TITLE,
-                    command: '/command/template/clear-diagnostic',
+                    command: getCommands().clearDiagnostic,
                     arguments: [uri, diagnostic.data],
                 },
             },
@@ -328,7 +328,7 @@ export class CodeActionService {
             } else {
                 codeAction.command = {
                     title: 'Track code action',
-                    command: TRACK_CODE_ACTION_ACCEPTED,
+                    command: getCommands().trackCodeAction,
                     arguments: [fix.actionType],
                 };
             }
@@ -597,7 +597,7 @@ export class CodeActionService {
                         params.textDocument.uri,
                         extractionResult.parameterName,
                         context.documentType,
-                        TRACK_CODE_ACTION_ACCEPTED,
+                        getCommands().trackCodeAction,
                         'extractToParameter',
                     ],
                 },
@@ -649,7 +649,7 @@ export class CodeActionService {
                         params.textDocument.uri,
                         extractionResult.parameterName,
                         context.documentType,
-                        TRACK_CODE_ACTION_ACCEPTED,
+                        getCommands().trackCodeAction,
                         'extractAllToParameter',
                     ],
                 },

tst/unit/protocol/LspCapabilities.test.ts

@@ -1,9 +1,16 @@
 import { describe, it, expect } from 'vitest';
 import { TextDocumentSyncKind, CodeActionKind } from 'vscode-languageserver';
-import { CLEAR_DIAGNOSTIC, TRACK_CODE_ACTION_ACCEPTED, UPDATE_REGION } from '../../../src/handlers/ExecutionHandler';
-import { LspCapabilities } from '../../../src/protocol/LspCapabilities';
+import {
+    CLEAR_DIAGNOSTIC,
+    TRACK_CODE_ACTION_ACCEPTED,
+    UPDATE_REGION,
+    createCommands,
+} from '../../../src/handlers/ExecutionHandler';
+import { buildCapabilities } from '../../../src/protocol/LspCapabilities';
 import { ExtensionName, ExtensionVersion } from '../../../src/utils/ExtensionConfig';
 
+const LspCapabilities = buildCapabilities(createCommands());
+
 describe('LspCapabilities', () => {
     describe('capabilities structure', () => {
         it('should have correct text document sync settings', () => {

tst/utils/TestExtension.ts

@@ -52,7 +52,8 @@ import { AwsCredentials } from '../../src/auth/AwsCredentials';
 import { UpdateCredentialsParams } from '../../src/auth/AwsLspAuthTypes';
 import { MultiDataStoreFactoryProvider } from '../../src/datastore/DataStore';
 import { featureFlagLocalFile, FeatureFlagProvider } from '../../src/featureFlag/FeatureFlagProvider';
-import { LspCapabilities } from '../../src/protocol/LspCapabilities';
+import { buildCapabilities } from '../../src/protocol/LspCapabilities';
+import { getCommands } from '../../src/handlers/ExecutionHandler';
 import { LspConnection } from '../../src/protocol/LspConnection';
 import { SchemaRetriever } from '../../src/schema/SchemaRetriever';
 import { SchemaStore } from '../../src/schema/SchemaStore';
@@ -173,7 +174,7 @@ export class TestExtension implements Closeable {
                         ),
                     });
                     this.server = new CfnServer(lsp, this.core, this.external, this.providers);
-                    return LspCapabilities;
+                    return buildCapabilities(getCommands());
                 },
                 onInitialized: (params) => this.server.initialized(params),
                 onShutdown: () => this.server.close(),