feat: add v2 relatedResources handler

src/handlers/RelatedResourcesHandler.ts

@@ -20,6 +20,34 @@ import {
 
 export function getAuthoredResourceTypesHandler(
     components: ServerComponents,
+): RequestHandler<TemplateUri, string[], void> {
+    return (rawParams) => {
+        try {
+            const templateUri = parseWithPrettyError(parseTemplateUriParams, rawParams);
+            const syntaxTree = components.syntaxTreeManager.getSyntaxTree(templateUri);
+            if (syntaxTree) {
+                const resourcesMap = getEntityMap(syntaxTree, TopLevelSection.Resources);
+                if (resourcesMap) {
+                    const resourceTypes = [...resourcesMap.values()]
+                        .map((context) => {
+                            const resource = context.entity as Resource;
+                            return resource?.Type;
+                        })
+                        .filter((type): type is string => type !== undefined && type !== null);
+
+                    return [...new Set(resourceTypes)];
+                }
+            }
+
+            return [];
+        } catch (error) {
+            handleLspError(error, 'Failed to get authored resource types');
+        }
+    };
+}
+
+export function getAuthoredResourceTypesHandlerV2(
+    components: ServerComponents,
 ): RequestHandler<TemplateUri, AuthoredResource[], void> {
     return (rawParams) => {
         try {

src/protocol/LspRelatedResourcesHandlers.ts

@@ -2,6 +2,7 @@ import { Connection, RequestHandler } from 'vscode-languageserver';
 import {
     AuthoredResource,
     GetAuthoredResourceTypesRequest,
+    GetAuthoredResourceTypesRequestV2,
     GetRelatedResourceTypesParams,
     GetRelatedResourceTypesRequest,
     InsertRelatedResourcesParams,
@@ -13,10 +14,14 @@ import {
 export class LspRelatedResourcesHandlers {
     constructor(private readonly connection: Connection) {}
 
-    onGetAuthoredResourceTypes(handler: RequestHandler<TemplateUri, AuthoredResource[], void>) {
+    onGetAuthoredResourceTypes(handler: RequestHandler<TemplateUri, string[], void>) {
         this.connection.onRequest(GetAuthoredResourceTypesRequest.method, handler);
     }
 
+    onGetAuthoredResourceTypesV2(handler: RequestHandler<TemplateUri, AuthoredResource[], void>) {
+        this.connection.onRequest(GetAuthoredResourceTypesRequestV2.method, handler);
+    }
+
     onGetRelatedResourceTypes(handler: RequestHandler<GetRelatedResourceTypesParams, string[], void>) {
         this.connection.onRequest(GetRelatedResourceTypesRequest.method, handler);
     }

src/protocol/RelatedResourcesProtocol.ts

@@ -25,10 +25,14 @@ export interface RelatedResourcesCodeAction extends CodeAction {
     };
 }
 
-export const GetAuthoredResourceTypesRequest = new RequestType<TemplateUri, AuthoredResource[], void>(
+export const GetAuthoredResourceTypesRequest = new RequestType<TemplateUri, string[], void>(
     'aws/cfn/template/resources/authored',
 );
 
+export const GetAuthoredResourceTypesRequestV2 = new RequestType<TemplateUri, AuthoredResource[], void>(
+    'aws/cfn/template/resources/authored/v2',
+);
+
 export const GetRelatedResourceTypesRequest = new RequestType<GetRelatedResourceTypesParams, string[], void>(
     'aws/cfn/template/resources/related',
 );

src/server/CfnServer.ts

@@ -13,6 +13,7 @@ import { hoverHandler } from '../handlers/HoverHandler';
 import { initializedHandler } from '../handlers/Initialize';
 import {
     getAuthoredResourceTypesHandler,
+    getAuthoredResourceTypesHandlerV2,
     getRelatedResourceTypesHandler,
     insertRelatedResourcesHandler,
 } from '../handlers/RelatedResourcesHandler';
@@ -246,6 +247,12 @@ export class CfnServer {
         this.lsp.relatedResourcesHandlers.onGetAuthoredResourceTypes(
             withTelemetryContext('Related.Resources.Get.Authored', getAuthoredResourceTypesHandler(this.components)),
         );
+        this.lsp.relatedResourcesHandlers.onGetAuthoredResourceTypesV2(
+            withTelemetryContext(
+                'Related.Resources.Get.Authored.V2',
+                getAuthoredResourceTypesHandlerV2(this.components),
+            ),
+        );
         this.lsp.relatedResourcesHandlers.onGetRelatedResourceTypes(
             withTelemetryContext('Related.Resources.Get.Related', getRelatedResourceTypesHandler(this.components)),
         );

tst/unit/handlers/RelatedResourcesHandler.test.ts

@@ -3,6 +3,7 @@ import { CancellationToken } from 'vscode-languageserver';
 import { getEntityMap } from '../../../src/context/SectionContextBuilder';
 import {
     getAuthoredResourceTypesHandler,
+    getAuthoredResourceTypesHandlerV2,
     getRelatedResourceTypesHandler,
     insertRelatedResourcesHandler,
 } from '../../../src/handlers/RelatedResourcesHandler';
@@ -37,9 +38,89 @@ describe('RelatedResourcesHandler', () => {
     });
 
     describe('getAuthoredResourceTypesHandler', () => {
-        it('should return authored resources with logical IDs and types', () => {
+        it('should return resource type strings', () => {
+            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const templateUri = 'file:///test/template.yaml';
+
+            syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any);
+            mockGetEntityMap.mockReturnValue(
+                new Map([
+                    ['Bucket1', { entity: { Type: 'AWS::S3::Bucket' } }],
+                    ['Function1', { entity: { Type: 'AWS::Lambda::Function' } }],
+                ]),
+            );
+
+            const result = handler(templateUri, mockToken);
+
+            expect(result).toEqual(['AWS::S3::Bucket', 'AWS::Lambda::Function']);
+        });
+
+        it('should return empty array when no syntax tree found', () => {
+            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const templateUri = 'file:///test/template.yaml';
+
+            syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns(undefined);
+
+            const result = handler(templateUri, mockToken);
+
+            expect(result).toEqual([]);
+        });
+
+        it('should return empty array when no resources found', () => {
+            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const templateUri = 'file:///test/template.yaml';
+
+            syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any);
+            mockGetEntityMap.mockReturnValue(undefined);
+
+            const result = handler(templateUri, mockToken);
+
+            expect(result).toEqual([]);
+        });
+
+        it('should filter out undefined and null resource types', () => {
+            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const templateUri = 'file:///test/template.yaml';
+
+            const mockResourceContext1 = {
+                entity: { Type: 'AWS::S3::Bucket' },
+            };
+            const mockResourceContext2 = {
+                entity: { Type: undefined as any },
+            };
+            const mockResourceContext3 = {
+                entity: { Type: null as any },
+            };
+
+            syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any);
+            mockGetEntityMap.mockReturnValue(
+                new Map([
+                    ['Bucket1', mockResourceContext1],
+                    ['Resource2', mockResourceContext2],
+                    ['Resource3', mockResourceContext3],
+                ]) as any,
+            );
+
+            const result = handler(templateUri, mockToken);
+
+            expect(result).toEqual(['AWS::S3::Bucket']);
+        });
+
+        it('should handle errors and rethrow them', () => {
             const handler = getAuthoredResourceTypesHandler(mockComponents);
             const templateUri = 'file:///test/template.yaml';
+            const error = new Error('Syntax tree error');
+
+            syntaxTreeManager.getSyntaxTree.withArgs(templateUri).throws(error);
+
+            expect(() => handler(templateUri, mockToken)).toThrow('Syntax tree error');
+        });
+    });
+
+    describe('getAuthoredResourceTypesHandlerV2', () => {
+        it('should return authored resources with logical IDs and types', () => {
+            const handler = getAuthoredResourceTypesHandlerV2(mockComponents);
+            const templateUri = 'file:///test/template.yaml';
 
             const mockResourceContext1 = {
                 entity: { Type: 'AWS::S3::Bucket' },
@@ -71,7 +152,7 @@ describe('RelatedResourcesHandler', () => {
         });
 
         it('should return empty array when no syntax tree found', () => {
-            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const handler = getAuthoredResourceTypesHandlerV2(mockComponents);
             const templateUri = 'file:///test/template.yaml';
 
             syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns(undefined);
@@ -82,7 +163,7 @@ describe('RelatedResourcesHandler', () => {
         });
 
         it('should return empty array when no resources found', () => {
-            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const handler = getAuthoredResourceTypesHandlerV2(mockComponents);
             const templateUri = 'file:///test/template.yaml';
 
             syntaxTreeManager.getSyntaxTree.withArgs(templateUri).returns({} as any);
@@ -94,7 +175,7 @@ describe('RelatedResourcesHandler', () => {
         });
 
         it('should filter out undefined and null resource types', () => {
-            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const handler = getAuthoredResourceTypesHandlerV2(mockComponents);
             const templateUri = 'file:///test/template.yaml';
 
             const mockResourceContext1 = {
@@ -122,7 +203,7 @@ describe('RelatedResourcesHandler', () => {
         });
 
         it('should handle errors and rethrow them', () => {
-            const handler = getAuthoredResourceTypesHandler(mockComponents);
+            const handler = getAuthoredResourceTypesHandlerV2(mockComponents);
             const templateUri = 'file:///test/template.yaml';
             const error = new Error('Syntax tree error');
 
@@ -133,6 +214,46 @@ describe('RelatedResourcesHandler', () => {
     });
 
     describe('getRelatedResourceTypesHandler', () => {
+        it('should return related resource types for a given resource type', () => {
+            const handler = getRelatedResourceTypesHandler(mockComponents);
+            const params = { parentResourceType: 'AWS::S3::Bucket' };
+
+            const relatedTypes = new Set(['AWS::Lambda::Function', 'AWS::IAM::Role']);
+            relationshipSchemaService.getAllRelatedResourceTypes.withArgs('AWS::S3::Bucket').returns(relatedTypes);
+
+            relationshipSchemaService.getRelationshipsForResourceType.withArgs('AWS::Lambda::Function').returns({
+                resourceType: 'AWS::Lambda::Function',
+                relationships: [
+                    {
+                        property: 'BucketName',
+                        relatedResourceTypes: [{ typeName: 'AWS::S3::Bucket', attribute: '/properties/BucketName' }],
+                    },
+                ],
+            });
+
+            relationshipSchemaService.getRelationshipsForResourceType.withArgs('AWS::IAM::Role').returns({
+                resourceType: 'AWS::IAM::Role',
+                relationships: [
+                    {
+                        property: 'BucketArn',
+                        relatedResourceTypes: [{ typeName: 'AWS::S3::Bucket', attribute: '/properties/Arn' }],
+                    },
+                ],
+            });
+
+            mockComponents.schemaRetriever.getDefault.returns({
+                schemas: new Map([
+                    ['AWS::Lambda::Function', { properties: { BucketName: { type: 'string' } } }],
+                    ['AWS::IAM::Role', { properties: { BucketArn: { type: 'string' } } }],
+                ]),
+            } as any);
+
+            const result = handler(params, mockToken);
+
+            expect(result).toEqual(['AWS::Lambda::Function', 'AWS::IAM::Role']);
+            expect(relationshipSchemaService.getAllRelatedResourceTypes.calledWith('AWS::S3::Bucket')).toBe(true);
+        });
+
         it('should return related resource types that have exactly one populatable relationship', () => {
             const handler = getRelatedResourceTypesHandler(mockComponents);
             const params = { parentResourceType: 'AWS::S3::Bucket' };
@@ -343,6 +464,40 @@ describe('RelatedResourcesHandler', () => {
     });
 
     describe('insertRelatedResourcesHandler', () => {
+        it('should insert related resources and return code action', () => {
+            const handler = insertRelatedResourcesHandler(mockComponents);
+            const params = {
+                templateUri: 'file:///test/template.yaml',
+                relatedResourceTypes: ['AWS::Lambda::Function', 'AWS::IAM::Role'],
+                parentResourceType: 'AWS::S3::Bucket',
+            };
+
+            const mockCodeAction = {
+                title: 'Insert 2 related resources',
+                kind: 'refactor',
+                edit: {
+                    changes: {
+                        'file:///test/template.yaml': [],
+                    },
+                },
+            };
+
+            mockComponents.relatedResourcesSnippetProvider.insertRelatedResources
+                .withArgs('file:///test/template.yaml', ['AWS::Lambda::Function', 'AWS::IAM::Role'], 'AWS::S3::Bucket')
+                .returns(mockCodeAction);
+
+            const result = handler(params, mockToken);
+
+            expect(result).toEqual(mockCodeAction);
+            expect(
+                mockComponents.relatedResourcesSnippetProvider.insertRelatedResources.calledWith(
+                    'file:///test/template.yaml',
+                    ['AWS::Lambda::Function', 'AWS::IAM::Role'],
+                    'AWS::S3::Bucket',
+                ),
+            ).toBe(true);
+        });
+
         it('should insert related resources and return code action without parentLogicalId', () => {
             const handler = insertRelatedResourcesHandler(mockComponents);
             const params = {

tst/unit/protocol/LspRelatedResourcesHandlers.test.ts

@@ -5,6 +5,7 @@ import { LspRelatedResourcesHandlers } from '../../../src/protocol/LspRelatedRes
 import {
     AuthoredResource,
     GetAuthoredResourceTypesRequest,
+    GetAuthoredResourceTypesRequestV2,
     GetRelatedResourceTypesParams,
     GetRelatedResourceTypesRequest,
     InsertRelatedResourcesParams,
@@ -23,13 +24,21 @@ describe('LspRelatedResourcesHandlers', () => {
     });
 
     it('should register onGetAuthoredResourceTypes handler', () => {
-        const mockHandler: RequestHandler<TemplateUri, AuthoredResource[], void> = vi.fn();
+        const mockHandler: RequestHandler<TemplateUri, string[], void> = vi.fn();
 
         relatedResourcesHandlers.onGetAuthoredResourceTypes(mockHandler);
 
         expect(connection.onRequest.calledWith(GetAuthoredResourceTypesRequest.method)).toBe(true);
     });
 
+    it('should register onGetAuthoredResourceTypesV2 handler', () => {
+        const mockHandler: RequestHandler<TemplateUri, AuthoredResource[], void> = vi.fn();
+
+        relatedResourcesHandlers.onGetAuthoredResourceTypesV2(mockHandler);
+
+        expect(connection.onRequest.calledWith(GetAuthoredResourceTypesRequestV2.method)).toBe(true);
+    });
+
     it('should register onGetRelatedResourceTypes handler', () => {
         const mockHandler: RequestHandler<GetRelatedResourceTypesParams, string[], void> = vi.fn();
 

tst/unit/server/CfnServer.test.ts

@@ -48,6 +48,7 @@ describe('CfnServer', () => {
             expect(mockFeatures.stackHandlers.onListStackResources.calledOnce).toBe(true);
 
             expect(mockFeatures.relatedResourcesHandlers.onGetAuthoredResourceTypes.calledOnce).toBe(true);
+            expect(mockFeatures.relatedResourcesHandlers.onGetAuthoredResourceTypesV2.calledOnce).toBe(true);
             expect(mockFeatures.relatedResourcesHandlers.onGetRelatedResourceTypes.calledOnce).toBe(true);
             expect(mockFeatures.relatedResourcesHandlers.onInsertRelatedResources.calledOnce).toBe(true);
         });