fix(sdk): resolve React "Maximum update depth exceeded" warning by batching updates in a macrotask (#1802)

.changeset/great-masks-design.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph-sdk": minor
+---
+
+Add throttle option to `useStream` and batch updates in a macrotask to prevent `Maximum update depth exceeded` error

libs/sdk-validation/src/tests/stream.test.tsx

@@ -15,6 +15,7 @@ import {
   Runtime,
   interrupt,
   END,
+  pushMessage,
 } from "@langchain/langgraph";
 import { MemorySaver } from "@langchain/langgraph-checkpoint";
 import { FakeStreamingChatModel } from "@langchain/core/utils/testing";
@@ -92,33 +93,16 @@ const interruptAgent = new StateGraph(MessagesAnnotation)
 const removeMessageAgent = new StateGraph(MessagesAnnotation)
   .addSequence({
     step1: () => ({ messages: [new AIMessage("Step 1: To Remove")] }),
-    step2: async (state, config) => {
-      // Send message before persisting to state
-      // TODO: replace with `pushMessage` when part of 1.x
+    step2: async (state) => {
       const messages: BaseMessage[] = [
         ...state.messages
-          .filter((m) => m.getType() === "ai")
+          .filter((m) => AIMessage.isInstance(m))
           .map((m) => new RemoveMessage({ id: m.id! })),
         new AIMessage({ id: randomUUID(), content: "Step 2: To Keep" }),
       ];
 
-      const messagesHandler = (
-        config.callbacks as { handlers: object[] }
-      )?.handlers?.find(
-        (
-          cb
-        ): cb is {
-          _emit: (
-            chunk: [namespace: string[], metadata: Record<string, unknown>],
-            message: BaseMessage,
-            runId: string | undefined,
-            dedupe: boolean
-          ) => void;
-        } => "name" in cb && cb.name === "StreamMessagesHandler"
-      );
-
       for (const message of messages) {
-        messagesHandler?._emit([[], {}], message, undefined, false);
+        pushMessage(message, { stateKey: null });
       }
 
       return { messages };
@@ -1246,6 +1230,7 @@ describe("useStream", () => {
       const { submit, messages, isLoading } = useStream({
         assistantId: "removeMessageAgent",
         apiKey: "test-api-key",
+        throttle: false,
       });
 
       const rawMessages = messages.map((msg, i) => ({

libs/sdk/src/react/stream.custom.tsx

@@ -115,7 +115,10 @@ export function useStreamCustom<
 
   const [messageManager] = useState(() => new MessageTupleManager());
   const [stream] = useState(
-    () => new StreamManager<StateType, Bag>(messageManager)
+    () =>
+      new StreamManager<StateType, Bag>(messageManager, {
+        throttle: options.throttle ?? false,
+      })
   );
 
   useSyncExternalStore(

libs/sdk/src/react/stream.lgp.tsx

@@ -195,7 +195,10 @@ export function useStreamLGP<
 
   const [messageManager] = useState(() => new MessageTupleManager());
   const [stream] = useState(
-    () => new StreamManager<StateType, Bag>(messageManager)
+    () =>
+      new StreamManager<StateType, Bag>(messageManager, {
+        throttle: options.throttle ?? false,
+      })
   );
 
   useSyncExternalStore(

libs/sdk/src/react/types.tsx

@@ -287,6 +287,15 @@ export interface UseStreamOptions<
    * @experimental
    */
   experimental_thread?: UseStreamThread<StateType>;
+
+  /**
+   * Throttle the stream.
+   * If a number is provided, the stream will be throttled to the given number of milliseconds.
+   * If `true`, updates are batched in a single macrotask.
+   * If `false`, updates are not throttled or batched.
+   * @default true
+   */
+  throttle?: number | boolean;
 }
 
 interface RunMetadataStorage {
@@ -489,6 +498,7 @@ export type UseStreamCustomOptions<
   | "onTaskEvent"
   | "onStop"
   | "initialValues"
+  | "throttle"
 > & { transport: UseStreamTransport<StateType, Bag> };
 
 export type UseStreamCustom<

libs/sdk/src/ui/manager.ts

@@ -104,15 +104,21 @@ export class StreamManager<
 
   private listeners = new Set<() => void>();
 
+  private throttle: number | boolean;
+
   private state: {
     isLoading: boolean;
     values: [values: StateType, kind: "stream" | "stop"] | null;
     error: unknown;
   };
 
-  constructor(messages: MessageTupleManager) {
+  constructor(
+    messages: MessageTupleManager,
+    options: { throttle: number | boolean }
+  ) {
     this.messages = messages;
     this.state = { isLoading: false, values: null, error: undefined };
+    this.throttle = options.throttle;
   }
 
   private setState = (newState: Partial<typeof this.state>) => {
@@ -125,8 +131,27 @@ export class StreamManager<
   };
 
   subscribe = (listener: () => void): (() => void) => {
-    this.listeners.add(listener);
-    return () => this.listeners.delete(listener);
+    if (this.throttle === false) {
+      this.listeners.add(listener);
+      return () => this.listeners.delete(listener);
+    }
+
+    const timeoutMs = this.throttle === true ? 0 : this.throttle;
+    let timeoutId: NodeJS.Timeout | number | undefined;
+
+    const throttledListener = () => {
+      clearTimeout(timeoutId);
+      timeoutId = setTimeout(() => {
+        clearTimeout(timeoutId);
+        listener();
+      }, timeoutMs);
+    };
+
+    this.listeners.add(throttledListener);
+    return () => {
+      clearTimeout(timeoutId);
+      this.listeners.delete(throttledListener);
+    };
   };
 
   getSnapshot = () => this.state;