fix(amazonq): fix cache watcher on token file to avoid logout after migration

Author: opieter-aws

packages/core/src/auth/auth2.ts

@@ -41,7 +41,8 @@ import { LanguageClient } from 'vscode-languageclient'
 import { getLogger } from '../shared/logger/logger'
 import { ToolkitError } from '../shared/errors'
 import { useDeviceFlow } from './sso/ssoAccessTokenProvider'
-import { getCacheFileWatcher } from './sso/cache'
+import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName } from './sso/cache'
+import { VSCODE_EXTENSION_ID } from '../shared/extensions'
 
 export const notificationTypes = {
     updateBearerToken: new RequestType<UpdateCredentialsParams, ResponseMessage, Error>(
@@ -77,7 +78,7 @@ export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoToken
  * Handles auth requests to the Identity Server in the Amazon Q LSP.
  */
 export class LanguageClientAuth {
-    readonly #ssoCacheWatcher = getCacheFileWatcher()
+    readonly #ssoCacheWatcher = getCacheFileWatcher(getCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq))
 
     constructor(
         private readonly client: LanguageClient,

packages/core/src/auth/sso/cache.ts

@@ -45,8 +45,8 @@ export function getCache(directory = getCacheDir()): SsoCache {
     }
 }
 
-export function getCacheFileWatcher(directory = getCacheDir()) {
-    const watcher = vscode.workspace.createFileSystemWatcher(new vscode.RelativePattern(directory, '*.json'))
+export function getCacheFileWatcher(directory = getCacheDir(), file = '*.json') {
+    const watcher = vscode.workspace.createFileSystemWatcher(new vscode.RelativePattern(directory, file))
     globals.context.subscriptions.push(watcher)
     return watcher
 }
@@ -158,3 +158,19 @@ export function getRegistrationCacheFile(ssoCacheDir: string, key: RegistrationK
     const suffix = `${key.region}${key.scopes && key.scopes.length > 0 ? `-${hash(key.startUrl, key.scopes)}` : ''}`
     return path.join(ssoCacheDir, `aws-toolkit-vscode-client-id-${suffix}.json`)
 }
+
+/**
+ * Returns the cache file name that Flare identity server uses for SSO token and registration
+ *
+ * @param key - The key to use for the new registration cache file.
+ * See https://github.com/aws/language-servers/blob/c10819ea2c25ce564c75fb43a6792f3c919b757a/server/aws-lsp-identity/src/sso/cache/fileSystemSsoCache.ts
+ * @returns File name of the Flare cache file
+ */
+export function getFlareCacheFileName(key: string) {
+    const hash = (str: string) => {
+        const hasher = crypto.createHash('sha1')
+        return hasher.update(str).digest('hex')
+    }
+
+    return hash(key) + '.json'
+}

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

@@ -8,7 +8,6 @@ import * as localizedText from '../../shared/localizedText'
 import * as nls from 'vscode-nls'
 import { fs } from '../../shared/fs/fs'
 import * as path from 'path'
-import * as crypto from 'crypto'
 import { ToolkitError } from '../../shared/errors'
 import { AmazonQPromptSettings } from '../../shared/settings'
 import {
@@ -37,7 +36,7 @@ import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
 import { RegionProfileManager } from '../region/regionProfileManager'
 import { AuthFormId } from '../../login/webview/vue/types'
 import { getEnvironmentSpecificMemento } from '../../shared/utilities/mementos'
-import { getCacheDir, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache'
+import { getCacheDir, getFlareCacheFileName, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache'
 import { notifySelectDeveloperProfile } from '../region/utils'
 import { once } from '../../shared/utilities/functionUtils'
 
@@ -278,6 +277,7 @@ export class AuthUtil implements IAuthProvider {
     }
 
     private async cacheChangedHandler(event: cacheChangedEvent) {
+        getLogger().debug(`Auth: Cache change event received: ${event}`)
         if (event === 'delete') {
             await this.logout()
         } else if (event === 'create') {
@@ -427,25 +427,20 @@ export class AuthUtil implements IAuthProvider {
 
                 const cacheDir = getCacheDir()
 
-                const hash = (str: string) => {
-                    const hasher = crypto.createHash('sha1')
-                    return hasher.update(str).digest('hex')
-                }
-                const filePath = (str: string) => {
-                    return path.join(cacheDir, hash(str) + '.json')
-                }
-
                 const fromRegistrationFile = getRegistrationCacheFile(cacheDir, registrationKey)
-                const toRegistrationFile = filePath(
-                    JSON.stringify({
-                        region: toImport.ssoRegion,
-                        startUrl: toImport.startUrl,
-                        tool: clientName,
-                    })
+                const toRegistrationFile = path.join(
+                    cacheDir,
+                    getFlareCacheFileName(
+                        JSON.stringify({
+                            region: toImport.ssoRegion,
+                            startUrl: toImport.startUrl,
+                            tool: clientName,
+                        })
+                    )
                 )
 
                 const fromTokenFile = getTokenCacheFile(cacheDir, profileId)
-                const toTokenFile = filePath(this.profileName)
+                const toTokenFile = path.join(cacheDir, getFlareCacheFileName(this.profileName))
 
                 try {
                     await fs.rename(fromRegistrationFile, toRegistrationFile)