feat: related resources for resources already authored in template showed using inline completion suggestions

Author: Deep Furiya

src/autocomplete/InlineCompletionRouter.ts

@@ -1,6 +1,9 @@
 import { InlineCompletionList, InlineCompletionParams, InlineCompletionItem } from 'vscode-languageserver-protocol';
+import { Context } from '../context/Context';
 import { ContextManager } from '../context/ContextManager';
+import { DocumentManager } from '../document/DocumentManager';
 import { Closeable, Configurable, ServerComponents } from '../server/ServerComponents';
+import { RelationshipSchemaService } from '../services/RelationshipSchemaService';
 import {
     CompletionSettings,
     DefaultSettings,
@@ -9,8 +12,10 @@ import {
     SettingsSubscription,
 } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { InlineCompletionProvider } from './InlineCompletionProvider';
+import { RelatedResourcesInlineCompletionProvider } from './RelatedResourcesInlineCompletionProvider';
 
-export type InlineCompletionProviderType = 'ResourceBlock' | 'PropertyBlock' | 'TemplateSection' | 'AIGenerated';
+export type InlineCompletionProviderType = 'RelatedResources' | 'ResourceBlock' | 'PropertyBlock';
 type ReturnType = InlineCompletionList | InlineCompletionItem[] | null | undefined;
 
 export class InlineCompletionRouter implements Configurable, Closeable {
@@ -20,7 +25,10 @@ export class InlineCompletionRouter implements Configurable, Closeable {
     private editorSettingsSubscription?: SettingsSubscription;
     private readonly log = LoggerFactory.getLogger(InlineCompletionRouter);
 
-    constructor(private readonly contextManager: ContextManager) {}
+    constructor(
+        private readonly contextManager: ContextManager,
+        private readonly inlineCompletionProviderMap: Map<InlineCompletionProviderType, InlineCompletionProvider>,
+    ) {}
 
     getInlineCompletions(params: InlineCompletionParams): Promise<ReturnType> | ReturnType {
         if (!this.completionSettings.enabled) return;
@@ -31,6 +39,31 @@ export class InlineCompletionRouter implements Configurable, Closeable {
             return;
         }
 
+        this.log.debug(
+            {
+                position: params.position,
+                section: context.section,
+                propertyPath: context.propertyPath,
+            },
+            'Processing inline completion request',
+        );
+
+        // Check if we are authoring a new resource
+        if (this.isAuthoringNewResource(context)) {
+            const relatedResourcesProvider = this.inlineCompletionProviderMap.get('RelatedResources');
+            if (relatedResourcesProvider) {
+                const result = relatedResourcesProvider.getlineCompletion(context, params, this.editorSettings);
+
+                if (result instanceof Promise) {
+                    return result.then((items) => {
+                        return { items };
+                    });
+                } else if (result && Array.isArray(result)) {
+                    return { items: result };
+                }
+            }
+        }
+
         return;
     }
 
@@ -42,18 +75,14 @@ export class InlineCompletionRouter implements Configurable, Closeable {
             this.editorSettingsSubscription.unsubscribe();
         }
 
-        // Get initial settings
-        this.completionSettings = settingsManager.getCurrentSettings().completion;
-        this.editorSettings = settingsManager.getCurrentSettings().editor;
-
         // Subscribe to completion settings changes
         this.settingsSubscription = settingsManager.subscribe('completion', (newCompletionSettings) => {
-            this.onCompletionSettingsChanged(newCompletionSettings);
+            this.completionSettings = newCompletionSettings;
         });
 
         // Subscribe to editor settings changes
         this.editorSettingsSubscription = settingsManager.subscribe('editor', (newEditorSettings) => {
-            this.onEditorSettingsChanged(newEditorSettings);
+            this.editorSettings = newEditorSettings;
         });
     }
 
@@ -68,15 +97,36 @@ export class InlineCompletionRouter implements Configurable, Closeable {
         }
     }
 
-    private onCompletionSettingsChanged(settings: CompletionSettings): void {
-        this.completionSettings = settings;
-    }
-
-    private onEditorSettingsChanged(settings: EditorSettings): void {
-        this.editorSettings = settings;
+    private isAuthoringNewResource(context: Context): boolean {
+        // Only provide suggestions in Resources section when positioned for a new resource
+        return (
+            String(context.section) === 'Resources' &&
+            (context.propertyPath.length === 1 ||
+                (context.propertyPath.length === 2 &&
+                    context.hasLogicalId &&
+                    context.isKey() &&
+                    !context.atEntityKeyLevel()))
+        );
     }
 
     static create(components: ServerComponents) {
-        return new InlineCompletionRouter(components.contextManager);
+        return new InlineCompletionRouter(
+            components.contextManager,
+            createInlineCompletionProviders(components.documentManager, RelationshipSchemaService.getInstance()),
+        );
     }
 }
+
+export function createInlineCompletionProviders(
+    documentManager: DocumentManager,
+    relationshipSchemaService: RelationshipSchemaService,
+): Map<InlineCompletionProviderType, InlineCompletionProvider> {
+    const inlineCompletionProviderMap = new Map<InlineCompletionProviderType, InlineCompletionProvider>();
+
+    inlineCompletionProviderMap.set(
+        'RelatedResources',
+        new RelatedResourcesInlineCompletionProvider(relationshipSchemaService, documentManager),
+    );
+
+    return inlineCompletionProviderMap;
+}

src/autocomplete/RelatedResourcesInlineCompletionProvider.ts

@@ -0,0 +1,141 @@
+import { InlineCompletionItem, InlineCompletionParams } from 'vscode-languageserver-protocol';
+import { Context } from '../context/Context';
+import { DocumentManager } from '../document/DocumentManager';
+import { RelationshipSchemaService } from '../services/RelationshipSchemaService';
+import { EditorSettings } from '../settings/Settings';
+import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { InlineCompletionProvider } from './InlineCompletionProvider';
+
+export class RelatedResourcesInlineCompletionProvider implements InlineCompletionProvider {
+    private readonly log = LoggerFactory.getLogger(RelatedResourcesInlineCompletionProvider);
+
+    constructor(
+        private readonly relationshipSchemaService: RelationshipSchemaService,
+        private readonly documentManager: DocumentManager,
+    ) {}
+
+    getlineCompletion(
+        context: Context,
+        params: InlineCompletionParams,
+        _editorSettings: EditorSettings,
+    ): Promise<InlineCompletionItem[]> | InlineCompletionItem[] | undefined {
+        this.log.debug(
+            {
+                provider: 'RelatedResourcesInlineCompletion',
+                position: params.position,
+                section: context.section,
+                propertyPath: context.propertyPath,
+            },
+            'Processing related resources inline completion request',
+        );
+
+        try {
+            const document = this.documentManager.get(params.textDocument.uri);
+            if (!document) {
+                return undefined;
+            }
+
+            const existingResourceTypes = this.relationshipSchemaService.extractResourceTypesFromTemplate(
+                document.getText(),
+            );
+
+            if (existingResourceTypes.length === 0) {
+                return undefined;
+            }
+
+            const relatedResourceTypes = this.getRelatedResourceTypes(existingResourceTypes);
+
+            if (relatedResourceTypes.length === 0) {
+                return undefined;
+            }
+
+            return this.generateInlineCompletionItems(relatedResourceTypes, params);
+        } catch (error) {
+            this.log.error({ error: String(error) }, 'Error generating related resources inline completion');
+            return undefined;
+        }
+    }
+
+    private getRelatedResourceTypes(existingResourceTypes: string[]): string[] {
+        const existingRelationships = new Map<string, Set<string>>();
+        const allRelatedTypes = new Set<string>();
+
+        for (const resourceType of existingResourceTypes) {
+            const relatedTypes = this.relationshipSchemaService.getAllRelatedResourceTypes(resourceType);
+            existingRelationships.set(resourceType, relatedTypes);
+            for (const type of relatedTypes) {
+                allRelatedTypes.add(type);
+            }
+        }
+
+        const existingTypesSet = new Set(existingResourceTypes);
+        const suggestedTypes = [...allRelatedTypes].filter((type) => !existingTypesSet.has(type));
+
+        return this.rankSuggestionsByFrequency(suggestedTypes, existingRelationships);
+    }
+
+    private rankSuggestionsByFrequency(
+        suggestedTypes: string[],
+        existingRelationships: Map<string, Set<string>>,
+    ): string[] {
+        const frequencyMap = new Map<string, number>();
+
+        for (const suggestedType of suggestedTypes) {
+            let frequency = 0;
+            for (const relatedTypes of existingRelationships.values()) {
+                if (relatedTypes.has(suggestedType)) {
+                    frequency++;
+                }
+            }
+            frequencyMap.set(suggestedType, frequency);
+        }
+
+        return suggestedTypes.sort((a, b) => {
+            const freqA = frequencyMap.get(a) ?? 0;
+            const freqB = frequencyMap.get(b) ?? 0;
+
+            if (freqA !== freqB) {
+                return freqB - freqA;
+            }
+
+            return a.localeCompare(b);
+        });
+    }
+
+    private generateInlineCompletionItems(
+        relatedResourceTypes: string[],
+        params: InlineCompletionParams,
+    ): InlineCompletionItem[] {
+        const completionItems: InlineCompletionItem[] = [];
+
+        const topSuggestions = relatedResourceTypes.slice(0, 5);
+
+        for (const resourceType of topSuggestions) {
+            const insertText = this.generatePropertySnippet(resourceType);
+
+            completionItems.push({
+                insertText,
+                range: {
+                    start: params.position,
+                    end: params.position,
+                },
+                filterText: `${resourceType}`,
+            });
+        }
+
+        this.log.debug(
+            {
+                suggestedCount: completionItems.length,
+                suggestions: completionItems.map((item) => item.insertText),
+            },
+            'Generated related resource inline completions',
+        );
+
+        return completionItems;
+    }
+
+    private generatePropertySnippet(resourceType: string): string {
+        // TODO: Convert AWS::Service::Resource to a complete resource type snippet
+        return `${resourceType}:`;
+    }
+}