fix(cloudformation): handle telemetry setting in upgrade path case where setting is not registered

packages/core/src/awsService/cloudformation/extension.ts

@@ -52,7 +52,7 @@ import { AwsCredentialsService, encryptionKey } from './auth/credentials'
 import { ExtensionId, ExtensionName, Version, CloudFormationTelemetrySettings } from './extensionConfig'
 import { commandKey } from './utils'
 import { CloudFormationExplorer } from './explorer/explorer'
-import { promptTelemetryOptInWithTimeout } from './telemetryOptIn'
+import { handleTelemetryOptIn } from './telemetryOptIn'
 
 import { refreshCommand, StacksManager } from './stacks/stacksManager'
 import { StackOverviewWebviewProvider } from './ui/stackOverviewWebviewProvider'
@@ -89,7 +89,7 @@ let clientDisposables: Disposable[] = []
 
 async function startClient(context: ExtensionContext) {
     const cfnTelemetrySettings = new CloudFormationTelemetrySettings()
-    const telemetryEnabled = await promptTelemetryOptInWithTimeout(context, cfnTelemetrySettings)
+    const telemetryEnabled = await handleTelemetryOptIn(context, cfnTelemetrySettings)
 
     const cfnLspConfig = {
         ...DevSettings.instance.getServiceConfig('cloudformationLsp', {}),

packages/core/src/awsService/cloudformation/telemetryOptIn.ts

@@ -7,74 +7,162 @@ import { ExtensionContext, env, Uri, window } from 'vscode'
 import { CloudFormationTelemetrySettings } from './extensionConfig'
 import { commandKey } from './utils'
 import { isAutomation } from '../../shared/vscode/env'
+import { getLogger } from '../../shared/logger/logger'
 
-export async function promptTelemetryOptInWithTimeout(
-    context: ExtensionContext,
-    cfnTelemetrySettings: CloudFormationTelemetrySettings
-): Promise<boolean> {
-    const promptPromise = promptTelemetryOptIn(context, cfnTelemetrySettings)
-    const timeoutPromise = new Promise<false>((resolve) => setTimeout(() => resolve(false), 2500))
+enum TelemetryChoice {
+    Allow = 'Yes, Allow',
+    Later = 'Not Now',
+    Never = 'Never',
+    LearnMore = 'Learn More',
+}
 
-    const result = await Promise.race([promptPromise, timeoutPromise])
+const telemetryKeys = {
+    hasResponded: commandKey('telemetry.hasResponded'),
+    lastPromptDate: commandKey('telemetry.lastPromptDate'),
+    unpersistedResponse: commandKey('telemetry.unpersistedResponse'),
+} as const
 
-    // Keep prompt alive in background
-    void promptPromise
+const telemetrySettings = {
+    enabled: 'enabled',
+} as const
 
-    return result
-}
+const thirtyDaysMs = 30 * 24 * 60 * 60 * 1000
+const promptTimeoutMs = 2500
+const telemetryDocsUrl = 'https://github.com/aws-cloudformation/cloudformation-languageserver/tree/main/src/telemetry'
 
 /* eslint-disable aws-toolkits/no-banned-usages */
-async function promptTelemetryOptIn(
+export async function handleTelemetryOptIn(
     context: ExtensionContext,
     cfnTelemetrySettings: CloudFormationTelemetrySettings
 ): Promise<boolean> {
-    const telemetryEnabled = cfnTelemetrySettings.get('enabled', false)
-    if (isAutomation()) {
-        return telemetryEnabled
+    // If previous choice failed to persist, persist it now and return
+    const unpersistedResponse = (await context.globalState.get(telemetryKeys.unpersistedResponse)) as string
+    const hasResponded = context.globalState.get<boolean>(telemetryKeys.hasResponded)
+    const lastPromptDate = context.globalState.get<number>(telemetryKeys.lastPromptDate)
+    if (unpersistedResponse) {
+        // May still raise popup if user lacks permission or file is corrupted
+        const didSave = await saveTelemetryResponse(unpersistedResponse, cfnTelemetrySettings)
+        await context.globalState.update(telemetryKeys.unpersistedResponse, undefined)
+        // If we still couldn't save, clear everything so they get asked again until the file/perms is fixed
+        if (!didSave) {
+            getLogger().warn(
+                'CloudFormation telemetry choice was not saved successfully after restart. Clearing related globalState keys for next restart'
+            )
+            await context.globalState.update(telemetryKeys.hasResponded, undefined)
+            await context.globalState.update(telemetryKeys.lastPromptDate, undefined)
+        }
+        return logAndReturnTelemetryChoice(
+            unpersistedResponse === TelemetryChoice.Allow.toString(),
+            hasResponded,
+            lastPromptDate
+        )
     }
 
-    const hasResponded = context.globalState.get<boolean>(commandKey('telemetry.hasResponded'), false)
-    const lastPromptDate = context.globalState.get<number>(commandKey('telemetry.lastPromptDate'), 0)
-    const now = Date.now()
-    const thirtyDaysMs = 30 * 24 * 60 * 60 * 1000
+    // Never throws because we provide a default
+    const telemetryEnabled = cfnTelemetrySettings.get(telemetrySettings.enabled, false)
+
+    if (isAutomation()) {
+        return logAndReturnTelemetryChoice(telemetryEnabled)
+    }
 
     // If user has permanently responded, use their choice
     if (hasResponded) {
-        return telemetryEnabled
+        return logAndReturnTelemetryChoice(telemetryEnabled, hasResponded)
     }
 
     // Check if we should show reminder (30 days since last prompt)
-    const shouldPrompt = lastPromptDate === 0 || now - lastPromptDate >= thirtyDaysMs
+    const shouldPrompt = lastPromptDate === undefined || Date.now() - lastPromptDate >= thirtyDaysMs
     if (!shouldPrompt) {
-        return telemetryEnabled
+        return logAndReturnTelemetryChoice(telemetryEnabled, hasResponded, lastPromptDate)
+    }
+
+    // Show prompt but set false if timeout
+    const promptPromise = promptTelemetryOptIn(context, cfnTelemetrySettings)
+    const timeoutPromise = new Promise<false>((resolve) => setTimeout(() => resolve(false), promptTimeoutMs))
+    const result = await Promise.race([promptPromise, timeoutPromise])
+
+    // Keep prompt alive in background
+    void promptPromise
+
+    return logAndReturnTelemetryChoice(result)
+}
+/**
+ * Updates the telemetry setting. In case of error, the update calls do not throw.
+ * They instead raise a popup and return false.
+ *
+ * @returns boolean whether the save/update was successful
+ */
+/* eslint-disable aws-toolkits/no-banned-usages */
+async function saveTelemetryResponse(
+    response: string | undefined,
+    cfnTelemetrySettings: CloudFormationTelemetrySettings
+): Promise<boolean> {
+    if (response === TelemetryChoice.Allow) {
+        return await cfnTelemetrySettings.update(telemetrySettings.enabled, true)
+    } else if (response === TelemetryChoice.Never) {
+        return await cfnTelemetrySettings.update(telemetrySettings.enabled, false)
+    } else if (response === TelemetryChoice.Later) {
+        return await cfnTelemetrySettings.update(telemetrySettings.enabled, false)
     }
+    return false
+}
 
+function logAndReturnTelemetryChoice(choice: boolean, hasResponded?: boolean, lastPromptDate?: number): boolean {
+    getLogger().info(
+        'CloudFormation telemetry: choice=%s, hasResponded=%s, lastPromptDate=%s',
+        choice,
+        hasResponded,
+        lastPromptDate
+    )
+    return choice
+}
+
+/* eslint-disable aws-toolkits/no-banned-usages */
+async function promptTelemetryOptIn(
+    context: ExtensionContext,
+    cfnTelemetrySettings: CloudFormationTelemetrySettings
+): Promise<boolean> {
     const message =
         'Help us improve the AWS CloudFormation Language Server by sharing anonymous telemetry data with AWS. You can change this preference at any time in aws.cloudformation Settings.'
 
-    const allow = 'Yes, Allow'
-    const later = 'Not Now'
-    const never = 'Never'
-    const learnMore = 'Learn More'
-    const response = await window.showInformationMessage(message, allow, later, never, learnMore)
+    const response = await window.showInformationMessage(
+        message,
+        TelemetryChoice.Allow,
+        TelemetryChoice.Later,
+        TelemetryChoice.Never,
+        TelemetryChoice.LearnMore
+    )
 
-    if (response === learnMore) {
-        await env.openExternal(
-            Uri.parse('https://github.com/aws-cloudformation/cloudformation-languageserver/tree/main/src/telemetry')
-        )
+    if (response === TelemetryChoice.LearnMore) {
+        await env.openExternal(Uri.parse(telemetryDocsUrl))
         return promptTelemetryOptIn(context, cfnTelemetrySettings)
     }
 
-    if (response === allow) {
-        await cfnTelemetrySettings.update('enabled', true)
-        await context.globalState.update(commandKey('telemetry.hasResponded'), true)
-    } else if (response === never) {
-        await cfnTelemetrySettings.update('enabled', false)
-        await context.globalState.update(commandKey('telemetry.hasResponded'), true)
-    } else if (response === later) {
-        await cfnTelemetrySettings.update('enabled', false)
-        await context.globalState.update(commandKey('telemetry.lastPromptDate'), now)
+    const now = Date.now()
+    await context.globalState.update(telemetryKeys.lastPromptDate, now)
+
+    // There's a chance our settings aren't registered yet from package.json, so we
+    // see if we can persist to settings first
+    try {
+        // Throws (with no popup) if setting is not registered
+        cfnTelemetrySettings.get(telemetrySettings.enabled)
+    } catch (err) {
+        getLogger().warn(err as Error)
+        // Save the choice in globalState and save to settings next time handleTelemetryOptIn is called
+        await context.globalState.update(telemetryKeys.unpersistedResponse, response)
+        if (response === TelemetryChoice.Allow) {
+            await context.globalState.update(telemetryKeys.hasResponded, true)
+            return true
+        } else if (response === TelemetryChoice.Never) {
+            await context.globalState.update(telemetryKeys.hasResponded, true)
+            return false
+        } else if (response === TelemetryChoice.Later) {
+            return false
+        }
     }
 
-    return cfnTelemetrySettings.get('enabled', false)
+    // At this point should be able to save and get successfully
+    await saveTelemetryResponse(response, cfnTelemetrySettings)
+    await context.globalState.update(telemetryKeys.hasResponded, response !== TelemetryChoice.Later)
+    return cfnTelemetrySettings.get(telemetrySettings.enabled, false)
 }

packages/core/src/test/awsService/cloudformation/telemetryOptIn.test.ts

@@ -0,0 +1,120 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import * as sinon from 'sinon'
+import { ExtensionContext } from 'vscode'
+import { handleTelemetryOptIn } from '../../../awsService/cloudformation/telemetryOptIn'
+import { CloudFormationTelemetrySettings } from '../../../awsService/cloudformation/extensionConfig'
+import { commandKey } from '../../../awsService/cloudformation/utils'
+
+describe('telemetryOptIn', function () {
+    let mockContext: ExtensionContext
+    let mockSettings: CloudFormationTelemetrySettings
+    let globalState: Map<string, any>
+
+    beforeEach(function () {
+        globalState = new Map()
+
+        mockContext = {
+            globalState: {
+                get: (key: string, defaultValue?: any) => globalState.get(key) ?? defaultValue,
+                update: async (key: string, value: any) => {
+                    globalState.set(key, value)
+                },
+            },
+        } as any
+
+        mockSettings = {
+            get: sinon.stub().returns(false),
+            update: sinon.stub().resolves(),
+        } as any
+    })
+
+    describe('promptTelemetryOptIn - automation mode', function () {
+        it('should return current setting without prompting in automation mode', async function () {
+            sinon.stub(require('../../../shared/vscode/env'), 'isAutomation').returns(true)
+            ;(mockSettings.get as sinon.SinonStub).returns(true)
+
+            const result = await handleTelemetryOptIn(mockContext, mockSettings)
+
+            assert.strictEqual(result, true)
+        })
+    })
+
+    describe('promptTelemetryOptIn - user has responded', function () {
+        it('should return current setting if user has permanently responded', async function () {
+            globalState.set(commandKey('telemetry.hasResponded'), true)
+            ;(mockSettings.get as sinon.SinonStub).returns(true)
+
+            const result = await handleTelemetryOptIn(mockContext, mockSettings)
+
+            assert.strictEqual(result, true)
+        })
+    })
+
+    describe('promptTelemetryOptIn - prompt timing', function () {
+        it('should not prompt if less than 30 days since last prompt', async function () {
+            const now = Date.now()
+            const twentyDaysAgo = now - 20 * 24 * 60 * 60 * 1000
+            globalState.set(commandKey('telemetry.lastPromptDate'), twentyDaysAgo)
+
+            const result = await handleTelemetryOptIn(mockContext, mockSettings)
+
+            assert.strictEqual(result, false)
+        })
+    })
+
+    describe('promptTelemetryOptIn - unpersisted response', function () {
+        it('should persist unpersisted Allow response', async function () {
+            globalState.set(commandKey('telemetry.unpersistedResponse'), 'Yes, Allow')
+
+            const result = await handleTelemetryOptIn(mockContext, mockSettings)
+
+            assert.strictEqual(result, true)
+            assert.ok((mockSettings.update as sinon.SinonStub).calledWith('enabled', true))
+            assert.strictEqual(globalState.get(commandKey('telemetry.unpersistedResponse')), undefined)
+        })
+
+        it('should persist unpersisted Never response', async function () {
+            globalState.set(commandKey('telemetry.unpersistedResponse'), 'Never')
+            ;(mockSettings.update as sinon.SinonStub).resolves(true)
+
+            const result = await handleTelemetryOptIn(mockContext, mockSettings)
+
+            assert.strictEqual(result, false)
+            assert.ok((mockSettings.update as sinon.SinonStub).calledWith('enabled', false))
+            assert.strictEqual(globalState.get(commandKey('telemetry.unpersistedResponse')), undefined)
+        })
+
+        it('should persist unpersisted Later response', async function () {
+            const lastPromptDate = Date.now() - 1000
+            globalState.set(commandKey('telemetry.unpersistedResponse'), 'Not Now')
+            globalState.set(commandKey('telemetry.lastPromptDate'), lastPromptDate)
+            ;(mockSettings.update as sinon.SinonStub).resolves(true)
+
+            const result = await handleTelemetryOptIn(mockContext, mockSettings)
+
+            assert.strictEqual(result, false)
+            assert.ok((mockSettings.update as sinon.SinonStub).calledWith('enabled', false))
+            assert.strictEqual(globalState.get(commandKey('telemetry.lastPromptDate')), lastPromptDate)
+            assert.strictEqual(globalState.get(commandKey('telemetry.unpersistedResponse')), undefined)
+        })
+
+        it('should clear all state if setting save fails', async function () {
+            globalState.set(commandKey('telemetry.unpersistedResponse'), 'Yes, Allow')
+            globalState.set(commandKey('telemetry.hasResponded'), true)
+            globalState.set(commandKey('telemetry.lastPromptDate'), Date.now())
+            ;(mockSettings.update as sinon.SinonStub).resolves(false)
+
+            const result = await handleTelemetryOptIn(mockContext, mockSettings)
+
+            assert.strictEqual(result, true)
+            assert.strictEqual(globalState.get(commandKey('telemetry.unpersistedResponse')), undefined)
+            assert.strictEqual(globalState.get(commandKey('telemetry.hasResponded')), undefined)
+            assert.strictEqual(globalState.get(commandKey('telemetry.lastPromptDate')), undefined)
+        })
+    })
+})

packages/toolkit/.changes/next-release/Bug Fix-6113ced5-5e07-4662-a641-db4588ddf8de.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "CloudFormation: Handle telemetry setting in upgrade path case where setting is not registered"
+}