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

Author: opieter-aws

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'
@@ -114,14 +115,14 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
 
     await initializeAuth(globals.loginManager)
 
+    await activateAmazonqLsp(context)
+
     const extContext = {
         extensionContext: 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/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 } 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 vscode.window.showInformationMessage(`Unable to launch amazonq language server: ${e.message}`)

packages/amazonq/src/lsp/client.ts

@@ -6,21 +6,46 @@
 import vscode, { env, version } from 'vscode'
 import * as nls from 'vscode-nls'
 import * as crypto from 'crypto'
+import * as jose from 'jose'
 import { LanguageClient, LanguageClientOptions } from 'vscode-languageclient'
 import { InlineCompletionManager } from '../app/inline/completion'
-import { AmazonQLspAuth, encryptionKey, notificationTypes } from './auth'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
-import { ConnectionMetadata } from '@aws/language-server-runtimes/protocol'
-import { Settings, oidcClientName, createServerOptions, globals, Experiments, Commands } from 'aws-core-vscode/shared'
+import {
+    ConnectionMetadata,
+    GetSsoTokenProgress,
+    GetSsoTokenProgressToken,
+    GetSsoTokenProgressType,
+    MessageActionItem,
+    ShowDocumentParams,
+    ShowDocumentRequest,
+    ShowDocumentResult,
+    ShowMessageRequest,
+    ShowMessageRequestParams,
+} from '@aws/language-server-runtimes/protocol'
+import {
+    Settings,
+    oidcClientName,
+    createServerOptions,
+    globals,
+    Experiments,
+    Commands,
+    openUrl,
+    getLogger,
+} from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
+import { auth2 } from 'aws-core-vscode/auth'
 
 const localize = nls.loadMessageBundle()
 
+export const clientId = 'amazonq'
+export const clientName = oidcClientName()
+export const encryptionKey = crypto.randomBytes(32)
+
 export async function startLanguageServer(
     extensionContext: vscode.ExtensionContext,
     resourcePaths: AmazonQResourcePaths
-) {
+): Promise<LanguageClient> {
     const toDispose = extensionContext.subscriptions
 
     const serverModule = resourcePaths.lsp
@@ -39,8 +64,6 @@ export async function startLanguageServer(
     })
 
     const documentSelector = [{ scheme: 'file', language: '*' }]
-
-    const clientId = 'amazonq'
     const traceServerEnabled = Settings.instance.isSet(`${clientId}.trace.server`)
 
     // Options to control the language client
@@ -80,57 +103,89 @@ export async function startLanguageServer(
               }),
     }
 
-    const client = new LanguageClient(
-        clientId,
-        localize('amazonq.server.name', 'Amazon Q Language Server'),
-        serverOptions,
-        clientOptions
-    )
+    const lspName = localize('amazonq.server.name', 'Amazon Q Language Server')
+    const client = new LanguageClient(clientId, lspName, serverOptions, clientOptions)
 
     const disposable = client.start()
     toDispose.push(disposable)
 
-    const auth = new AmazonQLspAuth(client)
+    await client.onReady()
 
