Add warning when streamables are not closed for a while during development (#1105)

Co-authored-by: Max Leiter <max.leiter@vercel.com>

Author: shuding

Diff for: .changeset/good-masks-suffer.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+ai/rsc: improve dev error and warnings by trying to detect hanging streams

Diff for: packages/core/rsc/constants.ts

@@ -1 +1,2 @@
 export const STREAMABLE_VALUE_TYPE = Symbol.for('ui.streamable.value');
+export const DEV_DEFAULT_STREAMABLE_WARNING_TIME = 15 * 1000;

Diff for: packages/core/rsc/streamable.tsx

@@ -6,7 +6,10 @@ import zodToJsonSchema from 'zod-to-json-schema';
 // TODO: This needs to be externalized.
 import { OpenAIStream } from '../streams';
 
-import { STREAMABLE_VALUE_TYPE } from './constants';
+import {
+  STREAMABLE_VALUE_TYPE,
+  DEV_DEFAULT_STREAMABLE_WARNING_TIME,
+} from './constants';
 import {
   createResolvablePromise,
   createSuspensedChunk,
@@ -22,26 +25,43 @@ export function createStreamableUI(initialValue?: React.ReactNode) {
   let closed = false;
   let { row, resolve, reject } = createSuspensedChunk(initialValue);
 
-  function assertStream() {
+  function assertStream(method: string) {
     if (closed) {
-      throw new Error('UI stream is already closed.');
+      throw new Error(method + ': UI stream is already closed.');
     }
   }
 
+  let warningTimeout: NodeJS.Timeout | undefined;
+  function warnUnclosedStream() {
+    if (process.env.NODE_ENV === 'development') {
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
+      warningTimeout = setTimeout(() => {
+        console.warn(
+          'The streamable UI has been slow to update. This may be a bug or a performance issue or you forgot to call `.done()`.',
+        );
+      }, DEV_DEFAULT_STREAMABLE_WARNING_TIME);
+    }
+  }
+  warnUnclosedStream();
+
   return {
     value: row,
     update(value: React.ReactNode) {
-      assertStream();
+      assertStream('.update()');
 
       const resolvable = createResolvablePromise();
       currentValue = value;
 
       resolve({ value: currentValue, done: false, next: resolvable.promise });
       resolve = resolvable.resolve;
       reject = resolvable.reject;
+
+      warnUnclosedStream();
     },
     append(value: React.ReactNode) {
-      assertStream();
+      assertStream('.append()');
 
       const resolvable = createResolvablePromise();
       if (typeof currentValue === 'string' && typeof value === 'string') {
@@ -58,16 +78,24 @@ export function createStreamableUI(initialValue?: React.ReactNode) {
       resolve({ value: currentValue, done: false, next: resolvable.promise });
       resolve = resolvable.resolve;
       reject = resolvable.reject;
+
+      warnUnclosedStream();
     },
     error(error: any) {
-      assertStream();
+      assertStream('.error()');
 
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
       closed = true;
       reject(error);
     },
     done(...args: any) {
-      assertStream();
+      assertStream('.done()');
 
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
       closed = true;
       if (args.length) {
         resolve({ value: args[0], done: true });
@@ -83,15 +111,29 @@ export function createStreamableUI(initialValue?: React.ReactNode) {
  * On the client side, the value can be accessed via the useStreamableValue() hook.
  */
 export function createStreamableValue<T = any>(initialValue?: T) {
-  // let currentValue = initialValue
   let closed = false;
   let { promise, resolve, reject } = createResolvablePromise();
 
-  function assertStream() {
+  function assertStream(method: string) {
     if (closed) {
-      throw new Error('Value stream is already closed.');
+      throw new Error(method + ': Value stream is already closed.');
+    }
+  }
+
+  let warningTimeout: NodeJS.Timeout | undefined;
+  function warnUnclosedStream() {
+    if (process.env.NODE_ENV === 'development') {
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
+      warningTimeout = setTimeout(() => {
+        console.warn(
+          'The streamable UI has been slow to update. This may be a bug or a performance issue or you forgot to call `.done()`.',
+        );
+      }, DEV_DEFAULT_STREAMABLE_WARNING_TIME);
     }
   }
+  warnUnclosedStream();
 
   function createWrapped(val: T | undefined, initial?: boolean) {
     if (initial) {
@@ -111,7 +153,7 @@ export function createStreamableValue<T = any>(initialValue?: T) {
   return {
     value: createWrapped(initialValue, true),
     update(value: T) {
-      assertStream();
+      assertStream('.update()');
 
       const resolvePrevious = resolve;
       const resolvable = createResolvablePromise();
@@ -121,17 +163,23 @@ export function createStreamableValue<T = any>(initialValue?: T) {
 
       resolvePrevious(createWrapped(value));
 
-      // currentValue = value
+      warnUnclosedStream();
     },
     error(error: any) {
-      assertStream();
+      assertStream('.error()');
 
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
       closed = true;
       reject(error);
     },
     done(...args: any) {
-      assertStream();
+      assertStream('.done()');
 
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
       closed = true;
 
       if (args.length) {

Diff for: packages/core/rsc/streamable.ui.test.tsx

@@ -319,4 +319,15 @@ describe('rsc - createStreamableUI()', () => {
     );
     expect(final).toMatchInlineSnapshot('"hello world!"');
   });
+
+  it('should error when updating a closed streamable', async () => {
+    const ui = createStreamableUI(<div>1</div>);
+    ui.done(<div>2</div>);
+
+    expect(() => {
+      ui.update(<div>3</div>);
+    }).toThrowErrorMatchingInlineSnapshot(
+      '".update(): UI stream is already closed."',
+    );
+  });
 });