From a54ea7754d443bea2eb6ecc92ec155261cf09041 Mon Sep 17 00:00:00 2001
From: Shu Ding <g@shud.in>
Date: Wed, 27 Mar 2024 13:40:10 +0100
Subject: [PATCH] feat(ai/rsc): Add `useStreamableValue` again (#1233)

---
 .changeset/smooth-rockets-sneeze.md           |   5 +
 packages/core/rsc/index.ts                    |   1 +
 packages/core/rsc/rsc-client.ts               |   1 +
 packages/core/rsc/rsc-shared.mts              |   1 +
 packages/core/rsc/shared-client/index.ts      |   2 +-
 .../core/rsc/shared-client/streamable.tsx     | 106 ++++++++++++++++--
 packages/core/rsc/streamable.tsx              |   2 +-
 packages/core/rsc/types.ts                    |   2 +-
 8 files changed, 110 insertions(+), 10 deletions(-)
 create mode 100644 .changeset/smooth-rockets-sneeze.md

diff --git a/.changeset/smooth-rockets-sneeze.md b/.changeset/smooth-rockets-sneeze.md
new file mode 100644
index 000000000000..e0bdd10e8dfa
--- /dev/null
+++ b/.changeset/smooth-rockets-sneeze.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat(ai/rsc): add `useStreamableValue`
diff --git a/packages/core/rsc/index.ts b/packages/core/rsc/index.ts
index 69b00a69af25..b197c69c87b4 100644
--- a/packages/core/rsc/index.ts
+++ b/packages/core/rsc/index.ts
@@ -9,6 +9,7 @@ export type {
 
 export type {
   readStreamableValue,
+  useStreamableValue,
   useUIState,
   useAIState,
   useActions,
diff --git a/packages/core/rsc/rsc-client.ts b/packages/core/rsc/rsc-client.ts
index af9b9413134c..266db2bc6504 100644
--- a/packages/core/rsc/rsc-client.ts
+++ b/packages/core/rsc/rsc-client.ts
@@ -1,5 +1,6 @@
 export {
   readStreamableValue,
+  useStreamableValue,
   useUIState,
   useAIState,
   useActions,
diff --git a/packages/core/rsc/rsc-shared.mts b/packages/core/rsc/rsc-shared.mts
index 936ba7dcdc16..212b42b63bc3 100644
--- a/packages/core/rsc/rsc-shared.mts
+++ b/packages/core/rsc/rsc-shared.mts
@@ -2,6 +2,7 @@
 
 export {
   readStreamableValue,
+  useStreamableValue,
   useUIState,
   useAIState,
   useActions,
diff --git a/packages/core/rsc/shared-client/index.ts b/packages/core/rsc/shared-client/index.ts
index 5130dbb7a6dd..c23f210b89a9 100644
--- a/packages/core/rsc/shared-client/index.ts
+++ b/packages/core/rsc/shared-client/index.ts
@@ -1,6 +1,6 @@
 'use client';
 
-export { readStreamableValue } from './streamable';
+export { readStreamableValue, useStreamableValue } from './streamable';
 export {
   useUIState,
   useAIState,
diff --git a/packages/core/rsc/shared-client/streamable.tsx b/packages/core/rsc/shared-client/streamable.tsx
index 2f640dd09e79..a947a8717f20 100644
--- a/packages/core/rsc/shared-client/streamable.tsx
+++ b/packages/core/rsc/shared-client/streamable.tsx
@@ -1,21 +1,38 @@
+import { startTransition, useLayoutEffect, useState } from 'react';
 import { STREAMABLE_VALUE_TYPE } from '../constants';
 import type { StreamableValue } from '../types';
 
+function hasReadableValueSignature(value: unknown): value is StreamableValue {
+  return !!(
+    value &&
+    typeof value === 'object' &&
+    'type' in value &&
+    value.type === STREAMABLE_VALUE_TYPE
+  );
+}
+
 function assertStreamableValue(
   value: unknown,
 ): asserts value is StreamableValue {
-  if (
-    !value ||
-    typeof value !== 'object' ||
-    !('type' in value) ||
-    value.type !== STREAMABLE_VALUE_TYPE
-  ) {
+  if (!hasReadableValueSignature(value)) {
     throw new Error(
-      'Invalid value: this hook only accepts values created via `createStreamableValue` from the server.',
+      'Invalid value: this hook only accepts values created via `createStreamableValue`.',
     );
   }
 }
 
+function isStreamableValue(value: unknown): value is StreamableValue {
+  const hasSignature = hasReadableValueSignature(value);
+
+  if (!hasSignature && typeof value !== 'undefined') {
+    throw new Error(
+      'Invalid value: this hook only accepts values created via `createStreamableValue`.',
+    );
+  }
+
+  return hasSignature;
+}
+
 /**
  * `readStreamableValue` takes a streamable value created via the `createStreamableValue().value` API,
  * and returns an async iterator.
@@ -122,3 +139,78 @@ export function readStreamableValue<T = unknown>(
     },
   };
 }
+
+/**
+ * `useStreamableValue` is a React hook that takes a streamable value created via the `createStreamableValue().value` API,
+ * and returns the current value, error, and pending state.
+ *
+ * This is useful for consuming streamable values received from a component's props. For example:
+ *
+ * ```js
+ * function MyComponent({ streamableValue }) {
+ *   const [data, error, pending] = useStreamableValue(streamableValue);
+ *
+ *   if (pending) return <div>Loading...</div>;
+ *   if (error) return <div>Error: {error.message}</div>;
+ *
+ *   return <div>Data: {data}</div>;
+ * }
+ * ```
+ */
+export function useStreamableValue<T = unknown, Error = unknown>(
+  streamableValue?: StreamableValue<T>,
+): [data: T | undefined, error: Error | undefined, pending: boolean] {
+  const [curr, setCurr] = useState<T | undefined>(
+    isStreamableValue(streamableValue) ? streamableValue.curr : undefined,
+  );
+  const [error, setError] = useState<Error | undefined>(
+    isStreamableValue(streamableValue) ? streamableValue.error : undefined,
+  );
+  const [pending, setPending] = useState<boolean>(
+    isStreamableValue(streamableValue) ? !!streamableValue.next : false,
+  );
+
+  useLayoutEffect(() => {
+    if (!isStreamableValue(streamableValue)) return;
+
+    let cancelled = false;
+
+    const iterator = readStreamableValue(streamableValue);
+    if (streamableValue.next) {
+      startTransition(() => {
+        if (cancelled) return;
+        setPending(true);
+      });
+    }
+
+    (async () => {
+      try {
+        for await (const value of iterator) {
+          if (cancelled) return;
+          startTransition(() => {
+            if (cancelled) return;
+            setCurr(value);
+          });
+        }
+      } catch (e) {
+        if (cancelled) return;
+        startTransition(() => {
+          if (cancelled) return;
+          setError(e as Error);
+        });
+      } finally {
+        if (cancelled) return;
+        startTransition(() => {
+          if (cancelled) return;
+          setPending(false);
+        });
+      }
+    })();
+
+    return () => {
+      cancelled = true;
+    };
+  }, [streamableValue]);
+
+  return [curr, error, pending];
+}
diff --git a/packages/core/rsc/streamable.tsx b/packages/core/rsc/streamable.tsx
index 09799b545d64..f5814c00f1f7 100644
--- a/packages/core/rsc/streamable.tsx
+++ b/packages/core/rsc/streamable.tsx
@@ -215,7 +215,7 @@ export function createStreamableValue<T = any, E = any>(initialValue?: T) {
     /**
      * 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` API.
+     * `readStreamableValue` or `useStreamableValue` APIs.
      */
     get value() {
       return createWrapped(true);
diff --git a/packages/core/rsc/types.ts b/packages/core/rsc/types.ts
index af0044971f46..50a60ec8e786 100644
--- a/packages/core/rsc/types.ts
+++ b/packages/core/rsc/types.ts
@@ -91,7 +91,7 @@ export type StreamablePatch = undefined | [0, string]; // Append string.
 
 /**
  * StreamableValue is a value that can be streamed over the network via AI Actions.
- * To read the streamed values, use the `readStreamableValue` API.
+ * To read the streamed values, use the `readStreamableValue` or `useStreamableValue` APIs.
  */
 export type StreamableValue<T = any, E = any> = {
   /**