fix(sdk): useStream keeping stale thread history when switching during stream, no cancel when disconnecting resumable stream (#1677)

.changeset/clean-meals-agree.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph-sdk": patch
+---
+
+Fix `useStream()` keeping stale thread history when switching threads mid-stream (#1632)

.changeset/weak-phones-grab.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph-sdk": patch
+---
+
+Fix `stop()` behavior when cancelling a resumable stream via `useStream()` (#1610)

libs/sdk/src/react/manager.ts

@@ -227,6 +227,8 @@ export class StreamManager<
         | Promise<StateType | null | undefined | void>;
 
       onError: (error: unknown) => void | Promise<void>;
+
+      onFinish?: () => void;
     }
   ): Promise<void> => {
     if (this.state.isLoading) return;
@@ -330,6 +332,7 @@ export class StreamManager<
     } finally {
       this.setState({ isLoading: false });
       this.abortRef = new AbortController();
+      options.onFinish?.();
     }
   };
 
@@ -343,16 +346,21 @@ export class StreamManager<
       }) => void;
     }
   ): Promise<void> => {
-    if (this.abortRef) {
-      this.abortRef.abort();
-      this.abortRef = new AbortController();
-    }
+    this.abortRef.abort();
+    this.abortRef = new AbortController();
 
     options.onStop?.({ mutate: this.getMutateFn("stop", historyValues) });
   };
 
   clear = () => {
+    // Cancel any running streams
+    this.abortRef.abort();
+    this.abortRef = new AbortController();
+
+    // Set the stream state to null
     this.setState({ error: undefined, values: null });
+
+    // Clear any pending messages
     this.messages.clear();
   };
 }

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

@@ -2,7 +2,7 @@
 
 "use client";
 
