feat(sdk-js): api improvements for gen ui (#3760)

- merge `typedUi.create` and `typedUi.write` into `typedUi.push`
- Add mutate function in `onCustomEvent`

libs/sdk-js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.50",
+  "version": "0.0.51",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",

libs/sdk-js/src/react-ui/index.ts

@@ -2,4 +2,8 @@ import { bootstrapUiContext } from "./client.js";
 bootstrapUiContext();
 
 export { useStreamContext, LoadExternalComponent } from "./client.js";
-export type { UIMessage, RemoveUIMessage } from "./types.js";
+export {
+  uiMessageReducer,
+  type UIMessage,
+  type RemoveUIMessage,
+} from "./types.js";

libs/sdk-js/src/react-ui/server/server.ts

@@ -2,6 +2,10 @@ import { v4 as uuidv4 } from "uuid";
 import type { ComponentPropsWithoutRef, ElementType } from "react";
 import type { RemoveUIMessage, UIMessage } from "../types.js";
 
+interface MessageLike {
+  id?: string;
+}
+
 export const typedUi = <Decl extends Record<string, ElementType>>(config: {
   writer?: (chunk: unknown) => void;
   runId?: string;
@@ -10,7 +14,7 @@ export const typedUi = <Decl extends Record<string, ElementType>>(config: {
   runName?: string;
 }) => {
   type PropMap = { [K in keyof Decl]: ComponentPropsWithoutRef<Decl[K]> };
-  let collect: (UIMessage | RemoveUIMessage)[] = [];
+  let items: (UIMessage | RemoveUIMessage)[] = [];
 
   const runId = (config.metadata?.run_id as string | undefined) ?? config.runId;
   if (!runId) throw new Error("run_id is required");
@@ -22,28 +26,37 @@ export const typedUi = <Decl extends Record<string, ElementType>>(config: {
     run_id: runId,
   };
 
-  const create = <K extends keyof PropMap & string>(
-    name: K,
-    props: PropMap[K],
-  ): UIMessage => ({
-    type: "ui" as const,
-    id: uuidv4(),
-    name,
-    content: props,
-    additional_kwargs: metadata,
-  });
-
-  const remove = (id: string): RemoveUIMessage => ({ type: "remove-ui", id });
-
-  return {
-    create,
-    remove,
-
-    collect,
-    write: <K extends keyof PropMap & string>(name: K, props: PropMap[K]) => {
-      const evt: UIMessage = create(name, props);
-      collect.push(evt);
-      config.writer?.(evt);
+  const handlePush = <K extends keyof PropMap & string>(
+    message: {
+      id?: string;
+      name: K;
+      content: PropMap[K];
+      additional_kwargs?: Record<string, unknown>;
     },
+    options?: { message?: MessageLike },
+  ): UIMessage => {
+    const evt: UIMessage = {
+      type: "ui" as const,
+      id: message?.id ?? uuidv4(),
+      name: message?.name,
+      content: message?.content,
+      additional_kwargs: {
+        ...metadata,
+        ...message?.additional_kwargs,
+        ...(options?.message ? { message_id: options.message.id } : null),
+      },
+    };
+    items.push(evt);
+    config.writer?.(evt);
+    return evt;
   };
+
+  const handleDelete = (id: string): RemoveUIMessage => {
+    const evt: RemoveUIMessage = { type: "remove-ui", id };
+    items.push(evt);
+    config.writer?.(evt);
+    return evt;
+  };
+
+  return { push: handlePush, delete: handleDelete, items };
 };

libs/sdk-js/src/react-ui/types.ts

@@ -6,6 +6,7 @@ export interface UIMessage {
   content: Record<string, unknown>;
   additional_kwargs: {
     run_id: string;
+    message_id?: string;
     [key: string]: unknown;
   };
 }

libs/sdk-js/src/react/stream.tsx

@@ -464,6 +464,11 @@ interface UseStreamOptions<
    */
   onCustomEvent?: (
     data: CustomStreamEvent<GetCustomEventType<Bag>>["data"],
+    options: {
+      mutate: (
+        update: Partial<StateType> | ((prev: StateType) => Partial<StateType>),
+      ) => void;
+    },
   ) => void;
 
   /**
@@ -834,7 +839,18 @@ export function useStream<
         }
 
         if (event === "updates") options.onUpdateEvent?.(data);
-        if (event === "custom") options.onCustomEvent?.(data);
+        if (event === "custom")
+          options.onCustomEvent?.(data, {
+            mutate: (update) =>
+              setStreamValues((prev) => {
+                // should not happen
+                if (prev == null) return prev;
+                return {
+                  ...prev,
+                  ...(typeof update === "function" ? update(prev) : update),
+                };
+              }),
+          });
         if (event === "metadata") options.onMetadataEvent?.(data);
 
         if (event === "values") setStreamValues(data);