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.
This commit is contained in:
MauritsBrinkman
2025-06-30 16:43:13 +02:00
committed by Tat Dat Duong
parent d4b4eebe4a
commit ac9b6c416e
+46 -9
View File
@@ -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);