fix: tool UI fallback should not override makeAssistantToolUI definit…

…ions (#753)
Yonom · Sep 6, 2024 · 0a4b8d7 · 0a4b8d7
1 parent b4eabee
commit 0a4b8d7
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 10 deletions.
diff --git a/.changeset/afraid-walls-thank.md b/.changeset/afraid-walls-thank.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: adjust the override order of tool UIs
diff --git a/.changeset/weak-singers-double.md b/.changeset/weak-singers-double.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+fix: tool UI fallback should not override makeAssistantToolUI definitions
diff --git a/packages/react/src/primitives/message/MessageContent.tsx b/packages/react/src/primitives/message/MessageContent.tsx
@@ -40,6 +40,18 @@ export type MessagePrimitiveContentProps = {
     | undefined;
 };
 
+const ToolUIDisplay = ({
+  UI,
+  ...props
+}: {
+  UI: ToolCallContentPartComponent | undefined;
+} & ToolCallContentPartProps) => {
+  const { useToolUIs } = useAssistantContext();
+  const Render = useToolUIs((s) => s.getToolUI(props.part.toolName)) ?? UI;
+  if (!Render) return null;
+  return <Render {...props} />;
+};
+
 const defaultComponents = {
   Text: () => (
     <p style={{ whiteSpace: "pre-line" }}>
@@ -51,14 +63,6 @@ const defaultComponents = {
   ),
   Image: () => <ContentPartPrimitiveImage />,
   UI: () => <ContentPartPrimitiveDisplay />,
-  tools: {
-    Fallback: (props) => {
-      const { useToolUIs } = useAssistantContext();
-      const Render = useToolUIs((s) => s.getToolUI(props.part.toolName));
-      if (!Render) return null;
-      return <Render {...props} />;
-    },
-  },
 } satisfies MessagePrimitiveContentProps["components"];
 
 type MessageContentPartComponentProps = {
@@ -71,7 +75,7 @@ const MessageContentPartComponent: FC<MessageContentPartComponentProps> = ({
     Text = defaultComponents.Text,
     Image = defaultComponents.Image,
     UI = defaultComponents.UI,
-    tools: { by_name = {}, Fallback = defaultComponents.tools.Fallback } = {},
+    tools: { by_name = {}, Fallback = undefined } = {},
   } = {},
 }) => {
   const { useThreadActions } = useThreadContext();
@@ -110,7 +114,14 @@ const MessageContentPartComponent: FC<MessageContentPartComponentProps> = ({
           toolCallId: part.toolCallId,
           result,
         });
-      return <Tool part={part} status={status} addResult={addResult} />;
+      return (
+        <ToolUIDisplay
+          UI={Tool}
+          part={part}
+          status={status}
+          addResult={addResult}
+        />
+      );
     }
     default:
       const unhandledType: never = type;