feat(ui): show login prompt for split sessions (#5182)

**Separate sessions commit**

- If user is logged out of toolkit due to extension session splitting, display a prompt that opens the sign in page when clicked.
- Dismissing or clicking the button dismisses it permanently.
- The prompt will continue to show across restarts until it is dismissed, the button is pressed, or the user signs into toolkit.
- The prompt will display only once, regardless of how many connections are logged out (forgotten).
- The sign in button will go to the login page with the proper scopes, i.e.g if codecatalyst was signed out, go to code catalyst login. Otherwise go to explorer only login.

Author: hayemaxi

packages/core/src/auth/auth.ts

@@ -59,7 +59,7 @@ import {
     scopesSsoAccountAccess,
     AwsConnection,
 } from './connection'
-import { isSageMaker, isCloud9 } from '../shared/extensionUtilities'
+import { isSageMaker, isCloud9, isAmazonQ } from '../shared/extensionUtilities'
 import { telemetry } from '../shared/telemetry/telemetry'
 import { randomUUID } from '../common/crypto'
 
@@ -300,6 +300,11 @@ export class Auth implements AuthService, ConnectionManager {
             metadata: { connectionState: 'unauthenticated' },
         })
 
+        // Remove the split session logout prompt, if it exists.
+        if (!isAmazonQ()) {
+            await globals.context.globalState.update(SessionSeparationPrompt.instance.dismissKey, true)
+        }
+
         try {
             ;(await tokenProvider.getToken()) ?? (await tokenProvider.createToken())
             const storedProfile = await this.store.addProfile(id, profile)
@@ -1000,6 +1005,7 @@ export class Auth implements AuthService, ConnectionManager {
         delete this._declaredConnections[conn.startUrl]
     }
 }
