|
| 1 | +/*! |
| 2 | + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +import { |
| 7 | + isValidAuthFollowUpType, |
| 8 | + INSERT_TO_CURSOR_POSITION, |
| 9 | + AUTH_FOLLOW_UP_CLICKED, |
| 10 | + CHAT_OPTIONS, |
| 11 | + COPY_TO_CLIPBOARD, |
| 12 | +} from '@aws/chat-client-ui-types' |
| 13 | +import { |
| 14 | + ChatResult, |
| 15 | + chatRequestType, |
| 16 | + ChatParams, |
| 17 | + followUpClickNotificationType, |
| 18 | + quickActionRequestType, |
| 19 | + QuickActionResult, |
| 20 | + QuickActionParams, |
| 21 | + insertToCursorPositionNotificationType, |
| 22 | +} from '@aws/language-server-runtimes/protocol' |
| 23 | +import { v4 as uuidv4 } from 'uuid' |
| 24 | +import { window } from 'vscode' |
| 25 | +import { Disposable, LanguageClient, Position, State, TextDocumentIdentifier } from 'vscode-languageclient' |
| 26 | +import * as jose from 'jose' |
| 27 | +import { AmazonQChatViewProvider } from './webviewProvider' |
| 28 | + |
| 29 | +export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) { |
| 30 | + languageClient.onDidChangeState(({ oldState, newState }) => { |
| 31 | + if (oldState === State.Starting && newState === State.Running) { |
| 32 | + languageClient.info( |
| 33 | + 'Language client received initializeResult from server:', |
| 34 | + JSON.stringify(languageClient.initializeResult) |
| 35 | + ) |
| 36 | + |
| 37 | + const chatOptions = languageClient.initializeResult?.awsServerCapabilities?.chatOptions |
| 38 | + |
| 39 | + void provider.webview?.postMessage({ |
| 40 | + command: CHAT_OPTIONS, |
| 41 | + params: chatOptions, |
| 42 | + }) |
| 43 | + } |
| 44 | + }) |
| 45 | + |
| 46 | + languageClient.onTelemetry((e) => { |
| 47 | + languageClient.info(`[VSCode Client] Received telemetry event from server ${JSON.stringify(e)}`) |
| 48 | + }) |
| 49 | +} |
| 50 | + |
| 51 | +export function registerMessageListeners( |
| 52 | + languageClient: LanguageClient, |
| 53 | + provider: AmazonQChatViewProvider, |
| 54 | + encryptionKey: Buffer |
| 55 | +) { |
| 56 | + provider.webview?.onDidReceiveMessage(async (message) => { |
| 57 | + languageClient.info(`[VSCode Client] Received ${JSON.stringify(message)} from chat`) |
| 58 | + |
| 59 | + switch (message.command) { |
| 60 | + case COPY_TO_CLIPBOARD: |
| 61 | + languageClient.info('[VSCode Client] Copy to clipboard event received') |
| 62 | + break |
| 63 | + case INSERT_TO_CURSOR_POSITION: { |
| 64 | + const editor = window.activeTextEditor |
| 65 | + let textDocument: TextDocumentIdentifier | undefined = undefined |
| 66 | + let cursorPosition: Position | undefined = undefined |
| 67 | + if (editor) { |
| 68 | + cursorPosition = editor.selection.active |
| 69 | + textDocument = { uri: editor.document.uri.toString() } |
| 70 | + } |
| 71 | + |
| 72 | + languageClient.sendNotification(insertToCursorPositionNotificationType.method, { |
| 73 | + ...message.params, |
| 74 | + cursorPosition, |
| 75 | + textDocument, |
| 76 | + }) |
| 77 | + break |
| 78 | + } |
| 79 | + case AUTH_FOLLOW_UP_CLICKED: |
| 80 | + languageClient.info('[VSCode Client] AuthFollowUp clicked') |
| 81 | + break |
| 82 | + case chatRequestType.method: { |
| 83 | + const partialResultToken = uuidv4() |
| 84 | + const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, (partialResult) => |
| 85 | + handlePartialResult<ChatResult>(partialResult, encryptionKey, provider, message.params.tabId) |
| 86 | + ) |
| 87 | + |
| 88 | + const editor = |
| 89 | + window.activeTextEditor || |
| 90 | + window.visibleTextEditors.find((editor) => editor.document.languageId !== 'Log') |
| 91 | + if (editor) { |
| 92 | + message.params.cursorPosition = [editor.selection.active] |
| 93 | + message.params.textDocument = { uri: editor.document.uri.toString() } |
| 94 | + } |
| 95 | + |
| 96 | + const chatRequest = await encryptRequest<ChatParams>(message.params, encryptionKey) |
| 97 | + const chatResult = (await languageClient.sendRequest(chatRequestType.method, { |
| 98 | + ...chatRequest, |
| 99 | + partialResultToken, |
| 100 | + })) as string | ChatResult |
| 101 | + void handleCompleteResult<ChatResult>( |
| 102 | + chatResult, |
| 103 | + encryptionKey, |
| 104 | + provider, |
| 105 | + message.params.tabId, |
| 106 | + chatDisposable |
| 107 | + ) |
| 108 | + break |
| 109 | + } |
| 110 | + case quickActionRequestType.method: { |
| 111 | + const quickActionPartialResultToken = uuidv4() |
| 112 | + const quickActionDisposable = languageClient.onProgress( |
| 113 | + quickActionRequestType, |
| 114 | + quickActionPartialResultToken, |
| 115 | + (partialResult) => |
| 116 | + handlePartialResult<QuickActionResult>( |
| 117 | + partialResult, |
| 118 | + encryptionKey, |
| 119 | + provider, |
| 120 | + message.params.tabId |
| 121 | + ) |
| 122 | + ) |
| 123 | + |
| 124 | + const quickActionRequest = await encryptRequest<QuickActionParams>(message.params, encryptionKey) |
| 125 | + const quickActionResult = (await languageClient.sendRequest(quickActionRequestType.method, { |
| 126 | + ...quickActionRequest, |
| 127 | + partialResultToken: quickActionPartialResultToken, |
| 128 | + })) as string | ChatResult |
| 129 | + void handleCompleteResult<ChatResult>( |
| 130 | + quickActionResult, |
| 131 | + encryptionKey, |
| 132 | + provider, |
| 133 | + message.params.tabId, |
| 134 | + quickActionDisposable |
| 135 | + ) |
| 136 | + break |
| 137 | + } |
| 138 | + case followUpClickNotificationType.method: |
| 139 | + if (!isValidAuthFollowUpType(message.params.followUp.type)) { |
| 140 | + languageClient.sendNotification(followUpClickNotificationType.method, message.params) |
| 141 | + } |
| 142 | + break |
| 143 | + default: |
| 144 | + if (isServerEvent(message.command)) { |
| 145 | + languageClient.sendNotification(message.command, message.params) |
| 146 | + } |
| 147 | + break |
| 148 | + } |
| 149 | + }, undefined) |
| 150 | +} |
| 151 | + |
| 152 | +function isServerEvent(command: string) { |
| 153 | + return command.startsWith('aws/chat/') || command === 'telemetry/event' |
| 154 | +} |
| 155 | + |
| 156 | +async function encryptRequest<T>(params: T, encryptionKey: Buffer): Promise<{ message: string } | T> { |
| 157 | + const payload = new TextEncoder().encode(JSON.stringify(params)) |
| 158 | + |
| 159 | + const encryptedMessage = await new jose.CompactEncrypt(payload) |
| 160 | + .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) |
| 161 | + .encrypt(encryptionKey) |
| 162 | + |
| 163 | + return { message: encryptedMessage } |
| 164 | +} |
| 165 | + |
| 166 | +async function decodeRequest<T>(request: string, key: Buffer): Promise<T> { |
| 167 | + const result = await jose.jwtDecrypt(request, key, { |
| 168 | + clockTolerance: 60, // Allow up to 60 seconds to account for clock differences |
| 169 | + contentEncryptionAlgorithms: ['A256GCM'], |
| 170 | + keyManagementAlgorithms: ['dir'], |
| 171 | + }) |
| 172 | + |
| 173 | + if (!result.payload) { |
| 174 | + throw new Error('JWT payload not found') |
| 175 | + } |
| 176 | + return result.payload as T |
| 177 | +} |
| 178 | + |
| 179 | +async function handlePartialResult<T extends ChatResult>( |
| 180 | + partialResult: string | T, |
| 181 | + encryptionKey: Buffer | undefined, |
| 182 | + provider: AmazonQChatViewProvider, |
| 183 | + tabId: string |
| 184 | +) { |
| 185 | + const decryptedMessage = |
| 186 | + typeof partialResult === 'string' && encryptionKey |
| 187 | + ? await decodeRequest<T>(partialResult, encryptionKey) |
| 188 | + : (partialResult as T) |
| 189 | + |
| 190 | + if (decryptedMessage.body) { |
| 191 | + void provider.webview?.postMessage({ |
| 192 | + command: chatRequestType.method, |
| 193 | + params: decryptedMessage, |
| 194 | + isPartialResult: true, |
| 195 | + tabId: tabId, |
| 196 | + }) |
| 197 | + } |
| 198 | +} |
| 199 | + |
| 200 | +async function handleCompleteResult<T>( |
| 201 | + result: string | T, |
| 202 | + encryptionKey: Buffer | undefined, |
| 203 | + provider: AmazonQChatViewProvider, |
| 204 | + tabId: string, |
| 205 | + disposable: Disposable |
| 206 | +) { |
| 207 | + const decryptedMessage = |
| 208 | + typeof result === 'string' && encryptionKey ? await decodeRequest(result, encryptionKey) : result |
| 209 | + |
| 210 | + void provider.webview?.postMessage({ |
| 211 | + command: chatRequestType.method, |
| 212 | + params: decryptedMessage, |
| 213 | + tabId: tabId, |
| 214 | + }) |
| 215 | + disposable.dispose() |
| 216 | +} |
0 commit comments