paginate resource and stack list

Author: Zee2413

src/handlers/ResourceHandler.ts

@@ -13,6 +13,8 @@ import {
     ResourceStateResult,
     ResourceSummary,
     ResourceIdentifier,
+    SearchResourceParams,
+    SearchResourceResult,
 } from '../resourceState/ResourceStateTypes';
 import { ResourceStackManagementResult } from '../resourceState/StackManagementInfoProvider';
 import { ServerComponents } from '../server/ServerComponents';
@@ -38,22 +40,26 @@ export function getResourceTypesHandler(
 
 export function listResourcesHandler(
     components: ServerComponents,
-): ServerRequestHandler<ListResourcesParams, ListResourcesResult, never, void> {
+): RequestHandler<ListResourcesParams, ListResourcesResult, void> {
     return async (params: ListResourcesParams): Promise<ListResourcesResult> => {
         try {
-            const resourceTypes = params.resourceTypes;
-            if (!resourceTypes || resourceTypes.length === 0) {
+            const resourceRequests = params.resources;
+            if (!resourceRequests || resourceRequests.length === 0) {
                 return { resources: [] };
             }
 
             const resources: ResourceSummary[] = [];
 
-            for (const typeName of resourceTypes) {
-                const resourceList = await components.resourceStateManager.listResources(typeName);
+            for (const request of resourceRequests) {
+                const resourceList = await components.resourceStateManager.listResources(
+                    request.resourceType,
+                    request.nextToken,
+                );
                 if (resourceList) {
                     resources.push({
                         typeName: resourceList.typeName,
                         resourceIdentifiers: resourceList.resourceIdentifiers,
+                        nextToken: resourceList.nextToken,
                     });
                 }
             }
@@ -83,17 +89,41 @@ export function refreshResourceListHandler(
                 setTimeout(() => reject(new Error('Resource list refresh timed out')), 30_000),
             );
 
-            return await Promise.race([
-                components.resourceStateManager.refreshResourceList(params.resourceTypes),
-                timeout,
-            ]);
+            const resourceTypes = params.resources.map((r) => r.resourceType);
+            return await Promise.race([components.resourceStateManager.refreshResourceList(resourceTypes), timeout]);
         } catch (error) {
             log.error(error, 'Failed to refresh resource list');
             throw new Error(`Failed to refresh resource list: ${extractErrorMessage(error)}`);
         }
     };
 }
 
+export function searchResourceHandler(
+    components: ServerComponents,
+): ServerRequestHandler<SearchResourceParams, SearchResourceResult, never, void> {
+    return async (params: SearchResourceParams): Promise<SearchResourceResult> => {
+        try {
+            const result = await components.resourceStateManager.searchResourceByIdentifier(
+                params.resourceType,
+                params.identifier,
+            );
+            return {
+                found: result.found,
+                resource: result.resourceList
+                    ? {
+                          typeName: result.resourceList.typeName,
+                          resourceIdentifiers: result.resourceList.resourceIdentifiers,
+                          nextToken: result.resourceList.nextToken,
+                      }
+                    : undefined,
+            };
+        } catch (error) {
+            log.error(error, 'Failed to search resource');
+            return { found: false };
+        }
+    };
+}
+
 export function getStackMgmtInfo(
     components: ServerComponents,
 ): ServerRequestHandler<ResourceIdentifier, ResourceStackManagementResult, never, void> {

src/handlers/StackHandler.ts

@@ -1,4 +1,4 @@
-import { ResponseError, ErrorCodes, RequestHandler } from 'vscode-languageserver';
+import { ErrorCodes, RequestHandler, ResponseError } from 'vscode-languageserver';
 import { TopLevelSection } from '../context/ContextType';
 import { getEntityMap } from '../context/SectionContextBuilder';
 import { Parameter, Resource } from '../context/semantic/Entity';
@@ -8,15 +8,15 @@ import { ServerComponents } from '../server/ServerComponents';
 import { analyzeCapabilities } from '../stacks/actions/CapabilityAnalyzer';
 import { parseStackActionParams, parseTemplateUriParams } from '../stacks/actions/StackActionParser';
 import {
-    GetCapabilitiesResult,
-    TemplateUri,
-    GetParametersResult,
     CreateStackActionParams,
     CreateStackActionResult,
-    GetStackActionStatusResult,
-    DescribeValidationStatusResult,
     DescribeDeploymentStatusResult,
+    DescribeValidationStatusResult,
+    GetCapabilitiesResult,
+    GetParametersResult,
+    GetStackActionStatusResult,
     GetTemplateResourcesResult,
+    TemplateUri,
 } from '../stacks/actions/StackActionRequestType';
 import { ListStacksParams, ListStacksResult } from '../stacks/StackRequestType';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
@@ -239,10 +239,14 @@ export function listStacksHandler(
             if (params.statusToInclude?.length && params.statusToExclude?.length) {
                 throw new Error('Cannot specify both statusToInclude and statusToExclude');
             }
-            return { stacks: await components.cfnService.listStacks(params.statusToInclude, params.statusToExclude) };
+            return await components.stackManager.listStacks(
+                params.statusToInclude,
+                params.statusToExclude,
+                params.loadMore,
+            );
         } catch (error) {
             log.error({ error: extractErrorMessage(error) }, 'Error listing stacks');
-            return { stacks: [] };
+            return { stacks: [], nextToken: undefined };
         }
     };
 }

src/protocol/LspResourceHandlers.ts

@@ -1,4 +1,5 @@
 import { Connection, ServerRequestHandler } from 'vscode-languageserver';
+import { RequestHandler } from 'vscode-languageserver/node';
 import {
     ResourceTypesResult,
     ResourceTypesRequest,
@@ -13,13 +14,16 @@ import {
     RefreshResourcesResult,
     StackMgmtInfoRequest,
     ResourceIdentifier,
+    SearchResourceRequest,
+    SearchResourceParams,
+    SearchResourceResult,
 } from '../resourceState/ResourceStateTypes';
 import { ResourceStackManagementResult } from '../resourceState/StackManagementInfoProvider';
 
 export class LspResourceHandlers {
     constructor(private readonly connection: Connection) {}
 
-    onListResources(handler: ServerRequestHandler<ListResourcesParams, ListResourcesResult, never, void>) {
+    onListResources(handler: RequestHandler<ListResourcesParams, ListResourcesResult, void>) {
         this.connection.onRequest(ListResourcesRequest.method, handler);
     }
 
@@ -38,4 +42,8 @@ export class LspResourceHandlers {
     onStackMgmtInfo(handler: ServerRequestHandler<ResourceIdentifier, ResourceStackManagementResult, never, void>) {
         this.connection.onRequest(StackMgmtInfoRequest.method, handler);
     }
+
+    onSearchResource(handler: ServerRequestHandler<SearchResourceParams, SearchResourceResult, never, void>) {
+        this.connection.onRequest(SearchResourceRequest.method, handler);
+    }
 }

src/resourceState/ResourceStateManager.ts

@@ -1,4 +1,4 @@
-import { GetResourceCommandOutput, ListResourcesOutput, ResourceNotFoundException } from '@aws-sdk/client-cloudcontrol';
+import { GetResourceCommandOutput, ResourceNotFoundException } from '@aws-sdk/client-cloudcontrol';
 import { DateTime } from 'luxon';
 import { SchemaRetriever } from '../schema/SchemaRetriever';
 import { CfnExternal } from '../server/CfnExternal';
@@ -21,6 +21,7 @@ export type ResourceState = {
 type ResourceList = {
     typeName: string;
     resourceIdentifiers: string[];
+    nextToken?: string;
     createdTimestamp: DateTime;
     lastUpdatedTimestamp: DateTime;
 };
@@ -79,19 +80,65 @@ export class ResourceStateManager implements SettingsConfigurable, Closeable {
         return value;
     }
 
-    public async listResources(typeName: string, updateFromLive?: boolean): Promise<ResourceList | undefined> {
-        const cachedResourceList = this.resourceListMap.get(typeName);
-        if (cachedResourceList && !updateFromLive) {
-            return cachedResourceList;
-        }
-        const resourceList = await this.retrieveResourceList(typeName);
-        if (!resourceList) {
+    public async listResources(typeName: string, nextToken?: string): Promise<ResourceList | undefined> {
+        const cached = this.resourceListMap.get(typeName);
+
+        if (!nextToken) {
+            // Initial request - fetch first page and cache it
+            const resourceList = await this.retrieveResourceList(typeName);
+            if (resourceList) {
+                this.resourceListMap.set(typeName, resourceList);
+                return resourceList;
+            }
             return;
         }
 
-        this.resourceListMap.set(typeName, resourceList);
+        // Pagination request - fetch next page and append to cache
+        const resourceListNextPage = await this.retrieveResourceList(typeName, nextToken);
+        if (resourceListNextPage && cached) {
+            // Deduplicate efficiently using Set for O(1) lookup
+            const cachedSet = new Set(cached.resourceIdentifiers);
+            const newIdentifiers = resourceListNextPage.resourceIdentifiers.filter((id) => !cachedSet.has(id));
+            cached.resourceIdentifiers.push(...newIdentifiers);
+            cached.nextToken = resourceListNextPage.nextToken;
+            cached.lastUpdatedTimestamp = DateTime.now();
+            return cached;
+        }
 
-        return resourceList;
+        return resourceListNextPage;
+    }
+
+    public async searchResourceByIdentifier(
+        typeName: string,
+        identifier: string,
+    ): Promise<{ found: boolean; resourceList?: ResourceList }> {
+        const resource = await this.getResource(typeName, identifier);
+        if (!resource) {
+            return { found: false };
+        }
+
+        // Add to cache
+        const cached = this.resourceListMap.get(typeName);
+        if (cached && !cached.resourceIdentifiers.includes(identifier)) {
+            cached.resourceIdentifiers.push(identifier);
+            cached.lastUpdatedTimestamp = DateTime.now();
+            return { found: true, resourceList: cached };
+        }
+
+        // Create new cache entry if doesn't exist
+        if (!cached) {
+            const newList: ResourceList = {
+                typeName,
+                resourceIdentifiers: [identifier],
+                nextToken: undefined,
+                createdTimestamp: DateTime.now(),
+                lastUpdatedTimestamp: DateTime.now(),
+            };
+            this.resourceListMap.set(typeName, newList);
+            return { found: true, resourceList: newList };
+        }
+
+        return { found: true, resourceList: cached };
     }
 
     public getResourceTypes(): string[] {
@@ -113,40 +160,42 @@ export class ResourceStateManager implements SettingsConfigurable, Closeable {
         return resourceIdToStateMap?.get(identifier);
     }
 
-    private async retrieveResourceList(typeName: string): Promise<ResourceList | undefined> {
-        let output: ListResourcesOutput | undefined = undefined;
-
+    private async retrieveResourceList(typeName: string, nextToken?: string): Promise<ResourceList | undefined> {
         try {
-            output = await this.ccapiService.listResources(typeName);
+            const output = await this.ccapiService.listResources(typeName, { nextToken });
+
+            const identifiers =
+                output.ResourceDescriptions?.map((desc) => desc.Identifier).filter(
+                    (id): id is string => id !== undefined,
+                ) ?? [];
+
+            const now = DateTime.now();
+
+            return {
+                typeName: typeName,
+                resourceIdentifiers: identifiers,
+                createdTimestamp: now,
+                lastUpdatedTimestamp: now,
+                nextToken: output.NextToken,
+            };
         } catch (error) {
             log.error(error, `CCAPI ListResource failed for type ${typeName}`);
             return;
         }
-
-        if (!output?.ResourceDescriptions) {
-            return;
-        }
-
-        const now = DateTime.now();
-
-        return {
-            typeName: typeName,
-            resourceIdentifiers: output.ResourceDescriptions.map((desc) => desc.Identifier).filter(
-                (id) => id !== undefined,
-            ),
-            createdTimestamp: now,
-            lastUpdatedTimestamp: now,
-        };
     }
 
     public async refreshResourceList(resourceTypes: string[]): Promise<RefreshResourcesResult> {
         if (this.isRefreshing) {
             // return cached resource list
             return {
-                resources: resourceTypes.map((resourceType) => ({
-                    typeName: resourceType,
-                    resourceIdentifiers: this.resourceListMap.get(resourceType)?.resourceIdentifiers ?? [],
-                })),
+                resources: resourceTypes.map((resourceType) => {
+                    const cached = this.resourceListMap.get(resourceType);
+                    return {
+                        typeName: resourceType,
+                        resourceIdentifiers: cached?.resourceIdentifiers ?? [],
+                        nextToken: cached?.nextToken,
+                    };
+                }),
                 refreshFailed: false,
             };
         }
@@ -158,32 +207,30 @@ export class ResourceStateManager implements SettingsConfigurable, Closeable {
         try {
             this.isRefreshing = true;
             const result: ListResourcesResult = { resources: [] };
-            const now = DateTime.now();
             let anyRefreshFailed = false;
 
             for (const resourceType of resourceTypes) {
-                const storedResourceList = this.resourceListMap.get(resourceType);
+                // Clear cache and fetch first page only
+                this.resourceListMap.delete(resourceType);
 
-                const newResourceList = await this.retrieveResourceList(resourceType);
-                if (!newResourceList) {
-                    // Failed to update this resource type
+                const response = await this.retrieveResourceList(resourceType);
+                if (!response) {
+                    anyRefreshFailed = true;
                     result.resources.push({
                         typeName: resourceType,
-                        resourceIdentifiers: storedResourceList?.resourceIdentifiers ?? [],
+                        resourceIdentifiers: [],
+                        nextToken: undefined,
                     });
-                    anyRefreshFailed = true;
                     continue;
                 }
 
-                this.resourceListMap.set(resourceType, {
-                    ...newResourceList,
-                    createdTimestamp: storedResourceList?.createdTimestamp ?? now,
-                    lastUpdatedTimestamp: now,
-                });
+                // Cache the first page
+                this.resourceListMap.set(resourceType, response);
 
                 result.resources.push({
                     typeName: resourceType,
-                    resourceIdentifiers: newResourceList.resourceIdentifiers,
+                    resourceIdentifiers: response.resourceIdentifiers,
+                    nextToken: response.nextToken,
                 });
             }
             return { ...result, refreshFailed: anyRefreshFailed };