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: feat: Add exit UI w/ stats (#924)

---
 .../src/ui/components/HistoryItemDisplay.test.tsx  |  23 +++++
 .../cli/src/ui/components/HistoryItemDisplay.tsx   |   4 +
 .../ui/components/SessionSummaryDisplay.test.tsx   |  52 ++++++++++
 .../src/ui/components/SessionSummaryDisplay.tsx    |  75 ++++++++++++++
 packages/cli/src/ui/components/Stats.test.tsx      |  78 ++++++++++++++
 packages/cli/src/ui/components/Stats.tsx           | 114 +++++++++++++++++++++
 packages/cli/src/ui/components/StatsDisplay.tsx    |  91 ++--------------
 .../SessionSummaryDisplay.test.tsx.snap            |  45 ++++++++
 .../components/__snapshots__/Stats.test.tsx.snap   |  49 +++++++++
 9 files changed, 447 insertions(+), 84 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

(limited to 'packages/cli/src/ui/components')

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('<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!');
+  });
 });
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<HistoryItemDisplayProps> = ({
         duration={item.duration}
       />
     )}
+    {item.type === 'quit' && (
+      <SessionSummaryDisplay stats={item.stats} duration={item.duration} />
+    )}
     {item.type === 'tool_group' && (
       <ToolGroupMessage
         toolCalls={item.tools}
diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx
new file mode 100644
index 00000000..14d8a277
--- /dev/null
+++ b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx
@@ -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();
+  });
+});
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<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>
+  );
+};
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('<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();
+  });
+});
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 }) => (
+  <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>
+);
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 }) => (
-  <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>
 
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[`<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     │
+│                                 │
+╰─────────────────────────────────╯"
+`;
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[`<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"
+`;
-- 
cgit v1.2.3