bug: Amazon Q sign in blocked by stale connection (#5379)

Problem:

Previously the AWS Toolkit and Amazon Q extensions shared Auth connections.
To do this they communicated their Auth connections to each other, with each
extension storing this in their auth connection cache.

Later we decided to stop sharing the Auth connections between the extensions,
instead having each extension create its own connection.

The problem is that when we created this separate auth connection change, we
had stale connections remaining in Amazon Q. This caused issues like users being
unable to sign in to a new connection since the cached connections showed
it already existed.

Solution:

On every startup clear the stale connections.
We have existing code that does this, but it only ran in a CodeCatalyst dev environment.
But it looks like we always want to run this check as it seems there are other cases
where this happens.

Signed-off-by: Nikolas Komonen <[email protected]>

Author: nkomonen-amazon

packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts

@@ -278,6 +278,26 @@ describe('AuthUtil', async function () {
         assert.strictEqual(authUtil.reformatStartUrl(expected + '/#/'), expected)
         assert.strictEqual(authUtil.reformatStartUrl(expected + '####'), expected)
     })
+
+    it(`clearExtraConnections()`, async function () {
+        const conn1 = await auth.createConnection(createBuilderIdProfile())
+        const conn2 = await auth.createConnection(createSsoProfile({ startUrl: enterpriseSsoStartUrl }))
+        const conn3 = await auth.createConnection(createSsoProfile({ startUrl: enterpriseSsoStartUrl + 1 }))
+        // validate listConnections shows all connections
+        assert.deepStrictEqual(
+            (await authUtil.auth.listConnections()).map((conn) => conn.id).sort((a, b) => a.localeCompare(b)),
+            [conn1, conn2, conn3].map((conn) => conn.id).sort((a, b) => a.localeCompare(b))
+        )
+        await authUtil.secondaryAuth.useNewConnection(conn3)
+
+        await authUtil.clearExtraConnections() // method under test
+
+        // Only the conn that AuthUtil is using is remaining
+        assert.deepStrictEqual(
+            (await authUtil.auth.listConnections()).map((conn) => conn.id),
+            [conn3.id]
+        )
+    })
 })
 
 describe('getChatAuthState()', function () {

packages/core/src/codewhisperer/activation.ts

@@ -70,7 +70,6 @@ import { securityScanLanguageContext } from './util/securityScanLanguageContext'
 import { registerWebviewErrorHandler } from '../webviews/server'
 import { logAndShowError, logAndShowWebviewError } from '../shared/utilities/logAndShowUtils'
 import { openSettings } from '../shared/settings'
-import { getCodeCatalystDevEnvId } from '../shared/vscode/env'
 
 let localize: nls.LocalizeFunc
 
@@ -296,16 +295,7 @@ export async function activate(context: ExtContext): Promise<void> {
     )
 
     await auth.restore()
-
-    // Amazon Q may have code catalyst only credentials stored because it used to import the credentials stored on disk in the environment.
-    if (getCodeCatalystDevEnvId() !== undefined) {
-        for (const conn of await auth.auth.listConnections()) {
-            if (conn.id !== auth.conn?.id) {
-                getLogger().debug('forgetting extra amazon q connection in CoCa dev env: %O', conn)
-                await auth.auth.forgetConnection(conn)
-            }
-        }
-    }
+    await auth.clearExtraConnections()
 
     if (auth.isConnectionExpired()) {
         auth.showReauthenticatePrompt().catch((e) => {

packages/core/src/codewhisperer/util/authUtil.ts

@@ -34,6 +34,7 @@ import { indent } from '../../shared/utilities/textUtilities'
 import { showReauthenticateMessage } from '../../shared/utilities/messages'
 import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthrough'
 import { setContext } from '../../shared/vscode/setContext'
+import { isInDevEnv } from '../../shared/vscode/env'
 
 /** Backwards compatibility for connections w pre-chat scopes */
 export const codeWhispererCoreScopes = [...scopesCodeWhispererCore]
@@ -409,6 +410,23 @@ export class AuthUtil {
 
         return state
     }
+
+    /**
+     * Edge Case: Due to a change in behaviour/functionality, there are potential extra
+     * auth connections that the Amazon Q extension has cached. We need to remove these
+     * as they are irrelevant to the Q extension and can cause issues.
+     */
+    public async clearExtraConnections(): Promise<void> {
+        const currentQConn = this.conn
+        // Q currently only maintains 1 connection at a time, so we assume everything else is extra.
+        // IMPORTANT: In the case Q starts to manage multiple connections, this implementation will need to be updated.
+        const allOtherConnections = (await this.auth.listConnections()).filter((c) => c.id !== currentQConn?.id)
+        for (const conn of allOtherConnections) {
+            getLogger().warn(`forgetting extra amazon q connection: %O`, conn)
+            // in a Dev Env the connection may be used by code catalyst, so we forget instead of fully deleting
+            isInDevEnv() ? await this.auth.forgetConnection(conn) : await this.auth.deleteConnection(conn)
+        }
+    }
 }
 
 /**