Add CloudFormation capabilities detection handler

src/handlers/TemplateHandler.ts

@@ -1,15 +1,17 @@
-import { ServerRequestHandler, ResponseError, ErrorCodes } from 'vscode-languageserver';
+import { ResponseError, ErrorCodes, RequestHandler } from 'vscode-languageserver';
 import { TopLevelSection } from '../context/ContextType';
 import { getEntityMap } from '../context/SectionContextBuilder';
 import { Parameter } from '../context/semantic/Entity';
 import { parseIdentifiable } from '../protocol/LspParser';
 import { Identifiable } from '../protocol/LspTypes';
 import { ServerComponents } from '../server/ServerComponents';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
-import { parseTemplateActionParams, parseGetParametersParams } from '../templates/TemplateParser';
+import { analyzeCapabilities } from '../templates/CapabilityAnalyzer';
+import { parseTemplateActionParams, parseTemplateMetadataParams } from '../templates/TemplateParser';
 import {
-    GetParametersParams,
+    TemplateMetadataParams,
     GetParametersResult,
+    GetCapabilitiesResult,
     TemplateActionParams,
     TemplateActionResult,
     TemplateStatusResult,
@@ -21,12 +23,12 @@ const log = LoggerFactory.getLogger('TemplateHandler');
 
 export function templateParametersHandler(
     components: ServerComponents,
-): ServerRequestHandler<GetParametersParams, GetParametersResult, never, void> {
-    return (rawParams, _token, _workDoneProgress, _resultProgress) => {
+): RequestHandler<TemplateMetadataParams, GetParametersResult, void> {
+    return (rawParams) => {
         log.debug({ Handler: 'TemplateParameters', rawParams });
 
         try {
-            const params = parseWithPrettyError(parseGetParametersParams, rawParams);
+            const params = parseWithPrettyError(parseTemplateMetadataParams, rawParams);
             const syntaxTree = components.syntaxTreeManager.getSyntaxTree(params.uri);
             if (syntaxTree) {
                 const parametersMap = getEntityMap(syntaxTree, TopLevelSection.Parameters);
@@ -49,7 +51,7 @@ export function templateParametersHandler(
 
 export function templateValidationCreateHandler(
     components: ServerComponents,
-): ServerRequestHandler<TemplateActionParams, TemplateActionResult, never, void> {
+): RequestHandler<TemplateActionParams, TemplateActionResult, void> {
     return async (rawParams) => {
         log.debug({ Handler: 'TemplateValidationCreate', rawParams });
 
@@ -64,7 +66,7 @@ export function templateValidationCreateHandler(
 
 export function templateDeploymentCreateHandler(
     components: ServerComponents,
-): ServerRequestHandler<TemplateActionParams, TemplateActionResult, never, void> {
+): RequestHandler<TemplateActionParams, TemplateActionResult, void> {
     return async (rawParams) => {
         log.debug({ Handler: 'TemplateDeploymentCreate', rawParams });
 
@@ -79,7 +81,7 @@ export function templateDeploymentCreateHandler(
 
 export function templateValidationStatusHandler(
     components: ServerComponents,
-): ServerRequestHandler<Identifiable, TemplateStatusResult, never, void> {
+): RequestHandler<Identifiable, TemplateStatusResult, void> {
     return (rawParams) => {
         log.debug({ Handler: 'TemplateValidationStatus', rawParams });
 
@@ -94,7 +96,7 @@ export function templateValidationStatusHandler(
 
 export function templateDeploymentStatusHandler(
     components: ServerComponents,
-): ServerRequestHandler<Identifiable, TemplateStatusResult, never, void> {
+): RequestHandler<Identifiable, TemplateStatusResult, void> {
     return (rawParams) => {
         log.debug({ Handler: 'TemplateDeploymentStatus', rawParams });
 
@@ -107,6 +109,28 @@ export function templateDeploymentStatusHandler(
     };
 }
 
+export function templateCapabilitiesHandler(
+    components: ServerComponents,
+): RequestHandler<TemplateMetadataParams, GetCapabilitiesResult, void> {
+    return async (rawParams) => {
+        log.debug({ Handler: 'TemplateCapabilities', rawParams });
+
+        try {
+            const params = parseWithPrettyError(parseTemplateMetadataParams, rawParams);
+            const document = components.documentManager.get(params.uri);
+            if (!document) {
+                throw new ResponseError(ErrorCodes.InvalidRequest, 'Template body document not available');
+            }
+
+            const capabilities = await analyzeCapabilities(document, components.cfnService);
+
+            return { capabilities };
+        } catch (error) {
+            handleTemplateError(error, 'Failed to analyze template capabilities');
+        }
+    };
+}
+
 function handleTemplateError(error: unknown, contextMessage: string): never {
     if (error instanceof ResponseError) {
         throw error;

src/protocol/LspTemplateHandlers.ts

@@ -1,4 +1,4 @@
-import { Connection, ServerRequestHandler } from 'vscode-languageserver';
+import { Connection, RequestHandler } from 'vscode-languageserver';
 import {
     TemplateActionParams,
     TemplateActionResult,
@@ -7,32 +7,38 @@ import {
     TemplateDeploymentStatusRequest,
     TemplateStatusResult,
     TemplateValidationStatusRequest,
-    GetParametersParams,
+    TemplateMetadataParams,
     GetParametersRequest,
     GetParametersResult,
+    GetCapabilitiesRequest,
+    GetCapabilitiesResult,
 } from '../templates/TemplateRequestType';
 import { Identifiable } from './LspTypes';
 
 export class LspTemplateHandlers {
     constructor(private readonly connection: Connection) {}
 
-    onTemplateValidationCreate(handler: ServerRequestHandler<TemplateActionParams, TemplateActionResult, never, void>) {
+    onTemplateValidationCreate(handler: RequestHandler<TemplateActionParams, TemplateActionResult, void>) {
         this.connection.onRequest(TemplateValidationCreateRequest.method, handler);
     }
 
-    onTemplateDeploymentCreate(handler: ServerRequestHandler<TemplateActionParams, TemplateActionResult, never, void>) {
+    onTemplateDeploymentCreate(handler: RequestHandler<TemplateActionParams, TemplateActionResult, void>) {
         this.connection.onRequest(TemplateDeploymentCreateRequest.method, handler);
     }
 
-    onTemplateValidationStatus(handler: ServerRequestHandler<Identifiable, TemplateStatusResult, never, void>) {
+    onTemplateValidationStatus(handler: RequestHandler<Identifiable, TemplateStatusResult, void>) {
         this.connection.onRequest(TemplateValidationStatusRequest.method, handler);
     }
 
-    onTemplateDeploymentStatus(handler: ServerRequestHandler<Identifiable, TemplateStatusResult, never, void>) {
+    onTemplateDeploymentStatus(handler: RequestHandler<Identifiable, TemplateStatusResult, void>) {
         this.connection.onRequest(TemplateDeploymentStatusRequest.method, handler);
     }
 
-    onGetParameters(handler: ServerRequestHandler<GetParametersParams, GetParametersResult, never, void>) {
+    onGetParameters(handler: RequestHandler<TemplateMetadataParams, GetParametersResult, void>) {
         this.connection.onRequest(GetParametersRequest.method, handler);
     }
+
+    onGetCapabilities(handler: RequestHandler<TemplateMetadataParams, GetCapabilitiesResult, void>) {
+        this.connection.onRequest(GetCapabilitiesRequest.method, handler);
+    }
 }

src/server/CfnServer.ts

@@ -30,6 +30,7 @@ import {
     templateValidationStatusHandler,
     templateDeploymentStatusHandler,
     templateParametersHandler,
+    templateCapabilitiesHandler,
 } from '../handlers/TemplateHandler';
 import { LspFeatures } from '../protocol/LspConnection';
 import { ServerComponents } from './ServerComponents';
@@ -69,6 +70,7 @@ export class CfnServer {
         this.features.authHandlers.onSsoTokenChanged(ssoTokenChangedHandler(this.components));
 
         this.features.templateHandlers.onGetParameters(templateParametersHandler(this.components));
+        this.features.templateHandlers.onGetCapabilities(templateCapabilitiesHandler(this.components));
         this.features.templateHandlers.onTemplateValidationCreate(templateValidationCreateHandler(this.components));
         this.features.templateHandlers.onTemplateDeploymentCreate(templateDeploymentCreateHandler(this.components));
         this.features.templateHandlers.onTemplateValidationStatus(templateValidationStatusHandler(this.components));

src/services/CfnService.ts

@@ -34,6 +34,9 @@ import {
     StackResourceDriftStatus,
     Parameter,
     RegistryType,
+    ValidateTemplateCommand,
+    ValidateTemplateInput,
+    ValidateTemplateOutput,
     Visibility,
     TypeSummary,
     DescribeTypeOutput,
@@ -275,6 +278,10 @@ export class CfnService {
         });
     }
 
+    public async validateTemplate(params: ValidateTemplateInput): Promise<ValidateTemplateOutput> {
+        return await this.withClient((client) => client.send(new ValidateTemplateCommand(params)));
+    }
+
     static create(components: ServerComponents) {
         return new CfnService(components.awsClient);
     }

src/templates/CapabilityAnalyzer.ts

@@ -0,0 +1,26 @@
+import { Capability } from '@aws-sdk/client-cloudformation';
+import { Document } from '../document/Document';
+import { CfnService } from '../services/CfnService';
+import { LoggerFactory } from '../telemetry/LoggerFactory';
+
+const log = LoggerFactory.getLogger('CapabilityAnalyzer');
+
+export async function analyzeCapabilities(document: Document, cfnService: CfnService): Promise<Capability[]> {
+    try {
+        const validationResult = await cfnService.validateTemplate({ TemplateBody: document.getText() });
+
+        if (!validationResult.Capabilities) {
+            return [];
+        }
+
+        // ValidateTemplate cannot process transforms, assume all capabilities are required if a transform is detected
+        if (validationResult.Capabilities.includes(Capability.CAPABILITY_AUTO_EXPAND)) {
+            return [Capability.CAPABILITY_IAM, Capability.CAPABILITY_NAMED_IAM, Capability.CAPABILITY_AUTO_EXPAND];
+        }
+
+        return validationResult.Capabilities;
+    } catch (error) {
+        log.warn({ error }, 'Capability Analysis failed, assuming all capabilities are required');
+        return [Capability.CAPABILITY_IAM, Capability.CAPABILITY_NAMED_IAM, Capability.CAPABILITY_AUTO_EXPAND];
+    }
+}

src/templates/TemplateParser.ts

@@ -1,6 +1,6 @@
 import { Capability } from '@aws-sdk/client-cloudformation';
 import { z } from 'zod';
-import { TemplateActionParams, GetParametersParams } from './TemplateRequestType';
+import { TemplateActionParams, TemplateMetadataParams } from './TemplateRequestType';
 
 const CapabilitySchema = z.enum([
     Capability.CAPABILITY_AUTO_EXPAND,
@@ -31,6 +31,6 @@ export function parseTemplateActionParams(input: unknown): TemplateActionParams
     return TemplateActionParamsSchema.parse(input);
 }
 
-export function parseGetParametersParams(input: unknown): GetParametersParams {
+export function parseTemplateMetadataParams(input: unknown): TemplateMetadataParams {
     return GetParametersParamsSchema.parse(input);
 }

src/templates/TemplateRequestType.ts

@@ -16,7 +16,7 @@ export type TemplateActionResult = Identifiable & {
     stackName: string;
 };
 
-export type GetParametersParams = {
+export type TemplateMetadataParams = {
     uri: string;
 };
 
@@ -76,6 +76,14 @@ export const TemplateDeploymentStatusRequest = new RequestType<Identifiable, Tem
     'aws/cfn/template/deployment/status',
 );
 
-export const GetParametersRequest = new RequestType<GetParametersParams, GetParametersResult, void>(
+export const GetParametersRequest = new RequestType<TemplateMetadataParams, GetParametersResult, void>(
     'aws/cfn/template/parameters',
 );
+
+export type GetCapabilitiesResult = {
+    capabilities: Capability[];
+};
+
+export const GetCapabilitiesRequest = new RequestType<TemplateMetadataParams, GetCapabilitiesResult, void>(
+    'aws/cfn/template/capabilities',
+);