From 62e688bb3737c5831ed3d2279ca458c02ca304ab Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 15 May 2025 16:20:38 -0700
Subject: [PATCH 01/12] Add rejoining of streams

---
 libs/sdk-js/src/react/stream.tsx | 247 +++++++++++++++++++++++--------
 1 file changed, 184 insertions(+), 63 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 69c4a27a2..84f4bcf92 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -502,6 +502,9 @@ export interface UseStreamOptions<
    * Callback that is called when the thread ID is updated (ie when a new thread is created).
    */
   onThreadId?: (threadId: string) => void;
+
+  /** Will rejoin the stream on mount */
+  joinOnMount?: boolean;
 }
 
 export interface UseStream<
@@ -647,7 +650,7 @@ export function useStream<
     | ErrorStreamEvent
     | FeedbackStreamEvent;
 
-  let { assistantId, messagesKey, onError, onFinish } = options;
+  let { assistantId, messagesKey, onError, onFinish, joinOnMount } = options;
   messagesKey ??= "messages";
 
   const client = useMemo(
@@ -805,10 +808,10 @@ export function useStream<
     abortRef.current = null;
   }, []);
 
-  const submit = async (
-    values: UpdateType | null | undefined,
-    submitOptions?: SubmitOptions<StateType, ConfigurableType>,
-  ) => {
+  const join = async (metadata: {
+    runId: string;
+    threadId?: string | undefined | null;
+  }) => {
     try {
       setIsLoading(true);
       setStreamError(undefined);
@@ -816,65 +819,8 @@ export function useStream<
       submittingRef.current = true;
       abortRef.current = new AbortController();
 
-      // Unbranch things
-      const newPath = submitOptions?.checkpoint?.checkpoint_id
-        ? branchByCheckpoint[submitOptions?.checkpoint?.checkpoint_id]?.branch
-        : undefined;
-
-      if (newPath != null) setBranch(newPath ?? "");
-
-      // Assumption: we're setting the initial value
-      // Used for instant feedback
-      setStreamValues(() => {
-        const values = { ...historyValues };
-
-        if (submitOptions?.optimisticValues != null) {
-          return {
-            ...values,
-            ...(typeof submitOptions.optimisticValues === "function"
-              ? submitOptions.optimisticValues(values)
-              : submitOptions.optimisticValues),
-          };
-        }
-
-        return values;
-      });
-
-      let usableThreadId = threadId;
-      if (!usableThreadId) {
-        const thread = await client.threads.create();
-        onThreadId(thread.thread_id);
-        usableThreadId = thread.thread_id;
-      }
-
-      const streamMode = unique([
-        ...(submitOptions?.streamMode ?? []),
-        ...trackStreamModeRef.current,
-        ...callbackStreamMode,
-      ]);
-
-      const checkpoint =
-        submitOptions?.checkpoint ?? threadHead?.checkpoint ?? undefined;
-      // @ts-expect-error
-      if (checkpoint != null) delete checkpoint.thread_id;
-
-      const run = client.runs.stream(usableThreadId, assistantId, {
-        input: values as Record<string, unknown>,
-        config: submitOptions?.config,
-        command: submitOptions?.command,
-
-        interruptBefore: submitOptions?.interruptBefore,
-        interruptAfter: submitOptions?.interruptAfter,
-        metadata: submitOptions?.metadata,
-        multitaskStrategy: submitOptions?.multitaskStrategy,
-        onCompletion: submitOptions?.onCompletion,
-        onDisconnect: submitOptions?.onDisconnect ?? "cancel",
-
+      const run = client.runs.joinStream(metadata.threadId, metadata.runId, {
         signal: abortRef.current.signal,
-
-        checkpoint,
-        streamMode,
-        streamSubgraphs: submitOptions?.streamSubgraphs,
       }) as AsyncGenerator<EventStreamEvent>;
 
       let streamError: StreamError | undefined;
@@ -929,6 +875,181 @@ export function useStream<
         }
       }
 
+      // TODO: stream created checkpoints to avoid an unnecessary network request
+      if (metadata.threadId) {
+        const result = await history.mutate(metadata.threadId);
+        setStreamValues(null);
+
+        if (streamError != null) throw streamError;
+
+        const lastHead = result.at(0);
+        if (lastHead) onFinish?.(lastHead);
+      }
+    } catch (error) {
+      if (
+        !(
+          error instanceof Error &&
+          (error.name === "AbortError" || error.name === "TimeoutError")
+        )
+      ) {
+        console.error(error);
+        setStreamError(error);
+        onError?.(error);
+      }
+    } finally {
+      setIsLoading(false);
+
+      // Assumption: messages are already handled, we can clear the manager
+      messageManagerRef.current.clear();
+      submittingRef.current = false;
+      abortRef.current = null;
+    }
+  };
+
+  const joinRef = useRef<typeof join>(join);
+  joinRef.current = join;
+
+  const joinKey = useMemo(() => {
+    if (joinOnMount) return undefined;
+    return { runId: localStorage.get(`lg:rejoin:${threadId}`), threadId };
+  }, [joinOnMount, threadId]);
+
+  useEffect(() => {
+    if (joinKey) joinRef.current?.(joinKey);
+  }, [joinKey]);
+
+  const submit = async (
+    values: UpdateType | null | undefined,
+    submitOptions?: SubmitOptions<StateType, ConfigurableType>,
+  ) => {
+    try {
+      setIsLoading(true);
+      setStreamError(undefined);
+
+      submittingRef.current = true;
+      abortRef.current = new AbortController();
+
+      // Unbranch things
+      const newPath = submitOptions?.checkpoint?.checkpoint_id
+        ? branchByCheckpoint[submitOptions?.checkpoint?.checkpoint_id]?.branch
+        : undefined;
+
+      if (newPath != null) setBranch(newPath ?? "");
+
+      // Assumption: we're setting the initial value
+      // Used for instant feedback
+      setStreamValues(() => {
+        const values = { ...historyValues };
+
+        if (submitOptions?.optimisticValues != null) {
+          return {
+            ...values,
+            ...(typeof submitOptions.optimisticValues === "function"
+              ? submitOptions.optimisticValues(values)
+              : submitOptions.optimisticValues),
+          };
+        }
+
+        return values;
+      });
+
+      let usableThreadId = threadId;
+      if (!usableThreadId) {
+        const thread = await client.threads.create();
+        onThreadId(thread.thread_id);
+        usableThreadId = thread.thread_id;
+      }
+
+      const streamMode = unique([
+        ...(submitOptions?.streamMode ?? []),
+        ...trackStreamModeRef.current,
+        ...callbackStreamMode,
+      ]);
+
+      const checkpoint =
+        submitOptions?.checkpoint ?? threadHead?.checkpoint ?? undefined;
+      // @ts-expect-error
+      if (checkpoint != null) delete checkpoint.thread_id;
+      let rejoinKey: string | undefined;
+
+      const run = client.runs.stream(usableThreadId, assistantId, {
+        input: values as Record<string, unknown>,
+        config: submitOptions?.config,
+        command: submitOptions?.command,
+
+        interruptBefore: submitOptions?.interruptBefore,
+        interruptAfter: submitOptions?.interruptAfter,
+        metadata: submitOptions?.metadata,
+        multitaskStrategy: submitOptions?.multitaskStrategy,
+        onCompletion: submitOptions?.onCompletion,
+        onDisconnect: submitOptions?.onDisconnect ?? "cancel",
+
+        signal: abortRef.current.signal,
+
+        checkpoint,
+        streamMode,
+        streamSubgraphs: submitOptions?.streamSubgraphs,
+
+        onRunCreated(params) {
+          // rejoin the stream if needed
+          rejoinKey = `lg:rejoin:${params.thread_id ?? "temporary"}`;
+          window.localStorage.setItem(rejoinKey, params.run_id);
+        },
+      }) as AsyncGenerator<EventStreamEvent>;
+
+      let streamError: StreamError | undefined;
+      for await (const { event, data } of run) {
+        if (event === "error") {
+          streamError = new StreamError(data);
+          break;
+        }
+
+        if (event === "updates") options.onUpdateEvent?.(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 === "events") options.onLangChainEvent?.(data);
+        if (event === "debug") options.onDebugEvent?.(data);
+
+        if (event === "values") setStreamValues(data);
+        if (event === "messages") {
+          const [serialized] = data;
+
+          const messageId = messageManagerRef.current.add(serialized);
+          if (!messageId) {
+            console.warn(
+              "Failed to add message to manager, no message ID found",
+            );
+            continue;
+          }
+
+          setStreamValues((streamValues) => {
+            const values = { ...historyValues, ...streamValues };
+
+            // Assumption: we're concatenating the message
+            const messages = getMessages(values).slice();
+            const { chunk, index } =
+              messageManagerRef.current.get(messageId, messages.length) ?? {};
+
+            if (!chunk || index == null) return values;
+            messages[index] = toMessageDict(chunk);
+
+            return { ...values, [messagesKey!]: messages };
+          });
+        }
+      }
+      if (rejoinKey) window.localStorage.removeItem(rejoinKey);
+
       // TODO: stream created checkpoints to avoid an unnecessary network request
       const result = await history.mutate(usableThreadId);
       setStreamValues(null);

From 5c7ef9a4fcf53116d06d4949667a4c7a6862aebc Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Sat, 17 May 2025 21:30:27 +0100
Subject: [PATCH 02/12] Fix race condition

---
 libs/sdk-js/src/react/stream.tsx | 33 +++++++++++++++++++++++---------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 84f4bcf92..9acd62d56 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -808,10 +808,13 @@ export function useStream<
     abortRef.current = null;
   }, []);
 
-  const join = async (metadata: {
-    runId: string;
-    threadId?: string | undefined | null;
-  }) => {
+  const join = async (
+    metadata: {
+      runId: string;
+      threadId?: string | undefined | null;
+    },
+    lastEventId: (string & {}) | "-1",
+  ) => {
     try {
       setIsLoading(true);
       setStreamError(undefined);
@@ -819,8 +822,11 @@ export function useStream<
       submittingRef.current = true;
       abortRef.current = new AbortController();
 
+      const rejoinKey = `lg:rejoin:${metadata.threadId ?? "temporary"}`;
+
       const run = client.runs.joinStream(metadata.threadId, metadata.runId, {
         signal: abortRef.current.signal,
+        lastEventId,
       }) as AsyncGenerator<EventStreamEvent>;
 
       let streamError: StreamError | undefined;
@@ -875,6 +881,8 @@ export function useStream<
         }
       }
 
+      window.localStorage.removeItem(rejoinKey);
+
       // TODO: stream created checkpoints to avoid an unnecessary network request
       if (metadata.threadId) {
         const result = await history.mutate(metadata.threadId);
@@ -910,12 +918,17 @@ export function useStream<
   joinRef.current = join;
 
   const joinKey = useMemo(() => {
-    if (joinOnMount) return undefined;
-    return { runId: localStorage.get(`lg:rejoin:${threadId}`), threadId };
-  }, [joinOnMount, threadId]);
+    if (!joinOnMount || isLoading) return undefined;
+    if (typeof window === "undefined") return undefined;
+    const runId = window.localStorage.getItem(`lg:rejoin:${threadId}`);
+
+    if (!runId) return undefined;
+
+    return { runId, threadId };
+  }, [joinOnMount, isLoading, threadId]);
 
   useEffect(() => {
-    if (joinKey) joinRef.current?.(joinKey);
+    if (joinKey) joinRef.current?.(joinKey, "-1");
   }, [joinKey]);
 
   const submit = async (
@@ -982,7 +995,9 @@ export function useStream<
         metadata: submitOptions?.metadata,
         multitaskStrategy: submitOptions?.multitaskStrategy,
         onCompletion: submitOptions?.onCompletion,
-        onDisconnect: submitOptions?.onDisconnect ?? "cancel",
+        onDisconnect:
+          submitOptions?.onDisconnect ??
+          (options.joinOnMount ? "continue" : "cancel"),
 
         signal: abortRef.current.signal,
 

From 49746ab3a301e884b653140fb46fef20072239f7 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 17:29:54 +0200
Subject: [PATCH 03/12] Code cleanup

---
 libs/sdk-js/src/react/stream.tsx | 165 ++++++++++---------------------
 1 file changed, 52 insertions(+), 113 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 9acd62d56..40ab365d0 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -808,13 +808,12 @@ export function useStream<
     abortRef.current = null;
   }, []);
 
