From df37b294b8acb8ecc07c15a1b1908ce0be59107a Mon Sep 17 00:00:00 2001
From: Oliver <opopadich@squareup.com>
Date: Mon, 5 May 2025 10:49:00 -0500
Subject: [PATCH] ToolCall UI Update (#2429)

---
 ui/desktop/src/components/GooseMessage.tsx    |   8 +-
 .../src/components/ToolCallArguments.tsx      |  30 +--
 .../src/components/ToolCallWithResponse.tsx   | 245 ++++++++++--------
 ui/desktop/src/components/ui/Dot.tsx          |  24 ++
 ui/desktop/src/components/ui/Expand.tsx       |   9 +
 5 files changed, 185 insertions(+), 131 deletions(-)
 create mode 100644 ui/desktop/src/components/ui/Dot.tsx
 create mode 100644 ui/desktop/src/components/ui/Expand.tsx

diff --git a/ui/desktop/src/components/GooseMessage.tsx b/ui/desktop/src/components/GooseMessage.tsx
index a4fc0cb9..3e25097c 100644
--- a/ui/desktop/src/components/GooseMessage.tsx
+++ b/ui/desktop/src/components/GooseMessage.tsx
@@ -103,9 +103,7 @@ export default function GooseMessage({
       <div className="flex flex-col w-full">
         {textContent && (
           <div className="flex flex-col group">
-            <div
-              className={`goose-message-content bg-bgSubtle rounded-2xl px-4 py-2 ${toolRequests.length > 0 ? 'rounded-b-none' : 'rounded-bl-none'}`}
-            >
+            <div className={`goose-message-content py-2`}>
               <div ref={contentRef}>{<MarkdownContent content={textContent} />}</div>
             </div>
             {/* Only show MessageCopyLink if there's text content and no tool requests/responses */}
@@ -126,9 +124,7 @@ export default function GooseMessage({
 
         {toolRequests.length > 0 && (
           <div className="relative flex flex-col w-full">
-            <div
-              className={`goose-message-tool bg-bgApp border border-borderSubtle dark:border-gray-700 ${textContent ? '' : 'rounded-t-2xl'} rounded-b-2xl rounded-bl-none px-4 pt-4 pb-2`}
-            >
+            <div className={`goose-message-tool bg-bgSubtle rounded px-2 py-2 mt-2`}>
               {toolRequests.map((toolRequest) => (
                 <ToolCallWithResponse
                   // If the message is resumed and not matched tool response, it means the tool is broken or cancelled.
diff --git a/ui/desktop/src/components/ToolCallArguments.tsx b/ui/desktop/src/components/ToolCallArguments.tsx
index dbff369e..8335ac82 100644
--- a/ui/desktop/src/components/ToolCallArguments.tsx
+++ b/ui/desktop/src/components/ToolCallArguments.tsx
@@ -1,6 +1,6 @@
-import { ChevronUp } from 'lucide-react';
 import React, { useState } from 'react';
 import MarkdownContent from './MarkdownContent';
+import Expand from './ui/Expand';
 
 type ToolCallArgumentValue =
   | string
@@ -28,34 +28,32 @@ export function ToolCallArguments({ args }: ToolCallArgumentsProps) {
 
       if (!needsExpansion) {
         return (
-          <div className="mb-2">
+          <div className="text-sm mb-2">
             <div className="flex flex-row">
-              <span className="text-sm font-medium text-textSubtle min-w-[140px]">{key}</span>
-              <span className="text-sm text-textStandard">{value}</span>
+              <span className="text-textSubtle min-w-[140px]">{key}</span>
+              <span className="text-textPlaceholder">{value}</span>
             </div>
           </div>
         );
       }
 
       return (
-        <div className="mb-2">
+        <div className="text-sm mb-2">
           <div className="flex flex-row">
-            <span className="text-sm font-medium text-textSubtle min-w-[140px]">{key}</span>
-            <div className="flex items-center">
+            <span className="text-textSubtle min-w-[140px]">{key}</span>
+            <div className="w-full flex justify-between items-start">
               {isExpanded ? (
-                <div className="mt-2">
-                  <MarkdownContent content={value} />
+                <div className="">
+                  <MarkdownContent content={value} className="text-sm text-textPlaceholder" />
                 </div>
               ) : (
-                <span className="text-sm text-textStandard mr-2">{value.slice(0, 60)}...</span>
+                <span className="text-textPlaceholder mr-2">{value.slice(0, 60)}...</span>
               )}
               <button
                 onClick={() => toggleKey(key)}
-                className="text-sm hover:opacity-75 text-textStandard"
+                className="hover:opacity-75 text-textPlaceholder pr-2"
               >
-                <ChevronUp
-                  className={`h-5 w-5 transition-all origin-center ${!isExpanded ? 'rotate-180' : ''}`}
-                />
+                <Expand size={5} isExpanded={isExpanded} />
               </button>
             </div>
           </div>
@@ -73,8 +71,8 @@ export function ToolCallArguments({ args }: ToolCallArgumentsProps) {
     return (
       <div className="mb-2">
         <div className="flex flex-row">
-          <span className="font-medium mr- text-textStandard min-w-[140px]2">{key}:</span>
-          <pre className="whitespace-pre-wrap text-textStandard">{content}</pre>
+          <span className="mr- text-textPlaceholder min-w-[140px]2">{key}:</span>
+          <pre className="whitespace-pre-wrap text-textPlaceholder">{content}</pre>
         </div>
       </div>
     );
diff --git a/ui/desktop/src/components/ToolCallWithResponse.tsx b/ui/desktop/src/components/ToolCallWithResponse.tsx
index ab9146e1..327cf16d 100644
--- a/ui/desktop/src/components/ToolCallWithResponse.tsx
+++ b/ui/desktop/src/components/ToolCallWithResponse.tsx
@@ -1,12 +1,11 @@
 import React from 'react';
 import { Card } from './ui/card';
-import Box from './ui/Box';
 import { ToolCallArguments } from './ToolCallArguments';
 import MarkdownContent from './MarkdownContent';
-import { LoadingPlaceholder } from './LoadingPlaceholder';
-import { ChevronUp } from 'lucide-react';
 import { Content, ToolRequestMessageContent, ToolResponseMessageContent } from '../types/message';
 import { snakeToTitleCase } from '../utils';
+import Dot, { LoadingStatus } from './ui/Dot';
+import Expand from './ui/Expand';
 
 interface ToolCallWithResponseProps {
   isCancelledMessage: boolean;
@@ -20,138 +19,166 @@ export default function ToolCallWithResponse({
   toolResponse,
 }: ToolCallWithResponseProps) {
   const toolCall = toolRequest.toolCall.status === 'success' ? toolRequest.toolCall.value : null;
-
   if (!toolCall) {
     return null;
   }
 
   return (
-    <div className="w-full">
+    <div className={'w-full text-textSubtle text-sm'}>
       <Card className="">
-        <ToolCallView toolCall={toolCall} />
-        {!isCancelledMessage ? (
-          toolResponse ? (
-            <ToolResultView
-              result={
-                toolResponse.toolResult.status === 'success'
-                  ? toolResponse.toolResult.value
-                  : undefined
-              }
-            />
-          ) : (
-            <LoadingPlaceholder />
-          )
-        ) : undefined}
+        <ToolCallView {...{ isCancelledMessage, toolCall, toolResponse }} />
       </Card>
     </div>
   );
 }
 
+interface ToolCallExpandableProps {
+  label: string | React.ReactNode;
+  defaultExpanded?: boolean;
+  forceExpand?: boolean;
+  children: React.ReactNode;
+  className?: string;
+}
+
+function ToolCallExpandable({
+  label,
+  defaultExpanded = false,
+  forceExpand,
+  children,
+  className = '',
+}: ToolCallExpandableProps) {
+  const [isExpanded, setIsExpanded] = React.useState(defaultExpanded);
+  const toggleExpand = () => setIsExpanded((prev) => !prev);
+  React.useEffect(() => {
+    if (forceExpand) setIsExpanded(true);
+  }, [forceExpand]);
+
+  return (
+    <div className={className}>
+      <button onClick={toggleExpand} className="w-full flex justify-between items-center pr-2">
+        <span className="flex items-center">{label}</span>
+        <Expand size={5} isExpanded={isExpanded} />
+      </button>
+      {isExpanded && <div>{children}</div>}
+    </div>
+  );
+}
+
 interface ToolCallViewProps {
+  isCancelledMessage: boolean;
+  toolCall: {
+    name: string;
+    arguments: Record<string, unknown>;
+  };
+  toolResponse?: ToolResponseMessageContent;
+}
+
+function ToolCallView({ isCancelledMessage, toolCall, toolResponse }: ToolCallViewProps) {
+  const isToolDetails = Object.entries(toolCall?.arguments).length > 0;
+  const loadingStatus: LoadingStatus = !toolResponse?.toolResult.status
+    ? 'loading'
+    : toolResponse?.toolResult.status;
+
+  const toolResults: { result: Content; defaultExpanded: boolean }[] =
+    loadingStatus === 'success' && Array.isArray(toolResponse?.toolResult.value)
+      ? toolResponse.toolResult.value
+          .filter((item) => {
+            const audience = item.annotations?.audience as string[] | undefined;
+            return !audience || audience.includes('user');
+          })
+          .map((item) => ({
+            result: item,
+            defaultExpanded: ((item.annotations?.priority as number | undefined) ?? -1) >= 0.5,
+          }))
+      : [];
+
+  const shouldExpand = toolResults.some((v) => v.defaultExpanded);
+
+  return (
+    <ToolCallExpandable
+      defaultExpanded={shouldExpand}
+      forceExpand={shouldExpand}
+      label={
+        <>
+          <Dot size={2} loadingStatus={loadingStatus} />
+          <span className="ml-[10px]">
+            {snakeToTitleCase(toolCall.name.substring(toolCall.name.lastIndexOf('__') + 2))}
+          </span>
+        </>
+      }
+    >
+      {/* Tool Details */}
+      {isToolDetails && (
+        <div className="bg-bgStandard rounded-t mt-1">
+          <ToolDetailsView toolCall={toolCall} />
+        </div>
+      )}
+
+      {/* Tool Output */}
+      {!isCancelledMessage && (
+        <>
+          {toolResults.map(({ result, defaultExpanded }, index) => {
+            const isLast = index === toolResults.length - 1;
+            return (
+              <div
+                key={index}
+                className={`bg-bgStandard mt-1 ${isToolDetails ? 'rounded-t-none' : ''} ${isLast ? 'rounded-b' : ''}`}
+              >
+                <ToolResultView result={result} defaultExpanded={defaultExpanded} />
+              </div>
+            );
+          })}
+        </>
+      )}
+    </ToolCallExpandable>
+  );
+}
+
+interface ToolDetailsViewProps {
   toolCall: {
     name: string;
     arguments: Record<string, unknown>;
   };
 }
 
-function ToolCallView({ toolCall }: ToolCallViewProps) {
+function ToolDetailsView({ toolCall }: ToolDetailsViewProps) {
   return (
-    <div>
-      <div className="flex items-center mb-4">
-        <Box size={16} />
-        <span className="ml-[8px] text-textStandard">
-          {snakeToTitleCase(toolCall.name.substring(toolCall.name.lastIndexOf('__') + 2))}
-        </span>
-      </div>
-
+    <ToolCallExpandable label="Tool Details" className="pl-[19px] py-1">
       {toolCall.arguments && <ToolCallArguments args={toolCall.arguments} />}
-
-      <div className="self-stretch h-px my-[10px] -mx-4 bg-borderSubtle dark:bg-gray-700" />
-    </div>
+    </ToolCallExpandable>
   );
 }
 
 interface ToolResultViewProps {
-  result?: Content[];
+  result: Content;
+  defaultExpanded: boolean;
 }
 
-function ToolResultView({ result }: ToolResultViewProps) {
-  // State to track expanded items
-  const [expandedItems, setExpandedItems] = React.useState<number[]>([]);
-
-  // If no result info, don't show anything
-  if (!result) return null;
-
-  // Find results where either audience is not set, or it's set to a list that includes user
-  const filteredResults = result.filter((item) => {
-    // Check audience (which may not be in the type)
-    const audience = item.annotations?.audience;
-
-    return !audience || audience.includes('user');
-  });
-
-  if (filteredResults.length === 0) return null;
-
-  const toggleExpand = (index: number) => {
-    setExpandedItems((prev) =>
-      prev.includes(index) ? prev.filter((i) => i !== index) : [...prev, index]
-    );
-  };
-
-  const shouldShowExpanded = (item: Content, index: number) => {
-    return (
-      (item.annotations &&
-        item.annotations.priority !== undefined &&
-        item.annotations.priority >= 0.5) ||
-      expandedItems.includes(index)
-    );
-  };
-
+function ToolResultView({ result, defaultExpanded }: ToolResultViewProps) {
   return (
-    <div className="">
-      {filteredResults.map((item, index) => {
-        const isExpanded = shouldShowExpanded(item, index);
-        const shouldMinimize =
-          !item.annotations ||
-          item.annotations.priority === undefined ||
-          item.annotations.priority < 0.5;
-        return (
-          <div key={index} className="relative">
-            {shouldMinimize && (
-              <button
-                onClick={() => toggleExpand(index)}
-                className="mb-1 flex items-center text-textStandard"
-              >
-                <span className="mr-2 text-sm">Output</span>
-                <ChevronUp
-                  className={`h-5 w-5 transition-all origin-center ${!isExpanded ? 'rotate-180' : ''}`}
-                />
-              </button>
-            )}
-            {(isExpanded || !shouldMinimize) && (
-              <>
-                {item.type === 'text' && item.text && (
-                  <MarkdownContent
-                    content={item.text}
-                    className="whitespace-pre-wrap p-2 max-w-full overflow-x-auto"
-                  />
-                )}
-                {item.type === 'image' && (
-                  <img
-                    src={`data:${item.mimeType};base64,${item.data}`}
-                    alt="Tool result"
-                    className="max-w-full h-auto rounded-md my-2"
-                    onError={(e) => {
-                      console.error('Failed to load image: Invalid MIME-type encoded image data');
-                      e.currentTarget.style.display = 'none';
-                    }}
-                  />
-                )}
-              </>
-            )}
-          </div>
-        );
-      })}
-    </div>
+    <ToolCallExpandable
+      label={<span className="pl-[19px] py-1">Output</span>}
+      defaultExpanded={defaultExpanded}
+    >
+      <div className="bg-bgApp rounded-b pl-[19px] pr-2 py-4">
+        {result.type === 'text' && result.text && (
+          <MarkdownContent
+            content={result.text}
+            className="whitespace-pre-wrap p-2 max-w-full overflow-x-auto"
+          />
+        )}
+        {result.type === 'image' && (
+          <img
+            src={`data:${result.mimeType};base64,${result.data}`}
+            alt="Tool result"
+            className="max-w-full h-auto rounded-md my-2"
+            onError={(e) => {
+              console.error('Failed to load image');
+              e.currentTarget.style.display = 'none';
+            }}
+          />
+        )}
+      </div>
+    </ToolCallExpandable>
   );
 }
diff --git a/ui/desktop/src/components/ui/Dot.tsx b/ui/desktop/src/components/ui/Dot.tsx
new file mode 100644
index 00000000..c6583b8e
--- /dev/null
+++ b/ui/desktop/src/components/ui/Dot.tsx
@@ -0,0 +1,24 @@
+export type LoadingStatus = 'loading' | 'success' | 'error';
+export default function Dot({
+  size,
+  loadingStatus,
+}: {
+  size: number;
+  loadingStatus: LoadingStatus;
+}) {
+  const backgroundColor =
+    {
+      loading: '#2693FF',
+      success: 'var(--icon-extra-subtle)',
+      error: '#CC0023',
+    }[loadingStatus] ?? 'var(--icon-extra-subtle)';
+
+  return (
+    <div
+      className={`w-${size} h-${size} rounded-full`}
+      style={{
+        backgroundColor: backgroundColor,
+      }}
+    />
+  );
+}
diff --git a/ui/desktop/src/components/ui/Expand.tsx b/ui/desktop/src/components/ui/Expand.tsx
new file mode 100644
index 00000000..0529248b
--- /dev/null
+++ b/ui/desktop/src/components/ui/Expand.tsx
@@ -0,0 +1,9 @@
+import { ChevronUp } from 'lucide-react';
+
+export default function Expand({ size, isExpanded }: { size: number; isExpanded: boolean }) {
+  return (
+    <ChevronUp
+      className={`w-${size} h-${size} text-textPlaceholder transition-all origin-center ${isExpanded ? 'rotate-180' : 'rotate-90'}`}
+    />
+  );
+}