feat(amazonq): Use new auth class using flare identity server for Amazon Q

docs/lsp.md

@@ -48,7 +48,6 @@ sequenceDiagram
 3. Enable the lsp experiment:
     ```
     "aws.experiments": {
-      "amazonqLSP": true,
       "amazonqLSPInline": true, // optional: enables inline completion from flare
       "amazonqLSPChat": true // optional: enables chat from flare
     }

packages/amazonq/src/api.ts

@@ -8,6 +8,7 @@ import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseReques
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { ChatSession } from 'aws-core-vscode/codewhispererChat'
 import { api } from 'aws-core-vscode/amazonq'
+import { getLogger } from 'aws-core-vscode/shared'
 
 export default {
     chatApi: {
@@ -26,8 +27,25 @@ export default {
                 await AuthUtil.instance.showReauthenticatePrompt()
             }
         },
+        /**
+         * @deprecated use getAuthState() instead
+         *
+         * Legacy function for callers who expect auth state to be granular amongst Q features.
+         * Auth state is consistent between features, so getAuthState() can be consumed safely for all features.
+         *
+         */
         async getChatAuthState() {
-            return AuthUtil.instance.getChatAuthState()
+            getLogger().warn('Warning: getChatAuthState() is deprecated. Use getAuthState() instead.')
+            const state = AuthUtil.instance.getAuthState()
+            const convertedState = state === 'notConnected' ? 'disconnected' : state
+            return {
+                codewhispererCore: convertedState,
+                codewhispererChat: convertedState,
+                amazonQ: convertedState,
+            }
+        },
+        getAuthState() {
+            return AuthUtil.instance.getAuthState()
         },
     },
 } satisfies api

packages/amazonq/src/app/chat/activation.ts

@@ -27,7 +27,7 @@ export async function activate(context: ExtensionContext) {
 
     const setupLsp = funcUtil.debounce(async () => {
         void amazonq.LspController.instance.trySetupLsp(context, {
-            startUrl: AuthUtil.instance.startUrl,
+            startUrl: AuthUtil.instance.connection?.startUrl,
             maxIndexSize: CodeWhispererSettings.instance.getMaxIndexSize(),
             isVectorIndexEnabled: CodeWhispererSettings.instance.isLocalIndexEnabled(),
         })
@@ -44,8 +44,6 @@ export async function activate(context: ExtensionContext) {
         amazonq.walkthroughSecurityScanExample.register(),
         amazonq.openAmazonQWalkthrough.register(),
         amazonq.listCodeWhispererCommandsWalkthrough.register(),
-        amazonq.focusAmazonQPanel.register(),
-        amazonq.focusAmazonQPanelKeybinding.register(),
         amazonq.tryChatCodeLensCommand.register(),
         vscode.workspace.onDidChangeConfiguration(async (configurationChangeEvent) => {
             if (configurationChangeEvent.affectsConfiguration('amazonQ.workspaceIndex')) {

packages/amazonq/src/extension.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
+import { CredentialsStore, LoginManager, authUtils, initializeAuth } from 'aws-core-vscode/auth'
 import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer'
 import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
 import { CommonAuthWebview } from 'aws-core-vscode/login'
@@ -33,6 +33,7 @@ import {
     maybeShowMinVscodeWarning,
     Experiments,
     isSageMaker,
+    Commands,
 } from 'aws-core-vscode/shared'
 import { ExtStartUpSources } from 'aws-core-vscode/telemetry'
 import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
@@ -117,11 +118,11 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     const extContext = {
         extensionContext: context,
     }
+
+    await activateAmazonqLsp(context)
+
     // This contains every lsp agnostic things (auth, security scan, code scan)
     await activateCodeWhisperer(extContext as ExtContext)
-    if (Experiments.instance.get('amazonqLSP', false)) {
-        await activateAmazonqLsp(context)
-    }
 
     if (!Experiments.instance.get('amazonqLSPInline', false)) {
         await activateInlineCompletion()
@@ -130,6 +131,10 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     // Generic extension commands
     registerGenericCommands(context, amazonQContextPrefix)
 
+    // Create status bar and reference log UI elements
+    void Commands.tryExecute('aws.amazonq.refreshStatusBar')
+    void Commands.tryExecute('aws.amazonq.updateReferenceLog')
+
     // Amazon Q specific commands
     registerCommands(context)
 
@@ -156,7 +161,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     // reload webviews
     await vscode.commands.executeCommand('workbench.action.webview.reloadWebviewAction')
 
-    if (AuthUtils.ExtensionUse.instance.isFirstUse()) {
+    if (authUtils.ExtensionUse.instance.isFirstUse()) {
         // Give time for the extension to finish initializing.
         globals.clock.setTimeout(async () => {
             CommonAuthWebview.authSource = ExtStartUpSources.firstStartUp
@@ -166,7 +171,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
 
     context.subscriptions.push(
         Experiments.instance.onDidChange(async (event) => {
-            if (event.key === 'amazonqLSP' || event.key === 'amazonqChatLSP' || event.key === 'amazonqLSPInline') {
+            if (event.key === 'amazonqChatLSP' || event.key === 'amazonqLSPInline') {
                 await vscode.window
                     .showInformationMessage(
                         'Amazon Q LSP setting has changed. Reload VS Code for the changes to take effect.',

packages/amazonq/src/extensionNode.ts

@@ -7,29 +7,22 @@ import * as vscode from 'vscode'
 import { activateAmazonQCommon, amazonQContextPrefix, deactivateCommon } from './extension'
 import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
 import { activate as activateQGumby } from 'aws-core-vscode/amazonqGumby'
-import {
-    ExtContext,
-    globals,
-    CrashMonitoring,
-    getLogger,
-    isNetworkError,
-    isSageMaker,
-    Experiments,
-} from 'aws-core-vscode/shared'
+import { ExtContext, globals, CrashMonitoring /* getLogger, isSageMaker,*/, Experiments } from 'aws-core-vscode/shared'
 import { filetypes, SchemaService } from 'aws-core-vscode/sharedNode'
 import { updateDevMode } from 'aws-core-vscode/dev'
 import { CommonAuthViewProvider } from 'aws-core-vscode/login'
 import { isExtensionActive, VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
 import { registerSubmitFeedback } from 'aws-core-vscode/feedback'
 import { DevOptions } from 'aws-core-vscode/dev'
-import { Auth, AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection } from 'aws-core-vscode/auth'
+import { Auth /* , AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection*/ } from 'aws-core-vscode/auth'
 import api from './api'
 import { activate as activateCWChat } from './app/chat/activation'
 import { activate as activateInlineChat } from './inlineChat/activation'
 import { beta } from 'aws-core-vscode/dev'
-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'
+import * as amazonq from 'aws-core-vscode/amazonq'
+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'
 
 export async function activate(context: vscode.ExtensionContext) {
     // IMPORTANT: No other code should be added to this function. Place it in one of the following 2 functions where appropriate.
@@ -58,6 +51,8 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
     }
     activateInlineChat(context)
 
+    context.subscriptions.push(amazonq.focusAmazonQPanel.register(), amazonq.focusAmazonQPanelKeybinding.register())
+
     const authProvider = new CommonAuthViewProvider(
         context,
         amazonQContextPrefix,
@@ -78,49 +73,34 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
     await setupDevMode(context)
     await beta.activate(context)
 
-    // TODO: Should probably emit for web as well.
+    // TODO: @opieter Fix telemetry
     // Will the web metric look the same?
-    telemetry.auth_userState.emit({
-        passive: true,
-        result: 'Succeeded',
-        source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(),
-        ...(await getAuthState()),
-    })
-
-    void activateNotifications(context, getAuthState)
+    // telemetry.auth_userState.emit({
+    //     passive: true,
+    //     result: 'Succeeded',
+    //     source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(),
+    //     ...(await getAuthState()),
+    // })
+
+    // void activateNotifications(context, getAuthState)
 }
 
-async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
-    let authState: AuthState = 'disconnected'
-    try {
-        // May call connection validate functions that try to refresh the token.
-        // This could result in network errors.
-        authState = (await AuthUtil.instance._getChatAuthState(false)).codewhispererChat
-    } catch (err) {
-        if (
-            isNetworkError(err) &&
-            AuthUtil.instance.conn &&
-            AuthUtil.instance.auth.getConnectionState(AuthUtil.instance.conn) === 'valid'
-        ) {
-            authState = 'connectedWithNetworkError'
-        } else {
-            throw err
-        }
-    }
-    const currConn = AuthUtil.instance.conn
-    if (currConn !== undefined && !(isAnySsoConnection(currConn) || isSageMaker())) {
-        getLogger().error(`Current Amazon Q connection is not SSO, type is: %s`, currConn?.type)
-    }
-
-    return {
-        authStatus:
-            authState === 'connected' || authState === 'expired' || authState === 'connectedWithNetworkError'
-                ? authState
-                : 'notConnected',
-        authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','),
-        ...(await getTelemetryMetadataForConn(currConn)),
-    }
-}
+// async function getAuthState(): Promise<Omit<auth2.AuthState, 'source'>> {
+//     let authState: AuthState = 'disconnected'
+//     authState = AuthUtil.instance.getAuthState()
+
+//     if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
+//         getLogger().error('Current Amazon Q connection is not SSO')
+//     }
+
+//     return {
+//         authStatus:
+//             authState === 'connected' || authState === 'expired'
+//                 ? authState
+//                 : 'notConnected',
+//         ...(await getTelemetryMetadataForConn(currConn)),
+//     }
+// }
 
 /**
  * Some parts of this do not work in Web mode so we need to set Dev Mode up here.

packages/amazonq/src/lsp/activation.ts

@@ -4,16 +4,20 @@
  */
 
 import vscode from 'vscode'
-import { startLanguageServer } from './client'
+import { clientId, encryptionKey, startLanguageServer } from './client'
 import { AmazonQLspInstaller } from './lspInstaller'
-import { lspSetupStage, ToolkitError, messages } from 'aws-core-vscode/shared'
+import { lspSetupStage, ToolkitError } from 'aws-core-vscode/shared'
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import { auth2 } from 'aws-core-vscode/auth'
 
-export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
+export async function activate(ctx: vscode.ExtensionContext) {
     try {
-        await lspSetupStage('all', async () => {
+        const client = await lspSetupStage('all', async () => {
             const installResult = await new AmazonQLspInstaller().resolve()
-            await lspSetupStage('launch', async () => await startLanguageServer(ctx, installResult.resourcePaths))
+            return await lspSetupStage('launch', () => startLanguageServer(ctx, installResult.resourcePaths))
         })
+        AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
+        await AuthUtil.instance.restore()
     } catch (err) {
         const e = err as ToolkitError
         void messages.showViewLogsMessage(`Failed to launch Amazon Q language server: ${e.message}`)