-  const join = async (
-    metadata: {
-      runId: string;
-      threadId?: string | undefined | null;
-    },
-    lastEventId: (string & {}) | "-1",
-  ) => {
+  async function consumeStream(
+    action: (signal: AbortSignal) => Promise<{
+      onFinish: () => Promise<ThreadState<StateType>[]>;
+      stream: AsyncGenerator<EventStreamEvent>;
+    }>,
+  ) {
     try {
       setIsLoading(true);
       setStreamError(undefined);
@@ -822,15 +821,10 @@ export function useStream<
       submittingRef.current = true;
       abortRef.current = new AbortController();
 
-      const rejoinKey = `lg:rejoin:${metadata.threadId ?? "temporary"}`;
-
-      const run = client.runs.joinStream(metadata.threadId, metadata.runId, {
-        signal: abortRef.current.signal,
-        lastEventId,
-      }) as AsyncGenerator<EventStreamEvent>;
+      const run = await action(abortRef.current.signal);
 
       let streamError: StreamError | undefined;
-      for await (const { event, data } of run) {
+      for await (const { event, data } of run.stream) {
         if (event === "error") {
           streamError = new StreamError(data);
           break;
@@ -879,15 +873,11 @@ export function useStream<
             return { ...values, [messagesKey!]: messages };
           });
         }
-      }
 
-      window.localStorage.removeItem(rejoinKey);
+        // TODO: stream created checkpoints to avoid an unnecessary network request
+        const result = await run.onFinish();
 
-      // TODO: stream created checkpoints to avoid an unnecessary network request
-      if (metadata.threadId) {
-        const result = await history.mutate(metadata.threadId);
         setStreamValues(null);
-
         if (streamError != null) throw streamError;
 
         const lastHead = result.at(0);
@@ -912,6 +902,33 @@ export function useStream<
       submittingRef.current = false;
       abortRef.current = null;
     }
+  }
+
+  const join = async (
+    metadata: {
+      runId: string;
+      threadId?: string | undefined | null;
+    },
+    lastEventId: (string & {}) | "-1",
+  ) => {
+    await consumeStream(async (signal: AbortSignal) => {
+      const usableThreadId = metadata.threadId;
+      const rejoinKey = `lg:rejoin:${usableThreadId ?? "temporary"}`;
+
+      const stream = client.runs.joinStream(usableThreadId, metadata.runId, {
+        signal,
+        lastEventId,
+      }) as AsyncGenerator<EventStreamEvent>;
+
+      return {
+        onFinish: async () => {
+          if (rejoinKey) window.localStorage.removeItem(rejoinKey);
+          if (!usableThreadId) return [];
+          return history.mutate(usableThreadId);
+        },
+        stream,
+      };
+    });
   };
 
   const joinRef = useRef<typeof join>(join);
@@ -935,13 +952,7 @@ export function useStream<
     values: UpdateType | null | undefined,
     submitOptions?: SubmitOptions<StateType, ConfigurableType>,
   ) => {
-    try {
-      setIsLoading(true);
-      setStreamError(undefined);
-
-      submittingRef.current = true;
-      abortRef.current = new AbortController();
-
+    await consumeStream(async (signal: AbortSignal) => {
       // Unbranch things
       const newPath = submitOptions?.checkpoint?.checkpoint_id
         ? branchByCheckpoint[submitOptions?.checkpoint?.checkpoint_id]?.branch
@@ -985,7 +996,7 @@ export function useStream<
       if (checkpoint != null) delete checkpoint.thread_id;
       let rejoinKey: string | undefined;
 
-      const run = client.runs.stream(usableThreadId, assistantId, {
+      const stream = client.runs.stream(usableThreadId, assistantId, {
         input: values as Record<string, unknown>,
         config: submitOptions?.config,
         command: submitOptions?.command,
@@ -999,99 +1010,27 @@ export function useStream<
           submitOptions?.onDisconnect ??
           (options.joinOnMount ? "continue" : "cancel"),
 
-        signal: abortRef.current.signal,
+        signal,
 
         checkpoint,
         streamMode,
         streamSubgraphs: submitOptions?.streamSubgraphs,
-
         onRunCreated(params) {
-          // rejoin the stream if needed
-          rejoinKey = `lg:rejoin:${params.thread_id ?? "temporary"}`;
-          window.localStorage.setItem(rejoinKey, params.run_id);
+          if (options.joinOnMount) {
+            rejoinKey = `lg:rejoin:${params.thread_id ?? "temporary"}`;
+            window.localStorage.setItem(rejoinKey, params.run_id);
+          }
         },
       }) as AsyncGenerator<EventStreamEvent>;
 
-      let streamError: StreamError | undefined;
-      for await (const { event, data } of run) {
-        if (event === "error") {
-          streamError = new StreamError(data);
-          break;
-        }
-
-        if (event === "updates") options.onUpdateEvent?.(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 === "events") options.onLangChainEvent?.(data);
-        if (event === "debug") options.onDebugEvent?.(data);
-
-        if (event === "values") setStreamValues(data);
-        if (event === "messages") {
-          const [serialized] = data;
-
-          const messageId = messageManagerRef.current.add(serialized);
-          if (!messageId) {
-            console.warn(
-              "Failed to add message to manager, no message ID found",
-            );
-            continue;
-          }
-
-          setStreamValues((streamValues) => {
-            const values = { ...historyValues, ...streamValues };
-
-            // Assumption: we're concatenating the message
-            const messages = getMessages(values).slice();
-            const { chunk, index } =
-              messageManagerRef.current.get(messageId, messages.length) ?? {};
-
-            if (!chunk || index == null) return values;
-            messages[index] = toMessageDict(chunk);
-
-            return { ...values, [messagesKey!]: messages };
-          });
-        }
-      }
-      if (rejoinKey) window.localStorage.removeItem(rejoinKey);
-
-      // TODO: stream created checkpoints to avoid an unnecessary network request
-      const result = await history.mutate(usableThreadId);
-      setStreamValues(null);
-
-      if (streamError != null) throw streamError;
-
-      const lastHead = result.at(0);
-      if (lastHead) onFinish?.(lastHead);
-    } catch (error) {
-      if (
-        !(
-          error instanceof Error &&
-          (error.name === "AbortError" || error.name === "TimeoutError")
-        )
-      ) {
-        console.error(error);
-        setStreamError(error);
-        onError?.(error);
-      }
-    } finally {
-      setIsLoading(false);
-
-      // Assumption: messages are already handled, we can clear the manager
-      messageManagerRef.current.clear();
-      submittingRef.current = false;
-      abortRef.current = null;
-    }
+      return {
+        stream,
+        onFinish: () => {
+          if (rejoinKey) window.localStorage.removeItem(rejoinKey);
+          return history.mutate(usableThreadId);
+        },
+      };
+    });
   };
 
   const error = streamError ?? historyError;

From 0db618ae75d1113264eeff472196dbf15616080b Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 17:37:37 +0200
Subject: [PATCH 04/12] Further cleanup

---
 libs/sdk-js/src/react/stream.tsx | 50 ++++++++++++++------------------
 1 file changed, 22 insertions(+), 28 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 40ab365d0..e2610974d 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -725,6 +725,7 @@ export function useStream<
 
   // TODO: this should be done on the server to avoid pagination
   // TODO: should we permit adapter? SWR / React Query?
+  // TODO: make this only when branching is expected
   const history = useThreadHistory<StateType>(
     threadId,
     client,
@@ -904,18 +905,14 @@ export function useStream<
     }
   }
 
-  const join = async (
-    metadata: {
-      runId: string;
-      threadId?: string | undefined | null;
-    },
+  const joinStream = async (
+    runId: string,
     lastEventId: (string & {}) | "-1",
   ) => {
+    if (!threadId) return;
     await consumeStream(async (signal: AbortSignal) => {
-      const usableThreadId = metadata.threadId;
-      const rejoinKey = `lg:rejoin:${usableThreadId ?? "temporary"}`;
-
-      const stream = client.runs.joinStream(usableThreadId, metadata.runId, {
+      const rejoinKey = `lg:rejoin:${threadId}`;
+      const stream = client.runs.joinStream(threadId, runId, {
         signal,
         lastEventId,
       }) as AsyncGenerator<EventStreamEvent>;
@@ -923,31 +920,13 @@ export function useStream<
       return {
         onFinish: async () => {
           if (rejoinKey) window.localStorage.removeItem(rejoinKey);
-          if (!usableThreadId) return [];
-          return history.mutate(usableThreadId);
+          return history.mutate(threadId);
         },
         stream,
       };
     });
   };
 
-  const joinRef = useRef<typeof join>(join);
-  joinRef.current = join;
-
-  const joinKey = useMemo(() => {
-    if (!joinOnMount || isLoading) return undefined;
-    if (typeof window === "undefined") return undefined;
-    const runId = window.localStorage.getItem(`lg:rejoin:${threadId}`);
-
-    if (!runId) return undefined;
-
-    return { runId, threadId };
-  }, [joinOnMount, isLoading, threadId]);
-
-  useEffect(() => {
-    if (joinKey) joinRef.current?.(joinKey, "-1");
-  }, [joinKey]);
-
   const submit = async (
     values: UpdateType | null | undefined,
     submitOptions?: SubmitOptions<StateType, ConfigurableType>,
@@ -1033,6 +1012,21 @@ export function useStream<
     });
   };
 
+  const joinStreamRef = useRef<typeof joinStream>(joinStream);
+  joinStreamRef.current = joinStream;
+
+  const joinKey = useMemo(() => {
+    if (!joinOnMount || isLoading) return undefined;
+    if (typeof window === "undefined") return undefined;
+    const runId = window.localStorage.getItem(`lg:rejoin:${threadId}`);
+    if (!runId) return undefined;
+    return { runId, threadId };
+  }, [joinOnMount, isLoading, threadId]);
+
+  useEffect(() => {
+    if (joinKey) joinStreamRef.current?.(joinKey.runId, "-1");
+  }, [joinKey]);
+
   const error = streamError ?? historyError;
   const values = streamValues ?? historyValues;
 

From 9d71d941fd00ff61cfda5c09edbe1ec2235392cb Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 17:44:37 +0200
Subject: [PATCH 05/12] resumable

---
 libs/sdk-js/src/react/stream.tsx | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index e2610974d..6d3f96b77 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -874,16 +874,16 @@ export function useStream<
             return { ...values, [messagesKey!]: messages };
           });
         }
