feat(CodeCatalyst): handle partial scope expiration #4557

Problem:
Extended session duration (90 days) for CodeWhisperer is coming soon, but other scopes on the same SSO connection will not be affected, so there must be handling for "partial" expiration. Partial expiration is relevant when an SSO connection is authorized for both Q/CodeWhisperer and CodeCatalyst.

Solution:
When an AccessDeniedException (ADE) is encountered when making a CodeCatalyst call, it will clear the bearer token. When a bearer token is refreshed, it will call `CodeCatalystClientInternal.verifySession()` before making further calls. If the token was able to refresh, this means that the entire token was not expired, and that specific scopes are expired.

1. If `verifySession` fails due to an ADE, fire an EventEmitter. EventEmitters were used to prevent circular dependencies between `CodeCatalystAuthenticationProvider` and `CodeCatalystClientInternal`. I think it might be better long-term to have the client consume the auth provider, but it would have required a decent amount of refactoring.
2.  `CodeCatalystAuthenticationProvider` subscribes to this event, handling the ADE. It will determine if the scope is partially expired, or if all scopes in the connection are expired. It will make a call to CodeWhisperer to determine this.
3. The CodeCatalyst context will be updated, displaying a "Re-Authorize" node in the explorer tree that matches the existing behavior in CodeWhisperer.
4. An information box will display informing the user that CodeCatalyst is expired. If the connection is partially expired, then it will also mention that CodeWhisperer is still useable. This notification will only display at most once per session, and has a "Don't show again" option.
6. On CC activation, check if the connection is shared with CodeWhisperer, and trigger a `verifySession` if it is. This shouldn't impact activation time of users that are not using CC and CW on the same connection.

Additional Changes:
1. Moved `isInDevEnv` from `codecatalyst/utils` to `shared/vscode/env` to prevent circular dependencies.
2. Call `verifySession` during CodeCatalyst activation to detect scope expiration on startup
3. Created `showReauthenticateMessage` function in `utils/messages.ts` for shared functionality between CodeWhisperer and CodeCatalyst

UX Impact:
When the user is not using CodeWhisperer and CodeCatalyst on the same SSO profile, there will not be a UX impact. When the connection is still "valid" (able to refresh the access token), but code catalyst receives an access denied exception while attempting to call `verifySession`, the CC scope is considered expired. Anytime CC scope is expired, it will update the CodeCatalyst explorer to show this. If a user action triggered the CC call, a notification will appear telling the user to authenticate. If the connection was previously expired when starting VSCode, there will not be a notification.

Testing:

1. Added debug commands to expire and un-expire tokens instantly by changing the bearerToken to an invalid token during request signing.
2. Performed testing on a feature-enabled IdC service with the upcoming extended duration changes.
2. Verified that all behavior is expected for just CodeWhisperer, CodeWhisperer & CodeCatalyst, and just CodeCatalyst

Author: aarolieb

.changes/next-release/Feature-a5f37aec-e586-4057-9834-7d58c55c5695.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Support extended CodeWhisperer session durations when using CodeCatalyst on the same SSO connection"
+}

packages/core/package.json

@@ -209,6 +209,10 @@
                         "amazonQWelcomePage": {
                             "type": "boolean",
                             "default": false
+                        },
+                        "codeCatalystConnectionExpired": {
+                            "type": "boolean",
+                            "default": false
                         }
                     },
                     "additionalProperties": false

packages/core/src/auth/activation.ts

@@ -12,7 +12,7 @@ import { placeholder } from '../shared/vscode/commands2'
 import { getLogger } from '../shared/logger'
 import { ExtensionUse } from './utils'
 import { isCloud9 } from '../shared/extensionUtilities'
-import { isInDevEnv } from '../codecatalyst/utils'
+import { isInDevEnv } from '../shared/vscode/env'
 import { showManageConnections } from './ui/vue/show'
 import { isWeb } from '../common/webUtils'
 

packages/core/src/codecatalyst/activation.ts

@@ -18,7 +18,7 @@ import { PromptSettings } from '../shared/settings'
 import { dontShow } from '../shared/localizedText'
 import { getIdeProperties, isCloud9 } from '../shared/extensionUtilities'
 import { Commands, placeholder } from '../shared/vscode/commands2'
-import { getCodeCatalystConfig } from '../shared/clients/codecatalystClient'
+import { createClient, getCodeCatalystConfig } from '../shared/clients/codecatalystClient'
 import { isDevenvVscode } from './utils'
 import { getThisDevEnv } from './model'
 import { getLogger } from '../shared/logger/logger'
