feat(amazonq): add experiment for basic e2e flow of inline chat through Flare. (#7235)

## Problem
This is the initial set of work required to get inline chat running
through flare.

## Solution
- **add a feature flag for inline chat**: this allows testing of the two
implementations side-by-side by flipping the feature flag.
- **move general utils out of chat**: stuff like encryption and
editorState can all be reused.
- **render full diff response from inline chat**: this does not include
progress updates from the language server.

## Testing and Verification


https://github.com/user-attachments/assets/0dff58b7-40f7-487d-9f9e-d58610201041



## Future Work / Next Steps
- ensure telemetry is still being emitted. 
- add tests for new flow. (there aren't any for the existing one)
- handle partial events from the language server. 

## Known Bugs
- selecting part of a line will cause the text to insert mid-line 
- running inline-chat without a selection causes the entire file to be
copied (This is in JB, Eclipse Prod, but IMO it makes the feature
unusable).


---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: Hweinstock

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,

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)
 }

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<void>
     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
         }

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<void>()
     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<InlineChatResult> {
+        // TODO: handle partial responses.
+        getLogger().info('Making inline chat request with message %O', message)
+        const params = this.getCurrentEditorParams(message.message ?? '')
+        const inlineChatRequest = await encryptRequest<InlineChatParams>(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')

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<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
  */

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 () => {

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<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 }
+}
+
+export 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
+}

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,
+            },
+        },
+    }))
+}

packages/core/src/shared/settings-toolkit.gen.ts

@@ -44,6 +44,7 @@ export const toolkitSettings = {
         "jsonResourceModification": {},
         "amazonqLSP": {},
         "amazonqLSPInline": {},
+        "amazonqLSPInlineChat": {},
         "amazonqChatLSP": {}
     },
     "aws.resources.enabledResources": {},

packages/toolkit/package.json

@@ -255,6 +255,10 @@
                             "type": "boolean",
                             "default": false
                         },
+                        "amazonqLSPInlineChat": {
+                            "type": "boolean",
+                            "default": false
+                        },
                         "amazonqChatLSP": {
                             "type": "boolean",
                             "default": true