feat(amazonq): initialize Q services when credentials arrive

server/aws-lsp-codewhisperer/package.json

@@ -38,7 +38,7 @@
         "@aws-sdk/util-arn-parser": "^3.723.0",
         "@aws-sdk/util-retry": "^3.374.0",
         "@aws/chat-client-ui-types": "0.1.71",
-        "@aws/language-server-runtimes": "^0.3.19",
+        "@aws/language-server-runtimes": "^0.3.21",
         "@aws/lsp-core": "^0.0.21",
         "@modelcontextprotocol/sdk": "^1.23.0",
         "@mozilla/readability": "^0.6.0",

server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts

@@ -116,15 +116,14 @@ export const AmazonQServiceServerFactory =
             amazonQServiceManager.handleOnCredentialsDeleted(type)
         })
 
-        // Add credentials update handler to track when credentials are updated
-        if ('onCredentialsUpdated' in credentialsProvider) {
-            ;(credentialsProvider as any).onCredentialsUpdated((type: CredentialsType) => {
-                log(`Received onCredentialsUpdated notification for type: ${type}`)
-                if ('handleOnCredentialsUpdated' in amazonQServiceManager) {
-                    ;(amazonQServiceManager as any).handleOnCredentialsUpdated(type)
-                }
-            })
-        }
+        // Optional chaining rather than a capability check: the event was added to the runtime after
+        // this server, so a build running on an older runtime simply never receives it. Before that
+        // runtime change this block was an `in` check plus an `any` cast that could never fire, which
+        // is why service construction waited for whichever consumer asked first.
+        credentialsProvider.onCredentialsUpdated?.((type: CredentialsType) => {
+            log(`Received onCredentialsUpdated notification for type: ${type}`)
+            amazonQServiceManager.handleOnCredentialsUpdated(type)
+        })
 
         logging.log('Amazon Q Service server has been initialised')
         return () => {}

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts

@@ -2,6 +2,7 @@ import * as assert from 'assert'
 import sinon, { StubbedInstance, stubInterface } from 'ts-sinon'
 import { AmazonQTokenServiceManager } from './AmazonQTokenServiceManager'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
+import { CredentialsType } from '@aws/language-server-runtimes/server-interface'
 import { CodeWhispererServiceToken, GenerateSuggestionsRequest } from '../codeWhispererService'
 import {
     AmazonQServiceInitializationError,
@@ -299,6 +300,74 @@ describe('AmazonQTokenServiceManager', () => {
         })
     })
 
