fix(telemetry): toolkit replacing amazon q clientId

Problem: There is a function called `setupTelemetryId` which syncs the client Id between the two extensions when both are installed. However, there is a race condition that causes toolkit to generate the client ID before calling this setup function. The result is that amazon q is signalled to update its client id with a fresh one from toolkit. This results in telemetry showing we have a new user even though that is not the case.

Solution: Bypass this race condition by putting the logic in the client id generation function. We will also change how we compute the client ID so that extensions more intuitively  converge to the same client id.
See code for computation docs.

Cases:
1. New user installs an extension -> client id is generated once and stored in global state
2. Existing extension user installs another extension -> new extension will fetch the client id from the first extension via env variables, and store it itself
3. Existing users of both extensions update to the newer versions -> client ids were already synced between extensions, so they will store their identical client ids to themselves and env variables

There is an edge case that will still result in a failure:

4. Existing extension user installs another extension manually while vscode is closed -> the new extension loads first and stores its fresh client id in environment variables. The first extension will see this and update its client id to the new one, resulting in the original issue.

This is unavoidable because:
- Extensions can load in arbitrary orders
  - There is no way to know which extension's client ID is the one to use unless we check which one has something stored in its global state already.
  - This requires the new extension to ask the old extension, so it must be initialized already.
  - While waiting for this request, either extension could already start emitting telemetry. It has also already initialized a bunch of internal memory states with the other client id.
  - Creating a local file will only help if the file pre-exists. Users updating to this version will have no file to determine from.
- We want to preserve current client IDs

Author: hayemaxi

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')
     })