fix(ui): hide model quota in /stats and refactor quota display (#24206)

This commit is contained in:
Dan Zaharia
2026-04-02 16:49:14 -04:00
committed by GitHub
parent e5adeaca80
commit 29caa52bb7
15 changed files with 986 additions and 748 deletions
@@ -9,10 +9,7 @@ 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,
type RetrieveUserQuotaResponse,
} from '@google/gemini-cli-core';
import { ToolCallDecision } from '@google/gemini-cli-core';
// Mock the context to provide controlled data for testing
vi.mock('../contexts/SessionContext.js', async (importOriginal) => {
@@ -124,10 +121,13 @@ describe('<StatsDisplay />', () => {
const { lastFrame } = await renderWithMockedStats(metrics);
const output = lastFrame();
expect(output).toContain('gemini-2.5-pro');
expect(output).toContain('gemini-2.5-flash');
expect(output).toContain('15,000');
expect(output).toContain('10,000');
expect(output).toContain('Performance');
expect(output).toContain('Interaction Summary');
expect(output).toContain('Model Usage');
expect(output).toContain('Reqs');
expect(output).toContain('Input Tokens');
expect(output).toContain('Cache Reads');
expect(output).toContain('Output Tokens');
expect(output).toMatchSnapshot();
});
@@ -182,7 +182,7 @@ describe('<StatsDisplay />', () => {
expect(output).toContain('Performance');
expect(output).toContain('Interaction Summary');
expect(output).toContain('User Agreement');
expect(output).toContain('gemini-2.5-pro');
expect(output).toContain('Model Usage');
expect(output).toMatchSnapshot();
});
@@ -406,174 +406,6 @@ describe('<StatsDisplay />', () => {
});
});
describe('Quota Display', () => {
it('renders quota information when quotas are provided', async () => {
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: {
input: 50,
prompt: 100,
candidates: 100,
total: 250,
cached: 50,
thoughts: 0,
tool: 0,
},
roles: {},
},
},
});
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',
remainingAmount: '75',
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 } = await renderWithProviders(
<StatsDisplay duration="1s" quotas={quotas} />,
{ width: 100 },
);
const output = lastFrame();
expect(output).toContain('Model usage');
expect(output).toContain('25%');
expect(output).toContain('Usage resets');
expect(output).toMatchSnapshot();
vi.useRealTimers();
});
it('renders pooled quota information for auto mode', async () => {
const now = new Date('2025-01-01T12:00:00Z');
vi.useFakeTimers();
vi.setSystemTime(now);
const metrics = createTestMetrics();
const quotas: RetrieveUserQuotaResponse = {
buckets: [
{
modelId: 'gemini-2.5-pro',
remainingAmount: '10',
remainingFraction: 0.1, // limit = 100
},
{
modelId: 'gemini-2.5-flash',
remainingAmount: '700',
remainingFraction: 0.7, // limit = 1000
},
],
};
useSessionStatsMock.mockReturnValue({
stats: {
sessionId: 'test-session-id',
sessionStartTime: new Date(),
metrics,
lastPromptTokenCount: 0,
promptCount: 5,
},
getPromptCount: () => 5,
startNewPrompt: vi.fn(),
});
const { lastFrame } = await renderWithProviders(
<StatsDisplay
duration="1s"
quotas={quotas}
currentModel="auto"
quotaStats={{
remaining: 710,
limit: 1100,
}}
/>,
{ width: 100 },
);
const output = lastFrame();
// (1 - 710/1100) * 100 = 35.5%
expect(output).toContain('35%');
expect(output).toContain('Usage limit: 1,100');
expect(output).toMatchSnapshot();
vi.useRealTimers();
});
it('renders quota information for unused models', async () => {
const now = new Date('2025-01-01T12:00:00Z');
vi.useFakeTimers();
vi.setSystemTime(now);
// No models in metrics, but a quota for gemini-2.5-flash
const metrics = createTestMetrics();
const resetTime = new Date(now.getTime() + 1000 * 60 * 120).toISOString(); // 2 hours from now
const quotas: RetrieveUserQuotaResponse = {
buckets: [
{
modelId: 'gemini-2.5-flash',
remainingAmount: '50',
remainingFraction: 0.5,
resetTime,
},
],
};
useSessionStatsMock.mockReturnValue({
stats: {
sessionId: 'test-session-id',
sessionStartTime: new Date(),
metrics,
lastPromptTokenCount: 0,
promptCount: 5,
},
getPromptCount: () => 5,
startNewPrompt: vi.fn(),
});
const { lastFrame } = await renderWithProviders(
<StatsDisplay duration="1s" quotas={quotas} />,
{ width: 100 },
);
const output = lastFrame();
expect(output).toContain('gemini-2.5-flash');
expect(output).toContain('-'); // for requests
expect(output).toContain('50%');
expect(output).toContain('Usage resets');
expect(output).toMatchSnapshot();
vi.useRealTimers();
});
});
describe('User Identity Display', () => {
it('renders User row with Auth Method and Tier', async () => {
const metrics = createTestMetrics();