Cancel on unmount + clearData method

packages/react/src/hooks/use-stream.ts

@@ -30,6 +30,10 @@ const resolveListener = (id: string) => {
     return listeners.get(id)!;
 };
 
+const hasListeners = (id: string) => {
+    return listeners.has(id) && listeners.get(id)?.length;
+};
+
 const addListener = (id: string, listener: StreamListenerCallback) => {
     resolveListener(id).push(listener);
 
@@ -38,6 +42,11 @@ const addListener = (id: string, listener: StreamListenerCallback) => {
             id,
             resolveListener(id).filter((l) => l !== listener),
         );
+
+        if (!hasListeners(id)) {
+            streams.delete(id);
+            listeners.delete(id);
+        }
     };
 };
 
@@ -95,6 +104,12 @@ export const useStream = (url: string, options: StreamOptions = {}) => {
         });
     }, [isFetching, isStreaming]);
 
+    const clearData = useCallback(() => {
+        updateStream({
+            data: "",
+        });
+    }, []);
+
     const makeRequest = useCallback(
         (body: Record<string, any> = {}) => {
             const controller = new AbortController();
@@ -150,9 +165,7 @@ export const useStream = (url: string, options: StreamOptions = {}) => {
     const send = useCallback((body: Record<string, any>) => {
         cancel();
         makeRequest(body);
-        updateStream({
-            data: "",
-        });
+        clearData();
     }, []);
 
     const read = useCallback(
@@ -200,6 +213,10 @@ export const useStream = (url: string, options: StreamOptions = {}) => {
 
         return () => {
             stopListening();
+
+            if (!hasListeners(id.current)) {
+                cancel();
+            }
         };
     }, []);
 
@@ -224,5 +241,6 @@ export const useStream = (url: string, options: StreamOptions = {}) => {
         id: id.current,
         send,
         cancel,
+        clearData,
     };
 };

packages/react/tests/use-stream.test.ts

@@ -75,6 +75,12 @@ describe("useStream", () => {
 
         expect(result.current.isStreaming).toBe(false);
         expect(result.current.data).toBe("chunk1chunk2");
+
+        await act(() => {
+            result.current.clearData();
+        });
+
+        expect(result.current.data).toBe("");
     });
 
     it("can send data back to the endpoint", async () => {
@@ -364,4 +370,25 @@ describe("useStream", () => {
 
         expect(capturedHeaders.get("X-STREAM-ID")).toBe(id);
     });
+
+    it.skip("should cancel stream when component unmounts", async () => {
+        const onCancel = vi.fn();
+        const { unmount, result } = renderHook(() =>
+            useStream(url, { onCancel }),
+        );
+
+        await act(() => {
+            result.current.send({
+                test: "ok",
+            });
+        });
+
+        await waitFor(() => expect(result.current.isStreaming).toBe(true));
+
+        unmount();
+
+        await waitFor(() => expect(result.current.isStreaming).toBe(false));
+
+        expect(onCancel).toHaveBeenCalled();
+    });
 });

packages/vue/src/composables/useStream.ts

@@ -30,6 +30,10 @@ const resolveListener = (id: string) => {
     return listeners.get(id)!;
 };
 
+const hasListeners = (id: string) => {
+    return listeners.has(id) && listeners.get(id)?.length;
+};
+
 const addListener = (id: string, listener: StreamListenerCallback) => {
     resolveListener(id).push(listener);
 
@@ -38,6 +42,11 @@ const addListener = (id: string, listener: StreamListenerCallback) => {
             id,
             resolveListener(id).filter((l) => l !== listener),
         );
+
+        if (!hasListeners(id)) {
+            streams.delete(id);
+            listeners.delete(id);
+        }
     };
 };
 
@@ -145,6 +154,10 @@ export const useStream = (url: string, options: StreamOptions = {}) => {
     const send = (body: Record<string, any>) => {
         cancel();
         makeRequest(body);
+        clearData();
+    };
+
+    const clearData = () => {
         updateStream({
             data: "",
         });
@@ -197,6 +210,10 @@ export const useStream = (url: string, options: StreamOptions = {}) => {
     onUnmounted(() => {
         stopListening();
         window.removeEventListener("beforeunload", cancel);
+
+        if (!hasListeners(id)) {
+            cancel();
+        }
     });
 
     return {
@@ -206,5 +223,6 @@ export const useStream = (url: string, options: StreamOptions = {}) => {
         id,
         send,
         cancel,
+        clearData,
     };
 };

packages/vue/tests/useStream.test.ts

@@ -87,6 +87,10 @@ describe("useStream", () => {
         await vi.waitFor(() => expect(result.isStreaming.value).toBe(false));
 
         expect(result.data.value).toBe("chunk1chunk2");
+
+        result.clearData();
+
+        expect(result.data.value).toBe("");
     });
 
     it("can send data to the endpoint", async () => {