-
-        // TODO: stream created checkpoints to avoid an unnecessary network request
-        const result = await run.onFinish();
-
-        setStreamValues(null);
-        if (streamError != null) throw streamError;
-
-        const lastHead = result.at(0);
-        if (lastHead) onFinish?.(lastHead);
       }
+
+      // TODO: stream created checkpoints to avoid an unnecessary network request
+      const result = await run.onFinish();
+
+      setStreamValues(null);
+      if (streamError != null) throw streamError;
+
+      const lastHead = result.at(0);
+      if (lastHead) onFinish?.(lastHead);
     } catch (error) {
       if (
         !(
@@ -994,6 +994,7 @@ export function useStream<
         checkpoint,
         streamMode,
         streamSubgraphs: submitOptions?.streamSubgraphs,
+        streamResumable: !!options.joinOnMount,
         onRunCreated(params) {
           if (options.joinOnMount) {
             rejoinKey = `lg:rejoin:${params.thread_id ?? "temporary"}`;

From f2c32727e83b97e6583810efa93a4d35b6fe2301 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 18:09:16 +0200
Subject: [PATCH 06/12] Fix invalid states

---
 libs/sdk-js/src/react/stream.tsx | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 6d3f96b77..3ffdf7b84 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -804,14 +804,21 @@ export function useStream<
     );
   })();
 
-  const stop = useCallback(() => {
+  const stop = () => {
     if (abortRef.current != null) abortRef.current.abort();
     abortRef.current = null;
-  }, []);
+
+    if (options.joinOnMount && threadId) {
+      const rejoinKey = `lg:rejoin:${threadId}`;
+      const runId = window.localStorage.getItem(rejoinKey);
+      if (runId) client.runs.cancel(threadId, runId);
+      window.localStorage.removeItem(rejoinKey);
+    }
+  };
 
   async function consumeStream(
     action: (signal: AbortSignal) => Promise<{
-      onFinish: () => Promise<ThreadState<StateType>[]>;
+      onSuccess: () => Promise<ThreadState<StateType>[]>;
       stream: AsyncGenerator<EventStreamEvent>;
     }>,
   ) {
@@ -877,7 +884,7 @@ export function useStream<
       }
 
       // TODO: stream created checkpoints to avoid an unnecessary network request
-      const result = await run.onFinish();
+      const result = await run.onSuccess();
 
       setStreamValues(null);
       if (streamError != null) throw streamError;
@@ -918,8 +925,8 @@ export function useStream<
       }) as AsyncGenerator<EventStreamEvent>;
 
       return {
-        onFinish: async () => {
-          if (rejoinKey) window.localStorage.removeItem(rejoinKey);
+        onSuccess: () => {
+          window.localStorage.removeItem(rejoinKey);
           return history.mutate(threadId);
         },
         stream,
@@ -1005,7 +1012,7 @@ export function useStream<
 
       return {
         stream,
-        onFinish: () => {
+        onSuccess: () => {
           if (rejoinKey) window.localStorage.removeItem(rejoinKey);
           return history.mutate(usableThreadId);
         },
@@ -1014,6 +1021,7 @@ export function useStream<
   };
 
   const joinStreamRef = useRef<typeof joinStream>(joinStream);
+  const autoJoinRef = useRef(options.joinOnMount);
   joinStreamRef.current = joinStream;
 
   const joinKey = useMemo(() => {
@@ -1025,7 +1033,10 @@ export function useStream<
   }, [joinOnMount, isLoading, threadId]);
 
   useEffect(() => {
-    if (joinKey) joinStreamRef.current?.(joinKey.runId, "-1");
+    if (joinKey && autoJoinRef.current) {
+      autoJoinRef.current = false;
+      joinStreamRef.current?.(joinKey.runId, "-1");
+    }
   }, [joinKey]);
 
   const error = streamError ?? historyError;

From 5823a659fc538e8c8c59e2763fd9a6f6d9e498a4 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 18:11:38 +0200
Subject: [PATCH 07/12] use sessionStorage instead

---
 libs/sdk-js/src/react/stream.tsx | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 3ffdf7b84..61f046e1e 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -810,9 +810,9 @@ export function useStream<
 
     if (options.joinOnMount && threadId) {
       const rejoinKey = `lg:rejoin:${threadId}`;
-      const runId = window.localStorage.getItem(rejoinKey);
+      const runId = window.sessionStorage.getItem(rejoinKey);
       if (runId) client.runs.cancel(threadId, runId);
-      window.localStorage.removeItem(rejoinKey);
+      window.sessionStorage.removeItem(rejoinKey);
     }
   };
 
@@ -926,7 +926,7 @@ export function useStream<
 
       return {
         onSuccess: () => {
-          window.localStorage.removeItem(rejoinKey);
+          window.sessionStorage.removeItem(rejoinKey);
           return history.mutate(threadId);
         },
         stream,
@@ -1005,7 +1005,7 @@ export function useStream<
         onRunCreated(params) {
           if (options.joinOnMount) {
             rejoinKey = `lg:rejoin:${params.thread_id ?? "temporary"}`;
-            window.localStorage.setItem(rejoinKey, params.run_id);
+            window.sessionStorage.setItem(rejoinKey, params.run_id);
           }
         },
       }) as AsyncGenerator<EventStreamEvent>;
@@ -1013,7 +1013,7 @@ export function useStream<
       return {
         stream,
         onSuccess: () => {
-          if (rejoinKey) window.localStorage.removeItem(rejoinKey);
+          if (rejoinKey) window.sessionStorage.removeItem(rejoinKey);
           return history.mutate(usableThreadId);
         },
       };
@@ -1027,7 +1027,7 @@ export function useStream<
   const joinKey = useMemo(() => {
     if (!joinOnMount || isLoading) return undefined;
     if (typeof window === "undefined") return undefined;
-    const runId = window.localStorage.getItem(`lg:rejoin:${threadId}`);
+    const runId = window.sessionStorage.getItem(`lg:rejoin:${threadId}`);
     if (!runId) return undefined;
     return { runId, threadId };
   }, [joinOnMount, isLoading, threadId]);

From f5a2410ddcae4747f219941292c4d097391616ce Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 18:38:09 +0200
Subject: [PATCH 08/12] Make useStream more hackable

---
 libs/sdk-js/src/react/stream.tsx | 58 +++++++++++++++++++++++---------
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 61f046e1e..897095b78 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -457,6 +457,11 @@ export interface UseStreamOptions<
    */
   onFinish?: (state: ThreadState<StateType>) => void;
 
+  /**
+   * Callback that is called when a new stream is created.
+   */
+  onCreated?: (run: { run_id: string; thread_id: string }) => void;
+
   /**
    * Callback that is called when an update event is received.
    */
@@ -504,7 +509,13 @@ export interface UseStreamOptions<
   onThreadId?: (threadId: string) => void;
 
   /** Will rejoin the stream on mount */
-  joinOnMount?: boolean;
+  joinOnMount?: boolean | (() => RunMetadataStorage);
+}
+
+interface RunMetadataStorage {
+  getItem(key: `lg:rejoin:${string}`): string | null;
+  setItem(key: `lg:rejoin:${string}`, value: string): void;
+  removeItem(key: `lg:rejoin:${string}`): void;
 }
 
 export interface UseStream<
