Update references and fix unit tests

Author: opieter-aws

packages/amazonq/src/app/amazonqScan/chat/controller/controller.ts

@@ -33,7 +33,6 @@ import {
 } from '../../models/constants'
 import path from 'path'
 import { telemetry } from 'aws-core-vscode/telemetry'
-import { auth2 } from 'aws-core-vscode/auth'
 
 export class ScanController {
     private readonly messenger: Messenger

packages/amazonq/src/extensionNode.ts

@@ -78,49 +78,49 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
     await setupDevMode(context)
     await beta.activate(context)
 
+    // TODO: @opieter add telemetry
+
     // TODO: Should probably emit for web as well.
     // 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<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 = AuthUtil.instance.getAuthState()
+//     } catch (err) {
+//         if (
+//             isNetworkError(err) &&
+//             AuthUtil.instance.isConnected()
+//         ) {
+//             authState = 'connectedWithNetworkError'
+//         } else {
+//             throw err
+//         }
+//     }
+//     if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
+//         getLogger().error('Current Amazon Q connection is not SSO')
+//     }
+
+//     return {
+//         authStatus:
+//             authState === 'connected' || authState === 'expired' || authState === 'connectedWithNetworkError'
+//                 ? authState
+//                 : 'notConnected',
+//         authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','),
+//         ...(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/inlineChat/provider/inlineChatProvider.ts

@@ -25,7 +25,6 @@ import { codeWhispererClient } from 'aws-core-vscode/codewhisperer'
 import type { InlineChatEvent } from 'aws-core-vscode/codewhisperer'
 import { InlineTask } from '../controller/inlineTask'
 import { extractAuthFollowUp } from 'aws-core-vscode/amazonq'
-import { auth2 } from 'aws-core-vscode/auth'
 
 export class InlineChatProvider {
     private readonly editorContextExtractor: EditorContextExtractor

packages/amazonq/src/lsp/client.ts

@@ -8,15 +8,17 @@ import * as nls from 'vscode-nls'
 import * as crypto from 'crypto'
 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 { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
+import { auth2 } from 'aws-core-vscode/auth'
 
 const localize = nls.loadMessageBundle()
 
+export const encryptionKey = crypto.randomBytes(32)
+
 export async function startLanguageServer(
     extensionContext: vscode.ExtensionContext,
     resourcePaths: AmazonQResourcePaths
@@ -90,18 +92,15 @@ export async function startLanguageServer(
     const disposable = client.start()
     toDispose.push(disposable)
 
-    const auth = new AmazonQLspAuth(client)
-
     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, () => {
+        client.onRequest<ConnectionMetadata, Error>(auth2.notificationTypes.getConnectionMetadata.method, () => {
             return {
                 sso: {
-                    startUrl: AuthUtil.instance.auth.startUrl,
+                    startUrl: AuthUtil.instance.connection?.startUrl,
                 },
             }
         })
-        await auth.refreshConnection()
 
         if (Experiments.instance.get('amazonqLSPInline', false)) {
             const inlineManager = new InlineCompletionManager(client)
@@ -120,17 +119,5 @@ export async function startLanguageServer(
         if (Experiments.instance.get('amazonqChatLSP', false)) {
             activate(client, encryptionKey, resourcePaths.ui)
         }
-
-        const refreshInterval = auth.startTokenRefreshInterval()
-
-        toDispose.push(
-            AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
-                await auth.refreshConnection()
-            }),
-            AuthUtil.instance.auth.onDidDeleteConnection(async () => {
-                client.sendNotification(notificationTypes.deleteBearerToken.method)
-            }),
-            { dispose: () => clearInterval(refreshInterval) }
-        )
     })
 }

packages/amazonq/test/e2e/amazonq/utils/setup.ts

@@ -5,9 +5,9 @@
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 
 export async function loginToIdC() {
-    const authState = await AuthUtil.instance.getChatAuthState()
+    const authState = AuthUtil.instance.getAuthState()
     if (process.env['AWS_TOOLKIT_AUTOMATION'] === 'local') {
-        if (authState.amazonQ !== 'connected') {
+        if (authState !== 'connected') {
             throw new Error('You will need to login manually before running tests.')
         }
         return
@@ -22,5 +22,5 @@ export async function loginToIdC() {
         )
     }
 
-    await AuthUtil.instance.connectToEnterpriseSso(startUrl, region)
+    await AuthUtil.instance.login(startUrl, region)
 }