diff --git a/packages/amazonq/src/extensionNode.ts b/packages/amazonq/src/extensionNode.ts index 8224b9ce310..d42fafea058 100644 --- a/packages/amazonq/src/extensionNode.ts +++ b/packages/amazonq/src/extensionNode.ts @@ -25,7 +25,6 @@ import { DevOptions } from 'aws-core-vscode/dev' import { Auth, AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection } from 'aws-core-vscode/auth' import api from './api' import { activate as activateCWChat } from './app/chat/activation' -import { activate as activateInlineChat } from './inlineChat/activation' import { beta } from 'aws-core-vscode/dev' import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications' import { AuthState, AuthUtil } from 'aws-core-vscode/codewhisperer' @@ -73,7 +72,6 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) { } activateAgents() await activateTransformationHub(extContext as ExtContext) - activateInlineChat(context) const authProvider = new CommonAuthViewProvider( context, diff --git a/packages/amazonq/src/inlineChat/activation.ts b/packages/amazonq/src/inlineChat/activation.ts index a42dfdb3e02..01e9f420c05 100644 --- a/packages/amazonq/src/inlineChat/activation.ts +++ b/packages/amazonq/src/inlineChat/activation.ts @@ -5,8 +5,9 @@ import * as vscode from 'vscode' import { InlineChatController } from './controller/inlineChatController' import { registerInlineCommands } from './command/registerInlineCommands' +import { LanguageClient } from 'vscode-languageclient' -export function activate(context: vscode.ExtensionContext) { - const inlineChatController = new InlineChatController(context) +export function activate(context: vscode.ExtensionContext, client: LanguageClient, encryptionKey: Buffer) { + const inlineChatController = new InlineChatController(context, client, encryptionKey) registerInlineCommands(context, inlineChatController) } diff --git a/packages/amazonq/src/inlineChat/controller/inlineChatController.ts b/packages/amazonq/src/inlineChat/controller/inlineChatController.ts index 7ace8d0095e..4eb7c0a7c26 100644 --- a/packages/amazonq/src/inlineChat/controller/inlineChatController.ts +++ b/packages/amazonq/src/inlineChat/controller/inlineChatController.ts @@ -14,6 +14,7 @@ import { CodelensProvider } from '../codeLenses/codeLenseProvider' import { PromptMessage, ReferenceLogController } from 'aws-core-vscode/codewhispererChat' import { CodeWhispererSettings } from 'aws-core-vscode/codewhisperer' import { UserWrittenCodeTracker } from 'aws-core-vscode/codewhisperer' +import { LanguageClient } from 'vscode-languageclient' import { codicon, getIcon, @@ -23,6 +24,7 @@ import { Timeout, textDocumentUtil, isSageMaker, + Experiments, } from 'aws-core-vscode/shared' import { InlineLineAnnotationController } from '../decorations/inlineLineAnnotationController' @@ -33,14 +35,18 @@ export class InlineChatController { private readonly codeLenseProvider: CodelensProvider private readonly referenceLogController = new ReferenceLogController() private readonly inlineLineAnnotationController: InlineLineAnnotationController + private readonly computeDiffAndRenderOnEditor: (query: string) => Promise private userQuery: string | undefined private listeners: vscode.Disposable[] = [] - constructor(context: vscode.ExtensionContext) { - this.inlineChatProvider = new InlineChatProvider() + constructor(context: vscode.ExtensionContext, client: LanguageClient, encryptionKey: Buffer) { + this.inlineChatProvider = new InlineChatProvider(client, encryptionKey) this.inlineChatProvider.onErrorOccured(() => this.handleError()) this.codeLenseProvider = new CodelensProvider(context) this.inlineLineAnnotationController = new InlineLineAnnotationController(context) + this.computeDiffAndRenderOnEditor = Experiments.instance.get('amazonqLSPInlineChat', false) + ? this.computeDiffAndRenderOnEditorLSP.bind(this) + : this.computeDiffAndRenderOnEditorLocal.bind(this) } public async createTask( @@ -206,7 +212,7 @@ export class InlineChatController { await textDocumentUtil.addEofNewline(editor) this.task = await this.createTask(query, editor.document, editor.selection) await this.inlineLineAnnotationController.disable(editor) - await this.computeDiffAndRenderOnEditor(query, editor.document).catch(async (err) => { + await this.computeDiffAndRenderOnEditor(query).catch(async (err) => { getLogger().error('computeDiffAndRenderOnEditor error: %s', (err as Error)?.message) if (err instanceof Error) { void vscode.window.showErrorMessage(`Amazon Q: ${err.message}`) @@ -218,7 +224,46 @@ export class InlineChatController { }) } - private async computeDiffAndRenderOnEditor(query: string, document: vscode.TextDocument) { + private async computeDiffAndRenderOnEditorLSP(query: string) { + if (!this.task) { + return + } + + await this.updateTaskAndLenses(this.task, TaskState.InProgress) + getLogger().info(`inline chat query:\n${query}`) + const uuid = randomUUID() + const message: PromptMessage = { + message: query, + messageId: uuid, + command: undefined, + userIntent: undefined, + tabID: uuid, + } + + const response = await this.inlineChatProvider.processPromptMessageLSP(message) + + // TODO: add tests for this case. + if (!response.body) { + getLogger().warn('Empty body in inline chat response') + await this.handleError() + return + } + + // Update inline diff view + const textDiff = computeDiff(response.body, this.task, false) + const decorations = computeDecorations(this.task) + this.task.decorations = decorations + await this.applyDiff(this.task, textDiff ?? []) + this.decorator.applyDecorations(this.task) + + // Update Codelenses + await this.updateTaskAndLenses(this.task, TaskState.WaitingForDecision) + await setContext('amazonq.inline.codelensShortcutEnabled', true) + this.undoListener(this.task) + } + + // TODO: remove this implementation in favor of LSP + private async computeDiffAndRenderOnEditorLocal(query: string) { if (!this.task) { return } diff --git a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts index e6534d65532..86fe0ac2ade 100644 --- a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts +++ b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts @@ -8,6 +8,8 @@ import { CodeWhispererStreamingServiceException, GenerateAssistantResponseCommandOutput, } from '@amzn/codewhisperer-streaming' +import { LanguageClient } from 'vscode-languageclient' +import { inlineChatRequestType } from '@aws/language-server-runtimes/protocol' import { AuthUtil, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' import { ChatSessionStorage, @@ -25,6 +27,9 @@ import { codeWhispererClient } from 'aws-core-vscode/codewhisperer' import type { InlineChatEvent } from 'aws-core-vscode/codewhisperer' import { InlineTask } from '../controller/inlineTask' import { extractAuthFollowUp } from 'aws-core-vscode/amazonq' +import { InlineChatParams, InlineChatResult } from '@aws/language-server-runtimes-types' +import { decodeRequest, encryptRequest } from '../../lsp/encryption' +import { getCursorState } from '../../lsp/utils' export class InlineChatProvider { private readonly editorContextExtractor: EditorContextExtractor @@ -34,13 +39,52 @@ export class InlineChatProvider { private errorEmitter = new vscode.EventEmitter() public onErrorOccured = this.errorEmitter.event - public constructor() { + public constructor( + private readonly client: LanguageClient, + private readonly encryptionKey: Buffer + ) { this.editorContextExtractor = new EditorContextExtractor() this.userIntentRecognizer = new UserIntentRecognizer() this.sessionStorage = new ChatSessionStorage() this.triggerEventsStorage = new TriggerEventsStorage() } + private getCurrentEditorParams(prompt: string): InlineChatParams { + const editor = vscode.window.activeTextEditor + if (!editor) { + throw new ToolkitError('No active editor') + } + + const documentUri = editor.document.uri.toString() + const cursorState = getCursorState(editor.selections) + return { + prompt: { + prompt, + }, + cursorState, + textDocument: { + uri: documentUri, + }, + } + } + + public async processPromptMessageLSP(message: PromptMessage): Promise { + // TODO: handle partial responses. + getLogger().info('Making inline chat request with message %O', message) + const params = this.getCurrentEditorParams(message.message ?? '') + const inlineChatRequest = await encryptRequest(params, this.encryptionKey) + const response = await this.client.sendRequest(inlineChatRequestType.method, inlineChatRequest) + const decryptedMessage = + typeof response === 'string' && this.encryptionKey + ? await decodeRequest(response, this.encryptionKey) + : response + const result: InlineChatResult = decryptedMessage as InlineChatResult + this.client.info(`Logging response for inline chat ${JSON.stringify(decryptedMessage)}`) + + return result + } + + // TODO: remove in favor of LSP implementation. public async processPromptMessage(message: PromptMessage) { return this.editorContextExtractor .extractContextForTrigger('ChatMessage') diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts index 9578858b708..93ede65fc9a 100644 --- a/packages/amazonq/src/lsp/chat/messages.ts +++ b/packages/amazonq/src/lsp/chat/messages.ts @@ -55,7 +55,6 @@ import { import { v4 as uuidv4 } from 'uuid' import * as vscode from 'vscode' import { Disposable, LanguageClient, Position, TextDocumentIdentifier } from 'vscode-languageclient' -import * as jose from 'jose' import { AmazonQChatViewProvider } from './webviewProvider' import { AuthUtil, ReferenceLogViewProvider } from 'aws-core-vscode/codewhisperer' import { amazonQDiffScheme, AmazonQPromptSettings, messages, openUrl } from 'aws-core-vscode/shared' @@ -68,6 +67,8 @@ import { } from 'aws-core-vscode/amazonq' import { telemetry, TelemetryBase } from 'aws-core-vscode/telemetry' import { isValidResponseError } from './error' +import { decodeRequest, encryptRequest } from '../encryption' +import { getCursorState } from '../utils' export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) { languageClient.info( @@ -99,21 +100,6 @@ export function registerLanguageServerEventListener(languageClient: LanguageClie }) } -function getCursorState(selection: readonly vscode.Selection[]) { - return selection.map((s) => ({ - range: { - start: { - line: s.start.line, - character: s.start.character, - }, - end: { - line: s.end.line, - character: s.end.character, - }, - }, - })) -} - export function registerMessageListeners( languageClient: LanguageClient, provider: AmazonQChatViewProvider, @@ -487,29 +473,6 @@ function isServerEvent(command: string) { return command.startsWith('aws/chat/') || command === 'telemetry/event' } -async function encryptRequest(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(request: string, key: Buffer): Promise { - 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 */ diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index a8cb8d76a40..2db8ab43b0c 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -39,6 +39,7 @@ import { processUtils } from 'aws-core-vscode/shared' import { activate } from './chat/activation' import { AmazonQResourcePaths } from './lspInstaller' import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config' +import { activate as activateInlineChat } from '../inlineChat/activation' const localize = nls.loadMessageBundle() const logger = getLogger('amazonqLsp.lspClient') @@ -182,6 +183,8 @@ export async function startLanguageServer( await activate(client, encryptionKey, resourcePaths.ui) } + activateInlineChat(extensionContext, client, encryptionKey) + const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond) const sendProfileToLsp = async () => { diff --git a/packages/amazonq/src/lsp/encryption.ts b/packages/amazonq/src/lsp/encryption.ts new file mode 100644 index 00000000000..213ee3c1553 --- /dev/null +++ b/packages/amazonq/src/lsp/encryption.ts @@ -0,0 +1,28 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import * as jose from 'jose' + +export async function encryptRequest(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 } +} + +export async function decodeRequest(request: string, key: Buffer): Promise { + 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 +} diff --git a/packages/amazonq/src/lsp/utils.ts b/packages/amazonq/src/lsp/utils.ts new file mode 100644 index 00000000000..f5b010c536b --- /dev/null +++ b/packages/amazonq/src/lsp/utils.ts @@ -0,0 +1,26 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import * as vscode from 'vscode' +import { CursorState } from '@aws/language-server-runtimes-types' + +/** + * Convert from vscode selection type to the general CursorState expected by the AmazonQLSP. + * @param selection + * @returns + */ +export function getCursorState(selection: readonly vscode.Selection[]): CursorState[] { + return selection.map((s) => ({ + range: { + start: { + line: s.start.line, + character: s.start.character, + }, + end: { + line: s.end.line, + character: s.end.character, + }, + }, + })) +} diff --git a/packages/core/src/shared/settings-toolkit.gen.ts b/packages/core/src/shared/settings-toolkit.gen.ts index 59a637a4870..10020cf51f9 100644 --- a/packages/core/src/shared/settings-toolkit.gen.ts +++ b/packages/core/src/shared/settings-toolkit.gen.ts @@ -44,6 +44,7 @@ export const toolkitSettings = { "jsonResourceModification": {}, "amazonqLSP": {}, "amazonqLSPInline": {}, + "amazonqLSPInlineChat": {}, "amazonqChatLSP": {} }, "aws.resources.enabledResources": {}, diff --git a/packages/toolkit/package.json b/packages/toolkit/package.json index aa59eb2f0d0..1b1a1823164 100644 --- a/packages/toolkit/package.json +++ b/packages/toolkit/package.json @@ -255,6 +255,10 @@ "type": "boolean", "default": false }, + "amazonqLSPInlineChat": { + "type": "boolean", + "default": false + }, "amazonqChatLSP": { "type": "boolean", "default": true