Merge branch 'feature/amazonqLSP-auth' into feature/amazonqLSP-auth

Author: opieter-aws

package-lock.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Users might be bound to a customization which they dont have access with the selected profile and it causes service throwing 403 when using inline suggestion and chat features"
+}

packages/amazonq/.changes/next-release/Bug Fix-a6e9ce99-842b-4d64-b9b0-967383d7acb9.json

@@ -11,6 +11,7 @@ import {
     COPY_TO_CLIPBOARD,
     AuthFollowUpType,
     DISCLAIMER_ACKNOWLEDGED,
+    UiMessageResultParams,
 } from '@aws/chat-client-ui-types'
 import {
     ChatResult,
@@ -21,30 +22,42 @@ import {
     QuickActionResult,
     QuickActionParams,
     insertToCursorPositionNotificationType,
+    ErrorCodes,
+    ResponseError,
+    openTabRequestType,
+    getSerializedChatRequestType,
+    listConversationsRequestType,
+    conversationClickRequestType,
+    ShowSaveFileDialogRequestType,
+    ShowSaveFileDialogParams,
+    LSPErrorCodes,
+    tabBarActionRequestType,
 } from '@aws/language-server-runtimes/protocol'
 import { v4 as uuidv4 } from 'uuid'
 import * as vscode from 'vscode'
-import { Disposable, LanguageClient, Position, State, TextDocumentIdentifier } from 'vscode-languageclient'
+import { Disposable, LanguageClient, Position, TextDocumentIdentifier } from 'vscode-languageclient'
 import * as jose from 'jose'
 import { AmazonQChatViewProvider } from './webviewProvider'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { AmazonQPromptSettings, messages } from 'aws-core-vscode/shared'
 
 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)
-            )
+    languageClient.info(
+        'Language client received initializeResult from server:',
+        JSON.stringify(languageClient.initializeResult)
+    )
 
-            const chatOptions = languageClient.initializeResult?.awsServerCapabilities?.chatOptions
+    const chatOptions = languageClient.initializeResult?.awsServerCapabilities?.chatOptions
 
