fix(chat): Resend auth token on server restart

Problem:

When the server restarted due to a crash, the auth token was not
sent again.

Note that server restart is triggered automatically by the LanguageClient, I think.

Solution:

Detect when the server is restarted and manually resend the bearer token again.

Note, this solution needs to be revisited since there may be other initialization logic
that needs to run on server restart, aside from just the bearer token.

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

Author: nkomonen-amazon

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