diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 9c9fc8a888..ffecab0b6f 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -30,9 +30,13 @@ import { ShowDocumentResult, ResponseError, 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, @@ -173,22 +177,22 @@ export async function startLanguageServer( const disposable = client.start() toDispose.push(disposable) - await client.onReady() 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) + await initializeAuth(client, toDispose) - await postStartLanguageServer(client, resourcePaths, toDispose) + await onLanguageServerReady(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 @@ -200,20 +204,25 @@ 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 */ - 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) }) - }) + ) - // 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() } } -async function postStartLanguageServer( +async function onLanguageServerReady( client: LanguageClient, resourcePaths: AmazonQResourcePaths, toDispose: vscode.Disposable[] @@ -260,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) => { @@ -297,7 +289,7 @@ async function postStartLanguageServer( } // send profile to lsp once. - void sendProfileToLsp() + void sendDeveloperProfileToLsp(client) void vscode.window.withProgress( { @@ -329,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' @@ -390,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. diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts index 952edb4ac4..44efc46212 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