diff --git a/packages/cli/src/ui/commands/statsCommand.ts b/packages/cli/src/ui/commands/statsCommand.ts index 224b05f40f..fbc8e9b6c1 100644 --- a/packages/cli/src/ui/commands/statsCommand.ts +++ b/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); }, }, { diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.tsx index a5f3668345..36920c0592 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.tsx @@ -115,7 +115,10 @@ export const HistoryItemDisplay: React.FC = ({ )} {itemForDisplay.type === 'stats' && ( - + )} {itemForDisplay.type === 'model_stats' && } {itemForDisplay.type === 'tool_stats' && } diff --git a/packages/cli/src/ui/components/StatsDisplay.test.tsx b/packages/cli/src/ui/components/StatsDisplay.test.tsx index 927b9c7615..de6a33caf9 100644 --- a/packages/cli/src/ui/components/StatsDisplay.test.tsx +++ b/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('', () => { 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( + , + ); + 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(); + }); + }); }); diff --git a/packages/cli/src/ui/components/StatsDisplay.tsx b/packages/cli/src/ui/components/StatsDisplay.tsx index 964f98dfcb..7290640ba6 100644 --- a/packages/cli/src/ui/components/StatsDisplay.tsx +++ b/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 = ({ title, children }) => ( ); +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; 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 ( @@ -103,6 +130,13 @@ const ModelUsageTable: React.FC<{ Output Tokens + {quotas && ( + + + Usage limit remaining + + + )} {/* Divider */} {/* Rows */} - {Object.entries(models).map(([name, modelMetrics]) => ( - - - {name.replace('-001', '')} + {Object.entries(models).map(([name, modelMetrics]) => { + const modelName = name.replace('-001', ''); + const bucket = quotas?.buckets?.find((b) => b.modelId === modelName); + + return ( + + + {modelName} + + + + {modelMetrics.api.totalRequests} + + + + + {modelMetrics.tokens.prompt.toLocaleString()} + + + + + {modelMetrics.tokens.candidates.toLocaleString()} + + + + {bucket && + bucket.remainingFraction != null && + bucket.resetTime && ( + + {(bucket.remainingFraction * 100).toFixed(1)}%{' '} + {formatResetTime(bucket.resetTime)} + + )} + - - - {modelMetrics.api.totalRequests} - - - - - {modelMetrics.tokens.prompt.toLocaleString()} - - - - - {modelMetrics.tokens.candidates.toLocaleString()} - - - - ))} + ); + })} {cacheEfficiency > 0 && ( @@ -145,11 +200,19 @@ const ModelUsageTable: React.FC<{ {totalCachedTokens.toLocaleString()} ({cacheEfficiency.toFixed(1)} %) of input tokens were served from the cache, reducing costs. - + + )} + {models && ( + <> + + + {`Usage limits span all sessions and reset daily.\n/auth to upgrade or switch to API key.`} + + » Tip: For a full token breakdown, run `/stats model`. - + )} ); @@ -158,11 +221,13 @@ const ModelUsageTable: React.FC<{ interface StatsDisplayProps { duration: string; title?: string; + quotas?: RetrieveUserQuotaResponse; } export const StatsDisplay: React.FC = ({ duration, title, + quotas, }) => { const { stats } = useSessionStats(); const { metrics } = stats; @@ -276,6 +341,7 @@ export const StatsDisplay: React.FC = ({ models={models} totalCachedTokens={computed.totalCachedTokens} cacheEfficiency={computed.cacheEfficiency} + quotas={quotas} /> )} diff --git a/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap index 2912ac155d..c71c770df4 100644 --- a/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap +++ b/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap @@ -24,6 +24,10 @@ exports[` > 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\`. │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" diff --git a/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap index 8106d1f5dc..ded84b667f 100644 --- a/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap +++ b/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap @@ -122,6 +122,12 @@ exports[` > Conditional Rendering Tests > hides Efficiency secti │ ─────────────────────────────────────────────────────────────── │ │ gemini-2.5-pro 1 100 100 │ │ │ +│ 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\`. │ +│ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; @@ -145,6 +151,38 @@ exports[` > Conditional Rendering Tests > hides User Agreement w ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; +exports[` > Quota Display > renders quota information when quotas are provided 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ │ +│ Session Stats │ +│ │ +│ Interaction Summary │ +│ Session ID: test-session-id │ +│ Tool Calls: 0 ( ✓ 0 x 0 ) │ +│ Success Rate: 0.0% │ +│ │ +│ Performance │ +│ Wall Time: 1s │ +│ Agent Active: 100ms │ +│ » API Time: 100ms (100.0%) │ +│ » Tool Time: 0s (0.0%) │ +│ │ +│ │ +│ Model Usage Reqs Input Tokens Output Tokens Usage limit remaining │ +│ ───────────────────────────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-pro 1 100 100 75.0% (Resets in 1h 30m) │ +│ │ +│ Savings Highlight: 50 (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\`. │ +│ │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; + exports[` > Title Rendering > renders the custom title when a title prop is provided 1`] = ` "╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ │ │ @@ -209,6 +247,10 @@ exports[` > renders a table with two models correctly 1`] = ` │ │ │ Savings Highlight: 10,500 (40.4%) 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\`. │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" @@ -238,6 +280,10 @@ exports[` > renders all sections when all data is present 1`] = │ │ │ Savings Highlight: 50 (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\`. │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 491a1eede1..d8ec6f62a1 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -12,6 +12,7 @@ import type { ToolCallConfirmationDetails, ToolConfirmationOutcome, ToolResultDisplay, + RetrieveUserQuotaResponse, } from '@google/gemini-cli-core'; import type { PartListUnion } from '@google/genai'; import { type ReactNode } from 'react'; @@ -142,6 +143,7 @@ export type HistoryItemHelp = HistoryItemBase & { export type HistoryItemStats = HistoryItemBase & { type: 'stats'; duration: string; + quotas?: RetrieveUserQuotaResponse; }; export type HistoryItemModelStats = HistoryItemBase & { diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts index 099c05ad5e..281b46549d 100644 --- a/packages/core/src/code_assist/server.test.ts +++ b/packages/core/src/code_assist/server.test.ts @@ -380,4 +380,38 @@ describe('CodeAssistServer', () => { }); expect(response).toEqual(mockResponse); }); + + it('should call the retrieveUserQuota endpoint', async () => { + const client = new OAuth2Client(); + const server = new CodeAssistServer( + client, + 'test-project', + {}, + 'test-session', + UserTierId.FREE, + ); + const mockResponse = { + buckets: [ + { + modelId: 'gemini-2.5-pro', + tokenType: 'REQUESTS', + remainingFraction: 0.75, + resetTime: '2025-10-22T16:01:15Z', + }, + ], + }; + const requestPostSpy = vi + .spyOn(server, 'requestPost') + .mockResolvedValue(mockResponse); + + const req = { + project: 'projects/my-cloudcode-project', + userAgent: 'CloudCodePlugin/1.0 (gaghosh)', + }; + + const response = await server.retrieveUserQuota(req); + + expect(requestPostSpy).toHaveBeenCalledWith('retrieveUserQuota', req); + expect(response).toEqual(mockResponse); + }); }); diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index 8670fda289..9cb2a8ce0f 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -14,6 +14,8 @@ import type { OnboardUserRequest, SetCodeAssistGlobalUserSettingRequest, ClientMetadata, + RetrieveUserQuotaRequest, + RetrieveUserQuotaResponse, } from './types.js'; import type { ListExperimentsRequest, @@ -171,6 +173,15 @@ export class CodeAssistServer implements ContentGenerator { ); } + async retrieveUserQuota( + req: RetrieveUserQuotaRequest, + ): Promise { + return await this.requestPost( + 'retrieveUserQuota', + req, + ); + } + async requestPost( method: string, req: object, diff --git a/packages/core/src/code_assist/types.ts b/packages/core/src/code_assist/types.ts index e59dd3ac40..36e2f3f2fb 100644 --- a/packages/core/src/code_assist/types.ts +++ b/packages/core/src/code_assist/types.ts @@ -200,3 +200,20 @@ export interface GoogleRpcResponse { interface GoogleRpcErrorInfo { reason?: string; } + +export interface RetrieveUserQuotaRequest { + project: string; + userAgent?: string; +} + +export interface BucketInfo { + remainingAmount?: string; + remainingFraction?: number; + resetTime?: string; + tokenType?: string; + modelId?: string; +} + +export interface RetrieveUserQuotaResponse { + buckets?: BucketInfo[]; +}