feat(amazonq): Use codewhisperer language server for completions

Author: jpinkney-aws

package-lock.json

@@ -0,0 +1,117 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+    CancellationToken,
+    InlineCompletionContext,
+    InlineCompletionItem,
+    InlineCompletionItemProvider,
+    InlineCompletionList,
+    Position,
+    TextDocument,
+    commands,
+    languages,
+} from 'vscode'
+import { LanguageClient } from 'vscode-languageclient'
+import {
+    InlineCompletionItemWithReferences,
+    InlineCompletionListWithReferences,
+    InlineCompletionWithReferencesParams,
+    inlineCompletionWithReferencesRequestType,
+    logInlineCompletionSessionResultsNotificationType,
+    LogInlineCompletionSessionResultsParams,
+} from '@aws/language-server-runtimes/protocol'
+
+export const CodewhispererInlineCompletionLanguages = [
+    { scheme: 'file', language: 'typescript' },
+    { scheme: 'file', language: 'javascript' },
+    { scheme: 'file', language: 'json' },
+    { scheme: 'file', language: 'yaml' },
+    { scheme: 'file', language: 'java' },
+    { scheme: 'file', language: 'go' },
+    { scheme: 'file', language: 'php' },
+    { scheme: 'file', language: 'rust' },
+    { scheme: 'file', language: 'kotlin' },
+    { scheme: 'file', language: 'terraform' },
+    { scheme: 'file', language: 'ruby' },
+    { scheme: 'file', language: 'shellscript' },
+    { scheme: 'file', language: 'dart' },
+    { scheme: 'file', language: 'lua' },
+    { scheme: 'file', language: 'powershell' },
+    { scheme: 'file', language: 'r' },
+    { scheme: 'file', language: 'swift' },
+    { scheme: 'file', language: 'systemverilog' },
+    { scheme: 'file', language: 'scala' },
+    { scheme: 'file', language: 'vue' },
+    { scheme: 'file', language: 'csharp' },
+]
+
+export function registerInlineCompletion(languageClient: LanguageClient) {
+    const inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(languageClient)
+    languages.registerInlineCompletionItemProvider(CodewhispererInlineCompletionLanguages, inlineCompletionProvider)
+
+    const onInlineAcceptance = async (
+        sessionId: string,
+        itemId: string,
+        requestStartTime: number,
+        firstCompletionDisplayLatency?: number
+    ) => {
+        const params: LogInlineCompletionSessionResultsParams = {
+            sessionId: sessionId,
+            completionSessionResult: {
+                [itemId]: {
+                    seen: true,
+                    accepted: true,
+                    discarded: false,
+                },
+            },
+            totalSessionDisplayTime: Date.now() - requestStartTime,
+            firstCompletionDisplayLatency: firstCompletionDisplayLatency,
+        }
+        languageClient.sendNotification(logInlineCompletionSessionResultsNotificationType as any, params)
+    }
+    commands.registerCommand('aws.sample-vscode-ext-amazonq.accept', onInlineAcceptance)
+}
+
+export class AmazonQInlineCompletionItemProvider implements InlineCompletionItemProvider {
+    constructor(private readonly languageClient: LanguageClient) {}
+
+    async provideInlineCompletionItems(
+        document: TextDocument,
+        position: Position,
+        context: InlineCompletionContext,
+        token: CancellationToken
+    ): Promise<InlineCompletionItem[] | InlineCompletionList> {
+        const requestStartTime = Date.now()
+        const request: InlineCompletionWithReferencesParams = {
+            textDocument: {
+                uri: document.uri.toString(),
+            },
+            position,
+            context,
+        }
+
+        const response = await this.languageClient.sendRequest(
+            inlineCompletionWithReferencesRequestType as any,
+            request,
+            token
+        )
+
+        const list: InlineCompletionListWithReferences = response as InlineCompletionListWithReferences
+        this.languageClient.info(`Client: Received ${list.items.length} suggestions`)
+        const firstCompletionDisplayLatency = Date.now() - requestStartTime
+
+        // Add completion session tracking and attach onAcceptance command to each item to record used decision
+        list.items.forEach((item: InlineCompletionItemWithReferences) => {
+            item.command = {
+                command: 'aws.sample-vscode-ext-amazonq.accept',
+                title: 'On acceptance',
+                arguments: [list.sessionId, item.itemId, requestStartTime, firstCompletionDisplayLatency],
+            }
+        })
+
+        return list as InlineCompletionList
+    }
+}

packages/amazonq/src/inline/completion.ts

@@ -4,7 +4,9 @@
  */
 
 import vscode from 'vscode'
+import path from 'path'
 import { AmazonQLSPDownloader } from './download'
+import { startLanguageServer } from './client'
 
 export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
     const serverPath = ctx.asAbsolutePath('resources/qdeveloperserver')
@@ -18,4 +20,5 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
      *
      * TODO: actually hook up the language server
      */
+    await startLanguageServer(ctx, path.join(serverPath, 'aws-lsp-codewhisperer.js'))
 }

packages/amazonq/src/lsp/activation.ts

@@ -0,0 +1,88 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { window } from 'vscode'
+import { RequestType, ResponseMessage } from '@aws/language-server-runtimes/protocol'
+import * as jose from 'jose'
+import * as crypto from 'crypto'
+import { LanguageClient } from 'vscode-languageclient'
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import { Writable } from 'stream'
+
+const encryptionKey = crypto.randomBytes(32)
+
+/**
+ * Sends a json payload to the language server, who is waiting to know what the encryption key is.
+ * Code reference: https://github.com/aws/language-servers/blob/7da212185a5da75a72ce49a1a7982983f438651a/client/vscode/src/credentialsActivation.ts#L77
+ */
+export function writeEncryptionInit(stream: Writable): void {
+    const request = {
+        version: '1.0',
+        mode: 'JWT',
+        key: encryptionKey.toString('base64'),
+    }
+    stream.write(JSON.stringify(request))
+    stream.write('\n')
+}
+
+/**
+ * Request for custom notifications that Update Credentials and tokens.
+ * See core\aws-lsp-core\src\credentials\updateCredentialsRequest.ts for details
+ */
+export interface UpdateCredentialsRequest {
+    /**
+     * Encrypted token (JWT or PASETO)
+     * The token's contents differ whether IAM or Bearer token is sent
+     */
+    data: string
+    /**
+     * Used by the runtime based language servers.
+     * Signals that this client will encrypt its credentials payloads.
+     */
+    encrypted: boolean
+}
+
+const notificationTypes = {
+    updateBearerToken: new RequestType<UpdateCredentialsRequest, ResponseMessage, Error>(
+        'aws/credentials/token/update'
+    ),
+}
+
+export class AmazonQLSPAuth {
+    constructor(private readonly client: LanguageClient) {}
+
+    async init() {
+        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)
+            void window.showErrorMessage(`Updated bearer token`)
+        }
+    }
+
+    private async updateBearerToken(token: string) {
+        const request = await this.createUpdateCredentialsRequest({
+            token,
+        })
+
+        await this.client.sendRequest(notificationTypes.updateBearerToken.method, request)
+
+        this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)
+    }
+
+    private async createUpdateCredentialsRequest(data: any) {
+        const payload = new TextEncoder().encode(JSON.stringify({ data }))
+
+        const jwt = await new jose.CompactEncrypt(payload)
+            .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
+            .encrypt(encryptionKey)
+
+        return {
+            data: jwt,
+            encrypted: true,
+        }
+    }
+}