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)