fix(auth): token refresh rapidly called unexpectedly (#6479)

## Problem:

getChatAuthState() is called in many places by the Q features
simultaneously,
this eventually triggers multiple calls to getToken() and if needed
refreshToken().

This resulted in refreshToken being spammed and the Identity team seeing
spikes in token refreshes
from clients.

## Solution:

Throttle getChatAuthState().

Throttling w/ leading: true, allows us to instantly return
a fresh result OR a cached result in the case we are throttled. Debounce
on the
other hand would cause callers to hang since they have to wait for
debounce to timeout.

Also, we put a debounce on getToken() before in #6282 but this did not
work since a new
SsoAccessToken instance is created each time the offending code flow
triggered (we could
look to cache the instance instead which would enable the getToken()
debounce to be useful.

### Testing

To test the difference after adding the throttle:
- Add log statements to `getToken()`
- Set an expired date in the SSO cache for both token expiration +
client registration expiration
- Use chat

What would happen is that without throttle it would trigger
getChatAuthState() many times, likely due to the connection
becoming invalid and sending an event to all Q features, causing each of
them to call getChatAuthState() at the same time.

But when the throttle was added, the amount of these calls dropped to at
most 2.

Signed-off-by: nkomonen-amazon <[email protected]>


---

- 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.

---------

Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

packages/amazonq/src/extensionNode.ts

@@ -82,7 +82,7 @@ async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
     try {
         // May call connection validate functions that try to refresh the token.
         // This could result in network errors.
-        authState = (await AuthUtil.instance.getChatAuthState(false)).codewhispererChat
+        authState = (await AuthUtil.instance._getChatAuthState(false)).codewhispererChat
     } catch (err) {
         if (
             isNetworkError(err) &&

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

@@ -33,7 +33,7 @@ import { randomUUID } from '../../shared/crypto'
 import { getExtRuntimeContext } from '../../shared/vscode/env'
 import { showInputBox } from '../../shared/ui/inputPrompter'
 import { AmazonQPromptSettings, DevSettings, PromptSettings, ToolkitPromptSettings } from '../../shared/settings'
-import { debounce, onceChanged } from '../../shared/utilities/functionUtils'
+import { onceChanged } from '../../shared/utilities/functionUtils'
 import { NestedMap } from '../../shared/utilities/map'
 import { asStringifiedStack } from '../../shared/telemetry/spans'
 import { showViewLogsMessage } from '../../shared/utilities/messages'
@@ -97,20 +97,7 @@ export abstract class SsoAccessTokenProvider {
         this.reAuthState.set(this.profile, { reAuthReason: `invalidate():${reason}` })
     }
 
-    /**
-     * Sometimes we get many calls at once and this
-     * can trigger redundant disk reads, or token refreshes.
-     * We debounce to avoid this.
-     *
-     * NOTE: The property {@link getTokenDebounced()} does not work with being stubbed for tests, so
-     * this redundant function was created to work around that.
-     */
     public async getToken(): Promise<SsoToken | undefined> {
-        return this.getTokenDebounced()
-    }
-    private getTokenDebounced = debounce(() => this._getToken(), 50)
-    /** Exposed for testing purposes only */
-    public async _getToken(): Promise<SsoToken | undefined> {
         const data = await this.cache.token.load(this.tokenCacheKey)
         SsoAccessTokenProvider.logIfChanged(
             indent(

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

@@ -44,6 +44,7 @@ import { telemetry } from '../../shared/telemetry/telemetry'
 import { asStringifiedStack } from '../../shared/telemetry/spans'
 import { withTelemetryContext } from '../../shared/telemetry/util'
 import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands'
+import { throttle } from 'lodash'
 
 /** Backwards compatibility for connections w pre-chat scopes */
 export const codeWhispererCoreScopes = [...scopesCodeWhispererCore]
@@ -402,9 +403,23 @@ export class AuthUtil {
      *
      * By default, network errors are ignored when determining auth state since they may be silently
      * recoverable later.
+     *
+     * THROTTLE: This function is called in rapid succession by Amazon Q features and can lead to
+     *           a barrage of disk access and/or token refreshes. We throttle to deal with this.
+     *
+     *           Note we do an explicit cast of the return type due to Lodash types incorrectly indicating
+     *           a FeatureAuthState or undefined can be returned. But since we set `leading: true`
+     *           it will always return FeatureAuthState
+     */
+    public getChatAuthState = throttle(() => this._getChatAuthState(), 2000, {
+        leading: true,
+    }) as () => Promise<FeatureAuthState>
+    /**
+     * IMPORTANT: Only use this if you do NOT want to swallow network errors, otherwise use {@link getChatAuthState()}
+     * @param ignoreNetErr swallows network errors
      */
     @withTelemetryContext({ name: 'getChatAuthState', class: authClassName })
-    public async getChatAuthState(ignoreNetErr: boolean = true): Promise<FeatureAuthState> {
+    public async _getChatAuthState(ignoreNetErr: boolean = true): 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.
         if (ignoreNetErr) {

packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts

@@ -156,20 +156,6 @@ describe('SsoAccessTokenProvider', function () {
             assert.strictEqual(cachedToken, undefined)
         })
 
-        it('concurrent calls are debounced', async function () {
-            const validToken = createToken(hourInMs)
-            await cache.token.save(startUrl, { region, startUrl, token: validToken })
-            const actualGetToken = sinon.spy(sut, '_getToken')
-
-            const result = await Promise.all([sut.getToken(), sut.getToken(), sut.getToken()])
-
-            // Subsequent other calls were debounced so this was only called once
-            assert.strictEqual(actualGetToken.callCount, 1)
-            for (const r of result) {
-                assert.deepStrictEqual(r, validToken)
-            }
-        })
-
         describe('Exceptions', function () {
             it('drops expired tokens if failure was a client-fault', async function () {
                 const exception = new UnauthorizedClientException({ message: '', $metadata: {} })