From 7a72d255d8effec1396170306cc6be57f598a6d8 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Wed, 11 Jun 2025 16:40:31 -0400 Subject: [PATCH] feat: Add exit UI w/ stats (#924) --- .../ui/components/HistoryItemDisplay.test.tsx | 23 ++++ .../src/ui/components/HistoryItemDisplay.tsx | 4 + .../components/SessionSummaryDisplay.test.tsx | 52 ++++++++ .../ui/components/SessionSummaryDisplay.tsx | 75 ++++++++++++ packages/cli/src/ui/components/Stats.test.tsx | 78 ++++++++++++ packages/cli/src/ui/components/Stats.tsx | 114 ++++++++++++++++++ .../cli/src/ui/components/StatsDisplay.tsx | 91 ++------------ .../SessionSummaryDisplay.test.tsx.snap | 45 +++++++ .../__snapshots__/Stats.test.tsx.snap | 49 ++++++++ .../ui/hooks/slashCommandProcessor.test.ts | 37 ++++++ .../cli/src/ui/hooks/slashCommandProcessor.ts | 23 +++- .../cli/src/ui/hooks/useReactToolScheduler.ts | 2 +- packages/cli/src/ui/types.ts | 17 ++- 13 files changed, 522 insertions(+), 88 deletions(-) create mode 100644 packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx create mode 100644 packages/cli/src/ui/components/SessionSummaryDisplay.tsx create mode 100644 packages/cli/src/ui/components/Stats.test.tsx create mode 100644 packages/cli/src/ui/components/Stats.tsx create mode 100644 packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap create mode 100644 packages/cli/src/ui/components/__snapshots__/Stats.test.tsx.snap diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx index 0fe739df..5999f0ad 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx @@ -73,4 +73,27 @@ describe('', () => { ); 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( + , + ); + expect(lastFrame()).toContain('Agent powering down. Goodbye!'); + }); }); diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.tsx index 8c4fede9..229672ec 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.tsx @@ -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 = ({ duration={item.duration} /> )} + {item.type === 'quit' && ( + + )} {item.type === 'tool_group' && ( ', () => { + 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( + , + ); + + 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( + , + ); + + expect(lastFrame()).toMatchSnapshot(); + }); +}); diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.tsx new file mode 100644 index 00000000..d3ee0f5f --- /dev/null +++ b/packages/cli/src/ui/components/SessionSummaryDisplay.tsx @@ -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 = ({ + 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 ( + + + {Colors.GradientColors ? ( + + {title} + + ) : ( + {title} + )} + + + + + + + + + + + + ); +}; diff --git a/packages/cli/src/ui/components/Stats.test.tsx b/packages/cli/src/ui/components/Stats.test.tsx new file mode 100644 index 00000000..1436d485 --- /dev/null +++ b/packages/cli/src/ui/components/Stats.test.tsx @@ -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('', () => { + it('renders a label and value', () => { + const { lastFrame } = render( + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); + + it('renders with a specific value color', () => { + const { lastFrame } = render( + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); +}); + +describe('', () => { + 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( + + + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); + + it('renders a stats column with a specific width', () => { + const { lastFrame } = render( + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); + + it('renders a cumulative stats column with percentages', () => { + const { lastFrame } = render( + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); +}); + +describe('', () => { + it('renders a duration column', () => { + const { lastFrame } = render( + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); +}); diff --git a/packages/cli/src/ui/components/Stats.tsx b/packages/cli/src/ui/components/Stats.tsx new file mode 100644 index 00000000..92fadd11 --- /dev/null +++ b/packages/cli/src/ui/components/Stats.tsx @@ -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 }) => ( + + {label} + {value} + +); + +/** + * 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 ( + + {title} + + {/* All StatRows below will now inherit the gap */} + + + + + + {/* Divider Line */} + + + {children} + + + ); +}; + +/** + * Renders a column for displaying duration information. + */ +export const DurationColumn: React.FC<{ + apiTime: string; + wallTime: string; +}> = ({ apiTime, wallTime }) => ( + + Duration + + + + + +); diff --git a/packages/cli/src/ui/components/StatsDisplay.tsx b/packages/cli/src/ui/components/StatsDisplay.tsx index be447595..76d48821 100644 --- a/packages/cli/src/ui/components/StatsDisplay.tsx +++ b/packages/cli/src/ui/components/StatsDisplay.tsx @@ -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 }) => ( - - {label} - {value} - -); - -/** - * 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 ( - - {title} - - - - - - - {/* Divider Line */} - - - - - ); -}; - // --- Main Component --- export const StatsDisplay: React.FC = ({ @@ -143,11 +61,16 @@ export const StatsDisplay: React.FC = ({ - + diff --git a/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap new file mode 100644 index 00000000..74b067b7 --- /dev/null +++ b/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap @@ -0,0 +1,45 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[` > 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[` > 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 │ +│ │ +╰─────────────────────────────────╯" +`; diff --git a/packages/cli/src/ui/components/__snapshots__/Stats.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/Stats.test.tsx.snap new file mode 100644 index 00000000..9b003891 --- /dev/null +++ b/packages/cli/src/ui/components/__snapshots__/Stats.test.tsx.snap @@ -0,0 +1,49 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[` > renders a duration column 1`] = ` +"Duration + +API Time 5s +Wall Time 10s" +`; + +exports[` > renders a label and value 1`] = `"Test Label Test Value"`; + +exports[` > renders with a specific value color 1`] = `"Test Label Test Value"`; + +exports[` > 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[` > 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[` > 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" +`; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index f16d3239..0c12d855 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -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(); diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 3699b4e9..8e2f2bd2 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -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( diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index 4e55cba4..8ae7ebfb 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -132,7 +132,7 @@ export function useReactToolScheduler( }); onComplete(completedToolCalls); }, - [onComplete], + [onComplete, config], ); const toolCallsUpdateHandler: ToolCallsUpdateHandler = useCallback( diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 5fae1568..728b3476 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -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 {