feat(sdk-js): add docs for new useStream hook (#3420)

- **Add basic docs**
- **Add docs**
- **feat(sdk-js): add docs, how-to guide**

docs/docs/cloud/how-tos/use_stream_react.md

@@ -0,0 +1,410 @@
+# How to integrate LangGraph into your React application
+
+!!! info "Prerequisites"
+    - [LangGraph Platform](../../concepts/langgraph_platform.md)
+    - [LangGraph Server](../../concepts/langgraph_server.md)
+
+The `useStream()` React hook provides a seamless way to integrate LangGraph into your React applications. It handles all the complexities of streaming, state management, and branching logic, letting you focus on building great chat experiences.
+
+Key features:
+
+- Messages streaming: Handle a stream of message chunks to form a complete message
+- Automatic state management for messages, loading states, and errors
+- Conversation branching: Create alternate conversation paths from any point in the chat history
+- UI-agnostic design - bring your own components and styling
+
+Let's explore how to use `useStream()` in your React application.
+
+The `useStream()` provides a solid foundation for creating bespoke chat experiences. For pre-built chat components and interfaces, we recommend checking out [CopilotKit](https://docs.copilotkit.ai/coagents/quickstart/langgraph) and [assistant-ui](https://github.com/langchain-ai/assistant-ui).
+
+## Example
+
+```tsx
+"use client";
+
+import { useStream } from "@langchain/langgraph-sdk/react";
+import type { Message } from "@langchain/langgraph-sdk";
+
+export default function App() {
+  const thread = useStream<{ messages: Message[] }>({
+    apiUrl: "http://localhost:2024",
+    assistantId: "agent",
+    messagesKey: "messages",
+  });
+
+  return (
+    <div>
+      <div>
+        {thread.messages.map((message) => (
+          <div key={message.id}>{message.content as string}</div>
+        ))}
+      </div>
+
+      <form
+        onSubmit={(e) => {
+          e.preventDefault();
+
+          const form = e.target as HTMLFormElement;
+          const message = new FormData(form).get("message") as string;
+
+          form.reset();
+          thread.submit({ messages: [{ type: "human", content: message }] });
+        }}
+      >
+        <input type="text" name="message" />
+
+        {thread.isLoading ? (
+          <button key="stop" type="button" onClick={() => thread.stop()}>
+            Stop
+          </button>
+        ) : (
+          <button key="submit" type="submit">
+            Send
+          </button>
+        )}
+      </form>
+    </div>
+  );
+}
+```
+
+## Customizing Your UI
+
+The `useStream()` hook takes care of all the complex state management behind the scenes, providing you with simple interfaces to build your UI. Here's what you get out of the box:
+
+- Thread state management
+- Loading and error states
+- Message handling and updates
+- Branching support
+
+Here are some examples on how to use these features effectively:
+
+### Loading States
+
+The `isLoading` property tells you when a stream is active, enabling you to:
+
+- Show a loading indicator
+- Disable input fields during processing
+- Display a cancel button
+
+```tsx
+export default function App() {
+  const { isLoading, stop } = useStream<{ messages: Message[] }>({
+    apiUrl: "http://localhost:2024",
+    assistantId: "agent",
+    messagesKey: "messages",
+  });
+
+  return (
+    <form>
+      {isLoading && (
+        <button key="stop" type="button" onClick={() => stop()}>
+          Stop
+        </button>
+      )}
+    </form>
+  );
+}
+```
+
+### Thread Management
+
+Keep track of conversations with built-in thread management. You can access the current thread ID and get notified when new threads are created:
+
+```tsx
+const [threadId, setThreadId] = useState<string | null>(null);
+
+const thread = useStream<{ messages: Message[] }>({
+  apiUrl: "http://localhost:2024",
+  assistantId: "agent",
+
+  threadId: threadId,
+  onThreadId: setThreadId,
+});
+```
+
+We recommend storing the `threadId` in your URL's query parameters to let users resume conversations after page refreshes.
+
+### Messages Handling
+
+To enable messages handling, you need to pass the `messagesKey` option to the `useStream()` hook. 
+
+When enabled, the `useStream()` hook will keep track of the message chunks received from the server and concatenate them together to form a complete message. The completed message chunks can be retrieved via the `messages` property.
+
+```tsx
+import type { Message } from "@langchain/langgraph-sdk";
+import { useStream } from "@langchain/langgraph-sdk/react";
+
+export default function HomePage() {
+  const thread = useStream<{ messages: Message[] }>({
+    apiUrl: "http://localhost:2024",
+    assistantId: "agent",
+    messagesKey: "messages",
+  });
+
+  return (
+    <div>
+      {thread.messages.map((message) => (
+        <div key={message.id}>{message.content as string}</div>
+      ))}
+    </div>
+  );
+}
+```
+
+### Branching Support
+
+To enable branching, you need to enable messages handling. Pass the `messagesKey` option to the `useStream()` hook. For each message, you can use `getMessagesMetadata()` to get the first checkpoint from which the message has been first seen. You can then create a new run from the checkpoint preceding the first seen checkpoint to create a new branch in a thread.
+
+A branch can be created in following ways:
+
+1. Edit a previous user message.
+2. Request a regeneration of a previous assistant message.
+
+```tsx
+/* eslint-disable @typescript-eslint/no-floating-promises */
+"use client";
+
+import type { Message } from "@langchain/langgraph-sdk";
+import { useStream } from "@langchain/langgraph-sdk/react";
+import {
+  Annotation,
+  MessagesAnnotation,
+  type StateType,
+  type UpdateType,
+} from "@langchain/langgraph/web";
+import { useState } from "react";
+
+const AgentState = Annotation.Root({
+  ...MessagesAnnotation.spec,
+});
+
+function BranchSwitcher({
+  branch,
+  branchOptions,
+  onSelect,
+}: {
+  branch: string | undefined;
+  branchOptions: string[] | undefined;
+  onSelect: (branch: string) => void;
+}) {
+  if (!branchOptions || !branch) return null;
+  const index = branchOptions.indexOf(branch);
+
+  return (
+    <div className="flex items-center gap-2">
+      <button
+        type="button"
+        onClick={() => {
+          const prevBranch = branchOptions[index - 1];
+          if (!prevBranch) return;
+          onSelect(prevBranch);
+        }}
+      >
+        Prev
+      </button>
+      <span>
+        {index + 1} / {branchOptions.length}
+      </span>
+      <button
+        type="button"
+        onClick={() => {
+          const nextBranch = branchOptions[index + 1];
+          if (!nextBranch) return;
+          onSelect(nextBranch);
+        }}
+      >
+        Next
+      </button>
+    </div>
+  );
+}
+
+function EditMessage({
+  message,
+  onEdit,
+}: {
+  message: Message;
+  onEdit: (message: Message) => void;
+}) {
+  const [editing, setEditing] = useState(false);
+
+  if (!editing) {
+    return (
+      <button type="button" onClick={() => setEditing(true)}>
+        Edit
+      </button>
+    );
+  }
+
+  return (
+    <form
+      onSubmit={(e) => {
+        e.preventDefault();
+        const form = e.target as HTMLFormElement;
+        const content = new FormData(form).get("content") as string;
+
+        form.reset();
+        onEdit({ type: "human", content });
+        setEditing(false);
+      }}
+    >
+      <input name="content" defaultValue={message.content as string} />
+      <button type="submit">Save</button>
+    </form>
+  );
+}
+
+export default function App() {
+  const thread = useStream<
+    StateType<typeof AgentState.spec>,
+    UpdateType<typeof AgentState.spec>
+  >({
+    apiUrl: "http://localhost:2024",
+    assistantId: "agent",
+    messagesKey: "messages",
+  });
+
+  return (
+    <div>
+      <div>
+        {thread.messages.map((message) => {
+          const meta = thread.getMessagesMetadata(message);
+          const parentCheckpoint = meta?.firstSeenState?.parent_checkpoint;
+
+          return (
+            <div key={message.id}>
+              <div>{message.content as string}</div>
+
+              {message.type === "human" && (
+                <EditMessage
+                  message={message}
+                  onEdit={(message) =>
+                    thread.submit(
+                      { messages: [message] },
+                      { checkpoint: parentCheckpoint }
+                    )
+                  }
+                />
+              )}
+
+              {message.type === "ai" && (
+                <button
+                  type="button"
+                  onClick={() =>
+                    thread.submit(undefined, { checkpoint: parentCheckpoint })
+                  }
+                >
+                  <span>Regenerate</span>
+                </button>
+              )}
+
+              <BranchSwitcher
+                branch={meta?.branch}
+                branchOptions={meta?.branchOptions}
+                onSelect={(branch) => thread.setBranch(branch)}
+              />
+            </div>
+          );
+        })}
+      </div>
+
+      <form
+        onSubmit={(e) => {
+          e.preventDefault();
+
+          const form = e.target as HTMLFormElement;
+          const message = new FormData(form).get("message") as string;
+
+          form.reset();
+          thread.submit({ messages: [message] });
+        }}
+      >
+        <input type="text" name="message" />
+
+        {thread.isLoading ? (
+          <button key="stop" type="button" onClick={() => thread.stop()}>
+            Stop
+          </button>
+        ) : (
+          <button key="submit" type="submit">
+            Send
+          </button>
+        )}
+      </form>
+    </div>
+  );
+}
+```
+
+### TypeScript
+
+The `useStream()` hook is fully typed to help catch errors early and provide better IDE support. You can specify types for:
+
+- State shape
+- Update format
+- Custom events
+
+```tsx
+// Define your types
+type State = {
+  messages: Message[];
+  context?: Record<string, unknown>;
+};
+
+type Update = {
+  messages: Message[] | Message;
+  context?: Record<string, unknown>;
+};
+
+type CustomEvent = {
+  type: "progress" | "debug";
+  payload: unknown;
+};
+
+// Use them with the hook
+const thread = useStream<State, Update, CustomEvent>({
+  apiUrl: "http://localhost:2024",
+  assistantId: "agent",
+  messagesKey: "messages",
+});
+```
+
+If you're using LangGraph.js, you can reuse your graph's annotation types:
+
+```tsx
+import {
+  Annotation,
+  MessagesAnnotation,
+  type StateType,
+  type UpdateType,
+} from "@langchain/langgraph/web";
+
+const AgentState = Annotation.Root({
+  ...MessagesAnnotation.spec,
+  context: Annotation.Optional(Annotation.Any()),
+});
+
+const thread = useStream<
+  StateType<typeof AgentState.spec>,
+  UpdateType<typeof AgentState.spec>
+>({
+  apiUrl: "http://localhost:2024",
+  assistantId: "agent",
+  messagesKey: "messages",
+});
+```
+
+## Event Handling
+
+The `useStream()` hook provides several callback options to help you respond to different events:
+
+- `onError`: Called when an error occurs.
+- `onFinish`: Called when the stream is finished.
+- `onUpdateEvent`: Called when an update event is received.
+- `onCustomEvent`: Called when a custom event is received. See [Custom events](../../concepts/streaming.md#custom) to learn how to stream custom events.
+- `onMetadataEvent`: Called when a metadata event is received.
+
+## Learn More
+
+- [JS/TS SDK Reference](../reference/sdk/js_ts_sdk_ref.md)

docs/docs/how-tos/index.md

@@ -204,6 +204,7 @@ Learn how to set up your app for deployment to LangGraph Platform:
 - [How to test locally](../cloud/deployment/test_locally.md)
 - [How to rebuild graph at runtime](../cloud/deployment/graph_rebuild.md)
 - [How to use LangGraph Platform to deploy CrewAI, AutoGen, and other frameworks](autogen-langgraph-platform.ipynb)
+- [How to integrate LangGraph into your React application](../cloud/how-tos/use_stream_react.md)
 
 ### Deployment
 

docs/mkdocs.yml

@@ -253,6 +253,7 @@ nav:
                 - cloud/how-tos/stream_events.md
                 - cloud/how-tos/stream_debug.md
                 - cloud/how-tos/stream_multiple.md
+                - cloud/how-tos/use_stream_react.md
             - Human-in-the-loop:
                 - Human-in-the-loop: how-tos#human-in-the-loop_1
                 - cloud/how-tos/human_in_the_loop_breakpoint.md

libs/sdk-js/.gitignore

@@ -13,3 +13,4 @@ react.d.cts
 node_modules
 dist
 .yarn
+docs

libs/sdk-js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.40",
+  "version": "0.0.41",
   "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"
   },

libs/sdk-js/src/react/stream.tsx

@@ -128,13 +128,145 @@ interface ValidSequence<StateType = any> {
 }
 
 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;
 };
 
+function getBranchSequence<StateType extends Record<string, unknown>>(
+  history: ThreadState<StateType>[],
+) {
+  const childrenMap: Record<string, ThreadState<StateType>[]> = {};
+
+  // First pass - collect nodes for each checkpoint
+  history.forEach((state) => {
+    const checkpointId = state.parent_checkpoint?.checkpoint_id ?? "$";
+    childrenMap[checkpointId] ??= [];
+    childrenMap[checkpointId].push(state);
+  });
+
+  // Second pass - create a tree of sequences
+  type Task = { id: string; sequence: Sequence; path: string[] };
+  const rootSequence: Sequence = { type: "sequence", items: [] };
+  const queue: Task[] = [{ id: "$", sequence: rootSequence, path: [] }];
+
+  const paths: string[][] = [];
+
+  const visited = new Set<string>();
+  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: Fork | undefined;
+    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<StateType extends Record<string, unknown>>(
+  sequence: Sequence<StateType>,
+  paths: string[][],
+  branch: string,
+) {
+  const path = branch.split(PATH_SEP);
+  const pathMap: Record<string, string[][]> = {};
+
+  for (const path of paths) {
+    const parent = path.at(-2) ?? ROOT_ID;
+    pathMap[parent] ??= [];
+    pathMap[parent].unshift(path);
+  }
+
+  const history: ThreadState<StateType>[] = [];
+  const branchByCheckpoint: Record<
+    string,
+    { branch: string | undefined; branchOptions: string[] | undefined }
+  > = {};
+
+  const forkStack = path.slice();
+  const queue: (Node<StateType> | Fork<StateType>)[] = [...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<StateType extends Record<string, unknown>>(
   client: Client,
   threadId: string,
@@ -179,29 +311,189 @@ function useThreadHistory<StateType extends Record<string, unknown>>(
   };
 }
 
+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 (typeof options?.threadId === "undefined") {
+    return [localThreadId, onThreadId];
+  }
+
+  return [options.threadId, onThreadId];
+};
+
+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"];
+
+  /**
+   * 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;
+}
+
+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;
+
+  /**
+   * 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 function useStream<
   StateType extends Record<string, unknown> = Record<string, unknown>,
   UpdateType extends Record<string, unknown> = Partial<StateType>,
   CustomType = unknown,
->(options: {
-  assistantId: string;
-
-  apiUrl: ClientConfig["apiUrl"];
-  apiKey?: ClientConfig["apiKey"];
-
-  withMessages?: string;
-
-  onError?: (error: unknown) => void;
-  onFinish?: (state: ThreadState<StateType>) => void;
-
-  onUpdateEvent?: (data: UpdatesStreamEvent<UpdateType>["data"]) => void;
-  onCustomEvent?: (data: CustomStreamEvent<CustomType>["data"]) => void;
-  onMetadataEvent?: (data: MetadataStreamEvent["data"]) => void;
-
-  // TODO: can we make threadId uncontrollable / controllable?
-  threadId?: string | null;
-  onThreadId?: (threadId: string) => void;
-}) {
+>(
+  options: UseStreamOptions<StateType, UpdateType, CustomType>,
+): UseStream<StateType, UpdateType> {
   type EventStreamEvent =
     | ValuesStreamEvent<StateType>
     | UpdatesStreamEvent<UpdateType>
@@ -214,15 +506,17 @@ export function useStream<
     | ErrorStreamEvent
     | FeedbackStreamEvent;
 
-  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 [threadId, onThreadId] = useControllableThreadId(options);
 
-  const [branchPath, setBranchPath] = useState<string[]>([]);
+  const [branch, setBranch] = useState<string>("");
   const [isLoading, setIsLoading] = useState(false);
-  const [_, setEvents] = useState<EventStreamEvent[]>([]);
 
   const [streamError, setStreamError] = useState<unknown>(undefined);
   const [streamValues, setStreamValues] = useState<StateType | null>(null);
@@ -269,118 +563,20 @@ export function useStream<
   );
 
   const getMessages = useMemo(() => {
-    if (withMessages == null) return undefined;
     return (value: StateType) =>
-      Array.isArray(value[withMessages])
-        ? (value[withMessages] as Message[])
+      Array.isArray(value[messagesKey])
+        ? (value[messagesKey] as Message[])
         : [];
-  }, [withMessages]);
+  }, [messagesKey]);
 
-  const [sequence, pathMap] = (() => {
-    const childrenMap: Record<string, ThreadState<StateType>[]> = {};
+  const { rootSequence, paths } = getBranchSequence(history.data);
+  const { history: flatHistory, branchByCheckpoint } = getBranchView(
+    rootSequence,
+    paths,
+    branch,
+  );
 
-    // First pass - collect nodes for each checkpoint
-    history.data.forEach((state) => {
-      const checkpointId = state.parent_checkpoint?.checkpoint_id ?? "$";
-      childrenMap[checkpointId] ??= [];
-      childrenMap[checkpointId].push(state);
-    });
-
-    // Second pass - create a tree of sequences
-    type Task = { id: string; sequence: Sequence; path: string[] };
-    const rootSequence: Sequence = { type: "sequence", items: [] };
-    const queue: Task[] = [{ id: "$", sequence: rootSequence, path: [] }];
-
-    const paths: string[][] = [];
-
-    const visited = new Set<string>();
-    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: Fork | undefined;
-      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: Record<string, string[][]> = {};
-    for (const path of paths) {
-      const parent = path.at(-2) ?? "$";
-      pathMap[parent] ??= [];
-      pathMap[parent].unshift(path);
-    }
-
-    return [rootSequence as ValidSequence, pathMap];
-  })();
-
-  const [flatValues, flatPaths] = (() => {
-    const result: ThreadState<StateType>[] = [];
-    const flatPaths: Record<
-      string,
-      { current: string[] | undefined; branches: string[][] | undefined }
-    > = {};
-
-    const forkStack = branchPath.slice();
-    const queue: (Node<StateType> | Fork<StateType>)[] = [...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: ThreadState<StateType> | undefined = flatValues.at(-1);
+  const threadHead: ThreadState<StateType> | undefined = flatHistory.at(-1);
   const historyValues = threadHead?.values ?? ({} as StateType);
   const historyError = (() => {
     const error = threadHead?.tasks?.at(-1)?.error;
@@ -399,8 +595,6 @@ export function useStream<
   })();
 
   const messageMetadata = (() => {
-    if (getMessages == null) return undefined;
-
     const alreadyShown = new Set<string>();
     return getMessages(historyValues).map(
       (message, idx): MessageMetadata<StateType> => {
@@ -416,13 +610,13 @@ export function useStream<
           | undefined;
 
         let branch = firstSeen
-          ? flatPaths[firstSeen.checkpoint.checkpoint_id!]
+          ? branchByCheckpoint[firstSeen.checkpoint.checkpoint_id!]
           : undefined;
 
-        if (!branch?.current?.length) branch = undefined;
+        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;
           alreadyShown.add(optionsShown);
@@ -431,8 +625,9 @@ export function useStream<
         return {
           messageId: messageId.toString(),
           firstSeenState: firstSeen,
-          branch: branch?.current?.join(">"),
-          branchOptions: branch?.branches?.map((b) => b.join(">")),
+
+          branch: branch?.branch,
+          branchOptions: branch?.branchOptions,
         };
       },
     );
@@ -445,22 +640,7 @@ export function useStream<
 
   const submit = async (
     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>);
-    },
+    submitOptions?: SubmitOptions<StateType>,
   ) => {
     try {
       setIsLoading(true);
@@ -472,7 +652,7 @@ export function useStream<
       let usableThreadId = threadId;
       if (!usableThreadId) {
         const thread = await client.threads.create();
-        options?.onThreadId?.(thread.thread_id);
+        onThreadId(thread.thread_id);
         usableThreadId = thread.thread_id;
       }
 
@@ -507,9 +687,10 @@ export function useStream<
 
       // 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 ?? []);
+
+      if (newPath != null) setBranch(newPath ?? "");
 
       // Assumption: we're setting the initial value
       // Used for instant feedback
@@ -530,32 +711,17 @@ export function useStream<
 
       let streamError: StreamError | undefined;
       for await (const { event, data } of run) {
-        setEvents((events) => [...events, { event, data } as EventStreamEvent]);
-
         if (event === "error") {
           streamError = new StreamError(data);
           break;
         }
 
-        if (event === "updates") {
-          options.onUpdateEvent?.(data);
-        }
-
-        if (event === "custom") {
-          options.onCustomEvent?.(data);
-        }
-
-        if (event === "metadata") {
-          options.onMetadataEvent?.(data);
-        }
-
-        if (event === "values") {
-          setStreamValues(data);
-        }
+        if (event === "updates") options.onUpdateEvent?.(data);
+        if (event === "custom") options.onCustomEvent?.(data);
+        if (event === "metadata") options.onMetadataEvent?.(data);
 
+        if (event === "values") setStreamValues(data);
         if (event === "messages") {
-          if (!getMessages) continue;
-
           const [serialized] = data;
 
           const messageId = messageManagerRef.current.add(serialized);
@@ -577,15 +743,13 @@ export function useStream<
             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;
@@ -615,11 +779,6 @@ export function useStream<
   const error = isLoading ? streamError : historyError;
   const values = streamValues ?? historyValues;
 
-  const setBranch = useCallback(
-    (path: string) => setBranchPath(path.split(">")),
-    [setBranchPath],
-  );
-
   return {
     get values() {
       trackStreamMode("values");
@@ -631,17 +790,15 @@ export function useStream<
 
     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);
     },
 
@@ -650,13 +807,6 @@ export function useStream<
       index?: number,
     ): MessageMetadata<StateType> | undefined {
       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),
       );

libs/sdk-js/typedoc.react.json

@@ -0,0 +1,5 @@
+{
+  "pageTitleTemplates": {
+    "index": "{projectName}/react"
+  }
+}

libs/sdk-js/yarn.lock

@@ -301,6 +301,15 @@
   resolved "https://registry.yarnpkg.com/@cfworker/json-schema/-/json-schema-4.1.1.tgz#4a2a3947ee9fa7b7c24be981422831b8674c3be6"
   integrity sha512-gAmrUZSGtKc3AiBL71iNWxDsyUC5uMaKKGdvzYsBoTW/xi42JQHl7eKV2OYzCUqvc+D2RCcf7EXY2iCyFIk6og==
 
+"@gerrit0/mini-shiki@^1.24.0":
+  version "1.27.2"
+  resolved "https://registry.yarnpkg.com/@gerrit0/mini-shiki/-/mini-shiki-1.27.2.tgz#cf2a9fcb08a6581c78fc94821f0c854ec4b9f899"
+  integrity sha512-GeWyHz8ao2gBiUW4OJnQDxXQnFgZQwwQk05t/CVVgNBN7/rK8XZ7xY6YhLVv9tH3VppWWmr9DCl3MwemB/i+Og==
+  dependencies:
+    "@shikijs/engine-oniguruma" "^1.27.2"
+    "@shikijs/types" "^1.27.2"
+    "@shikijs/vscode-textmate" "^10.0.1"
+
 "@isaacs/cliui@^8.0.2":
   version "8.0.2"
   resolved "https://registry.yarnpkg.com/@isaacs/cliui/-/cliui-8.0.2.tgz#b37667b7bc181c168782259bab42474fbf52b550"
@@ -806,10 +815,26 @@
   optionalDependencies:
     fsevents "~2.3.2"
 
-"@shikijs/core@1.9.0":
-  version "1.9.0"
-  resolved "https://registry.yarnpkg.com/@shikijs/core/-/core-1.9.0.tgz#ff717fef5e0e9882f0848272699fd8f04d6f9a07"
-  integrity sha512-cbSoY8P/jgGByG8UOl3jnP/CWg/Qk+1q+eAKWtcrU3pNoILF8wTsLB0jT44qUBV8Ce1SvA9uqcM9Xf+u3fJFBw==
+"@shikijs/engine-oniguruma@^1.27.2":
+  version "1.29.2"
+  resolved "https://registry.yarnpkg.com/@shikijs/engine-oniguruma/-/engine-oniguruma-1.29.2.tgz#d879717ced61d44e78feab16f701f6edd75434f1"
+  integrity sha512-7iiOx3SG8+g1MnlzZVDYiaeHe7Ez2Kf2HrJzdmGwkRisT7r4rak0e655AcM/tF9JG/kg5fMNYlLLKglbN7gBqA==
+  dependencies:
+    "@shikijs/types" "1.29.2"
+    "@shikijs/vscode-textmate" "^10.0.1"
+
+"@shikijs/types@1.29.2", "@shikijs/types@^1.27.2":
+  version "1.29.2"
+  resolved "https://registry.yarnpkg.com/@shikijs/types/-/types-1.29.2.tgz#a93fdb410d1af8360c67bf5fc1d1a68d58e21c4f"
+  integrity sha512-VJjK0eIijTZf0QSTODEXCqinjBn0joAHQ+aPSBzrv4O2d/QSbsMw+ZeSRx03kV34Hy7NzUvV/7NqfYGRLrASmw==
+  dependencies:
+    "@shikijs/vscode-textmate" "^10.0.1"
+    "@types/hast" "^3.0.4"
+
+"@shikijs/vscode-textmate@^10.0.1":
+  version "10.0.1"
+  resolved "https://registry.yarnpkg.com/@shikijs/vscode-textmate/-/vscode-textmate-10.0.1.tgz#d06d45b67ac5e9b0088e3f67ebd3f25c6c3d711a"
+  integrity sha512-fTIQwLF+Qhuws31iw7Ncl1R3HUDtGwIipiJ9iU+UsDUwMhegFcQKQHd51nZjb7CArq0MvON8rbgCGQYWHUKAdg==
 
 "@sinclair/typebox@^0.27.8":
   version "0.27.8"
@@ -910,6 +935,13 @@
   dependencies:
     "@types/node" "*"
 
+"@types/hast@^3.0.4":
+  version "3.0.4"
+  resolved "https://registry.yarnpkg.com/@types/hast/-/hast-3.0.4.tgz#1d6b39993b82cea6ad783945b0508c25903e15aa"
+  integrity sha512-WPs+bbQw5aCj+x6laNGWLH3wviHtoCv/P3+otBhbOhJgG8qtpdAMlTCxLtsTWA7LH1Oh/bFCHsBn0TPS5m30EQ==
+  dependencies:
+    "@types/unist" "*"
+
 "@types/istanbul-lib-coverage@*", "@types/istanbul-lib-coverage@^2.0.0", "@types/istanbul-lib-coverage@^2.0.1":
   version "2.0.6"
   resolved "https://registry.yarnpkg.com/@types/istanbul-lib-coverage/-/istanbul-lib-coverage-2.0.6.tgz#7739c232a1fee9b4d3ce8985f314c0c6d33549d7"
@@ -996,6 +1028,11 @@
   resolved "https://registry.yarnpkg.com/@types/stack-utils/-/stack-utils-2.0.3.tgz#6209321eb2c1712a7e7466422b8cb1fc0d9dd5d8"
   integrity sha512-9aEbYZ3TbYMznPdcdr3SmIrLXwC/AKZXQeCf9Pgao5CKb8CyHuEX5jzWPTkvregvhRJHcpRO6BFoGW9ycaOkYw==
 
+"@types/unist@*":
+  version "3.0.3"
+  resolved "https://registry.yarnpkg.com/@types/unist/-/unist-3.0.3.tgz#acaab0f919ce69cce629c2d4ed2eb4adc1b6c20c"
+  integrity sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q==
+
 "@types/unist@^2", "@types/unist@^2.0.0", "@types/unist@^2.0.2":
   version "2.0.10"
   resolved "https://registry.yarnpkg.com/@types/unist/-/unist-2.0.10.tgz#04ffa7f406ab628f7f7e97ca23e290cd8ab15efc"
@@ -3271,7 +3308,7 @@ minimatch@^5.0.1:
   dependencies:
     brace-expansion "^2.0.1"
 
-minimatch@^9.0.3:
+minimatch@^9.0.3, minimatch@^9.0.5:
   version "9.0.5"
   resolved "https://registry.yarnpkg.com/minimatch/-/minimatch-9.0.5.tgz#d74f9dd6b57d83d8e98cfb82133b03978bc929e5"
   integrity sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==
@@ -3857,13 +3894,6 @@ shebang-regex@^3.0.0:
   resolved "https://registry.yarnpkg.com/shebang-regex/-/shebang-regex-3.0.0.tgz#ae16f1644d873ecad843b0307b143362d4c42172"
   integrity sha512-7++dFhtcx3353uBaq8DDR4NuxBetBzC7ZQOhmTQInHEd6bSrXdiEyzCvG07Z44UYdLShWUyXt5M/yhz8ekcb1A==
 
-shiki@^1.9.0:
-  version "1.9.0"
-  resolved "https://registry.yarnpkg.com/shiki/-/shiki-1.9.0.tgz#e4d3a044d9c746aefbea47615e83323fdc3dc361"
-  integrity sha512-i6//Lqgn7+7nZA0qVjoYH0085YdNk4MC+tJV4bo+HgjgRMJ0JmkLZzFAuvVioJqLkcGDK5GAMpghZEZkCnwxpQ==
-  dependencies:
-    "@shikijs/core" "1.9.0"
-
 side-channel@^1.0.4:
   version "1.0.6"
   resolved "https://registry.yarnpkg.com/side-channel/-/side-channel-1.0.6.tgz#abd25fb7cd24baf45466406b1096b7831c9215f2"
@@ -4231,21 +4261,21 @@ typedarray.prototype.slice@^1.0.3:
     typed-array-buffer "^1.0.2"
     typed-array-byte-offset "^1.0.2"
 
-typedoc-plugin-markdown@^4.1.0:
-  version "4.1.0"
-  resolved "https://registry.yarnpkg.com/typedoc-plugin-markdown/-/typedoc-plugin-markdown-4.1.0.tgz#0969e82d9821c956145a4b8a9a70f4e00bde27e8"
-  integrity sha512-sUiEJVaa6+MOFShRy14j1OP/VXC5OLyHNecJ2nKeGuBy2M3YiMatSLoIiddFAqVptSuILJTZiJzCBIY6yzAVyg==
+typedoc-plugin-markdown@^4.4.2:
+  version "4.4.2"
+  resolved "https://registry.yarnpkg.com/typedoc-plugin-markdown/-/typedoc-plugin-markdown-4.4.2.tgz#fc31779595aa9bf00e66709f3894e048345bf7ed"
+  integrity sha512-kJVkU2Wd+AXQpyL6DlYXXRrfNrHrEIUgiABWH8Z+2Lz5Sq6an4dQ/hfvP75bbokjNDUskOdFlEEm/0fSVyC7eg==
 
-typedoc@^0.26.1:
-  version "0.26.1"
-  resolved "https://registry.yarnpkg.com/typedoc/-/typedoc-0.26.1.tgz#fc43108abdea64929a2e636877e250d5dea50957"
-  integrity sha512-APsVXqh93jTlpkLuw6+/IORx7n5LN8hzJV8nvMIrYYaIva0VCq0CoDN7Z3hsRThEYVExI/qoFHnAAxrhG+Wd7Q==
+typedoc@^0.27.7:
+  version "0.27.7"
+  resolved "https://registry.yarnpkg.com/typedoc/-/typedoc-0.27.7.tgz#09047ffb5c845f45765de26c68b77260867fe967"
+  integrity sha512-K/JaUPX18+61W3VXek1cWC5gwmuLvYTOXJzBvD9W7jFvbPnefRnCHQCEPw7MSNrP/Hj7JJrhZtDDLKdcYm6ucg==
   dependencies:
+    "@gerrit0/mini-shiki" "^1.24.0"
     lunr "^2.3.9"
     markdown-it "^14.1.0"
-    minimatch "^9.0.4"
-    shiki "^1.9.0"
-    yaml "^2.4.5"
+    minimatch "^9.0.5"
+    yaml "^2.6.1"
 
 typescript@^5.4.5:
   version "5.4.5"
@@ -4468,10 +4498,10 @@ yallist@^4.0.0:
   resolved "https://registry.yarnpkg.com/yallist/-/yallist-4.0.0.tgz#9bb92790d9c0effec63be73519e11a35019a3a72"
   integrity sha512-3wdGidZyq5PB084XLES5TpOSRA3wjXAlIWMhum2kRcv/41Sn2emQ0dycQW4uZXLejwKvg6EsvbdlVL+FYEct7A==
 
-yaml@^2.4.5:
-  version "2.4.5"
-  resolved "https://registry.yarnpkg.com/yaml/-/yaml-2.4.5.tgz#60630b206dd6d84df97003d33fc1ddf6296cca5e"
-  integrity sha512-aBx2bnqDzVOyNKfsysjA2ms5ZlnjSAW2eG3/L5G/CSujfjLJTJsEw1bGw8kCf04KodQWk1pxlGnZ56CRxiawmg==
+yaml@^2.6.1:
+  version "2.7.0"
+  resolved "https://registry.yarnpkg.com/yaml/-/yaml-2.7.0.tgz#aef9bb617a64c937a9a748803786ad8d3ffe1e98"
+  integrity sha512-+hSoy/QHluxmC9kCIJyL/uyFmLmc+e5CFR5Wa+bpIhIj85LVb9ZH2nVnqrHoSvKogwODv0ClqZkmiSSaIH5LTA==
 
 yargs-parser@^20.2.3:
   version "20.2.9"