Merge pull request #5378 from hayemaxi/client-id

fix(telemetry): toolkit replacing amazon q clientId

Author: jpinkney-aws

packages/core/src/shared/telemetry/activation.ts

@@ -14,12 +14,11 @@ import { DefaultTelemetryService } from './telemetryService'
 import { getLogger } from '../logger'
 import { getComputeRegion, isAmazonQ, isCloud9, productName } from '../extensionUtilities'
 import { openSettingsId, Settings } from '../settings'
-import { TelemetryConfig, setupTelemetryId } from './util'
+import { TelemetryConfig } from './util'
 import { isAutomation, isReleaseVersion } from '../vscode/env'
 import { AWSProduct } from './clienttelemetry'
 import { DefaultTelemetryClient } from './telemetryClient'
 import { telemetry } from './telemetry'
-import { Commands } from '../vscode/commands2'
 
 export const noticeResponseViewSettings = localize('AWS.telemetry.notificationViewSettings', 'Settings')
 export const noticeResponseOk = localize('AWS.telemetry.notificationOk', 'OK')
@@ -71,19 +70,11 @@ export async function activate(
             })
         )
 
-        if (isAmazonQExt) {
-            extensionContext.subscriptions.push(
-                Commands.register('aws.amazonq.setupTelemetryId', async () => {
-                    await setupTelemetryId(extensionContext)
-                })
-            )
-        }
-
         // Prompt user about telemetry if they haven't been
         if (!isCloud9() && !hasUserSeenTelemetryNotice()) {
             showTelemetryNotice()
         }
