refactor(bearertoken): reduce re-auth by increasing polling rate. (aws#6882)

## Problem
we have an [internal that runs every 5 minutes
](https://github.com/aws/aws-toolkit-vscode/blob/c3ea31dbd68cfc231b3f81795864f9dbc8a5e488/packages/amazonq/src/lsp/client.ts#L113)
to take the current bearer token in the extension and send it to the
language server. This is attempting to solve the problem of ensuring the
bearer token used by the codewhisperer-LSP has not expired.

The difficulty is that the codewhisperer-LSP does not know when it is
expired so it is the responsibility of the client(us) to ensure this
token has not expired.

We currently do not passively refresh this token, but actively refresh
it on request (see here:
https://github.com/aws/aws-toolkit-vscode/blob/c3ea31dbd68cfc231b3f81795864f9dbc8a5e488/packages/core/src/auth/sso/ssoAccessTokenProvider.ts#L100-L122).
Therefore, we do not know if it has refreshed until we request it.

The current solution is to get and send the bearer token to the lsp
every 5 minutes. However, if the token expires in the middle of these
intervals, there is then a 1-5 minute period where if the customer
interacts with anything using the codewhisperer-lsp, they will be
prompted to re-auth unnecessarily.

The issue can be observed by manually poisoning the bearerToken used by
the codewhisper lsp with something like:
```
void sleep(delay).then(async () => {

   getLogger().info('Updating bearer token to bad token')

   await auth.updateBearerToken('bad Token')

})
```
Then, any lsp interaction will prompt for re-auth (until next interval 5
minute interval hits and token is refreshed).
Note that manually expiring the token locally does not work, since it is
still valid for requests on the LSP.

## Solution
- increase polling frequency to 1 minute by making use of a local cache
to avoid excessive LSP requests. This means a call is made to `getToken`
(linked above) on a regular interval. However, if its cached, this is
very cheap operation.
- once Auth is using flare, we can implement a proper fix in Flare.
Early idea is to add notifications for when Flare refreshes tokens, and
use those as a signal to update the token used by the other LSP servers.
- Document the problem and potential long term fixes for once Auth LSP
is being used (WIP)
- Refactored `auth.init -> auth.refreshConnection` to be more specific. 

## Verification
- Using the code above to observe the issue, we can see that after
poisoning the bearer token used by the LSP, once the first
`UpdateBearerToken: {request}` log message shows up (from:
https://github.com/aws/aws-toolkit-vscode/blob/ce01fa85c7a97c71cd3dcc11770f6cb75b29e1c2/packages/amazonq/src/lsp/auth.ts#L83),
flare interactions work again.

## Alternative Solutions
- Add passive refreshing to current Auth. 
- This is not a great approach because it would be specialized to our
current Auth logic, which is being discarded in favor of Flare.
- Retry failed requests with refreshed token. 
- The implementation would be very awkward here and tightly couples all
lsp server interactions to our auth logic. This is also against the
current paradigm of the LSP responding with a re-auth prompt when
authentication fails.
- Update bearer token before every request. 
- This seems excessive and messy, and would likely be throwaway work.
The number of LSP requests would double, potentially increasing latency
for flare operations. Since this is an edge case, it does not seem a
worthy tradeoff to degrade the common use case.
---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

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,8 @@ export class AmazonQLspAuth {
         }
     }
 
-    private async updateBearerToken(token: string) {
+    public updateBearerToken = onceChanged(this._updateBearerToken.bind(this))
+    private async _updateBearerToken(token: string) {
         const request = await this.createUpdateCredentialsRequest({
             token,
         })
@@ -83,6 +86,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,15 +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,
-    Commands,
-} from 'aws-core-vscode/shared'
+import { Settings, oidcClientName, createServerOptions, globals, Experiments, Commands } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
 
@@ -109,7 +101,7 @@ export async function startLanguageServer(
                 },
             }
         })
-        await auth.init()
+        await auth.refreshConnection()
 
         if (Experiments.instance.get('amazonqLSPInline', false)) {
             const inlineManager = new InlineCompletionManager(client)
@@ -129,23 +121,16 @@ export async function startLanguageServer(
             activate(client, encryptionKey, resourcePaths.ui)
         }
 
-        // 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)
-            })
+            }),
+            { 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'