@langchain/langgraph-sdk 0.0.40 → 0.0.42

package/dist/react/stream.cjs

@@ -58,6 +58,92 @@ function findLastIndex(array, predicate) {
     }
     return -1;
 }
+function getBranchSequence(history) {
+    const childrenMap = {};
+    // First pass - collect nodes for each checkpoint
+    history.forEach((state) => {
+        const checkpointId = state.parent_checkpoint?.checkpoint_id ?? "$";
+        childrenMap[checkpointId] ??= [];
+        childrenMap[checkpointId].push(state);
+    });
+    const rootSequence = { type: "sequence", items: [] };
+    const queue = [{ id: "$", sequence: rootSequence, path: [] }];
+    const paths = [];
+    const visited = new Set();
+    while (queue.length > 0) {
+        const task = queue.shift();
+        if (visited.has(task.id))
+            continue;
+        visited.add(task.id);
+        const children = childrenMap[task.id];
+        if (children == null || children.length === 0)
+            continue;
+        // If we've encountered a fork (2+ children), push the fork
+        // to the sequence and add a new sequence for each child
+        let fork;
+        if (children.length > 1) {
+            fork = { type: "fork", items: [] };
+            task.sequence.items.push(fork);
+        }
+        for (const value of children) {
+            const id = value.checkpoint.checkpoint_id;
+            let sequence = task.sequence;
+            let path = task.path;
+            if (fork != null) {
+                sequence = { type: "sequence", items: [] };
+                fork.items.unshift(sequence);
+                path = path.slice();
+                path.push(id);
+                paths.push(path);
+            }
+            sequence.items.push({ type: "node", value, path });
+            queue.push({ id, sequence, path });
+        }
+    }
+    return { rootSequence, paths };
+}
+const PATH_SEP = ">";
+const ROOT_ID = "$";
+// Get flat view
+function getBranchView(sequence, paths, branch) {
+    const path = branch.split(PATH_SEP);
+    const pathMap = {};
+    for (const path of paths) {
+        const parent = path.at(-2) ?? ROOT_ID;
+        pathMap[parent] ??= [];
+        pathMap[parent].unshift(path);
+    }
+    const history = [];
+    const branchByCheckpoint = {};
+    const forkStack = path.slice();
+    const queue = [...sequence.items];
+    while (queue.length > 0) {
+        const item = queue.shift();
+        if (item.type === "node") {
+            history.push(item.value);
+            branchByCheckpoint[item.value.checkpoint.checkpoint_id] = {
+                branch: item.path.join(PATH_SEP),
+                branchOptions: (item.path.length > 0
+                    ? pathMap[item.path.at(-2) ?? ROOT_ID] ?? []
+                    : []).map((p) => p.join(PATH_SEP)),
+            };
+        }
+        if (item.type === "fork") {
+            const forkId = forkStack.shift();
+            const index = forkId != null
+                ? item.items.findIndex((value) => {
+                    const firstItem = value.items.at(0);
+                    if (!firstItem || firstItem.type !== "node")
+                        return false;
+                    return firstItem.value.checkpoint.checkpoint_id === forkId;
+                })
+                : -1;
+            const nextItems = item.items.at(index)?.items ?? [];
+            queue.push(...nextItems);
+        }
+    }
+    return { history, branchByCheckpoint };
+}
 function fetchHistory(client, threadId) {
     return client.threads.getHistory(threadId, { limit: 1000 });
 }
@@ -84,12 +170,26 @@ function useThreadHistory(threadId, client, clearCallbackRef, submittingRef) {
         mutate: (mutateId) => fetcher(mutateId ?? threadId),
     };
 }
