fix(patch): cherry-pick 9a74271 to release/v0.33.0-preview.5-pr-21236 [CONFLICTS] (#21714)

gemini-cli-robot · gsquared94 · web-flow · commit 8df4210cf248 · 2026-03-09T17:13:48.000Z
Co-authored-by: Gaurav &lt;39389231+gsquared94@users.noreply.github.com&gt;
diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts
@@ -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,
diff --git 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
@@ -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
@@ -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;
@@ -1534,6 +1537,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
@@ -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
@@ -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
@@ -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
@@ -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,
@@ -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;
diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts
@@ -2397,6 +2397,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
@@ -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 }
@@ -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;
   }
@@ -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,
