diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index ed5753e9d82..6a0953c98ab 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/packages/amazonq/src/lsp/auth.ts @@ -68,12 +68,15 @@ export const notificationTypes = { export class AmazonQLspAuth { constructor(private readonly client: LanguageClient) {} - async refreshConnection() { + /** + * @param force bypass memoization, and forcefully update the bearer token + */ + async refreshConnection(force: boolean = false) { const activeConnection = AuthUtil.instance.auth.activeConnection if (activeConnection?.type === 'sso') { // send the token to the language server const token = await AuthUtil.instance.getBearerToken() - await this.updateBearerToken(token) + await (force ? this._updateBearerToken(token) : this.updateBearerToken(token)) } } diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index fa45e269114..3bc5622d837 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -6,7 +6,7 @@ import vscode, { env, version } from 'vscode' import * as nls from 'vscode-nls' import * as crypto from 'crypto' -import { LanguageClient, LanguageClientOptions, RequestType } from 'vscode-languageclient' +import { LanguageClient, LanguageClientOptions, RequestType, State } from 'vscode-languageclient' import { InlineCompletionManager } from '../app/inline/completion' import { AmazonQLspAuth, encryptionKey, notificationTypes } from './auth' import { @@ -268,7 +268,25 @@ export async function startLanguageServer( }, } as DidChangeWorkspaceFoldersParams) }), - { dispose: () => clearInterval(refreshInterval) } + { dispose: () => clearInterval(refreshInterval) }, + // Set this inside onReady so that it only triggers on subsequent language server starts (not the first) + onServerRestartHandler(client, auth) ) }) } + +/** + * When the server restarts (likely due to a crash, then the LanguageClient automatically starts it again) + * we need to run some server intialization again. + */ +function onServerRestartHandler(client: LanguageClient, auth: AmazonQLspAuth) { + return client.onDidChangeState(async (e) => { + // Ensure we are in a "restart" state + if (!(e.oldState === State.Starting && e.newState === State.Running)) { + return + } + + // Need to set the auth token in the again + await auth.refreshConnection(true) + }) +}