@@ -622,6 +633,7 @@ interface SubmitOptions<
    * @default false
    */
   streamSubgraphs?: boolean;
+  streamResumable?: boolean;
 }
 
 export function useStream<
@@ -650,7 +662,17 @@ export function useStream<
     | ErrorStreamEvent
     | FeedbackStreamEvent;
 
-  let { assistantId, messagesKey, onError, onFinish, joinOnMount } = options;
+  let { assistantId, messagesKey, onCreated, onError, onFinish } = options;
+
+  const joinOnMountRef = useRef(options.joinOnMount);
+  const runMetadataStorage = useMemo(() => {
+    if (typeof window === "undefined") return null;
+    const joinOnMount = joinOnMountRef.current;
+    if (joinOnMount === true) return window.sessionStorage;
+    if (typeof joinOnMount === "function") return joinOnMount();
+    return null;
+  }, []);
+
   messagesKey ??= "messages";
 
   const client = useMemo(
@@ -808,11 +830,10 @@ export function useStream<
     if (abortRef.current != null) abortRef.current.abort();
     abortRef.current = null;
 
-    if (options.joinOnMount && threadId) {
-      const rejoinKey = `lg:rejoin:${threadId}`;
-      const runId = window.sessionStorage.getItem(rejoinKey);
+    if (runMetadataStorage && threadId) {
+      const runId = runMetadataStorage.getItem(`lg:rejoin:${threadId}`);
       if (runId) client.runs.cancel(threadId, runId);
-      window.sessionStorage.removeItem(rejoinKey);
+      runMetadataStorage.removeItem(`lg:rejoin:${threadId}`);
     }
   };
 
@@ -918,7 +939,6 @@ export function useStream<
   ) => {
     if (!threadId) return;
     await consumeStream(async (signal: AbortSignal) => {
-      const rejoinKey = `lg:rejoin:${threadId}`;
       const stream = client.runs.joinStream(threadId, runId, {
         signal,
         lastEventId,
@@ -926,7 +946,7 @@ export function useStream<
 
       return {
         onSuccess: () => {
-          window.sessionStorage.removeItem(rejoinKey);
+          runMetadataStorage?.removeItem(`lg:rejoin:${threadId}`);
           return history.mutate(threadId);
         },
         stream,
@@ -980,7 +1000,7 @@ export function useStream<
         submitOptions?.checkpoint ?? threadHead?.checkpoint ?? undefined;
       // @ts-expect-error
       if (checkpoint != null) delete checkpoint.thread_id;
-      let rejoinKey: string | undefined;
+      let rejoinKey: `lg:rejoin:${string}` | undefined;
 
       const stream = client.runs.stream(usableThreadId, assistantId, {
         input: values as Record<string, unknown>,
@@ -994,26 +1014,31 @@ export function useStream<
         onCompletion: submitOptions?.onCompletion,
         onDisconnect:
           submitOptions?.onDisconnect ??
-          (options.joinOnMount ? "continue" : "cancel"),
+          (runMetadataStorage ? "continue" : "cancel"),
 
         signal,
 
         checkpoint,
         streamMode,
         streamSubgraphs: submitOptions?.streamSubgraphs,
-        streamResumable: !!options.joinOnMount,
+        streamResumable: submitOptions?.streamResumable ?? !!runMetadataStorage,
         onRunCreated(params) {
-          if (options.joinOnMount) {
-            rejoinKey = `lg:rejoin:${params.thread_id ?? "temporary"}`;
-            window.sessionStorage.setItem(rejoinKey, params.run_id);
+          const runParams = {
+            run_id: params.run_id,
+            thread_id: params.thread_id ?? usableThreadId,
+          };
+          if (runMetadataStorage) {
+            rejoinKey = `lg:rejoin:${runParams.thread_id}`;
+            runMetadataStorage.setItem(rejoinKey, runParams.run_id);
           }
+          onCreated?.(runParams);
         },
       }) as AsyncGenerator<EventStreamEvent>;
 
       return {
         stream,
         onSuccess: () => {
-          if (rejoinKey) window.sessionStorage.removeItem(rejoinKey);
+          if (rejoinKey) runMetadataStorage?.removeItem(rejoinKey);
           return history.mutate(usableThreadId);
         },
       };
@@ -1021,7 +1046,8 @@ export function useStream<
   };
 
   const joinStreamRef = useRef<typeof joinStream>(joinStream);
-  const autoJoinRef = useRef(options.joinOnMount);
+  const joinOnMount = !!runMetadataStorage;
+  const autoJoinRef = useRef(joinOnMount);
   joinStreamRef.current = joinStream;
 
   const joinKey = useMemo(() => {

From 74dbccd408df7c63e00a2896b1b86053467161f5 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 18:42:38 +0200
Subject: [PATCH 09/12] Rename to reconnectOnMount

---
 libs/sdk-js/src/react/stream.tsx | 48 ++++++++++++++++----------------
 1 file changed, 24 insertions(+), 24 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 897095b78..9ef5def25 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -508,14 +508,14 @@ export interface UseStreamOptions<
    */
   onThreadId?: (threadId: string) => void;
 
-  /** Will rejoin the stream on mount */
-  joinOnMount?: boolean | (() => RunMetadataStorage);
+  /** Will reconnect the stream on mount */
+  reconnectOnMount?: boolean | (() => RunMetadataStorage);
 }
 
 interface RunMetadataStorage {
-  getItem(key: `lg:rejoin:${string}`): string | null;
-  setItem(key: `lg:rejoin:${string}`, value: string): void;
-  removeItem(key: `lg:rejoin:${string}`): void;
+  getItem(key: `lg:stream:${string}`): string | null;
+  setItem(key: `lg:stream:${string}`, value: string): void;
+  removeItem(key: `lg:stream:${string}`): void;
 }
 
 export interface UseStream<
@@ -664,12 +664,12 @@ export function useStream<
 
   let { assistantId, messagesKey, onCreated, onError, onFinish } = options;
 
-  const joinOnMountRef = useRef(options.joinOnMount);
+  const reconnectOnMountRef = useRef(options.reconnectOnMount);
   const runMetadataStorage = useMemo(() => {
     if (typeof window === "undefined") return null;
-    const joinOnMount = joinOnMountRef.current;
-    if (joinOnMount === true) return window.sessionStorage;
-    if (typeof joinOnMount === "function") return joinOnMount();
+    const storage = reconnectOnMountRef.current;
+    if (storage === true) return window.sessionStorage;
+    if (typeof storage === "function") return storage();
     return null;
   }, []);
 
@@ -831,9 +831,9 @@ export function useStream<
     abortRef.current = null;
 
     if (runMetadataStorage && threadId) {
-      const runId = runMetadataStorage.getItem(`lg:rejoin:${threadId}`);
+      const runId = runMetadataStorage.getItem(`lg:stream:${threadId}`);
       if (runId) client.runs.cancel(threadId, runId);
-      runMetadataStorage.removeItem(`lg:rejoin:${threadId}`);
+      runMetadataStorage.removeItem(`lg:stream:${threadId}`);
     }
   };
 
@@ -946,7 +946,7 @@ export function useStream<
 
       return {
         onSuccess: () => {
-          runMetadataStorage?.removeItem(`lg:rejoin:${threadId}`);
+          runMetadataStorage?.removeItem(`lg:stream:${threadId}`);
           return history.mutate(threadId);
         },
         stream,
@@ -1000,7 +1000,7 @@ export function useStream<
         submitOptions?.checkpoint ?? threadHead?.checkpoint ?? undefined;
       // @ts-expect-error
       if (checkpoint != null) delete checkpoint.thread_id;
-      let rejoinKey: `lg:rejoin:${string}` | undefined;
+      let rejoinKey: `lg:stream:${string}` | undefined;
 
       const stream = client.runs.stream(usableThreadId, assistantId, {
         input: values as Record<string, unknown>,
@@ -1028,7 +1028,7 @@ export function useStream<
             thread_id: params.thread_id ?? usableThreadId,
           };
           if (runMetadataStorage) {
-            rejoinKey = `lg:rejoin:${runParams.thread_id}`;
+            rejoinKey = `lg:stream:${runParams.thread_id}`;
             runMetadataStorage.setItem(rejoinKey, runParams.run_id);
           }
           onCreated?.(runParams);
@@ -1046,24 +1046,24 @@ export function useStream<
   };
 
   const joinStreamRef = useRef<typeof joinStream>(joinStream);
-  const joinOnMount = !!runMetadataStorage;
-  const autoJoinRef = useRef(joinOnMount);
+  const shouldReconnect = !!runMetadataStorage;
+  const reconnectRef = useRef(shouldReconnect);
   joinStreamRef.current = joinStream;
 
-  const joinKey = useMemo(() => {
-    if (!joinOnMount || isLoading) return undefined;
+  const reconnectKey = useMemo(() => {
+    if (!shouldReconnect || isLoading) return undefined;
     if (typeof window === "undefined") return undefined;
-    const runId = window.sessionStorage.getItem(`lg:rejoin:${threadId}`);
+    const runId = window.sessionStorage.getItem(`lg:stream:${threadId}`);
     if (!runId) return undefined;
     return { runId, threadId };
-  }, [joinOnMount, isLoading, threadId]);
+  }, [shouldReconnect, isLoading, threadId]);
 
   useEffect(() => {
-    if (joinKey && autoJoinRef.current) {
-      autoJoinRef.current = false;
-      joinStreamRef.current?.(joinKey.runId, "-1");
+    if (reconnectKey && reconnectRef.current) {
+      reconnectRef.current = false;
+      joinStreamRef.current?.(reconnectKey.runId, "-1");
     }
-  }, [joinKey]);
+  }, [reconnectKey]);
 
   const error = streamError ?? historyError;
   const values = streamValues ?? historyValues;

From 102b1f63f3313f5b2515d1e315f6c64cc14a9b1f Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 19:01:14 +0200
Subject: [PATCH 10/12] Fix assumption on sessionStorage

---
 libs/sdk-js/src/react/stream.tsx | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 9ef5def25..cb3e5f3bf 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -1045,22 +1045,30 @@ export function useStream<
     });
   };
 
-  const joinStreamRef = useRef<typeof joinStream>(joinStream);
-  const shouldReconnect = !!runMetadataStorage;
-  const reconnectRef = useRef(shouldReconnect);
-  joinStreamRef.current = joinStream;
-
   const reconnectKey = useMemo(() => {
-    if (!shouldReconnect || isLoading) return undefined;
+    if (!runMetadataStorage || isLoading) return undefined;
     if (typeof window === "undefined") return undefined;
-    const runId = window.sessionStorage.getItem(`lg:stream:${threadId}`);
+    const runId = runMetadataStorage?.getItem(`lg:stream:${threadId}`);
     if (!runId) return undefined;
     return { runId, threadId };
-  }, [shouldReconnect, isLoading, threadId]);
+  }, [runMetadataStorage, isLoading, threadId]);
+
+  const shouldReconnect = !!runMetadataStorage;
+  const reconnectRef = useRef({ threadId, shouldReconnect });
+
+  const joinStreamRef = useRef<typeof joinStream>(joinStream);
+  joinStreamRef.current = joinStream;
 
   useEffect(() => {
-    if (reconnectKey && reconnectRef.current) {
-      reconnectRef.current = false;
+    // reset shouldReconnect when switching threads
+    if (reconnectRef.current.threadId !== threadId) {
+      reconnectRef.current = { threadId, shouldReconnect };
+    }
+  }, [threadId, shouldReconnect]);
+
+  useEffect(() => {
+    if (reconnectKey && reconnectRef.current.shouldReconnect) {
+      reconnectRef.current.shouldReconnect = false;
       joinStreamRef.current?.(reconnectKey.runId, "-1");
     }
   }, [reconnectKey]);

From 1ff5e2966b7000bc1f8fa2bebfc432fd9cc7dce0 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 19:07:01 +0200
Subject: [PATCH 11/12] Add joinStream method

---
 libs/sdk-js/src/react/stream.tsx | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index cb3e5f3bf..ba65c50c9 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -604,6 +604,11 @@ export interface UseStream<
    * The ID of the assistant to use.
    */
   assistantId: string;
+
+  /**
+   * Join an active stream.
+   */
+  joinStream: (runId: string) => Promise<void>;
 }
 
 type ConfigWithConfigurable<ConfigurableType extends Record<string, unknown>> =
@@ -933,10 +938,8 @@ export function useStream<
     }
   }
 