-            void provider.webview?.postMessage({
-                command: CHAT_OPTIONS,
-                params: chatOptions,
-            })
-        }
+    // Enable the history/export feature flags
+    chatOptions.history = true
+    chatOptions.export = true
+
+    provider.onDidResolveWebview(() => {
+        void provider.webview?.postMessage({
+            command: CHAT_OPTIONS,
+            params: chatOptions,
+        })
     })
 
     languageClient.onTelemetry((e) => {
@@ -60,6 +73,7 @@ export function registerMessageListeners(
     provider.webview?.onDidReceiveMessage(async (message) => {
         languageClient.info(`[VSCode Client]  Received ${JSON.stringify(message)} from chat`)
 
+        const webview = provider.webview
         switch (message.command) {
             case COPY_TO_CLIPBOARD:
                 languageClient.info('[VSCode Client] Copy to clipboard event received')
@@ -172,6 +186,11 @@ export function registerMessageListeners(
                 )
                 break
             }
+            case listConversationsRequestType.method:
+            case conversationClickRequestType.method:
+            case tabBarActionRequestType.method:
+                await resolveChatResponse(message.command, message.params, languageClient, webview)
+                break
             case followUpClickNotificationType.method:
                 if (!isValidAuthFollowUpType(message.params.followUp.type)) {
                     languageClient.sendNotification(followUpClickNotificationType.method, message.params)
@@ -184,6 +203,89 @@ export function registerMessageListeners(
                 break
         }
     }, undefined)
+
+    const registerHandlerWithResponseRouter = (command: string) => {
+        const handler = async (params: any, _: any) => {
+            const mapErrorType = (type: string | undefined): number => {
+                switch (type) {
+                    case 'InvalidRequest':
+                        return ErrorCodes.InvalidRequest
+                    case 'InternalError':
+                        return ErrorCodes.InternalError
+                    case 'UnknownError':
+                    default:
+                        return ErrorCodes.UnknownErrorCode
+                }
+            }
+            const requestId = uuidv4()
+
+            void provider.webview?.postMessage({
+                requestId: requestId,
+                command: command,
+                params: params,
+            })
+            const responsePromise = new Promise<UiMessageResultParams | undefined>((resolve, reject) => {
+                const timeout = setTimeout(() => {
+                    disposable?.dispose()
+                    reject(new Error('Request timed out'))
+                }, 30000)
+
+                const disposable = provider.webview?.onDidReceiveMessage((message: any) => {
+                    if (message.requestId === requestId) {
+                        clearTimeout(timeout)
+                        disposable?.dispose()
+                        resolve(message.params)
+                    }
+                })
+            })
+
+            const result = await responsePromise
+
+            if (result?.success) {
+                return result.result
+            } else {
+                return new ResponseError(
+                    mapErrorType(result?.error.type),
+                    result?.error.message ?? 'No response from client'
+                )
+            }
+        }
+
+        languageClient.onRequest(command, handler)
+    }
+
+    registerHandlerWithResponseRouter(openTabRequestType.method)
+    registerHandlerWithResponseRouter(getSerializedChatRequestType.method)
+
+    languageClient.onRequest(ShowSaveFileDialogRequestType.method, async (params: ShowSaveFileDialogParams) => {
+        const filters: Record<string, string[]> = {}
+        const formatMappings = [
+            { format: 'markdown', key: 'Markdown', extensions: ['md'] },
+            { format: 'html', key: 'HTML', extensions: ['html'] },
+        ]
+
+        for (const format of params.supportedFormats ?? []) {
+            const mapping = formatMappings.find((m) => m.format === format)
+            if (mapping) {
+                filters[mapping.key] = mapping.extensions
+            }
+        }
+
+        const saveAtUri = params.defaultUri ? vscode.Uri.parse(params.defaultUri) : vscode.Uri.file('export-chat.md')
+        const targetUri = await vscode.window.showSaveDialog({
+            filters,
+            defaultUri: saveAtUri,
+            title: 'Export',
+        })
+
+        if (!targetUri) {
+            return new ResponseError(LSPErrorCodes.RequestFailed, 'Export failed')
+        }
+
+        return {
+            targetUri: targetUri.toString(),
+        }
+    })
 }
 
 function isServerEvent(command: string) {
@@ -258,3 +360,16 @@ async function handleCompleteResult<T>(
     })
     disposable.dispose()
 }
+
+async function resolveChatResponse(
+    requestMethod: string,
+    params: any,
+    languageClient: LanguageClient,
+    webview: vscode.Webview | undefined
+) {
+    const result = await languageClient.sendRequest(requestMethod, params)
+    void webview?.postMessage({
+        command: requestMethod,
+        params: result,
+    })
+}

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

@@ -26,6 +26,7 @@ describe('registerMessageListeners', () => {
             info: sandbox.stub(),
             error: errorStub,
             sendNotification: sandbox.stub(),
+            onRequest: sandbox.stub(),
         } as unknown as LanguageClient
 
         provider = {

packages/amazonq/test/unit/amazonq/lsp/chat/messages.test.ts

@@ -441,8 +441,8 @@
     "devDependencies": {
         "@aws-sdk/types": "^3.13.1",
         "@aws/chat-client-ui-types": "^0.1.12",
-        "@aws/language-server-runtimes": "^0.2.49",
-        "@aws/language-server-runtimes-types": "^0.1.10",
+        "@aws/language-server-runtimes": "^0.2.58",
+        "@aws/language-server-runtimes-types": "^0.1.13",
         "@cspotcode/source-map-support": "^0.8.1",
         "@sinonjs/fake-timers": "^10.0.2",
         "@types/adm-zip": "^0.4.34",

packages/core/package.json

@@ -11,6 +11,7 @@ import { DiffTreeFileInfo } from '../webview/ui/diffTree/types'
 import { Messenger } from './connector/baseMessenger'
 import { FeatureClient } from '../client/client'
 import { TelemetryHelper } from '../util/telemetryHelper'
+import { MynahUI } from '@aws/mynah-ui'
 
 export enum FollowUpTypes {
     // UnitTestGeneration
@@ -164,3 +165,10 @@ export enum MetricDataResult {
     Error = 'Error',
     LlmFailure = 'LLMFailure',
 }
+
+/**
+ * Note: Passing a reference around allows us to lazily inject mynah UI into
+ * connectors and handlers. This is done to supported "hybrid chat", which
+ * injects mynah UI _after_ the connector has already been created
+ */
+export type MynahUIRef = { mynahUI: MynahUI | undefined }

packages/core/src/amazonq/commons/types.ts

@@ -8,21 +8,21 @@ import { Connector } from '../connector'
 import { TabsStorage } from '../storages/tabsStorage'
 import { WelcomeFollowupType } from '../apps/amazonqCommonsConnector'
 import { AuthFollowUpType } from './generator'
-import { FollowUpTypes } from '../../../commons/types'
+import { FollowUpTypes, MynahUIRef } from '../../../commons/types'
 
 export interface FollowUpInteractionHandlerProps {
-    mynahUI: MynahUI
+    mynahUIRef: MynahUIRef
     connector: Connector
     tabsStorage: TabsStorage
 }
 
 export class FollowUpInteractionHandler {
-    private mynahUI: MynahUI
+    private mynahUIRef: MynahUIRef
     private connector: Connector
     private tabsStorage: TabsStorage
 
     constructor(props: FollowUpInteractionHandlerProps) {
-        this.mynahUI = props.mynahUI
+        this.mynahUIRef = props.mynahUIRef
         this.connector = props.connector
         this.tabsStorage = props.tabsStorage
     }
@@ -41,6 +41,11 @@ export class FollowUpInteractionHandler {
             this.connector.help(tabID)
             return
         }
+
+        if (!this.mynahUI) {
+            return
+        }
+
         // we need to check if there is a prompt
         // which will cause an api call
         // then we can set the loading state to true
@@ -70,7 +75,7 @@ export class FollowUpInteractionHandler {
         }
 
         const addChatItem = (tabID: string, messageId: string, options: any[]) => {
-            this.mynahUI.addChatItem(tabID, {
+            this.mynahUI?.addChatItem(tabID, {
                 type: ChatItemType.ANSWER_PART,
                 messageId,
                 followUp: {
@@ -150,7 +155,7 @@ export class FollowUpInteractionHandler {
 
     public onWelcomeFollowUpClicked(tabID: string, welcomeFollowUpType: WelcomeFollowupType) {
         if (welcomeFollowUpType === 'continue-to-chat') {
-            this.mynahUI.addChatItem(tabID, {
+            this.mynahUI?.addChatItem(tabID, {
                 type: ChatItemType.ANSWER,
                 body: 'Ok, please write your question below.',
             })
@@ -159,4 +164,8 @@ export class FollowUpInteractionHandler {
             return
         }
     }
+
+    private get mynahUI(): MynahUI | undefined {
+        return this.mynahUIRef.mynahUI
+    }
 }