From 80b04dc505bf1c784a54c5d80d971310b05144cc Mon Sep 17 00:00:00 2001
From: Taylor Mullen <ntaylormullen@google.com>
Date: Tue, 22 Apr 2025 07:48:12 -0400
Subject: [PATCH] Update UI of tool messages

- Bring tool messages in line with original envisioned UI of: https://screenshot.googleplex.com/9yZCX636LzpMrgc
  - In particular this represents more descriptive names. FWIW we already had this tech we just weren't passing around information correctly (`displayName` vs. `name`)
 - Add gray to our list of color pallete's and removed Background (unused)
 - Re-enabled representing canceled messages
 - Migrated back towards a cleaner tool message design of status symbols & border colors vs. overly verbose text.
 - Removed border from confirmation diffs.
Fixes https://b.corp.google.com/issues/412598909
---
 packages/cli/src/ui/colors.ts                 |   2 +-
 .../ui/components/messages/DiffRenderer.tsx   |   8 +-
 .../ui/components/messages/GeminiMessage.tsx  |  23 +---
 .../components/messages/ToolGroupMessage.tsx  |  13 ++-
 .../ui/components/messages/ToolMessage.tsx    | 105 ++++++++----------
 .../ui/components/messages/UserMessage.tsx    |   6 +-
 packages/cli/src/ui/hooks/useGeminiStream.ts  |   2 +-
 packages/cli/src/ui/types.ts                  |   2 +-
 8 files changed, 70 insertions(+), 91 deletions(-)

