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