From 34050bdefb55fb20eb538b12a7077100c26911f4 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Thu, 15 May 2025 16:23:48 -0400 Subject: [PATCH 1/4] minor: remove redundant onReady call Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/client.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 9c9fc8a8883..bd2e022f6ef 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -173,7 +173,6 @@ export async function startLanguageServer( const disposable = client.start() toDispose.push(disposable) - await client.onReady() await client.onReady() From bf63299183fea35eaac631337c84ecc45166dab1 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Thu, 15 May 2025 18:06:07 -0400 Subject: [PATCH 2/4] fix:minor: move region profile handler to better spot We ultimately want to bubble up all Q activation code in to packages/amazonq instead of packages/core. As core existed before packages/amazonq did, and no one has ported all amazonq code to packages/amazonq. This does that. Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/client.ts | 17 +++++++++++++---- packages/core/src/codewhisperer/activation.ts | 4 +--- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index bd2e022f6ef..903182a9c2b 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -32,7 +32,12 @@ import { LSPErrorCodes, updateConfigurationRequestType, } from '@aws/language-server-runtimes/protocol' -import { AuthUtil, CodeWhispererSettings, getSelectedCustomization } from 'aws-core-vscode/codewhisperer' +import { + AuthUtil, + CodeWhispererSettings, + getSelectedCustomization, + onProfileChangedListener, +} from 'aws-core-vscode/codewhisperer' import { Settings, createServerOptions, @@ -175,11 +180,12 @@ export async function startLanguageServer( toDispose.push(disposable) await client.onReady() - /** * We use the Flare Auth language server, and our Auth client depends on it. * Because of this we initialize our Auth client **immediately** after the language server is ready. * Doing this removes the chance of something else attempting to use the Auth client before it is ready. + * + * All other LSP initialization steps should happen after this. */ await initializeAuth(client) @@ -199,15 +205,18 @@ export async function startLanguageServer( getLogger().error(`Error while migration SSO connection to Amazon Q LSP: ${e}`) } - /** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */ + // Push region profile to the Q Language Server whenever it changes AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { void pushConfigUpdate(client, { type: 'profile', profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, }) }) + // Handle for Customization when the Developer Profile changes + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener) - // Try and restore a cached connection if exists + // THIS SHOULD BE LAST!!! + // This will result start the process of initializing the cached token (if it exists) await AuthUtil.instance.restore() } } diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts index 952edb4ac42..44efc462125 100644 --- a/packages/core/src/codewhisperer/activation.ts +++ b/packages/core/src/codewhisperer/activation.ts @@ -72,7 +72,6 @@ import { AuthUtil } from './util/authUtil' import { ImportAdderProvider } from './service/importAdderProvider' import { TelemetryHelper } from './util/telemetryHelper' import { openUrl } from '../shared/utilities/vsCodeUtils' -import { onProfileChangedListener } from './util/customizationUtil' import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands' import { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider' import { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider' @@ -334,8 +333,7 @@ export async function activate(context: ExtContext): Promise { [...CodeWhispererConstants.securityScanLanguageIds], SecurityIssueCodeActionProvider.instance ), - vscode.commands.registerCommand('aws.amazonq.openEditorAtRange', openEditorAtRange), - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener) + vscode.commands.registerCommand('aws.amazonq.openEditorAtRange', openEditorAtRange) ) // run the auth startup code with context for telemetry From 7d1afb1d4f1c9a46b32e6db4593446bb7daf6974 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Thu, 15 May 2025 19:01:03 -0400 Subject: [PATCH 3/4] refactor:minor: Move Developer Profile handler code to a single spot Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/client.ts | 57 ++++++++++++++---------------- 1 file changed, 27 insertions(+), 30 deletions(-) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 903182a9c2b..b7c57e4c168 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -30,7 +30,6 @@ import { ShowDocumentResult, ResponseError, LSPErrorCodes, - updateConfigurationRequestType, } from '@aws/language-server-runtimes/protocol' import { AuthUtil, @@ -187,13 +186,13 @@ export async function startLanguageServer( * * All other LSP initialization steps should happen after this. */ - await initializeAuth(client) + await initializeAuth(client, toDispose) await postStartLanguageServer(client, resourcePaths, toDispose) return client - async function initializeAuth(client: LanguageClient) { + async function initializeAuth(client: LanguageClient, toDispose: vscode.Disposable[]) { AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey)) // Migrate SSO connections from old Auth to the LSP identity server @@ -205,15 +204,17 @@ export async function startLanguageServer( getLogger().error(`Error while migration SSO connection to Amazon Q LSP: ${e}`) } - // Push region profile to the Q Language Server whenever it changes - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { - void pushConfigUpdate(client, { - type: 'profile', - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + // We set these handlers before the auth restore below since it may trigger these + toDispose.push( + // Push region profile to the Q Language Server whenever it changes + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { + await sendDeveloperProfileToLsp(client) + }), + // Handle for Customization when the Developer Profile changes + AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile((e) => { + onProfileChangedListener(e) }) - }) - // Handle for Customization when the Developer Profile changes - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener) + ) // THIS SHOULD BE LAST!!! // This will result start the process of initializing the cached token (if it exists) @@ -268,23 +269,6 @@ async function postStartLanguageServer( } ) - 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) - } - } - let promise: Promise | undefined let resolver: () => void = () => {} client.onProgress(GetSsoTokenProgressType, GetSsoTokenProgressToken, async (partialResult: GetSsoTokenProgress) => { @@ -305,7 +289,7 @@ async function postStartLanguageServer( } // send profile to lsp once. - void sendProfileToLsp() + void sendDeveloperProfileToLsp(client) void vscode.window.withProgress( { @@ -337,7 +321,6 @@ async function postStartLanguageServer( } toDispose.push( - AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp), vscode.commands.registerCommand('aws.amazonq.getWorkspaceId', async () => { const requestType = new RequestType( 'aws/getConfigurationFromServer' @@ -398,6 +381,20 @@ async function postStartLanguageServer( ) } +async function sendDeveloperProfileToLsp(client: LanguageClient) { + const profileArn = AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn + + try { + await pushConfigUpdate(client, { + type: 'profile', + profileArn, + }) + client.info(`DeveloperProfile: Successfully pushed Developer Profile, ${profileArn}, to Amazon Q LSP`) + } catch (err) { + client.error(`DeveloperProfile: Failed to push Developer Profile, ${profileArn}, to Amazon Q LSP`, err) + } +} + /** * When the server restarts (likely due to a crash, then the LanguageClient automatically starts it again) * we need to run some server intialization again. From c8e7926366f6d8ec0be22b3afdf7b0aaf0c28da8 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Thu, 15 May 2025 19:13:46 -0400 Subject: [PATCH 4/4] minor: rename function to be more appropriate Signed-off-by: nkomonen-amazon --- packages/amazonq/src/lsp/client.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index b7c57e4c168..ffecab0b6fb 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -188,7 +188,7 @@ export async function startLanguageServer( */ await initializeAuth(client, toDispose) - await postStartLanguageServer(client, resourcePaths, toDispose) + await onLanguageServerReady(client, resourcePaths, toDispose) return client @@ -222,7 +222,7 @@ export async function startLanguageServer( } } -async function postStartLanguageServer( +async function onLanguageServerReady( client: LanguageClient, resourcePaths: AmazonQResourcePaths, toDispose: vscode.Disposable[]