test: edge cases for server action

e2e/fixtures/rsc-basic/modules/ai/package.json

@@ -0,0 +1,21 @@
+{
+  "name": "ai",
+  "type": "module",
+  "version": "1.0.0",
+  "description": "Vercel AI mockup",
+  "exports": {
+    "./rsc": {
+      "types": "./src/index.d.ts",
+      "react-server": "./src/server.js",
+      "import": "./src/client.js"
+    }
+  },
+  "devDependencies": {
+    "react-dom": "^18",
+    "react-server-dom-webpack": "18.3.0-canary-eb33bd747-20240312"
+  },
+  "peerDependencies": {
+    "react": "^18 || ^19"
+  },
+  "private": true
+}

e2e/fixtures/rsc-basic/modules/ai/src/client.js

@@ -0,0 +1,8 @@
+'use client';
+import { useActions } from './shared.js';
+
+export { useActions };
+
+export function createAI() {
+  throw new Error('You should not call createAI in the client side');
+}

e2e/fixtures/rsc-basic/modules/ai/src/index.d.ts

@@ -0,0 +1,7 @@
+import type { ReactNode } from 'react';
+
+declare function createAI(options: {
+  actions: Record<string, any>;
+}): (props: { children: ReactNode }) => ReactNode;
+
+declare function useActions(): Record<string, any>;

e2e/fixtures/rsc-basic/modules/ai/src/server.js

@@ -0,0 +1,25 @@
+'use server';
+import { InternalProvider } from './shared.js';
+import { jsx } from 'react/jsx-runtime';
+
+async function innerAction({ action }, ...args) {
+  'use server';
+  return await action(...args);
+}
+
+function wrapAction(action) {
+  return innerAction.bind(null, { action });
+}
+
+export function createAI({ actions }) {
+  const wrappedActions = {};
+  for (const name in actions) {
+    wrappedActions[name] = wrapAction(actions[name]);
+  }
+  return function AI(props) {
+    return jsx(InternalProvider, {
+      actions: wrappedActions,
+      children: props.children,
+    });
+  };
+}

e2e/fixtures/rsc-basic/modules/ai/src/shared.js

@@ -0,0 +1,19 @@
+'use client';
+import { createContext, useContext } from 'react';
+import { jsx } from 'react/jsx-runtime';
+
+const ActionContext = createContext(null);
+
+export function useActions() {
+  return useContext(ActionContext);
+}
+
+export function InternalProvider(props) {
+  return jsx('div', {
+    'data-testid': 'ai-internal-provider',
+    children: jsx(ActionContext.Provider, {
+      value: props.actions,
+      children: props.children,
+    }),
+  });
+}

e2e/fixtures/rsc-basic/package.json

@@ -9,6 +9,7 @@
     "start": "waku start"
   },
   "dependencies": {
+    "ai": "link:./modules/ai",
     "react": "19.0.0-rc-3da26163a3-20240704",
     "react-dom": "19.0.0-rc-3da26163a3-20240704",
     "react-server-dom-webpack": "19.0.0-rc-3da26163a3-20240704",

e2e/fixtures/rsc-basic/src/components/App.tsx

@@ -4,6 +4,8 @@
 import { ClientCounter } from './ClientCounter.js';
 import { ServerPing } from './ServerPing/index.js';
 import { ServerBox } from './Box.js';
+import { ServerProvider } from './ServerAction/Server.js';
+import { ClientActionsConsumer } from './ServerAction/Client.js';
 
 const App = ({ name }: { name: string }) => {
   return (
@@ -12,6 +14,9 @@ const App = ({ name }: { name: string }) => {
       <p data-testid="app-name">{name}</p>
       <ClientCounter />
       <ServerPing />
+      <ServerProvider>
+        <ClientActionsConsumer />
+      </ServerProvider>
     </ServerBox>
   );
 };

e2e/fixtures/rsc-basic/src/components/ServerAction/Client.tsx

@@ -0,0 +1,12 @@
+'use client';
+
+import { useActions } from 'ai/rsc';
+import { useEffect } from 'react';
+
+export const ClientActionsConsumer = () => {
+  const actions = useActions();
+  useEffect(() => {
+    (globalThis as any).actions = actions;
+  }, [actions]);
+  return <div>{JSON.stringify(Object.keys(actions))}</div>;
+};

e2e/fixtures/rsc-basic/src/components/ServerAction/Server.tsx

@@ -0,0 +1,15 @@
+import type { ReactNode } from 'react';
+import { createAI } from 'ai/rsc';
+
+const AI = createAI({
+  actions: {
+    foo: async () => {
+      'use server';
+      return 0;
+    },
+  },
+});
+
+export function ServerProvider({ children }: { children: ReactNode }) {
+  return <AI>{children}</AI>;
+}

e2e/rsc-basic.spec.ts

@@ -81,5 +81,18 @@ for (const { build, command } of commands) {
         page.getByTestId('server-ping').getByTestId('counter'),
       ).toHaveText('2');
     });