+
 /**
  * Returns true if credentials are provided by the environment (ex. via ~/.aws/)
  *
@@ -1012,3 +1018,70 @@ export function hasVendedIamCredentials(isC9?: boolean, isSM?: boolean) {
     isSM ??= isSageMaker()
     return isSM || isC9
 }
+
+type LoginCommand = 'aws.toolkit.auth.manageConnections' | 'aws.codecatalyst.manageConnections'
+/**
+ * Temporary class that handles notifiting users who were logged out as part of
+ * splitting auth sessions between extensions.
+ *
+ * TODO: Remove after some time.
+ */
+export class SessionSeparationPrompt {
+    public readonly dismissKey = 'aws.toolkit.separationPromptDismissed'
+    private readonly loginCmdKey = 'aws.toolkit.separationPromptCommand'
+
+    // Local variable handles per session displays, e.g. we forgot a CodeCatalyst connection AND
+    // an Explorer only connection. We only want to display once in this case.
+    // However, we don't want to set this at the global state level until a user interacts with the
+    // notification in case they miss it the first time.
+    #separationPromptDisplayed = false
+
+    /**
+     * Open a prompt for that last used command name (or do nothing if no command name has ever been passed),
+     * which is useful to redisplay the prompt after reloads in case a user misses it.
+     */
+    public async showAnyPreviousPrompt() {
+        const cmd = globals.context.globalState.get<string>(this.loginCmdKey)
+        return cmd ? await this.showForCommand(cmd as LoginCommand) : undefined
+    }
+
+    /**
+     * Displays a sign in prompt to the user if they have been logged out of the Toolkit as part of
+     * separating auth sessions between extensions. It will executed the passed command for sign in,
+     * (e.g. codecatalyst sign in vs explorer)
+     */
+    public async showForCommand(cmd: LoginCommand) {
+        if (this.#separationPromptDisplayed || globals.context.globalState.get<boolean>(this.dismissKey)) {
+            return
+        }
+
+        await globals.context.globalState.update(this.loginCmdKey, cmd)
+
+        await telemetry.toolkit_showNotification.run(async () => {
+            telemetry.record({ id: 'sessionSeparation' })
+            this.#separationPromptDisplayed = true
+            void vscode.window
+                .showWarningMessage(
+                    'Amazon Q and AWS Toolkit no longer share connections. Please sign in again to use AWS Toolkit.',
+                    'Sign In'
+                )
+                .then(async resp => {
+                    await telemetry.toolkit_invokeAction.run(async () => {
+                        if (resp === 'Sign In') {
+                            telemetry.record({ action: 'signIn' })
+                            await vscode.commands.executeCommand(cmd)
+                        } else {
+                            telemetry.record({ action: 'dismiss' })
+                        }
+
+                        await globals.context.globalState.update(this.dismissKey, true)
+                    })
+                })
+        })
+    }
+
+    static #instance: SessionSeparationPrompt
+    public static get instance() {
+        return (this.#instance ??= new SessionSeparationPrompt())
+    }
+}

packages/core/src/codecatalyst/activation.ts

@@ -26,6 +26,7 @@ import { DevEnvActivityStarter } from './devEnv'
 import { learnMoreCommand, onboardCommand, reauth } from './explorer'
 import { isInDevEnv } from '../shared/vscode/env'
 import { hasExactScopes } from '../auth/connection'
+import { SessionSeparationPrompt } from '../auth/auth'
 
 const localize = nls.loadMessageBundle()
 
@@ -49,6 +50,7 @@ export async function activate(ctx: ExtContext): Promise<void> {
     // TODO: Remove after some time?
     if (authProvider.isConnected() && !hasExactScopes(authProvider.activeConnection!, defaultScopes)) {
         await authProvider.secondaryAuth.forgetConnection()
+        await SessionSeparationPrompt.instance.showForCommand('aws.codecatalyst.manageConnections')
     }
 
     ctx.extensionContext.subscriptions.push(

packages/core/src/extension.ts

@@ -48,7 +48,7 @@ import globals from './shared/extensionGlobals'
 import { Experiments, Settings, showSettingsFailedMsg } from './shared/settings'
 import { isReleaseVersion } from './shared/vscode/env'
 import { telemetry } from './shared/telemetry/telemetry'
-import { Auth } from './auth/auth'
+import { Auth, SessionSeparationPrompt } from './auth/auth'
 import { registerSubmitFeedback } from './feedback/vue/submitFeedback'
 import { activateCommon, deactivateCommon, emitUserState } from './extensionCommon'
 import { learnMoreAmazonQCommand, qExtensionPageCommand, dismissQTree } from './amazonq/explorer/amazonQChildrenNodes'
@@ -129,9 +129,14 @@ export async function activate(context: vscode.ExtensionContext) {
         for (const conn of await Auth.instance.listConnections()) {
             if (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes)) {
                 await Auth.instance.forgetConnection(conn)
+                await SessionSeparationPrompt.instance.showForCommand('aws.toolkit.auth.manageConnections')
             }
         }
 
+        // Display last prompt if connections were forgotten in prior sessions
+        // but the user did not interact or sign in again. Useful in case the user misses it the first time.
+        await SessionSeparationPrompt.instance.showAnyPreviousPrompt()
+
         await activateCloudFormationTemplateRegistry(context)
 
         // MUST restore CW/Q auth so that we can see if this user is already a Q user.

packages/core/src/test/techdebt.test.ts

@@ -10,6 +10,12 @@ import * as env from '../shared/vscode/env'
 // Checks project config and dependencies, to remind us to remove old things
 // when possible.
 describe('tech debt', function () {
+    function fixByDate(date: string, msg: string) {
+        const now = Date.now()
+        const cutoffDate = Date.parse(date)
+        assert.ok(now <= cutoffDate, msg)
+    }
+
     it('vscode minimum version', async function () {
         const minVscode = env.getMinVscodeVersion()
 
@@ -35,8 +41,12 @@ describe('tech debt', function () {
     })
 
     it('remove missing Amazon Q scopes edge case handling', async function () {
-        const now = Date.now()
-        const cutoffDate = Date.parse('2024-06-30')
-        assert.ok(now <= cutoffDate, 'Remove the edge case code from the commit that this test is a part of.')
+        fixByDate('2024-06-30', 'Remove the edge case code from the commit that this test is a part of.')
+    })
+
+    it('remove separate sessions login edge cases', async function () {
+        // src/auth/auth.ts:SessionSeparationPrompt
+        // forgetConnection() function and calls
+        fixByDate('2024-07-30', 'Remove the edge case code from the commit that this test is a part of.')
     })
 })