From a7abe4c6d2abb6acd7a661349baa9c0a3454e02d Mon Sep 17 00:00:00 2001 From: Laxman Reddy Aileni Date: Thu, 10 Sep 2026 20:33:06 +0000 Subject: [PATCH] fix: coalesce auth-triggered GetProfile checks in ProfileStatusMonitor Every auth success event ran an independent MCP configuration check, and each check issued its own GetProfile request. Bursts of token refreshes or repeated profile configuration updates therefore turned into bursts of GetProfile calls, made worse during backend errors because each check also retried. Share a single in-flight check between concurrent callers, and skip auth-triggered checks for the same profile within a one-minute cooldown. Explicit initial and periodic checks are unchanged, and a profile change still triggers an immediate check. --- .../tools/mcp/profileStatusMonitor.test.ts | 85 +++++++++++++++++++ .../tools/mcp/profileStatusMonitor.ts | 65 +++++++++++++- 2 files changed, 148 insertions(+), 2 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts index 4d61a78bc9..44d97ec4c1 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.test.ts @@ -321,6 +321,91 @@ describe('ProfileStatusMonitor', () => { }) }) + describe('GetProfile call coalescing', () => { + let mockServiceManager: any + let getProfileStub: sinon.SinonStub + + const profileResponse = { + profile: { + optInFeatures: { + mcpConfiguration: { + toggle: 'ON', + }, + }, + }, + } + + beforeEach(() => { + getProfileStub = sinon.stub().resolves(profileResponse) + mockServiceManager = { + getActiveProfileArn: sinon.stub().returns('arn:aws:iam::123456789012:profile/test'), + getCodewhispererService: sinon.stub().returns({ getProfile: getProfileStub }), + getConnectionType: sinon.stub().returns('builderId'), + } + + sinon + .stub(AmazonQTokenServiceManagerModule.AmazonQTokenServiceManager, 'getInstance') + .returns(mockServiceManager as any) + }) + + it('should share a single GetProfile call between concurrent checks', async () => { + const first = (profileStatusMonitor as any).isMcpEnabled() + const second = (profileStatusMonitor as any).isMcpEnabled() + + const results = await Promise.all([first, second]) + + expect(results).to.deep.equal([true, true]) + expect(getProfileStub.callCount).to.equal(1) + }) + + it('should allow a new check once the previous one has completed', async () => { + await (profileStatusMonitor as any).isMcpEnabled() + await (profileStatusMonitor as any).isMcpEnabled() + + expect(getProfileStub.callCount).to.equal(2) + }) + + it('should not repeat GetProfile for the same profile within the auth event cooldown', async () => { + await (profileStatusMonitor as any).onAuthSuccess() + await (profileStatusMonitor as any).onAuthSuccess() + await (profileStatusMonitor as any).onAuthSuccess() + + expect(getProfileStub.callCount).to.equal(1) + + clock.tick(ProfileStatusMonitor.AUTH_EVENT_MIN_INTERVAL_MS) + await (profileStatusMonitor as any).onAuthSuccess() + + expect(getProfileStub.callCount).to.equal(2) + }) + + it('should check immediately when the active profile changes', async () => { + await (profileStatusMonitor as any).onAuthSuccess() + expect(getProfileStub.callCount).to.equal(1) + + mockServiceManager.getActiveProfileArn.returns('arn:aws:iam::123456789012:profile/other') + await (profileStatusMonitor as any).onAuthSuccess() + + expect(getProfileStub.callCount).to.equal(2) + }) + + it('should apply the cooldown even when the check fails', async () => { + const serverError = Object.assign(new Error('Internal error'), { statusCode: 500 }) + getProfileStub.rejects(serverError) + + const firstAttempt = (profileStatusMonitor as any).onAuthSuccess() + // retryWithBackoff waits between attempts; advance the fake clock so it can finish + await clock.tickAsync(5000) + await firstAttempt + const callsAfterFirstEvent = getProfileStub.callCount + expect(callsAfterFirstEvent).to.be.greaterThan(0) + + await (profileStatusMonitor as any).onAuthSuccess() + + expect(getProfileStub.callCount).to.equal(callsAfterFirstEvent) + expect(mockLogging.debug.calledWith(sinon.match('checked recently'))).to.be.true + }) + }) + describe('isEnterpriseUser', () => { let mockServiceManager: any diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts index e69829b8e9..2227218bb0 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/profileStatusMonitor.ts @@ -20,6 +20,14 @@ export const AUTH_SUCCESS_EVENT = 'authSuccess' export class ProfileStatusMonitor { private intervalId?: NodeJS.Timeout private readonly CHECK_INTERVAL = 24 * 60 * 60 * 1000 // 24 hours + /** + * Minimum time between auth-triggered profile checks for the same profile. + * Auth success events can arrive in rapid succession (token refresh, repeated + * configuration updates); without this guard each one issued a GetProfile call. + */ + static readonly AUTH_EVENT_MIN_INTERVAL_MS = 60 * 1000 + private inFlightCheck?: Promise + private lastAuthEventCheck?: { profileArn: string; timestamp: number } private codeWhispererClient?: CodeWhispererServiceToken private static lastMcpState: boolean = true private static readonly MCP_CACHE_DIR = path.join(os.homedir(), '.aws', 'amazonq', 'mcpAdmin') @@ -40,10 +48,48 @@ export class ProfileStatusMonitor { // Listen for auth success events ProfileStatusMonitor.eventEmitter.on(AUTH_SUCCESS_EVENT, () => { - void this.isMcpEnabled() + void this.onAuthSuccess() }) } + /** + * Handles an auth success event. Skips the profile check when the same profile + * was already checked within AUTH_EVENT_MIN_INTERVAL_MS, so bursts of auth or + * configuration updates do not turn into bursts of GetProfile calls. + */ + private async onAuthSuccess(): Promise { + const profileArn = this.tryGetActiveProfileArn() + const now = Date.now() + + if ( + profileArn && + this.lastAuthEventCheck?.profileArn === profileArn && + now - this.lastAuthEventCheck.timestamp < ProfileStatusMonitor.AUTH_EVENT_MIN_INTERVAL_MS + ) { + this.logging.debug('Skipping MCP configuration check: profile was checked recently') + return + } + + if (profileArn) { + this.lastAuthEventCheck = { profileArn, timestamp: now } + } + + try { + await this.isMcpEnabled() + } catch { + // Already logged by isMcpEnabled; nothing else to do for an event-triggered check. + } + } + + private tryGetActiveProfileArn(): string | undefined { + try { + return this.getProfileArn(AmazonQTokenServiceManager.getInstance()) + } catch (error) { + this.logging.debug(`Service manager not available for profile check: ${error}`) + return undefined + } + } + async checkInitialState(): Promise { try { const isMcpEnabled = await this.isMcpEnabled() @@ -79,7 +125,22 @@ export class ProfileStatusMonitor { } } - private async isMcpEnabled(isPeriodicCheck: boolean = false): Promise { + /** + * Returns the in-flight check if one is running so concurrent callers share a + * single GetProfile request instead of each issuing their own. + */ + private isMcpEnabled(isPeriodicCheck: boolean = false): Promise { + if (this.inFlightCheck) { + return this.inFlightCheck + } + + this.inFlightCheck = this.checkMcpEnabled(isPeriodicCheck).finally(() => { + this.inFlightCheck = undefined + }) + return this.inFlightCheck + } + + private async checkMcpEnabled(isPeriodicCheck: boolean = false): Promise { try { const serviceManager = AmazonQTokenServiceManager.getInstance() const profileArn = this.getProfileArn(serviceManager)