From 5cc9644a16460b7ad8f6464fe4da1b89eca78adc Mon Sep 17 00:00:00 2001
From: Gaurav Ghosh <gaghosh@google.com>
Date: Wed, 25 Feb 2026 22:26:18 -0800
Subject: [PATCH] update: add Pro 3.1

---
 packages/cli/src/config/settingsSchema.ts         |  2 +-
 packages/cli/src/ui/commands/statsCommand.test.ts |  2 +-
 packages/cli/src/ui/commands/statsCommand.ts      |  2 +-
 packages/cli/src/ui/components/StatsDisplay.tsx   |  3 +--
 packages/cli/src/ui/hooks/creditsFlowHandler.ts   |  4 ++--
 packages/cli/src/ui/types.ts                      |  3 +--
 packages/core/src/billing/billing.test.ts         |  2 +-
 packages/core/src/billing/billing.ts              | 10 ++++++++--
 8 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts
index e98605bc9a..9238ceea85 100644
--- a/packages/cli/src/config/settingsSchema.ts
+++ b/packages/cli/src/config/settingsSchema.ts
@@ -835,7 +835,7 @@ const SETTINGS_SCHEMA = {
     requiresRestart: false,
     default: {},
     description: 'Billing and AI credits settings.',
-    showInDialog: true,
+    showInDialog: false,
     properties: {
       overageStrategy: {
         type: 'enum',
diff --git a/packages/cli/src/ui/commands/statsCommand.test.ts b/packages/cli/src/ui/commands/statsCommand.test.ts
index 04b09012f8..2f36c333b9 100644
--- a/packages/cli/src/ui/commands/statsCommand.test.ts
+++ b/packages/cli/src/ui/commands/statsCommand.test.ts
@@ -62,7 +62,7 @@ describe('statsCommand', () => {
       tier: undefined,
       userEmail: 'mock@example.com',
       currentModel: undefined,
-      creditBalance: null,
+      creditBalance: undefined,
     });
   });
 
diff --git a/packages/cli/src/ui/commands/statsCommand.ts b/packages/cli/src/ui/commands/statsCommand.ts
index 42b4bee8d8..1ded006618 100644
--- a/packages/cli/src/ui/commands/statsCommand.ts
+++ b/packages/cli/src/ui/commands/statsCommand.ts
@@ -31,7 +31,7 @@ function getUserIdentity(context: CommandContext) {
 
   const tier = context.services.config?.getUserTierName();
   const paidTier = context.services.config?.getUserPaidTier();
-  const creditBalance = getG1CreditBalance(paidTier);
+  const creditBalance = getG1CreditBalance(paidTier) ?? undefined;
 
   return { selectedAuthType, userEmail, tier, creditBalance };
 }
diff --git a/packages/cli/src/ui/components/StatsDisplay.tsx b/packages/cli/src/ui/components/StatsDisplay.tsx
index 9f734cbbea..65169f6d74 100644
--- a/packages/cli/src/ui/components/StatsDisplay.tsx
+++ b/packages/cli/src/ui/components/StatsDisplay.tsx
@@ -395,8 +395,7 @@ interface StatsDisplayProps {
   tier?: string;
   currentModel?: string;
   quotaStats?: QuotaStats;
-  /** G1 AI Credits balance, null if not eligible */
-  creditBalance?: number | null;
+  creditBalance?: number;
 }
 
 export const StatsDisplay: React.FC<StatsDisplayProps> = ({
diff --git a/packages/cli/src/ui/hooks/creditsFlowHandler.ts b/packages/cli/src/ui/hooks/creditsFlowHandler.ts
index b5bbfeeb98..497d4904e6 100644
--- a/packages/cli/src/ui/hooks/creditsFlowHandler.ts
+++ b/packages/cli/src/ui/hooks/creditsFlowHandler.ts
@@ -74,8 +74,8 @@ export async function handleCreditsFlow(
   const { overageStrategy } = args;
 
   // If credits are already auto-enabled (strategy='always'), the request
-  // that just failed already included credits — they didn't help.
-  // Fall through to ProQuotaDialog which offers the Flash downgrade.
+  // that just failed already included enabledCreditTypes — credits didn't
+  // help. Fall through to ProQuotaDialog which offers the Flash downgrade.
   if (shouldAutoUseCredits(overageStrategy, creditBalance)) {
     return null;
   }
diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts
index 21c432fac7..8032b35d66 100644
--- a/packages/cli/src/ui/types.ts
+++ b/packages/cli/src/ui/types.ts
@@ -201,8 +201,7 @@ export type HistoryItemStats = HistoryItemQuotaBase & {
   type: 'stats';
   duration: string;
   quotas?: RetrieveUserQuotaResponse;
-  /** G1 AI Credits balance, null if not eligible */
-  creditBalance?: number | null;
+  creditBalance?: number;
 };
 
 export type HistoryItemModelStats = HistoryItemQuotaBase & {
diff --git a/packages/core/src/billing/billing.test.ts b/packages/core/src/billing/billing.test.ts
index b668094add..e594061ad6 100644
--- a/packages/core/src/billing/billing.test.ts
+++ b/packages/core/src/billing/billing.test.ts
@@ -226,7 +226,7 @@ describe('billing', () => {
     });
 
     it('should return true for gemini-3.1-pro-preview', () => {
-      expect(isOverageEligibleModel('gemini-3.1-pro-preview')).toBe(false);
+      expect(isOverageEligibleModel('gemini-3.1-pro-preview')).toBe(true);
     });
 
     it('should return true for gemini-3.1-pro-preview-customtools', () => {
diff --git a/packages/core/src/billing/billing.ts b/packages/core/src/billing/billing.ts
index 7f21fb1692..19afe72e16 100644
--- a/packages/core/src/billing/billing.ts
+++ b/packages/core/src/billing/billing.ts
@@ -9,7 +9,10 @@ import type {
   CreditType,
   GeminiUserTier,
 } from '../code_assist/types.js';
-import { PREVIEW_GEMINI_MODEL } from '../config/models.js';
+import {
+  PREVIEW_GEMINI_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
+} from '../config/models.js';
 
 /**
  * Strategy for handling quota exhaustion when AI credits are available.
@@ -26,7 +29,10 @@ export const G1_CREDIT_TYPE: CreditType = 'GOOGLE_ONE_AI';
  * The set of models that support AI credits overage billing.
  * Only these models are eligible for the credits-based retry flow.
  */
-export const OVERAGE_ELIGIBLE_MODELS = new Set([PREVIEW_GEMINI_MODEL]);
+export const OVERAGE_ELIGIBLE_MODELS = new Set([
+  PREVIEW_GEMINI_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
+]);
 
 /**
  * Checks if a model is eligible for AI credits overage billing.