diff --git a/packages/cli/src/ui/colors.ts b/packages/cli/src/ui/colors.ts
index 946be512..f320d856 100644
--- a/packages/cli/src/ui/colors.ts
+++ b/packages/cli/src/ui/colors.ts
@@ -5,7 +5,6 @@
  */
 
 export const Colors = {
-  Background: '#1E1E2E',
   Foreground: 'white',
   LightBlue: '#CDD6F4',
   AccentBlue: '#89B4FA',
@@ -15,4 +14,5 @@ export const Colors = {
   AccentYellow: '#F9E2AF',
   AccentRed: '#F38BA8',
   SubtleComment: '#6C7086',
+  Gray: 'gray',
 };
diff --git a/packages/cli/src/ui/components/messages/DiffRenderer.tsx b/packages/cli/src/ui/components/messages/DiffRenderer.tsx
index 01cc4938..eb3133c3 100644
--- a/packages/cli/src/ui/components/messages/DiffRenderer.tsx
+++ b/packages/cli/src/ui/components/messages/DiffRenderer.tsx
@@ -138,11 +138,7 @@ export const DiffRenderer: React.FC<DiffRendererProps> = ({
   // --- End Modification ---
 
   return (
-    <Box
-      borderStyle="round"
-      borderColor={Colors.SubtleComment}
-      flexDirection="column"
-    >
+    <Box flexDirection="column">
       {/* Iterate over the lines that should be displayed (already normalized) */}
       {displayableLines.map((line, index) => {
         const key = `diff-line-${index}`;
@@ -179,7 +175,7 @@ export const DiffRenderer: React.FC<DiffRendererProps> = ({
         return (
           // Using your original rendering structure
           <Box key={key} flexDirection="row">
-            <Text color={Colors.SubtleComment}>{gutterNumStr} </Text>
+            <Text color={Colors.Foreground}>{gutterNumStr} </Text>
             <Text color={color} dimColor={dim}>
               {prefixSymbol}{' '}
             </Text>
diff --git a/packages/cli/src/ui/components/messages/GeminiMessage.tsx b/packages/cli/src/ui/components/messages/GeminiMessage.tsx
index 584c7729..11449b18 100644
--- a/packages/cli/src/ui/components/messages/GeminiMessage.tsx
+++ b/packages/cli/src/ui/components/messages/GeminiMessage.tsx
@@ -16,28 +16,13 @@ interface GeminiMessageProps {
 export const GeminiMessage: React.FC<GeminiMessageProps> = ({ text }) => {
   const prefix = '✦ ';
   const prefixWidth = prefix.length;
-
-  // Handle potentially null or undefined text gracefully
-  const safeText = text || '';
-
-  // Use the static render method from the MarkdownRenderer class
-  // Pass safeText which is guaranteed to be a string
-  const renderedBlocks = MarkdownRenderer.render(safeText);
-
-  // If the original text was actually empty/null, render the minimal state
-  if (!safeText && renderedBlocks.length === 0) {
-    return (
-      <Box flexDirection="row">
-        <Box width={prefixWidth}>
-          <Text color={Colors.AccentPurple}>{prefix}</Text>
-        </Box>
-        <Box flexGrow={1}></Box>
-      </Box>
-    );
-  }
+  const renderedBlocks = MarkdownRenderer.render(text);
 
   return (
     <Box flexDirection="row">
+      <Box width={prefixWidth}>
+        <Text color={Colors.AccentPurple}>{prefix}</Text>
+      </Box>
       <Box flexGrow={1} flexDirection="column">
         {renderedBlocks}
       </Box>
diff --git a/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx b/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx
index 448ed4c5..0675411f 100644
--- a/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx
+++ b/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx
@@ -22,11 +22,18 @@ export const ToolGroupMessage: React.FC<ToolGroupMessageProps> = ({
   toolCalls,
   onSubmit,
 }) => {
-  const hasPending = toolCalls.some((t) => t.status === ToolCallStatus.Pending);
-  const borderColor = hasPending ? Colors.AccentYellow : Colors.AccentBlue;
+  const hasPending = !toolCalls.every(
+    (t) => t.status === ToolCallStatus.Success,
+  );
+  const borderColor = hasPending ? Colors.AccentYellow : Colors.AccentCyan;
 
   return (
-    <Box flexDirection="column" borderStyle="round" borderColor={borderColor}>
+    <Box
+      flexDirection="column"
+      borderStyle="round"
+      borderDimColor={hasPending}
+      borderColor={borderColor}
+    >
       {toolCalls.map((tool) => (
         <React.Fragment key={tool.callId}>
           <ToolMessage
diff --git a/packages/cli/src/ui/components/messages/ToolMessage.tsx b/packages/cli/src/ui/components/messages/ToolMessage.tsx
index 53f31db2..ab590f53 100644
--- a/packages/cli/src/ui/components/messages/ToolMessage.tsx
+++ b/packages/cli/src/ui/components/messages/ToolMessage.tsx
@@ -9,74 +9,63 @@ import { Box, Text } from 'ink';
 import Spinner from 'ink-spinner';
 import { IndividualToolCallDisplay, ToolCallStatus } from '../../types.js';
 import { DiffRenderer } from './DiffRenderer.js';
-import { FileDiff, ToolResultDisplay } from '@gemini-code/server';
 import { Colors } from '../../colors.js';
+import { MarkdownRenderer } from '../../utils/MarkdownRenderer.js';
 
 export const ToolMessage: React.FC<IndividualToolCallDisplay> = ({
-  callId,
   name,
   description,
   resultDisplay,
   status,
 }) => {
-  const typedResultDisplay = resultDisplay as ToolResultDisplay | undefined;
-
-  let color = Colors.SubtleComment;
-  let prefix = '';
-  switch (status) {
-    case ToolCallStatus.Pending:
-      prefix = 'Pending:';
-      break;
-    case ToolCallStatus.Invoked:
-      prefix = 'Executing:';
-      break;
-    case ToolCallStatus.Confirming:
-      color = Colors.AccentYellow;
-      prefix = 'Confirm:';
-      break;
-    case ToolCallStatus.Success:
-      color = Colors.AccentGreen;
-      prefix = 'Success:';
-      break;
-    case ToolCallStatus.Error:
-      color = Colors.AccentRed;
-      prefix = 'Error:';
-      break;
-    default:
-      // Handle unexpected status if necessary, or just break
-      break;
-  }
-
-  const title = `${prefix} ${name}`;
-
+  const statusIndicatorWidth = 3;
+  const hasResult = resultDisplay && resultDisplay.toString().trim().length > 0;
   return (
-    <Box key={callId} flexDirection="column" paddingX={1}>
-      <Box>
-        {status === ToolCallStatus.Invoked && (
-          <Box marginRight={1}>
-            <Text color={Colors.AccentBlue}>
-              <Spinner type="dots" />
-            </Text>
-          </Box>
-        )}
-        <Text bold color={color}>
-          {title}
-        </Text>
-        <Text color={color}>
-          {status === ToolCallStatus.Error && typedResultDisplay
-            ? `: ${typedResultDisplay}`
-            : ` - ${description}`}
-        </Text>
-      </Box>
-      {status === ToolCallStatus.Success && typedResultDisplay && (
-        <Box flexDirection="column" marginLeft={2}>
-          {typeof typedResultDisplay === 'string' ? (
-            <Text>{typedResultDisplay}</Text>
-          ) : (
-            <DiffRenderer
-              diffContent={(typedResultDisplay as FileDiff).fileDiff}
-            />
+    <Box paddingX={1} paddingY={0} flexDirection="column">
+      <Box minHeight={1}>
+        {/* Status Indicator */}
+        <Box minWidth={statusIndicatorWidth}>
+          {status === ToolCallStatus.Pending && <Spinner type="dots" />}
+          {status === ToolCallStatus.Success && (
+            <Text color={Colors.AccentGreen}>✔</Text>
           )}
+          {status === ToolCallStatus.Confirming && (
+            <Text color={Colors.AccentPurple}>?</Text>
+          )}
+          {status === ToolCallStatus.Canceled && (
+            <Text color={Colors.AccentYellow} bold>
+              -
+            </Text>
+          )}
+          {status === ToolCallStatus.Error && (
+            <Text color={Colors.AccentRed} bold>
+              x
+            </Text>
+          )}
+        </Box>
+        <Box>
+          <Text
+            wrap="truncate-end"
+            strikethrough={status === ToolCallStatus.Canceled}
+          >
+            <Text bold>{name}</Text>{' '}
+            <Text color={Colors.SubtleComment}>{description}</Text>
+          </Text>
+        </Box>
+      </Box>
+      {hasResult && (
+        <Box paddingLeft={statusIndicatorWidth}>
+          <Box flexShrink={1} flexDirection="row">
+            {/* Use default text color (white) or gray instead of dimColor */}
+            {typeof resultDisplay === 'string' && (
+              <Box flexDirection="column">
+                {MarkdownRenderer.render(resultDisplay)}
+              </Box>
+            )}
+            {typeof resultDisplay === 'object' && (
+              <DiffRenderer diffContent={resultDisplay.fileDiff} />
+            )}
+          </Box>
         </Box>
       )}
     </Box>
diff --git a/packages/cli/src/ui/components/messages/UserMessage.tsx b/packages/cli/src/ui/components/messages/UserMessage.tsx
index e9771b82..a83ffea7 100644
--- a/packages/cli/src/ui/components/messages/UserMessage.tsx
+++ b/packages/cli/src/ui/components/messages/UserMessage.tsx
@@ -19,10 +19,12 @@ export const UserMessage: React.FC<UserMessageProps> = ({ text }) => {
   return (
     <Box flexDirection="row">
       <Box width={prefixWidth}>
-        <Text color={Colors.SubtleComment}>{prefix}</Text>
+        <Text color={Colors.Gray}>{prefix}</Text>
       </Box>
       <Box flexGrow={1}>
-        <Text wrap="wrap">{text}</Text>
+        <Text wrap="wrap" color={Colors.Gray}>
+          {text}
+        </Text>
       </Box>
     </Box>
   );
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index ffdd9967..5baebe09 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -269,7 +269,7 @@ export const useGeminiStream = (
             // Create the UI display object matching IndividualToolCallDisplay
             const toolCallDisplay: IndividualToolCallDisplay = {
               callId,
-              name,
+              name: cliTool.displayName,
               description,
               status: ToolCallStatus.Pending,
               resultDisplay: undefined,
diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts
index fe135909..70ecca67 100644
--- a/packages/cli/src/ui/types.ts
+++ b/packages/cli/src/ui/types.ts
@@ -25,7 +25,7 @@ export enum GeminiEventType {
 
 export enum ToolCallStatus {
   Pending = 'Pending',
-  Invoked = 'Invoked',
+  Canceled = 'Canceled',
   Confirming = 'Confirming',
   Success = 'Success',
   Error = 'Error',