feat(amazonq): re-add basic chat through a language server

packages/amazonq/src/extension.ts

@@ -157,7 +157,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
 
     context.subscriptions.push(
         Experiments.instance.onDidChange(async (event) => {
-            if (event.key === 'amazonqLSP') {
+            if (event.key === 'amazonqLSP' || event.key === 'amazonqChatLSP') {
                 await vscode.window
                     .showInformationMessage(
                         'Amazon Q LSP setting has changed. Reload VS Code for the changes to take effect.',

packages/amazonq/src/extensionNode.ts

@@ -7,7 +7,15 @@ import * as vscode from 'vscode'
 import { activateAmazonQCommon, amazonQContextPrefix, deactivateCommon } from './extension'
 import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
 import { activate as activateQGumby } from 'aws-core-vscode/amazonqGumby'
-import { ExtContext, globals, CrashMonitoring, getLogger, isNetworkError, isSageMaker } from 'aws-core-vscode/shared'
+import {
+    ExtContext,
+    globals,
+    CrashMonitoring,
+    getLogger,
+    isNetworkError,
+    isSageMaker,
+    Experiments,
+} from 'aws-core-vscode/shared'
 import { filetypes, SchemaService } from 'aws-core-vscode/sharedNode'
 import { updateDevMode } from 'aws-core-vscode/dev'
 import { CommonAuthViewProvider } from 'aws-core-vscode/login'
@@ -43,8 +51,10 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
         extensionContext: context,
     }
 
-    await activateCWChat(context)
-    await activateQGumby(extContext as ExtContext)
+    if (!Experiments.instance.get('amazonqChatLSP', false)) {
+        await activateCWChat(context)
+        await activateQGumby(extContext as ExtContext)
+    }
 
     const authProvider = new CommonAuthViewProvider(
         context,

packages/amazonq/src/lsp/chat/activation.ts

@@ -0,0 +1,34 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { window } from 'vscode'
+import { LanguageClient } from 'vscode-languageclient'
+import { AmazonQChatViewProvider } from './webviewProvider'
+import { registerCommands } from './commands'
+import { registerLanguageServerEventListener, registerMessageListeners } from './messages'
+import { globals } from 'aws-core-vscode/shared'
+
+export function activate(languageClient: LanguageClient, encryptionKey: Buffer, mynahUIPath: string) {
+    const provider = new AmazonQChatViewProvider(mynahUIPath)
+
+    globals.context.subscriptions.push(
+        window.registerWebviewViewProvider(AmazonQChatViewProvider.viewType, provider, {
+            webviewOptions: {
+                retainContextWhenHidden: true,
+            },
+        })
+    )
+
+    /**
+     * Commands are registered independent of the webview being open because when they're executed
+     * they focus the webview
+     **/
+    registerCommands(provider)
+    registerLanguageServerEventListener(languageClient, provider)
+
+    provider.onDidResolveWebview(() => {
+        registerMessageListeners(languageClient, provider, encryptionKey)
+    })
+}

packages/amazonq/src/lsp/chat/commands.ts

@@ -0,0 +1,79 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import { Commands, globals } from 'aws-core-vscode/shared'
+import { window } from 'vscode'
+import { AmazonQChatViewProvider } from './webviewProvider'
+
+export function registerCommands(provider: AmazonQChatViewProvider) {
+    globals.context.subscriptions.push(
+        registerGenericCommand('aws.amazonq.explainCode', 'Explain', provider),
+        registerGenericCommand('aws.amazonq.refactorCode', 'Refactor', provider),
+        registerGenericCommand('aws.amazonq.fixCode', 'Fix', provider),
+        registerGenericCommand('aws.amazonq.optimizeCode', 'Optimize', provider),
+        Commands.register('aws.amazonq.sendToPrompt', (data) => {
+            const triggerType = getCommandTriggerType(data)
+            const selection = getSelectedText()
+
+            void focusAmazonQPanel().then(() => {
+                void provider.webview?.postMessage({
+                    command: 'sendToPrompt',
+                    params: { selection: selection, triggerType },
+                })
+            })
+        }),
+        Commands.register('aws.amazonq.openTab', () => {
+            void focusAmazonQPanel().then(() => {
+                void provider.webview?.postMessage({
+                    command: 'aws/chat/openTab',
+                    params: {},
+                })
+            })
+        })
+    )
+}
+
+function getSelectedText(): string {
+    const editor = window.activeTextEditor
+    if (editor) {
+        const selection = editor.selection
+        const selectedText = editor.document.getText(selection)
+        return selectedText
+    }
+
+    return ' '
+}
+
+function getCommandTriggerType(data: any): string {
+    // data is undefined when commands triggered from keybinding or command palette. Currently no
+    // way to differentiate keybinding and command palette, so both interactions are recorded as keybinding
+    return data === undefined ? 'hotkeys' : 'contextMenu'
+}
+
+function registerGenericCommand(commandName: string, genericCommand: string, provider: AmazonQChatViewProvider) {
+    return Commands.register(commandName, (data) => {
+        const triggerType = getCommandTriggerType(data)
+        const selection = getSelectedText()
+
+        void focusAmazonQPanel().then(() => {
+            void provider.webview?.postMessage({
+                command: 'genericCommand',
+                params: { genericCommand, selection, triggerType },
+            })
+        })
+    })
+}
+
+/**
+ * Importing focusAmazonQPanel from aws-core-vscode/amazonq leads to several dependencies down the chain not resolving since AmazonQ chat
+ * is currently only activated on node, but the language server is activated on both web and node.
+ *
+ * Instead, we just create our own as a temporary solution
+ */
+async function focusAmazonQPanel() {
+    await vscode.commands.executeCommand('aws.amazonq.AmazonQChatView.focus')
+    await vscode.commands.executeCommand('aws.amazonq.AmazonCommonAuth.focus')
+}

packages/amazonq/src/lsp/chat/messages.ts

@@ -0,0 +1,225 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+    isValidAuthFollowUpType,
+    INSERT_TO_CURSOR_POSITION,
+    AUTH_FOLLOW_UP_CLICKED,
+    CHAT_OPTIONS,
+    COPY_TO_CLIPBOARD,
+} from '@aws/chat-client-ui-types'
+import {
+    ChatResult,
+    chatRequestType,
+    ChatParams,
+    followUpClickNotificationType,
+    quickActionRequestType,
+    QuickActionResult,
+    QuickActionParams,
+    insertToCursorPositionNotificationType,
+} from '@aws/language-server-runtimes/protocol'
+import { v4 as uuidv4 } from 'uuid'
+import { window } from 'vscode'
+import { Disposable, LanguageClient, Position, State, TextDocumentIdentifier } from 'vscode-languageclient'
+import * as jose from 'jose'
+import { AmazonQChatViewProvider } from './webviewProvider'
+
+export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) {
+    languageClient.onDidChangeState(({ oldState, newState }) => {
+        if (oldState === State.Starting && newState === State.Running) {
+            languageClient.info(
+                'Language client received initializeResult from server:',
+                JSON.stringify(languageClient.initializeResult)
+            )
+
+            const chatOptions = languageClient.initializeResult?.awsServerCapabilities?.chatOptions
+
+            void provider.webview?.postMessage({
+                command: CHAT_OPTIONS,
+                params: chatOptions,
+            })
+        }
+    })
+
+    languageClient.onTelemetry((e) => {
+        languageClient.info(`[VSCode Client] Received telemetry event from server ${JSON.stringify(e)}`)
+    })
+}
+
+export function registerMessageListeners(
+    languageClient: LanguageClient,
+    provider: AmazonQChatViewProvider,
+    encryptionKey: Buffer
+) {
+    provider.webview?.onDidReceiveMessage(async (message) => {
+        languageClient.info(`[VSCode Client]  Received ${JSON.stringify(message)} from chat`)
+
+        switch (message.command) {
+            case COPY_TO_CLIPBOARD:
+                // TODO see what we need to hook this up
+                languageClient.info('[VSCode Client] Copy to clipboard event received')
+                break
+            case INSERT_TO_CURSOR_POSITION: {
+                const editor = window.activeTextEditor
+                let textDocument: TextDocumentIdentifier | undefined = undefined
+                let cursorPosition: Position | undefined = undefined
+                if (editor) {
+                    cursorPosition = editor.selection.active
+                    textDocument = { uri: editor.document.uri.toString() }
+                }
+
+                languageClient.sendNotification(insertToCursorPositionNotificationType.method, {
+                    ...message.params,
+                    cursorPosition,
+                    textDocument,
+                })
+                break
+            }
+            case AUTH_FOLLOW_UP_CLICKED:
+                // TODO hook this into auth
+                languageClient.info('[VSCode Client] AuthFollowUp clicked')
+                break
+            case chatRequestType.method: {
+                const partialResultToken = uuidv4()
+                const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, (partialResult) =>
+                    handlePartialResult<ChatResult>(partialResult, encryptionKey, provider, message.params.tabId)
+                )
+
+                const editor =
+                    window.activeTextEditor ||
+                    window.visibleTextEditors.find((editor) => editor.document.languageId !== 'Log')
+                if (editor) {
+                    message.params.cursorPosition = [editor.selection.active]
+                    message.params.textDocument = { uri: editor.document.uri.toString() }
+                }
+
+                const chatRequest = await encryptRequest<ChatParams>(message.params, encryptionKey)
+                const chatResult = (await languageClient.sendRequest(chatRequestType.method, {
+                    ...chatRequest,
+                    partialResultToken,
+                })) as string | ChatResult
+                void handleCompleteResult<ChatResult>(
+                    chatResult,
+                    encryptionKey,
+                    provider,
+                    message.params.tabId,
+                    chatDisposable
+                )
+                break
+            }
+            case quickActionRequestType.method: {
+                const quickActionPartialResultToken = uuidv4()
+                const quickActionDisposable = languageClient.onProgress(
+                    quickActionRequestType,
+                    quickActionPartialResultToken,
+                    (partialResult) =>
+                        handlePartialResult<QuickActionResult>(
+                            partialResult,
+                            encryptionKey,
+                            provider,
+                            message.params.tabId
+                        )
+                )
+
+                const quickActionRequest = await encryptRequest<QuickActionParams>(message.params, encryptionKey)
+                const quickActionResult = (await languageClient.sendRequest(quickActionRequestType.method, {
+                    ...quickActionRequest,
+                    partialResultToken: quickActionPartialResultToken,
+                })) as string | ChatResult
+                void handleCompleteResult<ChatResult>(
+                    quickActionResult,
+                    encryptionKey,
+                    provider,
+                    message.params.tabId,
+                    quickActionDisposable
+                )
+                break
+            }
+            case followUpClickNotificationType.method:
+                if (!isValidAuthFollowUpType(message.params.followUp.type)) {
+                    languageClient.sendNotification(followUpClickNotificationType.method, message.params)
+                }
+                break
+            default:
+                if (isServerEvent(message.command)) {
+                    languageClient.sendNotification(message.command, message.params)
+                }
+                break
+        }
+    }, undefined)
+}
+
+function isServerEvent(command: string) {
+    return command.startsWith('aws/chat/') || command === 'telemetry/event'
+}
+
+async function encryptRequest<T>(params: T, encryptionKey: Buffer): Promise<{ message: string } | T> {
+    const payload = new TextEncoder().encode(JSON.stringify(params))
+
+    const encryptedMessage = await new jose.CompactEncrypt(payload)
+        .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
+        .encrypt(encryptionKey)
+
+    return { message: encryptedMessage }
+}
+
+async function decodeRequest<T>(request: string, key: Buffer): Promise<T> {
+    const result = await jose.jwtDecrypt(request, key, {
+        clockTolerance: 60, // Allow up to 60 seconds to account for clock differences
+        contentEncryptionAlgorithms: ['A256GCM'],
+        keyManagementAlgorithms: ['dir'],
+    })
+
+    if (!result.payload) {
+        throw new Error('JWT payload not found')
+    }
+    return result.payload as T
+}
+
+/**
+ * Decodes partial chat responses from the language server before sending them to mynah UI
+ */
+async function handlePartialResult<T extends ChatResult>(
+    partialResult: string | T,
+    encryptionKey: Buffer | undefined,
+    provider: AmazonQChatViewProvider,
+    tabId: string
+) {
+    const decryptedMessage =
+        typeof partialResult === 'string' && encryptionKey
+            ? await decodeRequest<T>(partialResult, encryptionKey)
+            : (partialResult as T)
+
+    if (decryptedMessage.body) {
+        void provider.webview?.postMessage({
+            command: chatRequestType.method,
+            params: decryptedMessage,
+            isPartialResult: true,
+            tabId: tabId,
+        })
+    }
+}
+
+/**
+ * Decodes the final chat responses from the language server before sending it to mynah UI.
+ * Once this is called the answer response is finished
+ */
+async function handleCompleteResult<T>(
+    result: string | T,
+    encryptionKey: Buffer | undefined,
+    provider: AmazonQChatViewProvider,
+    tabId: string,
+    disposable: Disposable
+) {
+    const decryptedMessage =
+        typeof result === 'string' && encryptionKey ? await decodeRequest(result, encryptionKey) : result
+
+    void provider.webview?.postMessage({
+        command: chatRequestType.method,
+        params: decryptedMessage,
+        tabId: tabId,
+    })
+    disposable.dispose()
+}