Repository: dai-shi/waku
Pull request #774 (2024-07-03 08:39:54)

fix: server action cases
Fix some RSC case, will fix https://github.com/dai-shi/waku/pull/745. In the future, I will migrate the logic to https://github.com/himself65/react-server and keep same output

Some issues still happening:

1. `export default async function () {"use server"}` cannot be RSCA
2. `export function a() { "use server" }` should throw error, because it's not async

vercel[bot]: [vc]: #2wo/w0BL2YdEeSPGweeK5WuBfm6bNU0SoxFS6r14aRQ=:eyJpc01vbm9yZXBvIjp0cnVlLCJ0eXBlIjoiZ2l0aHViIiwicHJvamVjdHMiOlt7Im5hbWUiOiJ3YWt1Iiwicm9vdERpcmVjdG9yeSI6bnVsbCwiaW5zcGVjdG9yVXJsIjoiaHR0cHM6Ly92ZXJjZWwuY29tL2RhaXNoaXMtcHJvamVjdHMvd2FrdS9FbVVVRGFyU2I3THluWnJWWnlUYzVTS0tFTXVQIiwicHJldmlld1VybCI6Indha3UtZ2l0LWZvcmstaGltc2VsZjY1LWhpbXNlbGY2NS0yMDI0MDcwM2JpbmQtZGFpc2hpcy1wcm9qZWN0cy52ZXJjZWwuYXBwIiwibmV4dENvbW1pdFN0YXR1cyI6IlBFTkRJTkcifV19
**The latest updates on your projects**. Learn more about [Vercel for Git ↗︎](https://vercel.link/github-learn-more)

| Name | Status | Preview | Updated (UTC) |
| :--- | :----- | :------ | :------ |
| **waku** | 🔄 Building ([Inspect](https://vercel.com/daishis-projects/waku/EmUUDarSb7LynZrVZyTc5SKKEMuP)) | [Visit Preview](https://waku-git-fork-himself65-himself65-20240703bind-daishis-projects.vercel.app) | Jul 3, 2024 8:39am |


codesandbox-ci[bot]: This pull request is automatically built and testable in [CodeSandbox](https://codesandbox.io).

  To see build info of the built libraries, click [here](https://ci.codesandbox.io/status/dai-shi/waku/pr/774/builds/519436) or the icon next to each commit SHA.
dai-shi: > 1. `export default async function () {"use server"}` cannot be RSCA

Yeah, I notice that too. Even with top-level "use server", it fails. Gotta fix it, but maybe later.

> 2. `export function a() { "use server" }` should throw error, because it's not async

Originally, I thought sync functions are supported, and I have such an example in 23_actions.
Does nextjs throw an error? I'll leave it to your `react-server`.

Patch:
diff --git a/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts b/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts
index 79969145..f4487e70 100644
--- a/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts
+++ b/packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts
@@ -91,10 +91,18 @@ const createCallExpression = (
   span: { start: 0, end: 0, ctxt: 0 },
 });
 
-const serverActionsInitCode = swc.parseSync(`
+const serverInitCode = swc.parseSync(`
 import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
 `).body;
 
+const findLastImportIndex = (mod: swc.Module) => {
+  const lastImportIndex = mod.body.findIndex(
+    (node) =>
+      node.type !== 'ExpressionStatement' && node.type !== 'ImportDeclaration',
+  );
+  return lastImportIndex === -1 ? 0 : lastImportIndex;
+};
+
 type FunctionWithBlockBody = (
   | swc.FunctionDeclaration
   | swc.FunctionExpression
@@ -323,62 +331,59 @@ const transformServerActions = (
   if (!serverActionIndex) {
     return;
   }
-  const serverActionsCode = [...serverActionsInitCode];
-  for (const [actionIndex, [actionFn, closureVars]] of serverActions) {
-    if (actionFn.type === 'FunctionDeclaration') {
-      const stmt1: swc.ExportDeclaration = {
-        type: 'ExportDeclaration',
-        declaration: prependArgsToFn(actionFn, closureVars),
-        span: { start: 0, end: 0, ctxt: 0 },
-      };
-      const stmt2: swc.ExpressionStatement = {
-        type: 'ExpressionStatement',
-        expression: createCallExpression(
-          createIdentifier('__waku_registerServerReference'),
-          [
-            createIdentifier(actionFn.identifier.value),
-            createStringLiteral(getActionId()),
-            createStringLiteral('__waku_action' + actionIndex),
-          ],
-        ),
-        span: { start: 0, end: 0, ctxt: 0 },
-      };
-      serverActionsCode.push(stmt1, stmt2);
-    } else {
-      const stmt: swc.ExportDeclaration = {
-        type: 'ExportDeclaration',
-        declaration: {
-          type: 'VariableDeclaration',
-          kind: 'const',
-          declare: false,
-          declarations: [
-            {
-              type: 'VariableDeclarator',
-              id: createIdentifier('__waku_action' + actionIndex),
-              init: createCallExpression(
-                createIdentifier('__waku_registerServerReference'),
-                [
-                  prependArgsToFn(actionFn, closureVars),
-                  createStringLiteral(getActionId()),
-                  createStringLiteral('__waku_action' + actionIndex),
-                ],
-              ),
-              definite: false,
-              span: { start: 0, end: 0, ctxt: 0 },
-            },
-          ],
+  const serverActionsCode = Array.from(serverActions).flatMap(
+    ([actionIndex, [actionFn, closureVars]]) => {
+      if (actionFn.type === 'FunctionDeclaration') {
+        const stmt1: swc.ExportDeclaration = {
+          type: 'ExportDeclaration',
+          declaration: prependArgsToFn(actionFn, closureVars),
           span: { start: 0, end: 0, ctxt: 0 },
-        },
-        span: { start: 0, end: 0, ctxt: 0 },
-      };
-      serverActionsCode.push(stmt);
-    }
-  }
-  const lastImportIndex = mod.body.findIndex(
-    (node) =>
-      node.type !== 'ExpressionStatement' && node.type !== 'ImportDeclaration',
+        };
+        const stmt2: swc.ExpressionStatement = {
+          type: 'ExpressionStatement',
+          expression: createCallExpression(
+            createIdentifier('__waku_registerServerReference'),
+            [
+              createIdentifier(actionFn.identifier.value),
+              createStringLiteral(getActionId()),
+              createStringLiteral('__waku_action' + actionIndex),
+            ],
+          ),
+          span: { start: 0, end: 0, ctxt: 0 },
+        };
+        return [stmt1, stmt2];
+      } else {
+        const stmt: swc.ExportDeclaration = {
+          type: 'ExportDeclaration',
+          declaration: {
+            type: 'VariableDeclaration',
+            kind: 'const',
+            declare: false,
+            declarations: [
+              {
+                type: 'VariableDeclarator',
+                id: createIdentifier('__waku_action' + actionIndex),
+                init: createCallExpression(
+                  createIdentifier('__waku_registerServerReference'),
+                  [
+                    prependArgsToFn(actionFn, closureVars),
+                    createStringLiteral(getActionId()),
+                    createStringLiteral('__waku_action' + actionIndex),
+                  ],
+                ),
+                definite: false,
+                span: { start: 0, end: 0, ctxt: 0 },
+              },
+            ],
+            span: { start: 0, end: 0, ctxt: 0 },
+          },
+          span: { start: 0, end: 0, ctxt: 0 },
+        };
+        return [stmt];
+      }
+    },
   );
-  mod.body.splice(lastImportIndex, 0, ...serverActionsCode);
+  mod.body.splice(findLastImportIndex(mod), 0, ...serverActionsCode);
   return mod;
 };
 
@@ -391,14 +396,18 @@ const transformServer = (
   const ext = extname(id);
   const mod = swc.parseSync(code, parseOpts(ext));
   let hasUseClient = false;
-  let hasUseServer: swc.ExpressionStatement | undefined;
-  for (const item of mod.body) {
+  let hasUseServer = false;
+  for (let i = 0; i < mod.body.length; ++i) {
+    const item = mod.body[i]!;
     if (item.type === 'ExpressionStatement') {
       if (item.expression.type === 'StringLiteral') {
         if (item.expression.value === 'use client') {
           hasUseClient = true;
+          break;
         } else if (item.expression.value === 'use server') {
-          hasUseServer = item;
+          hasUseServer = true;
+          mod.body.splice(i, 1); // remove this directive
+          break;
         }
       }
     } else {
@@ -417,42 +426,56 @@ export ${name === 'default' ? name : `const ${name} =`} registerClientReference(
 `;
     }
     return newCode;
-  } else if (hasUseServer) {
-    const useServerPosStart = hasUseServer.span.start - mod.span.start;
-    const useServerPosEnd = hasUseServer.span.end - mod.span.start;
-    const lastImportIndex = mod.body.findIndex(
-      (node) =>
-        node.type !== 'ExpressionStatement' &&
-        node.type !== 'ImportDeclaration',
-    );
-    let lastImportPos =
-      lastImportIndex === -1 ? 0 : mod.body[lastImportIndex]!.span.end;
-    if (lastImportIndex < useServerPosEnd) {
-      lastImportPos = useServerPosEnd;
-    }
+  }
+  if (hasUseServer) {
     const exportNames = collectExportNames(mod);
-    const newCode = [
-      code.slice(0, useServerPosStart),
-      code.slice(useServerPosEnd, lastImportPos),
-      `
-import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
-`,
-      code.slice(lastImportPos),
-      ...[...exportNames].map(
-        (name) => `
-if (typeof ${name} === 'function') {
-  __waku_registerServerReference(${name}, '${getServerId(id)}', '${name}');
-}
-`,
-      ),
-    ].join('');
-    return newCode;
+    const serverActionsCode = Array.from(exportNames).map((name) => {
+      const blockStmt: swc.BlockStatement = {
+        type: 'BlockStatement',
+        stmts: [
+          {
+            type: 'ExpressionStatement',
+            expression: createCallExpression(
+              createIdentifier('__waku_registerServerReference'),
+              [
+                createIdentifier(name),
+                createStringLiteral(getServerId(id)),
+                createStringLiteral(name),
+              ],
+            ),
+            span: { start: 0, end: 0, ctxt: 0 },
+          },
+        ],
+        span: { start: 0, end: 0, ctxt: 0 },
+      };
+      const ifStmt: swc.IfStatement = {
+        type: 'IfStatement',
+        test: {
+          type: 'BinaryExpression',
+          operator: '===',
+          left: {
+            type: 'UnaryExpression',
+            operator: 'typeof',
+            argument: createIdentifier(name),
+            span: { start: 0, end: 0, ctxt: 0 },
+          },
+          right: createStringLiteral('function'),
+          span: { start: 0, end: 0, ctxt: 0 },
+        },
+        consequent: blockStmt,
+        span: { start: 0, end: 0, ctxt: 0 },
+      };
+      return ifStmt;
+    });
+    mod.body.push(...serverActionsCode);
   }
   // transform server actions in server components
   const newMod =
-    code.includes('use server') &&
-    transformServerActions(mod, () => getServerId(id));
+    (code.includes('use server') &&
+      transformServerActions(mod, () => getServerId(id))) ||
+    (hasUseServer && mod);
   if (newMod) {
+    newMod.body.splice(findLastImportIndex(newMod), 0, ...serverInitCode);
     const newCode = swc.printSync(newMod).code;
     return newCode;
   }


Test patch:
diff --git a/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts b/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts
index a480f487..b94752f7 100644
--- a/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts
+++ b/packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts
@@ -53,16 +53,113 @@ export const log = (mesg) => {
 `;
     expect(await transform(code, '/src/App.tsx', { ssr: true }))
       .toMatchInlineSnapshot(`
+        "import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
+        export const log = (mesg)=>{
+            console.log(mesg);
+        };
+        if (typeof log === "function") {
+            __waku_registerServerReference(log, "/src/App.tsx", "log");
+        }
         "
+      `);
+  });
+
+  test('server action in object', async () => {
+    const code = `
+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>;
+}
+`;
+    expect(await transform(code, '/src/App.tsx', { ssr: true }))
+      .toMatchInlineSnapshot(`
+        "import type { ReactNode } from 'react';
+        import { createAI } from 'ai/rsc';
         import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
-        ;
+        export const __waku_action1 = __waku_registerServerReference(async ()=>{
+            return 0;
+        }, "/src/App.tsx", "__waku_action1");
+        const AI = createAI({
+            actions: {
+                foo: __waku_action1.bind(null)
+            }
+        });
+        export function ServerProvider({ children }: {
+            children: ReactNode;
+        }) {
+            return <AI>{children}</AI>;
+        }
+        "
+      `);
+  });
 
-        export const log = (mesg) => {
-          console.log(mesg);
-        };
+  test('create server with bind', async () => {
+    const code = `
+'use server';
+import { InternalProvider } from './shared.js';
+import { jsx } from 'react/jsx-runtime';
 
-        if (typeof log === 'function') {
-          __waku_registerServerReference(log, '/src/App.tsx', 'log');
+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,
+    });
+  };
+}`;
+    expect(await transform(code, '/src/App.tsx', { ssr: true }))
+      .toMatchInlineSnapshot(`
+        "import { InternalProvider } from './shared.js';
+        import { jsx } from 'react/jsx-runtime';
+        import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
+        export async function __waku_action1({ action }, ...args) {
+            return await action(...args);
+        }
+        __waku_registerServerReference(__waku_action1, "/src/App.tsx", "__waku_action1");
+        const innerAction = __waku_action1.bind(null);
+        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
+                });
+            };
+        }
+        if (typeof createAI === "function") {
+            __waku_registerServerReference(createAI, "/src/App.tsx", "createAI");
         }
         "
       `);
@@ -71,7 +168,7 @@ export const log = (mesg) => {
   test('inline use server (function declaration)', async () => {
     const code = `
 export default function App({ a }) {
-  function log(mesg) {
+  async function log(mesg) {
     'use server';
     console.log(mesg, a);
   }
@@ -80,17 +177,17 @@ export default function App({ a }) {
 `;
     expect(await transform(code, '/src/App.tsx', { ssr: true }))
       .toMatchInlineSnapshot(`
-        "import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
-        export function __waku_action1(a, mesg) {
-            console.log(mesg, a);
-        }
-        __waku_registerServerReference(__waku_action1, "/src/App.tsx", "__waku_action1");
-        export default function App({ a }) {
-            const log = __waku_action1.bind(null, a);
-            return <Hello log={log}/>;
-        }
-        "
-      `);
+      "import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
+      export async function __waku_action1(a, mesg) {
+          console.log(mesg, a);
+      }
+      __waku_registerServerReference(__waku_action1, "/src/App.tsx", "__waku_action1");
+      export default function App({ a }) {
+          const log = __waku_action1.bind(null, a);
+          return <Hello log={log}/>;
+      }
+      "
+    `);
   });
 
   test('inline use server (const function expression)', async () => {
@@ -148,7 +245,7 @@ export default function App() {
   test('inline use server (in an object)', async () => {
     const code = `
 const actions = {
-  log: (mesg) => {
+  log: async (mesg) => {
     'use server';
     console.log(mesg);
   },
@@ -160,7 +257,7 @@ export default function App() {
     expect(await transform(code, '/src/App.tsx', { ssr: true }))
       .toMatchInlineSnapshot(`
         "import { registerServerReference as __waku_registerServerReference } from 'react-server-dom-webpack/server.edge';
-        export const __waku_action1 = __waku_registerServerReference((mesg)=>{
+        export const __waku_action1 = __waku_registerServerReference(async (mesg)=>{
             console.log(mesg);
         }, "/src/App.tsx", "__waku_action1");
         const actions = {