-import { useState, useSyncExternalStore } from "react";
+import { useEffect, useRef, useState, useSyncExternalStore } from "react";
 import { EventStreamEvent, StreamManager } from "./manager.js";
 import type {
   BagTemplate,
@@ -21,6 +21,7 @@ import { MessageTupleManager } from "./messages.js";
 import { Interrupt } from "../schema.js";
 import { BytesLineDecoder, SSEDecoder } from "../utils/sse.js";
 import { IterableReadableStream } from "../utils/stream.js";
+import { useControllableThreadId } from "./thread.js";
 import { Command } from "../types.js";
 
 interface FetchStreamTransportOptions {
@@ -123,6 +124,17 @@ export function useStreamCustom<
     stream.getSnapshot
   );
 
+  const [threadId, onThreadId] = useControllableThreadId(options);
+  const threadIdRef = useRef<string | null>(threadId);
+
+  // Cancel the stream if thread ID has changed
+  useEffect(() => {
+    if (threadIdRef.current !== threadId) {
+      threadIdRef.current = threadId;
+      stream.clear();
+    }
+  }, [threadId, stream]);
+
   const getMessages = (value: StateType): Message[] => {
     const messagesKey = options.messagesKey ?? "messages";
     return Array.isArray(value[messagesKey]) ? value[messagesKey] : [];
@@ -142,6 +154,7 @@ export function useStreamCustom<
     submitOptions?: CustomSubmitOptions<StateType, ConfigurableType>
   ) => {
     let callbackMeta: RunCallbackMeta | undefined;
+    let usableThreadId = threadId;
 
     stream.setStreamValues(() => {
       if (submitOptions?.optimisticValues != null) {
@@ -157,15 +170,34 @@ export function useStreamCustom<
     });
 
     await stream.start(
-      async (signal: AbortSignal) =>
-        options.transport.stream({
+      async (signal: AbortSignal) => {
+        if (!usableThreadId) {
+          // generate random thread id
+          usableThreadId = crypto.randomUUID();
+          threadIdRef.current = usableThreadId;
+          onThreadId(usableThreadId);
+        }
+
+        if (!usableThreadId) {
+          throw new Error("Failed to obtain valid thread ID.");
+        }
+
+        return options.transport.stream({
           input: values,
           context: submitOptions?.context,
           command: submitOptions?.command,
           signal,
+          config: {
+            ...submitOptions?.config,
+            configurable: {
+              thread_id: usableThreadId,
+              ...submitOptions?.config?.configurable,
+            } as unknown as GetConfigurableType<Bag>,
+          },
         }) as Promise<
           AsyncGenerator<EventStreamEvent<StateType, UpdateType, CustomType>>
-        >,
+        >;
+      },
       {
         getMessages,
         setMessages,

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

@@ -32,6 +32,15 @@ import type { Message } from "../types.messages.js";
 import type { Interrupt, ThreadState } from "../schema.js";
 import type { StreamMode } from "../types.stream.js";
 import { MessageTupleManager } from "./messages.js";
+import { useControllableThreadId } from "./thread.js";
+
+function getFetchHistoryKey(
+  client: Client,
+  threadId: string | undefined | null,
+  limit: boolean | number
+) {
+  return [getClientConfigHash(client), threadId, limit].join(":");
+}
 
 function fetchHistory<StateType extends Record<string, unknown>>(
   client: Client,
@@ -50,99 +59,82 @@ function fetchHistory<StateType extends Record<string, unknown>>(
 }
 
 function useThreadHistory<StateType extends Record<string, unknown>>(
-  threadId: string | undefined | null,
   client: Client,
+  threadId: string | undefined | null,
   limit: boolean | number,
-  clearCallbackRef: RefObject<(() => void) | undefined>,
-  submittingRef: RefObject<boolean>,
-  onErrorRef: RefObject<((error: unknown) => void) | undefined>
+  options: {
+    submittingRef: RefObject<string | null>;
+    onError?: (error: unknown, run?: RunCallbackMeta) => void;
+  }
 ) {
-  const [history, setHistory] = useState<ThreadState<StateType>[] | undefined>(
-    undefined
-  );
-  const [isLoading, setIsLoading] = useState(() => {
-    if (threadId == null) return false;
-    return true;
-  });
-  const [error, setError] = useState<unknown | undefined>(undefined);
+  const key = getFetchHistoryKey(client, threadId, limit);
+  const [state, setState] = useState<{
+    data: ThreadState<StateType>[] | undefined;
+    error: unknown | undefined;
+    isLoading: boolean;
+  }>(() => ({
+    data: undefined,
+    error: undefined,
+    isLoading: threadId != null,
+  }));
 
-  const clientHash = getClientConfigHash(client);
   const clientRef = useRef(client);
   clientRef.current = client;
 
+  const onErrorRef = useRef(options?.onError);
+  onErrorRef.current = options?.onError;
+
   const fetcher = useCallback(
     (
-      threadId: string | undefined | null
+      threadId: string | undefined | null,
+      limit: boolean | number
     ): Promise<ThreadState<StateType>[]> => {
       if (threadId != null) {
         const client = clientRef.current;
 
-        setIsLoading(true);
-        return fetchHistory<StateType>(client, threadId, {
-          limit,
-        })
-          .then(
-            (history) => {
-              setHistory(history);
-              return history;
-            },
-            (error) => {
-              setError(error);
-              onErrorRef.current?.(error);
-              return Promise.reject(error);
-            }
-          )
-          .finally(() => {
-            setIsLoading(false);
-          });
+        setState((state) => ({ ...state, isLoading: true }));
+        return fetchHistory<StateType>(client, threadId, { limit }).then(
+          (data) => {
+            setState({ data, error: undefined, isLoading: false });
+            return data;
+          },
+          (error) => {
+            setState(({ data }) => ({ data, error, isLoading: false }));
+            onErrorRef.current?.(error);
+            return Promise.reject(error);
+          }
+        );
       }
 
-      setHistory(undefined);
-      setError(undefined);
-      setIsLoading(false);
-
-      clearCallbackRef.current?.();
+      setState({ data: undefined, error: undefined, isLoading: false });
       return Promise.resolve([]);
     },
-    [clearCallbackRef, onErrorRef, limit]
+    []
   );
 
   useEffect(() => {
-    if (submittingRef.current) return;
-    void fetcher(threadId);
-  }, [fetcher, submittingRef, clientHash, limit, threadId]);
+    // Skip if a stream is already in progress, no need to fetch history
+    if (
+      options.submittingRef.current != null &&
+      options.submittingRef.current === threadId
+    ) {
+      return;
+    }
+
+    void fetcher(threadId, limit);
+    // The `threadId` and `limit` arguments are already present in `key`
+    // Thus we don't need to include them in the dependency array
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [fetcher, key]);
 
   return {
-    data: history,
-    isLoading,
-    error,
-    mutate: (mutateId?: string) => fetcher(mutateId ?? threadId),
+    data: state.data,
+    error: state.error,
+    isLoading: state.isLoading,
+    mutate: (mutateId?: string) => fetcher(mutateId ?? threadId, limit),
   };
 }
 
-const useControllableThreadId = (options?: {
-  threadId?: string | null;
-  onThreadId?: (threadId: string) => void;
-}): [string | null, (threadId: string) => void] => {
-  const [localThreadId, _setLocalThreadId] = useState<string | null>(
-    options?.threadId ?? null
-  );
-
-  const onThreadIdRef = useRef(options?.onThreadId);
-  onThreadIdRef.current = options?.onThreadId;
-
-  const onThreadId = useCallback((threadId: string) => {
-    _setLocalThreadId(threadId);
-    onThreadIdRef.current?.(threadId);
-  }, []);
-
-  if (!options || !("threadId" in options)) {
-    return [localThreadId, onThreadId];
-  }
-
-  return [options.threadId ?? null, onThreadId];
-};
-
 export function useStreamLGP<
   StateType extends Record<string, unknown> = Record<string, unknown>,
   Bag extends {
@@ -236,13 +228,16 @@ export function useStreamLGP<
   const clearCallbackRef = useRef<() => void>(null!);
   clearCallbackRef.current = stream.clear;
 
-  const submittingRef = useRef(false);
-  submittingRef.current = stream.isLoading;
+  const threadIdRef = useRef<string | null>(threadId);
+  const threadIdStreamingRef = useRef<string | null>(null);
 
-  const onErrorRef = useRef<
-    ((error: unknown, run?: RunCallbackMeta) => void) | undefined
-  >(undefined);
-  onErrorRef.current = options.onError;
+  // Cancel the stream if thread ID has changed
+  useEffect(() => {
+    if (threadIdRef.current !== threadId) {
+      threadIdRef.current = threadId;
+      stream.clear();
+    }
+  }, [threadId, stream]);
 
   const historyLimit =
     typeof options.fetchStateHistory === "object" &&
@@ -250,14 +245,10 @@ export function useStreamLGP<
       ? options.fetchStateHistory.limit ?? false
       : options.fetchStateHistory ?? false;
 
-  const history = useThreadHistory<StateType>(
-    threadId,
-    client,
-    historyLimit,
-    clearCallbackRef,
-    submittingRef,
-    onErrorRef
-  );
+  const history = useThreadHistory<StateType>(client, threadId, historyLimit, {
+    submittingRef: threadIdStreamingRef,
+    onError: options.onError,
+  });
 
   const getMessages = (value: StateType): Message[] => {
     const messagesKey = options.messagesKey ?? "messages";
@@ -328,7 +319,18 @@ export function useStreamLGP<
     );
   })();
 
-  const stop = () => stream.stop(historyValues, { onStop: options.onStop });
+  const stop = () =>
+    stream.stop(historyValues, {
+      onStop: (args) => {
+        if (runMetadataStorage && threadId) {
+          const runId = runMetadataStorage.getItem(`lg:stream:${threadId}`);
+          if (runId) void client.runs.cancel(threadId, runId);
+          runMetadataStorage.removeItem(`lg:stream:${threadId}`);
+        }
+
+        options.onStop?.(args);
+      },
+    });
 
   // --- TRANSPORT ---
   const submit = async (
@@ -372,14 +374,24 @@ export function useStreamLGP<
             threadId: submitOptions?.threadId,
             metadata: submitOptions?.metadata,
           });
-          onThreadId(thread.thread_id);
+
           usableThreadId = thread.thread_id;
+
+          // Pre-emptively update the thread ID before
+          // stream cancellation is kicked off and thread
+          // is being refetched
+          threadIdRef.current = usableThreadId;
+          threadIdStreamingRef.current = usableThreadId;
+
+          onThreadId(usableThreadId);
         }
 
         if (!usableThreadId) {
           throw new Error("Failed to obtain valid thread ID.");
         }
 
+        threadIdStreamingRef.current = usableThreadId;
+
         const streamMode = unique([
           ...(submitOptions?.streamMode ?? []),
           ...trackStreamModeRef.current,
@@ -473,6 +485,9 @@ export function useStreamLGP<
         onError(error) {
           options.onError?.(error, callbackMeta);
         },
+        onFinish() {
+          threadIdStreamingRef.current = null;
+        },
       }
     );
   };
@@ -493,6 +508,7 @@ export function useStreamLGP<
 
     await stream.start(
       async (signal: AbortSignal) => {
+        threadIdStreamingRef.current = threadId;
         return client.runs.joinStream(threadId, runId, {
           signal,
           lastEventId,
@@ -516,6 +532,9 @@ export function useStreamLGP<
         onError(error) {
           options.onError?.(error, callbackMeta);
         },
+        onFinish() {
+          threadIdStreamingRef.current = null;
+        },
       }
     );
   };
@@ -547,7 +566,6 @@ export function useStreamLGP<
       void joinStreamRef.current?.(reconnectKey.runId);
     }
   }, [reconnectKey]);
-  // --- END TRANSPORT ---
 
   const error = stream.error ?? historyError ?? history.error;
   const values = stream.values ?? historyValues;

libs/sdk/src/react/thread.tsx

@@ -0,0 +1,26 @@
+"use client";
+
+import { useState, useRef, useCallback } from "react";
+
+export const useControllableThreadId = (options?: {
+  threadId?: string | null;
+  onThreadId?: (threadId: string) => void;
+}): [string | null, (threadId: string) => void] => {
+  const [localThreadId, _setLocalThreadId] = useState<string | null>(
+    options?.threadId ?? null
+  );
+
+  const onThreadIdRef = useRef(options?.onThreadId);
+  onThreadIdRef.current = options?.onThreadId;
+
+  const setThreadId = useCallback((threadId: string) => {
+    _setLocalThreadId(threadId);
+    onThreadIdRef.current?.(threadId);
+  }, []);
+
+  if (!options || !("threadId" in options)) {
+    return [localThreadId, setThreadId];
+  }
+
+  return [options.threadId ?? null, setThreadId];
+};

libs/sdk/src/react/types.tsx

@@ -449,6 +449,7 @@ export interface UseStreamTransport<
     input: GetUpdateType<Bag, StateType> | null | undefined;
     context: GetConfigurableType<Bag> | undefined;
     command: Command | undefined;
+    config: ConfigWithConfigurable<GetConfigurableType<Bag>> | undefined;
     signal: AbortSignal;
   }) => Promise<AsyncGenerator<{ id?: string; event: string; data: unknown }>>;
 }
@@ -459,6 +460,8 @@ export type UseStreamCustomOptions<
 > = Pick<
   UseStreamOptions<StateType, Bag>,
   | "messagesKey"
+  | "threadId"
+  | "onThreadId"
   | "onError"
   | "onCreated"
   | "onUpdateEvent"
@@ -490,5 +493,5 @@ export type CustomSubmitOptions<
   ConfigurableType extends Record<string, unknown> = Record<string, unknown>
 > = Pick<
   SubmitOptions<StateType, ConfigurableType>,
-  "optimisticValues" | "context" | "command"
+  "optimisticValues" | "context" | "command" | "config"
 >;