telemetry(lsp): Integrate language server/manifest resolver telemetry

package.json

@@ -40,7 +40,7 @@
         "mergeReports": "ts-node ./scripts/mergeReports.ts"
     },
     "devDependencies": {
-        "@aws-toolkits/telemetry": "^1.0.295",
+        "@aws-toolkits/telemetry": "^1.0.296",
         "@playwright/browser-chromium": "^1.43.1",
         "@stylistic/eslint-plugin": "^2.11.0",
         "@types/he": "^1.2.3",

packages/amazonq/src/lsp/activation.ts

@@ -6,12 +6,14 @@
 import vscode from 'vscode'
 import { startLanguageServer } from './client'
 import { AmazonQLSPResolver } from './lspInstaller'
-import { Commands, ToolkitError } from 'aws-core-vscode/shared'
+import { Commands, lspSetupStage, ToolkitError } from 'aws-core-vscode/shared'
 
 export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
     try {
-        const installResult = await new AmazonQLSPResolver().resolve()
-        await startLanguageServer(ctx, installResult.resourcePaths)
+        await lspSetupStage('all', async () => {
+            const installResult = await new AmazonQLSPResolver().resolve()
+            await lspSetupStage('launch', async () => await startLanguageServer(ctx, installResult.resourcePaths))
+        })
         ctx.subscriptions.push(
             Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
                 await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')

packages/core/package.nls.json

@@ -20,7 +20,7 @@
     "AWS.configuration.description.suppressPrompts": "Prompts which ask for confirmation. Checking an item suppresses the prompt.",
     "AWS.configuration.enableCodeLenses": "Enable SAM hints in source code and template.yaml files",
     "AWS.configuration.description.resources.enabledResources": "AWS resources to display in the 'Resources' portion of the explorer.",
-      "AWS.configuration.description.featureDevelopment.allowRunningCodeAndTests": "Allow /dev to run code and test commands",
+    "AWS.configuration.description.featureDevelopment.allowRunningCodeAndTests": "Allow /dev to run code and test commands",
     "AWS.configuration.description.experiments": "Try experimental features and give feedback. Note that experimental features may be removed at any time.\n * `jsonResourceModification` - Enables basic create, update, and delete support for cloud resources via the JSON Resources explorer component.",
     "AWS.stepFunctions.asl.format.enable.desc": "Enables the default formatter used with Amazon States Language files",
     "AWS.stepFunctions.asl.maxItemsComputed.desc": "The maximum number of outline symbols and folding regions computed (limited for performance reasons).",

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 '../../shared'
 
 export interface Chunk {
     readonly filePath: string
@@ -160,9 +161,7 @@ export class LspController {
         }
         setImmediate(async () => {
             try {
-                const installResult = await new WorkspaceLSPResolver().resolve()
-                await activateLsp(context, installResult.resourcePaths)
-                getLogger().info('LspController: LSP activated')
+                await this.setupLsp(context)
                 void LspController.instance.buildIndex(buildIndexConfig)
                 // log the LSP server CPU and Memory usage per 30 minutes.
                 globals.clock.setInterval(
@@ -183,4 +182,12 @@ export class LspController {
             }
         })
     }
+
+    private async setupLsp(context: vscode.ExtensionContext) {
+        await lspSetupStage('all', async () => {
+            const installResult = await new WorkspaceLSPResolver().resolve()
+            await lspSetupStage('launch', async () => activateLsp(context, installResult.resourcePaths))
+            getLogger().info('LspController: LSP activated')
+        })
+    }
 }

packages/core/src/shared/index.ts

@@ -64,6 +64,7 @@ export { TabTypeDataMap } from '../amazonq/webview/ui/tabs/constants'
 export * from './lsp/manifestResolver'
 export * from './lsp/lspResolver'
 export * from './lsp/types'
+export * from './lsp/utils/setupStage'
 export * from './lsp/utils/cleanup'
 export { default as request } from './request'
 export * from './lsp/utils/platform'

packages/core/src/shared/lsp/lspResolver.ts

@@ -12,6 +12,7 @@ import AdmZip from 'adm-zip'
 import { TargetContent, logger, LspResult, LspVersion, Manifest } from './types'
 import { getApplicationSupportFolder } from '../vscode/env'
 import { createHash } from '../crypto'
+import { lspSetupStage, StageResolver, tryStageResolvers } from './utils/setupStage'
 import { HttpResourceFetcher } from '../resourcefetcher/httpResourceFetcher'
 
 export class LanguageServerResolver {
@@ -30,63 +31,90 @@ export class LanguageServerResolver {
      * @throws ToolkitError if no compatible version can be found
      */
     async resolve() {
-        const result: LspResult = {
-            location: 'unknown',
-            version: '',
-            assetDirectory: '',
-        }
-
         const latestVersion = this.latestCompatibleLspVersion()
         const targetContents = this.getLSPTargetContents(latestVersion)
         const cacheDirectory = this.getDownloadDirectory(latestVersion.serverVersion)
 
-        if (await this.hasValidLocalCache(cacheDirectory, targetContents)) {
-            result.location = 'cache'
-            result.version = latestVersion.serverVersion
-            result.assetDirectory = cacheDirectory
-            return result
-        } else {
-            // Delete the cached directory since it's invalid
-            if (await fs.existsDir(cacheDirectory)) {
-                await fs.delete(cacheDirectory, {
-                    recursive: true,
-                })
-            }
-        }
-
-        if (await this.downloadRemoteTargetContent(targetContents, latestVersion.serverVersion)) {
-            result.location = 'remote'
-            result.version = latestVersion.serverVersion
-            result.assetDirectory = cacheDirectory
-            return result
-        } else {
-            // clean up any leftover content that may have been downloaded
-            if (await fs.existsDir(cacheDirectory)) {
-                await fs.delete(cacheDirectory, {
-                    recursive: true,
-                })
+        const serverResolvers: StageResolver<LspResult>[] = [
+            {
+                resolve: async () => await this.getLocalServer(cacheDirectory, latestVersion, targetContents),
+                telemetryMetadata: { id: this.lsName, languageServerLocation: 'cache' },
+            },
+            {
+                resolve: async () => await this.fetchRemoteServer(cacheDirectory, latestVersion, targetContents),
+                telemetryMetadata: { id: this.lsName, languageServerLocation: 'remote' },
+            },
+            {
+                resolve: async () => await this.getFallbackServer(latestVersion),
+                telemetryMetadata: { id: this.lsName, languageServerLocation: 'fallback' },
+            },
+        ]
+
+        return await tryStageResolvers('getServer', serverResolvers, getServerVersion)
+
+        function getServerVersion(result: LspResult) {
+            return {
+                languageServerVersion: result.version,
             }
         }
+    }
 
-        logger.info(
-            `Unable to download language server version ${latestVersion.serverVersion}. Attempting to fetch from fallback location`
-        )
-
+    private async getFallbackServer(latestVersion: LspVersion): Promise<LspResult> {
         const fallbackDirectory = await this.getFallbackDir(latestVersion.serverVersion)
         if (!fallbackDirectory) {
-            throw new ToolkitError('Unable to find a compatible version of the Language Server')
+            throw new ToolkitError('Unable to find a compatible version of the Language Server', {
+                code: 'IncompatibleVersion',
+            })
         }
 
         const version = path.basename(fallbackDirectory)
         logger.info(
             `Unable to install ${this.lsName} language server v${latestVersion.serverVersion}. Launching a previous version from ${fallbackDirectory}`
         )
 
-        result.location = 'fallback'
-        result.version = version
-        result.assetDirectory = fallbackDirectory
+        return {
+            location: 'fallback',
+            version: version,
+            assetDirectory: fallbackDirectory,
+        }
+    }
+
+    private async fetchRemoteServer(
+        cacheDirectory: string,
+        latestVersion: LspVersion,
+        targetContents: TargetContent[]
+    ): Promise<LspResult> {
+        if (await this.downloadRemoteTargetContent(targetContents, latestVersion.serverVersion)) {
+            return {
+                location: 'remote',
+                version: latestVersion.serverVersion,
+                assetDirectory: cacheDirectory,
+            }
+        } else {
+            throw new ToolkitError('Failed to download server from remote', { code: 'RemoteDownloadFailed' })
+        }
+    }
 
-        return result
+    private async getLocalServer(
+        cacheDirectory: string,
+        latestVersion: LspVersion,
+        targetContents: TargetContent[]
+    ): Promise<LspResult> {
+        if (await this.hasValidLocalCache(cacheDirectory, targetContents)) {
+            return {
+                location: 'cache',
+                version: latestVersion.serverVersion,
+                assetDirectory: cacheDirectory,
+            }
+        } else {
+            // Delete the cached directory since it's invalid
+            if (await fs.existsDir(cacheDirectory)) {
+                await fs.delete(cacheDirectory, {
+                    recursive: true,
+                })
+            }
+            throw new ToolkitError('Failed to retrieve server from cache', { code: 'InvalidCache' })
+        }
     }
 
     /**
@@ -164,25 +192,39 @@ export class LanguageServerResolver {
             await fs.mkdir(downloadDirectory)
         }
 
-        const downloadTasks = contents.map(async (content) => {
-            const res = await new HttpResourceFetcher(content.url, { showUrl: true }).get()
-            if (!res || !res.ok || !res.body) {
-                return false
+        const fetchTasks = contents.map(async (content) => {
+            return {
+                res: await new HttpResourceFetcher(content.url, { showUrl: true }).get(),
+                hash: content.hashes[0],
+                filename: content.filename,
+            }
+        })
+        const fetchResults = await Promise.all(fetchTasks)
+        const verifyTasks = fetchResults.flatMap(async (fetchResult) => {
+            if (!(fetchResult.res && fetchResult.res.ok && fetchResult.res.body)) {
+                return []
             }
 
-            const arrBuffer = await res.arrayBuffer()
+            const arrBuffer = await fetchResult.res.arrayBuffer()
             const data = Buffer.from(arrBuffer)
 
             const hash = createHash('sha384', data)
-            if (hash === content.hashes[0]) {
-                await fs.writeFile(`${downloadDirectory}/${content.filename}`, data)
-                return true
+            if (hash === fetchResult.hash) {
+                return [{ filename: fetchResult.filename, data }]
             }
-            return false
+            return []
         })
-        const downloadResults = await Promise.all(downloadTasks)
-        const downloadResult = downloadResults.every(Boolean)
-        return downloadResult && this.extractZipFilesFromRemote(downloadDirectory)
+        const filesToDownload = await lspSetupStage('validate', async () => (await Promise.all(verifyTasks)).flat())
+
+        if (filesToDownload.length !== contents.length) {
+            return false
+        }
+
+        for (const file of filesToDownload) {
+            await fs.writeFile(`${downloadDirectory}/${file.filename}`, file.data)
+        }
+
+        return this.extractZipFilesFromRemote(downloadDirectory)
     }
 
     private async extractZipFilesFromRemote(downloadDirectory: string) {
@@ -333,7 +375,9 @@ export class LanguageServerResolver {
     private getCompatibleLspTarget(version: LspVersion) {
         // TODO make this web friendly
         // TODO make this fully support windows
-        const platform = process.platform
+
+        // Workaround: Manifest platform field is `windows`, whereas node returns win32
+        const platform = process.platform === 'win32' ? 'windows' : process.platform
         const arch = process.arch
         return version.targets.find((x) => x.arch === arch && x.platform === platform)
     }

packages/core/src/shared/lsp/manifestResolver.ts

@@ -8,6 +8,7 @@ import { ToolkitError } from '../errors'
 import { Timeout } from '../utilities/timeoutUtils'
 import globals from '../extensionGlobals'
 import { Manifest } from './types'
+import { StageResolver, tryStageResolvers } from './utils/setupStage'
 import { HttpResourceFetcher } from '../resourcefetcher/httpResourceFetcher'
 
 const logger = getLogger('lsp')
@@ -32,10 +33,23 @@ export class ManifestResolver {
      * Fetches the latest manifest, falling back to local cache on failure
      */
     async resolve(): Promise<Manifest> {
-        try {
-            return await this.fetchRemoteManifest()
-        } catch (error) {
-            return await this.getLocalManifest()
+        const resolvers: StageResolver<Manifest>[] = [
+            {
+                resolve: async () => await this.fetchRemoteManifest(),
+                telemetryMetadata: { id: this.lsName, manifestLocation: 'remote' },
+            },
+            {
+                resolve: async () => await this.getLocalManifest(),
+                telemetryMetadata: { id: this.lsName, manifestLocation: 'cache' },
+            },
+        ]
+
+        return await tryStageResolvers('getManifest', resolvers, extractMetadata)
+
+        function extractMetadata(r: Manifest) {
+            return {
+                manifestSchemaVersion: r.manifestSchemaVersion,
+            }
         }
     }
 
@@ -52,7 +66,7 @@ export class ManifestResolver {
         const manifest = this.parseManifest(resp.content)
         await this.saveManifest(resp.eTag, resp.content)
         this.checkDeprecation(manifest)
-
+        manifest.location = 'remote'
         return manifest
     }
 
@@ -67,6 +81,7 @@ export class ManifestResolver {
 
         const manifest = this.parseManifest(manifestData.content)
         this.checkDeprecation(manifest)
+        manifest.location = 'cache'
         return manifest
     }
 

packages/core/src/shared/lsp/types.ts

@@ -4,13 +4,12 @@
  */
 
 import { getLogger } from '../logger/logger'
+import { LanguageServerLocation, ManifestLocation } from '../telemetry'
 
 export const logger = getLogger('lsp')
 
-type Location = 'remote' | 'cache' | 'override' | 'fallback' | 'unknown'
-
 export interface LspResult {
-    location: Location
+    location: LanguageServerLocation
     version: string
     assetDirectory: string
 }
@@ -53,6 +52,7 @@ export interface Manifest {
     artifactDescription: string
     isManifestDeprecated: boolean
     versions: LspVersion[]
+    location?: ManifestLocation
 }
 
 export interface VersionRange {