Refactor ServerComponents into separate classes to that Telemetry SDK can startup before other OTEL instrumentations are loaded

Author: satyakigh

src/ai/CfnAI.ts

@@ -1,13 +1,14 @@
 import { BaseMessage } from '@langchain/core/messages';
 import { StructuredTool } from '@langchain/core/tools';
 import { DocumentManager } from '../document/DocumentManager';
-import { Closeable, Configurable, ServerComponents } from '../server/ServerComponents';
+import { CfnExternal } from '../server/CfnExternal';
+import { CfnInfraCore } from '../server/CfnInfraCore';
 import { AwsClient } from '../services/AwsClient';
 import { RelationshipSchemaService } from '../services/RelationshipSchemaService';
 import { getFilteredScannedResources, formatScannedResourcesForAI } from '../services/ResourceScanService';
-import { ISettingsSubscriber } from '../settings/Settings';
-import { ClientMessage } from '../telemetry/ClientMessage';
+import { SettingsConfigurable, ISettingsSubscriber } from '../settings/ISettingsSubscriber';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Closeable } from '../utils/Closeable';
 import { extractErrorMessage } from '../utils/Errors';
 import { toString } from '../utils/String';
 import { Agent } from './Agent';
@@ -17,18 +18,16 @@ import { Prompts } from './Prompts';
 
 const logger = LoggerFactory.getLogger('CfnAI');
 
