From 74b3c475fb32eb3ca7b3628d9afd2cc34b1bfc8d Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Wed, 13 Nov 2024 14:10:36 -0500
Subject: [PATCH 1/9] refactor(notifications): remove unneeded
 InstalledExtensions criteria

We have both ActiveExtensions (checks what extensions are installed and active), and InstalledExtensions (just what extensions are installed).
This removes InstalledExtensions because most of the intention behind it is captured by ActiveExtensions. When would we only care about whats installed but not active?

Removing a criteria decreases complexity and confusion in both code and for those drafting notifications, so this is an easy removal.
---
 packages/core/src/notifications/rules.ts      |  6 ++---
 packages/core/src/notifications/types.ts      | 11 +-------
 .../src/test/notifications/controller.test.ts |  1 -
 .../core/src/test/notifications/rules.test.ts | 26 -------------------
 4 files changed, 3 insertions(+), 41 deletions(-)

diff --git a/packages/core/src/notifications/rules.ts b/packages/core/src/notifications/rules.ts
index 66b3a10d2b5..efe28808237 100644
--- a/packages/core/src/notifications/rules.ts
+++ b/packages/core/src/notifications/rules.ts
@@ -134,8 +134,6 @@ export class RuleEngine {
                 return hasAnyOfExpected(this.context.authStates)
             case 'AuthScopes':
                 return isEqualSetToExpected(this.context.authScopes)
-            case 'InstalledExtensions':
-                return isSuperSetOfExpected(this.context.installedExtensions)
             case 'ActiveExtensions':
                 return isSuperSetOfExpected(this.context.activeExtensions)
             default:
@@ -169,7 +167,6 @@ export async function getRuleContext(context: vscode.ExtensionContext, authState
         computeEnv: await getComputeEnvType(),
         authTypes: [...new Set(authTypes)],
         authScopes: authState.authScopes ? authState.authScopes?.split(',') : [],
-        installedExtensions: vscode.extensions.all.map((e) => e.id),
         activeExtensions: vscode.extensions.all.filter((e) => e.isActive).map((e) => e.id),
 
         // Toolkit (and eventually Q?) may have multiple connections with different regions and states.
@@ -177,7 +174,8 @@ export async function getRuleContext(context: vscode.ExtensionContext, authState
         authRegions: authState.awsRegion ? [authState.awsRegion] : [],
         authStates: [authState.authStatus],
     }
-    const { activeExtensions, installedExtensions, ...loggableRuleContext } = ruleContext
+
+    const { activeExtensions, ...loggableRuleContext } = ruleContext
     getLogger('notifications').debug('getRuleContext() determined rule context: %O', loggableRuleContext)
 
     return ruleContext
diff --git a/packages/core/src/notifications/types.ts b/packages/core/src/notifications/types.ts
index 596a5f2b6bd..c99eb849759 100644
--- a/packages/core/src/notifications/types.ts
+++ b/packages/core/src/notifications/types.ts
@@ -9,15 +9,7 @@ import { TypeConstructor } from '../shared/utilities/typeConstructors'
 import { AuthUserState, AuthStatus } from '../shared/telemetry/telemetry.gen'
 
 /** Types of information that we can use to determine whether to show a notification or not. */
-export type Criteria =
-    | 'OS'
-    | 'ComputeEnv'
-    | 'AuthType'
-    | 'AuthRegion'
-    | 'AuthState'
-    | 'AuthScopes'
-    | 'InstalledExtensions'
-    | 'ActiveExtensions'
+export type Criteria = 'OS' | 'ComputeEnv' | 'AuthType' | 'AuthRegion' | 'AuthState' | 'AuthScopes' | 'ActiveExtensions'
 
 /** Generic condition where the type determines how the values are evaluated. */
 export interface CriteriaCondition {
@@ -132,7 +124,6 @@ export interface RuleContext {
     readonly authRegions: string[]
     readonly authStates: AuthStatus[]
     readonly authScopes: string[]
-    readonly installedExtensions: string[]
     readonly activeExtensions: string[]
 }
 
diff --git a/packages/core/src/test/notifications/controller.test.ts b/packages/core/src/test/notifications/controller.test.ts
index e1ac14094b8..ea902fdbfe7 100644
--- a/packages/core/src/test/notifications/controller.test.ts
+++ b/packages/core/src/test/notifications/controller.test.ts
@@ -39,7 +39,6 @@ describe('Notifications Controller', function () {
         authRegions: ['us-east-1'],
         authStates: ['connected'],
         authScopes: ['codewhisperer:completions', 'codewhisperer:analysis'],
-        installedExtensions: ['ext1', 'ext2', 'ext3'],
         activeExtensions: ['ext1', 'ext2'],
     })
 
diff --git a/packages/core/src/test/notifications/rules.test.ts b/packages/core/src/test/notifications/rules.test.ts
index 6b9c7454869..867b2119546 100644
--- a/packages/core/src/test/notifications/rules.test.ts
+++ b/packages/core/src/test/notifications/rules.test.ts
@@ -24,7 +24,6 @@ describe('Notifications Rule Engine', function () {
         authRegions: ['us-east-1'],
         authStates: ['connected'],
         authScopes: ['codewhisperer:completions', 'codewhisperer:analysis'],
-        installedExtensions: ['ext1', 'ext2', 'ext3'],
         activeExtensions: ['ext1', 'ext2'],
     }
 
@@ -405,28 +404,6 @@ describe('Notifications Rule Engine', function () {
         )
     })
 
-    it('should display notification for InstalledExtensions criteria', function () {
-        assert.equal(
-            ruleEngine.shouldDisplayNotification(
-                buildNotification({
-                    additionalCriteria: [{ type: 'InstalledExtensions', values: ['ext1', 'ext2'] }],
-                })
-            ),
-            true
-        )
-    })
-
-    it('should NOT display notification for invalid InstalledExtensions criteria', function () {
-        assert.equal(
-            ruleEngine.shouldDisplayNotification(
-                buildNotification({
-                    additionalCriteria: [{ type: 'InstalledExtensions', values: ['ext1', 'ext2', 'unknownExtension'] }],
-                })
-            ),
-            false
-        )
-    })
-
     it('should display notification for ActiveExtensions criteria', function () {
         assert.equal(
             ruleEngine.shouldDisplayNotification(
@@ -479,7 +456,6 @@ describe('Notifications Rule Engine', function () {
                         { type: 'AuthRegion', values: ['us-east-1', 'us-west-2'] },
                         { type: 'AuthState', values: ['connected'] },
                         { type: 'AuthScopes', values: ['codewhisperer:completions', 'codewhisperer:analysis'] },
-                        { type: 'InstalledExtensions', values: ['ext1', 'ext2'] },
                         { type: 'ActiveExtensions', values: ['ext1', 'ext2'] },
                     ],
                 })
@@ -517,7 +493,6 @@ describe('Notifications Rule Engine', function () {
                         { type: 'AuthRegion', values: ['us-east-1', 'us-west-2'] },
                         { type: 'AuthState', values: ['connected'] },
                         { type: 'AuthScopes', values: ['codewhisperer:completions', 'codewhisperer:analysis'] },
-                        { type: 'InstalledExtensions', values: ['ex1', 'ext2'] },
                         { type: 'ActiveExtensions', values: ['ext1', 'ext2'] },
 
                         { type: 'ComputeEnv', values: ['ec2'] }, // no 'local'
@@ -576,7 +551,6 @@ describe('Notifications getRuleContext()', function () {
             authRegions: ['us-east-1'],
             authStates: ['connected'],
             authScopes: amazonQScopes,
-            installedExtensions: vscode.extensions.all.map((e) => e.id),
             activeExtensions: vscode.extensions.all.filter((e) => e.isActive).map((e) => e.id),
         })
     })

From 9a81077ca5b89d091e2aca3b678d30f00c53e2c7 Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Wed, 13 Nov 2024 14:49:09 -0500
Subject: [PATCH 2/9] types(notifications): add types to onClick, onReceive,
 actions

---
 packages/core/src/notifications/panelNode.ts           |  8 ++++++--
 packages/core/src/notifications/types.ts               | 10 +++++++---
 .../core/src/test/notifications/controller.test.ts     |  6 ++++--
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/packages/core/src/notifications/panelNode.ts b/packages/core/src/notifications/panelNode.ts
index 87ef872c59a..c162ab67c0a 100644
--- a/packages/core/src/notifications/panelNode.ts
+++ b/packages/core/src/notifications/panelNode.ts
@@ -8,7 +8,7 @@ import { ResourceTreeDataProvider, TreeNode } from '../shared/treeview/resourceT
 import { Command, Commands } from '../shared/vscode/commands2'
 import { Icon, IconPath, getIcon } from '../shared/icons'
 import { contextKey, setContext } from '../shared/vscode/setContext'
-import { NotificationType, ToolkitNotification, getNotificationTelemetryId } from './types'
+import { NotificationType, OnReceiveType, ToolkitNotification, getNotificationTelemetryId } from './types'
 import { ToolkitError } from '../shared/errors'
 import { isAmazonQ } from '../shared/extensionUtilities'
 import { getLogger } from '../shared/logger/logger'
@@ -165,7 +165,11 @@ export class NotificationsNode implements TreeNode {
      * Can be either a blocking modal or a bottom-right corner toast
      * Handles the button click actions based on the button type.
      */
-    private showInformationWindow(notification: ToolkitNotification, type: string = 'toast', passive: boolean = false) {
+    private showInformationWindow(
+        notification: ToolkitNotification,
+        type: OnReceiveType = 'toast',
+        passive: boolean = false
+    ) {
         const isModal = type === 'modal'
 
         // modal has to have defined actions (buttons)
diff --git a/packages/core/src/notifications/types.ts b/packages/core/src/notifications/types.ts
index c99eb849759..59717cf9079 100644
--- a/packages/core/src/notifications/types.ts
+++ b/packages/core/src/notifications/types.ts
@@ -38,6 +38,10 @@ export interface ExactMatch {
 
 export type ConditionalClause = Range | ExactMatch | OR
 
+export type OnReceiveType = 'toast' | 'modal'
+export type OnClickType = 'modal' | 'openTextDocument' | 'openUrl'
+export type ActionType = 'openUrl' | 'updateAndReload' | 'openTxt'
+
 /** How to display the notification. */
 export interface UIRenderInstructions {
     content: {
@@ -47,13 +51,13 @@ export interface UIRenderInstructions {
             toastPreview?: string // optional property for toast
         }
     }
-    onRecieve: string
+    onRecieve: OnReceiveType // TODO: typo
     onClick: {
-        type: string
+        type: OnClickType
         url?: string // optional property for 'openUrl'
     }
     actions?: Array<{
-        type: string
+        type: ActionType
         displayText: {
             [`en-US`]: string
         }
diff --git a/packages/core/src/test/notifications/controller.test.ts b/packages/core/src/test/notifications/controller.test.ts
index ea902fdbfe7..669700a9cde 100644
--- a/packages/core/src/test/notifications/controller.test.ts
+++ b/packages/core/src/test/notifications/controller.test.ts
@@ -534,7 +534,7 @@ describe('RemoteFetcher', function () {
     })
 })
 
-function getValidTestNotification(id: string) {
+function getValidTestNotification(id: string): ToolkitNotification {
     return {
         id,
         displayIf: {
@@ -556,7 +556,7 @@ function getValidTestNotification(id: string) {
     }
 }
 
-function getInvalidTestNotification(id: string) {
+function getInvalidTestNotification(id: string): ToolkitNotification {
     return {
         id,
         displayIf: {
@@ -570,6 +570,8 @@ function getInvalidTestNotification(id: string) {
                     description: 'test',
                 },
             },
+            onRecieve: 'toast',
+            onClick: { type: 'modal' },
         },
     }
 }

From fd00403e6085d30f4afdb20b7ba76b45324bbf50 Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Wed, 13 Nov 2024 14:49:50 -0500
Subject: [PATCH 3/9] fix(notifications): add timeout to update and reload

Telemetry needs time to be sent, maybe other cleanup routines as well.
Also, make sure to flush telemetry before reloading so it isn't lost.
---
 packages/core/src/notifications/panelNode.ts  | 20 ++++++++++----
 packages/core/src/notifications/rules.ts      | 27 ++++++++++++++++---
 .../src/shared/telemetry/telemetryService.ts  |  6 +++--
 .../src/test/notifications/rendering.test.ts  | 16 ++++++++++-
 4 files changed, 58 insertions(+), 11 deletions(-)

diff --git a/packages/core/src/notifications/panelNode.ts b/packages/core/src/notifications/panelNode.ts
index c162ab67c0a..ddb24fa8132 100644
--- a/packages/core/src/notifications/panelNode.ts
+++ b/packages/core/src/notifications/panelNode.ts
@@ -16,6 +16,9 @@ import { registerToolView } from '../awsexplorer/activationShared'
 import { readonlyDocument } from '../shared/utilities/textDocumentUtilities'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
 import { telemetry } from '../shared/telemetry/telemetry'
+import { globals } from '../shared'
+
+const logger = getLogger('notifications')
 
 /**
  * Controls the "Notifications" side panel/tree in each extension. It takes purely UX actions
@@ -132,7 +135,7 @@ export class NotificationsNode implements TreeNode {
         switch (notification.uiRenderInstructions.onClick.type) {
             case 'modal':
                 // Render blocking modal
-                getLogger('notifications').verbose(`rendering modal for notificaiton: ${notification.id} ...`)
+                logger.verbose(`rendering modal for notificaiton: ${notification.id} ...`)
                 await this.showInformationWindow(notification, 'modal', false)
                 break
             case 'openUrl':
@@ -140,7 +143,7 @@ export class NotificationsNode implements TreeNode {
                 if (!notification.uiRenderInstructions.onClick.url) {
                     throw new ToolkitError('No url provided for onclick open url')
                 }
-                getLogger('notifications').verbose(`opening url for notification: ${notification.id} ...`)
+                logger.verbose(`opening url for notification: ${notification.id} ...`)
                 await openUrl(
                     vscode.Uri.parse(notification.uiRenderInstructions.onClick.url),
                     getNotificationTelemetryId(notification)
@@ -148,7 +151,7 @@ export class NotificationsNode implements TreeNode {
                 break
             case 'openTextDocument':
                 // Display read-only txt document
-                getLogger('notifications').verbose(`showing txt document for notification: ${notification.id} ...`)
+                logger.verbose(`showing txt document for notification: ${notification.id} ...`)
                 await telemetry.toolkit_invokeAction.run(async () => {
                     telemetry.record({ source: getNotificationTelemetryId(notification), action: 'openTxt' })
                     await readonlyDocument.show(
@@ -207,7 +210,10 @@ export class NotificationsNode implements TreeNode {
                                     )
                                     break
                                 case 'updateAndReload':
-                                    await this.updateAndReload(notification.displayIf.extensionId)
+                                    // Give things time to finish executing.
+                                    globals.clock.setTimeout(() => {
+                                        void this.updateAndReload(notification.displayIf.extensionId)
+                                    }, 1000)
                                     break
                                 case 'openUrl':
                                     if (selectedButton.url) {
@@ -232,7 +238,11 @@ export class NotificationsNode implements TreeNode {
     }
 
     private async updateAndReload(id: string) {
-        getLogger('notifications').verbose('Updating and reloading the extension...')
+        logger.verbose('Updating and reloading the extension...')
+
+        // Publish pending telemetry before it is lost to the window reload.
+        await globals.telemetry.flushRecords()
+
         await vscode.commands.executeCommand('workbench.extensions.installExtension', id)
         await vscode.commands.executeCommand('workbench.action.reloadWindow')
     }
diff --git a/packages/core/src/notifications/rules.ts b/packages/core/src/notifications/rules.ts
index efe28808237..9e62adb7cc1 100644
--- a/packages/core/src/notifications/rules.ts
+++ b/packages/core/src/notifications/rules.ts
@@ -11,6 +11,7 @@ import { getComputeEnvType, getOperatingSystem } from '../shared/telemetry/util'
 import { AuthFormId } from '../login/webview/vue/types'
 import { getLogger } from '../shared/logger/logger'
 
+const logger = getLogger('notifications')
 /**
  * Evaluates if a given version fits into the parameters specified by a notification, e.g:
  *
@@ -72,22 +73,40 @@ export class RuleEngine {
     constructor(private readonly context: RuleContext) {}
 
     public shouldDisplayNotification(payload: ToolkitNotification) {
-        return this.evaluate(payload.displayIf)
+        return this.evaluate(payload.id, payload.displayIf)
     }
 
-    private evaluate(condition: DisplayIf): boolean {
+    private evaluate(id: string, condition: DisplayIf): boolean {
         const currentExt = globals.context.extension.id
         if (condition.extensionId !== currentExt) {
+            logger.verbose(
+                'notification id: (%s) did NOT pass extension id check, actual ext id: (%s), expected ext id: (%s)',
+                id,
+                currentExt,
+                condition.extensionId
+            )
             return false
         }
 
         if (condition.ideVersion) {
             if (!isValidVersion(this.context.ideVersion, condition.ideVersion)) {
+                logger.verbose(
+                    'notification id: (%s) did NOT pass IDE version check, actual version: (%s), expected version: (%s)',
+                    id,
+                    this.context.ideVersion,
+                    condition.ideVersion
+                )
                 return false
             }
         }
         if (condition.extensionVersion) {
             if (!isValidVersion(this.context.extensionVersion, condition.extensionVersion)) {
+                logger.verbose(
+                    'notification id: (%s) did NOT pass extension version check, actual ext version: (%s), expected ext version: (%s)',
+                    id,
+                    this.context.extensionVersion,
+                    condition.extensionVersion
+                )
                 return false
             }
         }
@@ -95,8 +114,10 @@ export class RuleEngine {
         if (condition.additionalCriteria) {
             for (const criteria of condition.additionalCriteria) {
                 if (!this.evaluateRule(criteria)) {
+                    logger.verbose('notification id: (%s) did NOT pass criteria check: %O', id, criteria)
                     return false
                 }
+                logger.debug('notification id: (%s) passed criteria check: %O', id, criteria)
             }
         }
 
@@ -176,7 +197,7 @@ export async function getRuleContext(context: vscode.ExtensionContext, authState
     }
 
     const { activeExtensions, ...loggableRuleContext } = ruleContext
-    getLogger('notifications').debug('getRuleContext() determined rule context: %O', loggableRuleContext)
+    logger.debug('getRuleContext() determined rule context: %O', loggableRuleContext)
 
     return ruleContext
 }
diff --git a/packages/core/src/shared/telemetry/telemetryService.ts b/packages/core/src/shared/telemetry/telemetryService.ts
index c56ba3d2a24..61e9559db7c 100644
--- a/packages/core/src/shared/telemetry/telemetryService.ts
+++ b/packages/core/src/shared/telemetry/telemetryService.ts
@@ -195,9 +195,11 @@ export class DefaultTelemetryService {
     }
 
     /**
-     * Publish metrics to the Telemetry Service.
+     * Publish metrics to the Telemetry Service. Usually it will automatically flush recent events
+     * on a regular interval. This should not be used unless you are interrupting this interval,
+     * e.g. via a forced window reload.
      */
-    private async flushRecords(): Promise<void> {
+    public async flushRecords(): Promise<void> {
         if (this.telemetryEnabled) {
             await this._flushRecords()
         }
diff --git a/packages/core/src/test/notifications/rendering.test.ts b/packages/core/src/test/notifications/rendering.test.ts
index ada48c2c64b..4df1caddc56 100644
--- a/packages/core/src/test/notifications/rendering.test.ts
+++ b/packages/core/src/test/notifications/rendering.test.ts
@@ -5,24 +5,35 @@
 
 import * as vscode from 'vscode'
 import * as sinon from 'sinon'
+import * as FakeTimers from '@sinonjs/fake-timers'
 import assert from 'assert'
 import { ToolkitNotification } from '../../notifications/types'
 import { panelNode } from './controller.test'
 import { getTestWindow } from '../shared/vscode/window'
 import * as VsCodeUtils from '../../shared/utilities/vsCodeUtils'
-import { assertTextEditorContains } from '../testUtil'
+import { assertTextEditorContains, installFakeClock } from '../testUtil'
 
 describe('Notifications Rendering', function () {
     let sandbox: sinon.SinonSandbox
+    let clock: FakeTimers.InstalledClock
+
+    before(function () {
+        clock = installFakeClock()
+    })
 
     beforeEach(function () {
         sandbox = sinon.createSandbox()
     })
 
     afterEach(function () {
+        clock.reset()
         sandbox.restore()
     })
 
+    after(function () {
+        clock.uninstall()
+    })
+
     // util to test txt pop-up under different senarios
     async function verifyTxtNotification(notification: ToolkitNotification) {
         const expectedContent = notification.uiRenderInstructions.content['en-US'].description
@@ -95,6 +106,9 @@ describe('Notifications Rendering', function () {
         const notification = getModalNotification()
         await panelNode.openNotification(notification)
 
+        // Update and Reload is put on a timer so that other methods (e.g. telemetry) can finish.
+        await clock.tickAsync(2000)
+
         assert.ok(excuteCommandStub.calledWith('workbench.extensions.installExtension', 'aws.toolkit.fake.extension'))
         assert.ok(excuteCommandStub.calledWith('workbench.action.reloadWindow'))
     })

From 0418b0022d72f910e5758c83f320cb9f5d019f03 Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Wed, 13 Nov 2024 17:25:06 -0500
Subject: [PATCH 4/9] fix(notifications): icons causing runtime error

Problem: Making the icon red for emergency notifications requires changing the color property. To do it in 1 line we use a ternary statement and unwrap. This causes runtime error similar to: `cannot access property of undefined: 0`.

Solution: Use object.assign instead
---
 packages/core/src/notifications/panelNode.ts | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/packages/core/src/notifications/panelNode.ts b/packages/core/src/notifications/panelNode.ts
index ddb24fa8132..3644fd035c1 100644
--- a/packages/core/src/notifications/panelNode.ts
+++ b/packages/core/src/notifications/panelNode.ts
@@ -6,7 +6,7 @@
 import * as vscode from 'vscode'
 import { ResourceTreeDataProvider, TreeNode } from '../shared/treeview/resourceTreeDataProvider'
 import { Command, Commands } from '../shared/vscode/commands2'
-import { Icon, IconPath, getIcon } from '../shared/icons'
+import { Icon, getIcon } from '../shared/icons'
 import { contextKey, setContext } from '../shared/vscode/setContext'
 import { NotificationType, OnReceiveType, ToolkitNotification, getNotificationTelemetryId } from './types'
 import { ToolkitError } from '../shared/errors'
@@ -81,10 +81,13 @@ export class NotificationsNode implements TreeNode {
 
     public getChildren() {
         const buildNode = (n: ToolkitNotification, type: NotificationType) => {
-            const icon: Icon | IconPath =
-                type === 'startUp'
-                    ? getIcon('vscode-question')
-                    : { ...getIcon('vscode-alert'), color: new vscode.ThemeColor('errorForeground') }
+            const icon: Icon =
+                type === 'emergency'
+                    ? Object.assign(getIcon('vscode-alert') as Icon, {
+                          color: new vscode.ThemeColor('errorForeground'),
+                      })
+                    : (getIcon('vscode-question') as Icon)
+
             return this.openNotificationCmd.build(n).asTreeNode({
                 label: n.uiRenderInstructions.content['en-US'].title,
                 iconPath: icon,

From 860b44190046537326e794640a87a6ad1c2401ee Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Thu, 14 Nov 2024 17:05:54 -0500
Subject: [PATCH 5/9] feat(notifications): adjust wording and add activity bar
 badge

- Remove description from the notifications panel - unneeded. It is obvious the panel is for notifications.
- Make hover text the full title of the notification in case user side bar width isn't enough to display the full title.
- Add a badge to the Q activity bar icon indicating how many notifications you have. Also adds this as (#) to the panel title.
---
 packages/core/src/notifications/panelNode.ts | 38 ++++++++++++++++----
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/packages/core/src/notifications/panelNode.ts b/packages/core/src/notifications/panelNode.ts
index 3644fd035c1..cfcc0ad5eac 100644
--- a/packages/core/src/notifications/panelNode.ts
+++ b/packages/core/src/notifications/panelNode.ts
@@ -4,6 +4,7 @@
  */
 
 import * as vscode from 'vscode'
+import * as nls from 'vscode-nls'
 import { ResourceTreeDataProvider, TreeNode } from '../shared/treeview/resourceTreeDataProvider'
 import { Command, Commands } from '../shared/vscode/commands2'
 import { Icon, getIcon } from '../shared/icons'
@@ -18,6 +19,7 @@ import { openUrl } from '../shared/utilities/vsCodeUtils'
 import { telemetry } from '../shared/telemetry/telemetry'
 import { globals } from '../shared'
 
+const localize = nls.loadMessageBundle()
 const logger = getLogger('notifications')
 
 /**
@@ -25,6 +27,8 @@ const logger = getLogger('notifications')
  * and does not determine what notifications to dispaly or how to fetch and store them.
  */
 export class NotificationsNode implements TreeNode {
+    public static readonly title = localize('AWS.notifications.title', 'Notifications')
+
     public readonly id = 'notifications'
     public readonly resource = this
     public provider?: ResourceTreeDataProvider
@@ -37,6 +41,7 @@ export class NotificationsNode implements TreeNode {
     private readonly showContextStr: contextKey
     private readonly startUpNodeContext: string
     private readonly emergencyNodeContext: string
+    private view: vscode.TreeView<TreeNode> | undefined
 
     static #instance: NotificationsNode
 
@@ -65,7 +70,7 @@ export class NotificationsNode implements TreeNode {
     }
 
     public getTreeItem() {
-        const item = new vscode.TreeItem('Notifications')
+        const item = new vscode.TreeItem(NotificationsNode.title)
         item.collapsibleState = vscode.TreeItemCollapsibleState.Collapsed
         item.contextValue = 'notifications'
 
@@ -73,8 +78,23 @@ export class NotificationsNode implements TreeNode {
     }
 
     public refresh(): void {
-        const hasNotifications = this.startUpNotifications.length > 0 || this.emergencyNotifications.length > 0
-        void setContext(this.showContextStr, hasNotifications)
+        if (!this.view) {
+            throw new ToolkitError('NotificationsNode: TreeView accessed without being registered.')
+        }
+
+        const totalNotifications = this.notificationCount()
+        if (totalNotifications > 0) {
+            this.view.badge = {
+                tooltip: `${totalNotifications} notification${totalNotifications > 1 ? 's' : ''}`,
+                value: totalNotifications,
+            }
+            this.view.title = `${NotificationsNode.title} (${totalNotifications})`
+            void setContext(this.showContextStr, true)
+        } else {
+            this.view.badge = undefined
+            this.view.title = NotificationsNode.title
+            void setContext(this.showContextStr, false)
+        }
 
         this.provider?.refresh()
     }
@@ -88,11 +108,12 @@ export class NotificationsNode implements TreeNode {
                       })
                     : (getIcon('vscode-question') as Icon)
 
+            const title = n.uiRenderInstructions.content['en-US'].title
             return this.openNotificationCmd.build(n).asTreeNode({
-                label: n.uiRenderInstructions.content['en-US'].title,
+                label: title,
+                tooltip: title,
                 iconPath: icon,
                 contextValue: type === 'startUp' ? this.startUpNodeContext : this.emergencyNodeContext,
-                tooltip: 'Click to open',
             })
         }
 
@@ -130,6 +151,10 @@ export class NotificationsNode implements TreeNode {
         return vscode.commands.executeCommand(this.focusCmdStr)
     }
 
+    private notificationCount() {
+        return this.startUpNotifications.length + this.emergencyNotifications.length
+    }
+
     /**
      * Fired when a notification is clicked on in the panel. It will run any rendering
      * instructions included in the notification. See {@link ToolkitNotification.uiRenderInstructions}.
@@ -275,7 +300,7 @@ export class NotificationsNode implements TreeNode {
     }
 
     registerView(context: vscode.ExtensionContext) {
-        const view = registerToolView(
+        this.view = registerToolView(
             {
                 nodes: [this],
                 view: isAmazonQ() ? 'aws.amazonq.notifications' : 'aws.toolkit.notifications',
@@ -283,6 +308,5 @@ export class NotificationsNode implements TreeNode {
             },
             context
         )
-        view.message = `New feature announcements and emergency notifications for ${isAmazonQ() ? 'Amazon Q' : 'AWS Toolkit'} will appear here.`
     }
 }

From 29dd9c2a3224c2017b44e9b787f6a4deb8c0bf43 Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Fri, 15 Nov 2024 11:06:53 -0500
Subject: [PATCH 6/9] test(notifications): fix and improve

- Fix async tests broken by recent commits.
- Add setContext and telemetry verification
---
 packages/core/src/notifications/controller.ts |  6 +--
 packages/core/src/notifications/panelNode.ts  | 27 ++++++-------
 .../src/test/notifications/controller.test.ts | 40 ++++++++++++++++++-
 .../src/test/notifications/rendering.test.ts  | 23 +++++------
 4 files changed, 64 insertions(+), 32 deletions(-)

diff --git a/packages/core/src/notifications/controller.ts b/packages/core/src/notifications/controller.ts
index 9d1636c7a24..735721ad953 100644
--- a/packages/core/src/notifications/controller.ts
+++ b/packages/core/src/notifications/controller.ts
@@ -40,8 +40,6 @@ import { telemetry } from '../shared/telemetry/telemetry'
 export class NotificationsController {
     public static readonly suggestedPollIntervalMs = 1000 * 60 * 10 // 10 minutes
 
-    public readonly storageKey: globalKey
-
     /** Internal memory state that is written to global state upon modification. */
     private readonly state: NotificationsState
 
@@ -49,7 +47,8 @@ export class NotificationsController {
 
     constructor(
         private readonly notificationsNode: NotificationsNode,
-        private readonly fetcher: NotificationFetcher = new RemoteFetcher()
+        private readonly fetcher: NotificationFetcher = new RemoteFetcher(),
+        public readonly storageKey: globalKey = 'aws.notifications'
     ) {
         if (!NotificationsController.#instance) {
             // Register on first creation only.
@@ -57,7 +56,6 @@ export class NotificationsController {
         }
         NotificationsController.#instance = this
 
-        this.storageKey = 'aws.notifications'
         this.state = globals.globalState.tryGet<NotificationsState>(this.storageKey, NotificationsStateConstructor, {
             startUp: {},
             emergency: {},
diff --git a/packages/core/src/notifications/panelNode.ts b/packages/core/src/notifications/panelNode.ts
index cfcc0ad5eac..d3e93109da9 100644
--- a/packages/core/src/notifications/panelNode.ts
+++ b/packages/core/src/notifications/panelNode.ts
@@ -78,24 +78,23 @@ export class NotificationsNode implements TreeNode {
     }
 
     public refresh(): void {
-        if (!this.view) {
-            throw new ToolkitError('NotificationsNode: TreeView accessed without being registered.')
-        }
-
         const totalNotifications = this.notificationCount()
-        if (totalNotifications > 0) {
-            this.view.badge = {
-                tooltip: `${totalNotifications} notification${totalNotifications > 1 ? 's' : ''}`,
-                value: totalNotifications,
+        if (this.view) {
+            if (totalNotifications > 0) {
+                this.view.badge = {
+                    tooltip: `${totalNotifications} notification${totalNotifications > 1 ? 's' : ''}`,
+                    value: totalNotifications,
+                }
+                this.view.title = `${NotificationsNode.title} (${totalNotifications})`
+            } else {
+                this.view.badge = undefined
+                this.view.title = NotificationsNode.title
             }
-            this.view.title = `${NotificationsNode.title} (${totalNotifications})`
-            void setContext(this.showContextStr, true)
         } else {
-            this.view.badge = undefined
-            this.view.title = NotificationsNode.title
-            void setContext(this.showContextStr, false)
+            logger.warn('NotificationsNode was refreshed but the view was not initialized!')
         }
 
+        void setContext(this.showContextStr, totalNotifications > 0)
         this.provider?.refresh()
     }
 
@@ -240,7 +239,7 @@ export class NotificationsNode implements TreeNode {
                                 case 'updateAndReload':
                                     // Give things time to finish executing.
                                     globals.clock.setTimeout(() => {
-                                        void this.updateAndReload(notification.displayIf.extensionId)
+                                        return this.updateAndReload(notification.displayIf.extensionId)
                                     }, 1000)
                                     break
                                 case 'openUrl':
diff --git a/packages/core/src/test/notifications/controller.test.ts b/packages/core/src/test/notifications/controller.test.ts
index 669700a9cde..1f097fe7b9f 100644
--- a/packages/core/src/test/notifications/controller.test.ts
+++ b/packages/core/src/test/notifications/controller.test.ts
@@ -8,7 +8,8 @@ import * as FakeTimers from '@sinonjs/fake-timers'
 import assert from 'assert'
 import sinon from 'sinon'
 import globals from '../../shared/extensionGlobals'
-import { randomUUID } from '../../shared'
+import { randomUUID } from '../../shared/crypto'
+import * as setContext from '../../shared/vscode/setContext'
 import { assertTelemetry, installFakeClock } from '../testUtil'
 import {
     NotificationFetcher,
@@ -89,7 +90,7 @@ describe('Notifications Controller', function () {
     beforeEach(async function () {
         panelNode.setNotifications([], [])
         fetcher = new TestFetcher()
-        controller = new NotificationsController(panelNode, fetcher)
+        controller = new NotificationsController(panelNode, fetcher, '_aws.test.notification' as any)
 
         ruleEngineSpy = sinon.spy(ruleEngine, 'shouldDisplayNotification')
         focusPanelSpy = sinon.spy(panelNode, 'focusPanel')
@@ -107,6 +108,10 @@ describe('Notifications Controller', function () {
     })
 
     it('can fetch and store startup notifications', async function () {
+        // There seems to be a race condition with having a global spy object that is reset after each
+        // test. Doesn't seem to affect the other ones, for some reason.
+        const setContextSpy = sinon.spy(setContext, 'setContext')
+
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -135,9 +140,14 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.startUpNotifications, [content.notifications[0]])
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(focusPanelSpy.callCount, 0)
+        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
+
+        setContextSpy.restore()
     })
 
     it('can fetch and store emergency notifications', async function () {
+        const setContextSpy = sinon.spy(setContext, 'setContext')
+
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -166,9 +176,14 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.emergencyNotifications, [content.notifications[0]])
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(focusPanelSpy.callCount, 1)
+        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
+
+        setContextSpy.restore()
     })
 
     it('can fetch and store both startup and emergency notifications', async function () {
+        const setContextSpy = sinon.spy(setContext, 'setContext')
+
         const eTag1 = randomUUID()
         const eTag2 = randomUUID()
         const startUpContent = {
@@ -224,9 +239,14 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.emergencyNotifications, [emergencyContent.notifications[0]])
         assert.equal(panelNode.getChildren().length, 2)
         assert.equal(focusPanelSpy.callCount, 1)
+        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
+
+        setContextSpy.restore()
     })
 
     it('dismisses a startup notification', async function () {
+        const setContextSpy = sinon.spy(setContext, 'setContext')
+
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -241,6 +261,7 @@ describe('Notifications Controller', function () {
 
         assert.equal(panelNode.getChildren().length, 2)
         assert.equal(panelNode.startUpNotifications.length, 2)
+        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
 
         assert.deepStrictEqual(await globals.globalState.get(controller.storageKey), {
             startUp: {
@@ -267,9 +288,13 @@ describe('Notifications Controller', function () {
 
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(panelNode.startUpNotifications.length, 1)
+
+        setContextSpy.restore()
     })
 
     it('does not redisplay dismissed notifications', async function () {
+        const setContextSpy = sinon.spy(setContext, 'setContext')
+
         const content = {
             schemaVersion: '1.x',
             notifications: [getValidTestNotification('id:startup1')],
@@ -281,9 +306,11 @@ describe('Notifications Controller', function () {
 
         await controller.pollForStartUp(ruleEngine)
         assert.equal(panelNode.getChildren().length, 1)
+        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
 
         await dismissNotification(content.notifications[0])
         assert.equal(panelNode.getChildren().length, 0)
+        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', false])
 
         content.notifications.push(getValidTestNotification('id:startup2'))
         fetcher.setStartUpContent({
@@ -305,6 +332,9 @@ describe('Notifications Controller', function () {
         })
 
         assert.equal(panelNode.getChildren().length, 1)
+        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
+
+        setContextSpy.restore()
     })
 
     it('does not refocus emergency notifications', async function () {
@@ -376,6 +406,8 @@ describe('Notifications Controller', function () {
     })
 
     it('cleans out dismissed state', async function () {
+        const setContextSpy = sinon.spy(setContext, 'setContext')
+
         const startUpContent = {
             schemaVersion: '1.x',
             notifications: [getValidTestNotification('id:startup1')],
@@ -434,6 +466,7 @@ describe('Notifications Controller', function () {
             newlyReceived: [],
         })
         assert.equal(panelNode.getChildren().length, 1)
+        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
 
         fetcher.setEmergencyContent({
             eTag: '1',
@@ -455,6 +488,9 @@ describe('Notifications Controller', function () {
         })
 
         assert.equal(panelNode.getChildren().length, 0)
+        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', false])
+
+        setContextSpy.restore()
     })
 
     it('does not rethrow errors when fetching', async function () {
diff --git a/packages/core/src/test/notifications/rendering.test.ts b/packages/core/src/test/notifications/rendering.test.ts
index 4df1caddc56..e446e53fdba 100644
--- a/packages/core/src/test/notifications/rendering.test.ts
+++ b/packages/core/src/test/notifications/rendering.test.ts
@@ -12,28 +12,20 @@ import { panelNode } from './controller.test'
 import { getTestWindow } from '../shared/vscode/window'
 import * as VsCodeUtils from '../../shared/utilities/vsCodeUtils'
 import { assertTextEditorContains, installFakeClock } from '../testUtil'
+import { waitUntil } from '../../shared/utilities/timeoutUtils'
+import globals from '../../shared/extensionGlobals'
 
 describe('Notifications Rendering', function () {
     let sandbox: sinon.SinonSandbox
-    let clock: FakeTimers.InstalledClock
-
-    before(function () {
-        clock = installFakeClock()
-    })
 
     beforeEach(function () {
         sandbox = sinon.createSandbox()
     })
 
     afterEach(function () {
-        clock.reset()
         sandbox.restore()
     })
 
-    after(function () {
-        clock.uninstall()
-    })
-
     // util to test txt pop-up under different senarios
     async function verifyTxtNotification(notification: ToolkitNotification) {
         const expectedContent = notification.uiRenderInstructions.content['en-US'].description
@@ -103,14 +95,21 @@ describe('Notifications Rendering', function () {
             message.selectItem('Update and Reload')
         })
         const excuteCommandStub = sandbox.stub(vscode.commands, 'executeCommand').resolves()
+        const telemetrySpy = sandbox.spy(globals.telemetry, 'flushRecords')
         const notification = getModalNotification()
-        await panelNode.openNotification(notification)
 
         // Update and Reload is put on a timer so that other methods (e.g. telemetry) can finish.
-        await clock.tickAsync(2000)
+        const clock: FakeTimers.InstalledClock = installFakeClock()
+        await panelNode.openNotification(notification)
+
+        await clock.tickAsync(1000)
+        clock.uninstall()
+
+        await waitUntil(async () => excuteCommandStub.called, { interval: 5, timeout: 5000 })
 
         assert.ok(excuteCommandStub.calledWith('workbench.extensions.installExtension', 'aws.toolkit.fake.extension'))
         assert.ok(excuteCommandStub.calledWith('workbench.action.reloadWindow'))
+        assert.ok(telemetrySpy.calledOnce)
     })
 
     it('executes openURL type button', async function () {

From 27dad6eab66bc1b99dd2e1c5d99bb7eb5dc55f16 Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Fri, 15 Nov 2024 12:49:40 -0500
Subject: [PATCH 7/9] feat(context): add getContext, update notifications tests

Add counterpart to setContext wrapper.
---
 packages/core/src/notifications/controller.ts |  4 +-
 packages/core/src/notifications/panelNode.ts  | 12 ++---
 packages/core/src/shared/vscode/setContext.ts | 20 ++++++--
 .../src/test/notifications/controller.test.ts | 48 +++++--------------
 4 files changed, 35 insertions(+), 49 deletions(-)

diff --git a/packages/core/src/notifications/controller.ts b/packages/core/src/notifications/controller.ts
index 735721ad953..a93835d89b1 100644
--- a/packages/core/src/notifications/controller.ts
+++ b/packages/core/src/notifications/controller.ts
@@ -92,7 +92,7 @@ export class NotificationsController {
             ruleEngine.shouldDisplayNotification(n)
         )
 
-        NotificationsNode.instance.setNotifications(startUp, emergency)
+        await NotificationsNode.instance.setNotifications(startUp, emergency)
 
         // Emergency notifications can't be dismissed, but if the user minimizes the panel then
         // we don't want to focus it each time we set the notification nodes.
@@ -126,7 +126,7 @@ export class NotificationsController {
         this.state.dismissed.push(notificationId)
         await this.writeState()
 
-        NotificationsNode.instance.dismissStartUpNotification(notificationId)
+        await NotificationsNode.instance.dismissStartUpNotification(notificationId)
     }
 
     /**
diff --git a/packages/core/src/notifications/panelNode.ts b/packages/core/src/notifications/panelNode.ts
index d3e93109da9..94e4f2031f1 100644
--- a/packages/core/src/notifications/panelNode.ts
+++ b/packages/core/src/notifications/panelNode.ts
@@ -77,7 +77,7 @@ export class NotificationsNode implements TreeNode {
         return item
     }
 
-    public refresh(): void {
+    public refresh() {
         const totalNotifications = this.notificationCount()
         if (this.view) {
             if (totalNotifications > 0) {
@@ -94,8 +94,8 @@ export class NotificationsNode implements TreeNode {
             logger.warn('NotificationsNode was refreshed but the view was not initialized!')
         }
 
-        void setContext(this.showContextStr, totalNotifications > 0)
         this.provider?.refresh()
+        return setContext(this.showContextStr, totalNotifications > 0)
     }
 
     public getChildren() {
@@ -126,10 +126,10 @@ export class NotificationsNode implements TreeNode {
      * Sets the current list of notifications. Nodes are generated for each notification.
      * No other processing is done, see NotificationController.
      */
-    public setNotifications(startUp: ToolkitNotification[], emergency: ToolkitNotification[]) {
+    public async setNotifications(startUp: ToolkitNotification[], emergency: ToolkitNotification[]) {
         this.startUpNotifications = startUp
         this.emergencyNotifications = emergency
-        this.refresh()
+        await this.refresh()
     }
 
     /**
@@ -138,9 +138,9 @@ export class NotificationsNode implements TreeNode {
      *
      * Only dismisses startup notifications.
      */
-    public dismissStartUpNotification(id: string) {
+    public async dismissStartUpNotification(id: string) {
         this.startUpNotifications = this.startUpNotifications.filter((n) => n.id !== id)
-        this.refresh()
+        await this.refresh()
     }
 
     /**
diff --git a/packages/core/src/shared/vscode/setContext.ts b/packages/core/src/shared/vscode/setContext.ts
index 57450126412..7786a620962 100644
--- a/packages/core/src/shared/vscode/setContext.ts
+++ b/packages/core/src/shared/vscode/setContext.ts
@@ -36,22 +36,34 @@ export type contextKey =
     | 'amazonq.inline.codelensShortcutEnabled'
     | 'aws.toolkit.lambda.walkthroughSelected'
 
+const contextMap: Partial<Record<contextKey, any>> = {}
+
 /**
  * Calls the vscode "setContext" command.
  *
  * This wrapper adds structure and traceability to the vscode "setContext". It also opens the door
- * for:
- * - validation
- * - getContext() (see also https://github.com/microsoft/vscode/issues/10471)
+ * for validation.
  *
  * Use "setContext" only as a last resort, to set flags that are detectable in package.json
  * declarations. Do not use it as a general way to store global state (which should be avoided
  * anyway).
  *
  * Warning: vscode context keys/values are NOT isolated to individual extensions. Other extensions
- * can read and modify them.
+ * can read and modify them. See also https://github.com/microsoft/vscode/issues/10471
  */
 export async function setContext(key: contextKey, val: any): Promise<void> {
     // eslint-disable-next-line aws-toolkits/no-banned-usages
     await vscode.commands.executeCommand('setContext', key, val)
+    contextMap[key] = val
+}
+
+/**
+ * Returns the value of a context key set via {@link setContext} wrapper for this session.
+ *
+ * Warning: this does not gaurantee the state of the context key in vscode because it may have
+ * been set via `vscode.commands.executeCommand('setContext')`. It has no connection the
+ * context keys stored in vscode itself because an API for this is not exposed.
+ */
+export function getContext(key: contextKey): any {
+    return contextMap[key]
 }
diff --git a/packages/core/src/test/notifications/controller.test.ts b/packages/core/src/test/notifications/controller.test.ts
index 1f097fe7b9f..c1e2d299463 100644
--- a/packages/core/src/test/notifications/controller.test.ts
+++ b/packages/core/src/test/notifications/controller.test.ts
@@ -9,7 +9,7 @@ import assert from 'assert'
 import sinon from 'sinon'
 import globals from '../../shared/extensionGlobals'
 import { randomUUID } from '../../shared/crypto'
-import * as setContext from '../../shared/vscode/setContext'
+import { getContext } from '../../shared/vscode/setContext'
 import { assertTelemetry, installFakeClock } from '../testUtil'
 import {
     NotificationFetcher,
@@ -88,7 +88,7 @@ describe('Notifications Controller', function () {
     }
 
     beforeEach(async function () {
-        panelNode.setNotifications([], [])
+        await panelNode.setNotifications([], [])
         fetcher = new TestFetcher()
         controller = new NotificationsController(panelNode, fetcher, '_aws.test.notification' as any)
 
@@ -108,10 +108,6 @@ describe('Notifications Controller', function () {
     })
 
     it('can fetch and store startup notifications', async function () {
-        // There seems to be a race condition with having a global spy object that is reset after each
-        // test. Doesn't seem to affect the other ones, for some reason.
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -140,14 +136,10 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.startUpNotifications, [content.notifications[0]])
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(focusPanelSpy.callCount, 0)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('can fetch and store emergency notifications', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -176,14 +168,10 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.emergencyNotifications, [content.notifications[0]])
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(focusPanelSpy.callCount, 1)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('can fetch and store both startup and emergency notifications', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag1 = randomUUID()
         const eTag2 = randomUUID()
         const startUpContent = {
@@ -239,14 +227,10 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.emergencyNotifications, [emergencyContent.notifications[0]])
         assert.equal(panelNode.getChildren().length, 2)
         assert.equal(focusPanelSpy.callCount, 1)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('dismisses a startup notification', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -261,7 +245,7 @@ describe('Notifications Controller', function () {
 
         assert.equal(panelNode.getChildren().length, 2)
         assert.equal(panelNode.startUpNotifications.length, 2)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
 
         assert.deepStrictEqual(await globals.globalState.get(controller.storageKey), {
             startUp: {
@@ -288,13 +272,9 @@ describe('Notifications Controller', function () {
 
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(panelNode.startUpNotifications.length, 1)
-
-        setContextSpy.restore()
     })
 
     it('does not redisplay dismissed notifications', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const content = {
             schemaVersion: '1.x',
             notifications: [getValidTestNotification('id:startup1')],
@@ -306,11 +286,11 @@ describe('Notifications Controller', function () {
 
         await controller.pollForStartUp(ruleEngine)
         assert.equal(panelNode.getChildren().length, 1)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
 
         await dismissNotification(content.notifications[0])
         assert.equal(panelNode.getChildren().length, 0)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', false])
+        assert.equal(getContext('aws.toolkit.notifications.show'), false)
 
         content.notifications.push(getValidTestNotification('id:startup2'))
         fetcher.setStartUpContent({
@@ -332,9 +312,7 @@ describe('Notifications Controller', function () {
         })
 
         assert.equal(panelNode.getChildren().length, 1)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('does not refocus emergency notifications', async function () {
@@ -406,8 +384,6 @@ describe('Notifications Controller', function () {
     })
 
     it('cleans out dismissed state', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const startUpContent = {
             schemaVersion: '1.x',
             notifications: [getValidTestNotification('id:startup1')],
@@ -466,7 +442,7 @@ describe('Notifications Controller', function () {
             newlyReceived: [],
         })
         assert.equal(panelNode.getChildren().length, 1)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
 
         fetcher.setEmergencyContent({
             eTag: '1',
@@ -488,9 +464,7 @@ describe('Notifications Controller', function () {
         })
 
         assert.equal(panelNode.getChildren().length, 0)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', false])
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), false)
     })
 
     it('does not rethrow errors when fetching', async function () {

From a4bc3eb0e24e7d94c6abbc00672efb9b6d0bb593 Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Fri, 15 Nov 2024 13:14:27 -0500
Subject: [PATCH 8/9] fix(notifications): await polling in activation

Prevents async functions from running at same time, which may cause notifications to fire twice, etc.
---
 packages/core/src/notifications/activation.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/notifications/activation.ts b/packages/core/src/notifications/activation.ts
index a563880d46f..2ac36bcd5bd 100644
--- a/packages/core/src/notifications/activation.ts
+++ b/packages/core/src/notifications/activation.ts
@@ -38,8 +38,8 @@ export async function activate(
     const controller = new NotificationsController(panelNode)
     const engine = new RuleEngine(await getRuleContext(context, initialState))
 
-    void controller.pollForStartUp(engine)
-    void controller.pollForEmergencies(engine)
+    await controller.pollForStartUp(engine)
+    await controller.pollForEmergencies(engine)
 
     globals.clock.setInterval(async () => {
         const ruleContext = await getRuleContext(context, await authStateFn())

From 2f07377e2f005087f05306e0b139ba0e99e24690 Mon Sep 17 00:00:00 2001
From: Maxim Hayes <149123719+hayemaxi@users.noreply.github.com>
Date: Fri, 15 Nov 2024 17:22:36 -0500
Subject: [PATCH 9/9] docs: update wording for getContext

Co-authored-by: Justin M. Keyes <jmkeyes@amazon.com>
---
 packages/core/src/shared/vscode/setContext.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/core/src/shared/vscode/setContext.ts b/packages/core/src/shared/vscode/setContext.ts
index 7786a620962..2e8d1a6e750 100644
--- a/packages/core/src/shared/vscode/setContext.ts
+++ b/packages/core/src/shared/vscode/setContext.ts
@@ -60,7 +60,7 @@ export async function setContext(key: contextKey, val: any): Promise<void> {
 /**
  * Returns the value of a context key set via {@link setContext} wrapper for this session.
  *
- * Warning: this does not gaurantee the state of the context key in vscode because it may have
+ * Warning: this does not guarantee the state of the context key in vscode because it may have
  * been set via `vscode.commands.executeCommand('setContext')`. It has no connection the
  * context keys stored in vscode itself because an API for this is not exposed.
  */