Autoupdate validation underline when template changed (#229)

rongyamazon · web-flow · commit 5956530053d4 · 2025-11-12T11:28:24.000-05:00
diff --git a/src/handlers/DocumentHandler.ts b/src/handlers/DocumentHandler.ts
@@ -10,6 +10,7 @@ import { LspDocuments } from '../protocol/LspDocuments';
 import { ServerComponents } from '../server/ServerComponents';
 import { LintTrigger } from '../services/cfnLint/CfnLintService';
 import { ValidationTrigger } from '../services/guard/GuardService';
+import { publishValidationDiagnostics } from '../stacks/actions/StackActionOperations';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 
 const log = LoggerFactory.getLogger('DocumentHandler');
@@ -123,6 +124,19 @@ export function didChangeHandler(
                 }
             });
 
+        // Republish validation diagnostics if available
+        const validationDetails = components.validationManager
+            .getLastValidationByUri(documentUri)
+            ?.getValidationDetails();
+        if (validationDetails) {
+            void publishValidationDiagnostics(
+                documentUri,
+                validationDetails,
+                components.syntaxTreeManager,
+                components.diagnosticCoordinator,
+            );
+        }
+
         components.documentManager.sendDocumentMetadata();
     };
 }
diff --git a/src/server/CfnInfraCore.ts b/src/server/CfnInfraCore.ts
@@ -8,6 +8,7 @@ import { DocumentMetadata } from '../document/DocumentProtocol';
 import { LspComponents } from '../protocol/LspComponents';
 import { DiagnosticCoordinator } from '../services/DiagnosticCoordinator';
 import { SettingsManager } from '../settings/SettingsManager';
+import { ValidationManager } from '../stacks/actions/ValidationManager';
 import { ClientMessage } from '../telemetry/ClientMessage';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { TelemetryService } from '../telemetry/TelemetryService';
