feat(billing): implement G1 AI credits overage flow with billing telemetry (#18590)

2026-03-20 19:11:23 -07:00 · 2026-02-27 10:15:06 -08:00
parent fdd844b405
commit b2d6844f9b
55 changed files with 3182 additions and 23 deletions
--- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
+++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
@@ -14,7 +14,8 @@ import {
  type Mock,
 } from 'vitest';
 import { act } from 'react';
-import { renderHook } from '../../test-utils/render.js';
+import { renderHook, mockSettings } from '../../test-utils/render.js';
+import { waitFor } from '../../test-utils/async.js';
 import {
  type Config,
  type FallbackModelHandler,
@@ -29,6 +30,12 @@ import {
  ModelNotFoundError,
  DEFAULT_GEMINI_MODEL,
  DEFAULT_GEMINI_FLASH_MODEL,
+  getG1CreditBalance,
+  shouldAutoUseCredits,
+  shouldShowOverageMenu,
+  shouldShowEmptyWalletMenu,
+  logBillingEvent,
+  G1_CREDIT_TYPE,
 } from '@google/gemini-cli-core';
 import { useQuotaAndFallback } from './useQuotaAndFallback.js';
 import type { UseHistoryManagerReturn } from './useHistoryManager.js';
@@ -37,6 +44,19 @@ import { MessageType } from '../types.js';
 // Use a type alias for SpyInstance as it's not directly exported
 type SpyInstance = ReturnType<typeof vi.spyOn>;

+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import('@google/gemini-cli-core')>();
+  return {
+    ...actual,
+    getG1CreditBalance: vi.fn(),
+    shouldAutoUseCredits: vi.fn(),
+    shouldShowOverageMenu: vi.fn(),
+    shouldShowEmptyWalletMenu: vi.fn(),
+    logBillingEvent: vi.fn(),
+  };
+});
+
 describe('useQuotaAndFallback', () => {
  let mockConfig: Config;
  let mockHistoryManager: UseHistoryManagerReturn;
@@ -74,10 +94,16 @@ describe('useQuotaAndFallback', () => {
    vi.spyOn(mockConfig, 'setModel');
    vi.spyOn(mockConfig, 'setActiveModel');
    vi.spyOn(mockConfig, 'activateFallbackMode');
+
+    // Mock billing utility functions
+    vi.mocked(getG1CreditBalance).mockReturnValue(0);
+    vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
+    vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
+    vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(false);
  });

  afterEach(() => {
-    vi.clearAllMocks();
+    vi.restoreAllMocks();
  });

  it('should register a fallback handler on initialization', () => {
@@ -88,6 +114,8 @@ describe('useQuotaAndFallback', () => {
        userTier: UserTierId.FREE,
        setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
        onShowAuthSelection: mockOnShowAuthSelection,
+        paidTier: null,
+        settings: mockSettings,
      }),
    );

@@ -109,6 +137,8 @@ describe('useQuotaAndFallback', () => {
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -140,6 +170,8 @@ describe('useQuotaAndFallback', () => {
            userTier: UserTierId.FREE,
            setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
            onShowAuthSelection: mockOnShowAuthSelection,
+            paidTier: null,
+            settings: mockSettings,
          }),
        );

@@ -193,6 +225,8 @@ describe('useQuotaAndFallback', () => {
            userTier: UserTierId.FREE,
            setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
            onShowAuthSelection: mockOnShowAuthSelection,
+            paidTier: null,
+            settings: mockSettings,
          }),
        );

@@ -232,6 +266,8 @@ describe('useQuotaAndFallback', () => {
            userTier: UserTierId.FREE,
            setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
            onShowAuthSelection: mockOnShowAuthSelection,
+            paidTier: null,
+            settings: mockSettings,
          }),
        );

@@ -264,6 +300,8 @@ describe('useQuotaAndFallback', () => {
            userTier: UserTierId.FREE,
            setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
            onShowAuthSelection: mockOnShowAuthSelection,
+            paidTier: null,
+            settings: mockSettings,
          }),
        );

@@ -330,6 +368,8 @@ describe('useQuotaAndFallback', () => {
              setModelSwitchedFromQuotaError:
                mockSetModelSwitchedFromQuotaError,
              onShowAuthSelection: mockOnShowAuthSelection,
+              paidTier: null,
+              settings: mockSettings,
            }),
          );

