reduce duplication across stage implementations

Hweinstock · Hweinstock · commit 45b023699f1b · 2025-01-17T11:16:58.000-05:00
diff --git a/packages/core/src/amazonq/lsp/lspClient.ts b/packages/core/src/amazonq/lsp/lspClient.ts
@@ -33,7 +33,6 @@ import {
 import { Writable } from 'stream'
 import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings'
 import { ResourcePaths, fs, getLogger, globals } from '../../shared'
-import { telemetry } from '../../shared/telemetry'
 
 const localize = nls.loadMessageBundle()
 
@@ -174,108 +173,103 @@ export class LspClient {
  * This function assumes the LSP server has already been downloaded.
  */
 export async function activate(extensionContext: ExtensionContext, resourcePaths: ResourcePaths) {
-    return await telemetry.languageServer_setup.run(async (span) => {
-        span.record({ languageServerSetupStage: 'launch' })
-        const startTime = performance.now()
-        LspClient.instance
-        const toDispose = extensionContext.subscriptions
+    LspClient.instance
+    const toDispose = extensionContext.subscriptions
 
-        let rangeFormatting: Disposable | undefined
-        // The debug options for the server
-        // --inspect=6009: runs the server in Node's Inspector mode so VS Code can attach to the server for debugging
-        const debugOptions = { execArgv: ['--nolazy', '--preserve-symlinks', '--stdio'] }
+    let rangeFormatting: Disposable | undefined
+    // The debug options for the server
+    // --inspect=6009: runs the server in Node's Inspector mode so VS Code can attach to the server for debugging
+    const debugOptions = { execArgv: ['--nolazy', '--preserve-symlinks', '--stdio'] }
 
-        const workerThreads = CodeWhispererSettings.instance.getIndexWorkerThreads()
-        const gpu = CodeWhispererSettings.instance.isLocalIndexGPUEnabled()
+    const workerThreads = CodeWhispererSettings.instance.getIndexWorkerThreads()
+    const gpu = CodeWhispererSettings.instance.isLocalIndexGPUEnabled()
 
-        if (gpu) {
-            process.env.Q_ENABLE_GPU = 'true'
-        } else {
-            delete process.env.Q_ENABLE_GPU
-        }
-        if (workerThreads > 0 && workerThreads < 100) {
-            process.env.Q_WORKER_THREADS = workerThreads.toString()
-        } else {
-            delete process.env.Q_WORKER_THREADS
-        }
+    if (gpu) {
+        process.env.Q_ENABLE_GPU = 'true'
+    } else {
+        delete process.env.Q_ENABLE_GPU
+    }
+    if (workerThreads > 0 && workerThreads < 100) {
+        process.env.Q_WORKER_THREADS = workerThreads.toString()
+    } else {
+        delete process.env.Q_WORKER_THREADS
+    }
 
-        const serverModule = resourcePaths.lsp
+    const serverModule = resourcePaths.lsp
 
-        const child = spawn(resourcePaths.node, [serverModule, ...debugOptions.execArgv])
-        // share an encryption key using stdin
-        // follow same practice of DEXP LSP server
-        writeEncryptionInit(child.stdin)
+    const child = spawn(resourcePaths.node, [serverModule, ...debugOptions.execArgv])
+    // share an encryption key using stdin
+    // follow same practice of DEXP LSP server
+    writeEncryptionInit(child.stdin)
 
-        // If the extension is launch in debug mode the debug server options are use
-        // Otherwise the run options are used
-        let serverOptions: ServerOptions = {
-            run: { module: serverModule, transport: TransportKind.ipc },
-            debug: { module: serverModule, transport: TransportKind.ipc, options: debugOptions },
-        }
+    // If the extension is launch in debug mode the debug server options are use
+    // Otherwise the run options are used
+    let serverOptions: ServerOptions = {
+        run: { module: serverModule, transport: TransportKind.ipc },
+        debug: { module: serverModule, transport: TransportKind.ipc, options: debugOptions },
+    }
 
-        serverOptions = () => Promise.resolve(child!)
+    serverOptions = () => Promise.resolve(child!)
 
-        const documentSelector = [{ scheme: 'file', language: '*' }]
+    const documentSelector = [{ scheme: 'file', language: '*' }]
 
-        // Options to control the language client
-        const clientOptions: LanguageClientOptions = {
-            // Register the server for json documents
-            documentSelector,
-            initializationOptions: {
-                handledSchemaProtocols: ['file', 'untitled'], // language server only loads file-URI. Fetching schemas with other protocols ('http'...) are made on the client.
-                provideFormatter: false, // tell the server to not provide formatting capability and ignore the `aws.stepfunctions.asl.format.enable` setting.
-                // this is used by LSP to determine index cache path, move to this folder so that when extension updates index is not deleted.
-                extensionPath: path.join(fs.getUserHomeDir(), '.aws', 'amazonq', 'cache'),
-            },
-            // Log to the Amazon Q Logs so everything is in a single channel
-            // TODO: Add prefix to the language server logs so it is easier to search
-            outputChannel: globals.logOutputChannel,
-        }
+    // Options to control the language client
+    const clientOptions: LanguageClientOptions = {
+        // Register the server for json documents
+        documentSelector,
+        initializationOptions: {
+            handledSchemaProtocols: ['file', 'untitled'], // language server only loads file-URI. Fetching schemas with other protocols ('http'...) are made on the client.
+            provideFormatter: false, // tell the server to not provide formatting capability and ignore the `aws.stepfunctions.asl.format.enable` setting.
+            // this is used by LSP to determine index cache path, move to this folder so that when extension updates index is not deleted.
+            extensionPath: path.join(fs.getUserHomeDir(), '.aws', 'amazonq', 'cache'),
+        },
+        // Log to the Amazon Q Logs so everything is in a single channel
+        // TODO: Add prefix to the language server logs so it is easier to search
+        outputChannel: globals.logOutputChannel,
+    }
 
-        // Create the language client and start the client.
-        LspClient.instance.client = new LanguageClient(
-            'amazonq',
-            localize('amazonq.server.name', 'Amazon Q Language Server'),
-            serverOptions,
-            clientOptions
-        )
-        LspClient.instance.client.registerProposedFeatures()
+    // Create the language client and start the client.
+    LspClient.instance.client = new LanguageClient(
+        'amazonq',
+        localize('amazonq.server.name', 'Amazon Q Language Server'),
+        serverOptions,
+        clientOptions
+    )
+    LspClient.instance.client.registerProposedFeatures()
 
-        const disposable = LspClient.instance.client.start()
-        toDispose.push(disposable)
+    const disposable = LspClient.instance.client.start()
+    toDispose.push(disposable)
 
-        let savedDocument: vscode.Uri | undefined = undefined
+    let savedDocument: vscode.Uri | undefined = undefined
 
-        toDispose.push(
-            vscode.workspace.onDidSaveTextDocument((document) => {
-                if (document.uri.scheme !== 'file') {
-                    return
-                }
-                savedDocument = document.uri
-            }),
-            vscode.window.onDidChangeActiveTextEditor((editor) => {
-                if (savedDocument && editor && editor.document.uri.fsPath !== savedDocument.fsPath) {
-                    void LspClient.instance.updateIndex([savedDocument.fsPath], 'update')
-                }
-            }),
-            vscode.workspace.onDidCreateFiles((e) => {
-                void LspClient.instance.updateIndex(
-                    e.files.map((f) => f.fsPath),
-                    'add'
-                )
-            }),
-            vscode.workspace.onDidDeleteFiles((e) => {
-                void LspClient.instance.updateIndex(
-                    e.files.map((f) => f.fsPath),
-                    'remove'
-                )
-            })
-        )
-        void LspClient.instance.client.onReady().then(() => {
-            const disposableFunc = { dispose: () => rangeFormatting?.dispose() as void }
-            toDispose.push(disposableFunc)
+    toDispose.push(
+        vscode.workspace.onDidSaveTextDocument((document) => {
+            if (document.uri.scheme !== 'file') {
+                return
+            }
+            savedDocument = document.uri
+        }),
+        vscode.window.onDidChangeActiveTextEditor((editor) => {
+            if (savedDocument && editor && editor.document.uri.fsPath !== savedDocument.fsPath) {
+                void LspClient.instance.updateIndex([savedDocument.fsPath], 'update')
+            }
+        }),
+        vscode.workspace.onDidCreateFiles((e) => {
+            void LspClient.instance.updateIndex(
+                e.files.map((f) => f.fsPath),
+                'add'
+            )
+        }),
+        vscode.workspace.onDidDeleteFiles((e) => {
+            void LspClient.instance.updateIndex(
+                e.files.map((f) => f.fsPath),
+                'remove'
+            )
         })
-        span.record({ duration: performance.now() - startTime })
+    )
+    void LspClient.instance.client.onReady().then(() => {
+        const disposableFunc = { dispose: () => rangeFormatting?.dispose() as void }
+        toDispose.push(disposableFunc)
     })
 }
 
diff --git a/packages/core/src/amazonq/lsp/lspController.ts b/packages/core/src/amazonq/lsp/lspController.ts
@@ -15,6 +15,7 @@ import { isCloud9 } from '../../shared/extensionUtilities'
 import globals, { isWeb } from '../../shared/extensionGlobals'
 import { isAmazonInternalOs } from '../../shared/vscode/env'
 import { WorkspaceLSPResolver } from './workspaceInstaller'
+import { lspSetupStage } from './util'
 
 export interface Chunk {
     readonly filePath: string
@@ -160,13 +161,10 @@ export class LspController {
         }
         setImmediate(async () => {
             try {
-                await telemetry.languageServer_setup.run(async (span) => {
-                    const startTime = performance.now()
-                    span.record({ languageServerSetupStage: 'final' })
+                await lspSetupStage('final', async () => {
                     const installResult = await new WorkspaceLSPResolver().resolve()
-                    await activateLsp(context, installResult.resourcePaths)
+                    await lspSetupStage('launch', async () => activateLsp(context, installResult.resourcePaths))
                     getLogger().info('LspController: LSP activated')
-                    span.record({ duration: performance.now() - startTime })
                 })
                 void LspController.instance.buildIndex(buildIndexConfig)
                 // log the LSP server CPU and Memory usage per 30 minutes.
diff --git a/packages/core/src/amazonq/lsp/util.ts b/packages/core/src/amazonq/lsp/util.ts
@@ -0,0 +1,16 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { LanguageServerSetupStage, telemetry } from '../../shared/telemetry'
+
+export async function lspSetupStage<T>(stageName: LanguageServerSetupStage, stage: () => Promise<T>) {
+    return await telemetry.languageServer_setup.run(async (span) => {
+        const startTime = performance.now()
+        const result = await stage()
+        span.record({ languageServerSetupStage: stageName })
+        span.record({ duration: performance.now() - startTime })
+        return result
+    })
+}
diff --git a/packages/core/src/amazonq/lsp/workspaceInstaller.ts b/packages/core/src/amazonq/lsp/workspaceInstaller.ts
@@ -11,6 +11,7 @@ import { Range } from 'semver'
 import { getNodeExecutableName } from '../../shared/lsp/utils/platform'
 import { fs } from '../../shared/fs/fs'
 import { telemetry } from '../../shared/telemetry'
+import { lspSetupStage } from './util'
 
 const manifestUrl = 'https://aws-toolkit-language-servers.amazonaws.com/q-context/manifest.json'
 // this LSP client in Q extension is only going to work with these LSP server versions
@@ -19,32 +20,25 @@ const supportedLspServerVersions = '0.1.32'
 export class WorkspaceLSPResolver implements LspResolver {
     async resolve(): Promise<LspResolution> {
         const name = 'AmazonQ-Workspace'
-        const manifest = await telemetry.languageServer_setup.run(async (span) => {
-            const startTime = performance.now()
-            span.record({ languageServerSetupStage: 'getManifest' })
+        const manifest = await lspSetupStage('getManifest', async () => {
             const result = await new ManifestResolver(manifestUrl, name).resolve()
-            span.record({
-                languageServerResourceLocation: result.location ?? 'unknown',
+            telemetry.record({
                 manifestVersion: result.manifestSchemaVersion,
-                duration: performance.now() - startTime,
+                languageServerResourceLocation: result.location ?? 'unknown',
             })
             return result
         })
         telemetry.record({
             manifestVersion: manifest.manifestSchemaVersion,
         })
-
-        const installationResult = await telemetry.languageServer_setup.run(async (span) => {
-            const startTime = performance.now()
-            span.record({ languageServerSetupStage: 'getServer' })
+        const installationResult = await lspSetupStage('getServer', async () => {
             const result = await new LanguageServerResolver(
                 manifest,
                 name,
                 new Range(supportedLspServerVersions)
             ).resolve()
-            span.record({
+            telemetry.record({
                 languageServerResourceLocation: result.location ?? 'unknown',
-                duration: performance.now() - startTime,
             })
             return result
         })
diff --git a/packages/core/src/shared/lsp/lspResolver.ts b/packages/core/src/shared/lsp/lspResolver.ts
@@ -13,7 +13,7 @@ import { TargetContent, logger, LspResult, LspVersion, Manifest } from './types'
 import { getApplicationSupportFolder } from '../vscode/env'
 import { createHash } from '../crypto'
 import request from '../request'
-import { telemetry } from '../telemetry'
+import { lspSetupStage } from '../../amazonq/lsp/util'
 
 export class LanguageServerResolver {
     constructor(
@@ -187,13 +187,7 @@ export class LanguageServerResolver {
             }
             return []
         })
-        const filesToDownload = await telemetry.languageServer_setup.run(async (span) => {
-            span.record({ languageServerSetupStage: 'validate' })
-            const startTime = performance.now()
-            const result = (await Promise.all(verifyTasks)).flat()
-            span.record({ duration: performance.now() - startTime })
-            return result
-        })
+        const filesToDownload = await lspSetupStage('validate', async () => (await Promise.all(verifyTasks)).flat())
 
         if (filesToDownload.length !== contents.length) {
             return false
