fix (ai/react): useObject stop & isLoading reset at end of stream (#2126

)
vercel · Jun 27, 2024 · 9b50003 · 9b50003
1 parent 83204b7
commit 9b50003
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/.changeset/gentle-dolphins-boil.md b/.changeset/gentle-dolphins-boil.md
@@ -0,0 +1,5 @@
+---
+'@ai-sdk/react': patch
+---
+
+fix (ai/react): useObject stop & isLoading reset at end of stream
diff --git a/packages/react/src/use-object.ts b/packages/react/src/use-object.ts
@@ -148,15 +148,17 @@ function useObject<RESULT, INPUT = any>({
               mutate(currentObject);
             }
           },
+
+          close() {
+            setIsLoading(false);
+            abortControllerRef.current = null;
+          },
         }),
       );
 
       setError(undefined);
     } catch (error) {
       setError(error);
-    } finally {
-      setIsLoading(false);
-      abortControllerRef.current = null;
     }
   };
 

diff --git a/packages/react/src/use-object.ui.test.tsx b/packages/react/src/use-object.ui.test.tsx
@@ -111,8 +111,10 @@ describe('text stream', () => {
       streamController.enqueue('"Hello, world!"}');
       streamController.close();
 
-      await screen.findByTestId('loading');
-      expect(screen.getByTestId('loading')).toHaveTextContent('false');
+      // wait for element "loading" to have text content "false":
+      await waitFor(() => {
+        expect(screen.getByTestId('loading')).toHaveTextContent('true');
+      });
     });
   });