feat: Add token stats in footer (#909)

2025-06-15 11:15:53 -07:00 · 2025-06-15 11:15:53 -07:00 · b3d89a1075
parent da09431be9
commit b3d89a1075
5 changed files with 160 additions and 57 deletions
--- a/packages/cli/src/ui/App.tsx
+++ b/packages/cli/src/ui/App.tsx
@ -53,7 +53,10 @@ import {
 } from '@gemini-cli/core';
 import { useLogger } from './hooks/useLogger.js';
 import { StreamingContext } from './contexts/StreamingContext.js';
-import { SessionStatsProvider } from './contexts/SessionContext.js';
+import {
+  SessionStatsProvider,
+  useSessionStats,
+} from './contexts/SessionContext.js';
 import { useGitBranchName } from './hooks/useGitBranchName.js';
 import { useTextBuffer } from './components/shared/text-buffer.js';
 import * as fs from 'fs';
@ -79,6 +82,7 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
    handleNewMessage,
    clearConsoleMessages: clearConsoleMessagesState,
  } = useConsoleMessages();
+  const { stats: sessionStats } = useSessionStats();
  const [staticNeedsRefresh, setStaticNeedsRefresh] = useState(false);
  const [staticKey, setStaticKey] = useState(0);
  const refreshStatic = useCallback(() => {
@ -648,6 +652,11 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
            showMemoryUsage={
              config.getDebugMode() || config.getShowMemoryUsage()
            }
+            promptTokenCount={sessionStats.currentResponse.promptTokenCount}
+            candidatesTokenCount={
+              sessionStats.currentResponse.candidatesTokenCount
+            }
+            totalTokenCount={sessionStats.currentResponse.totalTokenCount}
          />
        </Box>
      </Box>
--- a/packages/cli/src/ui/components/Footer.tsx
+++ b/packages/cli/src/ui/components/Footer.tsx
@ -7,7 +7,7 @@
 import React from 'react';
 import { Box, Text } from 'ink';
 import { Colors } from '../colors.js';
-import { shortenPath, tildeifyPath } from '@gemini-cli/core';
+import { shortenPath, tildeifyPath, tokenLimit } from '@gemini-cli/core';
 import { ConsoleSummaryDisplay } from './ConsoleSummaryDisplay.js';
 import process from 'node:process';
 import { MemoryUsageDisplay } from './MemoryUsageDisplay.js';
@ -22,6 +22,9 @@ interface FooterProps {
  errorCount: number;
  showErrorDetails: boolean;
  showMemoryUsage?: boolean;
+  promptTokenCount: number;
+  candidatesTokenCount: number;
+  totalTokenCount: number;
 }

 export const Footer: React.FC<FooterProps> = ({
@ -34,63 +37,75 @@ export const Footer: React.FC<FooterProps> = ({
  errorCount,
  showErrorDetails,
  showMemoryUsage,
-}) => (
-  <Box marginTop={1} justifyContent="space-between" width="100%">
-    <Box>
-      <Text color={Colors.LightBlue}>
-        {shortenPath(tildeifyPath(targetDir), 70)}
-        {branchName && <Text color={Colors.Gray}> ({branchName}*)</Text>}
-      </Text>
-      {debugMode && (
-        <Text color={Colors.AccentRed}>
-          {' ' + (debugMessage || '--debug')}
-        </Text>
-      )}
-    </Box>
+  totalTokenCount,
+}) => {
+  const limit = tokenLimit(model);
+  const percentage = totalTokenCount / limit;

-    {/* Middle Section: Centered Sandbox Info */}
-    <Box
-      flexGrow={1}
-      alignItems="center"
-      justifyContent="center"
-      display="flex"
-    >
-      {process.env.SANDBOX && process.env.SANDBOX !== 'sandbox-exec' ? (
-        <Text color="green">
-          {process.env.SANDBOX.replace(/^gemini-(?:cli-)?/, '')}
+  return (
+    <Box marginTop={1} justifyContent="space-between" width="100%">
+      <Box>
+        <Text color={Colors.LightBlue}>
+          {shortenPath(tildeifyPath(targetDir), 70)}
+          {branchName && <Text color={Colors.Gray}> ({branchName}*)</Text>}
        </Text>
-      ) : process.env.SANDBOX === 'sandbox-exec' ? (
-        <Text color={Colors.AccentYellow}>
-          MacOS Seatbelt{' '}
-          <Text color={Colors.Gray}>({process.env.SEATBELT_PROFILE})</Text>
-        </Text>
-      ) : (
-        <Text color={Colors.AccentRed}>
-          no sandbox <Text color={Colors.Gray}>(see docs)</Text>
-        </Text>
-      )}
-    </Box>
+        {debugMode && (
+          <Text color={Colors.AccentRed}>
+            {' ' + (debugMessage || '--debug')}
+          </Text>
+        )}
+      </Box>

-    {/* Right Section: Gemini Label and Console Summary */}
-    <Box alignItems="center">
-      <Text color={Colors.AccentBlue}> {model} </Text>
-      {corgiMode && (
-        <Text>
-          <Text color={Colors.Gray}>| </Text>
-          <Text color={Colors.AccentRed}>▼</Text>
-          <Text color={Colors.Foreground}>(´</Text>
-          <Text color={Colors.AccentRed}>ᴥ</Text>
-          <Text color={Colors.Foreground}>`)</Text>
-          <Text color={Colors.AccentRed}>▼ </Text>
+      {/* Middle Section: Centered Sandbox Info */}
+      <Box
+        flexGrow={1}
+        alignItems="center"
+        justifyContent="center"
+        display="flex"
+      >
+        {process.env.SANDBOX && process.env.SANDBOX !== 'sandbox-exec' ? (
+          <Text color="green">
+            {process.env.SANDBOX.replace(/^gemini-(?:cli-)?/, '')}
+          </Text>
+        ) : process.env.SANDBOX === 'sandbox-exec' ? (
+          <Text color={Colors.AccentYellow}>
+            MacOS Seatbelt{' '}
+            <Text color={Colors.Gray}>({process.env.SEATBELT_PROFILE})</Text>
+          </Text>
+        ) : (
+          <Text color={Colors.AccentRed}>
+            no sandbox <Text color={Colors.Gray}>(see docs)</Text>
+          </Text>
+        )}
+      </Box>
+
+      {/* Right Section: Gemini Label and Console Summary */}
+      <Box alignItems="center">
+        <Text color={Colors.AccentBlue}>
+          {' '}
+          {model}{' '}
+          <Text color={Colors.Gray}>
+            ({((1 - percentage) * 100).toFixed(0)}% context left)
+          </Text>
        </Text>
-      )}
-      {!showErrorDetails && errorCount > 0 && (
-        <Box>
-          <Text color={Colors.Gray}>| </Text>
-          <ConsoleSummaryDisplay errorCount={errorCount} />
-        </Box>
-      )}
-      {showMemoryUsage && <MemoryUsageDisplay />}
+        {corgiMode && (
+          <Text>
+            <Text color={Colors.Gray}>| </Text>
+            <Text color={Colors.AccentRed}>▼</Text>
+            <Text color={Colors.Foreground}>(´</Text>
+            <Text color={Colors.AccentRed}>ᴥ</Text>
+            <Text color={Colors.Foreground}>`)</Text>
+            <Text color={Colors.AccentRed}>▼ </Text>
+          </Text>
+        )}
+        {!showErrorDetails && errorCount > 0 && (
+          <Box>
+            <Text color={Colors.Gray}>| </Text>
+            <ConsoleSummaryDisplay errorCount={errorCount} />
+          </Box>
+        )}
+        {showMemoryUsage && <MemoryUsageDisplay />}
+      </Box>
    </Box>
-  </Box>
-);
+  );
+};
--- a/packages/cli/src/ui/contexts/SessionContext.test.tsx
+++ b/packages/cli/src/ui/contexts/SessionContext.test.tsx
@ -177,6 +177,51 @@ describe('SessionStatsContext', () => {
    expect(stats?.currentTurn.apiTimeMs).toBe(100 + 50);
  });

+  it('should overwrite currentResponse with each API call', () => {
+    const contextRef: MutableRefObject<
+      ReturnType<typeof useSessionStats> | undefined
+    > = { current: undefined };
+
+    render(
+      <SessionStatsProvider>
+        <TestHarness contextRef={contextRef} />
+      </SessionStatsProvider>,
+    );
+
+    // 1. First API call
+    act(() => {
+      contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
+    });
+
+    let stats = contextRef.current?.stats;
+
+    // currentResponse should match the first call
+    expect(stats?.currentResponse.totalTokenCount).toBe(300);
+    expect(stats?.currentResponse.apiTimeMs).toBe(100);
+
+    // 2. Second API call
+    act(() => {
+      contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
+    });
+
+    stats = contextRef.current?.stats;
+
+    // currentResponse should now match the second call
+    expect(stats?.currentResponse.totalTokenCount).toBe(30);
+    expect(stats?.currentResponse.apiTimeMs).toBe(50);
+
+    // 3. Start a new turn
+    act(() => {
+      contextRef.current?.startNewTurn();
+    });
+
+    stats = contextRef.current?.stats;
+
+    // currentResponse should be reset
+    expect(stats?.currentResponse.totalTokenCount).toBe(0);
+    expect(stats?.currentResponse.apiTimeMs).toBe(0);
+  });
+
  it('should throw an error when useSessionStats is used outside of a provider', () => {
    // Suppress the expected console error during this test.
    const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
--- a/packages/cli/src/ui/contexts/SessionContext.tsx
+++ b/packages/cli/src/ui/contexts/SessionContext.tsx
@ -31,6 +31,7 @@ interface SessionStatsState {
  sessionStartTime: Date;
  cumulative: CumulativeStats;
  currentTurn: CumulativeStats;
+  currentResponse: CumulativeStats;
 }

 // Defines the final "value" of our context, including the state
@ -97,6 +98,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
      thoughtsTokenCount: 0,
      apiTimeMs: 0,
    },
+    currentResponse: {
+      turnCount: 0,
+      promptTokenCount: 0,
+      candidatesTokenCount: 0,
+      totalTokenCount: 0,
+      cachedContentTokenCount: 0,
+      toolUsePromptTokenCount: 0,
+      thoughtsTokenCount: 0,
+      apiTimeMs: 0,
+    },
  });

  // A single, internal worker function to handle all metadata aggregation.
@ -107,15 +118,27 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
      setStats((prevState) => {
        const newCumulative = { ...prevState.cumulative };
        const newCurrentTurn = { ...prevState.currentTurn };
+        const newCurrentResponse = {
+          turnCount: 0,
+          promptTokenCount: 0,
+          candidatesTokenCount: 0,
+          totalTokenCount: 0,
+          cachedContentTokenCount: 0,
+          toolUsePromptTokenCount: 0,
+          thoughtsTokenCount: 0,
+          apiTimeMs: 0,
+        };

        // Add all tokens to the current turn's stats as well as cumulative stats.
        addTokens(newCurrentTurn, metadata);
        addTokens(newCumulative, metadata);
+        addTokens(newCurrentResponse, metadata);

        return {
          ...prevState,
          cumulative: newCumulative,
          currentTurn: newCurrentTurn,
+          currentResponse: newCurrentResponse,
        };
      });
    },
@ -139,6 +162,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
        thoughtsTokenCount: 0,
        apiTimeMs: 0,
      },
+      currentResponse: {
+        turnCount: 0,
+        promptTokenCount: 0,
+        candidatesTokenCount: 0,
+        totalTokenCount: 0,
+        cachedContentTokenCount: 0,
+        toolUsePromptTokenCount: 0,
+        thoughtsTokenCount: 0,
+        apiTimeMs: 0,
+      },
    }));
  }, []);

--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@ -13,6 +13,7 @@ export * from './core/contentGenerator.js';
 export * from './core/geminiChat.js';
 export * from './core/logger.js';
 export * from './core/prompts.js';
+export * from './core/tokenLimits.js';
 export * from './core/turn.js';
 export * from './core/geminiRequest.js';
 export * from './core/coreToolScheduler.js';