@@ -385,6 +425,8 @@ describe('useQuotaAndFallback', () => {
            userTier: UserTierId.FREE,
            setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
            onShowAuthSelection: mockOnShowAuthSelection,
+            paidTier: null,
+            settings: mockSettings,
          }),
        );

@@ -430,6 +472,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
            userTier: UserTierId.FREE,
            setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
            onShowAuthSelection: mockOnShowAuthSelection,
+            paidTier: null,
+            settings: mockSettings,
          }),
        );

@@ -464,6 +508,243 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
    });
  });

+  describe('G1 AI Credits Flow', () => {
+    const mockPaidTier = {
+      id: UserTierId.STANDARD,
+      userTier: UserTierId.STANDARD,
+      availableCredits: [
+        {
+          creditType: G1_CREDIT_TYPE,
+          creditAmount: '100',
+        },
+      ],
+    };
+
+    beforeEach(() => {
+      // Default to having credits
+      vi.mocked(getG1CreditBalance).mockReturnValue(100);
+    });
+
+    it('should fall through to ProQuotaDialog if credits are already active (strategy=always)', async () => {
+      // If shouldAutoUseCredits is true, credits were already active on the
+      // failed request — they didn't help. Fall through to ProQuotaDialog
+      // so the user can downgrade to Flash instead of retrying infinitely.
+      vi.mocked(shouldAutoUseCredits).mockReturnValue(true);
+
+      const { result } = renderHook(() =>
+        useQuotaAndFallback({
+          config: mockConfig,
+          historyManager: mockHistoryManager,
+          userTier: UserTierId.STANDARD,
+          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
+          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: mockPaidTier,
+          settings: mockSettings,
+        }),
+      );
+
+      const handler = setFallbackHandlerSpy.mock
+        .calls[0][0] as FallbackModelHandler;
+
+      const error = new TerminalQuotaError(
+        'pro quota',
+        mockGoogleApiError,
+        1000 * 60 * 5,
+      );
+
+      const intentPromise = handler(
+        PREVIEW_GEMINI_MODEL,
+        'gemini-flash',
+        error,
+      );
+
+      // Since credits didn't help, the ProQuotaDialog should be shown
+      await waitFor(() => {
+        expect(result.current.proQuotaRequest).not.toBeNull();
+      });
+
+      // Resolve it to verify the flow completes
+      act(() => {
+        result.current.handleProQuotaChoice('stop');
+      });
+
+      const intent = await intentPromise;
+      expect(intent).toBe('stop');
+    });
+
+    it('should show overage menu if balance > 0 and not auto-using', async () => {
+      vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
+      vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
+
+      const { result } = renderHook(() =>
+        useQuotaAndFallback({
+          config: mockConfig,
+          historyManager: mockHistoryManager,
+          userTier: UserTierId.STANDARD,
+          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
+          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: mockPaidTier,
+          settings: mockSettings,
+        }),
+      );
+
+      const handler = setFallbackHandlerSpy.mock
+        .calls[0][0] as FallbackModelHandler;
+
+      let promise: Promise<FallbackIntent | null>;
+      act(() => {
+        promise = handler(
+          PREVIEW_GEMINI_MODEL,
+          'gemini-flash',
+          new TerminalQuotaError('pro quota', mockGoogleApiError),
+        );
+      });
+
+      expect(result.current.overageMenuRequest).not.toBeNull();
+      expect(result.current.overageMenuRequest?.creditBalance).toBe(100);
+      expect(logBillingEvent).toHaveBeenCalled();
+
+      // Simulate choosing "Use Credits"
+      await act(async () => {
+        result.current.handleOverageMenuChoice('use_credits');
+        await promise!;
+      });
+
+      const intent = await promise!;
+      expect(intent).toBe('retry_with_credits');
+    });
+
+    it('should handle use_fallback from overage menu', async () => {
+      vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
+      vi.mocked(shouldShowOverageMenu).mockReturnValue(true);
+
+      const { result } = renderHook(() =>
+        useQuotaAndFallback({
+          config: mockConfig,
+          historyManager: mockHistoryManager,
+          userTier: UserTierId.STANDARD,
+          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
+          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: mockPaidTier,
+          settings: mockSettings,
+        }),
+      );
+
+      const handler = setFallbackHandlerSpy.mock
+        .calls[0][0] as FallbackModelHandler;
+
+      let promise: Promise<FallbackIntent | null>;
+      act(() => {
+        promise = handler(
+          PREVIEW_GEMINI_MODEL,
+          'gemini-flash',
+          new TerminalQuotaError('pro quota', mockGoogleApiError),
+        );
+      });
+
+      // Simulate choosing "Switch to fallback"
+      await act(async () => {
+        result.current.handleOverageMenuChoice('use_fallback');
+        await promise!;
+      });
+
+      const intent = await promise!;
+      expect(intent).toBe('retry_always');
+    });
+
+    it('should show empty wallet menu if balance is 0', async () => {
+      vi.mocked(getG1CreditBalance).mockReturnValue(0);
+      vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
+      vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
+      vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
+
+      const { result } = renderHook(() =>
+        useQuotaAndFallback({
+          config: mockConfig,
+          historyManager: mockHistoryManager,
+          userTier: UserTierId.STANDARD,
+          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
+          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: { ...mockPaidTier, availableCredits: [] },
+          settings: mockSettings,
+        }),
+      );
+
+      const handler = setFallbackHandlerSpy.mock
+        .calls[0][0] as FallbackModelHandler;
+
+      let promise: Promise<FallbackIntent | null>;
+      act(() => {
+        promise = handler(
+          PREVIEW_GEMINI_MODEL,
+          'gemini-flash',
+          new TerminalQuotaError('pro quota', mockGoogleApiError),
+        );
+      });
+
+      expect(result.current.emptyWalletRequest).not.toBeNull();
+      expect(logBillingEvent).toHaveBeenCalled();
+
+      // Simulate choosing "Stop"
+      await act(async () => {
+        result.current.handleEmptyWalletChoice('stop');
+        await promise!;
+      });
+
+      const intent = await promise!;
+      expect(intent).toBe('stop');
+    });
+
+    it('should add info message to history when get_credits is selected', async () => {
+      vi.mocked(getG1CreditBalance).mockReturnValue(0);
+      vi.mocked(shouldAutoUseCredits).mockReturnValue(false);
+      vi.mocked(shouldShowOverageMenu).mockReturnValue(false);
+      vi.mocked(shouldShowEmptyWalletMenu).mockReturnValue(true);
+
+      const { result } = renderHook(() =>
+        useQuotaAndFallback({
+          config: mockConfig,
+          historyManager: mockHistoryManager,
+          userTier: UserTierId.STANDARD,
+          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
+          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: { ...mockPaidTier, availableCredits: [] },
+          settings: mockSettings,
+        }),
+      );
+
+      const handler = setFallbackHandlerSpy.mock
+        .calls[0][0] as FallbackModelHandler;
+
+      let promise: Promise<FallbackIntent | null>;
+      act(() => {
+        promise = handler(
+          PREVIEW_GEMINI_MODEL,
+          'gemini-flash',
+          new TerminalQuotaError('pro quota', mockGoogleApiError),
+        );
+      });
+
+      expect(result.current.emptyWalletRequest).not.toBeNull();
+
+      // Simulate choosing "Get AI Credits"
+      await act(async () => {
+        result.current.handleEmptyWalletChoice('get_credits');
+        await promise!;
+      });
+
+      const intent = await promise!;
+      expect(intent).toBe('stop');
+      expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: MessageType.INFO,
+          text: expect.stringContaining('few minutes'),
+        }),
+        expect.any(Number),
+      );
+    });
+  });
+
  describe('handleProQuotaChoice', () => {
    it('should do nothing if there is no pending pro quota request', () => {
      const { result } = renderHook(() =>
@@ -473,6 +754,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -491,6 +774,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -522,6 +807,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -566,6 +853,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -602,6 +891,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -646,6 +937,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -661,6 +954,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -703,6 +998,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -745,6 +1042,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -775,6 +1074,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );

@@ -805,6 +1106,8 @@ Your admin might have disabled the access. Contact them to enable the Preview Re
          userTier: UserTierId.FREE,
          setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
          onShowAuthSelection: mockOnShowAuthSelection,
+          paidTier: null,
+          settings: mockSettings,
        }),
      );