+const useControllableThreadId = (options) => {
+    const [localThreadId, _setLocalThreadId] = (0, react_1.useState)(options?.threadId ?? null);
+    const onThreadIdRef = (0, react_1.useRef)(options?.onThreadId);
+    onThreadIdRef.current = options?.onThreadId;
+    const onThreadId = (0, react_1.useCallback)((threadId) => {
+        _setLocalThreadId(threadId);
+        onThreadIdRef.current?.(threadId);
+    }, []);
+    if (typeof options?.threadId === "undefined") {
+        return [localThreadId, onThreadId];
+    }
+    return [options.threadId, onThreadId];
+};
 function useStream(options) {
-    const { assistantId, threadId, withMessages, onError, onFinish } = options;
+    let { assistantId, messagesKey, onError, onFinish } = options;
+    messagesKey ??= "messages";
     const client = (0, react_1.useMemo)(() => new client_js_1.Client({ apiUrl: options.apiUrl, apiKey: options.apiKey }), [options.apiKey, options.apiUrl]);
-    const [branchPath, setBranchPath] = (0, react_1.useState)([]);
+    const [threadId, onThreadId] = useControllableThreadId(options);
+    const [branch, setBranch] = (0, react_1.useState)("");
     const [isLoading, setIsLoading] = (0, react_1.useState)(false);
-    const [_, setEvents] = (0, react_1.useState)([]);
     const [streamError, setStreamError] = (0, react_1.useState)(undefined);
     const [streamValues, setStreamValues] = (0, react_1.useState)(null);
     const messageManagerRef = (0, react_1.useRef)(new MessageTupleManager());
@@ -119,94 +219,13 @@ function useStream(options) {
     // TODO: should we permit adapter? SWR / React Query?
     const history = useThreadHistory(threadId, client, clearCallbackRef, submittingRef);
     const getMessages = (0, react_1.useMemo)(() => {
-        if (withMessages == null)
-            return undefined;
-        return (value) => Array.isArray(value[withMessages])
-            ? value[withMessages]
+        return (value) => Array.isArray(value[messagesKey])
+            ? value[messagesKey]
             : [];
-    }, [withMessages]);
-    const [sequence, pathMap] = (() => {
-        const childrenMap = {};
-        // First pass - collect nodes for each checkpoint
-        history.data.forEach((state) => {
-            const checkpointId = state.parent_checkpoint?.checkpoint_id ?? "$";
-            childrenMap[checkpointId] ??= [];
-            childrenMap[checkpointId].push(state);
-        });
-        const rootSequence = { type: "sequence", items: [] };
-        const queue = [{ id: "$", sequence: rootSequence, path: [] }];
-        const paths = [];
-        const visited = new Set();
-        while (queue.length > 0) {
-            const task = queue.shift();
-            if (visited.has(task.id))
-                continue;
-            visited.add(task.id);
-            const children = childrenMap[task.id];
-            if (children == null || children.length === 0)
-                continue;
-            // If we've encountered a fork (2+ children), push the fork
-            // to the sequence and add a new sequence for each child
-            let fork;
-            if (children.length > 1) {
-                fork = { type: "fork", items: [] };
-                task.sequence.items.push(fork);
-            }
-            for (const value of children) {
-                const id = value.checkpoint.checkpoint_id;
-                let sequence = task.sequence;
-                let path = task.path;
-                if (fork != null) {
-                    sequence = { type: "sequence", items: [] };
-                    fork.items.unshift(sequence);
-                    path = path.slice();
-                    path.push(id);
-                    paths.push(path);
-                }
-                sequence.items.push({ type: "node", value, path });
-                queue.push({ id, sequence, path });
-            }
-        }
-        // Third pass, create a map for available forks
-        const pathMap = {};
-        for (const path of paths) {
-            const parent = path.at(-2) ?? "$";
-            pathMap[parent] ??= [];
-            pathMap[parent].unshift(path);
-        }
-        return [rootSequence, pathMap];
-    })();
-    const [flatValues, flatPaths] = (() => {
-        const result = [];
-        const flatPaths = {};
-        const forkStack = branchPath.slice();
-        const queue = [...sequence.items];
-        while (queue.length > 0) {
-            const item = queue.shift();
-            if (item.type === "node") {
-                result.push(item.value);
-                flatPaths[item.value.checkpoint.checkpoint_id] = {
-                    current: item.path,
-                    branches: item.path.length > 0 ? pathMap[item.path.at(-2) ?? "$"] ?? [] : [],
-                };
-            }
-            if (item.type === "fork") {
-                const forkId = forkStack.shift();
-                const index = forkId != null
-                    ? item.items.findIndex((value) => {
-                        const firstItem = value.items.at(0);
-                        if (!firstItem || firstItem.type !== "node")
-                            return false;
-                        return firstItem.value.checkpoint.checkpoint_id === forkId;
-                    })
-                    : -1;
-                const nextItems = item.items.at(index)?.items ?? [];
-                queue.push(...nextItems);
-            }
-        }
-        return [result, flatPaths];
-    })();
-    const threadHead = flatValues.at(-1);
+    }, [messagesKey]);
+    const { rootSequence, paths } = getBranchSequence(history.data);
+    const { history: flatHistory, branchByCheckpoint } = getBranchView(rootSequence, paths, branch);
+    const threadHead = flatHistory.at(-1);
     const historyValues = threadHead?.values ?? {};
     const historyError = (() => {
         const error = threadHead?.tasks?.at(-1)?.error;
@@ -225,8 +244,6 @@ function useStream(options) {
         return error;
     })();
     const messageMetadata = (() => {
-        if (getMessages == null)
-            return undefined;
         const alreadyShown = new Set();
         return getMessages(historyValues).map((message, idx) => {
             const messageId = message.id ?? idx;
@@ -235,12 +252,12 @@ function useStream(options) {
                 .includes(messageId));
             const firstSeen = history.data[firstSeenIdx];
             let branch = firstSeen
-                ? flatPaths[firstSeen.checkpoint.checkpoint_id]
+                ? branchByCheckpoint[firstSeen.checkpoint.checkpoint_id]
                 : undefined;
-            if (!branch?.current?.length)
+            if (!branch?.branch?.length)
                 branch = undefined;
             // serialize branches
-            const optionsShown = branch?.branches?.flat(2).join(",");
+            const optionsShown = branch?.branchOptions?.flat(2).join(",");
             if (optionsShown) {
                 if (alreadyShown.has(optionsShown))
                     branch = undefined;
@@ -249,8 +266,8 @@ function useStream(options) {
             return {
                 messageId: messageId.toString(),
                 firstSeenState: firstSeen,
-                branch: branch?.current?.join(">"),
-                branchOptions: branch?.branches?.map((b) => b.join(">")),
+                branch: branch?.branch,
+                branchOptions: branch?.branchOptions,
             };
         });
     })();
@@ -268,7 +285,7 @@ function useStream(options) {
             let usableThreadId = threadId;
             if (!usableThreadId) {
                 const thread = await client.threads.create();
-                options?.onThreadId?.(thread.thread_id);
+                onThreadId(thread.thread_id);
                 usableThreadId = thread.thread_id;
             }
             const streamMode = unique([
@@ -296,10 +313,10 @@ function useStream(options) {
             }));
             // Unbranch things
             const newPath = submitOptions?.checkpoint?.checkpoint_id
-                ? flatPaths[submitOptions?.checkpoint?.checkpoint_id]?.current
+                ? branchByCheckpoint[submitOptions?.checkpoint?.checkpoint_id]?.branch
                 : undefined;
             if (newPath != null)
-                setBranchPath(newPath ?? []);
+                setBranch(newPath ?? "");
             // Assumption: we're setting the initial value
             // Used for instant feedback
             setStreamValues(() => {
@@ -316,26 +333,19 @@ function useStream(options) {
             });
             let streamError;
             for await (const { event, data } of run) {
-                setEvents((events) => [...events, { event, data }]);
                 if (event === "error") {
                     streamError = new StreamError(data);
                     break;
                 }
-                if (event === "updates") {
+                if (event === "updates")
                     options.onUpdateEvent?.(data);
-                }
-                if (event === "custom") {
+                if (event === "custom")
                     options.onCustomEvent?.(data);
-                }
-                if (event === "metadata") {
+                if (event === "metadata")
                     options.onMetadataEvent?.(data);
-                }
-                if (event === "values") {
+                if (event === "values")
                     setStreamValues(data);
-                }
                 if (event === "messages") {
-                    if (!getMessages)
-                        continue;
                     const [serialized] = data;
                     const messageId = messageManagerRef.current.add(serialized);
                     if (!messageId) {
@@ -350,13 +360,12 @@ function useStream(options) {
                         if (!chunk || index == null)
                             return values;
                         messages[index] = toMessageDict(chunk);
-                        return { ...values, [withMessages]: messages };
+                        return { ...values, [messagesKey]: messages };
                     });
                 }
             }
             // TODO: stream created checkpoints to avoid an unnecessary network request
             const result = await history.mutate(usableThreadId);
-            // TODO: write tests verifying that stream values are properly handled lifecycle-wise
             setStreamValues(null);
             if (streamError != null)
                 throw streamError;
@@ -381,7 +390,6 @@ function useStream(options) {
     };
     const error = isLoading ? streamError : historyError;
     const values = streamValues ?? historyValues;
-    const setBranch = (0, react_1.useCallback)((path) => setBranchPath(path.split(">")), [setBranchPath]);
     return {
         get values() {
             trackStreamMode("values");
@@ -391,19 +399,16 @@ function useStream(options) {
         isLoading,
         stop,
         submit,
+        branch,
         setBranch,
+        history: flatHistory,
+        experimental_branchTree: rootSequence,
         get messages() {
             trackStreamMode("messages-tuple");
-            if (getMessages == null) {
-                throw new Error("No messages key provided. Make sure that `useStream` contains the `messagesKey` property.");
-            }
             return getMessages(values);
         },
         getMessagesMetadata(message, index) {
             trackStreamMode("messages-tuple");
-            if (getMessages == null) {
-                throw new Error("No messages key provided. Make sure that `useStream` contains the `messagesKey` property.");
-            }
             return messageMetadata?.find((m) => m.messageId === (message.id ?? index));
         },
     };

package/dist/react/stream.d.ts

@@ -3,44 +3,153 @@ import type { Command, DisconnectMode, MultitaskStrategy, OnCompletionBehavior }
 import type { Message } from "../types.messages.js";
 import type { Checkpoint, Config, Metadata, ThreadState } from "../schema.js";
 import type { CustomStreamEvent, MetadataStreamEvent, StreamMode, UpdatesStreamEvent } from "../types.stream.js";
+interface Node<StateType = any> {
+    type: "node";
+    value: ThreadState<StateType>;
+    path: string[];
+}
+interface Fork<StateType = any> {
+    type: "fork";
+    items: Array<Sequence<StateType>>;
+}
+interface Sequence<StateType = any> {
+    type: "sequence";
+    items: Array<Node<StateType> | Fork<StateType>>;
+}
 export type MessageMetadata<StateType extends Record<string, unknown>> = {
+    /**
+     * The ID of the message used.
+     */
     messageId: string;
+    /**
+     * The first thread state the message was seen in.
+     */
     firstSeenState: ThreadState<StateType> | undefined;
+    /**
+     * The branch of the message.
+     */
     branch: string | undefined;
+    /**
+     * The list of branches this message is part of.
+     * This is useful for displaying branching controls.
+     */
     branchOptions: string[] | undefined;
 };
-export declare function useStream<StateType extends Record<string, unknown> = Record<string, unknown>, UpdateType extends Record<string, unknown> = Partial<StateType>, CustomType = unknown>(options: {
+interface UseStreamOptions<StateType extends Record<string, unknown> = Record<string, unknown>, UpdateType extends Record<string, unknown> = Partial<StateType>, CustomType = unknown> {
+    /**
+     * The ID of the assistant to use.
+     */
     assistantId: string;
+    /**
+     * The URL of the API to use.
+     */
     apiUrl: ClientConfig["apiUrl"];
+    /**
+     * The API key to use.
+     */
     apiKey?: ClientConfig["apiKey"];
-    withMessages?: string;
+    /**
+     * Specify the key within the state that contains messages.
+     * Defaults to "messages".
+     *
+     * @default "messages"
+     */
+    messagesKey?: string;
+    /**
+     * Callback that is called when an error occurs.
+     */
     onError?: (error: unknown) => void;
+    /**
+     * Callback that is called when the stream is finished.
+     */
     onFinish?: (state: ThreadState<StateType>) => void;
+    /**
+     * Callback that is called when an update event is received.
+     */
     onUpdateEvent?: (data: UpdatesStreamEvent<UpdateType>["data"]) => void;
+    /**
+     * Callback that is called when a custom event is received.
+     */
     onCustomEvent?: (data: CustomStreamEvent<CustomType>["data"]) => void;
+    /**
+     * Callback that is called when a metadata event is received.
+     */
     onMetadataEvent?: (data: MetadataStreamEvent["data"]) => void;
+    /**
+     * The ID of the thread to fetch history and current values from.
+     */
     threadId?: string | null;
+    /**
+     * Callback that is called when the thread ID is updated (ie when a new thread is created).
+     */
     onThreadId?: (threadId: string) => void;
-}): {
-    readonly values: StateType;
+}
+interface UseStream<StateType extends Record<string, unknown> = Record<string, unknown>, UpdateType extends Record<string, unknown> = Partial<StateType>> {
+    /**
+     * The current values of the thread.
+     */
+    values: StateType;
+    /**
+     * Last seen error from the thread or during streaming.
+     */
     error: unknown;
+    /**
+     * Whether the stream is currently running.
+     */
     isLoading: boolean;
+    /**
+     * Stops the stream.
+     */
     stop: () => void;
-    submit: (values: UpdateType | undefined, submitOptions?: {
-        config?: Config;
-        checkpoint?: Omit<Checkpoint, "thread_id"> | null;
-        command?: Command;
-        interruptBefore?: "*" | string[];
-        interruptAfter?: "*" | string[];
-        metadata?: Metadata;
-        multitaskStrategy?: MultitaskStrategy;
-        onCompletion?: OnCompletionBehavior;
-        onDisconnect?: DisconnectMode;
-        feedbackKeys?: string[];
-        streamMode?: Array<StreamMode>;
-        optimisticValues?: Partial<StateType> | ((prev: StateType) => Partial<StateType>);
-    }) => Promise<void>;
-    setBranch: (path: string) => void;
-    readonly messages: Message[];
-    getMessagesMetadata(message: Message, index?: number): MessageMetadata<StateType> | undefined;
-};
+    /**
+     * Create and stream a run to the thread.
+     */
+    submit: (values: UpdateType, options?: SubmitOptions<StateType>) => void;
+    /**
+     * The current branch of the thread.
+     */
+    branch: string;
+    /**
+     * Set the branch of the thread.
+     */
+    setBranch: (branch: string) => void;
+    /**
+     * Flattened history of thread states of a thread.
+     */
+    history: ThreadState<StateType>[];
+    /**
+     * Tree of all branches for the thread.
+     * @experimental
+     */
+    experimental_branchTree: Sequence<StateType>;
+    /**
+     * Messages inferred from the thread.
+     * Will automatically update with incoming message chunks.
+     */
+    messages: Message[];
+    /**
+     * Get the metadata for a message, such as first thread state the message
+     * was seen in and branch information.
+     
+     * @param message - The message to get the metadata for.
+     * @param index - The index of the message in the thread.
+     * @returns The metadata for the message.
+     */
+    getMessagesMetadata: (message: Message, index?: number) => MessageMetadata<StateType> | undefined;
+}
+interface SubmitOptions<StateType extends Record<string, unknown> = Record<string, unknown>> {
+    config?: Config;
+    checkpoint?: Omit<Checkpoint, "thread_id"> | null;
+    command?: Command;
+    interruptBefore?: "*" | string[];
+    interruptAfter?: "*" | string[];
+    metadata?: Metadata;
+    multitaskStrategy?: MultitaskStrategy;
+    onCompletion?: OnCompletionBehavior;
+    onDisconnect?: DisconnectMode;
+    feedbackKeys?: string[];
+    streamMode?: Array<StreamMode>;
+    optimisticValues?: Partial<StateType> | ((prev: StateType) => Partial<StateType>);
+}
+export declare function useStream<StateType extends Record<string, unknown> = Record<string, unknown>, UpdateType extends Record<string, unknown> = Partial<StateType>, CustomType = unknown>(options: UseStreamOptions<StateType, UpdateType, CustomType>): UseStream<StateType, UpdateType>;
+export {};

package/dist/react/stream.js

@@ -55,6 +55,92 @@ function findLastIndex(array, predicate) {
     }
     return -1;
 }
+function getBranchSequence(history) {
+    const childrenMap = {};
+    // First pass - collect nodes for each checkpoint
+    history.forEach((state) => {
+        const checkpointId = state.parent_checkpoint?.checkpoint_id ?? "$";
+        childrenMap[checkpointId] ??= [];
+        childrenMap[checkpointId].push(state);
+    });
+    const rootSequence = { type: "sequence", items: [] };
+    const queue = [{ id: "$", sequence: rootSequence, path: [] }];
+    const paths = [];
+    const visited = new Set();
+    while (queue.length > 0) {
+        const task = queue.shift();
+        if (visited.has(task.id))
+            continue;
+        visited.add(task.id);
+        const children = childrenMap[task.id];
+        if (children == null || children.length === 0)
+            continue;
+        // If we've encountered a fork (2+ children), push the fork
+        // to the sequence and add a new sequence for each child
+        let fork;
+        if (children.length > 1) {
+            fork = { type: "fork", items: [] };
+            task.sequence.items.push(fork);
+        }
+        for (const value of children) {
+            const id = value.checkpoint.checkpoint_id;
+            let sequence = task.sequence;
+            let path = task.path;
+            if (fork != null) {
+                sequence = { type: "sequence", items: [] };
+                fork.items.unshift(sequence);
+                path = path.slice();
+                path.push(id);
+                paths.push(path);
+            }
+            sequence.items.push({ type: "node", value, path });
+            queue.push({ id, sequence, path });
+        }
+    }
+    return { rootSequence, paths };
+}
+const PATH_SEP = ">";
+const ROOT_ID = "$";
+// Get flat view
+function getBranchView(sequence, paths, branch) {
+    const path = branch.split(PATH_SEP);
+    const pathMap = {};
+    for (const path of paths) {
+        const parent = path.at(-2) ?? ROOT_ID;
+        pathMap[parent] ??= [];
+        pathMap[parent].unshift(path);
+    }
+    const history = [];
+    const branchByCheckpoint = {};
+    const forkStack = path.slice();
+    const queue = [...sequence.items];
+    while (queue.length > 0) {
+        const item = queue.shift();
+        if (item.type === "node") {
+            history.push(item.value);
+            branchByCheckpoint[item.value.checkpoint.checkpoint_id] = {
+                branch: item.path.join(PATH_SEP),
+                branchOptions: (item.path.length > 0
+                    ? pathMap[item.path.at(-2) ?? ROOT_ID] ?? []
+                    : []).map((p) => p.join(PATH_SEP)),
+            };
+        }
+        if (item.type === "fork") {
+            const forkId = forkStack.shift();
+            const index = forkId != null
+                ? item.items.findIndex((value) => {
+                    const firstItem = value.items.at(0);
+                    if (!firstItem || firstItem.type !== "node")
+                        return false;
+                    return firstItem.value.checkpoint.checkpoint_id === forkId;
+                })
+                : -1;
+            const nextItems = item.items.at(index)?.items ?? [];
+            queue.push(...nextItems);
+        }
+    }
+    return { history, branchByCheckpoint };
+}
 function fetchHistory(client, threadId) {
     return client.threads.getHistory(threadId, { limit: 1000 });
 }
@@ -81,12 +167,26 @@ function useThreadHistory(threadId, client, clearCallbackRef, submittingRef) {
         mutate: (mutateId) => fetcher(mutateId ?? threadId),
     };
 }
+const useControllableThreadId = (options) => {
+    const [localThreadId, _setLocalThreadId] = useState(options?.threadId ?? null);
+    const onThreadIdRef = useRef(options?.onThreadId);
+    onThreadIdRef.current = options?.onThreadId;
+    const onThreadId = useCallback((threadId) => {
+        _setLocalThreadId(threadId);
+        onThreadIdRef.current?.(threadId);
+    }, []);
+    if (typeof options?.threadId === "undefined") {
+        return [localThreadId, onThreadId];
+    }
+    return [options.threadId, onThreadId];
+};
 export function useStream(options) {
-    const { assistantId, threadId, withMessages, onError, onFinish } = options;
+    let { assistantId, messagesKey, onError, onFinish } = options;
+    messagesKey ??= "messages";
     const client = useMemo(() => new Client({ apiUrl: options.apiUrl, apiKey: options.apiKey }), [options.apiKey, options.apiUrl]);
-    const [branchPath, setBranchPath] = useState([]);
+    const [threadId, onThreadId] = useControllableThreadId(options);
+    const [branch, setBranch] = useState("");
     const [isLoading, setIsLoading] = useState(false);
-    const [_, setEvents] = useState([]);
     const [streamError, setStreamError] = useState(undefined);
     const [streamValues, setStreamValues] = useState(null);
     const messageManagerRef = useRef(new MessageTupleManager());
@@ -116,94 +216,13 @@ export function useStream(options) {
     // TODO: should we permit adapter? SWR / React Query?
     const history = useThreadHistory(threadId, client, clearCallbackRef, submittingRef);
     const getMessages = useMemo(() => {
-        if (withMessages == null)
-            return undefined;
-        return (value) => Array.isArray(value[withMessages])
-            ? value[withMessages]
+        return (value) => Array.isArray(value[messagesKey])
+            ? value[messagesKey]
             : [];
-    }, [withMessages]);
-    const [sequence, pathMap] = (() => {
-        const childrenMap = {};
-        // First pass - collect nodes for each checkpoint
-        history.data.forEach((state) => {
-            const checkpointId = state.parent_checkpoint?.checkpoint_id ?? "$";
-            childrenMap[checkpointId] ??= [];
-            childrenMap[checkpointId].push(state);
-        });
-        const rootSequence = { type: "sequence", items: [] };
-        const queue = [{ id: "$", sequence: rootSequence, path: [] }];
-        const paths = [];
-        const visited = new Set();
-        while (queue.length > 0) {
-            const task = queue.shift();
-            if (visited.has(task.id))
-                continue;
-            visited.add(task.id);
-            const children = childrenMap[task.id];
-            if (children == null || children.length === 0)
-                continue;
-            // If we've encountered a fork (2+ children), push the fork
-            // to the sequence and add a new sequence for each child
-            let fork;
-            if (children.length > 1) {
-                fork = { type: "fork", items: [] };
-                task.sequence.items.push(fork);
-            }
-            for (const value of children) {
-                const id = value.checkpoint.checkpoint_id;
-                let sequence = task.sequence;
-                let path = task.path;
-                if (fork != null) {
-                    sequence = { type: "sequence", items: [] };
-                    fork.items.unshift(sequence);
-                    path = path.slice();
-                    path.push(id);
-                    paths.push(path);
-                }
-                sequence.items.push({ type: "node", value, path });
-                queue.push({ id, sequence, path });
-            }
-        }
-        // Third pass, create a map for available forks
-        const pathMap = {};
-        for (const path of paths) {
-            const parent = path.at(-2) ?? "$";
-            pathMap[parent] ??= [];
-            pathMap[parent].unshift(path);
-        }
-        return [rootSequence, pathMap];
-    })();
-    const [flatValues, flatPaths] = (() => {
-        const result = [];
-        const flatPaths = {};
-        const forkStack = branchPath.slice();
-        const queue = [...sequence.items];
-        while (queue.length > 0) {
-            const item = queue.shift();
-            if (item.type === "node") {
-                result.push(item.value);
-                flatPaths[item.value.checkpoint.checkpoint_id] = {
-                    current: item.path,
-                    branches: item.path.length > 0 ? pathMap[item.path.at(-2) ?? "$"] ?? [] : [],
-                };
-            }
-            if (item.type === "fork") {
-                const forkId = forkStack.shift();
-                const index = forkId != null
-                    ? item.items.findIndex((value) => {
-                        const firstItem = value.items.at(0);
-                        if (!firstItem || firstItem.type !== "node")
-                            return false;
-                        return firstItem.value.checkpoint.checkpoint_id === forkId;
-                    })
-                    : -1;
-                const nextItems = item.items.at(index)?.items ?? [];
-                queue.push(...nextItems);
-            }
-        }
-        return [result, flatPaths];
-    })();
-    const threadHead = flatValues.at(-1);
+    }, [messagesKey]);
+    const { rootSequence, paths } = getBranchSequence(history.data);
+    const { history: flatHistory, branchByCheckpoint } = getBranchView(rootSequence, paths, branch);
+    const threadHead = flatHistory.at(-1);
     const historyValues = threadHead?.values ?? {};
     const historyError = (() => {
         const error = threadHead?.tasks?.at(-1)?.error;
@@ -222,8 +241,6 @@ export function useStream(options) {
         return error;
     })();
     const messageMetadata = (() => {
-        if (getMessages == null)
-            return undefined;
         const alreadyShown = new Set();
         return getMessages(historyValues).map((message, idx) => {
             const messageId = message.id ?? idx;
@@ -232,12 +249,12 @@ export function useStream(options) {
                 .includes(messageId));
             const firstSeen = history.data[firstSeenIdx];
             let branch = firstSeen
-                ? flatPaths[firstSeen.checkpoint.checkpoint_id]
+                ? branchByCheckpoint[firstSeen.checkpoint.checkpoint_id]
                 : undefined;
-            if (!branch?.current?.length)
+            if (!branch?.branch?.length)
                 branch = undefined;
             // serialize branches
-            const optionsShown = branch?.branches?.flat(2).join(",");
+            const optionsShown = branch?.branchOptions?.flat(2).join(",");
             if (optionsShown) {
                 if (alreadyShown.has(optionsShown))
                     branch = undefined;
@@ -246,8 +263,8 @@ export function useStream(options) {
             return {
                 messageId: messageId.toString(),
                 firstSeenState: firstSeen,
-                branch: branch?.current?.join(">"),
-                branchOptions: branch?.branches?.map((b) => b.join(">")),
+                branch: branch?.branch,
+                branchOptions: branch?.branchOptions,
             };
         });
     })();
@@ -265,7 +282,7 @@ export function useStream(options) {
             let usableThreadId = threadId;
             if (!usableThreadId) {
                 const thread = await client.threads.create();
-                options?.onThreadId?.(thread.thread_id);
+                onThreadId(thread.thread_id);
                 usableThreadId = thread.thread_id;
             }
             const streamMode = unique([
@@ -293,10 +310,10 @@ export function useStream(options) {
             }));
             // Unbranch things
             const newPath = submitOptions?.checkpoint?.checkpoint_id
-                ? flatPaths[submitOptions?.checkpoint?.checkpoint_id]?.current
+                ? branchByCheckpoint[submitOptions?.checkpoint?.checkpoint_id]?.branch
                 : undefined;
             if (newPath != null)
-                setBranchPath(newPath ?? []);
+                setBranch(newPath ?? "");
             // Assumption: we're setting the initial value
             // Used for instant feedback
             setStreamValues(() => {
@@ -313,26 +330,19 @@ export function useStream(options) {
             });
             let streamError;
             for await (const { event, data } of run) {
-                setEvents((events) => [...events, { event, data }]);
                 if (event === "error") {
                     streamError = new StreamError(data);
                     break;
                 }
-                if (event === "updates") {
+                if (event === "updates")
                     options.onUpdateEvent?.(data);
-                }
-                if (event === "custom") {
+                if (event === "custom")
                     options.onCustomEvent?.(data);
-                }
-                if (event === "metadata") {
+                if (event === "metadata")
                     options.onMetadataEvent?.(data);
-                }
-                if (event === "values") {
+                if (event === "values")
                     setStreamValues(data);
-                }
                 if (event === "messages") {
-                    if (!getMessages)
-                        continue;
                     const [serialized] = data;
                     const messageId = messageManagerRef.current.add(serialized);
                     if (!messageId) {
@@ -347,13 +357,12 @@ export function useStream(options) {
                         if (!chunk || index == null)
                             return values;
                         messages[index] = toMessageDict(chunk);
-                        return { ...values, [withMessages]: messages };
+                        return { ...values, [messagesKey]: messages };
                     });
                 }
             }
             // TODO: stream created checkpoints to avoid an unnecessary network request
             const result = await history.mutate(usableThreadId);
-            // TODO: write tests verifying that stream values are properly handled lifecycle-wise
             setStreamValues(null);
             if (streamError != null)
                 throw streamError;
@@ -378,7 +387,6 @@ export function useStream(options) {
     };
     const error = isLoading ? streamError : historyError;
     const values = streamValues ?? historyValues;
-    const setBranch = useCallback((path) => setBranchPath(path.split(">")), [setBranchPath]);
     return {
         get values() {
             trackStreamMode("values");
@@ -388,19 +396,16 @@ export function useStream(options) {
         isLoading,
         stop,
         submit,
+        branch,
         setBranch,
+        history: flatHistory,
+        experimental_branchTree: rootSequence,
         get messages() {
             trackStreamMode("messages-tuple");
-            if (getMessages == null) {
-                throw new Error("No messages key provided. Make sure that `useStream` contains the `messagesKey` property.");
-            }
             return getMessages(values);
         },
         getMessagesMetadata(message, index) {
             trackStreamMode("messages-tuple");
-            if (getMessages == null) {
-                throw new Error("No messages key provided. Make sure that `useStream` contains the `messagesKey` property.");
-            }
             return messageMetadata?.find((m) => m.messageId === (message.id ?? index));
         },
     };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.40",
+  "version": "0.0.42",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",
@@ -10,7 +10,8 @@
     "prepublish": "yarn run build",
     "format": "prettier --write src",
     "lint": "prettier --check src && tsc --noEmit",
-    "test": "NODE_OPTIONS=--experimental-vm-modules jest --testPathIgnorePatterns=\\.int\\.test.ts"
+    "test": "NODE_OPTIONS=--experimental-vm-modules jest --testPathIgnorePatterns=\\.int\\.test.ts",
+    "typedoc": "typedoc && typedoc src/react/index.ts --out docs/react --options typedoc.react.json"
   },
   "main": "index.js",
   "license": "MIT",
@@ -33,8 +34,8 @@
     "jest": "^29.7.0",
     "prettier": "^3.2.5",
     "ts-jest": "^29.1.2",
-    "typedoc": "^0.26.1",
-    "typedoc-plugin-markdown": "^4.1.0",
+    "typedoc": "^0.27.7",
+    "typedoc-plugin-markdown": "^4.4.2",
     "typescript": "^5.4.5",
     "react": "^18.3.1"
   },