Add usage limit remaining in /stats (#13843)

Author: sehoon38

packages/cli/src/ui/commands/statsCommand.ts

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { CodeAssistServer, getCodeAssistServer } from '@google/gemini-cli-core';
 import type { HistoryItemStats } from '../types.js';
 import { MessageType } from '../types.js';
 import { formatDuration } from '../utils/formatters.js';
@@ -13,7 +14,7 @@ import {
   CommandKind,
 } from './types.js';
 
-function defaultSessionView(context: CommandContext) {
+async function defaultSessionView(context: CommandContext) {
   const now = new Date();
   const { sessionStartTime } = context.session.stats;
   if (!sessionStartTime) {
@@ -33,6 +34,16 @@ function defaultSessionView(context: CommandContext) {
     duration: formatDuration(wallDuration),
   };
 
+  if (context.services.config) {
+    const server = getCodeAssistServer(context.services.config);
+    if (server instanceof CodeAssistServer && server.projectId) {
+      const quota = await server.retrieveUserQuota({
+        project: server.projectId,
+      });
+      statsItem.quotas = quota;
+    }
+  }
+
   context.ui.addItem(statsItem, Date.now());
 }
 
@@ -41,16 +52,16 @@ export const statsCommand: SlashCommand = {
   altNames: ['usage'],
   description: 'Check session stats. Usage: /stats [session|model|tools]',
   kind: CommandKind.BUILT_IN,
-  action: (context: CommandContext) => {
-    defaultSessionView(context);
+  action: async (context: CommandContext) => {
+    await defaultSessionView(context);
   },
   subCommands: [
     {
       name: 'session',
       description: 'Show session-specific usage statistics',
       kind: CommandKind.BUILT_IN,
-      action: (context: CommandContext) => {
-        defaultSessionView(context);
+      action: async (context: CommandContext) => {
+        await defaultSessionView(context);
       },
     },
     {

packages/cli/src/ui/components/HistoryItemDisplay.tsx

@@ -115,7 +115,10 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
         <Help commands={commands} />
       )}
       {itemForDisplay.type === 'stats' && (
-        <StatsDisplay duration={itemForDisplay.duration} />
+        <StatsDisplay
+          duration={itemForDisplay.duration}
+          quotas={itemForDisplay.quotas}
+        />
       )}
       {itemForDisplay.type === 'model_stats' && <ModelStatsDisplay />}
       {itemForDisplay.type === 'tool_stats' && <ToolStatsDisplay />}

packages/cli/src/ui/components/StatsDisplay.test.tsx

@@ -9,7 +9,10 @@ import { describe, it, expect, vi } from 'vitest';
 import { StatsDisplay } from './StatsDisplay.js';
 import * as SessionContext from '../contexts/SessionContext.js';
 import type { SessionMetrics } from '../contexts/SessionContext.js';
-import { ToolCallDecision } from '@google/gemini-cli-core';
+import {
+  ToolCallDecision,
+  type RetrieveUserQuotaResponse,
+} from '@google/gemini-cli-core';
 
 // Mock the context to provide controlled data for testing
 vi.mock('../contexts/SessionContext.js', async (importOriginal) => {
@@ -387,4 +390,65 @@ describe('<StatsDisplay />', () => {
       expect(output).toMatchSnapshot();
     });
   });
+
+  describe('Quota Display', () => {
+    it('renders quota information when quotas are provided', () => {
+      const now = new Date('2025-01-01T12:00:00Z');
+      vi.useFakeTimers();
+      vi.setSystemTime(now);
+
+      const metrics = createTestMetrics({
+        models: {
+          'gemini-2.5-pro': {
+            api: { totalRequests: 1, totalErrors: 0, totalLatencyMs: 100 },
+            tokens: {
+              prompt: 100,
+              candidates: 100,
+              total: 250,
+              cached: 50,
+              thoughts: 0,
+              tool: 0,
+            },
+          },
+        },
+      });
+
+      const resetTime = new Date(now.getTime() + 1000 * 60 * 90).toISOString(); // 1 hour 30 minutes from now
+
+      const quotas: RetrieveUserQuotaResponse = {
+        buckets: [
+          {
+            modelId: 'gemini-2.5-pro',
+            remainingFraction: 0.75,
+            resetTime,
+          },
+        ],
+      };
+
+      useSessionStatsMock.mockReturnValue({
+        stats: {
+          sessionId: 'test-session-id',
+          sessionStartTime: new Date(),
+          metrics,
+          lastPromptTokenCount: 0,
+          promptCount: 5,
+        },
+
+        getPromptCount: () => 5,
+        startNewPrompt: vi.fn(),
+      });
+
+      const { lastFrame } = render(
+        <StatsDisplay duration="1s" quotas={quotas} />,
+      );
+      const output = lastFrame();
+
+      expect(output).toContain('Usage limit remaining');
+      expect(output).toContain('75.0%');
+      expect(output).toContain('(Resets in 1h 30m)');
+      expect(output).toMatchSnapshot();
+
+      vi.useRealTimers();
+    });
+  });
 });

packages/cli/src/ui/components/StatsDisplay.tsx

@@ -19,6 +19,7 @@ import {
   USER_AGREEMENT_RATE_MEDIUM,
 } from '../utils/displayUtils.js';
 import { computeSessionStats } from '../utils/computeStats.js';
+import type { RetrieveUserQuotaResponse } from '@google/gemini-cli-core';
 
 // A more flexible and powerful StatRow component
 interface StatRowProps {
@@ -69,15 +70,41 @@ const Section: React.FC<SectionProps> = ({ title, children }) => (
   </Box>
 );
 
+const formatResetTime = (resetTime: string): string => {
+  const diff = new Date(resetTime).getTime() - Date.now();
+  if (diff <= 0) return '';
+
+  const totalMinutes = Math.ceil(diff / (1000 * 60));
+  const hours = Math.floor(totalMinutes / 60);
+  const minutes = totalMinutes % 60;
+
+  const fmt = (val: number, unit: 'hour' | 'minute') =>
+    new Intl.NumberFormat('en', {
+      style: 'unit',
+      unit,
+      unitDisplay: 'narrow',
+    }).format(val);
+
+  if (hours > 0 && minutes > 0) {
+    return `(Resets in ${fmt(hours, 'hour')} ${fmt(minutes, 'minute')})`;
+  } else if (hours > 0) {
+    return `(Resets in ${fmt(hours, 'hour')})`;
+  }
+
+  return `(Resets in ${fmt(minutes, 'minute')})`;
+};
+
 const ModelUsageTable: React.FC<{
   models: Record<string, ModelMetrics>;
   totalCachedTokens: number;
   cacheEfficiency: number;
-}> = ({ models, totalCachedTokens, cacheEfficiency }) => {
+  quotas?: RetrieveUserQuotaResponse;
+}> = ({ models, totalCachedTokens, cacheEfficiency, quotas }) => {
   const nameWidth = 25;
   const requestsWidth = 8;
   const inputTokensWidth = 15;
   const outputTokensWidth = 15;
+  const usageLimitWidth = quotas ? 30 : 0;
 
   return (
     <Box flexDirection="column" marginTop={1}>
@@ -103,6 +130,13 @@ const ModelUsageTable: React.FC<{
             Output Tokens
           </Text>
         </Box>
+        {quotas && (
+          <Box width={usageLimitWidth} justifyContent="flex-end">
+            <Text bold color={theme.text.primary}>
+              Usage limit remaining
+            </Text>
+          </Box>
+        )}
       </Box>
       {/* Divider */}
       <Box
@@ -112,44 +146,73 @@ const ModelUsageTable: React.FC<{
         borderLeft={false}
         borderRight={false}
         borderColor={theme.border.default}
-        width={nameWidth + requestsWidth + inputTokensWidth + outputTokensWidth}
+        width={
+          nameWidth +
+          requestsWidth +
+          inputTokensWidth +
+          outputTokensWidth +
+          usageLimitWidth
+        }
       ></Box>
 
       {/* Rows */}
-      {Object.entries(models).map(([name, modelMetrics]) => (
-        <Box key={name}>
-          <Box width={nameWidth}>
-            <Text color={theme.text.primary}>{name.replace('-001', '')}</Text>
-          </Box>
-          <Box width={requestsWidth} justifyContent="flex-end">
-            <Text color={theme.text.primary}>
-              {modelMetrics.api.totalRequests}
-            </Text>
-          </Box>
-          <Box width={inputTokensWidth} justifyContent="flex-end">
-            <Text color={theme.status.warning}>
-              {modelMetrics.tokens.prompt.toLocaleString()}
-            </Text>
-          </Box>
-          <Box width={outputTokensWidth} justifyContent="flex-end">
-            <Text color={theme.status.warning}>
-              {modelMetrics.tokens.candidates.toLocaleString()}
-            </Text>
+      {Object.entries(models).map(([name, modelMetrics]) => {
+        const modelName = name.replace('-001', '');
+        const bucket = quotas?.buckets?.find((b) => b.modelId === modelName);
+
+        return (
+          <Box key={name}>
+            <Box width={nameWidth}>
+              <Text color={theme.text.primary}>{modelName}</Text>
+            </Box>
+            <Box width={requestsWidth} justifyContent="flex-end">
+              <Text color={theme.text.primary}>
+                {modelMetrics.api.totalRequests}
+              </Text>
+            </Box>
+            <Box width={inputTokensWidth} justifyContent="flex-end">
+              <Text color={theme.status.warning}>
+                {modelMetrics.tokens.prompt.toLocaleString()}
+              </Text>
+            </Box>
+            <Box width={outputTokensWidth} justifyContent="flex-end">
+              <Text color={theme.status.warning}>
+                {modelMetrics.tokens.candidates.toLocaleString()}
+              </Text>
+            </Box>
+            <Box width={usageLimitWidth} justifyContent="flex-end">
+              {bucket &&
+                bucket.remainingFraction != null &&
+                bucket.resetTime && (
+                  <Text color={theme.text.secondary}>
+                    {(bucket.remainingFraction * 100).toFixed(1)}%{' '}
+                    {formatResetTime(bucket.resetTime)}
+                  </Text>
+                )}
+            </Box>
           </Box>
-        </Box>
-      ))}
+        );
+      })}
       {cacheEfficiency > 0 && (
         <Box flexDirection="column" marginTop={1}>
           <Text color={theme.text.primary}>
             <Text color={theme.status.success}>Savings Highlight:</Text>{' '}
             {totalCachedTokens.toLocaleString()} ({cacheEfficiency.toFixed(1)}
             %) of input tokens were served from the cache, reducing costs.
           </Text>
-          <Box height={1} />
+        </Box>
+      )}
+      {models && (
+        <>
+          <Box marginTop={1} marginBottom={2}>
+            <Text color={theme.text.primary}>
+              {`Usage limits span all sessions and reset daily.\n/auth to upgrade or switch to API key.`}
+            </Text>
+          </Box>
           <Text color={theme.text.secondary}>
             » Tip: For a full token breakdown, run `/stats model`.
           </Text>
-        </Box>
+        </>
       )}
     </Box>
   );
@@ -158,11 +221,13 @@ const ModelUsageTable: React.FC<{
 interface StatsDisplayProps {
   duration: string;
   title?: string;
+  quotas?: RetrieveUserQuotaResponse;
 }
 
 export const StatsDisplay: React.FC<StatsDisplayProps> = ({
   duration,
   title,
+  quotas,
 }) => {
   const { stats } = useSessionStats();
   const { metrics } = stats;
@@ -276,6 +341,7 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
           models={models}
           totalCachedTokens={computed.totalCachedTokens}
           cacheEfficiency={computed.cacheEfficiency}
+          quotas={quotas}
         />
       )}
     </Box>

packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap

@@ -24,6 +24,10 @@ exports[`<SessionSummaryDisplay /> > renders the summary display with a title 1`
 │                                                                                                  │
 │  Savings Highlight: 500 (50.0%) of input tokens were served from the cache, reducing costs.      │
 │                                                                                                  │
+│  Usage limits span all sessions and reset daily.                                                 │
+│  /auth to upgrade or switch to API key.                                                          │
+│                                                                                                  │
+│                                                                                                  │
 │  » Tip: For a full token breakdown, run \`/stats model\`.                                          │
 │                                                                                                  │
 ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"