refactor: poll more frequently by leveraging cache

Author: Hweinstock

packages/amazonq/src/lsp/auth.ts

@@ -14,6 +14,8 @@ import * as crypto from 'crypto'
 import { LanguageClient } from 'vscode-languageclient'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { Writable } from 'stream'
+import { onceChanged } from 'aws-core-vscode/utils'
+import { getLogger, oneMinute } from 'aws-core-vscode/shared'
 
 export const encryptionKey = crypto.randomBytes(32)
 
@@ -64,7 +66,7 @@ export const notificationTypes = {
 export class AmazonQLspAuth {
     constructor(private readonly client: LanguageClient) {}
 
-    async init() {
+    async refreshConnection() {
         const activeConnection = AuthUtil.instance.auth.activeConnection
         if (activeConnection?.type === 'sso') {
             // send the token to the language server
@@ -73,7 +75,9 @@ export class AmazonQLspAuth {
         }
     }
 
-    private async updateBearerToken(token: string) {
+    public updateBearerToken = onceChanged(this._updateBearerToken.bind(this))
+    private async _updateBearerToken(token: string) {
+        getLogger('amazonqLsp').info('Sending updated bearer token to LSP')
         const request = await this.createUpdateCredentialsRequest({
             token,
         })
@@ -83,6 +87,16 @@ export class AmazonQLspAuth {
         this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)
     }
 
+    public startTokenRefreshInterval(pollingTime: number = oneMinute) {
+        const interval = setInterval(async () => {
+            await this.refreshConnection().catch((e) => {
+                getLogger('amazonqLsp').error('Unable to update bearer token: %s', (e as Error).message)
+                clearInterval(interval)
+            })
+        }, pollingTime)
+        return interval
+    }
+
     private async createUpdateCredentialsRequest(data: any) {
         const payload = new TextEncoder().encode(JSON.stringify({ data }))
 

packages/amazonq/src/lsp/client.ts

@@ -11,7 +11,7 @@ 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, getLogger } from 'aws-core-vscode/shared'
+import { Settings, oidcClientName, createServerOptions, globals, Experiments } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
 
@@ -93,7 +93,7 @@ export async function startLanguageServer(
     const auth = new AmazonQLspAuth(client)
 
     return client.onReady().then(async () => {
-        await auth.init()
+        await auth.refreshConnection()
         const inlineManager = new InlineCompletionManager(client)
         inlineManager.registerInlineCompletion()
         if (Experiments.instance.get('amazonqChatLSP', false)) {
@@ -109,24 +109,17 @@ export async function startLanguageServer(
             }
         })
 
-        // Temporary code for pen test. Will be removed when we switch to the real flare auth
-        const authInterval = setInterval(async () => {
-            try {
-                await auth.init()
-            } catch (e) {
-                getLogger('amazonqLsp').error('Unable to update bearer token: %s', (e as Error).message)
-                clearInterval(authInterval)
-            }
-        }, 300000) // every 5 minutes
+        const refreshInterval = auth.startTokenRefreshInterval()
 
         toDispose.push(
             AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
-                await auth.init()
+                await auth.refreshConnection()
             }),
             AuthUtil.instance.auth.onDidDeleteConnection(async () => {
                 client.sendNotification(notificationTypes.deleteBearerToken.method)
             }),
-            inlineManager
+            inlineManager,
+            { dispose: () => clearInterval(refreshInterval) }
         )
     })
 }

packages/amazonq/test/unit/amazonq/lsp/auth.test.ts

@@ -0,0 +1,33 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import assert from 'assert'
+import { AmazonQLspAuth } from '../../../../src/lsp/auth'
+import { LanguageClient } from 'vscode-languageclient'
+
+describe('AmazonQLspAuth', function () {
+    describe('updateBearerToken', function () {
+        it('makes request to LSP when token changes', async function () {
+            // Note: this token will be encrypted
+            let lastSentToken = {}
+            const auth = new AmazonQLspAuth({
+                sendRequest: (_method: string, param: any) => {
+                    lastSentToken = param
+                },
+                info: (_message: string, _data: any) => {},
+            } as LanguageClient)
+
+            await auth.updateBearerToken('firstToken')
+            assert.notDeepStrictEqual(lastSentToken, {})
+            const encryptedFirstToken = lastSentToken
+
+            await auth.updateBearerToken('secondToken')
+            assert.notDeepStrictEqual(lastSentToken, encryptedFirstToken)
+            const encryptedSecondToken = lastSentToken
+
+            await auth.updateBearerToken('secondToken')
+            assert.deepStrictEqual(lastSentToken, encryptedSecondToken)
+        })
+    })
+})

packages/core/src/shared/index.ts

@@ -72,3 +72,4 @@ export * from './lsp/utils/platform'
 export * as processUtils from './utilities/processUtils'
 export * as BaseLspInstaller from './lsp/baseLspInstaller'
 export * as collectionUtil from './utilities/collectionUtils'
+export * from './datetime'

packages/core/src/shared/utilities/index.ts

@@ -5,3 +5,4 @@
 
 export { isExtensionInstalled, isExtensionActive } from './vsCodeUtils'
 export { VSCODE_EXTENSION_ID } from '../extensions'
+export * from './functionUtils'