feat(amazonq lsp): add authFollowUpClicked implementation (#6886)

## Problem
There was no implementation for the `authFollowUpClicked` message
received from LSP for Amazon Q. Because of that, the `authenticate`
button that appears in case of expired SSO connection had no effect.

## Solution
I added the implementation as follows:
* If the received `authFollowUpType` is `'re-auth'` or
`'missing_scopes'`, initiate a re-authentication
* If the received `authFollowUpType` is `'full-auth'` or
`'use-supported-auth'`, delete the connection so a new full
authentication flow kicks off

I added a unit test file for this message type, which is extendable for
every command handled in `registerLanguageServerEventListener`

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: opieter-aws

packages/amazonq/src/lsp/chat/messages.ts

@@ -9,6 +9,7 @@ import {
     AUTH_FOLLOW_UP_CLICKED,
     CHAT_OPTIONS,
     COPY_TO_CLIPBOARD,
+    AuthFollowUpType,
 } from '@aws/chat-client-ui-types'
 import {
     ChatResult,
@@ -25,6 +26,7 @@ import { window } from 'vscode'
 import { Disposable, LanguageClient, Position, State, TextDocumentIdentifier } from 'vscode-languageclient'
 import * as jose from 'jose'
 import { AmazonQChatViewProvider } from './webviewProvider'
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 
 export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) {
     languageClient.onDidChangeState(({ oldState, newState }) => {
@@ -77,10 +79,33 @@ export function registerMessageListeners(
                 })
                 break
             }
-            case AUTH_FOLLOW_UP_CLICKED:
-                // TODO hook this into auth
+            case AUTH_FOLLOW_UP_CLICKED: {
                 languageClient.info('[VSCode Client] AuthFollowUp clicked')
+                const authType = message.params.authFollowupType
+                const reAuthTypes: AuthFollowUpType[] = ['re-auth', 'missing_scopes']
+                const fullAuthTypes: AuthFollowUpType[] = ['full-auth', 'use-supported-auth']
+
+                if (reAuthTypes.includes(authType)) {
+                    try {
+                        await AuthUtil.instance.reauthenticate()
+                    } catch (e) {
+                        languageClient.error(
+                            `[VSCode Client] Failed to re-authenticate after AUTH_FOLLOW_UP_CLICKED: ${(e as Error).message}`
+                        )
+                    }
+                }
+
+                if (fullAuthTypes.includes(authType)) {
+                    try {
+                        await AuthUtil.instance.secondaryAuth.deleteConnection()
+                    } catch (e) {
+                        languageClient.error(
+                            `[VSCode Client] Failed to authenticate after AUTH_FOLLOW_UP_CLICKED: ${(e as Error).message}`
+                        )
+                    }
+                }
                 break
+            }
             case chatRequestType.method: {
                 const partialResultToken = uuidv4()
                 const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, (partialResult) =>

packages/amazonq/test/unit/amazonq/lsp/chat/messages.test.ts

@@ -0,0 +1,129 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as sinon from 'sinon'
+import { LanguageClient } from 'vscode-languageclient'
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import { registerMessageListeners } from '../../../../../src/lsp/chat/messages'
+import { AmazonQChatViewProvider } from '../../../../../src/lsp/chat/webviewProvider'
+import { secondaryAuth, authConnection, AuthFollowUpType } from 'aws-core-vscode/amazonq'
+
+describe('registerMessageListeners', () => {
+    let languageClient: LanguageClient
+    let provider: AmazonQChatViewProvider
+    let sandbox: sinon.SinonSandbox
+    let messageHandler: (message: any) => void | Promise<void>
+    let errorStub: sinon.SinonStub
+
+    beforeEach(() => {
+        sandbox = sinon.createSandbox()
+        errorStub = sandbox.stub()
+
+        languageClient = {
+            info: sandbox.stub(),
+            error: errorStub,
+            sendNotification: sandbox.stub(),
+        } as unknown as LanguageClient
+
+        provider = {
+            webview: {
+                onDidReceiveMessage: (callback: (message: any) => void | Promise<void>) => {
+                    messageHandler = callback
+                    return {
+                        dispose: (): void => {},
+                    }
+                },
+            },
+        } as any
+
+        registerMessageListeners(languageClient, provider, Buffer.from('test-key'))
+    })
+
+    afterEach(() => {
+        sandbox.restore()
+    })
+
+    describe('AUTH_FOLLOW_UP_CLICKED', () => {
+        let mockAuthUtil: AuthUtil
+        let deleteConnectionStub: sinon.SinonStub
+        let reauthenticateStub: sinon.SinonStub
+
+        const authFollowUpClickedCommand = 'authFollowUpClicked'
+
+        interface TestCase {
+            authType: AuthFollowUpType
+            stubToReject: sinon.SinonStub
+            errorMessage: string
+        }
+
+        const testFailure = async (testCase: TestCase) => {
+            testCase.stubToReject.rejects(new Error())
+
+            await messageHandler({
+                command: authFollowUpClickedCommand,
+                params: {
+                    authFollowupType: testCase.authType,
+                },
+            })
+
+            sinon.assert.calledOnce(errorStub)
+            sinon.assert.calledWith(errorStub, sinon.match(testCase.errorMessage))
+        }
+
+        beforeEach(() => {
+            deleteConnectionStub = sandbox.stub().resolves()
+            reauthenticateStub = sandbox.stub().resolves()
+
+            mockAuthUtil = {
+                reauthenticate: reauthenticateStub,
+                secondaryAuth: {
+                    deleteConnection: deleteConnectionStub,
+                } as unknown as secondaryAuth.SecondaryAuth<authConnection.Connection>,
+            } as unknown as AuthUtil
+
+            sandbox.replaceGetter(AuthUtil, 'instance', () => mockAuthUtil)
+        })
+
+        it('handles re-authentication request', async () => {
+            await messageHandler({
+                command: authFollowUpClickedCommand,
+                params: {
+                    authFollowupType: 're-auth',
+                },
+            })
+
+            sinon.assert.calledOnce(reauthenticateStub)
+            sinon.assert.notCalled(deleteConnectionStub)
+        })
+
+        it('handles full authentication request', async () => {
+            await messageHandler({
+                command: authFollowUpClickedCommand,
+                params: {
+                    authFollowupType: 'full-auth',
+                },
+            })
+
+            sinon.assert.notCalled(reauthenticateStub)
+            sinon.assert.calledOnce(deleteConnectionStub)
+        })
+
+        it('logs error if re-authentication fails', async () => {
+            await testFailure({
+                authType: 're-auth',
+                stubToReject: reauthenticateStub,
+                errorMessage: 'Failed to re-authenticate',
+            })
+        })
+
+        it('logs error if full authentication fails', async () => {
+            await testFailure({
+                authType: 'full-auth',
+                stubToReject: deleteConnectionStub,
+                errorMessage: 'Failed to authenticate',
+            })
+        })
+    })
+})

packages/core/src/amazonq/index.ts

@@ -46,6 +46,8 @@ export { extractAuthFollowUp } from './util/authUtils'
 export { Messenger } from './commons/connector/baseMessenger'
 export * from './lsp/config'
 export * as WorkspaceLspInstaller from './lsp/workspaceInstaller'
+export * as secondaryAuth from '../auth/secondaryAuth'
+export * as authConnection from '../auth/connection'
 import { FeatureContext } from '../shared/featureConfig'
 
 /**