fix(amazonq): Add missing bearer token auth

Author: jpinkney-aws

packages/amazonq/src/lsp/auth.ts

@@ -3,7 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { RequestType, ResponseMessage } from '@aws/language-server-runtimes/protocol'
+import {
+    ConnectionMetadata,
+    NotificationType,
+    RequestType,
+    ResponseMessage,
+} from '@aws/language-server-runtimes/protocol'
 import * as jose from 'jose'
 import * as crypto from 'crypto'
 import { LanguageClient } from 'vscode-languageclient'
@@ -43,12 +48,19 @@ export interface UpdateCredentialsRequest {
     encrypted: boolean
 }
 
-const notificationTypes = {
+export const notificationTypes = {
     updateBearerToken: new RequestType<UpdateCredentialsRequest, ResponseMessage, Error>(
         'aws/credentials/token/update'
     ),
+    deleteBearerToken: new NotificationType('aws/credentials/token/delete'),
+    getConnectionMetadata: new RequestType<undefined, ConnectionMetadata, Error>(
+        'aws/credentials/getConnectionMetadata'
+    ),
 }
 
+/**
+ * Facade over our VSCode Auth that does crud operations on the language server auth
+ */
 export class AmazonQLSPAuth {
     constructor(private readonly client: LanguageClient) {}
 

packages/amazonq/src/lsp/client.ts

@@ -9,8 +9,9 @@ import * as cp from 'child_process'
 import * as crypto from 'crypto'
 import { LanguageClient, LanguageClientOptions, ServerOptions, TransportKind } from 'vscode-languageclient'
 import { registerInlineCompletion } from '../inline/completion'
-import { AmazonQLSPAuth, writeEncryptionInit } from './auth'
+import { AmazonQLSPAuth, notificationTypes, writeEncryptionInit } from './auth'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import { ConnectionMetadata } from '@aws/language-server-runtimes/protocol'
 
 const localize = nls.loadMessageBundle()
 
@@ -86,9 +87,21 @@ export function startLanguageServer(extensionContext: vscode.ExtensionContext, s
         await auth.init()
         registerInlineCompletion(client)
 
+        // Request handler for when the server wants to know about the clients auth connnection
+        client.onRequest<ConnectionMetadata, Error>(notificationTypes.getConnectionMetadata.method, () => {
+            return {
+                sso: {
+                    startUrl: AuthUtil.instance.auth.startUrl,
+                },
+            }
+        })
+
         toDispose.push(
             AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
                 await auth.init()
+            }),
+            AuthUtil.instance.auth.onDidDeleteConnection(async () => {
+                client.sendNotification(notificationTypes.deleteBearerToken.method)
             })
         )
     })