@ai-sdk/rsc 2.0.44 → 2.0.46

package/CHANGELOG.md

@@ -1,5 +1,18 @@
 # @ai-sdk/rsc
 
+## 2.0.46
+
+### Patch Changes
+
+- 8dc54db: chore: add src folders to package bundle
+  - ai@6.0.46
+
+## 2.0.45
+
+### Patch Changes
+
+- ai@6.0.45
+
 ## 2.0.44
 
 ### Patch Changes

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ai-sdk/rsc",
-  "version": "2.0.44",
+  "version": "2.0.46",
   "license": "Apache-2.0",
   "sideEffects": false,
   "main": "./dist/rsc-client.mjs",
@@ -17,12 +17,13 @@
   },
   "files": [
     "dist/**/*",
+    "src",
     "CHANGELOG.md",
     "README.md"
   ],
   "dependencies": {
     "jsondiffpatch": "0.7.3",
-    "ai": "6.0.44",
+    "ai": "6.0.46",
     "@ai-sdk/provider": "3.0.4",
     "@ai-sdk/provider-utils": "4.0.8"
   },

package/src/ai-state.test.ts

@@ -0,0 +1,146 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import {
+  withAIState,
+  getAIState,
+  getMutableAIState,
+  sealMutableAIState,
+  getAIStateDeltaPromise,
+} from './ai-state';
+
+describe('AI State Management', () => {
+  beforeEach(() => {
+    vi.resetAllMocks();
+  });
+
+  it('should get the current AI state', () => {
+    const initialState = { foo: 'bar' };
+    const result = withAIState({ state: initialState, options: {} }, () => {
+      return getAIState();
+    });
+    expect(result).toEqual(initialState);
+  });
+
+  it('should get a specific key from the AI state', () => {
+    const initialState = { foo: 'bar', baz: 'qux' };
+    const result = withAIState({ state: initialState, options: {} }, () => {
+      return getAIState('foo');
+    });
+    expect(result).toBe('bar');
+  });
+
+  it('should update the AI state', () => {
+    const initialState = { foo: 'bar' };
+    withAIState({ state: initialState, options: {} }, () => {
+      const mutableState = getMutableAIState();
+      mutableState.update({ foo: 'baz' });
+      expect(getAIState()).toEqual({ foo: 'baz' });
+    });
+  });
+
+  it('should update a specific key in the AI state', () => {
+    const initialState = { foo: 'bar', baz: 'qux' };
+    withAIState({ state: initialState, options: {} }, () => {
+      const mutableState = getMutableAIState('foo');
+      mutableState.update('newValue');
+      expect(getAIState()).toEqual({ foo: 'newValue', baz: 'qux' });
+    });
+  });
+
+  it('should throw an error when accessing AI state outside of withAIState', () => {
+    expect(() => getAIState()).toThrow(
+      '`getAIState` must be called within an AI Action.',
+    );
+  });
+
+  it('should throw an error when updating AI state after it has been sealed', () => {
+    withAIState({ state: {}, options: {} }, () => {
+      sealMutableAIState();
+      expect(() => getMutableAIState()).toThrow(
+        '`getMutableAIState` must be called before returning from an AI Action.',
+      );
+    });
+  });
+
+  it('should call onSetAIState when updating state', () => {
+    const onSetAIState = vi.fn();
+    const initialState = { foo: 'bar' };
+    withAIState({ state: initialState, options: { onSetAIState } }, () => {
+      const mutableState = getMutableAIState();
+      mutableState.update({ foo: 'baz' });
+      mutableState.done({ foo: 'baz' });
+    });
+    expect(onSetAIState).toHaveBeenCalledWith(
+      expect.objectContaining({
+        state: { foo: 'baz' },
+        done: true,
+      }),
+    );
+  });
+
+  it('should handle updates with and without key', async () => {
+    type Message = { role: string; content: string };
+
+    type AIState = {
+      chatId: string;
+      messages: Array<Message>;
+    };
+
+    const initialState: AIState = {
+      chatId: '123',
+      messages: [],
+    };
+
+    await withAIState({ state: initialState, options: {} }, async () => {
+      // Test with getMutableState()
+      const stateWithoutKey = getMutableAIState();
+
+      stateWithoutKey.update((current: AIState) => ({
+        ...current,
+        messages: [...current.messages, { role: 'user', content: 'Hello!' }],
+      }));
+
+      stateWithoutKey.done((current: AIState) => ({
+        ...current,
+        messages: [
+          ...current.messages,
+          { role: 'assistant', content: 'Hello! How can I assist you today?' },
+        ],
+      }));
+
+      const deltaWithoutKey = await getAIStateDeltaPromise();
+      expect(deltaWithoutKey).toBeDefined();
+      expect(getAIState()).toEqual({
+        chatId: '123',
+        messages: [
+          { role: 'user', content: 'Hello!' },
+          { role: 'assistant', content: 'Hello! How can I assist you today?' },
+        ],
+      });
+    });
+
+    await withAIState({ state: initialState, options: {} }, async () => {
+      // Test with getMutableState('messages')
+      const stateWithKey = getMutableAIState('messages');
+
+      stateWithKey.update((current: Array<Message>) => [
+        ...current,
+        { role: 'user', content: 'Hello!' },
+      ]);
+
+      stateWithKey.done((current: Array<Message>) => [
+        ...current,
+        { role: 'assistant', content: 'Hello! How can I assist you today?' },
+      ]);
+
+      const deltaWithKey = await getAIStateDeltaPromise();
+      expect(deltaWithKey).toBeDefined();
+      expect(getAIState()).toEqual({
+        chatId: '123',
+        messages: [
+          { role: 'user', content: 'Hello!' },
+          { role: 'assistant', content: 'Hello! How can I assist you today?' },
+        ],
+      });
+    });
+  });
+});

