fix(ai/rsc): Fix types for createStreamableValue and createStreamableUI (#1735)

Co-authored-by: Lars Grammel <lars.grammel@gmail.com>

Author: shuding

.changeset/kind-islands-watch.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix(ai/rsc): Fix types for createStreamableValue and createStreamableUI

packages/core/rsc/streamable.tsx

@@ -17,6 +17,52 @@ import {
 } from './utils';
 import type { StreamablePatch, StreamableValue } from './types';
 
+// It's necessary to define the type manually here, otherwise TypeScript compiler
+// will not be able to infer the correct return type as it's circular.
+type StreamableUIWrapper = {
+  /**
+   * The value of the streamable UI. This can be returned from a Server Action and received by the client.
+   */
+  readonly value: React.ReactNode;
+
+  /**
+   * This method updates the current UI node. It takes a new UI node and replaces the old one.
+   */
+  update(value: React.ReactNode): StreamableUIWrapper;
+
+  /**
+   * This method is used to append a new UI node to the end of the old one.
+   * Once appended a new UI node, the previous UI node cannot be updated anymore.
+   *
+   * @example
+   * ```jsx
+   * const ui = createStreamableUI(<div>hello</div>)
+   * ui.append(<div>world</div>)
+   *
+   * // The UI node will be:
+   * // <>
+   * //   <div>hello</div>
+   * //   <div>world</div>
+   * // </>
+   * ```
+   */
+  append(value: React.ReactNode): StreamableUIWrapper;
+
+  /**
+   * This method is used to signal that there is an error in the UI stream.
+   * It will be thrown on the client side and caught by the nearest error boundary component.
+   */
+  error(error: any): StreamableUIWrapper;
+
+  /**
+   * This method marks the UI node as finalized. You can either call it without any parameters or with a new UI node as the final state.
+   * Once called, the UI node cannot be updated or appended anymore.
+   *
+   * This method is always **required** to be called, otherwise the response will be stuck in a loading state.
+   */
+  done(...args: [React.ReactNode] | []): StreamableUIWrapper;
+};
+
 /**
  * Create a piece of changable UI that can be streamed to the client.
  * On the client side, it can be rendered as a normal React node.
@@ -47,14 +93,8 @@ function createStreamableUI(initialValue?: React.ReactNode) {
   }
   warnUnclosedStream();
 
-  const streamable = {
-    /**
-     * The value of the streamable UI. This can be returned from a Server Action and received by the client.
-     */
+  const streamable: StreamableUIWrapper = {
     value: row,
-    /**
-     * This method updates the current UI node. It takes a new UI node and replaces the old one.
-     */
     update(value: React.ReactNode) {
       assertStream('.update()');
 
@@ -75,22 +115,6 @@ function createStreamableUI(initialValue?: React.ReactNode) {
 
       return streamable;
     },
-    /**
-     * This method is used to append a new UI node to the end of the old one.
-     * Once appended a new UI node, the previous UI node cannot be updated anymore.
-     *
-     * @example
-     * ```jsx
-     * const ui = createStreamableUI(<div>hello</div>)
-     * ui.append(<div>world</div>)
-     *
-     * // The UI node will be:
-     * // <>
-     * //   <div>hello</div>
-     * //   <div>world</div>
-     * // </>
-     * ```
-     */
     append(value: React.ReactNode) {
       assertStream('.append()');
 
@@ -105,10 +129,6 @@ function createStreamableUI(initialValue?: React.ReactNode) {
 
       return streamable;
     },
-    /**
-     * This method is used to signal that there is an error in the UI stream.
-     * It will be thrown on the client side and caught by the nearest error boundary component.
-     */
     error(error: any) {
       assertStream('.error()');
 
@@ -120,12 +140,6 @@ function createStreamableUI(initialValue?: React.ReactNode) {
 
       return streamable;
     },
-    /**
-     * This method marks the UI node as finalized. You can either call it without any parameters or with a new UI node as the final state.
-     * Once called, the UI node cannot be updated or appended anymore.
-     *
-     * This method is always **required** to be called, otherwise the response will be stuck in a loading state.
-     */
     done(...args: [] | [React.ReactNode]) {
       assertStream('.done()');
 
@@ -209,6 +223,59 @@ function createStreamableValue<T = any, E = any>(
   return streamableValue;
 }
 
+// It's necessary to define the type manually here, otherwise TypeScript compiler
+// will not be able to infer the correct return type as it's circular.
+type StreamableValueWrapper<T, E> = {
+  /**
+   * The value of the streamable. This can be returned from a Server Action and
+   * received by the client. To read the streamed values, use the
+   * `readStreamableValue` or `useStreamableValue` APIs.
+   */
+  readonly value: StreamableValue<T, E>;
+
+  /**
+   * This method updates the current value with a new one.
+   */
+  update(value: T): StreamableValueWrapper<T, E>;
+
+  /**
+   * This method is used to append a delta string to the current value. It
+   * requires the current value of the streamable to be a string.
+   *
+   * @example
+   * ```jsx
+   * const streamable = createStreamableValue('hello');
+   * streamable.append(' world');
+   *
+   * // The value will be 'hello world'
+   * ```
+   */
+  append(value: T): StreamableValueWrapper<T, E>;
+
+  /**
+   * This method is used to signal that there is an error in the value stream.
+   * It will be thrown on the client side when consumed via
+   * `readStreamableValue` or `useStreamableValue`.
+   */
+  error(error: any): StreamableValueWrapper<T, E>;
+
+  /**
+   * This method marks the value as finalized. You can either call it without
+   * any parameters or with a new value as the final state.
+   * Once called, the value cannot be updated or appended anymore.
+   *
+   * This method is always **required** to be called, otherwise the response
+   * will be stuck in a loading state.
+   */
+  done(...args: [T] | []): StreamableValueWrapper<T, E>;
+
+  /**
+   * @internal This is an internal lock to prevent the value from being
+   * updated by the user.
+   */
+  [STREAMABLE_VALUE_INTERNAL_LOCK]: boolean;
+};
+
 function createStreamableValueImpl<T = any, E = any>(initialValue?: T) {
   let closed = false;
   let locked = false;
@@ -286,25 +353,13 @@ function createStreamableValueImpl<T = any, E = any>(initialValue?: T) {
     currentValue = value;
   }
 
-  const streamable = {
-    /**
-     * @internal This is an internal lock to prevent the value from being
-     * updated by the user.
-     */
+  const streamable: StreamableValueWrapper<T, E> = {
     set [STREAMABLE_VALUE_INTERNAL_LOCK](state: boolean) {
       locked = state;
     },
-    /**
-     * The value of the streamable. This can be returned from a Server Action and
-     * received by the client. To read the streamed values, use the
-     * `readStreamableValue` or `useStreamableValue` APIs.
-     */
     get value() {
       return createWrapped(true);
     },
-    /**
-     * This method updates the current value with a new one.
-     */
     update(value: T) {
       assertStream('.update()');
 
@@ -319,18 +374,6 @@ function createStreamableValueImpl<T = any, E = any>(initialValue?: T) {
 
       return streamable;
     },
-    /**
-     * This method is used to append a delta string to the current value. It
-     * requires the current value of the streamable to be a string.
-     *
-     * @example
-     * ```jsx
-     * const streamable = createStreamableValue('hello');
-     * streamable.append(' world');
-     *
-     * // The value will be 'hello world'
-     * ```
-     */
     append(value: T) {
       assertStream('.append()');
 
@@ -366,11 +409,6 @@ function createStreamableValueImpl<T = any, E = any>(initialValue?: T) {
 
       return streamable;
     },
-    /**
-     * This method is used to signal that there is an error in the value stream.
-     * It will be thrown on the client side when consumed via
-     * `readStreamableValue` or `useStreamableValue`.
-     */
     error(error: any) {
       assertStream('.error()');
 
@@ -385,14 +423,6 @@ function createStreamableValueImpl<T = any, E = any>(initialValue?: T) {
 
       return streamable;
     },
-    /**
-     * This method marks the value as finalized. You can either call it without
-     * any parameters or with a new value as the final state.
-     * Once called, the value cannot be updated or appended anymore.
-     *
-     * This method is always **required** to be called, otherwise the response
-     * will be stuck in a loading state.
-     */
     done(...args: [] | [T]) {
       assertStream('.done()');
 

packages/core/rsc/streamable.ui.test.tsx

@@ -103,7 +103,7 @@ async function recursiveResolve(val: any): Promise<any> {
   return val;
 }
 
-async function simulateFlightServerRender(node: React.ReactElement) {
+async function simulateFlightServerRender(node: React.ReactNode) {
   async function traverse(node: any): Promise<any> {
     if (!node) return {};
 
@@ -318,7 +318,8 @@ describe('rsc - createStreamableUI()', () => {
     ui.append(<div>2</div>);
     ui.append(<div>3</div>);
 
-    const currentRsolved = ui.value.props.children.props.n;
+    const currentRsolved = (ui.value as React.ReactElement).props.children.props
+      .n;
     const tryResolve1 = await Promise.race([currentRsolved, nextTick()]);
     expect(tryResolve1).toBeDefined();
     const tryResolve2 = await Promise.race([tryResolve1.next, nextTick()]);

packages/core/rsc/utils.tsx

@@ -57,7 +57,7 @@ export function createSuspensedChunk(initialValue: React.ReactNode) {
       <Suspense fallback={initialValue}>
         <R c={initialValue} n={promise} />
       </Suspense>
-    ),
+    ) as React.ReactNode,
     resolve,
     reject,
   };