mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-06-13 04:48:09 -07:00
fix(billing): fix overage strategy lifecycle and settings integration (#21236)
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
@@ -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', () => {
|
||||||
|
|||||||
@@ -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,
|
||||||
]);
|
]);
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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', () => {
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user