-        await setupTelemetryId(extensionContext)
+
         await globals.telemetry.start()
     } catch (e) {
         // Only throw in a production build because:

packages/core/src/shared/telemetry/util.ts

@@ -17,16 +17,14 @@ import { Result } from './telemetry.gen'
 import { MetricDatum } from './clienttelemetry'
 import { isValidationExemptMetric } from './exemptMetrics'
 import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities'
-import { isExtensionInstalled, VSCODE_EXTENSION_ID } from '../utilities'
 import { randomUUID } from '../crypto'
-import { activateExtension } from '../utilities/vsCodeUtils'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
 
 const legacySettingsTelemetryValueDisable = 'Disable'
 const legacySettingsTelemetryValueEnable = 'Enable'
 
 const TelemetryFlag = addTypeName('boolean', convertLegacy)
-const telemetryClientIdEnvKey = '__TELEMETRY_CLIENT_ID'
+export const telemetryClientIdEnvKey = '__TELEMETRY_CLIENT_ID'
 
 export class TelemetryConfig {
     private readonly _toolkitConfig
@@ -83,6 +81,19 @@ export function convertLegacy(value: unknown): boolean {
     }
 }
 
+/**
+ * Calculates the clientId for the current profile. This calculation is performed once
+ * on first call and the result is stored for the remainder of the session.
+ *
+ * Web mode will always compute to whatever is stored in global state or vscode.machineId.
+ * For normal use, the clientId is fetched from the first providing source:
+ * 1. clientId stored in process.env
+ * 2. clientId stored in current extension's global state.
+ * 3. a random UUID
+ *
+ * The clientId in the current extension's global state AND the clientId stored in process.env
+ * is updated to the result of above to allow other extensions to converge to the same clientId.
+ */
 export const getClientId = memoize(
     /**
      * @param nonce Dummy parameter to allow tests to defeat memoize().
@@ -100,11 +111,34 @@ export const getClientId = memoize(
             return '11111111-1111-1111-1111-111111111111'
         }
         try {
-            let clientId = globalState.tryGet('telemetryClientId', String)
-            if (!clientId) {
-                clientId = randomUUID()
-                globalState.tryUpdate('telemetryClientId', clientId)
+            const globalClientId = process.env[telemetryClientIdEnvKey] // truly global across all extensions
+            const localClientId = globalState.tryGet('telemetryClientId', String) // local to extension, despite accessing "global" state
+            let clientId: string
+
+            if (isWeb()) {
+                const machineId = vscode.env.machineId
+                clientId = localClientId ?? machineId
+                getLogger().debug(
+                    'getClientId: web mode determined clientId: %s, stored clientId was: %s, vscode.machineId was: %s',
+                    clientId,
+                    localClientId,
+                    machineId
+                )
+            } else {
+                clientId = globalClientId ?? localClientId ?? randomUUID()
+                getLogger().debug(
+                    'getClientId: determined clientId as: %s, process.env clientId was: %s, stored clientId was: %s',
+                    clientId,
+                    globalClientId,
+                    localClientId
+                )
+                if (!globalClientId) {
+                    getLogger().debug(`getClientId: setting clientId in process.env to: %s`, clientId)
+                    process.env[telemetryClientIdEnvKey] = clientId
+                }
             }
+
+            globalState.tryUpdate('telemetryClientId', clientId)
             return clientId
         } catch (e) {
             getLogger().error('getClientId: failed to create client id: %O', e)
@@ -205,53 +239,6 @@ export function validateMetricEvent(event: MetricDatum, fatal: boolean) {
     }
 }
 
-/**
- * Setup the telemetry client id at extension activation.
- * This function is designed to let AWS Toolkit and Amazon Q share
- * the same telemetry client id.
- */
-export async function setupTelemetryId(extensionContext: vscode.ExtensionContext) {
-    try {
-        if (isWeb()) {
-            await globals.globalState.update('telemetryClientId', vscode.env.machineId)
-        } else {
-            const currentClientId = globals.globalState.tryGet('telemetryClientId', String)
-            const storedClientId = process.env[telemetryClientIdEnvKey]
-            if (currentClientId && storedClientId) {
-                if (extensionContext.extension.id === VSCODE_EXTENSION_ID.awstoolkit) {
-                    getLogger().debug(`telemetry: Store telemetry client id to env ${currentClientId}`)
-                    process.env[telemetryClientIdEnvKey] = currentClientId
-                    // notify amazon q to use this stored client id
-                    // if amazon q activates first. Do not block on activate amazon q
-                    if (isExtensionInstalled(VSCODE_EXTENSION_ID.amazonq)) {
-                        void activateExtension(VSCODE_EXTENSION_ID.amazonq).then(async () => {
-                            getLogger().debug(`telemetry: notifying Amazon Q to adopt client id ${currentClientId}`)
-                            await vscode.commands.executeCommand('aws.amazonq.setupTelemetryId')
-                        })
-                    }
-                } else if (isAmazonQ()) {
-                    getLogger().debug(`telemetry: Set telemetry client id to ${storedClientId}`)
-                    await globals.globalState.update('telemetryClientId', storedClientId)
-                } else {
-                    getLogger().error(`Unexpected extension id ${extensionContext.extension.id}`)
-                }
-            } else if (!currentClientId && storedClientId) {
-                getLogger().debug(`telemetry: Write telemetry client id to global state ${storedClientId}`)
-                await globals.globalState.update('telemetryClientId', storedClientId)
-            } else if (currentClientId && !storedClientId) {
-                getLogger().debug(`telemetry: Write telemetry client id to env ${currentClientId}`)
-                process.env[telemetryClientIdEnvKey] = currentClientId
-            } else {
-                const clientId = getClientId(globals.globalState)
-                getLogger().debug(`telemetry: Setup telemetry client id ${clientId}`)
-                process.env[telemetryClientIdEnvKey] = clientId
-            }
-        }
-    } catch (err) {
-        getLogger().error(`Error while setting up telemetry id ${err}`)
-    }
-}
-
 /**
  * Potentially helpful values for the 'source' field in telemetry.
  */

packages/core/src/test/shared/telemetry/util.test.ts

@@ -6,10 +6,18 @@
 import assert from 'assert'
 import { Memento, ConfigurationTarget } from 'vscode'
 import { Settings } from '../../../shared/settings'
-import { convertLegacy, getClientId, getUserAgent, platformPair, TelemetryConfig } from '../../../shared/telemetry/util'
+import {
+    convertLegacy,
+    getClientId,
+    getUserAgent,
+    platformPair,
+    telemetryClientIdEnvKey,
+    TelemetryConfig,
+} from '../../../shared/telemetry/util'
 import { extensionVersion } from '../../../shared/vscode/env'
 import { FakeMemento } from '../../fakeExtensionContext'
 import { GlobalState } from '../../../shared/globalState'
+import { randomUUID } from 'crypto'
 
 describe('TelemetryConfig', function () {
     const settingKey = 'aws.telemetry'
@@ -100,16 +108,48 @@ describe('TelemetryConfig', function () {
 })
 
 describe('getClientId', function () {
-    it('should generate a unique id', async function () {
-        const c1 = getClientId(new GlobalState(new FakeMemento()), true, false, 'x1')
-        const c2 = getClientId(new GlobalState(new FakeMemento()), true, false, 'x2')
+    before(function () {
+        delete process.env[telemetryClientIdEnvKey]
+    })
+
+    afterEach(function () {
+        delete process.env[telemetryClientIdEnvKey]
+    })
+
+    function testGetClientId(globalState: GlobalState) {
+        return getClientId(globalState, true, false, randomUUID())
+    }
+
+    it('generates a unique id if no other id is available', function () {
+        const c1 = testGetClientId(new GlobalState(new FakeMemento()))
+        delete process.env[telemetryClientIdEnvKey]
+        const c2 = testGetClientId(new GlobalState(new FakeMemento()))
         assert.notStrictEqual(c1, c2)
     })
 
+    it('uses id stored in global state if an id is not found in process.env', async function () {
+        const expectedClientId = 'myId'
+
+        const memento = new GlobalState(new FakeMemento())
+        await memento.update('telemetryClientId', expectedClientId)
+
+        assert.strictEqual(testGetClientId(memento), expectedClientId)
+    })
+
+    it('uses the id stored in process.env if available', async function () {
+        const expectedClientId = 'myId'
+
+        const e = new GlobalState(new FakeMemento())
+        await e.update('telemetryClientId', randomUUID())
+        process.env[telemetryClientIdEnvKey] = expectedClientId
+
+        assert.strictEqual(testGetClientId(new GlobalState(new FakeMemento())), expectedClientId)
+    })
+
     it('returns the same value across the calls sequentially', async function () {
         const memento = new GlobalState(new FakeMemento())
-        const c1 = getClientId(memento, true, false, 'y1')
-        const c2 = getClientId(memento, true, false, 'y2')
+        const c1 = testGetClientId(memento)
+        const c2 = testGetClientId(memento)
         assert.strictEqual(c1, c2)
     })
 
@@ -123,7 +163,7 @@ describe('getClientId', function () {
                 throw new Error()
             },
         }
-        const clientId = getClientId(new GlobalState(memento), true, false, 'x3')
+        const clientId = testGetClientId(new GlobalState(memento))
         // XXX: `notStrictEqual` since getClientId() is now synchronous. Because memento.update() is async.
         assert.notStrictEqual(clientId, '00000000-0000-0000-0000-000000000000')
     })
@@ -141,7 +181,7 @@ describe('getClientId', function () {
                 return this.memento.get(key)
             }
         }
-        const clientId = getClientId(new FakeGlobalState(memento), true, false, 'x4')
+        const clientId = testGetClientId(new FakeGlobalState(memento))
         assert.strictEqual(clientId, '00000000-0000-0000-0000-000000000000')
     })