+    describe('Eager initialization on credentials update', () => {
+        const testRegion = 'some-region'
+        const testEndpoint = 'http://some-endpoint-in-some-region'
+
+        beforeEach(() => {
+            // Each test sets up its own manager: the IdC case needs developer profile support enabled,
+            // and the manager is a singleton that refuses a second initialization.
+            AWS_Q_ENDPOINTS.set(testRegion, testEndpoint)
+        })
+
+        it('initializes on bearer credentials without waiting for a consumer', () => {
+            setupServiceManager()
+            setCredentials('builderId')
+            assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
+
+            // No getCodewhispererService() call here: the point is that nothing has asked yet.
+            amazonQTokenServiceManager.handleOnCredentialsUpdated('bearer' as CredentialsType)
+
+            assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED')
+        })
+
+        it('is idempotent, so token refreshes do not rebuild live services', () => {
+            setupServiceManager()
+            setCredentials('builderId')
+            amazonQTokenServiceManager.handleOnCredentialsUpdated('bearer' as CredentialsType)
+            const firstService = amazonQTokenServiceManager.getCodewhispererService()
+
+            // Credentials update fires on every token refresh. Rebuilding here would drop in-flight
+            // requests and replace the service the observers are attached to.
+            amazonQTokenServiceManager.handleOnCredentialsUpdated('bearer' as CredentialsType)
+            amazonQTokenServiceManager.handleOnCredentialsUpdated('bearer' as CredentialsType)
+
+            assert.strictEqual(amazonQTokenServiceManager.getCodewhispererService(), firstService)
+            assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED')
+        })
+
+        it('ignores iam credentials', () => {
+            setupServiceManager()
+            setCredentials('builderId')
+
+            amazonQTokenServiceManager.handleOnCredentialsUpdated('iam' as CredentialsType)
+
+            assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
+        })
+
+        it('does not build services for IdC before a profile is chosen', () => {
+            // The whole point of bringing initialization forward is to make a request sooner. For IdC
+            // with developer profiles that would be a request with no profile, so this must stop at
+            // PENDING_Q_PROFILE and create nothing.
+            setupServiceManager(true)
+            setCredentials('identityCenter')
+
+            amazonQTokenServiceManager.handleOnCredentialsUpdated('bearer' as CredentialsType)
+
+            assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE')
+            assert.throws(() => amazonQTokenServiceManager.getCodewhispererService())
+        })
+
+        it('does not throw when there are no credentials yet', () => {
+            setupServiceManager()
+            // Fires inside the client's credentials request; throwing would make the client believe the
+            // credentials never landed.
+            assert.doesNotThrow(() =>
+                amazonQTokenServiceManager.handleOnCredentialsUpdated('bearer' as CredentialsType)
+            )
+        })
+    })
+
     describe('IdentityCenter support', () => {
         describe('Developer Profiles Support is disabled', () => {
             it('should be INITIALIZED with IdentityCenter Connection', async () => {

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts

@@ -164,12 +164,41 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
         ProfileStatusMonitor.resetMcpState()
     }
 
-    public handleOnCredentialsUpdated(type: CredentialsType): void {
+    /**
+     * Constructs the Q services as soon as bearer credentials arrive, rather than leaving it to
+     * whichever consumer asks first.
+     *
+     * That consumer is in practice the chat webview finishing boot, which measured ~0.66s after the
+     * credentials landed -- time in which nothing had contacted the service yet, so a rejected
+     * identity could not be discovered and the user sat in front of a chat view that was never going
+     * to work.
+     *
+     * Deliberately reuses handleSsoConnectionChange rather than constructing anything here, so this
+     * changes only *when* initialization happens and not *how*:
+     * - It is idempotent. createCodewhispererServiceInstances records the connection type, so the
+     *   consumer's later call short-circuits on "Connection type did not change" instead of resetting
+     *   live services. Token refreshes, which fire this event repeatedly, are no-ops for the same
+     *   reason.
+     * - IdC with developer profile support still stops at PENDING_Q_PROFILE and creates nothing, so no
+     *   request is made before a profile is chosen.
+     *
+     * Failures are logged rather than thrown: this runs inside the client's credentials request, and
+     * the lazy path still runs for whoever asks first, so a failure here costs the head start and
+     * nothing else.
+     */
+    public override handleOnCredentialsUpdated(type: CredentialsType): void {
         this.logging.log(`Received credentials update event for type: ${type}`)
 
-        if (type === ('bearer' as CredentialsType)) {
-            // Check Q credentials
-            const qCreds = this.features.credentialsProvider.getCredentials('bearer' as CredentialsType)
+        if (type !== ('bearer' as CredentialsType)) {
+            return
+        }
+
+        try {
+            this.handleSsoConnectionChange()
+        } catch (e) {
+            this.logging.log(
+                `Could not initialize Q services on credentials update, deferring to first use: ${(e as Error)?.message}`
+            )
         }
     }
 

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts

@@ -123,6 +123,15 @@ export abstract class BaseAmazonQServiceManager<
     }
 
     abstract handleOnCredentialsDeleted(type: CredentialsType): void
+
+    /**
+     * Called when credentials of the given type have been stored.
+     *
+     * Default is a no-op so a manager only implements it if it has work to bring forward. Overridden
+     * by the token manager to construct services as soon as bearer credentials arrive instead of
+     * waiting for the first consumer.
+     */
+    handleOnCredentialsUpdated(_type: CredentialsType): void {}
     abstract handleOnUpdateConfiguration(params: UpdateConfigurationParams, token: CancellationToken): Promise<void>
 
     public async handleDidChangeConfiguration(): Promise<void> {