feate(amazonq): only display open chat codelens if authenticated (#4955)

* feate(amazonq): only display open chat codelens if authenticated

* fix tests

* use sync method to fix tests

* Update packages/core/src/test/codewhispererChat/editor/codelens.test.ts

Co-authored-by: Nikolas Komonen <[email protected]>

---------

Co-authored-by: Nikolas Komonen <[email protected]>

Author: hayemaxi

packages/amazonq/.changes/next-release/Feature-5bbe5e13-b015-4b01-88a5-14648ba9a35d.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Inline Suggestions: Only display the 'Open Chat' CodeLens if the user is signed into Amazon Q."
+}

packages/core/src/auth/auth.ts

@@ -370,9 +370,12 @@ export class Auth implements AuthService, ConnectionManager {
      * Alternatively you can use the `getToken()` call on an SSO connection to do the same thing,
      * but it will additionally prompt for reauthentication if the connection is invalid.
      */
-    public async refreshConnectionState(connection: Pick<Connection, 'id'>): Promise<undefined> {
-        const profile = this.store.getProfile(connection.id)
+    public async refreshConnectionState(connection?: Pick<Connection, 'id'>): Promise<undefined> {
+        if (connection === undefined) {
+            return
+        }
 
+        const profile = this.store.getProfile(connection.id)
         if (profile === undefined) {
             return
         }

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

@@ -387,28 +387,29 @@ export class AuthUtil {
     }
 
     /**
-     * Returns a snapshot of the overall auth state of CodeWhisperer + Chat features.
-     *
-     * @param shouldRefresh (default true) validate and update the current connection state.
-     * If this setting is set to false, there is a risk that the evaluated state is outdated,
-     * but it is safe from modifying the state of the connection.
+     * Asynchronously returns a snapshot of the overall auth state of CodeWhisperer + Chat features.
+     * It guarantees the latest state is correct at the risk of modifying connection state.
+     * If this guarantee is not required, use sync method getChatAuthStateSync()
      */
-    public async getChatAuthState(shouldRefresh: boolean = true): Promise<FeatureAuthState> {
-        const currentConnection = this.conn
+    public async getChatAuthState(): Promise<FeatureAuthState> {
+        // The state of the connection may not have been properly validated
+        // and the current state we see may be stale, so refresh for latest state.
+        await this.auth.refreshConnectionState(this.conn)
+        return this.getChatAuthStateSync(this.conn)
+    }
 
-        if (currentConnection === undefined) {
+    /**
+     * Synchronously returns a snapshot of the overall auth state of CodeWhisperer + Chat features without
+     * validating or modifying the connection state. It is possible that the connection
+     * is invalid/valid, but the current state displays something else. To guarantee the true state,
+     * use async method getChatAuthState()
+     */
+    public getChatAuthStateSync(conn = this.conn): FeatureAuthState {
+        if (conn === undefined) {
             return buildFeatureAuthState(AuthStates.disconnected)
         }
-        if (!isSsoConnection(currentConnection)) {
-            throw new ToolkitError(
-                `Connection "${currentConnection.id}" is not a valid type: ${currentConnection.type}`
-            )
-        }
-
-        // The state of the connection may not have been properly validated
-        // and the current state we see may be stale, so refresh for latest state.
-        if (shouldRefresh) {
-            await this.auth.refreshConnectionState(currentConnection)
+        if (!isSsoConnection(conn)) {
+            throw new ToolkitError(`Connection "${conn.id}" is not a valid type: ${conn.type}`)
         }
 
         // default to expired to indicate reauth is needed if unmodified
@@ -418,11 +419,11 @@ export class AuthUtil {
             return state
         }
 
-        if (isBuilderIdConnection(currentConnection) || isIdcSsoConnection(currentConnection)) {
-            if (isValidCodeWhispererCoreConnection(currentConnection)) {
+        if (isBuilderIdConnection(conn) || isIdcSsoConnection(conn)) {
+            if (isValidCodeWhispererCoreConnection(conn)) {
                 state[Features.codewhispererCore] = AuthStates.connected
             }
-            if (isValidAmazonQConnection(currentConnection)) {
+            if (isValidAmazonQConnection(conn)) {
                 Object.values(Features).forEach(v => (state[v as Feature] = AuthStates.connected))
             }
         }

packages/core/src/codewhispererChat/editor/codelens.ts

@@ -8,6 +8,7 @@ import { ToolkitError } from '../../shared/errors'
 import { Commands, placeholder } from '../../shared/vscode/commands2'
 import { platform } from 'os'
 import { focusAmazonQPanel } from '../commands/registerCommands'
+import { AuthStates, AuthUtil } from '../../codewhisperer/util/authUtil'
 
 /** When the user clicks the CodeLens that prompts user to try Amazon Q chat */
 export const tryChatCodeLensCommand = Commands.declare(`_aws.amazonq.tryChatCodeLens`, () => async () => {
@@ -63,9 +64,13 @@ export class TryChatCodeLensProvider implements vscode.CodeLensProvider {
         document: vscode.TextDocument,
         token: vscode.CancellationToken
     ): vscode.ProviderResult<vscode.CodeLens[]> {
-        return new Promise(async resolve => {
+        return new Promise(resolve => {
             token.onCancellationRequested(() => resolve([]))
 
+            if (AuthUtil.instance.getChatAuthStateSync().amazonQ !== AuthStates.connected) {
+                return resolve([])
+            }
+
             if (this.count >= TryChatCodeLensProvider.maxCount) {
                 // We only want to show this code lens a certain amount of times
                 // to not annoy customers. The following ensures it is never shown again.

packages/core/src/test/codewhispererChat/editor/codelens.test.ts

@@ -14,6 +14,8 @@ import { assertTelemetry, installFakeClock, tryRegister } from '../../testUtil'
 import { InstalledClock } from '@sinonjs/fake-timers'
 import globals from '../../../shared/extensionGlobals'
 import { focusAmazonQPanel } from '../../../codewhispererChat/commands/registerCommands'
+import sinon from 'sinon'
+import { AuthState, AuthStates, AuthUtil, FeatureAuthState } from '../../../codewhisperer/util/authUtil'
 
 describe('TryChatCodeLensProvider', () => {
     let instance: TryChatCodeLensProvider = new TryChatCodeLensProvider()
@@ -39,9 +41,16 @@ describe('TryChatCodeLensProvider', () => {
         instance.dispose()
         cancellationTokenSource?.dispose()
         clock.uninstall()
+        sinon.restore()
     })
 
+    function stubConnection(state: AuthState) {
+        return sinon.stub(AuthUtil.instance, 'getChatAuthStateSync').returns({ amazonQ: state } as FeatureAuthState)
+    }
+
     it('keeps returning a code lense until it hits the max times it should show', async function () {
+        stubConnection('connected')
+
         let codeLensCount = 0
         const modifierKey = resolveModifierKey()
         while (codeLensCount < 10) {
@@ -83,6 +92,21 @@ describe('TryChatCodeLensProvider', () => {
         })
     })
 
+    it('does NOT show codelens if amazon Q is not connected', async function () {
+        const testConnection = async (state: AuthState) => {
+            const stub = stubConnection(state)
+
+            const emptyResult = await instance.provideCodeLenses({} as any, new vscode.CancellationTokenSource().token)
+            assert.deepStrictEqual(emptyResult, [], `codelens shown with state: ${state}`)
+            stub.restore()
+        }
+
+        const testStates = Object.values(AuthStates).filter(s => s !== AuthStates.connected)
+        for (const state of testStates) {
+            await testConnection(state)
+        }
+    })
+
     it('outputs expected telemetry', async function () {
         await tryChatCodeLensCommand.execute()
         assertTelemetry('vscode_executeCommand', { command: focusAmazonQPanel.id, source: 'codeLens' })