fix(chat): Resend auth token on server restart

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))
         }
     }
 

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)
+    })
+}