fix: 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(
+  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,28 @@
+'use server';
+import { InternalProvider } from './shared.js';
+import { jsx } from 'react/jsx-runtime';
+
+async function innerAction({ action, options }, state, ...args) {
+  'use server';
+  const result = await action(...args);
+  // eslint-disable-next-line no-undef
+  console.log('wrapped action', result);
+  return result;
+}
+
+function wrapAction(action, options) {
+  return innerAction.bind(null, { action, options });
+}
+
+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.0",
     "react-dom": "19.0.0-rc.0",
     "react-server-dom-webpack": "19.0.0-rc.0",

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,14 @@
+'use client';
+
+import { useActions } from 'ai/rsc';
+import { useEffect } from 'react';
+
+export const ClientActionsConsumer = () => {
+  const actions = useActions();
+  useEffect(() => {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-expect-error
+    globalThis.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,19 @@ 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(() => {
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-expect-error
+        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') {