Add ability to remove resource type

Author: Zee2413

src/handlers/ResourceHandler.ts

@@ -3,6 +3,7 @@ import { ServerRequestHandler } from 'vscode-languageserver';
 import { RequestHandler } from 'vscode-languageserver/node';
 import { TopLevelSection } from '../context/ContextType';
 import { getEntityMap } from '../context/SectionContextBuilder';
+import { parseResourceTypeName } from '../resourceState/ResourceStateParser';
 import {
     ResourceTypesResult,
     ListResourcesParams,
@@ -22,6 +23,7 @@ import { GetStackTemplateParams, GetStackTemplateResult } from '../stacks/StackR
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { TelemetryService } from '../telemetry/TelemetryService';
 import { extractErrorMessage } from '../utils/Errors';
+import { parseWithPrettyError } from '../utils/ZodErrorWrapper';
 
 const log = LoggerFactory.getLogger('ResourceHandler');
 
@@ -39,6 +41,13 @@ export function getResourceTypesHandler(
     };
 }
 
+export function removeResourceTypeHandler(components: ServerComponents): RequestHandler<string, void, void> {
+    return (rawParams: string) => {
+        const typeName = parseWithPrettyError(parseResourceTypeName, rawParams);
+        components.resourceStateManager.removeResourceType(typeName);
+    };
+}
+
 export function listResourcesHandler(
     components: ServerComponents,
 ): RequestHandler<ListResourcesParams, ListResourcesResult, void> {

src/protocol/LspResourceHandlers.ts

@@ -17,6 +17,7 @@ import {
     SearchResourceRequest,
     SearchResourceParams,
     SearchResourceResult,
+    RemoveResourceTypeRequest,
 } from '../resourceState/ResourceStateTypes';
 import { ResourceStackManagementResult } from '../resourceState/StackManagementInfoProvider';
 
@@ -35,6 +36,10 @@ export class LspResourceHandlers {
         this.connection.onRequest(ResourceTypesRequest.method, handler);
     }
 
+    onRemoveResourceType(handler: RequestHandler<string, void, void>) {
+        this.connection.onRequest(RemoveResourceTypeRequest.method, handler);
+    }
+
     onResourceStateImport(handler: ServerRequestHandler<ResourceStateParams, ResourceStateResult, never, void>) {
         this.connection.onRequest(ResourceStateRequest.method, handler);
     }

src/resourceState/ResourceStateManager.ts

@@ -12,7 +12,7 @@ import { ISettingsSubscriber, SettingsConfigurable, SettingsSubscription } from
 import { DefaultSettings, ProfileSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
-import { Telemetry, Measure } from '../telemetry/TelemetryDecorator';
+import { Telemetry, Measure, Count } from '../telemetry/TelemetryDecorator';
 import { Closeable } from '../utils/Closeable';
 import { handleLspError } from '../utils/Errors';
 import { NO_LIST_SUPPORT, REQUIRES_RESOURCE_MODEL } from './ListResourcesExclusionTypes';
@@ -170,6 +170,12 @@ export class ResourceStateManager implements SettingsConfigurable, Closeable {
         return [...listableTypes].filter((type) => !NO_LIST_SUPPORT.has(type) && !REQUIRES_RESOURCE_MODEL.has(type));
     }
 
+    @Count({ name: 'removeResourceType' })
+    public removeResourceType(typeName: string) {
+        this.resourceListMap.delete(typeName);
+        this.resourceStateMap.delete(typeName);
+    }
+
     private storeResourceState(typeName: ResourceType, id: ResourceId, state: ResourceState) {
         let resourceIdToStateMap = this.resourceStateMap.get(typeName);
         if (!resourceIdToStateMap) {

src/resourceState/ResourceStateParser.ts

@@ -0,0 +1,7 @@
+import { z } from 'zod';
+
+const ResourceTypeNameSchema = z.string().min(1);
+
+export function parseResourceTypeName(input: unknown): string {
+    return ResourceTypeNameSchema.parse(input);
+}

src/resourceState/ResourceStateTypes.ts

@@ -99,6 +99,8 @@ export const RefreshResourceListRequest = new RequestType<RefreshResourcesParams
     'aws/cfn/resources/refresh',
 );
 
+export const RemoveResourceTypeRequest = new RequestType<string, void, void>('aws/cfn/resources/list/remove');
+
 export const DeletionPolicyOnImport = 'Retain';
 
 export interface ResourceTemplateFormat {

src/server/CfnServer.ts

@@ -24,6 +24,7 @@ import {
     refreshResourceListHandler,
     searchResourceHandler,
     getStackMgmtInfo,
+    removeResourceTypeHandler,
 } from '../handlers/ResourceHandler';
 import { uploadFileToS3Handler } from '../handlers/S3Handler';
 import {
@@ -212,6 +213,9 @@ export class CfnServer {
         this.lsp.resourceHandlers.onGetResourceTypes(
             withTelemetryContext('Resource.Get.Types', getResourceTypesHandler(this.components)),
         );
+        this.lsp.resourceHandlers.onRemoveResourceType(
+            withTelemetryContext('Resource.Remove.Type', removeResourceTypeHandler(this.components)),
+        );
         this.lsp.resourceHandlers.onResourceStateImport(
             withTelemetryContext('Resource.State.Import', importResourceStateHandler(this.components)),
         );

tst/unit/handlers/ResourceHandler.test.ts

@@ -1,7 +1,10 @@
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { CancellationToken } from 'vscode-languageserver-protocol';
 import { getEntityMap } from '../../../src/context/SectionContextBuilder';
-import { getManagedResourceStackTemplateHandler } from '../../../src/handlers/ResourceHandler';
+import {
+    getManagedResourceStackTemplateHandler,
+    removeResourceTypeHandler,
+} from '../../../src/handlers/ResourceHandler';
 import { GetStackTemplateParams } from '../../../src/stacks/StackRequestType';
 import { createMockComponents } from '../../utils/MockServerComponents';
 
@@ -121,3 +124,35 @@ describe('ResourceHandler - getManagedResourceStackTemplateHandler', () => {
         await expect(handler(params, CancellationToken.None)).rejects.toThrow('AWS API Error');
     });
 });
+
+describe('ResourceHandler - removeResourceTypeHandler', () => {
+    let mockComponents: ReturnType<typeof createMockComponents>;
+    let handler: any;
+
+    beforeEach(() => {
+        vi.clearAllMocks();
+        mockComponents = createMockComponents();
+        handler = removeResourceTypeHandler(mockComponents);
+    });
+
+    it('should call resourceStateManager.removeResourceType with typeName', () => {
+        handler('AWS::S3::Bucket');
+
+        expect(mockComponents.resourceStateManager.removeResourceType.calledOnceWith('AWS::S3::Bucket')).toBe(true);
+    });
+
+    it('should handle multiple calls', () => {
+        handler('AWS::S3::Bucket');
+        handler('AWS::Lambda::Function');
+
+        expect(mockComponents.resourceStateManager.removeResourceType.callCount).toBe(2);
+        expect(mockComponents.resourceStateManager.removeResourceType.calledWith('AWS::S3::Bucket')).toBe(true);
+        expect(mockComponents.resourceStateManager.removeResourceType.calledWith('AWS::Lambda::Function')).toBe(true);
+    });
+
+    it('should throw error for invalid input', () => {
+        expect(() => handler('')).toThrow(TypeError);
+        expect(() => handler(null as any)).toThrow();
+        expect(() => handler(undefined as any)).toThrow();
+    });
+});

tst/unit/resourceState/ResourceStateManager.test.ts

@@ -520,4 +520,33 @@ describe('ResourceStateManager', () => {
             expect(result).not.toContain('MyOrg::Custom::Resource');
         });
     });
+
+    describe('removeResourceType()', () => {
+        it('should remove resource type from both maps', async () => {
+            const mockOutput: GetResourceCommandOutput = {
+                TypeName: 'AWS::S3::Bucket',
+                ResourceDescription: {
+                    Identifier: 'my-bucket',
+                    Properties: '{"BucketName": "my-bucket"}',
+                },
+                $metadata: {},
+            };
+            vi.mocked(mockCcapiService.getResource).mockResolvedValue(mockOutput);
+            vi.mocked(mockCcapiService.listResources).mockResolvedValue({
+                ResourceDescriptions: [{ Identifier: 'my-bucket' }],
+            });
+
+            await manager.getResource('AWS::S3::Bucket', 'my-bucket');
+            await manager.listResources('AWS::S3::Bucket');
+
+            manager.removeResourceType('AWS::S3::Bucket');
+
+            await manager.getResource('AWS::S3::Bucket', 'my-bucket');
+            expect(mockCcapiService.getResource).toHaveBeenCalledTimes(2);
+        });
+
+        it('should handle removing non-existent resource type', () => {
+            expect(() => manager.removeResourceType('AWS::DynamoDB::Table')).not.toThrow();
+        });
+    });
 });

tst/unit/resourceState/ResourceStateParser.test.ts

@@ -0,0 +1,22 @@
+import { describe, it, expect } from 'vitest';
+import { z } from 'zod';
+import { parseResourceTypeName } from '../../../src/resourceState/ResourceStateParser';
+
+describe('ResourceStateParser', () => {
+    describe('parseResourceTypeName', () => {
+        it('should parse valid resource type name', () => {
+            const result = parseResourceTypeName('AWS::S3::Bucket');
+            expect(result).toBe('AWS::S3::Bucket');
+        });
+
+        it('should throw error for empty string', () => {
+            expect(() => parseResourceTypeName('')).toThrow(z.ZodError);
+        });
+
+        it('should throw error for non-string input', () => {
+            expect(() => parseResourceTypeName(123)).toThrow(z.ZodError);
+            expect(() => parseResourceTypeName(null)).toThrow(z.ZodError);
+            expect(() => parseResourceTypeName(undefined)).toThrow(z.ZodError);
+        });
+    });
+});