Cleanup, make more robust

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
-          <LoadExternalComponent
+          key={customComponent.id}
-            stream={thread}
+          stream={thread}
-            message={customComponent}
+          message={customComponent}
-            meta={{ ui: customComponent, artifact }}
+          meta={{ ui: customComponent, artifact }}
-          />
+        />
-        </Fragment>
       ))}
     </Fragment>
   );

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) => {
+      if (isUIMessage(event) || isRemoveUIMessage(event)) {
-        const ui = uiMessageReducer(prev.ui ?? [], event);
+        options.mutate((prev) => {
-        return { ...prev, ui };
+          const ui = uiMessageReducer(prev.ui ?? [], event);
-      });
+          return { ...prev, ui };
+        });
+      }
     },
     onThreadId: (id) => {
       setThreadId(id);