Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/cli/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -875,6 +875,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,
Expand Down
8 changes: 0 additions & 8 deletions packages/cli/src/ui/hooks/creditsFlowHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,6 @@ async function handleOverageMenu(
isDialogPending,
setOverageMenuRequest,
setModelSwitchedFromQuotaError,
historyManager,
} = args;

logBillingEvent(
Expand Down Expand Up @@ -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':
Expand Down
1 change: 1 addition & 0 deletions packages/cli/src/ui/hooks/useGeminiStream.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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(() => ({
Expand Down
4 changes: 4 additions & 0 deletions packages/cli/src/ui/hooks/useGeminiStream.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1374,6 +1374,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;
Expand Down Expand Up @@ -1534,6 +1537,7 @@ export const useGeminiStream = (
setThought,
maybeAddSuppressedToolErrorNote,
maybeAddLowVerbosityFailureNote,
settings.merged.billing?.overageStrategy,
],
);

Expand Down
13 changes: 1 addition & 12 deletions packages/cli/src/ui/hooks/useQuotaAndFallback.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -209,7 +199,6 @@ export function useQuotaAndFallback({
userTier,
paidTier,
settings,
initialOverageStrategy,
setModelSwitchedFromQuotaError,
onShowAuthSelection,
errorVerbosity,
Expand Down
6 changes: 3 additions & 3 deletions packages/core/src/billing/billing.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -225,18 +225,18 @@
expect(isOverageEligibleModel('gemini-3-pro-preview')).toBe(true);
});

it('should return true for gemini-3.1-pro-preview', () => {

Check warning on line 228 in packages/core/src/billing/billing.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
expect(isOverageEligibleModel('gemini-3.1-pro-preview')).toBe(true);

Check warning on line 229 in packages/core/src/billing/billing.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
});

it('should return true for gemini-3.1-pro-preview-customtools', () => {
it('should return false for gemini-3.1-pro-preview-customtools', () => {

Check warning on line 232 in packages/core/src/billing/billing.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
expect(isOverageEligibleModel('gemini-3.1-pro-preview-customtools')).toBe(

Check warning on line 233 in packages/core/src/billing/billing.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
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', () => {
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/billing/billing.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import type {
import {
PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_3_1_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
} from '../config/models.js';

/**
Expand All @@ -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,
]);

/**
Expand Down
6 changes: 6 additions & 0 deletions packages/core/src/code_assist/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ import {
shouldAutoUseCredits,
} from '../billing/billing.js';
import { logBillingEvent } from '../telemetry/loggers.js';
import { coreEvents } from '../utils/events.js';
import { CreditsUsedEvent } from '../telemetry/billingEvents.js';
import type {
CaCountTokenResponse,
Expand Down Expand Up @@ -102,6 +103,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;
Expand Down
59 changes: 59 additions & 0 deletions packages/core/src/config/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2397,6 +2397,65 @@
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', () => {
Expand Down Expand Up @@ -2489,7 +2548,7 @@
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [
{
modelId: 'gemini-3.1-pro-preview',

Check warning on line 2551 in packages/core/src/config/config.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
remainingAmount: '100',
remainingFraction: 1.0,
},
Expand Down
15 changes: 15 additions & 0 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -671,6 +671,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 }
Expand Down Expand Up @@ -1438,6 +1439,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;
}
Expand All @@ -1450,6 +1457,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,
Expand Down
Loading