Updated diagnostic to only highlight the first line, not entire node (#125)

* Updated diagnostic to only highlight the first line, not entire node

* fix test failure

* Update to use getKeyRangeFromPath method to get line

Author: rongyamazon

src/server/CfnInfraCore.ts

@@ -54,7 +54,8 @@ export class CfnInfraCore implements Configurables, Closeable {
             overrides.awsCredentials ?? new AwsCredentials(lspComponents.authHandlers, this.settingsManager);
 
         this.diagnosticCoordinator =
-            overrides.diagnosticCoordinator ?? new DiagnosticCoordinator(lspComponents.diagnostics);
+            overrides.diagnosticCoordinator ??
+            new DiagnosticCoordinator(lspComponents.diagnostics, this.syntaxTreeManager);
     }
 
     configurables(): Configurable[] {

src/services/DiagnosticCoordinator.ts

@@ -1,4 +1,7 @@
-import { Diagnostic, PublishDiagnosticsParams } from 'vscode-languageserver';
+import { Diagnostic, PublishDiagnosticsParams, Range } from 'vscode-languageserver';
+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 { CFN_VALIDATION_SOURCE } from '../stacks/actions/ValidationWorkflow';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
@@ -16,7 +19,10 @@ export class DiagnosticCoordinator {
     private readonly urisToDiagnostics = new Map<string, SourceToDiagnostics>();
     private readonly log = LoggerFactory.getLogger(DiagnosticCoordinator);
 
-    constructor(private readonly lspDiagnostics: LspDiagnostics) {}
+    constructor(
+        private readonly lspDiagnostics: LspDiagnostics,
+        private readonly syntaxTreeManager: SyntaxTreeManager,
+    ) {}
 
     /**
      * Publish diagnostics from a specific source for a document.
@@ -164,4 +170,44 @@ export class DiagnosticCoordinator {
 
         return allDiagnostics;
     }
+
+    /**
+     * Extract the key range from a path using syntax tree directly
+     */
+    getKeyRangeFromPath(uri: string, path: string): Range | undefined {
+        // Parse paths like "/Resources/User/Properties/Policies"
+        // Remove leading slash if present and split by '/'
+        const pathSegments = path.startsWith('/') ? path.slice(1).split('/') : path.split('/');
+
+        const syntaxTree = this.syntaxTreeManager.getSyntaxTree(uri);
+        if (!syntaxTree) {
+            return undefined;
+        }
+
+        // Get the node at the path
+        const result = syntaxTree.getNodeByPath(pathSegments);
+        if (!result.node) {
+            return undefined;
+        }
+
+        // Check if this is a pair node (key/value pair)
+        if (NodeType.isPairNode(result.node, syntaxTree.type)) {
+            // Get the key node from the pair
+            const keyNode = result.node.childForFieldName(FieldNames.KEY);
+            if (keyNode) {
+                return {
+                    start: {
+                        line: keyNode.startPosition.row,
+                        character: keyNode.startPosition.column,
+                    },
+                    end: {
+                        line: keyNode.endPosition.row,
+                        character: keyNode.endPosition.column,
+                    },
+                };
+            }
+        }
+
+        return undefined;
+    }
 }

src/services/guard/GuardService.ts

@@ -1,8 +1,6 @@
 import { readFile } from 'fs/promises';
 import { Diagnostic, DiagnosticSeverity, Range } from 'vscode-languageserver';
 import { SyntaxTreeManager } from '../../context/syntaxtree/SyntaxTreeManager';
-import { NodeType } from '../../context/syntaxtree/utils/NodeType';
-import { FieldNames } from '../../context/syntaxtree/utils/TreeSitterTypes';
 import { CloudFormationFileType } from '../../document/Document';
 import { DocumentManager } from '../../document/DocumentManager';
 import { ServerComponents } from '../../server/ServerComponents';
@@ -259,14 +257,8 @@ export class GuardService implements SettingsConfigurable, Closeable {
     private getViolationRange(uri: string, violation: GuardViolation): Range {
         // If we have a CloudFormation path, try to resolve it to precise location
         if (violation.location.path && uri) {
-            // Guard provides paths like "/Resources/User/Properties/Policies"
-            // Remove leading slash if present and split by '/'
-            const pathSegments = violation.location.path.startsWith('/')
-                ? violation.location.path.slice(1).split('/')
-                : violation.location.path.split('/');
-
             // Try to get just the key part of the key/value pair using syntax tree directly
-            const keyRange = this.getKeyRangeFromPath(uri, pathSegments);
+            const keyRange = this.diagnosticCoordinator.getKeyRangeFromPath(uri, violation.location.path);
             if (keyRange) {
                 return keyRange;
             }
@@ -283,42 +275,6 @@ export class GuardService implements SettingsConfigurable, Closeable {
         };
     }
 
-    /**
-     * Extract the key range from a path using syntax tree directly
-     */
-    private getKeyRangeFromPath(uri: string, pathSegments: ReadonlyArray<string>): Range | undefined {
-        const syntaxTree = this.syntaxTreeManager.getSyntaxTree(uri);
-        if (!syntaxTree) {
-            return undefined;
-        }
-
-        // Get the node at the path
-        const result = syntaxTree.getNodeByPath(pathSegments);
-        if (!result.node) {
-            return undefined;
-        }
-
-        // Check if this is a pair node (key/value pair)
-        if (NodeType.isPairNode(result.node, syntaxTree.type)) {
-            // Get the key node from the pair
-            const keyNode = result.node.childForFieldName(FieldNames.KEY);
-            if (keyNode) {
-                return {
-                    start: {
-                        line: keyNode.startPosition.row,
-                        character: keyNode.startPosition.column,
-                    },
-                    end: {
-                        line: keyNode.endPosition.row,
-                        character: keyNode.endPosition.column,
-                    },
-                };
-            }
-        }
-
-        return undefined;
-    }
-
     /**
      * Validate a document with debouncing and proper trigger handling
      *

src/stacks/actions/StackActionOperations.ts

@@ -1,9 +1,8 @@
 import { Change, ChangeSetType, StackStatus } from '@aws-sdk/client-cloudformation';
 import { WaiterState } from '@smithy/util-waiter';
 import { DateTime } from 'luxon';
-import Parser from 'tree-sitter';
 import { v4 as uuidv4 } from 'uuid';
-import { ResponseError, ErrorCodes, Diagnostic, DiagnosticSeverity } from 'vscode-languageserver';
+import { ResponseError, ErrorCodes, Diagnostic, DiagnosticSeverity, Range } from 'vscode-languageserver';
 import { TopLevelSection } from '../../context/ContextType';
 import { getEntityMap } from '../../context/SectionContextBuilder';
 import { SyntaxTreeManager } from '../../context/syntaxtree/SyntaxTreeManager';
@@ -265,35 +264,31 @@ export async function publishValidationDiagnostics(
 
     const diagnostics: Diagnostic[] = [];
     for (const event of events) {
-        let startPosition: Parser.Point | undefined;
-        let endPosition: Parser.Point | undefined;
+        let range: Range | undefined;
 
         if (event.ResourcePropertyPath) {
             logger.debug({ event }, 'Getting property-specific start and end positions');
 
-            // Parse ValidationPath like "/Resources/S3Bucket" or "/Resources/S3Bucket/Properties/BucketName"
-            const pathSegments = event.ResourcePropertyPath.split('/').filter(Boolean);
-
-            const nodeByPath = syntaxTree.getNodeByPath(pathSegments);
-
-            startPosition = nodeByPath.node?.startPosition;
-            endPosition = nodeByPath.node?.endPosition;
+            range = diagnosticCoordinator.getKeyRangeFromPath(uri, event.ResourcePropertyPath);
         } else if (event.LogicalId) {
             // fall back to using LogicalId and underlining entire resource
             logger.debug({ event }, 'No ResourcePropertyPath found, falling back to using LogicalId');
             const resourcesMap = getEntityMap(syntaxTree, TopLevelSection.Resources);
 
-            startPosition = resourcesMap?.get(event.LogicalId)?.startPosition;
-            endPosition = resourcesMap?.get(event.LogicalId)?.endPosition;
+            const startPosition = resourcesMap?.get(event.LogicalId)?.startPosition;
+            const endPosition = resourcesMap?.get(event.LogicalId)?.endPosition;
+            if (startPosition && endPosition) {
+                range = {
+                    start: pointToPosition(startPosition),
+                    end: pointToPosition(endPosition),
+                };
+            }
         }
 
-        if (startPosition && endPosition) {
+        if (range) {
             diagnostics.push({
                 severity: event.Severity === 'ERROR' ? DiagnosticSeverity.Error : DiagnosticSeverity.Warning,
-                range: {
-                    start: pointToPosition(startPosition),
-                    end: pointToPosition(endPosition),
-                },
+                range: range,
                 message: event.Message,
                 source: CFN_VALIDATION_SOURCE,
                 data: uuidv4(),

tst/integration/DocumentHandlerCoordinator.test.ts

@@ -5,7 +5,7 @@ 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 } from '../utils/MockServerComponents';
+import { createMockComponents, createMockSyntaxTreeManager } from '../utils/MockServerComponents';
 
 function mockDocumentEvent(uri: DocumentUri, content: string): TextDocumentChangeEvent<TextDocument> {
     return {
@@ -28,7 +28,8 @@ describe('DocumentHandler Integration with DiagnosticCoordinator', () => {
         mockLspDiagnostics = {
             publishDiagnostics: vi.fn().mockResolvedValue(undefined),
         };
-        realCoordinator = new DiagnosticCoordinator(mockLspDiagnostics);
+        const mockSyntaxTreeManager = createMockSyntaxTreeManager();
+        realCoordinator = new DiagnosticCoordinator(mockLspDiagnostics, mockSyntaxTreeManager);
 
         // Create mock services with the real coordinator
         mockServices = createMockComponents({