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

[Hweinstock]

Problem

we have an internal that runs every 5 minutes 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:

aws-toolkit-vscode/packages/core/src/auth/sso/ssoAccessTokenProvider.ts

Lines 100 to 122 in c3ea31d

	public async getToken(): Promise<SsoToken | undefined> {
	const data = await this.cache.token.load(this.tokenCacheKey)
	SsoAccessTokenProvider.logIfChanged(
	indent(
	`current client registration id=${data?.registration?.clientId}
	expires at ${data?.registration?.expiresAt}
	key = ${this.tokenCacheKey}`,
	4,
	true
	)
	)
	if (!data || !isExpired(data.token)) {
	return data?.token
	}
	if (data.registration && !isExpired(data.registration) && hasProps(data.token, 'refreshToken')) {
	const refreshed = await this.refreshToken(data.token, data.registration)
	return refreshed.token
	} else {
	await this.invalidate('allCacheExpired')
	}
	}

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

  aws-toolkit-vscode/packages/amazonq/src/lsp/auth.ts

  Line 83 in ce01fa8

  this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)

  ), 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.
• License: I confirm that my contribution is made under the terms of the Apache 2.0 license.

[Hweinstock]

/runIntegrationTests