feat: working e2e flow of inline chat through language server

Hweinstock · Hweinstock · commit 47d98ca31a16 · 2025-05-06T13:27:19.000-04:00
diff --git 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
@@ -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
@@ -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,
@@ -38,8 +39,8 @@ export class InlineChatController {
     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)
@@ -239,102 +240,26 @@ export class InlineChatController {
             tabID: uuid,
         }
 
-        const requestStart = performance.now()
-        let responseStartLatency: number | undefined
-
-        const response = await this.inlineChatProvider.processPromptMessage(message)
-        this.task.requestId = response?.$metadata.requestId
+        const response = await this.inlineChatProvider.processPromptMessageLSP(message)
 
-        // Deselect all code
-        const editor = vscode.window.activeTextEditor
-        if (editor) {
-            const selection = editor.selection
-            if (!selection.isEmpty) {
-                const cursor = selection.active
-                const newSelection = new vscode.Selection(cursor, cursor)
-                editor.selection = newSelection
-            }
+        // TODO: add tests for this case.
+        if (!response.body) {
+            getLogger().warn('Empty body in inline chat response')
+            await this.handleError()
+            return
         }
 
-        if (response) {
-            let qSuggestedCodeResponse = ''
-            for await (const chatEvent of response.generateAssistantResponseResponse!) {
-                if (
-                    chatEvent.assistantResponseEvent?.content !== undefined &&
-                    chatEvent.assistantResponseEvent.content.length > 0
-                ) {
-                    if (responseStartLatency === undefined) {
-                        responseStartLatency = performance.now() - requestStart
-                    }
-
-                    qSuggestedCodeResponse += chatEvent.assistantResponseEvent.content
-
-                    const transformedResponse = responseTransformer(qSuggestedCodeResponse, this.task, false)
-                    if (transformedResponse) {
-                        const textDiff = computeDiff(transformedResponse, this.task, true)
-                        const decorations = computeDecorations(this.task)
-                        this.task.decorations = decorations
-                        await this.applyDiff(this.task!, textDiff ?? [], {
-                            undoStopBefore: false,
-                            undoStopAfter: false,
-                        })
-                        this.decorator.applyDecorations(this.task)
-                        this.task.previouseDiff = textDiff
-                    }
-                }
-                if (
-                    chatEvent.codeReferenceEvent?.references !== undefined &&
-                    chatEvent.codeReferenceEvent.references.length > 0
-                ) {
-                    this.task.codeReferences = this.task.codeReferences.concat(chatEvent.codeReferenceEvent?.references)
-                    // clear diff if user settings is off for code reference
-                    if (!CodeWhispererSettings.instance.isSuggestionsWithCodeReferencesEnabled()) {
-                        await this.rejectAllChanges(this.task, false)
-                        void vscode.window.showInformationMessage(
-                            'Your settings do not allow code generation with references.'
-                        )
-                        await this.updateTaskAndLenses(this.task, TaskState.Complete)
-                        return
-                    }
-                }
-                if (chatEvent.error) {
-                    getLogger().error('generateAssistantResponse stream error: %s', chatEvent.error)
-                    await this.rejectAllChanges(this.task, false)
-                    void vscode.window.showErrorMessage(`Amazon Q: ${chatEvent.error.message}`)
-                    await this.updateTaskAndLenses(this.task, TaskState.Complete)
-                    return
-                }
-            }
-
-            if (this.task) {
-                // Unclear why we need to check if task is defined, but occasionally an error occurs otherwise
-                this.task.responseStartLatency = responseStartLatency
-                this.task.responseEndLatency = performance.now() - requestStart
-            }
-            getLogger().info(`qSuggestedCodeResponse:\n${qSuggestedCodeResponse}`)
-            const transformedResponse = responseTransformer(qSuggestedCodeResponse, this.task, true)
-            if (transformedResponse) {
-                const textDiff = computeDiff(transformedResponse, this.task, false)
-                const decorations = computeDecorations(this.task)
-                this.task.decorations = decorations
-                await this.applyDiff(this.task, textDiff ?? [])
-                this.decorator.applyDecorations(this.task)
-                await this.updateTaskAndLenses(this.task, TaskState.WaitingForDecision)
-                await setContext('amazonq.inline.codelensShortcutEnabled', true)
-                this.undoListener(this.task)
-            } else {
-                void messages.showMessageWithCancel(
-                    'No suggestions from Q, please try different instructions.',
-                    new Timeout(5000)
-                )
-                await this.updateTaskAndLenses(this.task, TaskState.Complete)
-                await this.inlineQuickPick(this.userQuery)
-                await this.handleError()
-            }
-        }
+        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)
+        await this.updateTaskAndLenses(this.task, TaskState.WaitingForDecision)
+        await setContext('amazonq.inline.codelensShortcutEnabled', true)
+        this.undoListener(this.task)
     }
 
-    // TODO: remove this implementation in favor or LSP
+    // 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
@@ -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,51 @@ 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()
+
+        return {
+            prompt: {
+                prompt,
+            },
+            cursorState: getCursorState(editor.selections),
+            textDocument: {
+                uri: documentUri,
+            },
+        }
+    }
+
+    public async processPromptMessageLSP(message: PromptMessage): Promise<InlineChatResult> {
+        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')
diff --git 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 activateInline } 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)
         }
 
+        activateInline(extensionContext, client, encryptionKey)
+
         const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond)
 
         const sendProfileToLsp = async () => {
