From 9a7427197bd2d8df9bf4cf5f723b4932557c5e0e Mon Sep 17 00:00:00 2001 From: Gaurav <39389231+gsquared94@users.noreply.github.com> Date: Fri, 6 Mar 2026 19:14:44 -0800 Subject: [PATCH] fix(billing): fix overage strategy lifecycle and settings integration (#21236) --- packages/cli/src/config/config.ts | 1 + .../cli/src/ui/hooks/creditsFlowHandler.ts | 8 --- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 1 + packages/cli/src/ui/hooks/useGeminiStream.ts | 4 ++ .../cli/src/ui/hooks/useQuotaAndFallback.ts | 13 +--- packages/core/src/billing/billing.test.ts | 6 +- packages/core/src/billing/billing.ts | 2 + packages/core/src/code_assist/server.ts | 6 ++ packages/core/src/config/config.test.ts | 59 +++++++++++++++++++ packages/core/src/config/config.ts | 15 +++++ 10 files changed, 92 insertions(+), 23 deletions(-) diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 0d81fa39bc..a8c85975e9 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -805,6 +805,7 @@ export async function loadCliConfig( fakeResponses: argv.fakeResponses, recordResponses: argv.recordResponses, retryFetchErrors: settings.general?.retryFetchErrors, + billing: settings.billing, maxAttempts: settings.general?.maxAttempts, ptyInfo: ptyInfo?.name, disableLLMCorrection: settings.tools?.disableLLMCorrection, diff --git a/packages/cli/src/ui/hooks/creditsFlowHandler.ts b/packages/cli/src/ui/hooks/creditsFlowHandler.ts index 497d4904e6..91f0997873 100644 --- a/packages/cli/src/ui/hooks/creditsFlowHandler.ts +++ b/packages/cli/src/ui/hooks/creditsFlowHandler.ts @@ -110,7 +110,6 @@ async function handleOverageMenu( isDialogPending, setOverageMenuRequest, setModelSwitchedFromQuotaError, - historyManager, } = args; logBillingEvent( @@ -155,13 +154,6 @@ async function handleOverageMenu( setModelSwitchedFromQuotaError(false); config.setQuotaErrorOccurred(false); config.setOverageStrategy('always'); - historyManager.addItem( - { - type: MessageType.INFO, - text: `Using AI Credits for this request.`, - }, - Date.now(), - ); return 'retry_with_credits'; case 'use_fallback': diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index ec8ea0751a..cfffb28196 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -271,6 +271,7 @@ describe('useGeminiStream', () => { addHistory: vi.fn(), getSessionId: vi.fn(() => 'test-session-id'), setQuotaErrorOccurred: vi.fn(), + resetBillingTurnState: vi.fn(), getQuotaErrorOccurred: vi.fn(() => false), getModel: vi.fn(() => 'gemini-2.5-pro'), getContentGeneratorConfig: vi.fn(() => ({ diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 630566090b..b0b4f553a2 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -1376,6 +1376,9 @@ export const useGeminiStream = ( if (!options?.isContinuation) { setModelSwitchedFromQuotaError(false); config.setQuotaErrorOccurred(false); + config.resetBillingTurnState( + settings.merged.billing?.overageStrategy, + ); suppressedToolErrorCountRef.current = 0; suppressedToolErrorNoteShownRef.current = false; lowVerbosityFailureNoteShownRef.current = false; @@ -1536,6 +1539,7 @@ export const useGeminiStream = ( setThought, maybeAddSuppressedToolErrorNote, maybeAddLowVerbosityFailureNote, + settings.merged.billing?.overageStrategy, ], ); diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts index 40b1f68926..533eefa676 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts @@ -67,14 +67,6 @@ export function useQuotaAndFallback({ const isDialogPending = useRef(false); const isValidationPending = useRef(false); - // Initial overage strategy from settings; runtime value read from config at call time. - const initialOverageStrategy = - (settings.merged.billing?.overageStrategy as - | 'ask' - | 'always' - | 'never' - | undefined) ?? 'ask'; - // Set up Flash fallback handler useEffect(() => { const fallbackHandler: FallbackModelHandler = async ( @@ -109,9 +101,7 @@ export function useQuotaAndFallback({ ? getResetTimeMessage(error.retryDelayMs) : undefined; - const overageStrategy = - config.getBillingSettings().overageStrategy ?? - initialOverageStrategy; + const overageStrategy = config.getBillingSettings().overageStrategy; const creditsResult = await handleCreditsFlow({ config, @@ -209,7 +199,6 @@ export function useQuotaAndFallback({ userTier, paidTier, settings, - initialOverageStrategy, setModelSwitchedFromQuotaError, onShowAuthSelection, errorVerbosity, diff --git a/packages/core/src/billing/billing.test.ts b/packages/core/src/billing/billing.test.ts index e594061ad6..e38767c418 100644 --- a/packages/core/src/billing/billing.test.ts +++ b/packages/core/src/billing/billing.test.ts @@ -229,14 +229,14 @@ describe('billing', () => { expect(isOverageEligibleModel('gemini-3.1-pro-preview')).toBe(true); }); - it('should return true for gemini-3.1-pro-preview-customtools', () => { + it('should return false for gemini-3.1-pro-preview-customtools', () => { expect(isOverageEligibleModel('gemini-3.1-pro-preview-customtools')).toBe( false, ); }); - it('should return false for gemini-3-flash-preview', () => { - expect(isOverageEligibleModel('gemini-3-flash-preview')).toBe(false); + it('should return true for gemini-3-flash-preview', () => { + expect(isOverageEligibleModel('gemini-3-flash-preview')).toBe(true); }); it('should return false for gemini-2.5-pro', () => { diff --git a/packages/core/src/billing/billing.ts b/packages/core/src/billing/billing.ts index 19afe72e16..64fd791cfd 100644 --- a/packages/core/src/billing/billing.ts +++ b/packages/core/src/billing/billing.ts @@ -12,6 +12,7 @@ import type { import { PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_3_1_MODEL, + PREVIEW_GEMINI_FLASH_MODEL, } from '../config/models.js'; /** @@ -32,6 +33,7 @@ export const G1_CREDIT_TYPE: CreditType = 'GOOGLE_ONE_AI'; export const OVERAGE_ELIGIBLE_MODELS = new Set([ PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_3_1_MODEL, + PREVIEW_GEMINI_FLASH_MODEL, ]); /** diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index 114fa60092..52b01504d3 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -48,6 +48,7 @@ import { shouldAutoUseCredits, } from '../billing/billing.js'; import { logBillingEvent, logInvalidChunk } from '../telemetry/loggers.js'; +import { coreEvents } from '../utils/events.js'; import { CreditsUsedEvent } from '../telemetry/billingEvents.js'; import { fromCountTokenResponse, @@ -100,6 +101,11 @@ export class CodeAssistServer implements ContentGenerator { const modelIsEligible = isOverageEligibleModel(req.model); const shouldEnableCredits = modelIsEligible && autoUse; + if (shouldEnableCredits && !this.config?.getCreditsNotificationShown()) { + this.config?.setCreditsNotificationShown(true); + coreEvents.emitFeedback('info', 'Using AI Credits for this request.'); + } + const enabledCreditTypes = shouldEnableCredits ? ([G1_CREDIT_TYPE] as string[]) : undefined; diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index da30b13377..31e081c350 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2426,6 +2426,65 @@ describe('Availability Service Integration', () => { config.resetTurn(); expect(spy).toHaveBeenCalled(); }); + + it('resetTurn does NOT reset billing state', () => { + const config = new Config({ + ...baseParams, + billing: { overageStrategy: 'ask' }, + }); + + // Simulate accepting credits mid-turn + config.setOverageStrategy('always'); + config.setCreditsNotificationShown(true); + + // resetTurn should leave billing state intact + config.resetTurn(); + expect(config.getBillingSettings().overageStrategy).toBe('always'); + expect(config.getCreditsNotificationShown()).toBe(true); + }); + + it('resetBillingTurnState resets overageStrategy to configured value', () => { + const config = new Config({ + ...baseParams, + billing: { overageStrategy: 'ask' }, + }); + + config.setOverageStrategy('always'); + expect(config.getBillingSettings().overageStrategy).toBe('always'); + + config.resetBillingTurnState('ask'); + expect(config.getBillingSettings().overageStrategy).toBe('ask'); + }); + + it('resetBillingTurnState preserves overageStrategy when configured as always', () => { + const config = new Config({ + ...baseParams, + billing: { overageStrategy: 'always' }, + }); + + config.resetBillingTurnState('always'); + expect(config.getBillingSettings().overageStrategy).toBe('always'); + }); + + it('resetBillingTurnState defaults to ask when no strategy provided', () => { + const config = new Config({ + ...baseParams, + billing: { overageStrategy: 'always' }, + }); + + config.resetBillingTurnState(); + expect(config.getBillingSettings().overageStrategy).toBe('ask'); + }); + + it('resetBillingTurnState resets creditsNotificationShown', () => { + const config = new Config(baseParams); + + config.setCreditsNotificationShown(true); + expect(config.getCreditsNotificationShown()).toBe(true); + + config.resetBillingTurnState(); + expect(config.getCreditsNotificationShown()).toBe(false); + }); }); describe('Hooks configuration', () => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index e4c0fef6eb..e3201aa521 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -685,6 +685,7 @@ export class Config implements McpContext { fallbackModelHandler?: FallbackModelHandler; validationHandler?: ValidationHandler; private quotaErrorOccurred: boolean = false; + private creditsNotificationShown: boolean = false; private modelQuotas: Map< string, { remaining: number; limit: number; resetTime?: string } @@ -1454,6 +1455,12 @@ export class Config implements McpContext { this.modelAvailabilityService.resetTurn(); } + /** Resets billing state (overageStrategy, creditsNotificationShown) once per user prompt. */ + resetBillingTurnState(overageStrategy?: OverageStrategy): void { + this.creditsNotificationShown = false; + this.billing.overageStrategy = overageStrategy ?? 'ask'; + } + getMaxSessionTurns(): number { return this.maxSessionTurns; } @@ -1466,6 +1473,14 @@ export class Config implements McpContext { return this.quotaErrorOccurred; } + setCreditsNotificationShown(value: boolean): void { + this.creditsNotificationShown = value; + } + + getCreditsNotificationShown(): boolean { + return this.creditsNotificationShown; + } + setQuota( remaining: number | undefined, limit: number | undefined,