-  const joinStream = async (
-    runId: string,
-    lastEventId: (string & {}) | "-1",
-  ) => {
+  const joinStream = async (runId: string, lastEventId?: string) => {
+    lastEventId ??= "-1";
     if (!threadId) return;
     await consumeStream(async (signal: AbortSignal) => {
       const stream = client.runs.joinStream(threadId, runId, {
@@ -1069,7 +1072,7 @@ export function useStream<
   useEffect(() => {
     if (reconnectKey && reconnectRef.current.shouldReconnect) {
       reconnectRef.current.shouldReconnect = false;
-      joinStreamRef.current?.(reconnectKey.runId, "-1");
+      joinStreamRef.current?.(reconnectKey.runId);
     }
   }, [reconnectKey]);
 
@@ -1091,6 +1094,8 @@ export function useStream<
     stop,
     submit,
 
+    joinStream,
+
     branch,
     setBranch,
 

From 74b18acfc2d1e851e27ce8cb442941662c82355a Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 22 May 2025 19:08:39 +0200
Subject: [PATCH 12/12] Bump to 0.0.78

---
 libs/sdk-js/package.json | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/sdk-js/package.json b/libs/sdk-js/package.json
index d1e72bf60..795131974 100644
--- a/libs/sdk-js/package.json
+++ b/libs/sdk-js/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.77",
+  "version": "0.0.78",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",