diff --git a/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts b/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts index 57d4bcd91dad41..995ec1fa9a6d91 100644 --- a/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts +++ b/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts @@ -7,7 +7,7 @@ import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.j import { safeIntl } from '../../../../base/common/date.js'; import { localize } from '../../../../nls.js'; import { IContextKeyService } from '../../../../platform/contextkey/common/contextkey.js'; -import { IStorageService, StorageScope } from '../../../../platform/storage/common/storage.js'; +import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js'; import { IWorkbenchContribution } from '../../../common/contributions.js'; import { ChatEntitlement, IChatEntitlementService, IQuotaSnapshot, IRateLimitSnapshot } from '../../../services/chat/common/chatEntitlementService.js'; import { isSelectedModelCopilot, SELECTED_MODEL_STORAGE_KEY_PREFIX } from '../common/chatSelectedModel.js'; @@ -17,6 +17,14 @@ import { ChatInputNotificationSeverity, IChatInputNotification, IChatInputNotifi const QUOTA_NOTIFICATION_ID = 'copilot.quotaStatus'; const THRESHOLDS = [50, 75, 90, 95]; +/** + * Persisted flag remembering that the user dismissed the quota-exceeded + * notification. Kept until quota recovers (credit becomes available again) so + * the banner does not re-appear on every window reload while quota is still + * exhausted. + */ +const QUOTA_EXHAUSTED_DISMISSED_STORAGE_KEY = 'chat.quotaNotification.exhaustedDismissed'; + /** * Core-side workbench contribution that shows chat input notifications for * quota exhaustion and quota-approaching thresholds. @@ -69,6 +77,15 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo } })); + // Remember when the user dismisses the quota-exceeded notification so it + // does not re-appear on the next window reload while quota is still + // exhausted. The flag is cleared from `_update` once quota recovers. + this._register(this._chatInputNotificationService.onDidDismiss(id => { + if (id === QUOTA_NOTIFICATION_ID && this._showingExhausted) { + this._setExhaustedDismissed(); + } + })); + // Check initial state in case quota is already exhausted at startup this._update(); } @@ -101,6 +118,16 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo const entitlement = this._chatEntitlementService.entitlement; const isCopilot = this._isCopilotModelSelected(); + // Once quota recovers (credit is positively available again) drop any + // persisted dismissal so the quota-exceeded notification can show the next + // time quota runs out. Done before the Copilot/BYOK gate so a recovery is + // always observed, even while a BYOK model is selected. Guarded on a + // present snapshot so the transient "no quota data yet" state at + // startup/reload does not wipe the flag. + if (this._isQuotaKnownAvailable()) { + this._clearExhaustedDismissed(); + } + // Defer new notifications when a BYOK model is selected or the model // selection hasn't loaded yet — quota only applies to Copilot models. // Already-shown notifications stay visible. @@ -115,7 +142,9 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo // authoritative signal that the org has exceeded its budget, regardless of // overages or remaining quota. if (this._isManagedPlan(entitlement) && this._isManagedPlanBlocked()) { - this._showManagedPlanBlockedNotification(); + if (!this._isExhaustedDismissed()) { + this._showManagedPlanBlockedNotification(); + } return; } @@ -126,14 +155,16 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo const wasAdditionalUsageEnabled = this._prevAdditionalUsageEnabled; this._prevAdditionalUsageEnabled = additionalUsageEnabled; - if (additionalUsageEnabled) { - // Show overage notification on a live transition to 100%, - // or when overages are enabled while already at 100%. - if (this._prevQuotaPercentUsed !== undefined || wasAdditionalUsageEnabled === false) { - this._showOverageActivationNotification(); + if (!this._isExhaustedDismissed()) { + if (additionalUsageEnabled) { + // Show overage notification on a live transition to 100%, + // or when overages are enabled while already at 100%. + if (this._prevQuotaPercentUsed !== undefined || wasAdditionalUsageEnabled === false) { + this._showOverageActivationNotification(); + } + } else { + this._showExhaustedNotification(); } - } else { - this._showExhaustedNotification(); } // Keep the baseline up-to-date so that recovery from exhaustion @@ -410,4 +441,28 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo this._showingExhausted = false; this._chatInputNotificationService.deleteNotification(QUOTA_NOTIFICATION_ID); } + + // --- Exhausted dismissal persistence ------------------------------------ + + /** + * Returns `true` only when there is an actual quota snapshot indicating that + * credit is available (i.e. quota is not used up). Returns `false` when no + * snapshot has loaded yet, so the transient "no data" state at startup/reload + * is not mistaken for recovery. + */ + private _isQuotaKnownAvailable(): boolean { + return !!this._getRelevantSnapshot() && !this._isQuotaUsedUp(); + } + + private _isExhaustedDismissed(): boolean { + return this._storageService.getBoolean(QUOTA_EXHAUSTED_DISMISSED_STORAGE_KEY, StorageScope.APPLICATION, false); + } + + private _setExhaustedDismissed(): void { + this._storageService.store(QUOTA_EXHAUSTED_DISMISSED_STORAGE_KEY, true, StorageScope.APPLICATION, StorageTarget.MACHINE); + } + + private _clearExhaustedDismissed(): void { + this._storageService.remove(QUOTA_EXHAUSTED_DISMISSED_STORAGE_KEY, StorageScope.APPLICATION); + } } diff --git a/src/vs/workbench/contrib/chat/test/browser/chatQuotaNotification.test.ts b/src/vs/workbench/contrib/chat/test/browser/chatQuotaNotification.test.ts index 81108bd53e280a..99f9050e671368 100644 --- a/src/vs/workbench/contrib/chat/test/browser/chatQuotaNotification.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/chatQuotaNotification.test.ts @@ -127,6 +127,7 @@ function createMockNotificationService() { getNotification(): IChatInputNotification | undefined { return deleted || dismissed ? undefined : lastNotification; }, get wasDeleted() { return deleted; }, get setCount() { return setCount; }, + dismiss(id: string) { service.dismissNotification(id); }, reset() { lastNotification = undefined; deleted = false; dismissed = false; setCount = 0; }, }; } @@ -156,11 +157,11 @@ suite('ChatQuotaNotificationContribution', () => { const store = ensureNoDisposablesAreLeakedInTestSuite(); - function createContribution(entitlementOpts?: Parameters[0], modelOpts?: { vendor?: string }) { + function createContribution(entitlementOpts?: Parameters[0], modelOpts?: { vendor?: string }, sharedStorageService?: InMemoryStorageService) { const entitlementMock = createMockEntitlementService(entitlementOpts); const notificationMock = createMockNotificationService(); const contextKeyService = store.add(new MockContextKeyService()); - const storageService = store.add(new InMemoryStorageService()); + const storageService = sharedStorageService ?? store.add(new InMemoryStorageService()); const vendor = modelOpts?.vendor ?? 'copilot'; const isBYOK = vendor !== 'copilot'; // Persist model selection in storage (used by getSelectedModelVendor) @@ -273,6 +274,84 @@ suite('ChatQuotaNotificationContribution', () => { }); }); + // --- Exhausted dismissal persistence ------------------------------------ + + suite('exhausted dismissal persistence', () => { + test('does not re-show exhausted notification after reload when previously dismissed', () => { + const storageService = store.add(new InMemoryStorageService()); + + // First window: exhausted notification shown, then dismissed by the user. + const first = createContribution( + { quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(0) } }, + undefined, + storageService, + ); + const notification = first.notificationMock.getNotification(); + assert.ok(notification); + first.notificationMock.dismiss(notification!.id); + first.contribution.dispose(); + + // Reload: new contribution with the same (persisted) storage and still-exhausted quota. + const second = createContribution( + { quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(0) } }, + undefined, + storageService, + ); + assert.strictEqual(second.notificationMock.getNotification(), undefined); + }); + + test('re-shows exhausted notification after quota recovers and is exhausted again', () => { + const storageService = store.add(new InMemoryStorageService()); + + // Exhausted and dismissed. + const first = createContribution( + { quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(0) } }, + undefined, + storageService, + ); + first.notificationMock.dismiss(first.notificationMock.getNotification()!.id); + + // Quota recovers — persisted dismissal is cleared. + updateQuotas(first.entitlementMock, { premiumChat: makeQuotaSnapshot(50) }); + first.contribution.dispose(); + + // Reload while exhausted again — notification shows because the flag was cleared. + const second = createContribution( + { quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(0) } }, + undefined, + storageService, + ); + assert.ok(second.notificationMock.getNotification()); + assert.strictEqual(second.notificationMock.getNotification()!.message, 'Credit Limit Reached'); + }); + + test('keeps dismissal across reload when quota data is not loaded yet at startup', () => { + const storageService = store.add(new InMemoryStorageService()); + + // First window: exhausted notification shown, then dismissed by the user. + const first = createContribution( + { quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(0) } }, + undefined, + storageService, + ); + first.notificationMock.dismiss(first.notificationMock.getNotification()!.id); + first.contribution.dispose(); + + // Reload: quota snapshots have not been fetched yet (no relevant snapshot), + // so the dismissal must NOT be cleared by the transient "no data" state. + const second = createContribution( + { quotas: { usageBasedBilling: true, premiumChat: undefined } }, + undefined, + storageService, + ); + assert.strictEqual(second.notificationMock.getNotification(), undefined); + + // Quota data arrives showing it is still exhausted — banner stays suppressed. + updateQuotas(second.entitlementMock, { premiumChat: makeQuotaSnapshot(0) }); + assert.strictEqual(second.notificationMock.getNotification(), undefined); + }); + }); + // --- Exhausted notification descriptions -------------------------------- suite('exhausted notification descriptions', () => {