diff --git a/package-lock.json b/package-lock.json index c5e6016108..3f193aa49e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -3755,12 +3755,12 @@ } }, "node_modules/@aws/language-server-runtimes": { - "version": "0.3.19", - "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes/-/language-server-runtimes-0.3.19.tgz", - "integrity": "sha512-Cc5LnqfNRA6EAXL9DMwa/tgCyN5knlrMuYfo1JbVg3YPV8SoVGIejqz830A77fjFFkb4aPizp8wgiyB4VIfXiQ==", + "version": "0.3.21", + "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes/-/language-server-runtimes-0.3.21.tgz", + "integrity": "sha512-2Fyqi3/lgWS7UHO4qe2WUTWGuPK6Kd+9uReLhzibnWC4oAbenchtbJre0e0HCvSQ5obD9dDg+EryfaF+EWyxjA==", "license": "Apache-2.0", "dependencies": { - "@aws/language-server-runtimes-types": "^0.1.65", + "@aws/language-server-runtimes-types": "^0.1.66", "@opentelemetry/api": "^1.9.0", "@opentelemetry/api-logs": "^0.200.0", "@opentelemetry/core": "^2.0.0", @@ -3784,9 +3784,9 @@ } }, "node_modules/@aws/language-server-runtimes-types": { - "version": "0.1.65", - "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes-types/-/language-server-runtimes-types-0.1.65.tgz", - "integrity": "sha512-zVaeJ+IjWiovCDZuFOfOKKRFJk7jISp+ohyQRHgFT4Ek53iBnjEMJrpP+z2+9QgfF3p+k26ZsAZEfylcQJFo+Q==", + "version": "0.1.66", + "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes-types/-/language-server-runtimes-types-0.1.66.tgz", + "integrity": "sha512-+PlvurtSQQ4X27woztQcxBRKWY9XVrk3cLVoRNct1pX2jHSe6TUAXLZtOCJAm7R8VtlzNPiGHMXxBcCQAav0cQ==", "license": "Apache-2.0", "dependencies": { "vscode-languageserver-textdocument": "^1.0.12", @@ -30371,7 +30371,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", diff --git a/server/aws-lsp-codewhisperer/package.json b/server/aws-lsp-codewhisperer/package.json index 7f83033399..120744dc41 100644 --- a/server/aws-lsp-codewhisperer/package.json +++ b/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", diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts index 8ec4e19372..f4ce1ef494 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts +++ b/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 () => {} diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts index c93552a7d7..948ff2613f 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts +++ b/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 () => { diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts index 608d4ae2ed..c6d734d3a6 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts +++ b/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}` + ) } } diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts index 763a75dfdd..6764eaa752 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts +++ b/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 public async handleDidChangeConfiguration(): Promise {