diff --git a/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts b/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts index 57d4bcd91dad41..a42a2eb4d1e425 100644 --- a/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts +++ b/src/vs/workbench/contrib/chat/browser/chatQuotaNotification.ts @@ -3,11 +3,19 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js'; import { safeIntl } from '../../../../base/common/date.js'; +import { createMarkdownCommandLink, MarkdownString } from '../../../../base/common/htmlContent.js'; +import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js'; +import { Language } from '../../../../base/common/platform.js'; +import { URI } from '../../../../base/common/uri.js'; import { localize } from '../../../../nls.js'; +import { CommandsRegistry } from '../../../../platform/commands/common/commands.js'; import { IContextKeyService } from '../../../../platform/contextkey/common/contextkey.js'; -import { IStorageService, StorageScope } from '../../../../platform/storage/common/storage.js'; +import { ServicesAccessor } from '../../../../platform/instantiation/common/instantiation.js'; +import { IOpenerService } from '../../../../platform/opener/common/opener.js'; +import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js'; +import { ITelemetryService } from '../../../../platform/telemetry/common/telemetry.js'; +import { IWorkbenchAssignmentService } from '../../../services/assignment/common/assignmentService.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'; @@ -16,6 +24,45 @@ import { ChatInputNotificationSeverity, IChatInputNotification, IChatInputNotifi const QUOTA_NOTIFICATION_ID = 'copilot.quotaStatus'; const THRESHOLDS = [50, 75, 90, 95]; +const TRAJECTORY_NUDGE_SPEC = { + treatment: 'config.chatQuotaTrajectoryNudge', + shownStorageKey: 'chat.quotaTrajectory.shownPeriod', + averageDailyUsageThreshold: 4.5, + minimumPercentUsed: 10, + maximumPercentUsed: 35, + billingPeriodDays: 30, + msPerDay: 24 * 60 * 60 * 1000, + learnMoreUrl: 'https://aka.ms/token-usage-tips', + learnMoreCommandId: 'workbench.action.chat.learnMoreAboutCreditUsage', +} as const; + +type ChatQuotaTrajectoryNudgeEvent = { + severity: 'info'; + entitlement: string; + averageDailyUsage: number; + percentUsed: number; +}; + +type ChatQuotaTrajectoryNudgeClassification = { + owner: 'rfeltis'; + comment: 'Tracks when users click the chat quota trajectory nudge learn more link.'; + severity: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The severity of the quota trajectory nudge.' }; + entitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The user entitlement when the quota trajectory nudge event was logged.' }; + averageDailyUsage: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The average daily monthly quota usage percentage that caused the nudge.' }; + percentUsed: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The monthly quota percentage used when the quota trajectory nudge event was logged.' }; +}; + +type ChatQuotaTrajectoryNudgeEnrollmentEvent = { + treatment: boolean; + entitlement: string; +}; + +type ChatQuotaTrajectoryNudgeEnrollmentClassification = { + owner: 'rfeltis'; + comment: 'Tracks when a user is assigned to a flight for the chat quota trajectory nudge experiment, to measure experiment exposure.'; + treatment: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The treatment value assigned by the experiment service (true for the treatment arm, false for control).' }; + entitlement: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The user entitlement when the user was assigned to the experiment flight.' }; +}; /** * Core-side workbench contribution that shows chat input notifications for @@ -45,6 +92,9 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo private _prevAdditionalUsageEnabled: boolean | undefined; private _prevSessionPercentUsed: number | undefined; private _prevWeeklyPercentUsed: number | undefined; + private _trajectoryTreatment: boolean | undefined; + private _trajectoryAssignmentRequested = false; + private _activeTrajectoryTelemetryData: ChatQuotaTrajectoryNudgeEvent | undefined; constructor( @IChatEntitlementService private readonly _chatEntitlementService: IChatEntitlementService, @@ -52,12 +102,15 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo @IContextKeyService private readonly _contextKeyService: IContextKeyService, @ILanguageModelsService private readonly _languageModelsService: ILanguageModelsService, @IStorageService private readonly _storageService: IStorageService, + @IWorkbenchAssignmentService private readonly _assignmentService: IWorkbenchAssignmentService, + @ITelemetryService private readonly _telemetryService: ITelemetryService, ) { super(); this._register(this._chatEntitlementService.onDidChangeQuotaRemaining(() => this._update())); this._register(this._chatEntitlementService.onDidChangeQuotaExceeded(() => this._update())); this._register(this._chatEntitlementService.onDidChangeEntitlement(() => this._update())); + this._register(CommandsRegistry.registerCommand(TRAJECTORY_NUDGE_SPEC.learnMoreCommandId, (accessor: ServicesAccessor) => this._handleCreditEfficiencyLearnMoreCommand(accessor))); // Re-evaluate when the selected model changes (e.g. switching between Copilot and BYOK). // The chatModelId context key is widget-scoped and may not bubble to the global @@ -73,6 +126,27 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo this._update(); } + /** + * Enrolls the user in the trajectory experiment. Called lazily, only once + * the user has met every condition required to render the nudge, so that + * the experiment is not diluted with users who would never be exposed. + * + * Stores the raw treatment value. `undefined` means the user is not + * assigned to the flight (or assignments are not available); only a `true` + * treatment renders the nudge. We deliberately do not coerce a missing + * assignment into a synthetic "control" value, since that would assume an + * enrollment that may not exist. Enrollment telemetry is emitted only when + * the user is actually assigned to a flight. + */ + private async _resolveTrajectoryTreatment(): Promise { + const treatment = await this._assignmentService.getTreatment(TRAJECTORY_NUDGE_SPEC.treatment); + this._trajectoryTreatment = treatment; + if (treatment !== undefined) { + this._logQuotaTrajectoryNudgeEnrolled(treatment); + } + this._update(); + } + private _getRelevantSnapshot(): IQuotaSnapshot | undefined { const quotas = this._chatEntitlementService.quotas; const entitlement = this._chatEntitlementService.entitlement; @@ -148,6 +222,12 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo // Priority 2: Quota approaching threshold if (isQuotaNotificationEligible) { + const trajectoryWarning = this._computeQuotaTrajectoryWarning(); + if (trajectoryWarning) { + this._showQuotaTrajectoryWarning(trajectoryWarning); + return; + } + const quotaWarning = this._computeQuotaWarning(); if (quotaWarning) { this._showQuotaApproachingWarning(quotaWarning); @@ -186,6 +266,103 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo return undefined; } + private _computeQuotaTrajectoryWarning(): { averageDailyUsage: number; percentUsed: number } | undefined { + if (!Language.isDefaultVariant() || !this._isTrajectoryEligibleEntitlement() || this._isTrajectoryShownInCurrentPeriod()) { + return undefined; + } + + const snapshot = this._getRelevantSnapshot(); + if (!snapshot || snapshot.unlimited || snapshot.percentRemaining <= 0) { + return undefined; + } + + const resetDate = this._chatEntitlementService.quotas.resetDate; + if (!resetDate) { + return undefined; + } + + const resetTime = new Date(resetDate).getTime(); + if (!Number.isFinite(resetTime)) { + return undefined; + } + + const periodStartTime = resetTime - (TRAJECTORY_NUDGE_SPEC.billingPeriodDays * TRAJECTORY_NUDGE_SPEC.msPerDay); + const elapsedDays = Math.max(0, (Date.now() - periodStartTime) / TRAJECTORY_NUDGE_SPEC.msPerDay); + if (elapsedDays <= 0) { + return undefined; + } + + const percentUsed = 100 - snapshot.percentRemaining; + if (percentUsed < TRAJECTORY_NUDGE_SPEC.minimumPercentUsed || percentUsed > TRAJECTORY_NUDGE_SPEC.maximumPercentUsed) { + return undefined; + } + + const averageDailyUsage = percentUsed / elapsedDays; + if (averageDailyUsage < TRAJECTORY_NUDGE_SPEC.averageDailyUsageThreshold) { + return undefined; + } + + // Every render condition is met. Enroll in the experiment now + // (just-in-time) so treatment and control cohorts are assigned at the + // same point and users who would never see the nudge are not exposed. + // While assignment is pending, fall through; _update re-runs once it + // resolves and only the treatment cohort then renders the nudge. + if (!this._trajectoryAssignmentRequested) { + this._trajectoryAssignmentRequested = true; + void this._resolveTrajectoryTreatment(); + } + return this._trajectoryTreatment === true ? { averageDailyUsage, percentUsed } : undefined; + } + + private _showQuotaTrajectoryWarning(warning: { averageDailyUsage: number; percentUsed: number }): void { + this._showingExhausted = false; + this._activeTrajectoryTelemetryData = this._getQuotaTrajectoryNudgeTelemetryData(warning); + this._storeTrajectoryShown(); + const learnMoreLink = createMarkdownCommandLink({ + text: localize('quota.trajectory.learnMore', "Learn more about managing credits"), + id: TRAJECTORY_NUDGE_SPEC.learnMoreCommandId, + tooltip: localize('quota.trajectory.learnMoreTooltip', "Learn more about managing credits"), + }); + + this._setNotification({ + id: QUOTA_NOTIFICATION_ID, + severity: ChatInputNotificationSeverity.Info, + message: new MarkdownString(localize({ key: 'quota.trajectory.message', comment: ['{Locked="]({0})"}'] }, "Based on recent usage, your monthly allowance may run out before it resets. {0}", learnMoreLink), { isTrusted: { enabledCommands: [TRAJECTORY_NUDGE_SPEC.learnMoreCommandId] } }), + description: undefined, + actions: [], + dismissible: true, + autoDismissOnMessage: false, + }); + } + + private async _handleCreditEfficiencyLearnMoreCommand(accessor: ServicesAccessor): Promise { + if (this._activeTrajectoryTelemetryData) { + this._logQuotaTrajectoryNudgeLinkClicked(this._activeTrajectoryTelemetryData); + queueMicrotask(() => this._hideNotification()); + } + await accessor.get(IOpenerService).open(URI.parse(TRAJECTORY_NUDGE_SPEC.learnMoreUrl)); + } + + private _logQuotaTrajectoryNudgeLinkClicked(data: ChatQuotaTrajectoryNudgeEvent): void { + this._telemetryService.publicLog2('chatQuotaTrajectoryNudgeLinkClicked', data); + } + + private _logQuotaTrajectoryNudgeEnrolled(treatment: boolean): void { + this._telemetryService.publicLog2('chatQuotaTrajectoryNudgeEnrolled', { + treatment, + entitlement: ChatEntitlement[this._chatEntitlementService.entitlement], + }); + } + + private _getQuotaTrajectoryNudgeTelemetryData(warning: { averageDailyUsage: number; percentUsed: number }): ChatQuotaTrajectoryNudgeEvent { + return { + severity: 'info', + entitlement: ChatEntitlement[this._chatEntitlementService.entitlement], + averageDailyUsage: Math.round(warning.averageDailyUsage * 100) / 100, + percentUsed: Math.round(warning.percentUsed * 100) / 100, + }; + } + /** * Returns the highest threshold that was newly crossed, or `undefined`. */ @@ -206,6 +383,7 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo private _showExhaustedNotification(): void { this._showingExhausted = true; + this._activeTrajectoryTelemetryData = undefined; const entitlement = this._chatEntitlementService.entitlement; const quotas = this._chatEntitlementService.quotas; @@ -247,6 +425,7 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo private _showOverageActivationNotification(): void { this._showingExhausted = true; + this._activeTrajectoryTelemetryData = undefined; this._setNotification({ id: QUOTA_NOTIFICATION_ID, @@ -264,6 +443,7 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo private _showQuotaApproachingWarning(warning: { percentUsed: number }): void { this._showingExhausted = false; + this._activeTrajectoryTelemetryData = undefined; const entitlement = this._chatEntitlementService.entitlement; const quotas = this._chatEntitlementService.quotas; @@ -336,6 +516,7 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo private _showRateLimitWarning(warning: { percentUsed: number; type: 'session' | 'weekly'; resetDate: string | undefined }): void { this._showingExhausted = false; + this._activeTrajectoryTelemetryData = undefined; const message = warning.type === 'session' ? localize('rateLimit.session', "You've used {0}% of your session rate limit.", warning.percentUsed) @@ -372,6 +553,11 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo return entitlement === ChatEntitlement.Business || entitlement === ChatEntitlement.Enterprise; } + private _isTrajectoryEligibleEntitlement(): boolean { + const entitlement = this._chatEntitlementService.entitlement; + return entitlement === ChatEntitlement.Pro || entitlement === ChatEntitlement.ProPlus || entitlement === ChatEntitlement.Max; + } + private _isManagedPlanBlocked(): boolean { const snapshot = this._chatEntitlementService.quotas.premiumChat; return !!snapshot && snapshot.hasQuota === false; @@ -379,6 +565,7 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo private _showManagedPlanBlockedNotification(): void { this._showingExhausted = true; + this._activeTrajectoryTelemetryData = undefined; this._setNotification({ id: QUOTA_NOTIFICATION_ID, @@ -402,12 +589,37 @@ export class ChatQuotaNotificationContribution extends Disposable implements IWo ).value.format(resetDate); } + private _getTrajectoryPeriodKey(): string | undefined { + const resetDate = this._chatEntitlementService.quotas.resetDate; + if (!resetDate) { + return undefined; + } + const date = new Date(resetDate); + if (!Number.isFinite(date.getTime())) { + return undefined; + } + return `${date.getUTCFullYear()}-${date.getUTCMonth() + 1}`; + } + + private _isTrajectoryShownInCurrentPeriod(): boolean { + const periodKey = this._getTrajectoryPeriodKey(); + return !!periodKey && this._storageService.get(TRAJECTORY_NUDGE_SPEC.shownStorageKey, StorageScope.APPLICATION) === periodKey; + } + + private _storeTrajectoryShown(): void { + const periodKey = this._getTrajectoryPeriodKey(); + if (periodKey) { + this._storageService.store(TRAJECTORY_NUDGE_SPEC.shownStorageKey, periodKey, StorageScope.APPLICATION, StorageTarget.USER); + } + } + private _setNotification(notification: IChatInputNotification): void { this._chatInputNotificationService.setNotification(notification); } private _hideNotification(): void { this._showingExhausted = false; + this._activeTrajectoryTelemetryData = undefined; this._chatInputNotificationService.deleteNotification(QUOTA_NOTIFICATION_ID); } } diff --git a/src/vs/workbench/contrib/chat/browser/widget/input/media/chatInputNotificationWidget.css b/src/vs/workbench/contrib/chat/browser/widget/input/media/chatInputNotificationWidget.css index 5863b79faaa134..586a9f35180633 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/input/media/chatInputNotificationWidget.css +++ b/src/vs/workbench/contrib/chat/browser/widget/input/media/chatInputNotificationWidget.css @@ -106,6 +106,16 @@ margin: 0; } +.chat-input-notification .chat-input-notification-title-markdown a, +.chat-input-notification .chat-input-notification-title-markdown a:visited { + color: var(--vscode-textLink-foreground); +} + +.chat-input-notification .chat-input-notification-title-markdown a:hover, +.chat-input-notification .chat-input-notification-title-markdown a:active { + color: var(--vscode-textLink-activeForeground); +} + .chat-input-notification .chat-input-notification-title-markdown code { font-family: var(--monaco-monospace-font); font-size: var(--vscode-agents-fontSize-body2); 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..2a0bfda0ccebb7 100644 --- a/src/vs/workbench/contrib/chat/test/browser/chatQuotaNotification.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/chatQuotaNotification.test.ts @@ -4,20 +4,35 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import * as sinon from 'sinon'; import { Emitter, Event } from '../../../../../base/common/event.js'; +import { createMarkdownCommandLink } from '../../../../../base/common/htmlContent.js'; +import { Disposable } from '../../../../../base/common/lifecycle.js'; import { IObservable, observableValue } from '../../../../../base/common/observable.js'; +import { Language } from '../../../../../base/common/platform.js'; +import { URI } from '../../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { CommandsRegistry } from '../../../../../platform/commands/common/commands.js'; import { IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js'; +import { ServiceIdentifier, ServicesAccessor } from '../../../../../platform/instantiation/common/instantiation.js'; import { MockContextKeyService } from '../../../../../platform/keybinding/test/common/mockKeybindingService.js'; +import { IOpenerService } from '../../../../../platform/opener/common/opener.js'; import { InMemoryStorageService, StorageScope, StorageTarget } from '../../../../../platform/storage/common/storage.js'; +import { ITelemetryService } from '../../../../../platform/telemetry/common/telemetry.js'; +import { NullTelemetryServiceShape } from '../../../../../platform/telemetry/common/telemetryUtils.js'; +import { IWorkbenchAssignmentService } from '../../../../services/assignment/common/assignmentService.js'; import { ChatEntitlement, IChatEntitlementService, IChatSentiment, IQuotaSnapshot, IRateLimitSnapshot } from '../../../../services/chat/common/chatEntitlementService.js'; import { ChatQuotaNotificationContribution } from '../../browser/chatQuotaNotification.js'; import { ILanguageModelChatMetadata, ILanguageModelsService } from '../../common/languageModels.js'; -import { IChatInputNotification, IChatInputNotificationService } from '../../browser/widget/input/chatInputNotificationService.js'; +import { ChatInputNotificationSeverity, IChatInputNotification, IChatInputNotificationService } from '../../browser/widget/input/chatInputNotificationService.js'; + +const CREDIT_EFFICIENCY_LEARN_MORE_COMMAND_ID = 'workbench.action.chat.learnMoreAboutCreditUsage'; +const CREDIT_EFFICIENCY_LEARN_MORE_URL = 'https://aka.ms/token-usage-tips'; // --- Mock IChatEntitlementService ------------------------------------------- interface IMockQuotas { + resetDate?: string; usageBasedBilling?: boolean; chat?: IQuotaSnapshot; completions?: IQuotaSnapshot; @@ -45,6 +60,7 @@ function createMockEntitlementService(opts?: { onDidChangeQuotaRemaining: onDidChangeQuotaRemaining.event, onDidChangeUsageBasedBilling: Event.None, quotas: { + resetDate: opts?.quotas?.resetDate, usageBasedBilling: opts?.quotas?.usageBasedBilling ?? true, chat: opts?.quotas?.chat, completions: opts?.quotas?.completions, @@ -100,13 +116,15 @@ function createMockNotificationService() { setCount++; onDidChange.fire(); }, - deleteNotification(_id: string) { - deleted = true; - dismissed = false; - onDidChange.fire(); + deleteNotification(id: string) { + if (lastNotification?.id === id && !deleted) { + deleted = true; + dismissed = false; + onDidChange.fire(); + } }, dismissNotification(id: string) { - if (!lastNotification || lastNotification.id !== id || dismissed) { + if (!lastNotification || lastNotification.id !== id || deleted || dismissed) { return; } dismissed = true; @@ -127,10 +145,42 @@ function createMockNotificationService() { getNotification(): IChatInputNotification | undefined { return deleted || dismissed ? undefined : lastNotification; }, get wasDeleted() { return deleted; }, get setCount() { return setCount; }, + dismiss() { + if (lastNotification) { + service.dismissNotification(lastNotification.id); + } + }, reset() { lastNotification = undefined; deleted = false; dismissed = false; setCount = 0; }, }; } +function createMockAssignmentService(treatments?: Readonly>>) { + const onDidRefetchAssignments = new Emitter(); + const getTreatmentCalls: string[] = []; + const service: IWorkbenchAssignmentService = { + _serviceBrand: undefined, + onDidRefetchAssignments: onDidRefetchAssignments.event, + getTreatment(name: string) { + getTreatmentCalls.push(name); + return Promise.resolve(treatments?.[name]); + }, + getCurrentExperiments() { return Promise.resolve(undefined); }, + addTelemetryAssignmentFilter() { }, + } as unknown as IWorkbenchAssignmentService; + + return { service, onDidRefetchAssignments, getTreatmentCalls }; +} + +class TestTelemetryService extends NullTelemetryServiceShape { + readonly events: { name: string; data: unknown }[] = []; + + override publicLog2(eventName?: string, data?: unknown): void { + if (eventName) { + this.events.push({ name: eventName, data }); + } + } +} + // --- Helpers --------------------------------------------------------------- function makeQuotaSnapshot(percentRemaining: number, opts?: Partial): IQuotaSnapshot { @@ -141,6 +191,10 @@ function makeQuotaSnapshot(percentRemaining: number, opts?: Partial { + await new Promise(resolve => setTimeout(resolve, 0)); +} + function makeRateLimitSnapshot(percentRemaining: number, opts?: Partial): IRateLimitSnapshot { return { percentRemaining, @@ -150,15 +204,56 @@ function makeRateLimitSnapshot(percentRemaining: number, opts?: Partial { + const opened: URI[] = []; + const openerService: IOpenerService = { + _serviceBrand: undefined, + registerOpener: () => Disposable.None, + registerValidator: () => Disposable.None, + registerExternalUriResolver: () => Disposable.None, + setDefaultExternalOpener: () => { }, + registerExternalOpener: () => Disposable.None, + open: async resource => { + opened.push(URI.isUri(resource) ? resource : URI.parse(resource)); + return true; + }, + resolveExternalUri: async resource => ({ resolved: resource, dispose: () => { } }), + }; + const accessor: ServicesAccessor = { + get: (id: ServiceIdentifier): T => { + if (id === IOpenerService) { + return openerService as T; + } + throw new Error('Unexpected service'); + }, + }; + const command = CommandsRegistry.getCommand(CREDIT_EFFICIENCY_LEARN_MORE_COMMAND_ID); + assert.ok(command); + await command.handler(accessor); + return opened; +} + +function makeResetDate(daysUntilReset: number): string { + const resetDate = new Date(Date.now() + daysUntilReset * 24 * 60 * 60 * 1000); + return resetDate.toISOString(); +} + // --- Tests ----------------------------------------------------------------- suite('ChatQuotaNotificationContribution', () => { const store = ensureNoDisposablesAreLeakedInTestSuite(); - function createContribution(entitlementOpts?: Parameters[0], modelOpts?: { vendor?: string }) { + teardown(() => { + sinon.restore(); + }); + + function createContribution(entitlementOpts?: Parameters[0], modelOpts?: { vendor?: string; trajectoryTreatment?: boolean | Promise; telemetryService?: ITelemetryService }) { const entitlementMock = createMockEntitlementService(entitlementOpts); const notificationMock = createMockNotificationService(); + const assignmentMock = createMockAssignmentService({ + 'config.chatQuotaTrajectoryNudge': modelOpts?.trajectoryTreatment, + }); const contextKeyService = store.add(new MockContextKeyService()); const storageService = store.add(new InMemoryStorageService()); const vendor = modelOpts?.vendor ?? 'copilot'; @@ -179,6 +274,7 @@ suite('ChatQuotaNotificationContribution', () => { store.add(entitlementMock.onDidChangeQuotaRemaining); store.add(entitlementMock.onDidChangeQuotaExceeded); store.add(entitlementMock.onDidChangeEntitlement); + store.add(assignmentMock.onDidRefetchAssignments); const contribution = store.add(new ChatQuotaNotificationContribution( entitlementMock.service, @@ -186,9 +282,11 @@ suite('ChatQuotaNotificationContribution', () => { contextKeyService as IContextKeyService, languageModelsService, storageService, + assignmentMock.service, + modelOpts?.telemetryService ?? new NullTelemetryServiceShape(), )); - return { contribution, entitlementMock, notificationMock, storageService }; + return { contribution, entitlementMock, notificationMock, storageService, assignmentMock }; } function updateQuotas( @@ -359,31 +457,35 @@ suite('ChatQuotaNotificationContribution', () => { // --- Quota approaching threshold ---------------------------------------- suite('quota approaching threshold', () => { - test('first data arrival stores baseline without notification', () => { + test('first data arrival stores baseline without notification', async () => { const { notificationMock } = createContribution({ quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(25) }, // 75% used }); - // Initial _update runs in constructor but 75% is baseline, no crossing + await flushPromises(); + + // First data arrival stores 75% as the baseline without notifying. assert.strictEqual(notificationMock.getNotification(), undefined); }); - test('notifies when crossing 50% threshold', () => { + test('notifies when crossing 50% threshold', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(60) }, // 40% used baseline }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(50) }); // 50% used assert.ok(notificationMock.getNotification()); assert.strictEqual(notificationMock.getNotification()!.message, 'Credits at 50%'); }); - test('does not re-show the same threshold', () => { + test('does not re-show the same threshold', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(60) }, }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(50) }); assert.ok(notificationMock.getNotification()); @@ -394,11 +496,12 @@ suite('ChatQuotaNotificationContribution', () => { assert.strictEqual(notificationMock.getNotification(), undefined); }); - test('shows higher threshold when usage increases', () => { + test('shows higher threshold when usage increases', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(60) }, }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(50) }); // 50% assert.strictEqual(notificationMock.getNotification()!.message, 'Credits at 50%'); @@ -419,12 +522,13 @@ suite('ChatQuotaNotificationContribution', () => { assert.strictEqual(notificationMock.getNotification(), undefined); }); - test('does not show approaching notification for PRU user', () => { + test('does not show approaching notification for PRU user', async () => { const { entitlementMock, notificationMock } = createContribution({ entitlement: ChatEntitlement.Pro, quotas: { usageBasedBilling: false, premiumChat: makeQuotaSnapshot(60) }, }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(5) }); assert.strictEqual(notificationMock.getNotification(), undefined); }); @@ -479,14 +583,368 @@ suite('ChatQuotaNotificationContribution', () => { }); }); + // --- Quota trajectory warning -------------------------------------------- + + suite('quota trajectory warning', () => { + test('does not show when experiment treatment is disabled', async () => { + const { notificationMock } = createContribution({ + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }); + + await flushPromises(); + + assert.strictEqual(notificationMock.getNotification(), undefined); + }); + + test('does not enroll when UI language is not translated', async () => { + sinon.stub(Language, 'isDefaultVariant').returns(false); + const { assignmentMock, notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + assert.deepStrictEqual({ + treatments: assignmentMock.getTreatmentCalls, + notification: notificationMock.getNotification(), + }, { + treatments: [], + notification: undefined, + }); + }); + + test('does not show when user is eligible but not assigned to the experiment', async () => { + // No treatment configured: getTreatment resolves to undefined, i.e. + // the user is not in the flight. This must not be treated as control + // enrollment, but it should still attempt exposure since the user met + // every render condition. + const { assignmentMock, notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }); + + await flushPromises(); + + assert.deepStrictEqual({ + treatments: assignmentMock.getTreatmentCalls, + notification: notificationMock.getNotification(), + }, { + treatments: ['config.chatQuotaTrajectoryNudge'], + notification: undefined, + }); + }); + + test('does not show outside monthly usage window', async () => { + const results = []; + for (const percentRemaining of [91, 64]) { + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(percentRemaining), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + results.push(notificationMock.getNotification()?.message); + } + + assert.deepStrictEqual(results, [undefined, undefined]); + }); + + test('shows info notification when projected daily usage is above threshold', async () => { + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + const notification = notificationMock.getNotification(); + assert.ok(notification); + const message = notification.message; + const learnMoreLink = createMarkdownCommandLink({ + text: 'Learn more about managing credits', + id: CREDIT_EFFICIENCY_LEARN_MORE_COMMAND_ID, + tooltip: 'Learn more about managing credits', + }); + assert.deepStrictEqual({ + message: typeof message === 'string' ? message : message.value, + severity: notification.severity, + actions: notification.actions.length, + autoDismissOnMessage: notification.autoDismissOnMessage, + }, { + message: `Based on recent usage, your monthly allowance may run out before it resets. ${learnMoreLink}`, + severity: ChatInputNotificationSeverity.Info, + actions: 0, + autoDismissOnMessage: false, + }); + }); + + test('shows for Pro+ users', async () => { + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.ProPlus, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + assert.ok(notificationMock.getNotification()); + }); + + test('shows for Max users', async () => { + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Max, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + assert.ok(notificationMock.getNotification()); + }); + + test('does not show when projected daily usage is below threshold', async () => { + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(78), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + assert.strictEqual(notificationMock.getNotification(), undefined); + }); + + test('does not show when reset date implies no elapsed billing days', async () => { + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(31), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + assert.strictEqual(notificationMock.getNotification(), undefined); + }); + + test('shows trajectory nudge only after treatment resolves', async () => { + let resolveTreatment: ((value: boolean | undefined) => void) | undefined; + const trajectoryTreatment = new Promise(resolve => { + resolveTreatment = resolve; + }); + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment }); + + await flushPromises(); + assert.strictEqual(notificationMock.getNotification(), undefined); + + assert.ok(resolveTreatment); + resolveTreatment(true); + await flushPromises(); + + const notification = notificationMock.getNotification(); + assert.ok(notification); + const message = notification.message; + assert.ok(typeof message !== 'string' && message.value.includes('monthly allowance')); + }); + + test('logs link click telemetry', async () => { + const telemetryService = new TestTelemetryService(); + createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true, telemetryService }); + + await flushPromises(); + const opened = await runCreditEfficiencyLearnMoreCommand(); + + assert.deepStrictEqual({ + events: telemetryService.events, + opened: opened.map(uri => uri.toString()), + }, { + events: [ + { + name: 'chatQuotaTrajectoryNudgeEnrolled', + data: { + treatment: true, + entitlement: 'Pro', + }, + }, + { + name: 'chatQuotaTrajectoryNudgeLinkClicked', + data: { + severity: 'info', + entitlement: 'Pro', + averageDailyUsage: 4.67, + percentUsed: 28, + }, + }, + ], + opened: [CREDIT_EFFICIENCY_LEARN_MORE_URL], + }); + }); + + test('logs enrollment telemetry for control assignment without showing nudge', async () => { + const telemetryService = new TestTelemetryService(); + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: false, telemetryService }); + + await flushPromises(); + + assert.deepStrictEqual({ + events: telemetryService.events, + notification: notificationMock.getNotification(), + }, { + events: [{ + name: 'chatQuotaTrajectoryNudgeEnrolled', + data: { treatment: false, entitlement: 'Pro' }, + }], + notification: undefined, + }); + }); + + test('does not log enrollment telemetry when not assigned to a flight', async () => { + const telemetryService = new TestTelemetryService(); + const { notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { telemetryService }); // no treatment configured -> not assigned to the flight + + await flushPromises(); + + assert.deepStrictEqual({ + events: telemetryService.events, + notification: notificationMock.getNotification(), + }, { + events: [], + notification: undefined, + }); + }); + + test('action click dismisses trajectory nudge for the quota period', async () => { + const { entitlementMock, notificationMock } = createContribution({ + entitlement: ChatEntitlement.Pro, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + assert.ok(notificationMock.getNotification()); + + await runCreditEfficiencyLearnMoreCommand(); + await flushPromises(); + + assert.strictEqual(notificationMock.getNotification(), undefined); + + notificationMock.reset(); + entitlementMock.onDidChangeQuotaRemaining.fire(); + + assert.strictEqual(notificationMock.getNotification(), undefined); + }); + + test('remembers trajectory display for the quota period', async () => { + const { entitlementMock, notificationMock } = createContribution({ + entitlement: ChatEntitlement.ProPlus, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + assert.ok(notificationMock.getNotification()); + + notificationMock.reset(); + entitlementMock.onDidChangeQuotaRemaining.fire(); + + assert.strictEqual(notificationMock.getNotification(), undefined); + }); + + test('does not show for Edu, Business, Enterprise, Free, or Unknown users', async () => { + for (const entitlement of [ChatEntitlement.EDU, ChatEntitlement.Business, ChatEntitlement.Enterprise, ChatEntitlement.Free, ChatEntitlement.Unknown]) { + const { notificationMock } = createContribution({ + entitlement, + quotas: { + resetDate: makeResetDate(24), + usageBasedBilling: true, + premiumChat: makeQuotaSnapshot(72), + chat: makeQuotaSnapshot(72), + }, + }, { trajectoryTreatment: true }); + + await flushPromises(); + + assert.strictEqual(notificationMock.getNotification(), undefined, `Expected no trajectory notification for ${entitlement}`); + } + }); + }); + // --- Rate-limit warnings ------------------------------------------------ suite('rate-limit warnings', () => { - test('shows session rate limit warning on threshold crossing', () => { + test('shows session rate limit warning on threshold crossing', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, sessionRateLimit: makeRateLimitSnapshot(60) }, // baseline }); + await flushPromises(); updateQuotas(entitlementMock, { sessionRateLimit: makeRateLimitSnapshot(25) }); // 75% used assert.ok(notificationMock.getNotification()); @@ -494,11 +952,12 @@ suite('ChatQuotaNotificationContribution', () => { assert.ok((notificationMock.getNotification()!.message as string).includes('session')); }); - test('shows weekly rate limit warning on threshold crossing', () => { + test('shows weekly rate limit warning on threshold crossing', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, weeklyRateLimit: makeRateLimitSnapshot(60) }, // baseline }); + await flushPromises(); updateQuotas(entitlementMock, { weeklyRateLimit: makeRateLimitSnapshot(10) }); // 90% used assert.ok(notificationMock.getNotification()); @@ -506,11 +965,12 @@ suite('ChatQuotaNotificationContribution', () => { assert.ok((notificationMock.getNotification()!.message as string).includes('weekly')); }); - test('first rate limit data stores baseline without notification', () => { + test('first rate limit data stores baseline without notification', async () => { const { notificationMock } = createContribution({ quotas: { usageBasedBilling: true, sessionRateLimit: makeRateLimitSnapshot(10) }, // 90% used }); + await flushPromises(); assert.strictEqual(notificationMock.getNotification(), undefined); }); }); @@ -527,7 +987,7 @@ suite('ChatQuotaNotificationContribution', () => { assert.strictEqual(notificationMock.getNotification()!.message, 'Credit Limit Reached'); }); - test('approaching threshold takes priority over rate limit', () => { + test('approaching threshold takes priority over rate limit', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, @@ -536,6 +996,7 @@ suite('ChatQuotaNotificationContribution', () => { }, }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(10), // 90% — crosses threshold sessionRateLimit: makeRateLimitSnapshot(25), // 75% — crosses threshold @@ -549,46 +1010,50 @@ suite('ChatQuotaNotificationContribution', () => { // --- Approaching notification descriptions ------------------------------ suite('approaching notification descriptions', () => { - test('free user gets upgrade action', () => { + test('free user gets upgrade action', async () => { const { entitlementMock, notificationMock } = createContribution({ entitlement: ChatEntitlement.Free, quotas: { usageBasedBilling: true, chat: makeQuotaSnapshot(60) }, }); + await flushPromises(); updateQuotas(entitlementMock, { chat: makeQuotaSnapshot(50) }); assert.ok(notificationMock.getNotification()); assert.strictEqual(notificationMock.getNotification()!.description, 'Upgrade to continue past the limit.'); }); - test('managed plan user gets admin message', () => { + test('managed plan user gets admin message', async () => { const { entitlementMock, notificationMock } = createContribution({ entitlement: ChatEntitlement.Enterprise, quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(60) }, }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(50) }); assert.ok(notificationMock.getNotification()); assert.strictEqual(notificationMock.getNotification()!.description, 'Contact your admin to increase your limits.'); }); - test('paid user with overages enabled gets budget message', () => { + test('paid user with overages enabled gets budget message', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(60), additionalUsageEnabled: true }, }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(50) }); assert.ok(notificationMock.getNotification()); assert.strictEqual(notificationMock.getNotification()!.description, 'Additional budget is enabled to cover extra usage.'); }); - test('paid user without overages gets set budget action', () => { + test('paid user without overages gets set budget action', async () => { const { entitlementMock, notificationMock } = createContribution({ quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(60) }, }); + await flushPromises(); updateQuotas(entitlementMock, { premiumChat: makeQuotaSnapshot(50) }); assert.ok(notificationMock.getNotification()); @@ -624,6 +1089,7 @@ suite('ChatQuotaNotificationContribution', () => { quotas: { usageBasedBilling: true, premiumChat: makeQuotaSnapshot(0) }, }); const notificationMock = createMockNotificationService(); + const assignmentMock = createMockAssignmentService(); const contextKeyService = store.add(new MockContextKeyService()); const storageService = store.add(new InMemoryStorageService()); // Start with BYOK model @@ -642,6 +1108,7 @@ suite('ChatQuotaNotificationContribution', () => { store.add(entitlementMock.onDidChangeQuotaRemaining); store.add(entitlementMock.onDidChangeQuotaExceeded); store.add(entitlementMock.onDidChangeEntitlement); + store.add(assignmentMock.onDidRefetchAssignments); store.add(new ChatQuotaNotificationContribution( entitlementMock.service, @@ -649,6 +1116,8 @@ suite('ChatQuotaNotificationContribution', () => { contextKeyService as IContextKeyService, languageModelsService, storageService, + assignmentMock.service, + new NullTelemetryServiceShape(), )); // Initially deferred — BYOK model