feat: Add UI for /stats slash command (#883)
This commit is contained in:
Abhi
GitHub
parent
04e2fe0bff
commit
9c3f34890f
|
|
@ -0,0 +1,76 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { render } from 'ink-testing-library';
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { HistoryItemDisplay } from './HistoryItemDisplay.js';
|
||||
import { HistoryItem, MessageType } from '../types.js';
|
||||
import { CumulativeStats } from '../contexts/SessionContext.js';
|
||||
|
||||
// Mock child components
|
||||
vi.mock('./messages/ToolGroupMessage.js', () => ({
|
||||
ToolGroupMessage: () => <div />,
|
||||
}));
|
||||
|
||||
describe('<HistoryItemDisplay />', () => {
|
||||
const baseItem = {
|
||||
id: 1,
|
||||
timestamp: 12345,
|
||||
isPending: false,
|
||||
availableTerminalHeight: 100,
|
||||
};
|
||||
|
||||
it('renders UserMessage for "user" type', () => {
|
||||
const item: HistoryItem = {
|
||||
...baseItem,
|
||||
type: MessageType.USER,
|
||||
text: 'Hello',
|
||||
};
|
||||
const { lastFrame } = render(
|
||||
<HistoryItemDisplay {...baseItem} item={item} />,
|
||||
);
|
||||
expect(lastFrame()).toContain('Hello');
|
||||
});
|
||||
|
||||
it('renders StatsDisplay for "stats" type', () => {
|
||||
const stats: CumulativeStats = {
|
||||
turnCount: 1,
|
||||
promptTokenCount: 10,
|
||||
candidatesTokenCount: 20,
|
||||
totalTokenCount: 30,
|
||||
cachedContentTokenCount: 5,
|
||||
toolUsePromptTokenCount: 2,
|
||||
thoughtsTokenCount: 3,
|
||||
apiTimeMs: 123,
|
||||
};
|
||||
const item: HistoryItem = {
|
||||
...baseItem,
|
||||
type: MessageType.STATS,
|
||||
stats,
|
||||
lastTurnStats: stats,
|
||||
duration: '1s',
|
||||
};
|
||||
const { lastFrame } = render(
|
||||
<HistoryItemDisplay {...baseItem} item={item} />,
|
||||
);
|
||||
expect(lastFrame()).toContain('Stats');
|
||||
});
|
||||
|
||||
it('renders AboutBox for "about" type', () => {
|
||||
const item: HistoryItem = {
|
||||
...baseItem,
|
||||
type: MessageType.ABOUT,
|
||||
cliVersion: '1.0.0',
|
||||
osVersion: 'test-os',
|
||||
sandboxEnv: 'test-env',
|
||||
modelVersion: 'test-model',
|
||||
};
|
||||
const { lastFrame } = render(
|
||||
<HistoryItemDisplay {...baseItem} item={item} />,
|
||||
);
|
||||
expect(lastFrame()).toContain('About Gemini CLI');
|
||||
});
|
||||
});
|
||||
|
|
@ -15,6 +15,7 @@ import { ToolGroupMessage } from './messages/ToolGroupMessage.js';
|
|||
import { GeminiMessageContent } from './messages/GeminiMessageContent.js';
|
||||
import { Box } from 'ink';
|
||||
import { AboutBox } from './AboutBox.js';
|
||||
import { StatsDisplay } from './StatsDisplay.js';
|
||||
import { Config } from '@gemini-cli/core';
|
||||
|
||||
interface HistoryItemDisplayProps {
|
||||
|
|
@ -58,6 +59,13 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
|
|||
modelVersion={item.modelVersion}
|
||||
/>
|
||||
)}
|
||||
{item.type === 'stats' && (
|
||||
<StatsDisplay
|
||||
stats={item.stats}
|
||||
lastTurnStats={item.lastTurnStats}
|
||||
duration={item.duration}
|
||||
/>
|
||||
)}
|
||||
{item.type === 'tool_group' && (
|
||||
<ToolGroupMessage
|
||||
toolCalls={item.tools}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,71 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { render } from 'ink-testing-library';
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { StatsDisplay } from './StatsDisplay.js';
|
||||
import { type CumulativeStats } from '../contexts/SessionContext.js';
|
||||
|
||||
describe('<StatsDisplay />', () => {
|
||||
const mockStats: CumulativeStats = {
|
||||
turnCount: 10,
|
||||
promptTokenCount: 1000,
|
||||
candidatesTokenCount: 2000,
|
||||
totalTokenCount: 3500,
|
||||
cachedContentTokenCount: 500,
|
||||
toolUsePromptTokenCount: 200,
|
||||
thoughtsTokenCount: 300,
|
||||
apiTimeMs: 50234,
|
||||
};
|
||||
|
||||
const mockLastTurnStats: CumulativeStats = {
|
||||
turnCount: 1,
|
||||
promptTokenCount: 100,
|
||||
candidatesTokenCount: 200,
|
||||
totalTokenCount: 350,
|
||||
cachedContentTokenCount: 50,
|
||||
toolUsePromptTokenCount: 20,
|
||||
thoughtsTokenCount: 30,
|
||||
apiTimeMs: 1234,
|
||||
};
|
||||
|
||||
const mockDuration = '1h 23m 45s';
|
||||
|
||||
it('renders correctly with given stats and duration', () => {
|
||||
const { lastFrame } = render(
|
||||
<StatsDisplay
|
||||
stats={mockStats}
|
||||
lastTurnStats={mockLastTurnStats}
|
||||
duration={mockDuration}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(lastFrame()).toMatchSnapshot();
|
||||
});
|
||||
|
||||
it('renders zero state correctly', () => {
|
||||
const zeroStats: CumulativeStats = {
|
||||
turnCount: 0,
|
||||
promptTokenCount: 0,
|
||||
candidatesTokenCount: 0,
|
||||
totalTokenCount: 0,
|
||||
cachedContentTokenCount: 0,
|
||||
toolUsePromptTokenCount: 0,
|
||||
thoughtsTokenCount: 0,
|
||||
apiTimeMs: 0,
|
||||
};
|
||||
|
||||
const { lastFrame } = render(
|
||||
<StatsDisplay
|
||||
stats={zeroStats}
|
||||
lastTurnStats={zeroStats}
|
||||
duration="0s"
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(lastFrame()).toMatchSnapshot();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,174 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import React from 'react';
|
||||
import { Box, Text } from 'ink';
|
||||
import { Colors } from '../colors.js';
|
||||
import { formatDuration } from '../utils/formatters.js';
|
||||
import { CumulativeStats } from '../contexts/SessionContext.js';
|
||||
|
||||
// --- Constants ---
|
||||
|
||||
const COLUMN_WIDTH = '48%';
|
||||
|
||||
// --- Prop and Data Structures ---
|
||||
|
||||
interface StatsDisplayProps {
|
||||
stats: CumulativeStats;
|
||||
lastTurnStats: CumulativeStats;
|
||||
duration: string;
|
||||
}
|
||||
|
||||
interface FormattedStats {
|
||||
inputTokens: number;
|
||||
outputTokens: number;
|
||||
toolUseTokens: number;
|
||||
thoughtsTokens: number;
|
||||
cachedTokens: number;
|
||||
totalTokens: number;
|
||||
}
|
||||
|
||||
// --- Helper Components ---
|
||||
|
||||
/**
|
||||
* Renders a single row with a colored label on the left and a value on the right.
|
||||
*/
|
||||
const StatRow: React.FC<{
|
||||
label: string;
|
||||
value: string | number;
|
||||
valueColor?: string;
|
||||
}> = ({ label, value, valueColor }) => (
|
||||
<Box justifyContent="space-between">
|
||||
<Text color={Colors.LightBlue}>{label}</Text>
|
||||
<Text color={valueColor}>{value}</Text>
|
||||
</Box>
|
||||
);
|
||||
|
||||
/**
|
||||
* Renders a full column for either "Last Turn" or "Cumulative" stats.
|
||||
*/
|
||||
const StatsColumn: React.FC<{
|
||||
title: string;
|
||||
stats: FormattedStats;
|
||||
isCumulative?: boolean;
|
||||
}> = ({ title, stats, isCumulative = false }) => {
|
||||
const cachedDisplay =
|
||||
isCumulative && stats.totalTokens > 0
|
||||
? `${stats.cachedTokens.toLocaleString()} (${((stats.cachedTokens / stats.totalTokens) * 100).toFixed(1)}%)`
|
||||
: stats.cachedTokens.toLocaleString();
|
||||
|
||||
const cachedColor =
|
||||
isCumulative && stats.cachedTokens > 0 ? Colors.AccentGreen : undefined;
|
||||
|
||||
return (
|
||||
<Box flexDirection="column" width={COLUMN_WIDTH}>
|
||||
<Text bold>{title}</Text>
|
||||
<Box marginTop={1} flexDirection="column">
|
||||
<StatRow
|
||||
label="Input Tokens"
|
||||
value={stats.inputTokens.toLocaleString()}
|
||||
/>
|
||||
<StatRow
|
||||
label="Output Tokens"
|
||||
value={stats.outputTokens.toLocaleString()}
|
||||
/>
|
||||
<StatRow
|
||||
label="Tool Use Tokens"
|
||||
value={stats.toolUseTokens.toLocaleString()}
|
||||
/>
|
||||
<StatRow
|
||||
label="Thoughts Tokens"
|
||||
value={stats.thoughtsTokens.toLocaleString()}
|
||||
/>
|
||||
<StatRow
|
||||
label="Cached Tokens"
|
||||
value={cachedDisplay}
|
||||
valueColor={cachedColor}
|
||||
/>
|
||||
{/* Divider Line */}
|
||||
<Box
|
||||
borderTop={true}
|
||||
borderLeft={false}
|
||||
borderRight={false}
|
||||
borderBottom={false}
|
||||
borderStyle="single"
|
||||
/>
|
||||
<StatRow
|
||||
label="Total Tokens"
|
||||
value={stats.totalTokens.toLocaleString()}
|
||||
/>
|
||||
</Box>
|
||||
</Box>
|
||||
);
|
||||
};
|
||||
|
||||
// --- Main Component ---
|
||||
|
||||
export const StatsDisplay: React.FC<StatsDisplayProps> = ({
|
||||
stats,
|
||||
lastTurnStats,
|
||||
duration,
|
||||
}) => {
|
||||
const lastTurnFormatted: FormattedStats = {
|
||||
inputTokens: lastTurnStats.promptTokenCount,
|
||||
outputTokens: lastTurnStats.candidatesTokenCount,
|
||||
toolUseTokens: lastTurnStats.toolUsePromptTokenCount,
|
||||
thoughtsTokens: lastTurnStats.thoughtsTokenCount,
|
||||
cachedTokens: lastTurnStats.cachedContentTokenCount,
|
||||
totalTokens: lastTurnStats.totalTokenCount,
|
||||
};
|
||||
|
||||
const cumulativeFormatted: FormattedStats = {
|
||||
inputTokens: stats.promptTokenCount,
|
||||
outputTokens: stats.candidatesTokenCount,
|
||||
toolUseTokens: stats.toolUsePromptTokenCount,
|
||||
thoughtsTokens: stats.thoughtsTokenCount,
|
||||
cachedTokens: stats.cachedContentTokenCount,
|
||||
totalTokens: stats.totalTokenCount,
|
||||
};
|
||||
|
||||
return (
|
||||
<Box
|
||||
borderStyle="round"
|
||||
borderColor="gray"
|
||||
flexDirection="column"
|
||||
paddingY={1}
|
||||
paddingX={2}
|
||||
>
|
||||
<Text bold color={Colors.AccentPurple}>
|
||||
Stats
|
||||
</Text>
|
||||
|
||||
<Box flexDirection="row" justifyContent="space-between" marginTop={1}>
|
||||
<StatsColumn title="Last Turn" stats={lastTurnFormatted} />
|
||||
<StatsColumn
|
||||
title={`Cumulative (${stats.turnCount} Turns)`}
|
||||
stats={cumulativeFormatted}
|
||||
isCumulative={true}
|
||||
/>
|
||||
</Box>
|
||||
|
||||
<Box flexDirection="row" justifyContent="space-between" marginTop={1}>
|
||||
{/* Left column for "Last Turn" duration */}
|
||||
<Box width={COLUMN_WIDTH} flexDirection="column">
|
||||
<StatRow
|
||||
label="Turn Duration (API)"
|
||||
value={formatDuration(lastTurnStats.apiTimeMs)}
|
||||
/>
|
||||
</Box>
|
||||
|
||||
{/* Right column for "Cumulative" durations */}
|
||||
<Box width={COLUMN_WIDTH} flexDirection="column">
|
||||
<StatRow
|
||||
label="Total duration (API)"
|
||||
value={formatDuration(stats.apiTimeMs)}
|
||||
/>
|
||||
<StatRow label="Total duration (wall)" value={duration} />
|
||||
</Box>
|
||||
</Box>
|
||||
</Box>
|
||||
);
|
||||
};
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
|
||||
|
||||
exports[`<StatsDisplay /> > renders correctly with given stats and duration 1`] = `
|
||||
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
|
||||
│ │
|
||||
│ Stats │
|
||||
│ │
|
||||
│ Last Turn Cumulative (10 Turns) │
|
||||
│ │
|
||||
│ Input Tokens 100 Input Tokens 1,000 │
|
||||
│ Output Tokens 200 Output Tokens 2,000 │
|
||||
│ Tool Use Tokens 20 Tool Use Tokens 200 │
|
||||
│ Thoughts Tokens 30 Thoughts Tokens 300 │
|
||||
│ Cached Tokens 50 Cached Tokens 500 (14.3%) │
|
||||
│ ───────────────────────────────────────────── ───────────────────────────────────────────── │
|
||||
│ Total Tokens 350 Total Tokens 3,500 │
|
||||
│ │
|
||||
│ Turn Duration (API) 1.2s Total duration (API) 50.2s │
|
||||
│ Total duration (wall) 1h 23m 45s │
|
||||
│ │
|
||||
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
|
||||
`;
|
||||
|
||||
exports[`<StatsDisplay /> > renders zero state correctly 1`] = `
|
||||
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
|
||||
│ │
|
||||
│ Stats │
|
||||
│ │
|
||||
│ Last Turn Cumulative (0 Turns) │
|
||||
│ │
|
||||
│ Input Tokens 0 Input Tokens 0 │
|
||||
│ Output Tokens 0 Output Tokens 0 │
|
||||
│ Tool Use Tokens 0 Tool Use Tokens 0 │
|
||||
│ Thoughts Tokens 0 Thoughts Tokens 0 │
|
||||
│ Cached Tokens 0 Cached Tokens 0 │
|
||||
│ ───────────────────────────────────────────── ───────────────────────────────────────────── │
|
||||
│ Total Tokens 0 Total Tokens 0 │
|
||||
│ │
|
||||
│ Turn Duration (API) 0s Total duration (API) 0s │
|
||||
│ Total duration (wall) 0s │
|
||||
│ │
|
||||
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
|
||||
`;
|
||||
|
|
@ -59,7 +59,7 @@ describe('SessionStatsContext', () => {
|
|||
const stats = contextRef.current?.stats;
|
||||
|
||||
expect(stats?.sessionStartTime).toBeInstanceOf(Date);
|
||||
expect(stats?.lastTurn).toBeNull();
|
||||
expect(stats?.currentTurn).toBeDefined();
|
||||
expect(stats?.cumulative.turnCount).toBe(0);
|
||||
expect(stats?.cumulative.totalTokenCount).toBe(0);
|
||||
expect(stats?.cumulative.promptTokenCount).toBe(0);
|
||||
|
|
@ -81,6 +81,7 @@ describe('SessionStatsContext', () => {
|
|||
});
|
||||
|
||||
const stats = contextRef.current?.stats;
|
||||
expect(stats?.currentTurn.totalTokenCount).toBe(0);
|
||||
expect(stats?.cumulative.turnCount).toBe(1);
|
||||
// Ensure token counts are unaffected
|
||||
expect(stats?.cumulative.totalTokenCount).toBe(0);
|
||||
|
|
@ -98,7 +99,7 @@ describe('SessionStatsContext', () => {
|
|||
);
|
||||
|
||||
act(() => {
|
||||
contextRef.current?.addUsage(mockMetadata1);
|
||||
contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 123 });
|
||||
});
|
||||
|
||||
const stats = contextRef.current?.stats;
|
||||
|
|
@ -110,12 +111,16 @@ describe('SessionStatsContext', () => {
|
|||
expect(stats?.cumulative.promptTokenCount).toBe(
|
||||
mockMetadata1.promptTokenCount ?? 0,
|
||||
);
|
||||
expect(stats?.cumulative.apiTimeMs).toBe(123);
|
||||
|
||||
// Check that turn count is NOT incremented
|
||||
expect(stats?.cumulative.turnCount).toBe(0);
|
||||
|
||||
// Check that lastTurn is updated
|
||||
expect(stats?.lastTurn?.metadata).toEqual(mockMetadata1);
|
||||
// Check that currentTurn is updated
|
||||
expect(stats?.currentTurn?.totalTokenCount).toEqual(
|
||||
mockMetadata1.totalTokenCount,
|
||||
);
|
||||
expect(stats?.currentTurn?.apiTimeMs).toBe(123);
|
||||
});
|
||||
|
||||
it('should correctly track a full logical turn with multiple API calls', () => {
|
||||
|
|
@ -136,12 +141,12 @@ describe('SessionStatsContext', () => {
|
|||
|
||||
// 2. First API call (e.g., prompt with a tool request)
|
||||
act(() => {
|
||||
contextRef.current?.addUsage(mockMetadata1);
|
||||
contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
|
||||
});
|
||||
|
||||
// 3. Second API call (e.g., sending tool response back)
|
||||
act(() => {
|
||||
contextRef.current?.addUsage(mockMetadata2);
|
||||
contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
|
||||
});
|
||||
|
||||
const stats = contextRef.current?.stats;
|
||||
|
|
@ -149,18 +154,27 @@ describe('SessionStatsContext', () => {
|
|||
// Turn count should only be 1
|
||||
expect(stats?.cumulative.turnCount).toBe(1);
|
||||
|
||||
// --- Check Cumulative Stats ---
|
||||
// These fields should be the SUM of both calls
|
||||
expect(stats?.cumulative.totalTokenCount).toBe(330); // 300 + 30
|
||||
expect(stats?.cumulative.candidatesTokenCount).toBe(220); // 200 + 20
|
||||
expect(stats?.cumulative.thoughtsTokenCount).toBe(22); // 20 + 2
|
||||
expect(stats?.cumulative.totalTokenCount).toBe(300 + 30);
|
||||
expect(stats?.cumulative.candidatesTokenCount).toBe(200 + 20);
|
||||
expect(stats?.cumulative.thoughtsTokenCount).toBe(20 + 2);
|
||||
expect(stats?.cumulative.apiTimeMs).toBe(100 + 50);
|
||||
|
||||
// These fields should ONLY be from the FIRST call, because isNewTurnForAggregation was true
|
||||
expect(stats?.cumulative.promptTokenCount).toBe(100);
|
||||
expect(stats?.cumulative.cachedContentTokenCount).toBe(50);
|
||||
expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10);
|
||||
// These fields should be the SUM of both calls
|
||||
expect(stats?.cumulative.promptTokenCount).toBe(100 + 10);
|
||||
expect(stats?.cumulative.cachedContentTokenCount).toBe(50 + 5);
|
||||
expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10 + 1);
|
||||
|
||||
// Last turn should hold the metadata from the most recent call
|
||||
expect(stats?.lastTurn?.metadata).toEqual(mockMetadata2);
|
||||
// --- Check Current Turn Stats ---
|
||||
// All fields should be the SUM of both calls for the turn
|
||||
expect(stats?.currentTurn.totalTokenCount).toBe(300 + 30);
|
||||
expect(stats?.currentTurn.candidatesTokenCount).toBe(200 + 20);
|
||||
expect(stats?.currentTurn.thoughtsTokenCount).toBe(20 + 2);
|
||||
expect(stats?.currentTurn.promptTokenCount).toBe(100 + 10);
|
||||
expect(stats?.currentTurn.cachedContentTokenCount).toBe(50 + 5);
|
||||
expect(stats?.currentTurn.toolUsePromptTokenCount).toBe(10 + 1);
|
||||
expect(stats?.currentTurn.apiTimeMs).toBe(100 + 50);
|
||||
});
|
||||
|
||||
it('should throw an error when useSessionStats is used outside of a provider', () => {
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import { type GenerateContentResponseUsageMetadata } from '@google/genai';
|
|||
|
||||
// --- Interface Definitions ---
|
||||
|
||||
interface CumulativeStats {
|
||||
export interface CumulativeStats {
|
||||
turnCount: number;
|
||||
promptTokenCount: number;
|
||||
candidatesTokenCount: number;
|
||||
|
|
@ -24,18 +24,13 @@ interface CumulativeStats {
|
|||
cachedContentTokenCount: number;
|
||||
toolUsePromptTokenCount: number;
|
||||
thoughtsTokenCount: number;
|
||||
}
|
||||
|
||||
interface LastTurnStats {
|
||||
metadata: GenerateContentResponseUsageMetadata;
|
||||
// TODO(abhipatel12): Add apiTime, etc. here in a future step.
|
||||
apiTimeMs: number;
|
||||
}
|
||||
|
||||
interface SessionStatsState {
|
||||
sessionStartTime: Date;
|
||||
cumulative: CumulativeStats;
|
||||
lastTurn: LastTurnStats | null;
|
||||
isNewTurnForAggregation: boolean;
|
||||
currentTurn: CumulativeStats;
|
||||
}
|
||||
|
||||
// Defines the final "value" of our context, including the state
|
||||
|
|
@ -43,7 +38,9 @@ interface SessionStatsState {
|
|||
interface SessionStatsContextValue {
|
||||
stats: SessionStatsState;
|
||||
startNewTurn: () => void;
|
||||
addUsage: (metadata: GenerateContentResponseUsageMetadata) => void;
|
||||
addUsage: (
|
||||
metadata: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
|
||||
) => void;
|
||||
}
|
||||
|
||||
// --- Context Definition ---
|
||||
|
|
@ -52,6 +49,27 @@ const SessionStatsContext = createContext<SessionStatsContextValue | undefined>(
|
|||
undefined,
|
||||
);
|
||||
|
||||
// --- Helper Functions ---
|
||||
|
||||
/**
|
||||
* A small, reusable helper function to sum token counts.
|
||||
* It unconditionally adds all token values from the source to the target.
|
||||
* @param target The object to add the tokens to (e.g., cumulative, currentTurn).
|
||||
* @param source The metadata object from the API response.
|
||||
*/
|
||||
const addTokens = (
|
||||
target: CumulativeStats,
|
||||
source: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
|
||||
) => {
|
||||
target.candidatesTokenCount += source.candidatesTokenCount ?? 0;
|
||||
target.thoughtsTokenCount += source.thoughtsTokenCount ?? 0;
|
||||
target.totalTokenCount += source.totalTokenCount ?? 0;
|
||||
target.apiTimeMs += source.apiTimeMs ?? 0;
|
||||
target.promptTokenCount += source.promptTokenCount ?? 0;
|
||||
target.cachedContentTokenCount += source.cachedContentTokenCount ?? 0;
|
||||
target.toolUsePromptTokenCount += source.toolUsePromptTokenCount ?? 0;
|
||||
};
|
||||
|
||||
// --- Provider Component ---
|
||||
|
||||
export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
|
||||
|
|
@ -67,36 +85,37 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
|
|||
cachedContentTokenCount: 0,
|
||||
toolUsePromptTokenCount: 0,
|
||||
thoughtsTokenCount: 0,
|
||||
apiTimeMs: 0,
|
||||
},
|
||||
currentTurn: {
|
||||
turnCount: 0,
|
||||
promptTokenCount: 0,
|
||||
candidatesTokenCount: 0,
|
||||
totalTokenCount: 0,
|
||||
cachedContentTokenCount: 0,
|
||||
toolUsePromptTokenCount: 0,
|
||||
thoughtsTokenCount: 0,
|
||||
apiTimeMs: 0,
|
||||
},
|
||||
lastTurn: null,
|
||||
isNewTurnForAggregation: true,
|
||||
});
|
||||
|
||||
// A single, internal worker function to handle all metadata aggregation.
|
||||
const aggregateTokens = useCallback(
|
||||
(metadata: GenerateContentResponseUsageMetadata) => {
|
||||
(
|
||||
metadata: GenerateContentResponseUsageMetadata & { apiTimeMs?: number },
|
||||
) => {
|
||||
setStats((prevState) => {
|
||||
const { isNewTurnForAggregation } = prevState;
|
||||
const newCumulative = { ...prevState.cumulative };
|
||||
const newCurrentTurn = { ...prevState.currentTurn };
|
||||
|
||||
newCumulative.candidatesTokenCount +=
|
||||
metadata.candidatesTokenCount ?? 0;
|
||||
newCumulative.thoughtsTokenCount += metadata.thoughtsTokenCount ?? 0;
|
||||
newCumulative.totalTokenCount += metadata.totalTokenCount ?? 0;
|
||||
|
||||
if (isNewTurnForAggregation) {
|
||||
newCumulative.promptTokenCount += metadata.promptTokenCount ?? 0;
|
||||
newCumulative.cachedContentTokenCount +=
|
||||
metadata.cachedContentTokenCount ?? 0;
|
||||
newCumulative.toolUsePromptTokenCount +=
|
||||
metadata.toolUsePromptTokenCount ?? 0;
|
||||
}
|
||||
// Add all tokens to the current turn's stats as well as cumulative stats.
|
||||
addTokens(newCurrentTurn, metadata);
|
||||
addTokens(newCumulative, metadata);
|
||||
|
||||
return {
|
||||
...prevState,
|
||||
cumulative: newCumulative,
|
||||
lastTurn: { metadata },
|
||||
isNewTurnForAggregation: false,
|
||||
currentTurn: newCurrentTurn,
|
||||
};
|
||||
});
|
||||
},
|
||||
|
|
@ -110,7 +129,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
|
|||
...prevState.cumulative,
|
||||
turnCount: prevState.cumulative.turnCount + 1,
|
||||
},
|
||||
isNewTurnForAggregation: true,
|
||||
currentTurn: {
|
||||
turnCount: 0, // Reset for the new turn's accumulation.
|
||||
promptTokenCount: 0,
|
||||
candidatesTokenCount: 0,
|
||||
totalTokenCount: 0,
|
||||
cachedContentTokenCount: 0,
|
||||
toolUsePromptTokenCount: 0,
|
||||
thoughtsTokenCount: 0,
|
||||
apiTimeMs: 0,
|
||||
},
|
||||
}));
|
||||
}, []);
|
||||
|
||||
|
|
|
|||
|
|
@ -255,18 +255,19 @@ describe('useSlashCommandProcessor', () => {
|
|||
describe('/stats command', () => {
|
||||
it('should show detailed session statistics', async () => {
|
||||
// Arrange
|
||||
const cumulativeStats = {
|
||||
totalTokenCount: 900,
|
||||
promptTokenCount: 200,
|
||||
candidatesTokenCount: 400,
|
||||
cachedContentTokenCount: 100,
|
||||
turnCount: 1,
|
||||
toolUsePromptTokenCount: 50,
|
||||
thoughtsTokenCount: 150,
|
||||
};
|
||||
mockUseSessionStats.mockReturnValue({
|
||||
stats: {
|
||||
sessionStartTime: new Date('2025-01-01T00:00:00.000Z'),
|
||||
cumulative: {
|
||||
totalTokenCount: 900,
|
||||
promptTokenCount: 200,
|
||||
candidatesTokenCount: 400,
|
||||
cachedContentTokenCount: 100,
|
||||
turnCount: 1,
|
||||
toolUsePromptTokenCount: 50,
|
||||
thoughtsTokenCount: 150,
|
||||
},
|
||||
cumulative: cumulativeStats,
|
||||
},
|
||||
});
|
||||
|
||||
|
|
@ -280,24 +281,12 @@ describe('useSlashCommandProcessor', () => {
|
|||
});
|
||||
|
||||
// Assert
|
||||
const expectedContent = [
|
||||
` ⎿ Total duration (wall): 1h 2m 3s`,
|
||||
` Total Token usage:`,
|
||||
` Turns: 1`,
|
||||
` Total: 900`,
|
||||
` ├─ Input: 200`,
|
||||
` ├─ Output: 400`,
|
||||
` ├─ Cached: 100`,
|
||||
` └─ Overhead: 200`,
|
||||
` ├─ Model thoughts: 150`,
|
||||
` └─ Tool-use prompts: 50`,
|
||||
].join('\n');
|
||||
|
||||
expect(mockAddItem).toHaveBeenNthCalledWith(
|
||||
2, // Called after the user message
|
||||
expect.objectContaining({
|
||||
type: MessageType.INFO,
|
||||
text: expectedContent,
|
||||
type: MessageType.STATS,
|
||||
stats: cumulativeStats,
|
||||
duration: '1h 2m 3s',
|
||||
}),
|
||||
expect.any(Number),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ import { Message, MessageType, HistoryItemWithoutId } from '../types.js';
|
|||
import { useSessionStats } from '../contexts/SessionContext.js';
|
||||
import { createShowMemoryAction } from './useShowMemoryCommand.js';
|
||||
import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
|
||||
import { formatMemoryUsage } from '../utils/formatters.js';
|
||||
import { formatDuration, formatMemoryUsage } from '../utils/formatters.js';
|
||||
import { getCliVersion } from '../../utils/version.js';
|
||||
|
||||
export interface SlashCommandActionReturn {
|
||||
|
|
@ -69,6 +69,13 @@ export const useSlashCommandProcessor = (
|
|||
sandboxEnv: message.sandboxEnv,
|
||||
modelVersion: message.modelVersion,
|
||||
};
|
||||
} else if (message.type === MessageType.STATS) {
|
||||
historyItemContent = {
|
||||
type: 'stats',
|
||||
stats: message.stats,
|
||||
lastTurnStats: message.lastTurnStats,
|
||||
duration: message.duration,
|
||||
};
|
||||
} else {
|
||||
historyItemContent = {
|
||||
type: message.type as
|
||||
|
|
@ -152,41 +159,14 @@ export const useSlashCommandProcessor = (
|
|||
description: 'check session stats',
|
||||
action: (_mainCommand, _subCommand, _args) => {
|
||||
const now = new Date();
|
||||
const { sessionStartTime, cumulative } = session.stats;
|
||||
|
||||
const duration = now.getTime() - sessionStartTime.getTime();
|
||||
const durationInSeconds = Math.floor(duration / 1000);
|
||||
const hours = Math.floor(durationInSeconds / 3600);
|
||||
const minutes = Math.floor((durationInSeconds % 3600) / 60);
|
||||
const seconds = durationInSeconds % 60;
|
||||
|
||||
const durationString = [
|
||||
hours > 0 ? `${hours}h` : '',
|
||||
minutes > 0 ? `${minutes}m` : '',
|
||||
`${seconds}s`,
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join(' ');
|
||||
|
||||
const overheadTotal =
|
||||
cumulative.thoughtsTokenCount + cumulative.toolUsePromptTokenCount;
|
||||
|
||||
const statsContent = [
|
||||
` ⎿ Total duration (wall): ${durationString}`,
|
||||
` Total Token usage:`,
|
||||
` Turns: ${cumulative.turnCount.toLocaleString()}`,
|
||||
` Total: ${cumulative.totalTokenCount.toLocaleString()}`,
|
||||
` ├─ Input: ${cumulative.promptTokenCount.toLocaleString()}`,
|
||||
` ├─ Output: ${cumulative.candidatesTokenCount.toLocaleString()}`,
|
||||
` ├─ Cached: ${cumulative.cachedContentTokenCount.toLocaleString()}`,
|
||||
` └─ Overhead: ${overheadTotal.toLocaleString()}`,
|
||||
` ├─ Model thoughts: ${cumulative.thoughtsTokenCount.toLocaleString()}`,
|
||||
` └─ Tool-use prompts: ${cumulative.toolUsePromptTokenCount.toLocaleString()}`,
|
||||
].join('\n');
|
||||
const { sessionStartTime, cumulative, currentTurn } = session.stats;
|
||||
const wallDuration = now.getTime() - sessionStartTime.getTime();
|
||||
|
||||
addMessage({
|
||||
type: MessageType.INFO,
|
||||
content: statsContent,
|
||||
type: MessageType.STATS,
|
||||
stats: cumulative,
|
||||
lastTurnStats: currentTurn,
|
||||
duration: formatDuration(wallDuration),
|
||||
timestamp: new Date(),
|
||||
});
|
||||
},
|
||||
|
|
|
|||
|
|
@ -598,5 +598,18 @@ describe('useGeminiStream', () => {
|
|||
expect(mockStartNewTurn).toHaveBeenCalledTimes(1);
|
||||
expect(mockAddUsage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not call startNewTurn for a slash command', async () => {
|
||||
mockHandleSlashCommand.mockReturnValue(true);
|
||||
|
||||
const { result } = renderTestHook();
|
||||
|
||||
await act(async () => {
|
||||
await result.current.submitQuery('/stats');
|
||||
});
|
||||
|
||||
expect(mockStartNewTurn).not.toHaveBeenCalled();
|
||||
expect(mockSendMessageStream).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -432,10 +432,6 @@ export const useGeminiStream = (
|
|||
const userMessageTimestamp = Date.now();
|
||||
setShowHelp(false);
|
||||
|
||||
if (!options?.isContinuation) {
|
||||
startNewTurn();
|
||||
}
|
||||
|
||||
abortControllerRef.current = new AbortController();
|
||||
const abortSignal = abortControllerRef.current.signal;
|
||||
|
||||
|
|
@ -449,6 +445,10 @@ export const useGeminiStream = (
|
|||
return;
|
||||
}
|
||||
|
||||
if (!options?.isContinuation) {
|
||||
startNewTurn();
|
||||
}
|
||||
|
||||
if (!geminiClient) {
|
||||
const errorMsg = 'Gemini client is not available.';
|
||||
setInitError(errorMsg);
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import {
|
|||
ToolCallConfirmationDetails,
|
||||
ToolResultDisplay,
|
||||
} from '@gemini-cli/core';
|
||||
import { CumulativeStats } from './contexts/SessionContext.js';
|
||||
|
||||
// Only defining the state enum needed by the UI
|
||||
export enum StreamingState {
|
||||
|
|
@ -89,6 +90,13 @@ export type HistoryItemAbout = HistoryItemBase & {
|
|||
modelVersion: string;
|
||||
};
|
||||
|
||||
export type HistoryItemStats = HistoryItemBase & {
|
||||
type: 'stats';
|
||||
stats: CumulativeStats;
|
||||
lastTurnStats: CumulativeStats;
|
||||
duration: string;
|
||||
};
|
||||
|
||||
export type HistoryItemToolGroup = HistoryItemBase & {
|
||||
type: 'tool_group';
|
||||
tools: IndividualToolCallDisplay[];
|
||||
|
|
@ -111,7 +119,8 @@ export type HistoryItemWithoutId =
|
|||
| HistoryItemInfo
|
||||
| HistoryItemError
|
||||
| HistoryItemAbout
|
||||
| HistoryItemToolGroup;
|
||||
| HistoryItemToolGroup
|
||||
| HistoryItemStats;
|
||||
|
||||
export type HistoryItem = HistoryItemWithoutId & { id: number };
|
||||
|
||||
|
|
@ -121,6 +130,7 @@ export enum MessageType {
|
|||
ERROR = 'error',
|
||||
USER = 'user',
|
||||
ABOUT = 'about',
|
||||
STATS = 'stats',
|
||||
// Add GEMINI if needed by other commands
|
||||
}
|
||||
|
||||
|
|
@ -139,6 +149,14 @@ export type Message =
|
|||
sandboxEnv: string;
|
||||
modelVersion: string;
|
||||
content?: string; // Optional content, not really used for ABOUT
|
||||
}
|
||||
| {
|
||||
type: MessageType.STATS;
|
||||
timestamp: Date;
|
||||
stats: CumulativeStats;
|
||||
lastTurnStats: CumulativeStats;
|
||||
duration: string;
|
||||
content?: string;
|
||||
};
|
||||
|
||||
export interface ConsoleMessageItem {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,72 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { formatDuration, formatMemoryUsage } from './formatters.js';
|
||||
|
||||
describe('formatters', () => {
|
||||
describe('formatMemoryUsage', () => {
|
||||
it('should format bytes into KB', () => {
|
||||
expect(formatMemoryUsage(12345)).toBe('12.1 KB');
|
||||
});
|
||||
|
||||
it('should format bytes into MB', () => {
|
||||
expect(formatMemoryUsage(12345678)).toBe('11.8 MB');
|
||||
});
|
||||
|
||||
it('should format bytes into GB', () => {
|
||||
expect(formatMemoryUsage(12345678901)).toBe('11.50 GB');
|
||||
});
|
||||
});
|
||||
|
||||
describe('formatDuration', () => {
|
||||
it('should format milliseconds less than a second', () => {
|
||||
expect(formatDuration(500)).toBe('500ms');
|
||||
});
|
||||
|
||||
it('should format a duration of 0', () => {
|
||||
expect(formatDuration(0)).toBe('0s');
|
||||
});
|
||||
|
||||
it('should format an exact number of seconds', () => {
|
||||
expect(formatDuration(5000)).toBe('5.0s');
|
||||
});
|
||||
|
||||
it('should format a duration in seconds with one decimal place', () => {
|
||||
expect(formatDuration(12345)).toBe('12.3s');
|
||||
});
|
||||
|
||||
it('should format an exact number of minutes', () => {
|
||||
expect(formatDuration(120000)).toBe('2m');
|
||||
});
|
||||
|
||||
it('should format a duration in minutes and seconds', () => {
|
||||
expect(formatDuration(123000)).toBe('2m 3s');
|
||||
});
|
||||
|
||||
it('should format an exact number of hours', () => {
|
||||
expect(formatDuration(3600000)).toBe('1h');
|
||||
});
|
||||
|
||||
it('should format a duration in hours and seconds', () => {
|
||||
expect(formatDuration(3605000)).toBe('1h 5s');
|
||||
});
|
||||
|
||||
it('should format a duration in hours, minutes, and seconds', () => {
|
||||
expect(formatDuration(3723000)).toBe('1h 2m 3s');
|
||||
});
|
||||
|
||||
it('should handle large durations', () => {
|
||||
expect(formatDuration(86400000 + 3600000 + 120000 + 1000)).toBe(
|
||||
'25h 2m 1s',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle negative durations', () => {
|
||||
expect(formatDuration(-100)).toBe('0s');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -14,3 +14,50 @@ export const formatMemoryUsage = (bytes: number): string => {
|
|||
}
|
||||
return `${gb.toFixed(2)} GB`;
|
||||
};
|
||||
|
||||
/**
|
||||
* Formats a duration in milliseconds into a concise, human-readable string (e.g., "1h 5s").
|
||||
* It omits any time units that are zero.
|
||||
* @param milliseconds The duration in milliseconds.
|
||||
* @returns A formatted string representing the duration.
|
||||
*/
|
||||
export const formatDuration = (milliseconds: number): string => {
|
||||
if (milliseconds <= 0) {
|
||||
return '0s';
|
||||
}
|
||||
|
||||
if (milliseconds < 1000) {
|
||||
return `${milliseconds}ms`;
|
||||
}
|
||||
|
||||
const totalSeconds = milliseconds / 1000;
|
||||
|
||||
if (totalSeconds < 60) {
|
||||
return `${totalSeconds.toFixed(1)}s`;
|
||||
}
|
||||
|
||||
const hours = Math.floor(totalSeconds / 3600);
|
||||
const minutes = Math.floor((totalSeconds % 3600) / 60);
|
||||
const seconds = Math.floor(totalSeconds % 60);
|
||||
|
||||
const parts: string[] = [];
|
||||
|
||||
if (hours > 0) {
|
||||
parts.push(`${hours}h`);
|
||||
}
|
||||
if (minutes > 0) {
|
||||
parts.push(`${minutes}m`);
|
||||
}
|
||||
if (seconds > 0) {
|
||||
parts.push(`${seconds}s`);
|
||||
}
|
||||
|
||||
// If all parts are zero (e.g., exactly 1 hour), return the largest unit.
|
||||
if (parts.length === 0) {
|
||||
if (hours > 0) return `${hours}h`;
|
||||
if (minutes > 0) return `${minutes}m`;
|
||||
return `${seconds}s`;
|
||||
}
|
||||
|
||||
return parts.join(' ');
|
||||
};
|
||||
|
|
|
|||
|
|
@ -239,6 +239,8 @@ describe('Turn', () => {
|
|||
candidates: [{ content: { parts: [{ text: 'First response' }] } }],
|
||||
usageMetadata: mockMetadata1,
|
||||
} as unknown as GenerateContentResponse;
|
||||
// Add a small delay to ensure apiTimeMs is > 0
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
yield {
|
||||
functionCalls: [{ name: 'aTool' }],
|
||||
usageMetadata: mockMetadata2,
|
||||
|
|
@ -262,7 +264,10 @@ describe('Turn', () => {
|
|||
expect(metadataEvent.type).toBe(GeminiEventType.UsageMetadata);
|
||||
|
||||
// The value should be the *last* metadata object received.
|
||||
expect(metadataEvent.value).toEqual(mockMetadata2);
|
||||
expect(metadataEvent.value).toEqual(
|
||||
expect.objectContaining(mockMetadata2),
|
||||
);
|
||||
expect(metadataEvent.value.apiTimeMs).toBeGreaterThan(0);
|
||||
|
||||
// Also check the public getter
|
||||
expect(turn.getUsageMetadata()).toEqual(mockMetadata2);
|
||||
|
|
|
|||
|
|
@ -104,7 +104,7 @@ export type ServerGeminiChatCompressedEvent = {
|
|||
|
||||
export type ServerGeminiUsageMetadataEvent = {
|
||||
type: GeminiEventType.UsageMetadata;
|
||||
value: GenerateContentResponseUsageMetadata;
|
||||
value: GenerateContentResponseUsageMetadata & { apiTimeMs?: number };
|
||||
};
|
||||
|
||||
// The original union type, now composed of the individual types
|
||||
|
|
@ -137,6 +137,7 @@ export class Turn {
|
|||
req: PartListUnion,
|
||||
signal: AbortSignal,
|
||||
): AsyncGenerator<ServerGeminiStreamEvent> {
|
||||
const startTime = Date.now();
|
||||
try {
|
||||
const responseStream = await this.chat.sendMessageStream({
|
||||
message: req,
|
||||
|
|
@ -174,9 +175,10 @@ export class Turn {
|
|||
}
|
||||
|
||||
if (this.lastUsageMetadata) {
|
||||
const durationMs = Date.now() - startTime;
|
||||
yield {
|
||||
type: GeminiEventType.UsageMetadata,
|
||||
value: this.lastUsageMetadata,
|
||||
value: { ...this.lastUsageMetadata, apiTimeMs: durationMs },
|
||||
};
|
||||
}
|
||||
} catch (error) {
|
||||
|
|
|
|||
Loading…
Reference in New Issue