From d9ed1ef52ec3db945c6ae94a3f18ae8e9c65deab Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Tue, 11 Feb 2025 10:10:54 -0800
Subject: [PATCH] Branching

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

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 5eeafb6ed..a8df9e626 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -329,8 +329,6 @@ export function useStream<
 
   const [flatValues, flatPaths] = (() => {
     const result: ThreadState<StateType>[] = [];
-
-    // TODO: this is kinda ugly
     const flatPaths: Record<
       string,
       { current: string[] | undefined; branches: string[][] | undefined }
@@ -475,7 +473,6 @@ export function useStream<
       // @ts-expect-error
       if (checkpoint != null) delete checkpoint.thread_id;
 
-      // TODO: why non-existent assistant ID does not throw an error here?
       const run = (await client.runs.stream(usableThreadId, assistantId, {
         input: values as Record<string, unknown>,
         config: submitOptions?.config,
@@ -496,6 +493,12 @@ export function useStream<
         streamMode,
       })) as AsyncGenerator<EventStreamEvent>;
 
+      // Unbranch things
+      const newPath = submitOptions?.checkpoint?.checkpoint_id
+        ? flatPaths[submitOptions?.checkpoint?.checkpoint_id]?.current
+        : undefined;
+      if (newPath != null) setBranchPath(newPath ?? []);
+
       // Assumption: we're setting the initial value
       // Used for instant feedback
       setStreamValues(() => {