Speed up autocomplete intrinsic fns (#230)

Author: kddejong

src/autocomplete/CompletionRouter.ts

@@ -64,7 +64,7 @@ export class CompletionRouter implements SettingsConfigurable, Closeable {
         const triggerChar = params.context?.triggerCharacter ?? '';
 
         // Check for intrinsic function argument completions first
-        if (context.intrinsicContext.inIntrinsic()) {
+        if (context.intrinsicContext.inIntrinsic() && !this.shouldUseIntrinsicFunctionProvider(context)) {
             const doc = this.completionProviderMap.get('IntrinsicFunctionArgument')?.getCompletions(context, params);
 
             if (doc && !(doc instanceof Promise) && doc.length > 0) {

src/autocomplete/IntrinsicFunctionArgumentCompletionProvider.ts

@@ -279,23 +279,9 @@ export class IntrinsicFunctionArgumentCompletionProvider implements CompletionPr
 
     private getResourceAttributes(resourceType: string): string[] {
         const schema = this.schemaRetriever.getDefault().schemas.get(resourceType);
-        if (!schema?.readOnlyProperties || schema.readOnlyProperties.length === 0) {
-            return [];
-        }
+        if (!schema) return [];
 
-        return schema.readOnlyProperties
-            .map((propertyPath) => {
-                const match = propertyPath.match(/^\/properties\/(.+)$/);
-                return match ? match[1].replaceAll('/', '.') : undefined;
-            })
-            .filter((attr): attr is string => attr !== undefined)
-            .filter((attr) => {
-                const lastDotIndex = attr.lastIndexOf('.');
-                if (lastDotIndex === -1) return true;
-                const pathWithoutLastSegment = attr.slice(0, Math.max(0, lastDotIndex));
-                return !pathWithoutLastSegment.includes('*');
-            })
-            .filter((attr, index, array) => array.indexOf(attr) === index);
+        return schema.getAttributes().map((attr) => attr.name);
     }
 
     private getGetAttCompletions(syntaxTree: SyntaxTree, currentLogicalId?: string): CompletionItem[] {

src/schema/GetSchemaTaskManager.ts

@@ -29,6 +29,7 @@ export class GetSchemaTaskManager implements SettingsConfigurable, Closeable {
         private readonly getPublicSchemas: (region: AwsRegion) => Promise<SchemaFileType[]>,
         private readonly getPrivateResources: () => Promise<DescribeTypeOutput[]>,
         private profile: string = DefaultSettings.profile.profile,
+        private readonly onSchemaUpdate: (region?: string, profile?: string) => void,
     ) {
         this.privateTask = new GetPrivateSchemasTask(this.getPrivateResources, () => this.profile);
         this.samTask = new GetSamSchemaTask();
@@ -73,14 +74,18 @@ export class GetSchemaTaskManager implements SettingsConfigurable, Closeable {
     runPrivateTask() {
         this.privateTask
             .run(this.schemas.privateSchemas, this.log)
-            .then(() => this.schemas.invalidateCombinedSchemas())
+            .then(() => {
+                this.onSchemaUpdate(undefined, this.profile);
+            })
             .catch(() => {});
     }
 
     runSamTask() {
         this.samTask
             .run(this.schemas.samSchemas, this.log)
-            .then(() => this.schemas.invalidateCombinedSchemas())
+            .then(() => {
+                this.onSchemaUpdate(); // No params = SAM update
+            })
             .catch(() => {});
     }
 
@@ -99,7 +104,9 @@ export class GetSchemaTaskManager implements SettingsConfigurable, Closeable {
         const task = this.tasks.shift();
         if (task) {
             task.run(this.schemas.publicSchemas, this.log)
-                .then(() => this.schemas.invalidateCombinedSchemas())
+                .then(() => {
+                    this.onSchemaUpdate(task.region);
+                })
                 .catch(() => {
                     this.tasks.push(task);
                 })

src/schema/ResourceSchema.ts

@@ -48,6 +48,7 @@ export class ResourceSchema {
     public readonly writeOnlyProperties?: string[];
     public readonly createOnlyProperties?: string[];
     public readonly deprecatedProperties?: string[];
+    private _attributes?: Array<{ name: string; description: string }>;
     public readonly conditionalCreateOnlyProperties?: string[];
     public readonly nonPublicProperties?: string[];
     public readonly nonPublicDefinitions?: string[];
@@ -1041,6 +1042,42 @@ export class ResourceSchema {
             ...(this.oneOf && { oneOf: this.oneOf as PropertyType[] }),
         };
     }
+
+    public getAttributes(): Array<{ name: string; description: string }> {
+        this._attributes ??= this.computeAttributes();
+        return this._attributes;
+    }
+
+    private computeAttributes(): Array<{ name: string; description: string }> {
+        if (!this.readOnlyProperties) return [];
+
+        return this.readOnlyProperties
+            .filter((prop) => !prop.includes('/*/'))
+            .map((prop) => {
+                const match = prop.match(/^\/properties\/(.+)$/);
+                if (!match) return;
+
+                const name = match[1].replaceAll('/', '.');
+                const description = this.getAttributeDescription(prop);
+                return { name, description };
+            })
+            .filter((attr): attr is { name: string; description: string } => attr !== undefined);
+    }
+
+    private getAttributeDescription(propertyPath: string): string {
+        try {
+            const resolvedSchemas = this.resolveJsonPointerPath(propertyPath);
+            if (resolvedSchemas.length > 0 && resolvedSchemas[0].description) {
+                return resolvedSchemas[0].description;
+            }
+        } catch {
+            // Fall back to default description
+        }
+
+        const match = propertyPath.match(/^\/properties\/(.+)$/);
+        const attributeName = match ? match[1].replaceAll('/', '.') : propertyPath;
+        return `${attributeName} attribute of ${this.typeName}`;
+    }
 }
 
 export type PropertyType = {

src/schema/SchemaRetriever.ts

@@ -1,14 +1,16 @@
+import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation';
 import { DateTime } from 'luxon';
 import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
 import { DefaultSettings, ProfileSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
-import { Telemetry } from '../telemetry/TelemetryDecorator';
+import { Telemetry, Measure } from '../telemetry/TelemetryDecorator';
 import { Closeable } from '../utils/Closeable';
 import { AwsRegion, getRegion } from '../utils/Region';
 import { CombinedSchemas } from './CombinedSchemas';
 import { GetSchemaTaskManager } from './GetSchemaTaskManager';
-import { RegionalSchemasType } from './RegionalSchemas';
+import { PrivateSchemasType } from './PrivateSchemas';
+import { RegionalSchemasType, SchemaFileType } from './RegionalSchemas';
 import { SamSchemasType, SamStoreKey } from './SamSchemas';
 import { SchemaStore } from './SchemaStore';
 
@@ -19,14 +21,24 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable {
     private settingsSubscription?: SettingsSubscription;
     private settings: ProfileSettings = DefaultSettings.profile;
     private readonly log = LoggerFactory.getLogger(SchemaRetriever);
+    private readonly schemaTaskManager: GetSchemaTaskManager;
 
     @Telemetry()
     private readonly telemetry!: ScopedTelemetry;
 
     constructor(
-        private readonly schemaTaskManager: GetSchemaTaskManager,
         private readonly schemaStore: SchemaStore,
+        private readonly getPublicSchemas: (region: AwsRegion) => Promise<SchemaFileType[]>,
+        private readonly getPrivateResources: () => Promise<DescribeTypeOutput[]>,
     ) {
+        this.schemaTaskManager = new GetSchemaTaskManager(
+            this.schemaStore,
+            this.getPublicSchemas,
+            this.getPrivateResources,
+            this.settings.profile,
+            (region, profile) => this.rebuildAffectedCombinedSchemas(region, profile),
+        );
+
         this.telemetry.registerGaugeProvider('schema.public.maxAge', () => this.getPublicSchemaMaxAge(), {
             unit: 'ms',
         });
@@ -84,15 +96,31 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable {
         return this.get(this.settings.region, this.settings.profile);
     }
 
+    @Measure({ name: 'getSchemas' })
     get(region: AwsRegion, profile: string): CombinedSchemas {
+        // Check if combined schemas are already cached first
+        const cacheKey = `${region}:${profile}`;
+        const cachedCombined = this.schemaStore.combinedSchemas.get<CombinedSchemas>(cacheKey);
+
+        if (cachedCombined) {
+            return cachedCombined;
+        }
+
+        // Only do expensive regional check if no cached combined schemas
         const regionalSchemas = this.getRegionalSchemasFromStore(region);
         if (regionalSchemas) {
             this.availableRegions.add(region);
         } else {
             this.schemaTaskManager.addTask(region);
         }
 
-        return this.schemaStore.getCombinedSchemas(region, profile);
+        // Create and cache combined schemas
+        const privateSchemas = this.schemaStore.privateSchemas.get<PrivateSchemasType>(profile);
+        const samSchemas = this.schemaStore.samSchemas.get<SamSchemasType>(SamStoreKey);
+        const combinedSchemas = CombinedSchemas.from(regionalSchemas, privateSchemas, samSchemas);
+
+        void this.schemaStore.combinedSchemas.put(cacheKey, combinedSchemas);
+        return combinedSchemas;
     }
 
     updatePrivateSchemas() {
@@ -105,6 +133,34 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable {
         this.schemaStore.invalidateCombinedSchemas();
     }
 
+    // Proactively rebuild combined schemas to avoid lazy loading delays
+    @Measure({ name: 'rebuildCurrentSchemas' })
+    rebuildForCurrentSettings() {
+        this.schemaStore.invalidateCombinedSchemas();
+        this.get(this.settings.region, this.settings.profile);
+    }
+
+    // Surgically rebuild affected combined schemas
+    @Measure({ name: 'rebuildAffectedSchemas' })
+    rebuildAffectedCombinedSchemas(updatedRegion?: string, updatedProfile?: string) {
+        if (!updatedRegion && !updatedProfile) {
+            // SAM update - affects all schemas
+            this.schemaStore.invalidateCombinedSchemas();
+            this.rebuildForCurrentSettings();
+            return;
+        }
+
+        const keys = this.schemaStore.combinedSchemas.keys(1000);
+        for (const key of keys) {
+            const [region, profile] = key.split(':');
+            if ((updatedRegion && region === updatedRegion) || (updatedProfile && profile === updatedProfile)) {
+                // Invalidate and rebuild this specific combined schema
+                void this.schemaStore.combinedSchemas.remove(key);
+                this.get(region as AwsRegion, profile);
+            }
+        }
+    }
+
     private getRegionalSchemasFromStore(region: AwsRegion) {
         return this.schemaStore.publicSchemas.get<RegionalSchemasType>(region);
     }

src/schema/SchemaStore.ts

@@ -1,9 +1,4 @@
 import { DataStoreFactoryProvider, Persistence } from '../datastore/DataStore';
-import { AwsRegion } from '../utils/Region';
-import { CombinedSchemas } from './CombinedSchemas';
-import { PrivateSchemasType } from './PrivateSchemas';
-import { RegionalSchemasType } from './RegionalSchemas';
-import { SamSchemasType, SamStoreKey } from './SamSchemas';
 
 export class SchemaStore {
     public readonly publicSchemas = this.dataStoreFactory.get('public_schemas', Persistence.local);
@@ -13,22 +8,6 @@ export class SchemaStore {
 
     constructor(private readonly dataStoreFactory: DataStoreFactoryProvider) {}
 
-    getCombinedSchemas(region: AwsRegion, profile: string): CombinedSchemas {
-        const cacheKey = `${region}:${profile}`;
-        let cached = this.combinedSchemas.get<CombinedSchemas>(cacheKey);
-
-        if (!cached) {
-            const regionalSchemas = this.publicSchemas.get<RegionalSchemasType>(region);
-            const privateSchemas = this.privateSchemas.get<PrivateSchemasType>(profile);
-            const samSchemas = this.samSchemas.get<SamSchemasType>(SamStoreKey);
-
-            cached = CombinedSchemas.from(regionalSchemas, privateSchemas, samSchemas);
-            void this.combinedSchemas.put(cacheKey, cached);
-        }
-
-        return cached;
-    }
-
     invalidateCombinedSchemas() {
         void this.combinedSchemas.clear();
     }

src/server/CfnExternal.ts

@@ -1,7 +1,6 @@
 import { FeatureFlagProvider } from '../featureFlag/FeatureFlagProvider';
 import { LspComponents } from '../protocol/LspComponents';
 import { getRemotePrivateSchemas, getRemotePublicSchemas } from '../schema/GetSchemaTask';
-import { GetSchemaTaskManager } from '../schema/GetSchemaTaskManager';
 import { SchemaRetriever } from '../schema/SchemaRetriever';
 import { SchemaStore } from '../schema/SchemaStore';
 import { AwsClient } from '../services/AwsClient';
@@ -28,7 +27,6 @@ export class CfnExternal implements Configurables, Closeable {
     readonly s3Service: S3Service;
 
     readonly schemaStore: SchemaStore;
-    readonly schemaTaskManager: GetSchemaTaskManager;
     readonly schemaRetriever: SchemaRetriever;
 
     readonly cfnLintService: CfnLintService;
@@ -46,13 +44,11 @@ export class CfnExternal implements Configurables, Closeable {
         this.s3Service = overrides.s3Service ?? new S3Service(this.awsClient);
 
         this.schemaStore = overrides.schemaStore ?? new SchemaStore(core.dataStoreFactory);
-        this.schemaTaskManager =
-            overrides.schemaTaskManager ??
-            new GetSchemaTaskManager(this.schemaStore, getRemotePublicSchemas, () => {
-                return getRemotePrivateSchemas(core.awsCredentials, this.cfnService);
-            });
         this.schemaRetriever =
-            overrides.schemaRetriever ?? new SchemaRetriever(this.schemaTaskManager, this.schemaStore);
+            overrides.schemaRetriever ??
+            new SchemaRetriever(this.schemaStore, getRemotePublicSchemas, () =>
+                getRemotePrivateSchemas(core.awsCredentials, this.cfnService),
+            );
 
         this.cfnLintService =
             overrides.cfnLintService ??
@@ -66,14 +62,13 @@ export class CfnExternal implements Configurables, Closeable {
     }
 
     configurables(): Configurable[] {
-        return [this.schemaTaskManager, this.schemaRetriever, this.cfnLintService, this.guardService];
+        return [this.schemaRetriever, this.cfnLintService, this.guardService];
     }
 
     async close() {
         return await closeSafely(
             this.cfnLintService,
             this.guardService,
-            this.schemaTaskManager,
             this.schemaRetriever,
             this.onlineStatus,
             this.featureFlags,

tools/telemetry-generator.ts

@@ -200,9 +200,6 @@ function main() {
     const schemaStore = new SchemaStore(dataStoreFactory);
     const external = new CfnExternal(lsp, core, {
         schemaStore,
-        schemaTaskManager: new GetSchemaTaskManager(schemaStore, getRemotePublicSchemas, () =>
-            Promise.resolve(getTestPrivateSchemas()),
-        ),
         featureFlags: new FeatureFlagProvider(join(__dirname, '..', 'assets', 'featureFlag', 'alpha.json')),
     });
 

tst/unit/autocomplete/IntrinsicFunctionArgumentCompletionProvider.GetAtt.test.ts

@@ -629,51 +629,6 @@ describe('IntrinsicFunctionArgumentCompletionProvider - GetAtt Function', () =>
 
                 resolveJsonPointerPathSpy.mockRestore();
             });
-
-            it('should remove duplicate attribute names', () => {
-                // Setup mock schema with duplicates
-                const duplicateSchemaJson = JSON.stringify({
-                    typeName: 'AWS::S3::Bucket',
-                    description: 'Mock S3 Bucket schema',
-                    additionalProperties: false,
-                    primaryIdentifier: ['/properties/BucketName'],
-                    properties: {},
-                    readOnlyProperties: [
-                        '/properties/Arn',
-                        '/properties/Arn', // Duplicate
-                        '/properties/DomainName',
-                        '/properties/MetadataTableConfiguration/S3TablesDestination/TableNamespace',
-                        '/properties/MetadataTableConfiguration/S3TablesDestination/TableArn', // Same top-level
-                    ],
-                });
-                const duplicateSchema = new ResourceSchema(duplicateSchemaJson);
-
-                const duplicateSchemas = new Map([['AWS::S3::Bucket', duplicateSchema]]);
-                const duplicateCombinedSchemas = new CombinedSchemas();
-                (duplicateCombinedSchemas as any).schemas = duplicateSchemas;
-                const duplicateSchemaRetriever = createMockSchemaRetriever(duplicateCombinedSchemas);
-
-                provider = new IntrinsicFunctionArgumentCompletionProvider(
-                    mockSyntaxTreeManager,
-                    duplicateSchemaRetriever,
-                    mockDocumentManager,
-                );
-
-                setupResourceEntitiesWithSchema({ MyS3Bucket: { Type: 'AWS::S3::Bucket' } });
-
-                const mockContext = createMockGetAttContext('', 'MyS3Bucket.');
-
-                const result = provider.getCompletions(mockContext, createTestParams());
-
-                expect(result).toBeDefined();
-                expect(result!.length).toBe(4); // Arn, DomainName, and two nested MetadataTableConfiguration paths (no duplicates)
-
-                const labels = result!.map((item) => item.label);
-                expect(labels).toContain('Arn');
-                expect(labels).toContain('DomainName');
-                expect(labels).toContain('MetadataTableConfiguration.S3TablesDestination.TableNamespace');
-                expect(labels).toContain('MetadataTableConfiguration.S3TablesDestination.TableArn');
-            });
         });
     });
 });