From 54ffbde004b215eaef58179060dfae572211d5b2 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Wed, 14 May 2025 14:13:05 -0400 Subject: [PATCH 1/6] refactor: minor, move code in to separate function Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/client.ts | 215 +++++++++++++++-------------- 1 file changed, 113 insertions(+), 102 deletions(-) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 5559afb9f1d..089f8c2e641 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -163,117 +163,128 @@ export async function startLanguageServer( const auth = new AmazonQLspAuth(client) - return client.onReady().then(async () => { - await auth.refreshConnection() + await client.onReady() - if (Experiments.instance.get('amazonqLSPInline', false)) { - const inlineManager = new InlineCompletionManager(client) - inlineManager.registerInlineCompletion() - toDispose.push( - inlineManager, - Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => { - await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger') - }), - vscode.workspace.onDidCloseTextDocument(async () => { - await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion') - }) - ) - } + await postStartLanguageServer(auth, client, resourcePaths, toDispose) - if (Experiments.instance.get('amazonqChatLSP', true)) { - await activate(client, encryptionKey, resourcePaths.ui) - } + return client +} - const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond) +async function postStartLanguageServer( + auth: AmazonQLspAuth, + client: LanguageClient, + resourcePaths: AmazonQResourcePaths, + toDispose: vscode.Disposable[] +) { + await auth.refreshConnection() - const sendProfileToLsp = async () => { - try { - const result = await client.sendRequest(updateConfigurationRequestType.method, { - section: 'aws.q', - settings: { - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }, - }) - client.info( - `Client: Updated Amazon Q Profile ${AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn} to Amazon Q LSP`, - result - ) - } catch (err) { - client.error('Error when setting Q Developer Profile to Amazon Q LSP', err) - } + if (Experiments.instance.get('amazonqLSPInline', false)) { + const inlineManager = new InlineCompletionManager(client) + inlineManager.registerInlineCompletion() + toDispose.push( + inlineManager, + Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => { + await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger') + }), + vscode.workspace.onDidCloseTextDocument(async () => { + await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion') + }) + ) + } + + if (Experiments.instance.get('amazonqChatLSP', true)) { + await activate(client, encryptionKey, resourcePaths.ui) + } + + const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond) + + const sendProfileToLsp = async () => { + try { + const result = await client.sendRequest(updateConfigurationRequestType.method, { + section: 'aws.q', + settings: { + profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + }, + }) + client.info( + `Client: Updated Amazon Q Profile ${AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn} to Amazon Q LSP`, + result + ) + } catch (err) { + client.error('Error when setting Q Developer Profile to Amazon Q LSP', err) } + } - // send profile to lsp once. - void sendProfileToLsp() + // send profile to lsp once. + void sendProfileToLsp() - toDispose.push( - AuthUtil.instance.auth.onDidChangeActiveConnection(async () => { - await auth.refreshConnection() - }), - AuthUtil.instance.auth.onDidDeleteConnection(async () => { - client.sendNotification(notificationTypes.deleteBearerToken.method) - }), - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp), - vscode.commands.registerCommand('aws.amazonq.getWorkspaceId', async () => { - const requestType = new RequestType( - 'aws/getConfigurationFromServer' - ) - const workspaceIdResp = await client.sendRequest(requestType.method, { - section: 'aws.q.workspaceContext', - }) - return workspaceIdResp - }), - vscode.workspace.onDidCreateFiles((e) => { - client.sendNotification('workspace/didCreateFiles', { - files: e.files.map((it) => { - return { uri: it.fsPath } - }), - } as CreateFilesParams) - }), - vscode.workspace.onDidDeleteFiles((e) => { - client.sendNotification('workspace/didDeleteFiles', { - files: e.files.map((it) => { - return { uri: it.fsPath } + toDispose.push( + AuthUtil.instance.auth.onDidChangeActiveConnection(async () => { + await auth.refreshConnection() + }), + AuthUtil.instance.auth.onDidDeleteConnection(async () => { + client.sendNotification(notificationTypes.deleteBearerToken.method) + }), + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp), + vscode.commands.registerCommand('aws.amazonq.getWorkspaceId', async () => { + const requestType = new RequestType( + 'aws/getConfigurationFromServer' + ) + const workspaceIdResp = await client.sendRequest(requestType.method, { + section: 'aws.q.workspaceContext', + }) + return workspaceIdResp + }), + vscode.workspace.onDidCreateFiles((e) => { + client.sendNotification('workspace/didCreateFiles', { + files: e.files.map((it) => { + return { uri: it.fsPath } + }), + } as CreateFilesParams) + }), + vscode.workspace.onDidDeleteFiles((e) => { + client.sendNotification('workspace/didDeleteFiles', { + files: e.files.map((it) => { + return { uri: it.fsPath } + }), + } as DeleteFilesParams) + }), + vscode.workspace.onDidRenameFiles((e) => { + client.sendNotification('workspace/didRenameFiles', { + files: e.files.map((it) => { + return { oldUri: it.oldUri.fsPath, newUri: it.newUri.fsPath } + }), + } as RenameFilesParams) + }), + vscode.workspace.onDidSaveTextDocument((e) => { + client.sendNotification('workspace/didSaveTextDocument', { + textDocument: { + uri: e.uri.fsPath, + }, + } as DidSaveTextDocumentParams) + }), + vscode.workspace.onDidChangeWorkspaceFolders((e) => { + client.sendNotification('workspace/didChangeWorkspaceFolder', { + event: { + added: e.added.map((it) => { + return { + name: it.name, + uri: it.uri.fsPath, + } as WorkspaceFolder }), - } as DeleteFilesParams) - }), - vscode.workspace.onDidRenameFiles((e) => { - client.sendNotification('workspace/didRenameFiles', { - files: e.files.map((it) => { - return { oldUri: it.oldUri.fsPath, newUri: it.newUri.fsPath } + removed: e.removed.map((it) => { + return { + name: it.name, + uri: it.uri.fsPath, + } as WorkspaceFolder }), - } as RenameFilesParams) - }), - vscode.workspace.onDidSaveTextDocument((e) => { - client.sendNotification('workspace/didSaveTextDocument', { - textDocument: { - uri: e.uri.fsPath, - }, - } as DidSaveTextDocumentParams) - }), - vscode.workspace.onDidChangeWorkspaceFolders((e) => { - client.sendNotification('workspace/didChangeWorkspaceFolder', { - event: { - added: e.added.map((it) => { - return { - name: it.name, - uri: it.uri.fsPath, - } as WorkspaceFolder - }), - removed: e.removed.map((it) => { - return { - name: it.name, - uri: it.uri.fsPath, - } as WorkspaceFolder - }), - }, - } as DidChangeWorkspaceFoldersParams) - }), - { dispose: () => clearInterval(refreshInterval) }, - // Set this inside onReady so that it only triggers on subsequent language server starts (not the first) - onServerRestartHandler(client, auth) - ) - }) + }, + } as DidChangeWorkspaceFoldersParams) + }), + { dispose: () => clearInterval(refreshInterval) }, + // Set this inside onReady so that it only triggers on subsequent language server starts (not the first) + onServerRestartHandler(client, auth) + ) } /** From 32f37081f57a6d3e9e684cdd7bb9f319b76fe9e6 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Wed, 14 May 2025 14:30:03 -0400 Subject: [PATCH 2/6] refactor: minor, move function to appropriate file Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/chat/activation.ts | 47 +-------------------- packages/amazonq/src/lsp/config.ts | 47 +++++++++++++++++++++ 2 files changed, 48 insertions(+), 46 deletions(-) diff --git a/packages/amazonq/src/lsp/chat/activation.ts b/packages/amazonq/src/lsp/chat/activation.ts index 3a36377b9b5..fcfdb5804c2 100644 --- a/packages/amazonq/src/lsp/chat/activation.ts +++ b/packages/amazonq/src/lsp/chat/activation.ts @@ -12,10 +12,7 @@ import { Commands, getLogger, globals, undefinedIfEmpty } from 'aws-core-vscode/ import { activate as registerLegacyChatListeners } from '../../app/chat/activation' import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq' import { AuthUtil, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' -import { - DidChangeConfigurationNotification, - updateConfigurationRequestType, -} from '@aws/language-server-runtimes/protocol' +import { pushConfigUpdate } from '../config' export async function activate(languageClient: LanguageClient, encryptionKey: Buffer, mynahUIPath: string) { const disposables = globals.context.subscriptions @@ -99,45 +96,3 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu }) ) } - -/** - * Push a config value to the language server, effectively updating it with the - * latest configuration from the client. - * - * The issue is we need to push certain configs to different places, since there are - * different handlers for specific configs. So this determines the correct place to - * push the given config. - */ -async function pushConfigUpdate(client: LanguageClient, config: QConfigs) { - switch (config.type) { - case 'profile': - await client.sendRequest(updateConfigurationRequestType.method, { - section: 'aws.q', - settings: { profileArn: config.profileArn }, - }) - break - case 'customization': - client.sendNotification(DidChangeConfigurationNotification.type.method, { - section: 'aws.q', - settings: { customization: config.customization }, - }) - break - case 'logLevel': - client.sendNotification(DidChangeConfigurationNotification.type.method, { - section: 'aws.logLevel', - }) - break - } -} -type ProfileConfig = { - type: 'profile' - profileArn: string | undefined -} -type CustomizationConfig = { - type: 'customization' - customization: string | undefined -} -type LogLevelConfig = { - type: 'logLevel' -} -type QConfigs = ProfileConfig | CustomizationConfig | LogLevelConfig diff --git a/packages/amazonq/src/lsp/config.ts b/packages/amazonq/src/lsp/config.ts index bb6870cb561..1760fb51401 100644 --- a/packages/amazonq/src/lsp/config.ts +++ b/packages/amazonq/src/lsp/config.ts @@ -5,6 +5,11 @@ import * as vscode from 'vscode' import { DevSettings, getServiceEnvVarConfig } from 'aws-core-vscode/shared' import { LspConfig } from 'aws-core-vscode/amazonq' +import { LanguageClient } from 'vscode-languageclient' +import { + DidChangeConfigurationNotification, + updateConfigurationRequestType, +} from '@aws/language-server-runtimes/protocol' export interface ExtendedAmazonQLSPConfig extends LspConfig { ui?: string @@ -54,3 +59,45 @@ export function getAmazonQLspConfig(): ExtendedAmazonQLSPConfig { export function toAmazonQLSPLogLevel(logLevel: vscode.LogLevel): LspLogLevel { return lspLogLevelMapping.get(logLevel) ?? 'info' } + +/** + * Request/Notify a config value to the language server, effectively updating it with the + * latest configuration from the client. + * + * The issue is we need to push certain configs to different places, since there are + * different handlers for specific configs. So this determines the correct place to + * push the given config. + */ +export async function pushConfigUpdate(client: LanguageClient, config: QConfigs) { + switch (config.type) { + case 'profile': + await client.sendRequest(updateConfigurationRequestType.method, { + section: 'aws.q', + settings: { profileArn: config.profileArn }, + }) + break + case 'customization': + client.sendNotification(DidChangeConfigurationNotification.type.method, { + section: 'aws.q', + settings: { customization: config.customization }, + }) + break + case 'logLevel': + client.sendNotification(DidChangeConfigurationNotification.type.method, { + section: 'aws.logLevel', + }) + break + } +} +type ProfileConfig = { + type: 'profile' + profileArn: string | undefined +} +type CustomizationConfig = { + type: 'customization' + customization: string | undefined +} +type LogLevelConfig = { + type: 'logLevel' +} +type QConfigs = ProfileConfig | CustomizationConfig | LogLevelConfig From 9161c7d066b8a773aae3398f64330affd06df33d Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Wed, 14 May 2025 15:54:40 -0400 Subject: [PATCH 3/6] fix: Call auth setup in the correct order Problem: Profiles and Customizations were pushed to the LSP before Auth was initialized (and its bearer token pushed to the LPS), this caused errors since it does not make sense to push profile/customization if a bearer token does not exist. Solution: Ensure the bearer token is pushed first, then follow up with pushing the profile/customization after. Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/auth.ts | 5 ++-- packages/amazonq/src/lsp/chat/activation.ts | 27 +++++++++++---------- packages/amazonq/src/lsp/client.ts | 22 ++++++++++++----- packages/core/src/auth/index.ts | 1 + 4 files changed, 34 insertions(+), 21 deletions(-) diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index 816c4b09ab0..02888c0b556 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/packages/amazonq/src/lsp/auth.ts @@ -18,6 +18,7 @@ import { AuthUtil } from 'aws-core-vscode/codewhisperer' import { Writable } from 'stream' import { onceChanged } from 'aws-core-vscode/utils' import { getLogger, oneMinute } from 'aws-core-vscode/shared' +import { isSsoConnection } from 'aws-core-vscode/auth' export const encryptionKey = crypto.randomBytes(32) @@ -76,8 +77,8 @@ export class AmazonQLspAuth { * @param force bypass memoization, and forcefully update the bearer token */ async refreshConnection(force: boolean = false) { - const activeConnection = this.authUtil.auth.activeConnection - if (activeConnection?.state === 'valid' && activeConnection?.type === 'sso') { + const activeConnection = this.authUtil.conn + if (this.authUtil.isConnectionValid() && isSsoConnection(activeConnection)) { // send the token to the language server const token = await this.authUtil.getBearerToken() await (force ? this._updateBearerToken(token) : this.updateBearerToken(token)) diff --git a/packages/amazonq/src/lsp/chat/activation.ts b/packages/amazonq/src/lsp/chat/activation.ts index fcfdb5804c2..9ca94008f89 100644 --- a/packages/amazonq/src/lsp/chat/activation.ts +++ b/packages/amazonq/src/lsp/chat/activation.ts @@ -18,15 +18,20 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu const disposables = globals.context.subscriptions // Make sure we've sent an auth profile to the language server before even initializing the UI - await pushConfigUpdate(languageClient, { - type: 'profile', - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }) - // We need to push the cached customization on startup explicitly - await pushConfigUpdate(languageClient, { - type: 'customization', - customization: getSelectedCustomization(), - }) + // + // Ideally the handler for onDidChangeRegionProfile would trigger this, but because the handler is set up too late (due to the current structure), the initial event is missed. + // So we have to explicitly do it here on startup. + if (AuthUtil.instance.isConnectionValid()) { + await pushConfigUpdate(languageClient, { + type: 'profile', + profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + }) + + await pushConfigUpdate(languageClient, { + type: 'customization', + customization: getSelectedCustomization(), + }) + } const provider = new AmazonQChatViewProvider(mynahUIPath) @@ -76,10 +81,6 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu disposables.push( AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { - void pushConfigUpdate(languageClient, { - type: 'profile', - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }) await provider.refreshWebview() }), Commands.register('aws.amazonq.updateCustomizations', () => { diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 089f8c2e641..7d356d514e2 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -38,7 +38,7 @@ import { import { processUtils } from 'aws-core-vscode/shared' import { activate } from './chat/activation' import { AmazonQResourcePaths } from './lspInstaller' -import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config' +import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config' import { telemetry } from 'aws-core-vscode/telemetry' const localize = nls.loadMessageBundle() @@ -160,24 +160,34 @@ export async function startLanguageServer( const disposable = client.start() toDispose.push(disposable) - - const auth = new AmazonQLspAuth(client) - await client.onReady() + const auth = await initializeAuth(client) + await postStartLanguageServer(auth, client, resourcePaths, toDispose) return client } +async function initializeAuth(client: LanguageClient): Promise { + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { + void pushConfigUpdate(client, { + type: 'profile', + profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + }) + }) + + const auth = new AmazonQLspAuth(client) + await auth.refreshConnection(true) + return auth +} + async function postStartLanguageServer( auth: AmazonQLspAuth, client: LanguageClient, resourcePaths: AmazonQResourcePaths, toDispose: vscode.Disposable[] ) { - await auth.refreshConnection() - if (Experiments.instance.get('amazonqLSPInline', false)) { const inlineManager = new InlineCompletionManager(client) inlineManager.registerInlineCompletion() diff --git a/packages/core/src/auth/index.ts b/packages/core/src/auth/index.ts index 54dd17d702b..02a0067be45 100644 --- a/packages/core/src/auth/index.ts +++ b/packages/core/src/auth/index.ts @@ -18,6 +18,7 @@ export { isBuilderIdConnection, getTelemetryMetadataForConn, isIamConnection, + isSsoConnection, } from './connection' export { Auth } from './auth' export { CredentialsStore } from './credentials/store' From 04c91d9ed7c15b22bba2e40732e53ad4caac6bc2 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Wed, 14 May 2025 21:42:02 -0400 Subject: [PATCH 4/6] fix(lsp): startUrl was undefined when sent to Q LSP Problem: THe old way we got the startUrl did not return a value since the underlying activeConnection in the Auth class was empty. The activeConnection we actually want was in AuthUtil. Solution: Use the connection we get from AuthUtil as that is the most accurate one. Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/auth.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index 02888c0b556..b285102b513 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/packages/amazonq/src/lsp/auth.ts @@ -115,11 +115,14 @@ export class AmazonQLspAuth { .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) .encrypt(encryptionKey) + const conn = AuthUtil.instance.conn + const startUrl = conn?.type === 'sso' ? conn.startUrl : undefined + return { data: jwt, metadata: { sso: { - startUrl: AuthUtil.instance.auth.startUrl, + startUrl, }, }, encrypted: true, From ebceefffed158d3e86dbdeb43e605da759f11cf8 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Wed, 14 May 2025 22:20:42 -0400 Subject: [PATCH 5/6] changelog + comment fix Signed-off-by: nkomonen-amazon --- .../Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json | 4 ++++ .../Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json | 4 ++++ packages/amazonq/src/lsp/auth.ts | 5 +---- 3 files changed, 9 insertions(+), 4 deletions(-) create mode 100644 packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json create mode 100644 packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json diff --git a/packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json b/packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json new file mode 100644 index 00000000000..f7af0fbb1a4 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-2471c584-3904-4d90-bb7c-61efff219e43.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Fix Error: 'Amazon Q service is not signed in'" +} diff --git a/packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json b/packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json new file mode 100644 index 00000000000..e3a608296a0 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-91d391d4-3777-4053-9e71-15b36dfa1f67.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Fix Error: 'Amazon Q Profile is not selected for IDC connection type'" +} diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index b285102b513..d81f464d6a3 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/packages/amazonq/src/lsp/auth.ts @@ -115,14 +115,11 @@ export class AmazonQLspAuth { .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' }) .encrypt(encryptionKey) - const conn = AuthUtil.instance.conn - const startUrl = conn?.type === 'sso' ? conn.startUrl : undefined - return { data: jwt, metadata: { sso: { - startUrl, + startUrl: AuthUtil.instance.startUrl, }, }, encrypted: true, From 7a525ae25c13b64fa8d33ea69f1a7a820b37ac6b Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Thu, 15 May 2025 10:19:07 -0400 Subject: [PATCH 6/6] comment: Move Profile push up a level - Moves the profile/customization code up a level out of the LSP chat activation, since it doesn't make sense there - Remove the previous onDidChangeProfile handler from the auth setup since it is redundant due to us having it already in onLanguageServerReady() - Renamed function to onLanguageServerReady() to be more aligned with flare naming Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/chat/activation.ts | 16 ------- packages/amazonq/src/lsp/client.ts | 49 +++++++++------------ 2 files changed, 20 insertions(+), 45 deletions(-) diff --git a/packages/amazonq/src/lsp/chat/activation.ts b/packages/amazonq/src/lsp/chat/activation.ts index 9ca94008f89..f8e3ee16251 100644 --- a/packages/amazonq/src/lsp/chat/activation.ts +++ b/packages/amazonq/src/lsp/chat/activation.ts @@ -17,22 +17,6 @@ import { pushConfigUpdate } from '../config' export async function activate(languageClient: LanguageClient, encryptionKey: Buffer, mynahUIPath: string) { const disposables = globals.context.subscriptions - // Make sure we've sent an auth profile to the language server before even initializing the UI - // - // Ideally the handler for onDidChangeRegionProfile would trigger this, but because the handler is set up too late (due to the current structure), the initial event is missed. - // So we have to explicitly do it here on startup. - if (AuthUtil.instance.isConnectionValid()) { - await pushConfigUpdate(languageClient, { - type: 'profile', - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }) - - await pushConfigUpdate(languageClient, { - type: 'customization', - customization: getSelectedCustomization(), - }) - } - const provider = new AmazonQChatViewProvider(mynahUIPath) disposables.push( diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 7d356d514e2..549b0ac7dad 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -16,7 +16,6 @@ import { GetConfigurationFromServerParams, RenameFilesParams, ResponseMessage, - updateConfigurationRequestType, WorkspaceFolder, } from '@aws/language-server-runtimes/protocol' import { AuthUtil, CodeWhispererSettings, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' @@ -164,25 +163,18 @@ export async function startLanguageServer( const auth = await initializeAuth(client) - await postStartLanguageServer(auth, client, resourcePaths, toDispose) + await onLanguageServerReady(auth, client, resourcePaths, toDispose) return client } async function initializeAuth(client: LanguageClient): Promise { - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { - void pushConfigUpdate(client, { - type: 'profile', - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }) - }) - const auth = new AmazonQLspAuth(client) await auth.refreshConnection(true) return auth } -async function postStartLanguageServer( +async function onLanguageServerReady( auth: AmazonQLspAuth, client: LanguageClient, resourcePaths: AmazonQResourcePaths, @@ -208,25 +200,17 @@ async function postStartLanguageServer( const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond) - const sendProfileToLsp = async () => { - try { - const result = await client.sendRequest(updateConfigurationRequestType.method, { - section: 'aws.q', - settings: { - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, - }, - }) - client.info( - `Client: Updated Amazon Q Profile ${AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn} to Amazon Q LSP`, - result - ) - } catch (err) { - client.error('Error when setting Q Developer Profile to Amazon Q LSP', err) - } - } + // We manually push the cached values the first time since event handlers, which should push, may not have been setup yet. + // Execution order is weird and should be fixed in the flare implementation. + // TODO: Revisit if we need this if we setup the event handlers properly + if (AuthUtil.instance.isConnectionValid()) { + await sendProfileToLsp(client) - // send profile to lsp once. - void sendProfileToLsp() + await pushConfigUpdate(client, { + type: 'customization', + customization: getSelectedCustomization(), + }) + } toDispose.push( AuthUtil.instance.auth.onDidChangeActiveConnection(async () => { @@ -235,7 +219,7 @@ async function postStartLanguageServer( AuthUtil.instance.auth.onDidDeleteConnection(async () => { client.sendNotification(notificationTypes.deleteBearerToken.method) }), - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp), + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => sendProfileToLsp(client)), vscode.commands.registerCommand('aws.amazonq.getWorkspaceId', async () => { const requestType = new RequestType( 'aws/getConfigurationFromServer' @@ -295,6 +279,13 @@ async function postStartLanguageServer( // Set this inside onReady so that it only triggers on subsequent language server starts (not the first) onServerRestartHandler(client, auth) ) + + async function sendProfileToLsp(client: LanguageClient) { + await pushConfigUpdate(client, { + type: 'profile', + profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + }) + } } /**