fix(patch): cherry-pick 54de675 to release/v0.21.0-preview.2-pr-14961 [CONFLICTS] (#14968)

Co-authored-by: Jacob Richman <jacob314@gmail.com>

packages/cli/src/ui/components/StatsDisplay.tsx

@@ -17,9 +17,24 @@ import {
   TOOL_SUCCESS_RATE_MEDIUM,
   USER_AGREEMENT_RATE_HIGH,
   USER_AGREEMENT_RATE_MEDIUM,
+  CACHE_EFFICIENCY_HIGH,
+  CACHE_EFFICIENCY_MEDIUM,
 } from '../utils/displayUtils.js';
 import { computeSessionStats } from '../utils/computeStats.js';
 import type { RetrieveUserQuotaResponse } from '@google/gemini-cli-core';
+import {
+  DEFAULT_GEMINI_FLASH_LITE_MODEL,
+  DEFAULT_GEMINI_FLASH_MODEL,
+  DEFAULT_GEMINI_MODEL,
+  PREVIEW_GEMINI_MODEL,
+} from '@google/gemini-cli-core';
+
+const VALID_GEMINI_MODELS = new Set([
+  PREVIEW_GEMINI_MODEL,
+  DEFAULT_GEMINI_MODEL,
+  DEFAULT_GEMINI_FLASH_MODEL,
+  DEFAULT_GEMINI_FLASH_LITE_MODEL,
+]);
 
 // A more flexible and powerful StatRow component
 interface StatRowProps {
@@ -33,8 +48,7 @@ const StatRow: React.FC<StatRowProps> = ({ title, children }) => (
     <Box width={28}>
       <Text color={theme.text.link}>{title}</Text>
     </Box>
-    {/* FIX: Wrap children in a Box that can grow to fill remaining space */}
-    <Box flexGrow={1}>{children}</Box>
+    {children}
   </Box>
 );
 
@@ -50,8 +64,7 @@ const SubStatRow: React.FC<SubStatRowProps> = ({ title, children }) => (
     <Box width={26}>
       <Text color={theme.text.secondary}>» {title}</Text>
     </Box>
-    {/* FIX: Apply the same flexGrow fix here */}
-    <Box flexGrow={1}>{children}</Box>
+    {children}
   </Box>
 );
 
@@ -70,6 +83,57 @@ const Section: React.FC<SectionProps> = ({ title, children }) => (
   </Box>
 );
 
+// Logic for building the unified list of table rows
+const buildModelRows = (
+  models: Record<string, ModelMetrics>,
+  quotas?: RetrieveUserQuotaResponse,
+) => {
+  const getBaseModelName = (name: string) => name.replace('-001', '');
+  const usedModelNames = new Set(Object.keys(models).map(getBaseModelName));
+
+  // 1. Models with active usage
+  const activeRows = Object.entries(models).map(([name, metrics]) => {
+    const modelName = getBaseModelName(name);
+    const cachedTokens = metrics.tokens.cached;
+    const totalInputTokens = metrics.tokens.prompt;
+    const uncachedTokens = Math.max(0, totalInputTokens - cachedTokens);
+    return {
+      key: name,
+      modelName,
+      requests: metrics.api.totalRequests,
+      cachedTokens: cachedTokens.toLocaleString(),
+      uncachedTokens: uncachedTokens.toLocaleString(),
+      totalInputTokens: totalInputTokens.toLocaleString(),
+      outputTokens: metrics.tokens.candidates.toLocaleString(),
+      bucket: quotas?.buckets?.find((b) => b.modelId === modelName),
+      isActive: true,
+    };
+  });
+
+  // 2. Models with quota only
+  const quotaRows =
+    quotas?.buckets
+      ?.filter(
+        (b) =>
+          b.modelId &&
+          VALID_GEMINI_MODELS.has(b.modelId) &&
+          !usedModelNames.has(b.modelId),
+      )
+      .map((bucket) => ({
+        key: bucket.modelId!,
+        modelName: bucket.modelId!,
+        requests: '-',
+        cachedTokens: '-',
+        uncachedTokens: '-',
+        totalInputTokens: '-',
+        outputTokens: '-',
+        bucket,
+        isActive: false,
+      })) || [];
+
+  return [...activeRows, ...quotaRows];
+};
+
 const formatResetTime = (resetTime: string): string => {
   const diff = new Date(resetTime).getTime() - Date.now();
   if (diff <= 0) return '';
@@ -96,44 +160,91 @@ const formatResetTime = (resetTime: string): string => {
 
 const ModelUsageTable: React.FC<{
   models: Record<string, ModelMetrics>;
-  totalCachedTokens: number;
-  cacheEfficiency: number;
   quotas?: RetrieveUserQuotaResponse;
-}> = ({ models, totalCachedTokens, cacheEfficiency, quotas }) => {
+  cacheEfficiency: number;
+  totalCachedTokens: number;
+}> = ({ models, quotas, cacheEfficiency, totalCachedTokens }) => {
+  const rows = buildModelRows(models, quotas);
+
+  if (rows.length === 0) {
+    return null;
+  }
+
+  const showQuotaColumn = !!quotas && rows.some((row) => !!row.bucket);
+
   const nameWidth = 25;
-  const requestsWidth = 8;
-  const inputTokensWidth = 15;
+  const requestsWidth = 7;
+  const uncachedWidth = 15;
+  const cachedWidth = 14;
   const outputTokensWidth = 15;
-  const usageLimitWidth = quotas ? 30 : 0;
+  const usageLimitWidth = showQuotaColumn ? 28 : 0;
+
+  const cacheEfficiencyColor = getStatusColor(cacheEfficiency, {
+    green: CACHE_EFFICIENCY_HIGH,
+    yellow: CACHE_EFFICIENCY_MEDIUM,
+  });
 
   return (
     <Box flexDirection="column" marginTop={1}>
       {/* Header */}
-      <Box>
-        <Box width={nameWidth}>
-          <Text bold color={theme.text.primary}>
+      <Box alignItems="flex-end">
+        <Box width={nameWidth} flexGrow={1}>
+          <Text bold color={theme.text.primary} wrap="truncate-end">
             Model Usage
           </Text>
         </Box>
-        <Box width={requestsWidth} justifyContent="flex-end">
+        <Box
+          width={requestsWidth}
+          flexDirection="column"
+          alignItems="flex-end"
+          flexShrink={0}
+        >
           <Text bold color={theme.text.primary}>
             Reqs
           </Text>
         </Box>
-        <Box width={inputTokensWidth} justifyContent="flex-end">
-          <Text bold color={theme.text.primary}>
-            Input Tokens
-          </Text>
-        </Box>
-        <Box width={outputTokensWidth} justifyContent="flex-end">
-          <Text bold color={theme.text.primary}>
-            Output Tokens
-          </Text>
-        </Box>
-        {quotas && (
-          <Box width={usageLimitWidth} justifyContent="flex-end">
+        {!showQuotaColumn && (
+          <>
+            <Box
+              width={uncachedWidth}
+              flexDirection="column"
+              alignItems="flex-end"
+              flexShrink={0}
+            >
+              <Text bold color={theme.text.primary}>
+                Input Tokens
+              </Text>
+            </Box>
+            <Box
+              width={cachedWidth}
+              flexDirection="column"
+              alignItems="flex-end"
+              flexShrink={0}
+            >
+              <Text bold color={theme.text.primary}>
+                Cache Reads
+              </Text>
+            </Box>
+            <Box
+              width={outputTokensWidth}
+              flexDirection="column"
+              alignItems="flex-end"
+              flexShrink={0}
+            >
+              <Text bold color={theme.text.primary}>
+                Output Tokens
+              </Text>
+            </Box>
+          </>
+        )}
+        {showQuotaColumn && (
+          <Box
+            width={usageLimitWidth}
+            flexDirection="column"
+            alignItems="flex-end"
+          >
             <Text bold color={theme.text.primary}>
-              Usage limit remaining
+              Usage left
             </Text>
           </Box>
         )}
@@ -146,63 +257,99 @@ const ModelUsageTable: React.FC<{
         borderLeft={false}
         borderRight={false}
         borderColor={theme.border.default}
-        width={
-          nameWidth +
-          requestsWidth +
-          inputTokensWidth +
-          outputTokensWidth +
-          usageLimitWidth
-        }
+        width="100%"
       ></Box>
 
-      {/* Rows */}
-      {Object.entries(models).map(([name, modelMetrics]) => {
-        const modelName = name.replace('-001', '');
-        const bucket = quotas?.buckets?.find((b) => b.modelId === modelName);
-
-        return (
-          <Box key={name}>
-            <Box width={nameWidth}>
-              <Text color={theme.text.primary}>{modelName}</Text>
-            </Box>
-            <Box width={requestsWidth} justifyContent="flex-end">
-              <Text color={theme.text.primary}>
-                {modelMetrics.api.totalRequests}
-              </Text>
-            </Box>
-            <Box width={inputTokensWidth} justifyContent="flex-end">
-              <Text color={theme.status.warning}>
-                {modelMetrics.tokens.prompt.toLocaleString()}
-              </Text>
-            </Box>
-            <Box width={outputTokensWidth} justifyContent="flex-end">
-              <Text color={theme.status.warning}>
-                {modelMetrics.tokens.candidates.toLocaleString()}
-              </Text>
-            </Box>
-            <Box width={usageLimitWidth} justifyContent="flex-end">
-              {bucket &&
-                bucket.remainingFraction != null &&
-                bucket.resetTime && (
-                  <Text color={theme.text.secondary}>
-                    {(bucket.remainingFraction * 100).toFixed(1)}%{' '}
-                    {formatResetTime(bucket.resetTime)}
-                  </Text>
-                )}
-            </Box>
+      {rows.map((row) => (
+        <Box key={row.key}>
+          <Box width={nameWidth} flexGrow={1}>
+            <Text color={theme.text.primary} wrap="truncate-end">
+              {row.modelName}
+            </Text>
           </Box>
-        );
-      })}
-      {cacheEfficiency > 0 && (
+          <Box
+            width={requestsWidth}
+            flexDirection="column"
+            alignItems="flex-end"
+            flexShrink={0}
+          >
+            <Text
+              color={row.isActive ? theme.text.primary : theme.text.secondary}
+            >
+              {row.requests}
+            </Text>
+          </Box>
+          {!showQuotaColumn && (
+            <>
+              <Box
+                width={uncachedWidth}
+                flexDirection="column"
+                alignItems="flex-end"
+                flexShrink={0}
+              >
+                <Text
+                  color={
+                    row.isActive ? theme.text.primary : theme.text.secondary
+                  }
+                >
+                  {row.uncachedTokens}
+                </Text>
+              </Box>
+              <Box
+                width={cachedWidth}
+                flexDirection="column"
+                alignItems="flex-end"
+                flexShrink={0}
+              >
+                <Text color={theme.text.secondary}>{row.cachedTokens}</Text>
+              </Box>
+              <Box
+                width={outputTokensWidth}
+                flexDirection="column"
+                alignItems="flex-end"
+                flexShrink={0}
+              >
+                <Text
+                  color={
+                    row.isActive ? theme.text.primary : theme.text.secondary
+                  }
+                >
+                  {row.outputTokens}
+                </Text>
+              </Box>
+            </>
+          )}
+          <Box
+            width={usageLimitWidth}
+            flexDirection="column"
+            alignItems="flex-end"
+          >
+            {row.bucket &&
+              row.bucket.remainingFraction != null &&
+              row.bucket.resetTime && (
+                <Text color={theme.text.secondary} wrap="truncate-end">
+                  {(row.bucket.remainingFraction * 100).toFixed(1)}%{' '}
+                  {formatResetTime(row.bucket.resetTime)}
+                </Text>
+              )}
+          </Box>
+        </Box>
+      ))}
+
+      {cacheEfficiency > 0 && !showQuotaColumn && (
         <Box flexDirection="column" marginTop={1}>
           <Text color={theme.text.primary}>
             <Text color={theme.status.success}>Savings Highlight:</Text>{' '}
-            {totalCachedTokens.toLocaleString()} ({cacheEfficiency.toFixed(1)}
-            %) of input tokens were served from the cache, reducing costs.
+            {totalCachedTokens.toLocaleString()} (
+            <Text color={cacheEfficiencyColor}>
+              {cacheEfficiency.toFixed(1)}%
+            </Text>
+            ) of input tokens were served from the cache, reducing costs.
           </Text>
         </Box>
       )}
-      {models && (
+
+      {showQuotaColumn && (
         <>
           <Box marginTop={1} marginBottom={2}>
             <Text color={theme.text.primary}>
@@ -266,6 +413,7 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
       flexDirection="column"
       paddingY={1}
       paddingX={2}
+      overflow="hidden"
     >
       {renderTitle()}
       <Box height={1} />
@@ -335,15 +483,12 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
           </Text>
         </SubStatRow>
       </Section>
-
-      {Object.keys(models).length > 0 && (
-        <ModelUsageTable
-          models={models}
-          totalCachedTokens={computed.totalCachedTokens}
-          cacheEfficiency={computed.cacheEfficiency}
-          quotas={quotas}
-        />
-      )}
+      <ModelUsageTable
+        models={models}
+        quotas={quotas}
+        cacheEfficiency={computed.cacheEfficiency}
+        totalCachedTokens={computed.totalCachedTokens}
+      />
     </Box>
   );
 };