fix(chat): Resend auth token on server restart (#7156)

## Problem:

When the server restarted due to a crash, the auth token was not sent
again. This caused users to run in to the state where if it crashed and
restarted, when they sent a subsequent chat message they'd get stuck
with the server asking the user to Authenticate.

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.

## Repro Steps:

1. Ensure you do not have a workspace open, you can open a new vscode
window at the top left `File` > `New Window`
2. Make a random folder in your home directory
3. Make a couple typescript files in that folder
4. Send the prompt: `list all files in {folder}`
5. Accept the permissions
6. Click the toggle drop down from the response, and click the link to
the path (this is just to force a crash + restart)
7. Verify the server crashed in the logs, look for the message
`Connection to server got closed. Server will restart.`
8. ASSUMING this fix worked, the server will restart automatically and
you can continue using chat. Before this fix any subsequent messages
would ask the user to `Authenticate` again

---

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