fix(ai/rsc): Optimize streamable value (#1130)

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

Author: shuding

.changeset/large-rivers-tease.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+ai/rsc: optimize streamable value stream size

packages/core/rsc/shared-client/streamable.ui.test.tsx

@@ -1,6 +1,10 @@
 import { createStreamableValue } from '../streamable';
 import { readStreamableValue } from './streamable';
 
+function nextTick() {
+  return Promise.resolve();
+}
+
 describe('rsc - readStreamableValue()', () => {
   it('should return an async iterable', () => {
     const streamable = createStreamableValue();
@@ -11,20 +15,81 @@ describe('rsc - readStreamableValue()', () => {
     expect(result[Symbol.asyncIterator]).toBeDefined();
   });
 
-  it('should be able to read all values', async () => {
-    const streamable = createStreamableValue(0);
+  it('should directly emit the final value when reading .value', async () => {
+    const streamable = createStreamableValue('1');
+    streamable.update('2');
+    streamable.update('3');
+
+    expect(streamable.value).toMatchInlineSnapshot(`
+      {
+        "curr": "3",
+        "next": Promise {},
+        "type": Symbol(ui.streamable.value),
+      }
+    `);
+
+    streamable.done('4');
+
+    expect(streamable.value).toMatchInlineSnapshot(`
+      {
+        "curr": "4",
+        "type": Symbol(ui.streamable.value),
+      }
+    `);
+  });
+
+  it('should be able to stream any JSON values', async () => {
+    const streamable = createStreamableValue();
+    streamable.update({ v: 123 });
+
+    expect(streamable.value).toMatchInlineSnapshot(`
+      {
+        "curr": {
+          "v": 123,
+        },
+        "next": Promise {},
+        "type": Symbol(ui.streamable.value),
+      }
+    `);
+
+    streamable.done();
+  });
 
-    streamable.update(1);
-    streamable.update(2);
-    streamable.done(3);
+  it('should support .error()', async () => {
+    const streamable = createStreamableValue();
+    streamable.error('This is an error');
+
+    expect(streamable.value).toMatchInlineSnapshot(`
+      {
+        "error": "This is an error",
+        "type": Symbol(ui.streamable.value),
+      }
+    `);
+  });
+
+  it('should support reading streamed values and errors', async () => {
+    const streamable = createStreamableValue(1);
+    (async () => {
+      await nextTick();
+      streamable.update(2);
+      await nextTick();
+      streamable.update(3);
+      await nextTick();
+      streamable.error('This is an error');
+    })();
 
     const values = [];
-    for await (const v of readStreamableValue(streamable.value)) {
-      values.push(v);
+
+    try {
+      for await (const v of readStreamableValue(streamable.value)) {
+        values.push(v);
+      }
+    } catch (e) {
+      expect(e).toMatchInlineSnapshot(`"This is an error"`);
     }
+
     expect(values).toMatchInlineSnapshot(`
       [
-        0,
         1,
         2,
         3,

packages/core/rsc/streamable.tsx

@@ -113,7 +113,12 @@ export function createStreamableUI(initialValue?: React.ReactNode) {
  */
 export function createStreamableValue<T = any, E = any>(initialValue?: T) {
   let closed = false;
-  let { promise, resolve } = createResolvablePromise<StreamableValue<T, E>>();
+  let resolvable = createResolvablePromise<StreamableValue<T, E>>();
+
+  let currentValue = initialValue;
+  let currentError: E | undefined;
+  let currentPromise: typeof resolvable.promise | undefined =
+    resolvable.promise;
 
   function assertStream(method: string) {
     if (closed) {
@@ -136,35 +141,37 @@ export function createStreamableValue<T = any, E = any>(initialValue?: T) {
   }
   warnUnclosedStream();
 
-  function createWrapped(
-    val: T | undefined,
-    initial?: boolean,
-  ): StreamableValue<T, E> {
-    if (initial) {
-      return {
-        type: STREAMABLE_VALUE_TYPE,
-        curr: val,
-        next: promise,
-      };
+  function createWrapped(withType?: boolean): StreamableValue<T, E> {
+    // This makes the payload much smaller if there're mutative updates before the first read.
+    const init: Partial<StreamableValue<T, E>> =
+      currentError === undefined
+        ? { curr: currentValue }
+        : { error: currentError };
+
+    if (currentPromise) {
+      init.next = currentPromise;
     }
 
-    return {
-      curr: val,
-      next: promise,
-    };
+    if (withType) {
+      init.type = STREAMABLE_VALUE_TYPE;
+    }
+
+    return init;
   }
 
   return {
-    value: createWrapped(initialValue, true),
+    get value() {
+      return createWrapped(true);
+    },
     update(value: T) {
       assertStream('.update()');
 
-      const resolvePrevious = resolve;
-      const resolvable = createResolvablePromise();
-      promise = resolvable.promise;
-      resolve = resolvable.resolve;
+      const resolvePrevious = resolvable.resolve;
+      resolvable = createResolvablePromise();
 
-      resolvePrevious(createWrapped(value));
+      currentValue = value;
+      currentPromise = resolvable.promise;
+      resolvePrevious(createWrapped());
 
       warnUnclosedStream();
     },
@@ -175,7 +182,10 @@ export function createStreamableValue<T = any, E = any>(initialValue?: T) {
         clearTimeout(warningTimeout);
       }
       closed = true;
-      resolve({ error });
+      currentError = error;
+      currentPromise = undefined;
+
+      resolvable.resolve({ error });
     },
     done(...args: any) {
       assertStream('.done()');
@@ -184,13 +194,15 @@ export function createStreamableValue<T = any, E = any>(initialValue?: T) {
         clearTimeout(warningTimeout);
       }
       closed = true;
+      currentPromise = undefined;
 
       if (args.length) {
-        resolve({ curr: args[0] });
+        currentValue = args[0];
+        resolvable.resolve({ curr: args[0] });
         return;
       }
 
-      resolve({});
+      resolvable.resolve({});
     },
   };
 }