Rebase

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

@@ -119,9 +119,10 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
         extensionContext: context,
     }
 
+    await activateAmazonqLsp(context)
+
     // This contains every lsp agnostic things (auth, security scan, code scan)
     await activateCodeWhisperer(extContext as ExtContext)
-    await activateAmazonqLsp(context)
 
     if (!Experiments.instance.get('amazonqLSPInline', false)) {
         await activateInlineCompletion()

packages/amazonq/src/extension.ts

@@ -84,7 +84,7 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
         authEnabledConnections: (await AuthUtil.instance.getAuthFormIds()).join(','),
     })
 
-    void activateNotifications(context, getAuthState)
+    // void activateNotifications(context, getAuthState)
 }
 
 async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {

packages/amazonq/src/extensionNode.ts

@@ -6,7 +6,7 @@
 import vscode from 'vscode'
 import { clientId, encryptionKey, startLanguageServer } from './client'
 import { AmazonQLspInstaller } from './lspInstaller'
-import { lspSetupStage, ToolkitError } from 'aws-core-vscode/shared'
+import { lspSetupStage, messages, ToolkitError } from 'aws-core-vscode/shared'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { auth2 } from 'aws-core-vscode/auth'
 
@@ -20,6 +20,6 @@ export async function activate(ctx: vscode.ExtensionContext) {
         await AuthUtil.instance.restore()
     } catch (err) {
         const e = err as ToolkitError
-        void vscode.window.showInformationMessage(`Unable to launch amazonq language server: ${e.message}`)
+        void messages.showViewLogsMessage(`Failed to launch Amazon Q language server: ${e.message}`)
     }
 }

packages/amazonq/src/lsp/activation.ts

@@ -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

@@ -50,7 +50,7 @@ export class AmazonQChatViewProvider implements WebviewViewProvider {
     ) {
         this.webview = webviewView.webview
 
-        const lspDir = Uri.parse(LanguageServerResolver.defaultDir)
+        const lspDir = Uri.parse(LanguageServerResolver.defaultDir())
         webviewView.webview.options = {
             enableScripts: true,
             enableCommandUris: true,
@@ -64,7 +64,7 @@ export class AmazonQChatViewProvider implements WebviewViewProvider {
     }
 
     private async getWebviewContent(mynahUIPath: string) {
-        const disclaimerAcknowledged = AmazonQPromptSettings.instance.isPromptEnabled('amazonQChatDisclaimer')
+        const disclaimerAcknowledged = !AmazonQPromptSettings.instance.isPromptEnabled('amazonQChatDisclaimer')
         return `
         <!DOCTYPE html>
         <html lang="en">

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

@@ -31,12 +31,14 @@ import {
     Commands,
     openUrl,
     getLogger,
+    validateNodeExe,
 } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
 import { auth2 } from 'aws-core-vscode/auth'
 
 const localize = nls.loadMessageBundle()
+const logger = getLogger('amazonqLsp.lspClient')
 
 export const clientId = 'amazonq'
 export const clientName = oidcClientName()
@@ -50,22 +52,25 @@ export async function startLanguageServer(
 
     const serverModule = resourcePaths.lsp
 
+    const argv = [
+        '--nolazy',
+        '--preserve-symlinks',
+        '--stdio',
+        '--pre-init-encryption',
+        '--set-credentials-encryption-key',
+    ]
     const serverOptions = createServerOptions({
         encryptionKey,
         executable: resourcePaths.node,
         serverModule,
-        execArgv: [
-            '--nolazy',
-            '--preserve-symlinks',
-            '--stdio',
-            '--pre-init-encryption',
-            '--set-credentials-encryption-key',
-        ],
+        execArgv: argv,
     })
 
     const documentSelector = [{ scheme: 'file', language: '*' }]
     const traceServerEnabled = Settings.instance.isSet(`${clientId}.trace.server`)
 
+    await validateNodeExe(resourcePaths.node, resourcePaths.lsp, argv, logger)
+
     // Options to control the language client
     const clientOptions: LanguageClientOptions = {
         // Register the server for json documents