MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-10 22:21:22 -07:00
Code Issues Packages Projects Releases Wiki Activity

Add usage limit remaining in /stats (#13843)

Browse Source
This commit is contained in:
Sehoon Shon
2025-11-26 20:21:33 -05:00
committed by GitHub
parent 5949d56370
commit 69188c8538
10 changed files with 290 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
packages/cli/src/ui/commands/statsCommand.ts
View File

@@ -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);
},
},
{

5
packages/cli/src/ui/components/HistoryItemDisplay.tsx
View File

@@ -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 />}

66
packages/cli/src/ui/components/StatsDisplay.test.tsx
View File

@@ -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();
});
});
});

116
packages/cli/src/ui/components/StatsDisplay.tsx
View File

@@ -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,32 +146,53 @@ 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>
{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 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>
))}
);
})}
{cacheEfficiency > 0 && (
<Box flexDirection="column" marginTop={1}>
<Text color={theme.text.primary}>
@@ -145,11 +200,19 @@ const ModelUsageTable: React.FC<{
{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>

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

@@ -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\`. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"

46
packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap
View File

@@ -122,6 +122,12 @@ exports[`<StatsDisplay /> > 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[`<StatsDisplay /> > Conditional Rendering Tests > hides User Agreement w
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`;
exports[`<StatsDisplay /> > 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[`<StatsDisplay /> > Title Rendering > renders the custom title when a title prop is provided 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
@@ -209,6 +247,10 @@ exports[`<StatsDisplay /> > 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[`<StatsDisplay /> > 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\`. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"

2
packages/cli/src/ui/types.ts
View File

@@ -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 & {

34
packages/core/src/code_assist/server.test.ts
View File

@@ -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);
});
});

11
packages/core/src/code_assist/server.ts
View File

@@ -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<RetrieveUserQuotaResponse> {
return await this.requestPost<RetrieveUserQuotaResponse>(
'retrieveUserQuota',
req,
);
}
async requestPost<T>(
method: string,
req: object,

17
packages/core/src/code_assist/types.ts
View File

@@ -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[];
}