-    return client.onReady().then(async () => {
-        // Request handler for when the server wants to know about the clients auth connnection. Must be registered before the initial auth init call
-        client.onRequest<ConnectionMetadata, Error>(notificationTypes.getConnectionMetadata.method, () => {
-            return {
-                sso: {
-                    startUrl: AuthUtil.instance.auth.startUrl,
-                },
-            }
-        })
-        await auth.refreshConnection()
-
-        if (Experiments.instance.get('amazonqLSPInline', false)) {
-            const inlineManager = new InlineCompletionManager(client)
-            inlineManager.registerInlineCompletion()
-            toDispose.push(
-                inlineManager,
-                Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
-                    await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
-                }),
-                vscode.workspace.onDidCloseTextDocument(async () => {
-                    await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
-                })
-            )
+    // Request handler for when the server wants to know about the clients auth connnection. Must be registered before the initial auth init call
+    client.onRequest<ConnectionMetadata, Error>(auth2.notificationTypes.getConnectionMetadata.method, () => {
+        return {
+            sso: {
+                startUrl: AuthUtil.instance.connection?.startUrl,
+            },
         }
+    })
 
-        if (Experiments.instance.get('amazonqChatLSP', false)) {
-            activate(client, encryptionKey, resourcePaths.ui)
+    client.onRequest<ShowDocumentResult, Error>(ShowDocumentRequest.method, async (params: ShowDocumentParams) => {
+        try {
+            return { success: await openUrl(vscode.Uri.parse(params.uri), lspName) }
+        } catch (err: any) {
+            getLogger().error(`Failed to open document for LSP: ${lspName}, error: %s`, err)
+            return { success: false }
         }
+    })
+
+    client.onRequest<MessageActionItem | null, Error>(
+        ShowMessageRequest.method,
+        async (params: ShowMessageRequestParams) => {
+            const actions = params.actions?.map((a) => a.title) ?? []
+            const response = await vscode.window.showInformationMessage(params.message, { modal: true }, ...actions)
+            return params.actions?.find((a) => a.title === response) ?? (undefined as unknown as null)
+        }
+    )
 
-        const refreshInterval = auth.startTokenRefreshInterval()
+    let promise: Promise<void> | undefined
+    let resolver: () => void = () => {}
+    client.onProgress(GetSsoTokenProgressType, GetSsoTokenProgressToken, async (partialResult: GetSsoTokenProgress) => {
+        const decryptedKey = await jose.compactDecrypt(partialResult as unknown as string, encryptionKey)
+        const val: GetSsoTokenProgress = JSON.parse(decryptedKey.plaintext.toString())
 
+        if (val.state === 'InProgress') {
+            if (promise) {
+                resolver()
+            }
+            promise = new Promise<void>((resolve) => {
+                resolver = resolve
+            })
+        } else {
+            resolver()
+            promise = undefined
+            return
+        }
+
+        void vscode.window.withProgress(
+            {
+                cancellable: true,
+                location: vscode.ProgressLocation.Notification,
+                title: val.message,
+            },
+            async (_) => {
+                await promise
+            }
+        )
+    })
+
+    if (Experiments.instance.get('amazonqLSPInline', false)) {
+        const inlineManager = new InlineCompletionManager(client)
+        inlineManager.registerInlineCompletion()
         toDispose.push(
-            AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
-                await auth.refreshConnection()
+            inlineManager,
+            Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
             }),
-            AuthUtil.instance.auth.onDidDeleteConnection(async () => {
-                client.sendNotification(notificationTypes.deleteBearerToken.method)
-            }),
-            { dispose: () => clearInterval(refreshInterval) }
+            vscode.workspace.onDidCloseTextDocument(async () => {
+                await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
+            })
         )
-    })
+    }
+
+    if (Experiments.instance.get('amazonqChatLSP', false)) {
+        activate(client, encryptionKey, resourcePaths.ui)
+    }
+
+    return client
 }

packages/core/src/amazonq/extApi.ts

@@ -7,9 +7,14 @@ import vscode from 'vscode'
 import { VSCODE_EXTENSION_ID } from '../shared/extensions'
 import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
 import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
-import { FeatureAuthState } from '../codewhisperer/util/authUtil'
+import { auth2 } from 'aws-core-vscode/auth'
 import { ToolkitError } from '../shared/errors'
 
+/**
+ * @deprecated, for backwards comaptibility only.
+ */
+type OldAuthState = 'disconnected' | 'expired' | 'connected'
+
 /**
  * This interface is used and exported by the amazon q extension. If you make a change here then
  * update the corresponding api implementation in packages/amazonq/src/api.ts
@@ -21,7 +26,15 @@ export interface api {
     }
     authApi: {
         reauthIfNeeded(): Promise<void>
-        getChatAuthState(): Promise<FeatureAuthState>
+        /**
+         * @deprecated, for backwards comaptibility only.
+         */
+        getChatAuthState(): Promise<{
+            codewhispererCore: OldAuthState
+            codewhispererChat: OldAuthState
+            amazonQ: OldAuthState
+        }>
+        getAuthState(): auth2.AuthState
     }
 }