+
+    test('server action', async ({ page }) => {
+      await page.goto(`http://localhost:${port}/`);
+      await expect(page.getByTestId('app-name')).toHaveText('Waku');
+      await expect(page.getByTestId('ai-internal-provider')).toHaveText(
+        '["foo"]',
+      );
+      const result = await page.evaluate(() => {
+        // @ts-expect-error no types
+        return globalThis.actions.foo();
+      });
+      expect(result).toBe(0);
+    });
   });
 }

packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts

@@ -8,7 +8,18 @@ import { parseOpts } from '../utils/swc.js';
 const collectExportNames = (mod: swc.Module) => {
   const exportNames = new Set<string>();
   for (const item of mod.body) {
-    if (item.type === 'ExportDeclaration') {
+    if (item.type === 'FunctionDeclaration') {
+      const rscDeclaration = item.body?.stmts[0];
+      if (
+        rscDeclaration?.type === 'ExpressionStatement' &&
+        rscDeclaration.expression.type === 'StringLiteral' &&
+        rscDeclaration.expression.value === 'use server'
+      ) {
+        exportNames.add(item.identifier.value);
+      }
+    }
+    // fixme: this might be incorrect, not all exports from server file can be registered as server references
+    else if (item.type === 'ExportDeclaration') {
       if (item.declaration.type === 'FunctionDeclaration') {
         exportNames.add(item.declaration.identifier.value);
       } else if (item.declaration.type === 'VariableDeclaration') {

packages/waku/src/lib/renderers/rsc-renderer.ts

@@ -12,6 +12,16 @@
 import { streamToString } from '../utils/stream.js';
 import { decodeActionId } from '../renderers/utils.js';
 
+// HACK for react-server-dom-webpack without webpack
+(globalThis as any).__webpack_module_loading__ ||= new Map();
+(globalThis as any).__webpack_module_cache__ ||= new Map();
+(globalThis as any).__webpack_chunk_load__ ||= async (id: string) =>
+  (globalThis as any).__webpack_module_loading__.get(id);
+(globalThis as any).__webpack_require__ ||= (id: string) =>
+  (globalThis as any).__webpack_module_cache__.get(id);
+const moduleLoading = (globalThis as any).__webpack_module_loading__;
+const moduleCache = (globalThis as any).__webpack_module_cache__;
+
 export const SERVER_MODULE_MAP = {
   'rsdw-server': 'react-server-dom-webpack/server.edge',
   'waku-server': 'waku/server',
@@ -189,7 +199,46 @@
     ) {
       // XXX This doesn't support streaming unlike busboy
       const formData = parseFormData(bodyStr, contentType);
-      args = await decodeReply(formData);
+      const moduleMap = new Proxy({} as Record<string, ImportManifestEntry>, {
+        get(_target, rsfId: string): ImportManifestEntry {
+          const [fileId, name] = rsfId.split('#') as [string, string];
+          // fixme: race condition, server actions are not initialized in the first time
+          if (!moduleLoading.has(fileId)) {
+            if (!opts.isDev) {
+              moduleLoading.set(
+                fileId,
+                import(
+                  /* @vite-ignore */
+                  fileId
+                ).then((m: any) => {
+                  moduleCache.set(
+                    fileId,
+                    Object.fromEntries(m['__waku_serverActions']),
+                  );
+                }),
+              );
+            } else {
+              moduleLoading.set(
+                fileId,
+                opts
+                  .loadServerFile(filePathToFileURL(fileId))
+                  .then((m: any) => {
+                    moduleCache.set(
+                      fileId,
+                      Object.fromEntries(m['__waku_serverActions']),
+                    );
+                  }),
+              );
+            }
+          }
+          return {
+            id: fileId,
+            chunks: [],
+            name,
+          };
+        },
+      });
+      args = await decodeReply(formData, moduleMap);
     } else if (bodyStr) {
       args = await decodeReply(bodyStr);
     }