@langchain/langgraph-sdk 0.0.112 → 0.1.0

package/dist/react/stream.js

@@ -1,190 +1,12 @@
 /* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
 "use client";
-import { useCallback, useEffect, useMemo, useRef, useState, } from "react";
-import { coerceMessageLikeToMessage, convertToChunk, isBaseMessageChunk, } from "@langchain/core/messages";
+import { useCallback, useEffect, useMemo, useRef, useState, useSyncExternalStore, } from "react";
+import { findLast, unique } from "./utils.js";
+import { StreamError } from "./errors.js";
+import { getBranchContext } from "./branching.js";
+import { StreamManager } from "./manager.js";
 import { Client, getClientConfigHash } from "../client.js";
-class StreamError extends Error {
-    constructor(data) {
-        super(data.message);
-        this.name = data.name ?? data.error ?? "StreamError";
-    }
-    static isStructuredError(error) {
-        return typeof error === "object" && error != null && "message" in error;
-    }
-}
-function tryConvertToChunk(message) {
-    try {
-        return convertToChunk(message);
-    }
-    catch {
-        return null;
-    }
-}
-class MessageTupleManager {
-    constructor() {
-        Object.defineProperty(this, "chunks", {
-            enumerable: true,
-            configurable: true,
-            writable: true,
-            value: {}
-        });
-        this.chunks = {};
-    }
-    add(serialized, metadata) {
-        // TODO: this is sometimes sent from the API
-        // figure out how to prevent this or move this to LC.js
-        if (serialized.type.endsWith("MessageChunk")) {
-            // eslint-disable-next-line no-param-reassign
-            serialized.type = serialized.type
-                .slice(0, -"MessageChunk".length)
-                .toLowerCase();
-        }
-        const message = coerceMessageLikeToMessage(serialized);
-        const chunk = tryConvertToChunk(message);
-        const { id } = chunk ?? message;
-        if (!id) {
-            console.warn("No message ID found for chunk, ignoring in state", serialized);
-            return null;
-        }
-        this.chunks[id] ??= {};
-        this.chunks[id].metadata = metadata ?? this.chunks[id].metadata;
-        if (chunk) {
-            const prev = this.chunks[id].chunk;
-            this.chunks[id].chunk =
-                (isBaseMessageChunk(prev) ? prev : null)?.concat(chunk) ?? chunk;
-        }
-        else {
-            this.chunks[id].chunk = message;
-        }
-        return id;
-    }
-    clear() {
-        this.chunks = {};
-    }
-    get(id, defaultIndex) {
-        if (this.chunks[id] == null)
-            return null;
-        if (defaultIndex != null)
-            this.chunks[id].index ??= defaultIndex;
-        return this.chunks[id];
-    }
-}
-const toMessageDict = (chunk) => {
-    const { type, data } = chunk.toDict();
-    return { ...data, type };
-};
-function unique(array) {
-    return [...new Set(array)];
-}
-function findLastIndex(array, predicate) {
-    for (let i = array.length - 1; i >= 0; i -= 1) {
-        if (predicate(array[i]))
-            return i;
-    }
-    return -1;
-}
-function getBranchSequence(history) {
-    const childrenMap = {};
-    // Short circuit if there's only a singular one state
-    // TODO: I think we can make this more generalizable for all `fetchStateHistory` values.
-    if (history.length <= 1) {
-        return {
-            rootSequence: {
-                type: "sequence",
-                items: history.map((value) => ({ type: "node", value, path: [] })),
-            },
-            paths: [],
-        };
-    }
-    // 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;
-            if (id == null)
-                continue;
-            let { sequence } = task;
-            let { path } = task;
-            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);
-            const checkpointId = item.value.checkpoint?.checkpoint_id;
-            if (checkpointId == null)
-                continue;
-            branchByCheckpoint[checkpointId] = {
-                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 };
-}
+import { MessageTupleManager } from "./messages.js";
 function fetchHistory(client, threadId, options) {
     if (options?.limit === false) {
         return client.threads.getState(threadId).then((state) => {
@@ -193,7 +15,7 @@ function fetchHistory(client, threadId, options) {
             return [state];
         });
     }
-    const limit = typeof options?.limit === "number" ? options.limit : 1000;
+    const limit = typeof options?.limit === "number" ? options.limit : 10;
     return client.threads.getHistory(threadId, { limit });
 }
 function useThreadHistory(threadId, client, limit, clearCallbackRef, submittingRef, onErrorRef) {
@@ -257,39 +79,7 @@ const useControllableThreadId = (options) => {
     }
     return [options.threadId ?? null, onThreadId];
 };
-function useStreamValuesState() {
-    const [values, setValues] = useState(null);
-    const setStreamValues = useCallback((values, kind = "stream") => {
-        if (typeof values === "function") {
-            setValues((prevTuple) => {
-                const [prevValues, prevKind] = prevTuple ?? [null, "stream"];
-                const next = values(prevValues, prevKind);
-                if (next == null)
-                    return null;
-                return [next, kind];
-            });
-            return;
-        }
-        if (values == null)
-            setValues(null);
-        setValues([values, kind]);
-    }, []);
-    const mutate = useCallback((kind, serverValues) => (update) => {
-        setStreamValues((clientValues) => {
-            const prev = { ...serverValues, ...clientValues };
-            const next = typeof update === "function" ? update(prev) : update;
-            return { ...prev, ...next };
-        }, kind);
-    }, [setStreamValues]);
-    return [values?.[0] ?? null, setStreamValues, mutate];
-}
 export function useStream(options) {
-    const matchEventType = (expected, actual, _data) => {
-        return expected === actual || actual.startsWith(`${expected}|`);
-    };
-    let { messagesKey } = options;
-    const { assistantId, fetchStateHistory } = options;
-    const { onCreated, onError, onFinish } = options;
     const reconnectOnMountRef = useRef(options.reconnectOnMount);
     const runMetadataStorage = useMemo(() => {
         if (typeof window === "undefined")
@@ -301,7 +91,6 @@ export function useStream(options) {
             return storage();
         return null;
     }, []);
-    messagesKey ??= "messages";
     const client = useMemo(() => options.client ??
         new Client({
             apiUrl: options.apiUrl,
@@ -315,20 +104,16 @@ export function useStream(options) {
         options.callerOptions,
         options.defaultHeaders,
     ]);
+    const [messageManager] = useState(() => new MessageTupleManager());
+    const [stream] = useState(() => new StreamManager(messageManager));
+    useSyncExternalStore(stream.subscribe, stream.getSnapshot, stream.getSnapshot);
     const [threadId, onThreadId] = useControllableThreadId(options);
-    const [branch, setBranch] = useState("");
-    const [isLoading, setIsLoading] = useState(false);
-    const [streamError, setStreamError] = useState(undefined);
-    const [streamValues, setStreamValues, getMutateFn] = useStreamValuesState();
-    const messageManagerRef = useRef(new MessageTupleManager());
-    const submittingRef = useRef(false);
-    const abortRef = useRef(null);
     const trackStreamModeRef = useRef([]);
     const trackStreamMode = useCallback((...mode) => {
+        const ref = trackStreamModeRef.current;
         for (const m of mode) {
-            if (!trackStreamModeRef.current.includes(m)) {
-                trackStreamModeRef.current.push(m);
-            }
+            if (!ref.includes(m))
+                ref.push(m);
         }
     }, []);
     const hasUpdateListener = options.onUpdateEvent != null;
@@ -361,35 +146,37 @@ export function useStream(options) {
         hasTaskListener,
     ]);
     const clearCallbackRef = useRef(null);
-    clearCallbackRef.current = () => {
-        setStreamError(undefined);
-        setStreamValues(null);
-        messageManagerRef.current.clear();
-    };
+    clearCallbackRef.current = stream.clear;
+    const submittingRef = useRef(false);
+    submittingRef.current = stream.isLoading;
     const onErrorRef = useRef(undefined);
     onErrorRef.current = options.onError;
-    const historyLimit = typeof fetchStateHistory === "object" && fetchStateHistory != null
-        ? fetchStateHistory.limit ?? true
-        : fetchStateHistory ?? true;
+    const historyLimit = typeof options.fetchStateHistory === "object" &&
+        options.fetchStateHistory != null
+        ? options.fetchStateHistory.limit ?? false
+        : options.fetchStateHistory ?? false;
     const history = useThreadHistory(threadId, client, historyLimit, clearCallbackRef, submittingRef, onErrorRef);
-    const getMessages = useMemo(() => {
-        return (value) => Array.isArray(value[messagesKey])
-            ? value[messagesKey]
-            : [];
-    }, [messagesKey]);
-    const { rootSequence, paths } = getBranchSequence(history.data ?? []);
-    const { history: flatHistory, branchByCheckpoint } = getBranchView(rootSequence, paths, branch);
-    const threadHead = flatHistory.at(-1);
-    const historyValues = threadHead?.values ?? options.initialValues ?? {};
-    const historyValueError = (() => {
-        const error = threadHead?.tasks?.at(-1)?.error;
+    const getMessages = (value) => {
+        const messagesKey = options.messagesKey ?? "messages";
+        return Array.isArray(value[messagesKey]) ? value[messagesKey] : [];
+    };
+    const setMessages = (current, messages) => {
+        const messagesKey = options.messagesKey ?? "messages";
+        return { ...current, [messagesKey]: messages };
+    };
+    const [branch, setBranch] = useState("");
+    const branchContext = getBranchContext(branch, history.data);
+    const historyValues = branchContext.threadHead?.values ??
+        options.initialValues ??
+        {};
+    const historyError = (() => {
+        const error = branchContext.threadHead?.tasks?.at(-1)?.error;
         if (error == null)
             return undefined;
         try {
             const parsed = JSON.parse(error);
-            if (StreamError.isStructuredError(parsed)) {
+            if (StreamError.isStructuredError(parsed))
                 return new StreamError(parsed);
-            }
             return parsed;
         }
         catch {
@@ -401,16 +188,13 @@ export function useStream(options) {
         const alreadyShown = new Set();
         return getMessages(historyValues).map((message, idx) => {
             const messageId = message.id ?? idx;
-            const streamMetadata = message.id != null
-                ? messageManagerRef.current.get(message.id)?.metadata ?? undefined
-                : undefined;
-            const firstSeenIdx = findLastIndex(history.data ?? [], (state) => getMessages(state.values)
+            // Find the first checkpoint where the message was seen
+            const firstSeenState = findLast(history.data ?? [], (state) => getMessages(state.values)
                 .map((m, idx) => m.id ?? idx)
                 .includes(messageId));
-            const firstSeen = history.data?.[firstSeenIdx];
-            const checkpointId = firstSeen?.checkpoint?.checkpoint_id;
-            let branch = firstSeen && checkpointId != null
-                ? branchByCheckpoint[checkpointId]
+            const checkpointId = firstSeenState?.checkpoint?.checkpoint_id;
+            let branch = checkpointId != null
+                ? branchContext.branchByCheckpoint[checkpointId]
                 : undefined;
             if (!branch?.branch?.length)
                 branch = undefined;
@@ -423,158 +207,38 @@ export function useStream(options) {
             }
             return {
                 messageId: messageId.toString(),
-                firstSeenState: firstSeen,
+                firstSeenState,
                 branch: branch?.branch,
                 branchOptions: branch?.branchOptions,
-                streamMetadata,
             };
         });
     })();
-    const stop = () => {
-        if (abortRef.current != null)
-            abortRef.current.abort();
-        abortRef.current = null;
-        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?.({ mutate: getMutateFn("stop", historyValues) });
-    };
-    async function consumeStream(action) {
-        let getCallbackMeta;
-        try {
-            setIsLoading(true);
-            setStreamError(undefined);
-            submittingRef.current = true;
-            abortRef.current = new AbortController();
-            const run = await action(abortRef.current.signal);
-            getCallbackMeta = run.getCallbackMeta;
-            let streamError;
-            for await (const { event, data } of run.stream) {
-                if (event === "error") {
-                    streamError = new StreamError(data);
-                    break;
-                }
-                const namespace = event.includes("|")
-                    ? event.split("|").slice(1)
-                    : undefined;
-                const mutate = getMutateFn("stream", historyValues);
-                if (event === "metadata")
-                    options.onMetadataEvent?.(data);
-                if (event === "events")
-                    options.onLangChainEvent?.(data);
-                if (matchEventType("updates", event, data)) {
-                    options.onUpdateEvent?.(data, { namespace, mutate });
-                }
-                if (matchEventType("custom", event, data)) {
-                    options.onCustomEvent?.(data, { namespace, mutate });
-                }
-                if (matchEventType("checkpoints", event, data)) {
-                    options.onCheckpointEvent?.(data, { namespace });
-                }
-                if (matchEventType("tasks", event, data)) {
-                    options.onTaskEvent?.(data, { namespace });
-                }
-                if (matchEventType("debug", event, data)) {
-                    options.onDebugEvent?.(data, { namespace });
-                }
-                if (event === "values") {
-                    // don't update values on interrupt values event
-                    if ("__interrupt__" in data)
-                        continue;
-                    setStreamValues(data);
-                }
-                // Consume subgraph messages as well
-                if (matchEventType("messages", event, data)) {
-                    const [serialized, metadata] = data;
-                    const messageId = messageManagerRef.current.add(serialized, metadata);
-                    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 };
-                    });
-                }
-            }
-            // TODO: stream created checkpoints to avoid an unnecessary network request
-            const result = await run.onSuccess();
-            setStreamValues((values, kind) => {
-                // Do not clear out the user values set on `stop`.
-                if (kind === "stop")
-                    return values;
-                return null;
-            });
-            if (streamError != null)
-                throw streamError;
-            const lastHead = result.at(0);
-            if (lastHead)
-                onFinish?.(lastHead, getCallbackMeta?.());
-        }
-        catch (error) {
-            if (!(error instanceof Error && // eslint-disable-line no-instanceof/no-instanceof
-                (error.name === "AbortError" || error.name === "TimeoutError"))) {
-                console.error(error);
-                setStreamError(error);
-                onError?.(error, getCallbackMeta?.());
+    const stop = () => stream.stop(historyValues, { onStop: options.onStop });
+    // --- TRANSPORT ---
+    const submit = async (values, submitOptions) => {
+        // Unbranch things
+        const checkpointId = submitOptions?.checkpoint?.checkpoint_id;
+        setBranch(checkpointId != null
+            ? branchContext.branchByCheckpoint[checkpointId]?.branch ?? ""
+            : "");
+        stream.setStreamValues(() => {
+            if (submitOptions?.optimisticValues != null) {
+                return {
+                    ...historyValues,
+                    ...(typeof submitOptions.optimisticValues === "function"
+                        ? submitOptions.optimisticValues(historyValues)
+                        : submitOptions.optimisticValues),
+                };
             }
-        }
-        finally {
-            setIsLoading(false);
-            submittingRef.current = false;
-            abortRef.current = null;
-        }
-    }
-    const joinStream = async (runId, lastEventId, options) => {
-        // eslint-disable-next-line no-param-reassign
-        lastEventId ??= "-1";
-        if (!threadId)
-            return;
-        await consumeStream(async (signal) => {
-            const stream = client.runs.joinStream(threadId, runId, {
-                signal,
-                lastEventId,
-                streamMode: options?.streamMode,
-            });
-            return {
-                onSuccess: () => {
-                    runMetadataStorage?.removeItem(`lg:stream:${threadId}`);
-                    return history.mutate(threadId);
-                },
-                stream,
-                getCallbackMeta: () => ({ thread_id: threadId, run_id: runId }),
-            };
+            return { ...historyValues };
         });
-    };
-    const submit = async (values, submitOptions) => {
-        await consumeStream(async (signal) => {
-            // Unbranch things
-            const newPath = submitOptions?.checkpoint?.checkpoint_id
-                ? branchByCheckpoint[submitOptions?.checkpoint?.checkpoint_id]?.branch
-                : undefined;
-            if (newPath != null)
-                setBranch(newPath ?? "");
-            setStreamValues(() => {
-                if (submitOptions?.optimisticValues != null) {
-                    return {
-                        ...historyValues,
-                        ...(typeof submitOptions.optimisticValues === "function"
-                            ? submitOptions.optimisticValues(historyValues)
-                            : submitOptions.optimisticValues),
-                    };
-                }
-                return { ...historyValues };
-            });
-            let usableThreadId = threadId;
+        // When `fetchStateHistory` is requested, thus we assume that branching
+        // is enabled. We then need to include the implicit branch.
+        const includeImplicitBranch = historyLimit === true || typeof historyLimit === "number";
+        let callbackMeta;
+        let rejoinKey;
+        let usableThreadId = threadId;
+        await stream.start(async (signal) => {
             if (!usableThreadId) {
                 const thread = await client.threads.create({
                     threadId: submitOptions?.threadId,
@@ -583,26 +247,28 @@ export function useStream(options) {
                 onThreadId(thread.thread_id);
                 usableThreadId = thread.thread_id;
             }
-            if (!usableThreadId)
+            if (!usableThreadId) {
                 throw new Error("Failed to obtain valid thread ID.");
+            }
             const streamMode = unique([
                 ...(submitOptions?.streamMode ?? []),
                 ...trackStreamModeRef.current,
                 ...callbackStreamMode,
             ]);
-            let checkpoint = submitOptions?.checkpoint ?? threadHead?.checkpoint ?? undefined;
+            let checkpoint = submitOptions?.checkpoint ??
+                (includeImplicitBranch
+                    ? branchContext.threadHead?.checkpoint
+                    : undefined) ??
+                undefined;
             // Avoid specifying a checkpoint if user explicitly set it to null
-            if (submitOptions?.checkpoint === null) {
+            if (submitOptions?.checkpoint === null)
                 checkpoint = undefined;
-            }
             // eslint-disable-next-line @typescript-eslint/ban-ts-comment
             // @ts-expect-error
             if (checkpoint != null)
                 delete checkpoint.thread_id;
-            let rejoinKey;
-            let callbackMeta;
             const streamResumable = submitOptions?.streamResumable ?? !!runMetadataStorage;
-            const stream = client.runs.stream(usableThreadId, assistantId, {
+            return client.runs.stream(usableThreadId, options.assistantId, {
                 input: values,
                 config: submitOptions?.config,
                 context: submitOptions?.context,
@@ -626,25 +292,77 @@ export function useStream(options) {
                         thread_id: params.thread_id ?? usableThreadId,
                     };
                     if (runMetadataStorage) {
-                        rejoinKey = `lg:stream:${callbackMeta.thread_id}`;
+                        rejoinKey = `lg:stream:${usableThreadId}`;
                         runMetadataStorage.setItem(rejoinKey, callbackMeta.run_id);
                     }
-                    onCreated?.(callbackMeta);
+                    options.onCreated?.(callbackMeta);
                 },
             });
-            return {
-                stream,
-                getCallbackMeta: () => callbackMeta,
-                onSuccess: () => {
-                    if (rejoinKey)
-                        runMetadataStorage?.removeItem(rejoinKey);
-                    return history.mutate(usableThreadId);
-                },
-            };
+        }, {
+            getMessages,
+            setMessages,
+            initialValues: historyValues,
+            callbacks: options,
+            async onSuccess() {
+                if (rejoinKey)
+                    runMetadataStorage?.removeItem(rejoinKey);
+                const shouldRefetch = 
+                // We're expecting the whole thread state in onFinish
+                options.onFinish != null ||
+                    // We're fetching history, thus we need the latest checkpoint
+                    // to ensure we're not accidentally submitting to a wrong branch
+                    includeImplicitBranch;
+                if (shouldRefetch) {
+                    const newHistory = await history.mutate(usableThreadId);
+                    const lastHead = newHistory.at(0);
+                    if (lastHead) {
+                        // We now have the latest update from /history
+                        // Thus we can clear the local stream state
+                        options.onFinish?.(lastHead, callbackMeta);
+                        return null;
+                    }
+                }
+                return undefined;
+            },
+            onError(error) {
+                options.onError?.(error, callbackMeta);
+            },
+        });
+    };
+    const joinStream = async (runId, lastEventId, joinOptions) => {
+        // eslint-disable-next-line no-param-reassign
+        lastEventId ??= "-1";
+        if (!threadId)
+            return;
+        const callbackMeta = {
+            thread_id: threadId,
+            run_id: runId,
+        };
+        await stream.start(async (signal) => {
+            return client.runs.joinStream(threadId, runId, {
+                signal,
+                lastEventId,
+                streamMode: joinOptions?.streamMode,
+            });
+        }, {
+            getMessages,
+            setMessages,
+            initialValues: historyValues,
+            callbacks: options,
+            async onSuccess() {
+                runMetadataStorage?.removeItem(`lg:stream:${threadId}`);
+                const newHistory = await history.mutate(threadId);
+                const lastHead = newHistory.at(0);
+                if (lastHead)
+                    options.onFinish?.(lastHead, callbackMeta);
+            },
+            onError(error) {
+                options.onError?.(error, callbackMeta);
+            },
         });
     };
     const reconnectKey = useMemo(() => {
-        if (!runMetadataStorage || isLoading)
+        if (!runMetadataStorage || stream.isLoading)
             return undefined;
         if (typeof window === "undefined")
             return undefined;
@@ -652,7 +370,7 @@ export function useStream(options) {
         if (!runId)
             return undefined;
         return { runId, threadId };
-    }, [runMetadataStorage, isLoading, threadId]);
+    }, [runMetadataStorage, stream.isLoading, threadId]);
     const shouldReconnect = !!runMetadataStorage;
     const reconnectRef = useRef({ threadId, shouldReconnect });
     const joinStreamRef = useRef(joinStream);
@@ -669,38 +387,44 @@ export function useStream(options) {
             void joinStreamRef.current?.(reconnectKey.runId);
         }
     }, [reconnectKey]);
-    const error = streamError ?? historyValueError ?? history.error;
-    const values = streamValues ?? historyValues;
+    // --- END TRANSPORT ---
+    const error = stream.error ?? historyError ?? history.error;
+    const values = stream.values ?? historyValues;
     return {
         get values() {
             trackStreamMode("values");
             return values;
         },
         client,
-        assistantId,
+        assistantId: options.assistantId,
         error,
-        isLoading,
+        isLoading: stream.isLoading,
         stop,
-        submit, // eslint-disable-line @typescript-eslint/no-misused-promises
+        submit,
         joinStream,
         branch,
         setBranch,
-        history: flatHistory,
+        get history() {
+            if (historyLimit === false) {
+                throw new Error("`fetchStateHistory` must be set to `true` to use `history`");
+            }
+            return branchContext.flatHistory;
+        },
         isThreadLoading: history.isLoading && history.data == null,
         get experimental_branchTree() {
             if (historyLimit === false) {
-                throw new Error("`experimental_branchTree` is not available when `fetchStateHistory` is set to `false`");
+                throw new Error("`fetchStateHistory` must be set to `true` to use `experimental_branchTree`");
             }
-            return rootSequence;
+            return branchContext.branchTree;
         },
         get interrupt() {
             // Don't show the interrupt if the stream is loading
-            if (isLoading)
+            if (stream.isLoading)
                 return undefined;
-            const interrupts = threadHead?.tasks?.at(-1)?.interrupts;
+            const interrupts = branchContext.threadHead?.tasks?.at(-1)?.interrupts;
             if (interrupts == null || interrupts.length === 0) {
                 // check if there's a next task present
-                const next = threadHead?.next ?? [];
+                const next = branchContext.threadHead?.next ?? [];
                 if (!next.length || error != null)
                     return undefined;
                 return { when: "breakpoint" };
@@ -714,7 +438,15 @@ export function useStream(options) {
         },
         getMessagesMetadata(message, index) {
             trackStreamMode("values");
-            return messageMetadata?.find((m) => m.messageId === (message.id ?? index));
+            const streamMetadata = messageManager.get(message.id)?.metadata;
+            const historyMetadata = messageMetadata?.find((m) => m.messageId === (message.id ?? index));
+            if (streamMetadata != null || historyMetadata != null) {
+                return {
+                    ...historyMetadata,
+                    streamMetadata,
+                };
+            }
+            return undefined;
         },
     };
 }