From ac9b6c416eeffa7b97f52af5eb86633bbaaf72f0 Mon Sep 17 00:00:00 2001
From: MauritsBrinkman <maurits@captide.co>
Date: Tue, 17 Jun 2025 14:46:16 +0100
Subject: [PATCH] feat: add onStop callback to useStream for custom stop
 behavior

Add onStop callback to useStream hook enabling developers to customize
UI behavior when streams are stopped. This is especially useful for
UI messages with loading states that need to show "stopped" status
instead of remaining in infinite loading state.

The callback provides the same mutate function as onCustomEvent for
immediate local state updates, while users can optionally update
server thread state using the threads client.

Example usage:
```typescript
const stream = useStream({
  assistantId: "my-assistant",
  onStop: async ({ mutate }) => {
    // Immediate UI update - stop loading components
    mutate((prev) => ({
      ...prev,
      ui: prev.ui?.map(component =>
        component.props?.isLoading
          ? {
              ...component,
              props: {
                ...component.props,
                isLoading: false,
                isStopped: true
              }
            }
          : component
      )
    }));

    // Optional server thread state update
    if (stream.threadId) {
      await stream.client.threads.updateState(stream.threadId, {
        values: {
          ui: prev.ui // persist stopped state to server
        }
      });
    }
  }
});
```

This is especially useful for cases where gen UI components have loading states,
where we don't want the loading state to persist on cancellation.
---
 libs/sdk-js/src/react/stream.tsx | 55 ++++++++++++++++++++++++++------
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx
index 73288dc54..21da1c6ec 100644
--- a/libs/sdk-js/src/react/stream.tsx
+++ b/libs/sdk-js/src/react/stream.tsx
@@ -511,6 +511,31 @@ export interface UseStreamOptions<
    */
   onDebugEvent?: (data: DebugStreamEvent["data"]) => void;
 
+  /**
+   * Callback that is called when the stream is stopped by the user.
+   * Provides a mutate function to update the stream state immediately
+   * without requiring a server roundtrip.
+   * 
+   * @example
+   * ```typescript
+   * onStop: ({ mutate }) => {
+   *   mutate((prev) => ({
+   *     ...prev,
+   *     ui: prev.ui?.map(component => 
+   *       component.props.isLoading
+   *         ? { ...component, props: { ...component.props, stopped: true, isLoading: false }}
+   *         : component
+   *     )
+   *   }));
+   * }
+   * ```
+   */
+  onStop?: (options: {
+    mutate: (
+      update: Partial<StateType> | ((prev: StateType) => Partial<StateType>),
+    ) => void;
+  }) => void;
+
   /**
    * The ID of the thread to fetch history and current values from.
    */
@@ -862,6 +887,21 @@ export function useStream<
     );
   })();
 
+  // Create a reusable mutate function for both onCustomEvent and onStop callbacks
+  const mutateStreamValues = useCallback(
+    (update: Partial<StateType> | ((prev: StateType) => Partial<StateType>)) => {
+      setStreamValues((prev) => {
+        // should not happen
+        if (prev == null) return prev;
+        return {
+          ...prev,
+          ...(typeof update === "function" ? update(prev) : update),
+        };
+      });
+    },
+    [],
+  );
+
   const stop = () => {
     if (abortRef.current != null) abortRef.current.abort();
     abortRef.current = null;
@@ -871,6 +911,11 @@ export function useStream<
       if (runId) client.runs.cancel(threadId, runId);
       runMetadataStorage.removeItem(`lg:stream:${threadId}`);
     }
+
+    // Call onStop callback with mutate function
+    if (options.onStop) {
+      options.onStop({ mutate: mutateStreamValues });
+    }
   };
 
   async function consumeStream(
@@ -903,15 +948,7 @@ export function useStream<
         if (event === "updates") options.onUpdateEvent?.(data);
         if (event === "custom")
           options.onCustomEvent?.(data, {
-            mutate: (update) =>
-              setStreamValues((prev) => {
-                // should not happen
-                if (prev == null) return prev;
-                return {
-                  ...prev,
-                  ...(typeof update === "function" ? update(prev) : update),
-                };
-              }),
+            mutate: mutateStreamValues,
           });
         if (event === "metadata") options.onMetadataEvent?.(data);
         if (event === "events") options.onLangChainEvent?.(data);