fix(billing): fix overage strategy lifecycle and settings integration (#21236)

This commit is contained in:
Gaurav
2026-03-06 19:14:44 -08:00
committed by GitHub
parent 33be30ab04
commit 9a7427197b
10 changed files with 92 additions and 23 deletions
+1
View File
@@ -805,6 +805,7 @@ export async function loadCliConfig(
fakeResponses: argv.fakeResponses, fakeResponses: argv.fakeResponses,
recordResponses: argv.recordResponses, recordResponses: argv.recordResponses,
retryFetchErrors: settings.general?.retryFetchErrors, retryFetchErrors: settings.general?.retryFetchErrors,
billing: settings.billing,
maxAttempts: settings.general?.maxAttempts, maxAttempts: settings.general?.maxAttempts,
ptyInfo: ptyInfo?.name, ptyInfo: ptyInfo?.name,
disableLLMCorrection: settings.tools?.disableLLMCorrection, disableLLMCorrection: settings.tools?.disableLLMCorrection,
@@ -110,7 +110,6 @@ async function handleOverageMenu(
isDialogPending, isDialogPending,
setOverageMenuRequest, setOverageMenuRequest,
setModelSwitchedFromQuotaError, setModelSwitchedFromQuotaError,
historyManager,
} = args; } = args;
logBillingEvent( logBillingEvent(
@@ -155,13 +154,6 @@ async function handleOverageMenu(
setModelSwitchedFromQuotaError(false); setModelSwitchedFromQuotaError(false);
config.setQuotaErrorOccurred(false); config.setQuotaErrorOccurred(false);
config.setOverageStrategy('always'); config.setOverageStrategy('always');
historyManager.addItem(
{
type: MessageType.INFO,
text: `Using AI Credits for this request.`,
},
Date.now(),
);
return 'retry_with_credits'; return 'retry_with_credits';
case 'use_fallback': case 'use_fallback':
@@ -271,6 +271,7 @@ describe('useGeminiStream', () => {
addHistory: vi.fn(), addHistory: vi.fn(),
getSessionId: vi.fn(() => 'test-session-id'), getSessionId: vi.fn(() => 'test-session-id'),
setQuotaErrorOccurred: vi.fn(), setQuotaErrorOccurred: vi.fn(),
resetBillingTurnState: vi.fn(),
getQuotaErrorOccurred: vi.fn(() => false), getQuotaErrorOccurred: vi.fn(() => false),
getModel: vi.fn(() => 'gemini-2.5-pro'), getModel: vi.fn(() => 'gemini-2.5-pro'),
getContentGeneratorConfig: vi.fn(() => ({ getContentGeneratorConfig: vi.fn(() => ({
@@ -1376,6 +1376,9 @@ export const useGeminiStream = (
if (!options?.isContinuation) { if (!options?.isContinuation) {
setModelSwitchedFromQuotaError(false); setModelSwitchedFromQuotaError(false);
config.setQuotaErrorOccurred(false); config.setQuotaErrorOccurred(false);
config.resetBillingTurnState(
settings.merged.billing?.overageStrategy,
);
suppressedToolErrorCountRef.current = 0; suppressedToolErrorCountRef.current = 0;
suppressedToolErrorNoteShownRef.current = false; suppressedToolErrorNoteShownRef.current = false;
lowVerbosityFailureNoteShownRef.current = false; lowVerbosityFailureNoteShownRef.current = false;
@@ -1536,6 +1539,7 @@ export const useGeminiStream = (
setThought, setThought,
maybeAddSuppressedToolErrorNote, maybeAddSuppressedToolErrorNote,
maybeAddLowVerbosityFailureNote, maybeAddLowVerbosityFailureNote,
settings.merged.billing?.overageStrategy,
], ],
); );
@@ -67,14 +67,6 @@ export function useQuotaAndFallback({
const isDialogPending = useRef(false); const isDialogPending = useRef(false);
const isValidationPending = 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 // Set up Flash fallback handler
useEffect(() => { useEffect(() => {
const fallbackHandler: FallbackModelHandler = async ( const fallbackHandler: FallbackModelHandler = async (
@@ -109,9 +101,7 @@ export function useQuotaAndFallback({
? getResetTimeMessage(error.retryDelayMs) ? getResetTimeMessage(error.retryDelayMs)
: undefined; : undefined;
const overageStrategy = const overageStrategy = config.getBillingSettings().overageStrategy;
config.getBillingSettings().overageStrategy ??
initialOverageStrategy;
const creditsResult = await handleCreditsFlow({ const creditsResult = await handleCreditsFlow({
config, config,
@@ -209,7 +199,6 @@ export function useQuotaAndFallback({
userTier, userTier,
paidTier, paidTier,
settings, settings,
initialOverageStrategy,
setModelSwitchedFromQuotaError, setModelSwitchedFromQuotaError,
onShowAuthSelection, onShowAuthSelection,
errorVerbosity, errorVerbosity,
+3 -3
View File
@@ -229,14 +229,14 @@ describe('billing', () => {
expect(isOverageEligibleModel('gemini-3.1-pro-preview')).toBe(true); 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( expect(isOverageEligibleModel('gemini-3.1-pro-preview-customtools')).toBe(
false, false,
); );
}); });
it('should return false for gemini-3-flash-preview', () => { it('should return true for gemini-3-flash-preview', () => {
expect(isOverageEligibleModel('gemini-3-flash-preview')).toBe(false); expect(isOverageEligibleModel('gemini-3-flash-preview')).toBe(true);
}); });
it('should return false for gemini-2.5-pro', () => { it('should return false for gemini-2.5-pro', () => {
+2
View File
@@ -12,6 +12,7 @@ import type {
import { import {
PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_3_1_MODEL, PREVIEW_GEMINI_3_1_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
} from '../config/models.js'; } from '../config/models.js';
/** /**
@@ -32,6 +33,7 @@ export const G1_CREDIT_TYPE: CreditType = 'GOOGLE_ONE_AI';
export const OVERAGE_ELIGIBLE_MODELS = new Set([ export const OVERAGE_ELIGIBLE_MODELS = new Set([
PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_3_1_MODEL, PREVIEW_GEMINI_3_1_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
]); ]);
/** /**
+6
View File
@@ -48,6 +48,7 @@ import {
shouldAutoUseCredits, shouldAutoUseCredits,
} from '../billing/billing.js'; } from '../billing/billing.js';
import { logBillingEvent, logInvalidChunk } from '../telemetry/loggers.js'; import { logBillingEvent, logInvalidChunk } from '../telemetry/loggers.js';
import { coreEvents } from '../utils/events.js';
import { CreditsUsedEvent } from '../telemetry/billingEvents.js'; import { CreditsUsedEvent } from '../telemetry/billingEvents.js';
import { import {
fromCountTokenResponse, fromCountTokenResponse,
@@ -100,6 +101,11 @@ export class CodeAssistServer implements ContentGenerator {
const modelIsEligible = isOverageEligibleModel(req.model); const modelIsEligible = isOverageEligibleModel(req.model);
const shouldEnableCredits = modelIsEligible && autoUse; 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 const enabledCreditTypes = shouldEnableCredits
? ([G1_CREDIT_TYPE] as string[]) ? ([G1_CREDIT_TYPE] as string[])
: undefined; : undefined;
+59
View File
@@ -2426,6 +2426,65 @@ describe('Availability Service Integration', () => {
config.resetTurn(); config.resetTurn();
expect(spy).toHaveBeenCalled(); 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', () => { describe('Hooks configuration', () => {
+15
View File
@@ -685,6 +685,7 @@ export class Config implements McpContext {
fallbackModelHandler?: FallbackModelHandler; fallbackModelHandler?: FallbackModelHandler;
validationHandler?: ValidationHandler; validationHandler?: ValidationHandler;
private quotaErrorOccurred: boolean = false; private quotaErrorOccurred: boolean = false;
private creditsNotificationShown: boolean = false;
private modelQuotas: Map< private modelQuotas: Map<
string, string,
{ remaining: number; limit: number; resetTime?: string } { remaining: number; limit: number; resetTime?: string }
@@ -1454,6 +1455,12 @@ export class Config implements McpContext {
this.modelAvailabilityService.resetTurn(); 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 { getMaxSessionTurns(): number {
return this.maxSessionTurns; return this.maxSessionTurns;
} }
@@ -1466,6 +1473,14 @@ export class Config implements McpContext {
return this.quotaErrorOccurred; return this.quotaErrorOccurred;
} }
setCreditsNotificationShown(value: boolean): void {
this.creditsNotificationShown = value;
}
getCreditsNotificationShown(): boolean {
return this.creditsNotificationShown;
}
setQuota( setQuota(
remaining: number | undefined, remaining: number | undefined,
limit: number | undefined, limit: number | undefined,