refactor(amazonq): Minor Auth refactors to make code more readable

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<void> | 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<GetConfigurationFromServerParams, ResponseMessage, Error>(
                 '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.

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<void> {
             [...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