telemetry: Amazon Q indicate if it can auto connect (aws#4846)

* rename: findUsableConnection()

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

* telemetry: Amazon Q indicate if it can auto connect

In the telemetry event auth_userState we need to indicate
our authStatus field as 'connected' if the Amazon Q extension does
not have a connection, but we detect one exists already in the
AWS Toolkit extension.

This commit adds extra logic which checks if an autoconnectable
connection can be found in AWS Toolkit and then indicates that we
are connected. This assumes that the separate Login auto-connect logic
will eventually do the auto connection.

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

---------

Signed-off-by: Nikolas Komonen <[email protected]>
Co-authored-by: Maxim Hayes <[email protected]>

Author: nkomonen-amazon

packages/amazonq/src/auth/util.ts

@@ -0,0 +1,61 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import { AuthStatus } from '../../../core/dist/src/shared/telemetry/telemetry'
+import { AwsConnection } from 'aws-core-vscode/auth'
+import { activateExtension, getLogger } from 'aws-core-vscode/shared'
+import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
+
+/** Provides the status of the Auth connection for Amazon Q, specifically for telemetry purposes. */
+export async function getAuthStatus(): Promise<AuthStatus> {
+    // Get auth state from the Amazon Q extension
+    const authState = (await AuthUtil.instance.getChatAuthState()).codewhispererChat
+    let authStatus: AuthStatus = authState === 'connected' || authState === 'expired' ? authState : 'notConnected'
+
+    // If the Q extension does not have its own connection, it will fallback and check
+    // if the Toolkit extension can provide a connection that works with Q
+    if (authStatus === 'notConnected') {
+        let autoConnectConn: AwsConnection | undefined = undefined
+        try {
+            autoConnectConn = await getAutoConnectableConnection()
+        } catch (e) {
+            getLogger().error(`Failed ${getAutoConnectableConnection.name}:\n\n%s`, JSON.stringify(e))
+        }
+
+        // Determine the status of the Toolkit connection we will autoconnect to
+        if (autoConnectConn) {
+            authStatus = autoConnectConn.state === 'valid' ? 'connected' : 'expired'
+        }
+    }
+
+    return authStatus
+}
+
+/**
+ * Returns a connection from the standalone Toolkit extension that
+ * the Amazon Q extension can use. Otherwise it returns undefined.
+ *
+ * HACK: Our frontend Login ui for Amazon Q will auto connect if required/possible,
+ * but we cannot detect this at the end of Amazon Q extension activation.
+ * So we reuse the same {@link findUsableQConnection}()
+ * and assume that the frontend will have the same result and auto connect.
+ */
+async function getAutoConnectableConnection(): Promise<AwsConnection | undefined> {
+    const extension = await activateExtension<any>(VSCODE_EXTENSION_ID.awstoolkit)
+    if (!extension) {
+        return undefined
+    }
+    const importedApis = extension.exports.getApi(VSCODE_EXTENSION_ID.awstoolkit)
+
+    const listConnections: () => Promise<AwsConnection[]> = importedApis?.listConnections
+    if (!listConnections) {
+        // Either the user has an older toolkit version w/o the API, or the API has changed
+        // and this needs to be updated.
+        return undefined
+    }
+
+    return AuthUtil.instance.findUsableQConnection(await listConnections())
+}

packages/amazonq/src/extensionShared.ts

@@ -11,7 +11,6 @@ import {
     activate as activateCodeWhisperer,
     shutdown as codewhispererShutdown,
     amazonQDismissedKey,
-    AuthUtil,
 } from 'aws-core-vscode/codewhisperer'
 import {
     ExtContext,
@@ -36,6 +35,7 @@ import { isExtensionActive, VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
 import { registerSubmitFeedback } from 'aws-core-vscode/feedback'
 import { telemetry, ExtStartUpSources } from 'aws-core-vscode/telemetry'
 import { DevFunction, updateDevMode } from 'aws-core-vscode/dev'
+import { getAuthStatus } from './auth/util'
 
 export async function activateShared(context: vscode.ExtensionContext, isWeb: boolean) {
     initialize(context, isWeb)
@@ -156,16 +156,18 @@ export async function activateShared(context: vscode.ExtensionContext, isWeb: bo
             telemetry.record({ source: ExtStartUpSources.reload })
         }
 
-        const authState = (await AuthUtil.instance.getChatAuthState()).codewhispererChat
         const authKinds: AuthUtils.AuthSimpleId[] = []
         if (await AuthUtils.hasBuilderId('codewhisperer')) {
             authKinds.push('builderIdCodeWhisperer')
         }
         if (await AuthUtils.hasSso('codewhisperer')) {
             authKinds.push('identityCenterCodeWhisperer')
         }
+
+        const authStatus = await getAuthStatus()
+
         telemetry.record({
-            authStatus: authState === 'connected' || authState === 'expired' ? authState : 'notConnected',
+            authStatus,
             authEnabledConnections: authKinds.join(','),
         })
     })

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

@@ -429,6 +429,21 @@ export class AuthUtil {
 
         return state
     }
+
+    /**
+     * From the given connections, returns a connection that works with Amazon Q.
+     */
+    findUsableQConnection(connections: AwsConnection[]): AwsConnection | undefined {
+        const hasQScopes = (c: AwsConnection) => amazonQScopes.every(s => c.scopes?.includes(s))
+        const score = (c: AwsConnection) => Number(hasQScopes(c)) * 10 + Number(c.state === 'valid')
+        connections.sort(function (a, b) {
+            return score(b) - score(a)
+        })
+        if (hasQScopes(connections[0])) {
+            return connections[0]
+        }
+        return undefined
+    }
 }
 
 /**

packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts

@@ -54,16 +54,8 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
      * @param connections List of AWS Toolkit Connections
      * @returns Amazon Q connection, or undefined if none of the given connections have scopes required for Amazon Q.
      */
-    findConnection(connections: AwsConnection[]): AwsConnection | undefined {
-        const hasQScopes = (c: AwsConnection) => amazonQScopes.every(s => c.scopes?.includes(s))
-        const score = (c: AwsConnection) => Number(hasQScopes(c)) * 10 + Number(c.state === 'valid')
-        connections.sort(function (a, b) {
-            return score(b) - score(a)
-        })
-        if (hasQScopes(connections[0])) {
-            return connections[0]
-        }
-        return undefined
+    findUsableConnection(connections: AwsConnection[]): AwsConnection | undefined {
+        return AuthUtil.instance.findUsableQConnection(connections)
     }
 
     async useConnection(connectionId: string, auto: boolean): Promise<AuthError | undefined> {

packages/core/src/login/webview/vue/backend.ts

@@ -164,7 +164,7 @@ export abstract class CommonAuthWebview extends VueWebview {
      */
     abstract useConnection(connectionId: string, auto: boolean): Promise<AuthError | undefined>
 
-    abstract findConnection(connections: AwsConnection[]): AwsConnection | undefined
+    abstract findUsableConnection(connections: AwsConnection[]): AwsConnection | undefined
 
     async errorNotification(e: AuthError) {
         void vscode.window.showInformationMessage(`${e.text}`)

packages/core/src/login/webview/vue/login.vue

@@ -492,7 +492,7 @@ export default defineComponent({
             // If Amazon Q has no connections while Toolkit has connections
             // Auto connect Q using toolkit connection.
             if (connections.length === 0 && sharedConnections && sharedConnections.length > 0) {
-                const conn = await client.findConnection(sharedConnections)
+                const conn = await client.findUsableConnection(sharedConnections)
                 if (conn) {
                     await client.useConnection(conn.id, true)
                 }

packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts

@@ -146,7 +146,7 @@ export class ToolkitLoginWebview extends CommonAuthWebview {
         })
     }
 
-    findConnection(connections: AwsConnection[]): AwsConnection | undefined {
+    findUsableConnection(connections: AwsConnection[]): AwsConnection | undefined {
         return undefined
     }
 

packages/core/src/shared/index.ts

@@ -15,3 +15,4 @@ export { RegionProvider } from './regions/regionProvider'
 export { Commands } from './vscode/commands2'
 export { getMachineId } from './vscode/env'
 export { getLogger } from './logger/logger'
+export { activateExtension } from './utilities/vsCodeUtils'