-export class CfnAI implements Closeable, Configurable {
+export class CfnAI implements SettingsConfigurable, Closeable {
     private readonly llmConfig: LLMConfig;
     private agent?: Agent;
     private mcpTools?: McpTools;
 
     constructor(
-        components: ServerComponents,
-        private readonly documentManager: DocumentManager = components.documentManager,
-        private readonly clientMessage: ClientMessage = components.clientMessage,
-        private readonly awsClient: AwsClient = components.awsClient,
+        private readonly documentManager: DocumentManager,
+        private readonly awsClient: AwsClient,
     ) {
-        this.llmConfig = LLMConfig.create(components);
+        this.llmConfig = LLMConfig.create();
     }
 
     configure(settingsManager: ISettingsSubscriber): void {
@@ -40,7 +39,7 @@ export class CfnAI implements Closeable, Configurable {
             const config = this.llmConfig.get();
             if (config !== undefined) {
                 this.agent = Agent.create(config);
-                this.mcpTools = new McpTools(this.clientMessage);
+                this.mcpTools = new McpTools();
             }
         }
     }
@@ -162,7 +161,7 @@ export class CfnAI implements Closeable, Configurable {
         return this.mcpTools?.close();
     }
 
-    static create(components: ServerComponents): CfnAI {
-        return new CfnAI(components);
+    static create(core: CfnInfraCore, external: CfnExternal): CfnAI {
+        return new CfnAI(core.documentManager, external.awsClient);
     }
 }

src/ai/McpTools.ts

@@ -1,23 +1,20 @@
 import { StructuredTool } from '@langchain/core/tools';
 import { MultiServerMCPClient } from '@langchain/mcp-adapters';
-import { Closeable, Configurable } from '../server/ServerComponents';
-import { DefaultSettings, ISettingsSubscriber, ProfileSettings, SettingsSubscription } from '../settings/Settings';
-import { ClientMessage } from '../telemetry/ClientMessage';
+import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
+import { DefaultSettings, ProfileSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Closeable } from '../utils/Closeable';
 import { extractErrorMessage } from '../utils/Errors';
 
 const logger = LoggerFactory.getLogger('McpTools');
 
-export class McpTools implements Configurable, Closeable {
+export class McpTools implements SettingsConfigurable, Closeable {
     private mcpClient?: MultiServerMCPClient;
     private settingsSubscription?: SettingsSubscription;
     private initializationPromise?: Promise<void>;
     private pendingSettingsChange = false;
 
-    constructor(
-        private readonly clientMessage: ClientMessage,
-        private profile: ProfileSettings = DefaultSettings.profile,
-    ) {
+    constructor(private profile: ProfileSettings = DefaultSettings.profile) {
         void this.initializeMCP();
     }
 

src/ai/llm/LLMConfig.ts

@@ -2,7 +2,6 @@ import { readFileSync, existsSync } from 'fs';
 import { homedir } from 'os';
 import { join } from 'path';
 import { DeepReadonly } from 'ts-essentials';
-import { ServerComponents } from '../../server/ServerComponents';
 import { LoggerFactory } from '../../telemetry/LoggerFactory';
 import { extractErrorMessage } from '../../utils/Errors';
 import { parseWithPrettyError } from '../../utils/ZodErrorWrapper';
@@ -36,7 +35,7 @@ export class LLMConfig {
         return this.config === undefined ? undefined : structuredClone(this.config);
     }
 
-    static create(_components: ServerComponents) {
+    static create() {
         return new LLMConfig();
     }
 }

src/app/standalone.ts

@@ -2,14 +2,13 @@ import { createConnection, InitializeParams, ProposedFeatures } from 'vscode-lan
 import { InitializedParams } from 'vscode-languageserver-protocol';
 import { LspCapabilities } from '../protocol/LspCapabilities';
 import { LspConnection } from '../protocol/LspConnection';
+import { CfnInfraCore } from '../server/CfnInfraCore';
 import { CfnServer } from '../server/CfnServer';
-import { LoggerFactory } from '../telemetry/LoggerFactory';
-import { TelemetryService } from '../telemetry/TelemetryService';
 
 let server: CfnServer | undefined;
 
 function onInitialize(params: InitializeParams) {
-    server = new CfnServer(lsp.features, params);
+    server = new CfnServer(lsp.components, new CfnInfraCore(lsp.components, params));
     return LspCapabilities;
 }
 
@@ -23,10 +22,6 @@ function onShutdown() {
 
 function onExit() {}
 
-// Startup telemetry and loggers before LSP creation
-const _logger = LoggerFactory.instance; // eslint-disable-line @typescript-eslint/no-unused-vars
-const _telemetry = TelemetryService.instance; // eslint-disable-line @typescript-eslint/no-unused-vars
-
 const lsp = new LspConnection(createConnection(ProposedFeatures.all), {
     onInitialize,
     onInitialized,

src/auth/AwsCredentials.ts

@@ -2,7 +2,6 @@ import { AwsCredentialIdentity } from '@aws-sdk/types';
 import { DeepReadonly } from 'ts-essentials';
 import { MessageType } from 'vscode-languageserver-protocol/lib/common/protocol';
 import { LspAuthHandlers } from '../protocol/LspAuthHandlers';
-import { ServerComponents } from '../server/ServerComponents';
 import { DefaultSettings } from '../settings/Settings';
 import { SettingsManager } from '../settings/SettingsManager';
 import { parseProfile } from '../settings/SettingsParser';
@@ -46,7 +45,9 @@ export class AwsCredentials {
         private readonly awsHandlers: LspAuthHandlers,
         private readonly settingsManager: SettingsManager,
         private readonly clientMessage: ClientMessage,
-        private readonly getIAMFromSdk: (profile: string) => Promise<DeepReadonly<AwsCredentialIdentity>>,
+        private readonly getIAMFromSdk: (
+            profile: string,
+        ) => Promise<DeepReadonly<AwsCredentialIdentity>> = sdkIAMCredentials,
     ) {}
 
     getIAM(): Promise<DeepReadonly<AwsCredentialIdentity>> {
@@ -211,13 +212,4 @@ export class AwsCredentials {
             this.logger.error({ error }, 'Error handling SSO token change');
         }
     }
-
-    static create(components: ServerComponents) {
-        return new AwsCredentials(
-            components.authHandlers,
-            components.settingsManager,
-            components.clientMessage,
-            sdkIAMCredentials,
-        );
-    }
 }

src/autocomplete/CompletionRouter.ts

@@ -7,9 +7,13 @@ import { Entity, Output, Parameter } from '../context/semantic/Entity';
 import { EntityType } from '../context/semantic/SemanticTypes';
 import { DocumentType } from '../document/Document';
 import { DocumentManager } from '../document/DocumentManager';
-import { Closeable, Configurable, ServerComponents } from '../server/ServerComponents';
-import { CompletionSettings, DefaultSettings, ISettingsSubscriber, SettingsSubscription } from '../settings/Settings';
+import { CfnExternal } from '../server/CfnExternal';
+import { CfnInfraCore } from '../server/CfnInfraCore';
+import { CfnLspProviders } from '../server/CfnLspProviders';
+import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
+import { CompletionSettings, DefaultSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Closeable } from '../utils/Closeable';
 import { CompletionFormatter } from './CompletionFormatter';
 import { CompletionProvider } from './CompletionProvider';
 import { ConditionCompletionProvider } from './ConditionCompletionProvider';
@@ -28,7 +32,7 @@ export type CompletionProviderType =
     | EntityType;
 const Condition = 'Condition';
 
-export class CompletionRouter implements Configurable, Closeable {
+export class CompletionRouter implements SettingsConfigurable, Closeable {
     private completionSettings: CompletionSettings = DefaultSettings.completion;
     private settingsSubscription?: SettingsSubscription;
     private readonly log = LoggerFactory.getLogger(CompletionRouter);
@@ -268,32 +272,34 @@ export class CompletionRouter implements Configurable, Closeable {
         }
     }
 
-    static create(components: ServerComponents) {
+    static create(core: CfnInfraCore, external: CfnExternal, providers: CfnLspProviders) {
         return new CompletionRouter(
-            components.contextManager,
-            createCompletionProviders(components),
-            components.documentManager,
+            core.contextManager,
+            createCompletionProviders(core, external, providers),
+            core.documentManager,
         );
     }
 }
 
 export function createCompletionProviders(
-    components: ServerComponents,
+    core: CfnInfraCore,
+    external: CfnExternal,
+    providers: CfnLspProviders,
 ): Map<CompletionProviderType, CompletionProvider> {
     const completionProviderMap = new Map<CompletionProviderType, CompletionProvider>();
     completionProviderMap.set(
         'TopLevelSection',
-        new TopLevelSectionCompletionProvider(components.syntaxTreeManager, components.documentManager),
+        new TopLevelSectionCompletionProvider(core.syntaxTreeManager, core.documentManager),
     );
-    completionProviderMap.set(EntityType.Resource, new ResourceSectionCompletionProvider(components));
-    completionProviderMap.set(EntityType.Condition, new ConditionCompletionProvider(components.syntaxTreeManager));
+    completionProviderMap.set(EntityType.Resource, new ResourceSectionCompletionProvider(core, external, providers));
+    completionProviderMap.set(EntityType.Condition, new ConditionCompletionProvider(core.syntaxTreeManager));
     completionProviderMap.set('IntrinsicFunction', new IntrinsicFunctionCompletionProvider());
     completionProviderMap.set(
         'IntrinsicFunctionArgument',
         new IntrinsicFunctionArgumentCompletionProvider(
-            components.syntaxTreeManager,
-            components.schemaRetriever,
-            components.documentManager,
+            core.syntaxTreeManager,
+            external.schemaRetriever,
+            core.documentManager,
         ),
     );
     completionProviderMap.set('ParameterTypeValue', new ParameterTypeValueCompletionProvider());

src/autocomplete/InlineCompletionRouter.ts

@@ -2,17 +2,19 @@ import { InlineCompletionList, InlineCompletionParams, InlineCompletionItem } fr
 import { Context } from '../context/Context';
 import { ContextManager } from '../context/ContextManager';
 import { DocumentManager } from '../document/DocumentManager';
-import { Closeable, Configurable, ServerComponents } from '../server/ServerComponents';
+import { CfnInfraCore } from '../server/CfnInfraCore';
 import { RelationshipSchemaService } from '../services/RelationshipSchemaService';
-import { CompletionSettings, DefaultSettings, ISettingsSubscriber, SettingsSubscription } from '../settings/Settings';
+import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
+import { CompletionSettings, DefaultSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { Closeable } from '../utils/Closeable';
 import { InlineCompletionProvider } from './InlineCompletionProvider';
 import { RelatedResourcesInlineCompletionProvider } from './RelatedResourcesInlineCompletionProvider';
 
 export type InlineCompletionProviderType = 'RelatedResources' | 'ResourceBlock' | 'PropertyBlock';
 type ReturnType = InlineCompletionList | InlineCompletionItem[] | null | undefined;
 
-export class InlineCompletionRouter implements Configurable, Closeable {
+export class InlineCompletionRouter implements SettingsConfigurable, Closeable {
     private completionSettings: CompletionSettings = DefaultSettings.completion;
     private settingsSubscription?: SettingsSubscription;
     private readonly log = LoggerFactory.getLogger(InlineCompletionRouter);
@@ -94,11 +96,11 @@ export class InlineCompletionRouter implements Configurable, Closeable {
         );
     }
 
-    static create(components: ServerComponents) {
+    static create(core: CfnInfraCore) {
         return new InlineCompletionRouter(
-            components.contextManager,
-            createInlineCompletionProviders(components.documentManager, RelationshipSchemaService.getInstance()),
-            components.documentManager,
+            core.contextManager,
+            createInlineCompletionProviders(core.documentManager, RelationshipSchemaService.getInstance()),
+            core.documentManager,
         );
     }
 }

src/autocomplete/ResourceEntityCompletionProvider.ts

@@ -5,7 +5,6 @@ import { DocumentType } from '../document/Document';
 import { DocumentManager } from '../document/DocumentManager';
 import { ResourceSchema } from '../schema/ResourceSchema';
 import { SchemaRetriever } from '../schema/SchemaRetriever';
-import { ServerComponents } from '../server/ServerComponents';
 import { getFuzzySearchFunction } from '../utils/FuzzySearchUtil';
 import { applySnippetIndentation } from '../utils/IndentationUtils';
 import { CompletionFormatter, ExtendedCompletionItem } from './CompletionFormatter';
@@ -110,8 +109,4 @@ export class ResourceEntityCompletionProvider implements CompletionProvider {
 
         return applySnippetIndentation(snippet, documentSpecificSettings, documentType);
     }
-
-    static create(components: ServerComponents) {
-        return new ResourceEntityCompletionProvider(components.schemaRetriever, components.documentManager);
-    }
 }

src/autocomplete/ResourceSectionCompletionProvider.ts

@@ -1,7 +1,9 @@
 import { CompletionItem, CompletionParams, CompletionTriggerKind } from 'vscode-languageserver';
 import { Context } from '../context/Context';
 import { Resource } from '../context/semantic/Entity';
-import { ServerComponents } from '../server/ServerComponents';
+import { CfnExternal } from '../server/CfnExternal';
+import { CfnInfraCore } from '../server/CfnInfraCore';
+import { CfnLspProviders } from '../server/CfnLspProviders';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { CompletionProvider } from './CompletionProvider';
 import { ResourceEntityCompletionProvider } from './ResourceEntityCompletionProvider';
@@ -20,8 +22,10 @@ export class ResourceSectionCompletionProvider implements CompletionProvider {
     private readonly log = LoggerFactory.getLogger(ResourceSectionCompletionProvider);
 
     constructor(
-        components: ServerComponents,
-        private readonly resourceProviders = createResourceCompletionProviders(components),
+        core: CfnInfraCore,
+        external: CfnExternal,
+        providers: CfnLspProviders,
+        private readonly resourceProviders = createResourceCompletionProviders(core, external, providers),
     ) {}
 
     getCompletions(
@@ -74,27 +78,26 @@ export class ResourceSectionCompletionProvider implements CompletionProvider {
 }
 
 export function createResourceCompletionProviders(
-    components: ServerComponents,
+    core: CfnInfraCore,
+    external: CfnExternal,
+    providers: CfnLspProviders,
 ): Map<ResourceCompletionType, CompletionProvider> {
     const resourceProviderMap = new Map<ResourceCompletionType, CompletionProvider>();
     resourceProviderMap.set(
         ResourceCompletionType.Entity,
-        new ResourceEntityCompletionProvider(components.schemaRetriever, components.documentManager),
-    );
-    resourceProviderMap.set(
-        ResourceCompletionType.Type,
-        new ResourceTypeCompletionProvider(components.schemaRetriever),
+        new ResourceEntityCompletionProvider(external.schemaRetriever, core.documentManager),
     );
+    resourceProviderMap.set(ResourceCompletionType.Type, new ResourceTypeCompletionProvider(external.schemaRetriever));
     resourceProviderMap.set(
         ResourceCompletionType.Property,
-        new ResourcePropertyCompletionProvider(components.schemaRetriever),
+        new ResourcePropertyCompletionProvider(external.schemaRetriever),
     );
     resourceProviderMap.set(
         ResourceCompletionType.State,
         new ResourceStateCompletionProvider(
-            components.resourceStateManager,
-            components.documentManager,
-            components.schemaRetriever,
+            providers.resourceStateManager,
+            core.documentManager,
+            external.schemaRetriever,
         ),
     );
     return resourceProviderMap;

src/autocomplete/ResourceStateCompletionProvider.ts

@@ -17,7 +17,9 @@ import { ResourceStatePurpose } from '../resourceState/ResourceStateTypes';
 import { ResourceSchema } from '../schema/ResourceSchema';
 import { SchemaRetriever } from '../schema/SchemaRetriever';
 import { TransformersUtil } from '../schema/transformers/TransformersUtil';
-import { ServerComponents } from '../server/ServerComponents';
+import { CfnExternal } from '../server/CfnExternal';
+import { CfnInfraCore } from '../server/CfnInfraCore';
+import { CfnLspProviders } from '../server/CfnLspProviders';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { CompletionProvider } from './CompletionProvider';
 import { createCompletionItem, handleSnippetJsonQuotes } from './CompletionUtils';
@@ -168,11 +170,11 @@ export class ResourceStateCompletionProvider implements CompletionProvider {
         textEdit.range.start.character = textEdit.range.start.character - tabSize;
     }
 
-    static create(components: ServerComponents) {
+    static create(core: CfnInfraCore, external: CfnExternal, providers: CfnLspProviders) {
         return new ResourceStateCompletionProvider(
-            components.resourceStateManager,
-            components.documentManager,
-            components.schemaRetriever,
+            providers.resourceStateManager,
+            core.documentManager,
+            external.schemaRetriever,
         );
     }
 }