@@ -31,6 +32,7 @@ export class CfnInfraCore implements Configurables, Closeable {
     readonly fileContextManager: FileContextManager;
 
     readonly awsCredentials: AwsCredentials;
+    readonly validationManager: ValidationManager;
     readonly diagnosticCoordinator: DiagnosticCoordinator;
     readonly cloudformationEndpoint?: string;
 
@@ -62,9 +64,11 @@ export class CfnInfraCore implements Configurables, Closeable {
                 initializeParams.initializationOptions?.aws?.encryption?.key,
             );
 
+        this.validationManager = overrides.validationManager ?? new ValidationManager();
+
         this.diagnosticCoordinator =
             overrides.diagnosticCoordinator ??
-            new DiagnosticCoordinator(lspComponents.diagnostics, this.syntaxTreeManager);
+            new DiagnosticCoordinator(lspComponents.diagnostics, this.syntaxTreeManager, this.validationManager);
     }
 
     configurables(): Configurable[] {
diff --git a/src/server/CfnLspProviders.ts b/src/server/CfnLspProviders.ts
@@ -65,8 +65,8 @@ export class CfnLspProviders implements Configurables, Closeable {
         this.validationWorkflowService =
             overrides.validationWorkflowService ??
             (localCfnClientExists()
-                ? ValidationWorkflowV2.create(core, external)
-                : ValidationWorkflow.create(core, external));
+                ? ValidationWorkflowV2.create(core, external, core.validationManager)
+                : ValidationWorkflow.create(core, external, core.validationManager));
         this.deploymentWorkflowService =
             overrides.deploymentWorkflowService ?? DeploymentWorkflow.create(core, external);
         this.changeSetDeletionWorkflowService =
diff --git a/src/services/DiagnosticCoordinator.ts b/src/services/DiagnosticCoordinator.ts
@@ -3,6 +3,7 @@ import { SyntaxTreeManager } from '../context/syntaxtree/SyntaxTreeManager';
 import { NodeType } from '../context/syntaxtree/utils/NodeType';
 import { FieldNames } from '../context/syntaxtree/utils/TreeSitterTypes';
 import { LspDiagnostics } from '../protocol/LspDiagnostics';
+import { ValidationManager } from '../stacks/actions/ValidationManager';
 import { CFN_VALIDATION_SOURCE } from '../stacks/actions/ValidationWorkflow';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 
@@ -21,6 +22,7 @@ export class DiagnosticCoordinator {
     constructor(
         private readonly lspDiagnostics: LspDiagnostics,
         private readonly syntaxTreeManager: SyntaxTreeManager,
+        private readonly validationManager: ValidationManager,
     ) {}
 
     /**
@@ -98,6 +100,7 @@ export class DiagnosticCoordinator {
         const sourceDiagnostics = collection.get(CFN_VALIDATION_SOURCE);
         if (!sourceDiagnostics) return;
 
+        this.validationManager.getLastValidationByUri(uri)?.removeValidationDetailByDiagnosticId(diagnosticId);
         const filteredDiagnostics = sourceDiagnostics.filter((d) => d.data !== diagnosticId);
         collection.set(CFN_VALIDATION_SOURCE, filteredDiagnostics);
 
diff --git a/src/stacks/actions/StackActionOperations.ts b/src/stacks/actions/StackActionOperations.ts
@@ -387,13 +387,15 @@ export async function publishValidationDiagnostics(
         }
 
         if (range) {
+            const diagnosticId = uuidv4();
             diagnostics.push({
                 severity: event.Severity === 'ERROR' ? DiagnosticSeverity.Error : DiagnosticSeverity.Warning,
                 range: range,
                 message: event.Message,
                 source: CFN_VALIDATION_SOURCE,
-                data: uuidv4(),
+                data: diagnosticId,
             });
+            event.diagnosticId = diagnosticId;
         }
     }
 
diff --git a/src/stacks/actions/StackActionRequestType.ts b/src/stacks/actions/StackActionRequestType.ts
@@ -128,6 +128,7 @@ export type ValidationDetail = {
     Timestamp: DateTime;
     Severity: 'INFO' | 'ERROR';
     Message: string;
+    diagnosticId?: string;
 };
 
 export type DeploymentEvent = {
diff --git a/src/stacks/actions/Validation.ts b/src/stacks/actions/Validation.ts
@@ -1,5 +1,5 @@
 import { Parameter, Capability } from '@aws-sdk/client-cloudformation';
-import { StackActionPhase, StackChange } from './StackActionRequestType';
+import { StackActionPhase, StackChange, ValidationDetail } from './StackActionRequestType';
 
 export class Validation {
     private readonly uri: string;
@@ -11,6 +11,7 @@ export class Validation {
     private readonly s3Key?: string;
     private phase: StackActionPhase | undefined;
     private changes: StackChange[] | undefined;
+    private validationDetails: ValidationDetail[] | undefined;
 
     constructor(
         uri: string,
@@ -77,4 +78,18 @@ export class Validation {
     getS3Key(): string | undefined {
         return this.s3Key;
     }
+
+    getValidationDetails(): ValidationDetail[] | undefined {
+        return this.validationDetails;
+    }
+
+    setValidationDetails(validationDetails: ValidationDetail[]) {
+        this.validationDetails = validationDetails;
+    }
+
+    removeValidationDetailByDiagnosticId(diagnosticId: string): void {
+        if (this.validationDetails) {
+            this.validationDetails = this.validationDetails.filter((vd) => vd.diagnosticId !== diagnosticId);
+        }
+    }
 }
diff --git a/src/stacks/actions/ValidationManager.ts b/src/stacks/actions/ValidationManager.ts
@@ -3,15 +3,21 @@ import { Validation } from './Validation';
 
 export class ValidationManager {
     private readonly validations: Map<string, Validation> = new Map();
+    private readonly uriToValidation: Map<string, Validation> = new Map();
 
     add(validation: Validation): void {
         this.validations.set(validation.getStackName(), validation);
+        this.uriToValidation.set(validation.getUri(), validation);
     }
 
     get(stackName: string): Validation | undefined {
         return this.validations.get(stackName);
     }
 
+    getLastValidationByUri(uri: string): Validation | undefined {
+        return this.uriToValidation.get(uri);
+    }
+
     setChanges(stackName: string, changes: StackChange[]): void {
         const validation = this.validations.get(stackName);
         if (validation) {
@@ -20,10 +26,15 @@ export class ValidationManager {
     }
 
     remove(stackName: string): boolean {
+        const validation = this.validations.get(stackName);
+        if (validation) {
+            this.uriToValidation.delete(validation.getUri());
+        }
         return this.validations.delete(stackName);
     }
 
     clear(): void {
         this.validations.clear();
+        this.uriToValidation.clear();
     }
 }
diff --git a/src/stacks/actions/ValidationWorkflow.ts b/src/stacks/actions/ValidationWorkflow.ts
@@ -198,9 +198,6 @@ export class ValidationWorkflow implements StackActionWorkflow<CreateValidationP
     }
 
     protected async handleCleanup(params: CreateValidationParams, existingWorkflow: StackActionWorkflowState) {
-        // Cleanup validation object to prevent memory leaks
-        this.validationManager.remove(params.stackName);
-
         if (!params.keepChangeSet) {
             try {
                 if (await isStackInReview(params.stackName, this.cfnService)) {
@@ -214,13 +211,13 @@ export class ValidationWorkflow implements StackActionWorkflow<CreateValidationP
         }
     }
 
-    static create(core: CfnInfraCore, external: CfnExternal): ValidationWorkflow {
+    static create(core: CfnInfraCore, external: CfnExternal, validationManager: ValidationManager): ValidationWorkflow {
         return new ValidationWorkflow(
             external.cfnService,
             core.documentManager,
             core.diagnosticCoordinator,
             core.syntaxTreeManager,
-            new ValidationManager(),
+            validationManager,
             external.s3Service,
         );
     }
diff --git a/src/stacks/actions/ValidationWorkflowV2.ts b/src/stacks/actions/ValidationWorkflowV2.ts
@@ -51,11 +51,13 @@ export class ValidationWorkflowV2 extends ValidationWorkflow {
             const result = await waitForChangeSetValidation(this.cfnServiceV2, changeSetName, stackName);
 
             const validation = this.validationManager.get(stackName);
-            if (validation) {
-                validation.setPhase(result.phase);
-                if (result.changes) {
-                    validation.setChanges(result.changes);
-                }
+            if (!validation) {
+                throw new Error(`No validation found for stack: ${stackName}`);
+            }
+
+            validation.setPhase(result.phase);
+            if (result.changes) {
+                validation.setChanges(result.changes);
             }
 
             existingWorkflow = processWorkflowUpdates(this.workflows, existingWorkflow, {
@@ -82,6 +84,7 @@ export class ValidationWorkflowV2 extends ValidationWorkflow {
                     validationDetails: validationDetails,
                 });
 
+                validation.setValidationDetails(validationDetails);
                 await publishValidationDiagnostics(
                     uri,
                     validationDetails,
@@ -107,13 +110,17 @@ export class ValidationWorkflowV2 extends ValidationWorkflow {
         }
     }
 
-    static override create(core: CfnInfraCore, external: CfnExternal): ValidationWorkflowV2 {
+    static override create(
+        core: CfnInfraCore,
+        external: CfnExternal,
+        validationManager: ValidationManager,
+    ): ValidationWorkflowV2 {
         return new ValidationWorkflowV2(
             new CfnServiceV2(external.awsClient),
             core.documentManager,
             core.diagnosticCoordinator,
             core.syntaxTreeManager,
-            new ValidationManager(),
+            validationManager,
             core.fileContextManager,
             external.featureFlags.getTargeted<string>('EnhancedDryRun'),
             core.awsCredentials,
diff --git a/tst/integration/DocumentHandlerCoordinator.test.ts b/tst/integration/DocumentHandlerCoordinator.test.ts
@@ -5,7 +5,11 @@ import { TextDocument } from 'vscode-languageserver-textdocument';
 import { DocumentUri } from 'vscode-languageserver-textdocument/lib/esm/main';
 import { didCloseHandler } from '../../src/handlers/DocumentHandler';
 import { DiagnosticCoordinator } from '../../src/services/DiagnosticCoordinator';
-import { createMockComponents, createMockSyntaxTreeManager } from '../utils/MockServerComponents';
+import {
+    createMockComponents,
+    createMockSyntaxTreeManager,
+    createMockValidationManager,
+} from '../utils/MockServerComponents';
 
 function mockDocumentEvent(uri: DocumentUri, content: string): TextDocumentChangeEvent<TextDocument> {
     return {
@@ -29,7 +33,8 @@ describe('DocumentHandler Integration with DiagnosticCoordinator', () => {
             publishDiagnostics: vi.fn().mockResolvedValue(undefined),
         };
         const mockSyntaxTreeManager = createMockSyntaxTreeManager();
-        realCoordinator = new DiagnosticCoordinator(mockLspDiagnostics, mockSyntaxTreeManager);
+        const mockValidationManager = createMockValidationManager();
+        realCoordinator = new DiagnosticCoordinator(mockLspDiagnostics, mockSyntaxTreeManager, mockValidationManager);
 
         // Create mock services with the real coordinator
         mockServices = createMockComponents({
diff --git a/tst/unit/services/DiagnosticCoordinator.test.ts b/tst/unit/services/DiagnosticCoordinator.test.ts
@@ -5,7 +5,11 @@ import { Diagnostic, DiagnosticSeverity, Range, Position } from 'vscode-language
 import { SyntaxTree } from '../../../src/context/syntaxtree/SyntaxTree';
 import { DiagnosticCoordinator } from '../../../src/services/DiagnosticCoordinator';
 import { CFN_VALIDATION_SOURCE } from '../../../src/stacks/actions/ValidationWorkflow';
-import { createMockLspDiagnostics, createMockSyntaxTreeManager } from '../../utils/MockServerComponents';
+import {
+    createMockLspDiagnostics,
+    createMockSyntaxTreeManager,
+    createMockValidationManager,
+} from '../../utils/MockServerComponents';
 
 // Mock NodeType module
 vi.mock('../../../src/context/syntaxtree/utils/NodeType', () => ({
@@ -44,7 +48,8 @@ describe('DiagnosticCoordinator', () => {
         mockLspDiagnostics.publishDiagnostics.callsFake(mockPublishDiagnostics);
 
         mockSyntaxTreeManager = createMockSyntaxTreeManager();
-        coordinator = new DiagnosticCoordinator(mockLspDiagnostics, mockSyntaxTreeManager);
+        const mockValidationManager = createMockValidationManager();
+        coordinator = new DiagnosticCoordinator(mockLspDiagnostics, mockSyntaxTreeManager, mockValidationManager);
     });
 
     afterEach(() => {
diff --git a/tst/unit/stackActions/StackActionWorkflowOperations.test.ts b/tst/unit/stackActions/StackActionWorkflowOperations.test.ts
@@ -533,6 +533,7 @@ describe('StackActionWorkflowOperations', () => {
             );
 
             expect(mockDiagnosticCoordinator.publishDiagnostics.calledOnce).toBe(true);
+            expect(validationDetails[0].diagnosticId).toBeDefined();
         });
     });
 
diff --git a/tst/unit/stackActions/Validation.test.ts b/tst/unit/stackActions/Validation.test.ts
@@ -1,3 +1,4 @@
+import { DateTime } from 'luxon';
 import { describe, it, expect, beforeEach } from 'vitest';
 import { StackActionPhase } from '../../../src/stacks/actions/StackActionRequestType';
 import { Validation } from '../../../src/stacks/actions/Validation';
@@ -126,4 +127,77 @@ describe('Validation', () => {
             expect(validation.getCapabilities()).toBeUndefined();
         });
     });
+
+    describe('ValidationDetails management', () => {
+        it('should get and set validation details', () => {
+            const validationDetails = [
+                {
+                    ValidationName: 'test',
+                    LogicalId: 'TestResource',
+                    ResourcePropertyPath: '/Resources/TestResource',
+                    Timestamp: DateTime.now(),
+                    Severity: 'ERROR' as const,
+                    Message: 'Test error',
+                    diagnosticId: 'test-id-1',
+                },
+            ];
+
+            validation.setValidationDetails(validationDetails);
+            expect(validation.getValidationDetails()).toEqual(validationDetails);
+        });
+
+        it('should remove validation detail by diagnosticId', () => {
+            const validationDetails = [
+                {
+                    ValidationName: 'test1',
+                    LogicalId: 'TestResource1',
+                    ResourcePropertyPath: '/Resources/TestResource1',
+                    Timestamp: DateTime.now(),
+                    Severity: 'ERROR' as const,
+                    Message: 'Test error 1',
+                    diagnosticId: 'test-id-1',
+                },
+                {
+                    ValidationName: 'test2',
+                    LogicalId: 'TestResource2',
+                    ResourcePropertyPath: '/Resources/TestResource2',
+                    Timestamp: DateTime.now(),
+                    Severity: 'ERROR' as const,
+                    Message: 'Test error 2',
+                    diagnosticId: 'test-id-2',
+                },
+            ];
+
+            validation.setValidationDetails(validationDetails);
+            validation.removeValidationDetailByDiagnosticId('test-id-1');
+
+            const remaining = validation.getValidationDetails();
+            expect(remaining).toHaveLength(1);
+            expect(remaining?.[0].diagnosticId).toBe('test-id-2');
+        });
+
+        it('should handle removing non-existent diagnosticId', () => {
+            const validationDetails = [
+                {
+                    ValidationName: 'test',
+                    LogicalId: 'TestResource',
+                    ResourcePropertyPath: '/Resources/TestResource',
+                    Timestamp: DateTime.now(),
+                    Severity: 'ERROR' as const,
+                    Message: 'Test error',
+                    diagnosticId: 'test-id-1',
+                },
+            ];
+
+            validation.setValidationDetails(validationDetails);
+            validation.removeValidationDetailByDiagnosticId('non-existent-id');
+
+            expect(validation.getValidationDetails()).toHaveLength(1);
+        });
+
+        it('should handle removing from empty validation details', () => {
+            validation.removeValidationDetailByDiagnosticId('test-id');
+            expect(validation.getValidationDetails()).toBeUndefined();
+        });
+    });
 });
diff --git a/tst/unit/stackActions/ValidationManager.test.ts b/tst/unit/stackActions/ValidationManager.test.ts
diff --git a/tst/unit/stackActions/ValidationWorkflow.test.ts b/tst/unit/stackActions/ValidationWorkflow.test.ts
diff --git a/tst/unit/stackActions/ValidationWorkflowEnhanced.test.ts b/tst/unit/stackActions/ValidationWorkflowEnhanced.test.ts
diff --git a/tst/unit/stackActions/ValidationWorkflowV2.test.ts b/tst/unit/stackActions/ValidationWorkflowV2.test.ts
diff --git a/tst/utils/MockServerComponents.ts b/tst/utils/MockServerComponents.ts

Original file line number	Diff line number	Diff line change
`@@ -387,13 +387,15 @@ export async function publishValidationDiagnostics(`
`387`	`387`	`}`
`388`	`388`
`389`	`389`	`if (range) {`
	`390`	`+ const diagnosticId = uuidv4();`
`390`	`391`	`diagnostics.push({`
`391`	`392`	`severity: event.Severity === 'ERROR' ? DiagnosticSeverity.Error : DiagnosticSeverity.Warning,`
`392`	`393`	`range: range,`
`393`	`394`	`message: event.Message,`
`394`	`395`	`source: CFN_VALIDATION_SOURCE,`
`395`		`- data: uuidv4(),`
	`396`	`+ data: diagnosticId,`
`396`	`397`	`});`
	`398`	`+ event.diagnosticId = diagnosticId;`
`397`	`399`	`}`
`398`	`400`	`}`
`399`	`401`
Original file line number	Diff line number	Diff line change
`@@ -3,15 +3,21 @@ import { Validation } from './Validation';`
`3`	`3`
`4`	`4`	`export class ValidationManager {`
`5`	`5`	`private readonly validations: Map<string, Validation> = new Map();`
	`6`	`+ private readonly uriToValidation: Map<string, Validation> = new Map();`
`6`	`7`
`7`	`8`	`add(validation: Validation): void {`
`8`	`9`	`this.validations.set(validation.getStackName(), validation);`
	`10`	`+ this.uriToValidation.set(validation.getUri(), validation);`
`9`	`11`	`}`
`10`	`12`
`11`	`13`	`get(stackName: string): Validation \| undefined {`
`12`	`14`	`return this.validations.get(stackName);`
`13`	`15`	`}`
`14`	`16`
	`17`	`+ getLastValidationByUri(uri: string): Validation \| undefined {`
	`18`	`+ return this.uriToValidation.get(uri);`
	`19`	`+ }`
	`20`	`+`
`15`	`21`	`setChanges(stackName: string, changes: StackChange[]): void {`
`16`	`22`	`const validation = this.validations.get(stackName);`
`17`	`23`	`if (validation) {`
`@@ -20,10 +26,15 @@ export class ValidationManager {`
`20`	`26`	`}`
`21`	`27`
`22`	`28`	`remove(stackName: string): boolean {`
	`29`	`+ const validation = this.validations.get(stackName);`
	`30`	`+ if (validation) {`
	`31`	`+ this.uriToValidation.delete(validation.getUri());`
	`32`	`+ }`
`23`	`33`	`return this.validations.delete(stackName);`
`24`	`34`	`}`
`25`	`35`
`26`	`36`	`clear(): void {`
`27`	`37`	`this.validations.clear();`
	`38`	`+ this.uriToValidation.clear();`
`28`	`39`	`}`
`29`	`40`	`}`
Original file line number	Diff line number	Diff line change
`@@ -198,9 +198,6 @@ export class ValidationWorkflow implements StackActionWorkflow<CreateValidationP`
`198`	`198`	`}`
`199`	`199`
`200`	`200`	`protected async handleCleanup(params: CreateValidationParams, existingWorkflow: StackActionWorkflowState) {`
`201`		`- // Cleanup validation object to prevent memory leaks`
`202`		`- this.validationManager.remove(params.stackName);`
`203`		`-`
`204`	`201`	`if (!params.keepChangeSet) {`
`205`	`202`	`try {`
`206`	`203`	`if (await isStackInReview(params.stackName, this.cfnService)) {`
`@@ -214,13 +211,13 @@ export class ValidationWorkflow implements StackActionWorkflow<CreateValidationP`
`214`	`211`	`}`
`215`	`212`	`}`
`216`	`213`
`217`		`- static create(core: CfnInfraCore, external: CfnExternal): ValidationWorkflow {`
	`214`	`+ static create(core: CfnInfraCore, external: CfnExternal, validationManager: ValidationManager): ValidationWorkflow {`
`218`	`215`	`return new ValidationWorkflow(`
`219`	`216`	`external.cfnService,`
`220`	`217`	`core.documentManager,`
`221`	`218`	`core.diagnosticCoordinator,`
`222`	`219`	`core.syntaxTreeManager,`
`223`		`- new ValidationManager(),`
	`220`	`+ validationManager,`
`224`	`221`	`external.s3Service,`
`225`	`222`	`);`
`226`	`223`	`}`