diff --git a/src/components/thread/messages/ai.tsx b/src/components/thread/messages/ai.tsx
index 828ea29..af5eac2 100644
--- a/src/components/thread/messages/ai.tsx
+++ b/src/components/thread/messages/ai.tsx
@@ -32,13 +32,12 @@ function CustomComponent({
   return (
     <Fragment key={message.id}>
       {customComponents.map((customComponent) => (
-        <Fragment key={customComponent.id}>
-          <LoadExternalComponent
-            stream={thread}
-            message={customComponent}
-            meta={{ ui: customComponent, artifact }}
-          />
-        </Fragment>
+        <LoadExternalComponent
+          key={customComponent.id}
+          stream={thread}
+          message={customComponent}
+          meta={{ ui: customComponent, artifact }}
+        />
       ))}
     </Fragment>
   );
diff --git a/src/providers/Stream.tsx b/src/providers/Stream.tsx
index c0c471c..5215494 100644
--- a/src/providers/Stream.tsx
+++ b/src/providers/Stream.tsx
@@ -9,6 +9,8 @@ import { useStream } from "@langchain/langgraph-sdk/react";
 import { type Message } from "@langchain/langgraph-sdk";
 import {
   uiMessageReducer,
+  isUIMessage,
+  isRemoveUIMessage,
   type UIMessage,
   type RemoveUIMessage,
 } from "@langchain/langgraph-sdk/react-ui";
@@ -82,10 +84,12 @@ const StreamSession = ({
     assistantId,
     threadId: threadId ?? null,
     onCustomEvent: (event, options) => {
-      options.mutate((prev) => {
-        const ui = uiMessageReducer(prev.ui ?? [], event);
-        return { ...prev, ui };
-      });
+      if (isUIMessage(event) || isRemoveUIMessage(event)) {
+        options.mutate((prev) => {
+          const ui = uiMessageReducer(prev.ui ?? [], event);
+          return { ...prev, ui };
+        });
+      }
     },
     onThreadId: (id) => {
       setThreadId(id);