feat: RSC (#1008)

.changeset/nine-trainers-return.md

@@ -0,0 +1,5 @@
+---
+'ai': major
+---
+
+add ai/rsc

packages/core/package.json

@@ -12,11 +12,12 @@
     "svelte/dist/**/*",
     "vue/dist/**/*",
     "solid/dist/**/*",
-    "prompts/dist/**/*"
+    "prompts/dist/**/*",
+    "rsc/dist/**/*"
   ],
   "scripts": {
     "build": "tsup && cat react/dist/index.server.d.ts >> react/dist/index.d.ts",
-    "clean": "rm -rf dist && rm -rf react/dist && rm -rf svelte/dist && rm -rf vue/dist && rm -rf solid/dist",
+    "clean": "rm -rf dist && rm -rf react/dist && rm -rf svelte/dist && rm -rf vue/dist && rm -rf solid/dist && rm -rf rsc/dist",
     "dev": "tsup --watch",
     "lint": "eslint \"./**/*.ts*\"",
     "type-check": "tsc --noEmit",
@@ -42,6 +43,12 @@
       "module": "./dist/index.mjs",
       "require": "./dist/index.js"
     },
+    "./rsc": {
+      "types": "./rsc/dist/rsc-types.d.ts",
+      "react-server": "./rsc/dist/rsc-server.mjs",
+      "import": "./rsc/dist/rsc-client.mjs",
+      "module": "./rsc/dist/rsc-client.mjs"
+    },
     "./prompts": {
       "types": "./prompts/dist/index.d.ts",
       "import": "./prompts/dist/index.mjs",
@@ -76,12 +83,14 @@
   },
   "dependencies": {
     "eventsource-parser": "1.0.0",
+    "jsondiffpatch": "^0.6.0",
     "nanoid": "3.3.6",
     "solid-swr-store": "0.10.7",
     "sswr": "2.0.0",
     "swr": "2.2.0",
     "swr-store": "0.10.6",
-    "swrv": "1.0.4"
+    "swrv": "1.0.4",
+    "zod-to-json-schema": "^3.22.4"
   },
   "devDependencies": {
     "@anthropic-ai/sdk": "0.12.0",
@@ -118,7 +127,8 @@
     "react": "^18.2.0",
     "solid-js": "^1.7.7",
     "svelte": "^3.0.0 || ^4.0.0",
-    "vue": "^3.3.4"
+    "vue": "^3.3.4",
+    "zod": "^3.0.0"
   },
   "peerDependenciesMeta": {
     "react": {
@@ -132,6 +142,9 @@
     },
     "solid-js": {
       "optional": true
+    },
+    "zod": {
+      "optional": true
     }
   },
   "engines": {

packages/core/rsc/ai-state.tsx

@@ -0,0 +1,207 @@
+import { AsyncLocalStorage } from 'async_hooks';
+import * as jsondiffpatch from 'jsondiffpatch';
+import { createResolvablePromise, isFunction } from './utils';
+import type {
+  AIProvider,
+  InternalAIStateStorageOptions,
+  InferAIState,
+  MutableAIState,
+  ValueOrUpdater,
+} from './types';
+
+// It is possible that multiple AI requests get in concurrently, for different
+// AI instances. So ALS is necessary here for a simpler API.
+const asyncAIStateStorage = new AsyncLocalStorage<{
+  currentState: any;
+  originalState: any;
+  sealed: boolean;
+  options: InternalAIStateStorageOptions;
+  mutationDeltaPromise?: Promise<any>;
+  mutationDeltaResolve?: (v: any) => void;
+}>();
+
+function getAIStateStoreOrThrow(message: string) {
+  const store = asyncAIStateStorage.getStore();
+  if (!store) {
+    throw new Error(message);
+  }
+  return store;
+}
+
+export function withAIState<S, T>(
+  { state, options }: { state: S; options: InternalAIStateStorageOptions },
+  fn: () => T,
+): T {
+  return asyncAIStateStorage.run(
+    {
+      currentState: state,
+      originalState: state,
+      sealed: false,
+      options,
+    },
+    fn,
+  );
+}
+
+export function getAIStateDeltaPromise() {
+  const store = getAIStateStoreOrThrow('Internal error occurred.');
+  return store.mutationDeltaPromise;
+}
+
+// Internal method. This will be called after the AI Action has been returned
+// and you can no longer call `getMutableAIState()` inside any async callbacks
+// created by that Action.
+export function sealMutableAIState() {
+  const store = getAIStateStoreOrThrow('Internal error occurred.');
+  store.sealed = true;
+}
+
+/**
+ * Get the current AI state.
+ * If `key` is provided, it will return the value of the specified key in the
+ * AI state, if it's an object. If it's not an object, it will throw an error.
+ *
+ * @example const state = getAIState() // Get the entire AI state
+ * @example const field = getAIState('key') // Get the value of the key
+ */
+function getAIState<AI extends AIProvider = any>(): InferAIState<AI, any>;
+function getAIState<AI extends AIProvider = any>(
+  key: keyof InferAIState<AI, any>,
+): InferAIState<AI, any>[typeof key];
+function getAIState<AI extends AIProvider = any>(
+  ...args: [] | [key: keyof InferAIState<AI, any>]
+) {
+  const store = getAIStateStoreOrThrow(
+    '`getAIState` must be called within an AI Action.',
+  );
+
+  if (args.length > 0) {
+    const key = args[0];
+    if (typeof store.currentState !== 'object') {
+      throw new Error(
+        `You can't get the "${String(
+          key,
+        )}" field from the AI state because it's not an object.`,
+      );
+    }
+    return store.currentState[key as keyof typeof store.currentState];
+  }
+
+  return store.currentState;
+}
+
+/**
+ * Get the mutable AI state. Note that you must call `.close()` when finishing
+ * updating the AI state.
+ *
+ * @example
+ * ```tsx
+ * const state = getMutableAIState()
+ * state.update({ ...state.get(), key: 'value' })
+ * state.update((currentState) => ({ ...currentState, key: 'value' }))
+ * state.done()
+ * ```
+ *
+ * @example
+ * ```tsx
+ * const state = getMutableAIState()
+ * state.done({ ...state.get(), key: 'value' }) // Done with a new state
+ * ```
+ */
+function getMutableAIState<AI extends AIProvider = any>(): MutableAIState<
+  InferAIState<AI, any>
+>;
+function getMutableAIState<AI extends AIProvider = any>(
+  key: keyof InferAIState<AI, any>,
+): MutableAIState<InferAIState<AI, any>[typeof key]>;
+function getMutableAIState<AI extends AIProvider = any>(
+  ...args: [] | [key: keyof InferAIState<AI, any>]
+) {
+  type AIState = InferAIState<AI, any>;
+  type AIStateWithKey = typeof args extends [key: keyof AIState]
+    ? AIState[(typeof args)[0]]
+    : AIState;
+  type NewStateOrUpdater = ValueOrUpdater<AIStateWithKey>;
+
+  const store = getAIStateStoreOrThrow(
+    '`getMutableAIState` must be called within an AI Action.',
+  );
+
+  if (store.sealed) {
+    throw new Error(
+      "`getMutableAIState` must be called before returning from an AI Action. Please move it to the top level of the Action's function body.",
+    );
+  }
+
+  if (!store.mutationDeltaPromise) {
+    const { promise, resolve } = createResolvablePromise();
+    store.mutationDeltaPromise = promise;
+    store.mutationDeltaResolve = resolve;
+  }
+
+  function doUpdate(newState: NewStateOrUpdater, done: boolean) {
+    if (args.length > 0) {
+      if (typeof store.currentState !== 'object') {
+        const key = args[0];
+        throw new Error(
+          `You can't modify the "${String(
+            key,
+          )}" field of the AI state because it's not an object.`,
+        );
+      }
+    }
+
+    if (isFunction(newState)) {
+      if (args.length > 0) {
+        store.currentState[args[0]] = newState(store.currentState[args[0]]);
+      } else {
+        store.currentState = newState(store.currentState);
+      }
+    } else {
+      if (args.length > 0) {
+        store.currentState[args[0]] = newState;
+      } else {
+        store.currentState = newState;
+      }
+    }
+
+    store.options.onSetAIState?.({
+      key: args.length > 0 ? args[0] : undefined,
+      state: store.currentState,
+      done,
+    });
+  }
+
+  const mutableState = {
+    get: () => {
+      if (args.length > 0) {
+        const key = args[0];
+        if (typeof store.currentState !== 'object') {
+          throw new Error(
+            `You can't get the "${String(
+              key,
+            )}" field from the AI state because it's not an object.`,
+          );
+        }
+        return store.currentState[key];
+      }
+
+      return store.currentState as AIState;
+    },
+    update: function update(newAIState: NewStateOrUpdater) {
+      doUpdate(newAIState, false);
+    },
+    done: function done(...doneArgs: [] | [NewStateOrUpdater]) {
+      if (doneArgs.length > 0) {
+        doUpdate(doneArgs[0] as NewStateOrUpdater, true);
+      }
+
+      const delta = jsondiffpatch.diff(store.originalState, store.currentState);
+      store.mutationDeltaResolve!(delta);
+    },
+  };
+
+  return mutableState;
+}
+
+export { getAIState, getMutableAIState };

packages/core/rsc/constants.ts

@@ -0,0 +1 @@
+export const STREAMABLE_VALUE_TYPE = Symbol.for('ui.streamable.value');

packages/core/rsc/package.json

@@ -0,0 +1,14 @@
+{
+  "types": "./dist/rsc-types.d.ts",
+  "exports": {
+    "types": "./dist/rsc-types.d.ts",
+    "react-server": "./dist/rsc-server.mjs",
+    "import": "./dist/rsc-client.mjs",
+    "module": "./dist/rsc-client.mjs"
+  },
+  "private": true,
+  "peerDependencies": {
+    "react": ">=18",
+    "zod": ">=3"
+  }
+}