diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx index 6d3f96b77..3ffdf7b84 100644 --- a/libs/sdk-js/src/react/stream.tsx +++ b/libs/sdk-js/src/react/stream.tsx @@ -804,14 +804,21 @@ export function useStream< ); })(); - const stop = useCallback(() => { + const stop = () => { if (abortRef.current != null) abortRef.current.abort(); abortRef.current = null; - }, []); + + if (options.joinOnMount && threadId) { + const rejoinKey = `lg:rejoin:${threadId}`; + const runId = window.localStorage.getItem(rejoinKey); + if (runId) client.runs.cancel(threadId, runId); + window.localStorage.removeItem(rejoinKey); + } + }; async function consumeStream( action: (signal: AbortSignal) => Promise<{ - onFinish: () => Promise[]>; + onSuccess: () => Promise[]>; stream: AsyncGenerator; }>, ) { @@ -877,7 +884,7 @@ export function useStream< } // TODO: stream created checkpoints to avoid an unnecessary network request - const result = await run.onFinish(); + const result = await run.onSuccess(); setStreamValues(null); if (streamError != null) throw streamError; @@ -918,8 +925,8 @@ export function useStream< }) as AsyncGenerator; return { - onFinish: async () => { - if (rejoinKey) window.localStorage.removeItem(rejoinKey); + onSuccess: () => { + window.localStorage.removeItem(rejoinKey); return history.mutate(threadId); }, stream, @@ -1005,7 +1012,7 @@ export function useStream< return { stream, - onFinish: () => { + onSuccess: () => { if (rejoinKey) window.localStorage.removeItem(rejoinKey); return history.mutate(usableThreadId); }, @@ -1014,6 +1021,7 @@ export function useStream< }; const joinStreamRef = useRef(joinStream); + const autoJoinRef = useRef(options.joinOnMount); joinStreamRef.current = joinStream; const joinKey = useMemo(() => { @@ -1025,7 +1033,10 @@ export function useStream< }, [joinOnMount, isLoading, threadId]); useEffect(() => { - if (joinKey) joinStreamRef.current?.(joinKey.runId, "-1"); + if (joinKey && autoJoinRef.current) { + autoJoinRef.current = false; + joinStreamRef.current?.(joinKey.runId, "-1"); + } }, [joinKey]); const error = streamError ?? historyError;