test: add notification and reset state dev options (#6105)

- Add a generic dev menu item to reset the state of a feature. Simply
extend with a function handler to add an item to the menu. The initial
item resets in-ide notification state.

![image](https://github.com/user-attachments/assets/042bdad2-2bac-4694-a7ff-801a857bddd3)

![image](https://github.com/user-attachments/assets/43fa90aa-af58-4b6f-af0d-ca38fb01f5b7)

- Add a way to test notifications locally without modifying code. Using
a dev menu item you can add a notification to global state and receive
it in your IDE for visualization.

![image](https://github.com/user-attachments/assets/2a0f267c-e792-4535-a1a4-170f8f6f32ca)

![image](https://github.com/user-attachments/assets/dd08402c-c4db-422c-9b60-20cd4f9340a2)

- Refactor notifications controller to not require a RuleEngine when
polling is requested.

---

<!--- REMINDER: Ensure that your PR meets the guidelines in
CONTRIBUTING.md -->

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: hayemaxi

packages/amazonq/src/extensionNode.ts

@@ -18,7 +18,7 @@ import { Auth, AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection } from
 import api from './api'
 import { activate as activateCWChat } from './app/chat/activation'
 import { beta } from 'aws-core-vscode/dev'
-import { activate as activateNotifications } from 'aws-core-vscode/notifications'
+import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications'
 import { AuthState, AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry'
 
@@ -67,15 +67,14 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
 
     // TODO: Should probably emit for web as well.
     // Will the web metric look the same?
-    const authState = await getAuthState()
     telemetry.auth_userState.emit({
         passive: true,
         result: 'Succeeded',
         source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(),
-        ...authState,
+        ...(await getAuthState()),
     })
 
-    void activateNotifications(context, authState, getAuthState)
+    void activateNotifications(context, getAuthState)
 }
 
 async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
@@ -122,13 +121,16 @@ async function setupDevMode(context: vscode.ExtensionContext) {
 
     const devOptions: DevOptions = {
         context,
-        auth: Auth.instance,
+        auth: () => Auth.instance,
+        notificationsController: () => NotificationsController.instance,
         menuOptions: [
             'editStorage',
+            'resetState',
             'showEnvVars',
             'deleteSsoConnections',
             'expireSsoConnections',
             'editAuthConnections',
+            'notificationsSend',
             'forceIdeCrash',
         ],
     }

packages/core/src/dev/activation.ts

@@ -21,6 +21,9 @@ import { getEnvironmentSpecificMemento } from '../shared/utilities/mementos'
 import { setContext } from '../shared'
 import { telemetry } from '../shared/telemetry'
 import { getSessionId } from '../shared/telemetry/util'
+import { NotificationsController } from '../notifications/controller'
+import { DevNotificationsState } from '../notifications/types'
+import { QuickPickItem } from 'vscode'
 
 interface MenuOption {
     readonly label: string
@@ -34,21 +37,25 @@ export type DevFunction =
     | 'openTerminal'
     | 'deleteDevEnv'
     | 'editStorage'
+    | 'resetState'
     | 'showEnvVars'
     | 'deleteSsoConnections'
     | 'expireSsoConnections'
     | 'editAuthConnections'
+    | 'notificationsSend'
     | 'forceIdeCrash'
 
 export type DevOptions = {
     context: vscode.ExtensionContext
-    auth: Auth
+    auth: () => Auth
+    notificationsController: () => NotificationsController
     menuOptions?: DevFunction[]
 }
 
 let targetContext: vscode.ExtensionContext
 let globalState: vscode.Memento
 let targetAuth: Auth
+let targetNotificationsController: NotificationsController
 
 /**
  * Defines AWS Toolkit developer tools.
@@ -83,6 +90,11 @@ const menuOptions: () => Record<DevFunction, MenuOption> = () => {
             detail: 'Shows all globalState values, or edit a globalState/secret item',
             executor: openStorageFromInput,
         },
+        resetState: {
+            label: 'Reset feature state',
+            detail: 'Quick reset the state of extension components or features',
+            executor: resetState,
+        },
         showEnvVars: {
             label: 'Show Environment Variables',
             description: 'AWS Toolkit',
@@ -104,6 +116,11 @@ const menuOptions: () => Record<DevFunction, MenuOption> = () => {
             detail: 'Opens editor to all Auth Connections the extension is using.',
             executor: editSsoConnections,
         },
+        notificationsSend: {
+            label: 'Notifications: Send Notifications',
+            detail: 'Send JSON notifications for testing.',
+            executor: editNotifications,
+        },
         forceIdeCrash: {
             label: 'Crash: Force IDE ExtHost Crash',
             detail: `Will SIGKILL ExtHost, { pid: ${process.pid}, sessionId: '${getSessionId().slice(0, 8)}-...' }, but the IDE itself will not crash.`,
@@ -156,14 +173,19 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
         vscode.workspace.registerTextDocumentContentProvider('aws-dev2', new DevDocumentProvider()),
         // "AWS (Developer): Open Developer Menu"
         vscode.commands.registerCommand('aws.dev.openMenu', async () => {
-            await vscode.commands.executeCommand('_aws.dev.invokeMenu', { context: ctx, auth: Auth.instance })
+            await vscode.commands.executeCommand('_aws.dev.invokeMenu', {
+                context: ctx,
+                auth: () => Auth.instance,
+                notificationsController: () => NotificationsController.instance,
+            })
         }),
         // Internal command to open dev menu for a specific context and options
         vscode.commands.registerCommand('_aws.dev.invokeMenu', (opts: DevOptions) => {
             targetContext = opts.context
             // eslint-disable-next-line aws-toolkits/no-banned-usages
             globalState = targetContext.globalState
-            targetAuth = opts.auth
+            targetAuth = opts.auth()
+            targetNotificationsController = opts.notificationsController()
             const options = menuOptions()
             void openMenu(
                 entries(options)
@@ -302,7 +324,7 @@ class ObjectEditor {
         vscode.workspace.registerFileSystemProvider(ObjectEditor.scheme, this.fs)
     }
 
-    public async openStorage(type: 'globalsView' | 'globals' | 'secrets' | 'auth', key: string): Promise<void> {
+    public async openStorage(type: 'globalsView' | 'globals' | 'secrets' | 'auth', key: string) {
         switch (type) {
             case 'globalsView':
                 return showState('globalstate')
@@ -316,17 +338,19 @@ class ObjectEditor {
         }
     }
 
-    private async openState(storage: vscode.Memento | vscode.SecretStorage, key: string): Promise<void> {
+    private async openState(storage: vscode.Memento | vscode.SecretStorage, key: string) {
         const uri = this.uriFromKey(key, storage)
         const tab = this.tabs.get(this.fs.uriToKey(uri))
 
         if (tab) {
             tab.virtualFile.refresh()
             await vscode.window.showTextDocument(tab.editor.document)
+            return tab.virtualFile
         } else {
             const newTab = await this.createTab(storage, key)
             const newKey = this.fs.uriToKey(newTab.editor.document.uri)
             this.tabs.set(newKey, newTab)
+            return newTab.virtualFile
         }
     }
 
@@ -417,6 +441,62 @@ async function openStorageFromInput() {
     }
 }
 
+type ResettableFeature = {
+    name: string
+    executor: () => Promise<void> | void
+} & QuickPickItem
+
+/**
+ * Extend this array with features that may need state resets often for
+ * testing purposes. It will appear as an entry in the "Reset feature state" menu.
+ */
+const resettableFeatures: readonly ResettableFeature[] = [
+    {
+        name: 'notifications',
+        label: 'Notifications',
+        detail: 'Resets memory/global state for the notifications panel (includes dismissed, onReceive).',
+        executor: resetNotificationsState,
+    },
+] as const
+
+// TODO this is *somewhat* similar to `openStorageFromInput`. If we need another
+// one of these prompters, can we make it generic?
+async function resetState() {
+    const wizard = new (class extends Wizard<{ target: string; key: string }> {
+        constructor() {
+            super()
+
+            this.form.target.bindPrompter(() =>
+                createQuickPick(
+                    resettableFeatures.map((f) => {
+                        return {
+                            data: f.name,
+                            label: f.label,
+                            detail: f.detail,
+                        }
+                    }),
+                    {
+                        title: 'Select a feature/component to reset',
+                    }
+                )
+            )
+
+            this.form.key.bindPrompter(({ target }) => {
+                if (target && resettableFeatures.some((f) => f.name === target)) {
+                    return new SkipPrompter('')
+                }
+                throw new Error('invalid feature target')
+            })
+        }
+    })()
+
+    const response = await wizard.run()
+
+    if (response) {
+        return resettableFeatures.find((f) => f.name === response.target)?.executor()
+    }
+}
+
 async function editSsoConnections() {
     void openStorageCommand.execute('auth', 'auth.profiles')
 }
@@ -460,3 +540,41 @@ export const openStorageCommand = Commands.from(ObjectEditor).declareOpenStorage
 export async function updateDevMode() {
     await setContext('aws.isDevMode', DevSettings.instance.isDevMode())
 }
+
+async function resetNotificationsState() {
+    await targetNotificationsController.reset()
+}
+
+async function editNotifications() {
+    const storageKey = 'aws.notifications.dev'
+    const current = globalState.get(storageKey) ?? {}
+    const isValid = (item: any) => {
+        if (typeof item !== 'object' || !Array.isArray(item.startUp) || !Array.isArray(item.emergency)) {
+            return false
+        }
+        return true
+    }
+    if (!isValid(current)) {
+        // Set a default state if the developer does not have it or it's malformed.
+        await globalState.update(storageKey, { startUp: [], emergency: [] } as DevNotificationsState)
+    }
+
+    // Monitor for when the global state is updated.
+    // A notification will be sent based on the contents.
+    const virtualFile = await openStorageCommand.execute('globals', storageKey)
+    virtualFile?.onDidChange(async () => {
+        const val = globalState.get(storageKey) as DevNotificationsState
+        if (!isValid(val)) {
+            void vscode.window.showErrorMessage(
+                'Dev mode: invalid notification object provided. State data must take the form: { "startUp": ToolkitNotification[], "emergency": ToolkitNotification[] }'
+            )
+            return
+        }
+
+        // This relies on the controller being built with DevFetcher, as opposed to
+        // the default RemoteFetcher. DevFetcher will check for notifications in the
+        // global state, which was just modified.
+        await targetNotificationsController.pollForStartUp()
+        await targetNotificationsController.pollForEmergencies()
+    })
+}

packages/core/src/extensionNode.ts

@@ -238,15 +238,14 @@ export async function activate(context: vscode.ExtensionContext) {
 
         // TODO: Should probably emit for web as well.
         // Will the web metric look the same?
-        const authState = await getAuthState()
         telemetry.auth_userState.emit({
             passive: true,
             result: 'Succeeded',
             source: ExtensionUse.instance.sourceForTelemetry(),
-            ...authState,
+            ...(await getAuthState()),
         })
 
-        void activateNotifications(context, authState, getAuthState)
+        void activateNotifications(context, getAuthState)
     } catch (error) {
         const stacktrace = (error as Error).stack?.split('\n')
         // truncate if the stacktrace is unusually long

packages/core/src/notifications/activation.ts

@@ -4,9 +4,10 @@
  */
 
 import * as vscode from 'vscode'
-import { NotificationsController } from './controller'
+import { DevSettings } from '../shared/settings'
+import { DevFetcher, NotificationsController, RemoteFetcher } from './controller'
 import { NotificationsNode } from './panelNode'
-import { RuleEngine, getRuleContext } from './rules'
+import { getRuleContext } from './rules'
 import globals from '../shared/extensionGlobals'
 import { AuthState } from './types'
 import { getLogger } from '../shared/logger/logger'
@@ -24,25 +25,26 @@ const emergencyPollTime = oneMinute * 10
  * @param initialState initial auth state
  * @param authStateFn fn to get current auth state
  */
-export async function activate(
-    context: vscode.ExtensionContext,
-    initialState: AuthState,
-    authStateFn: () => Promise<AuthState>
-) {
+export async function activate(context: vscode.ExtensionContext, authStateFn: () => Promise<AuthState>) {
     try {
         const panelNode = NotificationsNode.instance
         panelNode.registerView(context)
 
-        const controller = new NotificationsController(panelNode)
-        const engine = new RuleEngine(await getRuleContext(context, initialState))
-
-        await controller.pollForStartUp(engine)
-        await controller.pollForEmergencies(engine)
-
-        globals.clock.setInterval(async () => {
-            const ruleContext = await getRuleContext(context, await authStateFn())
-            await controller.pollForEmergencies(new RuleEngine(ruleContext))
-        }, emergencyPollTime)
+        const controller = new NotificationsController(
+            panelNode,
+            async () => await getRuleContext(context, await authStateFn()),
+            DevSettings.instance.isDevMode() ? new DevFetcher() : new RemoteFetcher()
+        )
+
+        await controller.pollForStartUp()
+        await controller.pollForEmergencies()
+
+        globals.clock.setInterval(
+            async () => {
+                await controller.pollForEmergencies()
+            },
+            DevSettings.instance.get('notificationsPollInterval', emergencyPollTime)
+        )
 
         logger.debug('Activated in-IDE notifications polling module')
     } catch (err) {