From 0a4b8d740021a63ebefc026eb478445ed5cdab51 Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Fri, 6 Sep 2024 16:58:16 -0700
Subject: [PATCH] fix: tool UI fallback should not override makeAssistantToolUI
 definitions (#753)

---
 .changeset/afraid-walls-thank.md              |  5 +++
 .changeset/weak-singers-double.md             |  5 +++
 .../src/primitives/message/MessageContent.tsx | 31 +++++++++++++------
 3 files changed, 31 insertions(+), 10 deletions(-)
 create mode 100644 .changeset/afraid-walls-thank.md
 create mode 100644 .changeset/weak-singers-double.md

diff --git a/.changeset/afraid-walls-thank.md b/.changeset/afraid-walls-thank.md
new file mode 100644
index 000000000..2610a5b52
--- /dev/null
+++ 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
new file mode 100644
index 000000000..2ce0a6486
--- /dev/null
+++ 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
index 6fb9b159d..257dbe0d6 100644
--- 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;