Update CodeWhisperer server to use bearer tokens instead of IAM credentials

lsp/.eslintignore

@@ -1,2 +1,4 @@
 node_modules/**
 out/**
+server/aws-lsp-codewhisperer/src/client/codewhispererclient.d.ts
+server/aws-lsp-codewhisperer/src/client/token/codewhispererclient.d.ts

lsp/.prettierignore

@@ -4,3 +4,5 @@ node_modules/
 out/
 **/bin/
 **/obj/
+server/aws-lsp-codewhisperer/src/client/codewhispererclient.d.ts
+server/aws-lsp-codewhisperer/src/client/token/codewhispererclient.d.ts

lsp/.vscode/launch.json

@@ -46,7 +46,8 @@
             "outFiles": ["${workspaceFolder}/client/vscode/out/**/*.js"],
             "env": {
                 "LSP_SERVER": "${workspaceFolder}/app/aws-lsp-codewhisperer-binary/out/index.js",
-                "ENABLE_INLINE_COMPLETION": "true"
+                "ENABLE_INLINE_COMPLETION": "true",
+                "ENABLE_TOKEN_PROVIDER": "true"
             },
             "preLaunchTask": "npm: compile"
         },
@@ -59,7 +60,8 @@
             "outFiles": ["${workspaceFolder}/client/vscode/out/**/*.js"],
             "env": {
                 "LSP_SERVER": "${workspaceFolder}/app/aws-lsp-s3-binary/out/index.js",
-                "ENABLE_IAM_PROVIDER": "true"
+                "ENABLE_IAM_PROVIDER": "true",
+                "ENABLE_TOKEN_PROVIDER": "false"
             },
             "preLaunchTask": "npm: compile"
         },

lsp/app/aws-lsp-codewhisperer-binary/src/index.ts

@@ -4,20 +4,46 @@ import {
     CodeWhispererServiceProps,
     createCodeWhispererService,
 } from '@lsp-placeholder/aws-lsp-codewhisperer'
+import {
+    AwsInitializationOptions,
+    IdeCredentialsProvider,
+    readEncryptionInitialization,
+    shouldWaitForEncryptionKey,
+} from '@lsp-placeholder/aws-lsp-core'
 import { ProposedFeatures, createConnection } from 'vscode-languageserver/node'
 
-const connection = createConnection(ProposedFeatures.all)
+const lspConnection = createConnection(ProposedFeatures.all)
 
-const serviceProps: CodeWhispererServiceProps = {
-    displayName: CodeWhispererServer.serverId,
-    connection,
+if (shouldWaitForEncryptionKey()) {
+    // Before starting the language server, accept encryption initialization details
+    // directly from the host. This avoids writing the key to the same channel used
+    // to send encrypted data.
+    // Contract: Only read up to (and including) the first newline (\n).
+    readEncryptionInitialization(process.stdin).then(encryptionInit => {
+        createServer(lspConnection, encryptionInit.key)
+    })
+} else {
+    createServer(lspConnection)
 }
 
-const service = createCodeWhispererService(serviceProps)
+function createServer(connection: any, key?: string): CodeWhispererServer {
+    const credentialsProvider = new IdeCredentialsProvider(connection, key)
 
-const props: CodeWhispererServerProps = {
-    connection,
-    codeWhispererService: service,
-}
+    const serviceProps: CodeWhispererServiceProps = {
+        displayName: CodeWhispererServer.serverId,
+        connection,
+        credentialsProvider,
+    }
+
+    const service = createCodeWhispererService(serviceProps)
 
-export const server = new CodeWhispererServer(props)
+    const props: CodeWhispererServerProps = {
+        connection,
+        codeWhispererService: service,
+        onInitialize: (props: AwsInitializationOptions) => {
+            credentialsProvider.initialize(props)
+        },
+    }
+
+    return new CodeWhispererServer(props)
+}

lsp/app/aws-lsp-s3-binary/src/index.ts

@@ -1,12 +1,10 @@
 import {
     AwsInitializationOptions,
-    EncryptionInitialization,
     IdeCredentialsProvider,
+    readEncryptionInitialization,
     shouldWaitForEncryptionKey,
-    validateEncryptionDetails,
 } from '@lsp-placeholder/aws-lsp-core'
 import { S3Server, S3ServerProps, S3ServiceProps, createS3ervice } from '@lsp-placeholder/aws-lsp-s3'
-import { Readable } from 'stream'
 import { ProposedFeatures, createConnection } from 'vscode-languageserver/node'
 
 const lspConnection = createConnection(ProposedFeatures.all)
@@ -16,50 +14,13 @@ if (shouldWaitForEncryptionKey()) {
     // directly from the host. This avoids writing the key to the same channel used
     // to send encrypted data.
     // Contract: Only read up to (and including) the first newline (\n).
-    readLine(process.stdin).then(input => {
-        const encryptionDetails = JSON.parse(input) as EncryptionInitialization
-
-        validateEncryptionDetails(encryptionDetails)
-
-        createServer(lspConnection, encryptionDetails.key)
+    readEncryptionInitialization(process.stdin).then(encryptionInit => {
+        createServer(lspConnection, encryptionInit.key)
     })
 } else {
     createServer(lspConnection)
 }
 
-/**
- * Read from the given stream, stopping after the first newline (\n).
- * Return the string consumed from the stream.
- */
-function readLine(stream: Readable): Promise<string> {
-    return new Promise<string>((resolve, reject) => {
-        let contents = ''
-
-        // Fires when the stream has contents that can be read
-        const onStreamIsReadable = () => {
-            while (true) {
-                const byteRead: Buffer = process.stdin.read(1)
-                if (byteRead == null) {
-                    // wait for more content to arrive on the stream
-                    break
-                }
-
-                const nextChar = byteRead.toString('utf-8')
-                contents += nextChar
-
-                if (nextChar == '\n') {
-                    // Stop reading this stream, we have read a line from it
-                    stream.removeListener('readable', onStreamIsReadable)
-                    resolve(contents)
-                    break
-                }
-            }
-        }
-
-        stream.on('readable', onStreamIsReadable)
-    })
-}
-
 function createServer(connection: any, key?: string): S3Server {
     const credentialsProvider = new IdeCredentialsProvider(connection, key)