feat: upgrade to 0.0.52

src/components/thread/messages/ai.tsx

@@ -7,8 +7,8 @@ import { MarkdownText } from "../markdown-text";
 import { LoadExternalComponent } from "@langchain/langgraph-sdk/react-ui";
 import { cn } from "@/lib/utils";
 import { ToolCalls, ToolResult } from "./tool-calls";
-import { StringParam, useQueryParam } from "use-query-params";
 import { MessageContentComplex } from "@langchain/core/messages";
+import { Fragment } from "react/jsx-runtime";
 
 function CustomComponent({
   message,
@@ -17,32 +17,28 @@ function CustomComponent({
   message: Message;
   thread: ReturnType<typeof useStreamContext>;
 }) {
-  const [apiUrl] = useQueryParam("apiUrl", StringParam);
   const meta = thread.getMessagesMetadata(message);
   const seenState = meta?.firstSeenState;
   const customComponents = seenState?.values.ui
     ?.slice()
-    .filter(
-      ({ additional_kwargs }) =>
-        additional_kwargs.run_id === seenState.metadata?.run_id &&
-        (!additional_kwargs.message_id ||
-          additional_kwargs.message_id === message.id),
+    .filter(({ additional_kwargs }) =>
+      !additional_kwargs.message_id
+        ? additional_kwargs.run_id === seenState.metadata?.run_id
+        : additional_kwargs.message_id === message.id,
     );
 
   if (!customComponents?.length) return null;
   return (
-    <div key={message.id}>
+    <Fragment key={message.id}>
       {customComponents.map((customComponent) => (
         <LoadExternalComponent
           key={customComponent.id}
-          apiUrl={apiUrl ?? undefined}
-          assistantId="agent"
           stream={thread}
           message={customComponent}
           meta={{ ui: customComponent }}
         />
       ))}
-    </div>
+    </Fragment>
   );
 }
 

src/providers/Stream.tsx

@@ -1,9 +1,10 @@
 import React, { createContext, useContext, ReactNode, useState } from "react";
 import { useStream } from "@langchain/langgraph-sdk/react";
 import { type Message } from "@langchain/langgraph-sdk";
-import type {
-  UIMessage,
-  RemoveUIMessage,
+import {
+  uiMessageReducer,
+  type UIMessage,
+  type RemoveUIMessage,
 } from "@langchain/langgraph-sdk/react-ui";
 import { useQueryParam, StringParam } from "use-query-params";
 import { Input } from "@/components/ui/input";
@@ -24,7 +25,7 @@ const useTypedStream = useStream<
       messages?: Message[] | Message | string;
       ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
     };
-    CustomUpdateType: UIMessage | RemoveUIMessage;
+    CustomEventType: UIMessage | RemoveUIMessage;
   }
 >;
 
@@ -53,6 +54,12 @@ const StreamSession = ({
     apiKey: apiKey ?? undefined,
     assistantId,
     threadId: threadId ?? null,
+    onCustomEvent: (event, options) => {
+      options.mutate((prev) => {
+        const ui = uiMessageReducer(prev.ui ?? [], event);
+        return { ...prev, ui };
+      });
+    },
     onThreadId: (id) => {
       setThreadId(id);
       // Refetch threads list when thread ID changes.