package/src/ai-state.tsx

@@ -0,0 +1,210 @@
+import * as jsondiffpatch from 'jsondiffpatch';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { createResolvablePromise } from './util/create-resolvable-promise';
+import { isFunction } from './util/is-function';
+import type {
+  AIProvider,
+  InferAIState,
+  InternalAIStateStorageOptions,
+  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: JSON.parse(JSON.stringify(state)), // deep clone object
+      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>(): Readonly<
+  InferAIState<AI, any>
+>;
+function getAIState<AI extends AIProvider = any>(
+  key: keyof InferAIState<AI, any>,
+): Readonly<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 `.done()` 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] as Readonly<AIStateWithKey>;
+      }
+
+      return store.currentState as Readonly<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 };

package/src/index.ts

@@ -0,0 +1,20 @@
+export {
+  getAIState,
+  getMutableAIState,
+  createStreamableUI,
+  createStreamableValue,
+  streamUI,
+  createAI,
+} from './rsc-server';
+
+export {
+  readStreamableValue,
+  useStreamableValue,
+  useUIState,
+  useAIState,
+  useActions,
+  useSyncUIState,
+} from './rsc-client';
+
+export type { StreamableValue } from './streamable-value/streamable-value';
+export * from './types';

package/src/provider.tsx

@@ -0,0 +1,149 @@
+// This file provides the AI context to all AI Actions via AsyncLocalStorage.
+
+import * as React from 'react';
+import { InternalAIProvider } from './rsc-shared.mjs';
+import {
+  withAIState,
+  getAIStateDeltaPromise,
+  sealMutableAIState,
+} from './ai-state';
+import type {
+  ServerWrappedActions,
+  AIAction,
+  AIActions,
+  AIProvider,
+  InternalAIStateStorageOptions,
+  OnSetAIState,
+  OnGetUIState,
+} from './types';
+
+async function innerAction<T>(
+  {
+    action,
+    options,
+  }: { action: AIAction; options: InternalAIStateStorageOptions },
+  state: T,
+  ...args: unknown[]
+) {
+  'use server';
+  return await withAIState(
+    {
+      state,
+      options,
+    },
+    async () => {
+      const result = await action(...args);
+      sealMutableAIState();
+      return [getAIStateDeltaPromise() as Promise<T>, result];
+    },
+  );
+}
+
+function wrapAction<T = unknown>(
+  action: AIAction,
+  options: InternalAIStateStorageOptions,
+) {
+  return innerAction.bind(null, { action, options }) as AIAction<T>;
+}
+
+export function createAI<
+  AIState = any,
+  UIState = any,
+  Actions extends AIActions = {},
+>({
+  actions,
+  initialAIState,
+  initialUIState,
+
+  onSetAIState,
+  onGetUIState,
+}: {
+  actions: Actions;
+  initialAIState?: AIState;
+  initialUIState?: UIState;
+
+  /**
+   * This function is called whenever the AI state is updated by an Action.
+   * You can use this to persist the AI state to a database, or to send it to a
+   * logging service.
+   */
+  onSetAIState?: OnSetAIState<AIState>;
+
+  /**
+   * This function is used to retrieve the UI state based on the AI state.
+   * For example, to render the initial UI state based on a given AI state, or
+   * to sync the UI state when the application is already loaded.
+   *
+   * If returning `undefined`, the client side UI state will not be updated.
+   *
+   * This function must be annotated with the `"use server"` directive.
+   *
+   * @example
+   * ```tsx
+   * onGetUIState: async () => {
+   *   'use server';
+   *
+   *   const currentAIState = getAIState();
+   *   const externalAIState = await loadAIStateFromDatabase();
+   *
+   *   if (currentAIState === externalAIState) return undefined;
+   *
+   *   // Update current AI state and return the new UI state
+   *   const state = getMutableAIState()
+   *   state.done(externalAIState)
+   *
+   *   return <div>...</div>;
+   * }
+   * ```
+   */
+  onGetUIState?: OnGetUIState<UIState>;
+}) {
+  // Wrap all actions with our HoC.
+  const wrappedActions: ServerWrappedActions = {};
+  for (const name in actions) {
+    wrappedActions[name] = wrapAction(actions[name], {
+      onSetAIState,
+    });
+  }
+
+  const wrappedSyncUIState = onGetUIState
+    ? wrapAction(onGetUIState, {})
+    : undefined;
+
+  const AI: AIProvider<AIState, UIState, Actions> = async props => {
+    if ('useState' in React) {
+      // This file must be running on the React Server layer.
+      // Ideally we should be using `import "server-only"` here but we can have a
+      // more customized error message with this implementation.
+      throw new Error(
+        'This component can only be used inside Server Components.',
+      );
+    }
+
+    let uiState = props.initialUIState ?? initialUIState;
+    let aiState = props.initialAIState ?? initialAIState;
+    let aiStateDelta = undefined;
+
+    if (wrappedSyncUIState) {
+      const [newAIStateDelta, newUIState] = await wrappedSyncUIState(aiState);
+      if (newUIState !== undefined) {
+        aiStateDelta = newAIStateDelta;
+        uiState = newUIState;
+      }
+    }
+
+    return (
+      <InternalAIProvider
+        wrappedActions={wrappedActions}
+        wrappedSyncUIState={wrappedSyncUIState}
+        initialUIState={uiState}
+        initialAIState={aiState}
+        initialAIStatePatch={aiStateDelta}
+      >
+        {props.children}
+      </InternalAIProvider>
+    );
+  };
+
+  return AI;
+}

package/src/rsc-client.ts

@@ -0,0 +1,8 @@
+export {
+  readStreamableValue,
+  useStreamableValue,
+  useUIState,
+  useAIState,
+  useActions,
+  useSyncUIState,
+} from './rsc-shared.mjs';

package/src/rsc-server.ts

@@ -0,0 +1,5 @@
+export { getAIState, getMutableAIState } from './ai-state';
+export { createAI } from './provider';
+export { streamUI } from './stream-ui';
+export { createStreamableUI } from './streamable-ui/create-streamable-ui';
+export { createStreamableValue } from './streamable-value/create-streamable-value';

package/src/rsc-shared.mts

@@ -0,0 +1,11 @@
+'use client';
+
+export {
+  readStreamableValue,
+  useStreamableValue,
+  useUIState,
+  useAIState,
+  useActions,
+  useSyncUIState,
+  InternalAIProvider,
+} from './shared-client';