@@ -37,6 +37,17 @@ export async function activate(ctx: ExtContext): Promise<void> {
 
     await authProvider.restore()
 
+    // if connection is shared with CodeWhisperer, check if CodeCatalyst scopes are expired
+    if (authProvider.activeConnection && authProvider.isSharedConn()) {
+        try {
+            await createClient(authProvider.activeConnection, undefined, undefined, undefined, {
+                showReauthPrompt: false,
+            })
+        } catch (err) {
+            getLogger().info('codecatalyst: createClient failed during activation: %s', err)
+        }
+    }
+
     ctx.extensionContext.subscriptions.push(
         uriHandlers.register(ctx.uriHandler, CodeCatalystCommands.declared),
         ...Object.values(CodeCatalystCommands.declared).map(c => c.register(commands)),

packages/core/src/codecatalyst/auth.ts

@@ -4,7 +4,7 @@
  */
 
 import * as vscode from 'vscode'
-import { CodeCatalystClient, createClient } from '../shared/clients/codecatalystClient'
+import { onAccessDeniedException, CodeCatalystClient, createClient } from '../shared/clients/codecatalystClient'
 import { Auth } from '../auth/auth'
 import { getSecondaryAuth } from '../auth/secondaryAuth'
 import { getLogger } from '../shared/logger'
@@ -24,6 +24,9 @@ import {
 } from '../auth/connection'
 import { createBuilderIdConnection } from '../auth/utils'
 import { builderIdStartUrl } from '../auth/sso/model'
+import { codeWhispererClient } from '../codewhisperer/client/codewhisperer'
+import { AuthUtil as CodeWhispererAuth } from '../codewhisperer/util/authUtil'
+import { showReauthenticateMessage } from '../shared/utilities/messages'
 
 // Secrets stored on the macOS keychain appear as individual entries for each key
 // This is fine so long as the user has only a few accounts. Otherwise this should
@@ -51,12 +54,23 @@ export function setCodeCatalystConnectedContext(isConnected: boolean) {
     return vscode.commands.executeCommand('setContext', 'aws.codecatalyst.connected', isConnected)
 }
 
+type ConnectionState = {
+    onboarded: boolean
+    scopeExpired: boolean
+}
+
 export class CodeCatalystAuthenticationProvider {
     public readonly onDidChangeActiveConnection = this.secondaryAuth.onDidChangeActiveConnection
+    public readonly onAccessDeniedException = onAccessDeniedException
+    private readonly onDidChangeEmitter = new vscode.EventEmitter<void>()
+    public readonly onDidChange = this.onDidChangeEmitter.event
+
+    private readonly mementoKey = 'codecatalyst.connections'
 
     public constructor(
         protected readonly storage: CodeCatalystAuthStorage,
         protected readonly memento: vscode.Memento,
+
         public readonly auth = Auth.instance,
         public readonly secondaryAuth = getSecondaryAuth(
             auth,
@@ -65,8 +79,17 @@ export class CodeCatalystAuthenticationProvider {
             isValidCodeCatalystConnection
         )
     ) {
-        this.secondaryAuth.onDidChangeActiveConnection(async () => {
+        this.onDidChangeActiveConnection(async () => {
+            if (this.activeConnection) {
+                await this.setScopeExpired(this.activeConnection, false)
+            }
             await setCodeCatalystConnectedContext(this.isConnectionValid())
+            this.onDidChangeEmitter.fire()
+        })
+
+        this.onAccessDeniedException(async (showReauthPrompt: boolean) => {
+            await this.accessDeniedExceptionHandler(showReauthPrompt)
+            this.onDidChangeEmitter.fire()
         })
 
         // set initial context in case event does not trigger
@@ -81,8 +104,27 @@ export class CodeCatalystAuthenticationProvider {
         return this.secondaryAuth.hasSavedConnection
     }
 
+    public async setScopeExpired(conn: SsoConnection, isExpired: boolean) {
+        await this.updateConnectionState(conn, { scopeExpired: isExpired })
+    }
+
+    public isScopeExpired(conn: SsoConnection): boolean {
+        return this.getConnectionState(conn).scopeExpired
+    }
+
+    public isSharedConn(): boolean {
+        return (
+            this.secondaryAuth.activeConnection !== undefined &&
+            this.secondaryAuth.activeConnection?.id === CodeWhispererAuth.instance.secondaryAuth.activeConnection?.id
+        )
+    }
+
     public isConnectionValid(): boolean {
-        return this.activeConnection !== undefined && !this.secondaryAuth.isConnectionExpired
+        return (
+            this.activeConnection !== undefined &&
+            !this.secondaryAuth.isConnectionExpired &&
+            !this.isScopeExpired(this.activeConnection)
+        )
     }
 
     // Get rid of this? Not sure where to put PAT code.
@@ -118,6 +160,57 @@ export class CodeCatalystAuthenticationProvider {
         await this.secondaryAuth.restoreConnection()
     }
 
+    private async accessDeniedExceptionHandler(showReauthPrompt: boolean = true) {
+        if (!this.isConnectionValid()) {
+            return
+        }
+
+        // Check if CodeWhisper and CodeCatalyst share the same connection
+        const isSharedConn = this.isSharedConn()
+
+        if (isSharedConn) {
+            await codeWhispererClient.listAvailableCustomizations()
+        }
+
+        /*
+         * Partial Expiration occurs when CodeWhisperer and CodeCatalyst are using
+         * the same SSO connection, but CodeCatalyst scopes have expired before CodeWhisperer.
+         */
+        const isPartialExpiration = isSharedConn && !CodeWhispererAuth.instance.isConnectionExpired()
+
+        getLogger().info(
+            'auth: CodeCatalyst scopes are expired. shared=%s, partialExpiration=%s, showReauthPrompt=%s',
+            isSharedConn,
+            isPartialExpiration,
+            showReauthPrompt
+        )
+
+        await this.setScopeExpired(this.activeConnection!, true)
+        await setCodeCatalystConnectedContext(this.isConnectionValid())
+
+        // showReauthPrompt is true primarily when a user interaction triggered the ADE
+        if (showReauthPrompt) {
+            void this.showReauthenticationPrompt(this.activeConnection!, isPartialExpiration)
+        }
+    }
+
+    public async showReauthenticationPrompt(conn: Connection, isPartialExpiration?: boolean): Promise<void> {
+        const expiredMessage =
+            'Connection expired. To continue using CodeCatalyst, connect with AWS Builder ID or AWS IAM Identity center.'
+        const partiallyExpiredMessage =
+            'CodeCatalyst connection has expired. Amazon Q/CodeWhisperer is still connected.'
+
+        await showReauthenticateMessage({
+            message: isPartialExpiration ? partiallyExpiredMessage : expiredMessage,
+            connect: 'Connect with AWS',
+            doNotShow: "Don't Show Again",
+            suppressId: 'codeCatalystConnectionExpired',
+            reauthFunc: async () => {
+                await this.auth.reauthenticate(conn)
+            },
+        })
+    }
+
     public async promptOnboarding(): Promise<void> {
         const message = `Using CodeCatalyst requires onboarding with a Space. Sign up with CodeCatalyst to get started.`
         const openBrowser = 'Open Browser'
@@ -280,28 +373,40 @@ export class CodeCatalystAuthenticationProvider {
         }
     }
 
-    public async isConnectionOnboarded(conn: SsoConnection, recheck = false) {
-        const mementoKey = 'codecatalyst.connections'
-        const getState = () => this.memento.get(mementoKey, {} as Record<string, { onboarded: boolean }>)
-        const updateState = (state: { onboarded: boolean }) =>
-            this.memento.update(mementoKey, {
-                ...getState(),
-                [conn.id]: state,
-            })
+    private getState(): Record<string, ConnectionState> {
+        return this.memento.get(this.mementoKey, {} as Record<string, ConnectionState>)
+    }
+
+    public getConnectionState(conn: SsoConnection): ConnectionState {
+        return this.getState()[conn.id]
+    }
+
+    private async setConnectionState(conn: SsoConnection, state: ConnectionState) {
+        await this.memento.update(this.mementoKey, {
+            ...this.getState(),
+            [conn.id]: state,
+        })
+    }
 
-        const state = getState()[conn.id]
+    private async updateConnectionState(conn: SsoConnection, state: Partial<ConnectionState>) {
+        const initial = this.getConnectionState(conn)
+        await this.setConnectionState(conn, { ...initial, ...state })
+    }
+
+    public async isConnectionOnboarded(conn: SsoConnection, recheck = false) {
+        const state = this.getConnectionState(conn)
         if (state !== undefined && !recheck) {
             return state.onboarded
         }
 
         try {
             await createClient(conn)
-            await updateState({ onboarded: true })
+            await this.updateConnectionState(conn, { onboarded: true })
 
             return true
         } catch (e) {
             if (isOnboardingException(e) && this.auth.getConnectionState(conn) === 'valid') {
-                await updateState({ onboarded: false })
+                await this.updateConnectionState(conn, { onboarded: false })
 
                 return false
             }

packages/core/src/codecatalyst/explorer.ts

@@ -44,7 +44,18 @@ async function getLocalCommands(auth: CodeCatalystAuthenticationProvider) {
         iconPath: getIcon('vscode-question'),
     })
 
-    if (!auth.activeConnection || !auth.isConnectionValid()) {
+    // There is a connection, but it is expired, or CodeCatalyst scopes are expired.
+    if (auth.activeConnection && !auth.isConnectionValid()) {
+        return [
+            reauth.build(auth.activeConnection, auth).asTreeNode({
+                label: 'Re-authenticate to connect',
+                iconPath: addColor(getIcon('vscode-debug-disconnect'), 'notificationsErrorIcon.foreground'),
+            }),
+            learnMoreNode,
+        ]
+    }
+
+    if (!auth.activeConnection) {
         return [
             showManageConnections.build(placeholder, 'codecatalystDeveloperTools', 'codecatalyst').asTreeNode({
                 label: 'Sign in to get started',
@@ -129,10 +140,9 @@ export class CodeCatalystRootNode implements TreeNode {
 
     public constructor(private readonly authProvider: CodeCatalystAuthenticationProvider) {
         this.addRefreshEmitter(() => this.onDidChangeEmitter.fire())
-        this.authProvider.onDidChangeActiveConnection(() => {
-            for (const fire of this.refreshEmitters) {
-                fire()
-            }
+
+        this.authProvider.onDidChange(() => {
+            this.refreshEmitters.forEach(fire => fire())
         })
     }
 

packages/core/src/codecatalyst/utils.ts

@@ -7,7 +7,6 @@ import { Ides } from 'aws-sdk/clients/codecatalyst'
 import * as vscode from 'vscode'
 import { CodeCatalystResource, getCodeCatalystConfig } from '../shared/clients/codecatalystClient'
 import { pushIf } from '../shared/utilities/collectionUtils'
-import { getCodeCatalystDevEnvId } from '../shared/vscode/env'
 import { getLogger } from '../shared/logger'
 
 /**
@@ -59,10 +58,3 @@ export function openCodeCatalystUrl(o: CodeCatalystResource) {
 export function isDevenvVscode(ides: Ides | undefined): boolean {
     return ides !== undefined && ides.findIndex(ide => ide.name === 'VSCode') !== -1
 }
-
-/**
- * Returns true if we are in a dev env
- */
-export function isInDevEnv(): boolean {
-    return !!getCodeCatalystDevEnvId()
-}

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

@@ -9,7 +9,6 @@ import { Auth } from '../../auth/auth'
 import { ToolkitError } from '../../shared/errors'
 import { getSecondaryAuth } from '../../auth/secondaryAuth'
 import { isCloud9, isSageMaker } from '../../shared/extensionUtilities'
-import { PromptSettings } from '../../shared/settings'
 import {
     scopesCodeWhispererCore,
     createBuilderIdProfile,
@@ -33,6 +32,7 @@ import { GlobalState } from '../../shared/globalState'
 import { vsCodeState } from '../models/model'
 import { onceChanged } from '../../shared/utilities/functionUtils'
 import { indent } from '../../shared/utilities/textUtilities'
+import { showReauthenticateMessage } from '../../shared/utilities/messages'
 
 /** Backwards compatibility for connections w pre-chat scopes */
 export const codeWhispererCoreScopes = [...scopesSsoAccountAccess, ...scopesCodeWhispererCore]
@@ -362,25 +362,20 @@ export class AuthUtil {
     }
 
     public async showReauthenticatePrompt(isAutoTrigger?: boolean) {
-        const settings = PromptSettings.instance
-        const shouldShow = await settings.isPromptEnabled('codeWhispererConnectionExpired')
-        if (!shouldShow || (isAutoTrigger && this.reauthenticatePromptShown)) {
+        if (isAutoTrigger && this.reauthenticatePromptShown) {
             return
         }
 
-        await vscode.window
-            .showInformationMessage(
-                CodeWhispererConstants.connectionExpired,
-                CodeWhispererConstants.connectWithAWSBuilderId,
-                CodeWhispererConstants.DoNotShowAgain
-            )
-            .then(async resp => {
-                if (resp === CodeWhispererConstants.connectWithAWSBuilderId) {
-                    await this.reauthenticate()
-                } else if (resp === CodeWhispererConstants.DoNotShowAgain) {
-                    await settings.disablePrompt('codeWhispererConnectionExpired')
-                }
-            })
+        await showReauthenticateMessage({
+            message: CodeWhispererConstants.connectionExpired,
+            connect: CodeWhispererConstants.connectWithAWSBuilderId,
+            doNotShow: CodeWhispererConstants.DoNotShowAgain,
+            suppressId: 'codeWhispererConnectionExpired',
+            reauthFunc: async () => {
+                await this.reauthenticate()
+            },
+        })
+
         if (isAutoTrigger) {
             this.reauthenticatePromptShown = true
         }
@@ -457,7 +452,7 @@ export const AuthStates = {
     /**
      * The current connection exists, but needs to be reauthenticated for this feature to work
      *
-     * Look to use {@link AuthUtil.reauthenticate}
+     * Look to use {@link AuthUtil.reauthenticate()}
      */
     expired: 'expired',
     /**