feat: Add exit UI w/ stats (#924)

This commit is contained in:
Abhi 2025-06-11 16:40:31 -04:00 committed by GitHub
parent 4160d904da
commit 7a72d255d8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 522 additions and 88 deletions

View File

@ -73,4 +73,27 @@ describe('<HistoryItemDisplay />', () => {
);
expect(lastFrame()).toContain('About Gemini CLI');
});
it('renders SessionSummaryDisplay for "quit" 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: 'quit',
stats,
duration: '1s',
};
const { lastFrame } = render(
<HistoryItemDisplay {...baseItem} item={item} />,
);
expect(lastFrame()).toContain('Agent powering down. Goodbye!');
});
});

View File

@ -16,6 +16,7 @@ import { GeminiMessageContent } from './messages/GeminiMessageContent.js';
import { Box } from 'ink';
import { AboutBox } from './AboutBox.js';
import { StatsDisplay } from './StatsDisplay.js';
import { SessionSummaryDisplay } from './SessionSummaryDisplay.js';
import { Config } from '@gemini-cli/core';
interface HistoryItemDisplayProps {
@ -66,6 +67,9 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
duration={item.duration}
/>
)}
{item.type === 'quit' && (
<SessionSummaryDisplay stats={item.stats} duration={item.duration} />
)}
{item.type === 'tool_group' && (
<ToolGroupMessage
toolCalls={item.tools}

View File

@ -0,0 +1,52 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from 'ink-testing-library';
import { describe, it, expect } from 'vitest';
import { SessionSummaryDisplay } from './SessionSummaryDisplay.js';
import { type CumulativeStats } from '../contexts/SessionContext.js';
describe('<SessionSummaryDisplay />', () => {
const mockStats: CumulativeStats = {
turnCount: 10,
promptTokenCount: 1000,
candidatesTokenCount: 2000,
totalTokenCount: 3500,
cachedContentTokenCount: 500,
toolUsePromptTokenCount: 200,
thoughtsTokenCount: 300,
apiTimeMs: 50234,
};
const mockDuration = '1h 23m 45s';
it('renders correctly with given stats and duration', () => {
const { lastFrame } = render(
<SessionSummaryDisplay stats={mockStats} 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(
<SessionSummaryDisplay stats={zeroStats} duration="0s" />,
);
expect(lastFrame()).toMatchSnapshot();
});
});

View File

@ -0,0 +1,75 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import React from 'react';
import { Box, Text } from 'ink';
import Gradient from 'ink-gradient';
import { Colors } from '../colors.js';
import { formatDuration } from '../utils/formatters.js';
import { CumulativeStats } from '../contexts/SessionContext.js';
import { FormattedStats, StatRow, StatsColumn } from './Stats.js';
// --- Prop and Data Structures ---
interface SessionSummaryDisplayProps {
stats: CumulativeStats;
duration: string;
}
// --- Main Component ---
export const SessionSummaryDisplay: React.FC<SessionSummaryDisplayProps> = ({
stats,
duration,
}) => {
const cumulativeFormatted: FormattedStats = {
inputTokens: stats.promptTokenCount,
outputTokens: stats.candidatesTokenCount,
toolUseTokens: stats.toolUsePromptTokenCount,
thoughtsTokens: stats.thoughtsTokenCount,
cachedTokens: stats.cachedContentTokenCount,
totalTokens: stats.totalTokenCount,
};
const title = 'Agent powering down. Goodbye!';
return (
<Box
borderStyle="round"
borderColor="gray"
flexDirection="column"
paddingY={1}
paddingX={2}
alignSelf="flex-start"
>
<Box marginBottom={1} flexDirection="column">
{Colors.GradientColors ? (
<Gradient colors={Colors.GradientColors}>
<Text bold>{title}</Text>
</Gradient>
) : (
<Text bold>{title}</Text>
)}
</Box>
<Box marginTop={1}>
<StatsColumn
title={`Cumulative Stats (${stats.turnCount} Turns)`}
stats={cumulativeFormatted}
isCumulative={true}
>
<Box marginTop={1} flexDirection="column">
<StatRow
label="Total duration (API)"
value={formatDuration(stats.apiTimeMs)}
/>
<StatRow label="Total duration (wall)" value={duration} />
</Box>
</StatsColumn>
</Box>
</Box>
);
};

View File

@ -0,0 +1,78 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from 'ink-testing-library';
import { describe, it, expect } from 'vitest';
import {
StatRow,
StatsColumn,
DurationColumn,
FormattedStats,
} from './Stats.js';
import { Colors } from '../colors.js';
describe('<StatRow />', () => {
it('renders a label and value', () => {
const { lastFrame } = render(
<StatRow label="Test Label" value="Test Value" />,
);
expect(lastFrame()).toMatchSnapshot();
});
it('renders with a specific value color', () => {
const { lastFrame } = render(
<StatRow
label="Test Label"
value="Test Value"
valueColor={Colors.AccentGreen}
/>,
);
expect(lastFrame()).toMatchSnapshot();
});
});
describe('<StatsColumn />', () => {
const mockStats: FormattedStats = {
inputTokens: 100,
outputTokens: 200,
toolUseTokens: 50,
thoughtsTokens: 25,
cachedTokens: 10,
totalTokens: 385,
};
it('renders a stats column with children', () => {
const { lastFrame } = render(
<StatsColumn title="Test Stats" stats={mockStats}>
<StatRow label="Child Prop" value="Child Value" />
</StatsColumn>,
);
expect(lastFrame()).toMatchSnapshot();
});
it('renders a stats column with a specific width', () => {
const { lastFrame } = render(
<StatsColumn title="Test Stats" stats={mockStats} width="50%" />,
);
expect(lastFrame()).toMatchSnapshot();
});
it('renders a cumulative stats column with percentages', () => {
const { lastFrame } = render(
<StatsColumn title="Cumulative Stats" stats={mockStats} isCumulative />,
);
expect(lastFrame()).toMatchSnapshot();
});
});
describe('<DurationColumn />', () => {
it('renders a duration column', () => {
const { lastFrame } = render(
<DurationColumn apiTime="5s" wallTime="10s" />,
);
expect(lastFrame()).toMatchSnapshot();
});
});

View File

@ -0,0 +1,114 @@
/**
* @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';
// --- Prop and Data Structures ---
export 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.
*/
export const StatRow: React.FC<{
label: string;
value: string | number;
valueColor?: string;
}> = ({ label, value, valueColor }) => (
<Box justifyContent="space-between" gap={2}>
<Text color={Colors.LightBlue}>{label}</Text>
<Text color={valueColor}>{value}</Text>
</Box>
);
/**
* Renders a full column for either "Last Turn" or "Cumulative" stats.
*/
export const StatsColumn: React.FC<{
title: string;
stats: FormattedStats;
isCumulative?: boolean;
width?: string | number;
children?: React.ReactNode;
}> = ({ title, stats, isCumulative = false, width, children }) => {
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={width}>
<Text bold>{title}</Text>
<Box marginTop={1} flexDirection="column">
{/* All StatRows below will now inherit the gap */}
<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()}
/>
{children}
</Box>
</Box>
);
};
/**
* Renders a column for displaying duration information.
*/
export const DurationColumn: React.FC<{
apiTime: string;
wallTime: string;
}> = ({ apiTime, wallTime }) => (
<Box flexDirection="column" width={'48%'}>
<Text bold>Duration</Text>
<Box marginTop={1} flexDirection="column">
<StatRow label="API Time" value={apiTime} />
<StatRow label="Wall Time" value={wallTime} />
</Box>
</Box>
);

View File

@ -9,6 +9,7 @@ import { Box, Text } from 'ink';
import { Colors } from '../colors.js';
import { formatDuration } from '../utils/formatters.js';
import { CumulativeStats } from '../contexts/SessionContext.js';
import { FormattedStats, StatRow, StatsColumn } from './Stats.js';
// --- Constants ---
@ -22,89 +23,6 @@ interface StatsDisplayProps {
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> = ({
@ -143,11 +61,16 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
</Text>
<Box flexDirection="row" justifyContent="space-between" marginTop={1}>
<StatsColumn title="Last Turn" stats={lastTurnFormatted} />
<StatsColumn
title="Last Turn"
stats={lastTurnFormatted}
width={COLUMN_WIDTH}
/>
<StatsColumn
title={`Cumulative (${stats.turnCount} Turns)`}
stats={cumulativeFormatted}
isCumulative={true}
width={COLUMN_WIDTH}
/>
</Box>

View File

@ -0,0 +1,45 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`<SessionSummaryDisplay /> > renders correctly with given stats and duration 1`] = `
"╭─────────────────────────────────────╮
│ │
│ Agent powering down. Goodbye! │
│ │
│ │
│ Cumulative Stats (10 Turns) │
│ │
│ Input Tokens 1,000 │
│ Output Tokens 2,000 │
│ Tool Use Tokens 200 │
│ Thoughts Tokens 300 │
│ Cached Tokens 500 (14.3%) │
│ ───────────────────────────────── │
│ Total Tokens 3,500 │
│ │
│ Total duration (API) 50.2s │
│ Total duration (wall) 1h 23m 45s │
│ │
╰─────────────────────────────────────╯"
`;
exports[`<SessionSummaryDisplay /> > renders zero state correctly 1`] = `
"╭─────────────────────────────────╮
│ │
│ Agent powering down. Goodbye! │
│ │
│ │
│ Cumulative Stats (0 Turns) │
│ │
│ Input Tokens 0 │
│ Output Tokens 0 │
│ Tool Use Tokens 0 │
│ Thoughts Tokens 0 │
│ Cached Tokens 0 │
│ ────────────────────────── │
│ Total Tokens 0 │
│ │
│ Total duration (API) 0s │
│ Total duration (wall) 0s │
│ │
╰─────────────────────────────────╯"
`;

View File

@ -0,0 +1,49 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`<DurationColumn /> > renders a duration column 1`] = `
"Duration
API Time 5s
Wall Time 10s"
`;
exports[`<StatRow /> > renders a label and value 1`] = `"Test Label Test Value"`;
exports[`<StatRow /> > renders with a specific value color 1`] = `"Test Label Test Value"`;
exports[`<StatsColumn /> > renders a cumulative stats column with percentages 1`] = `
"Cumulative Stats
Input Tokens 100
Output Tokens 200
Tool Use Tokens 50
Thoughts Tokens 25
Cached Tokens 10 (2.6%)
────────────────────────────────────────────────────────────────────────────────────────────────────
Total Tokens 385"
`;
exports[`<StatsColumn /> > renders a stats column with a specific width 1`] = `
"Test Stats
Input Tokens 100
Output Tokens 200
Tool Use Tokens 50
Thoughts Tokens 25
Cached Tokens 10
──────────────────────────────────────────────────
Total Tokens 385"
`;
exports[`<StatsColumn /> > renders a stats column with children 1`] = `
"Test Stats
Input Tokens 100
Output Tokens 200
Tool Use Tokens 50
Thoughts Tokens 25
Cached Tokens 10
────────────────────────────────────────────────────────────────────────────────────────────────────
Total Tokens 385
Child Prop Child Value"
`;

View File

@ -396,6 +396,43 @@ Add any other context about the problem here.
});
});
describe('/quit and /exit commands', () => {
beforeEach(() => {
vi.useFakeTimers();
});
afterEach(() => {
vi.useRealTimers();
});
it.each([['/quit'], ['/exit']])(
'should handle %s, add a quit message, and exit the process',
async (command) => {
const { handleSlashCommand } = getProcessor();
const mockDate = new Date('2025-01-01T01:02:03.000Z');
vi.setSystemTime(mockDate);
await act(async () => {
handleSlashCommand(command);
});
expect(mockAddItem).toHaveBeenCalledTimes(2);
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
type: MessageType.QUIT,
duration: '1h 2m 3s',
}),
expect.any(Number),
);
// Fast-forward timers to trigger process.exit
vi.advanceTimersByTime(100);
expect(mockProcessExit).toHaveBeenCalledWith(0);
},
);
});
describe('Unknown command', () => {
it('should show an error and return true for a general unknown command', async () => {
const { handleSlashCommand } = getProcessor();

View File

@ -97,6 +97,12 @@ export const useSlashCommandProcessor = (
lastTurnStats: message.lastTurnStats,
duration: message.duration,
};
} else if (message.type === MessageType.QUIT) {
historyItemContent = {
type: 'quit',
stats: message.stats,
duration: message.duration,
};
} else {
historyItemContent = {
type: message.type as
@ -594,8 +600,20 @@ Add any other context about the problem here.
altName: 'exit',
description: 'exit the cli',
action: async (_mainCommand, _subCommand, _args) => {
onDebugMessage('Quitting. Good-bye.');
process.exit(0);
const now = new Date();
const { sessionStartTime, cumulative } = session.stats;
const wallDuration = now.getTime() - sessionStartTime.getTime();
addMessage({
type: MessageType.QUIT,
stats: cumulative,
duration: formatDuration(wallDuration),
timestamp: new Date(),
});
setTimeout(() => {
process.exit(0);
}, 100);
},
},
];
@ -721,6 +739,7 @@ Add any other context about the problem here.
session,
gitService,
loadHistory,
addItem,
]);
const handleSlashCommand = useCallback(

View File

@ -132,7 +132,7 @@ export function useReactToolScheduler(
});
onComplete(completedToolCalls);
},
[onComplete],
[onComplete, config],
);
const toolCallsUpdateHandler: ToolCallsUpdateHandler = useCallback(

View File

@ -97,6 +97,12 @@ export type HistoryItemStats = HistoryItemBase & {
duration: string;
};
export type HistoryItemQuit = HistoryItemBase & {
type: 'quit';
stats: CumulativeStats;
duration: string;
};
export type HistoryItemToolGroup = HistoryItemBase & {
type: 'tool_group';
tools: IndividualToolCallDisplay[];
@ -120,7 +126,8 @@ export type HistoryItemWithoutId =
| HistoryItemError
| HistoryItemAbout
| HistoryItemToolGroup
| HistoryItemStats;
| HistoryItemStats
| HistoryItemQuit;
export type HistoryItem = HistoryItemWithoutId & { id: number };
@ -131,6 +138,7 @@ export enum MessageType {
USER = 'user',
ABOUT = 'about',
STATS = 'stats',
QUIT = 'quit',
GEMINI = 'gemini',
}
@ -157,6 +165,13 @@ export type Message =
lastTurnStats: CumulativeStats;
duration: string;
content?: string;
}
| {
type: MessageType.QUIT;
timestamp: Date;
stats: CumulativeStats;
duration: string;
content?: string;
};
export interface ConsoleMessageItem {