feat: server actions in server components

examples/11_form/src/components/App.tsx

@@ -1,5 +1,6 @@
 import { Counter } from './Counter';
 import { Form } from './Form';
+import { ServerForm } from './ServerForm';
 import { getMessage, greet, increment } from './funcs';
 
 type ServerFunction<T> = T extends (...args: infer A) => infer R
@@ -19,6 +20,7 @@ const App = ({ name }: { name: string }) => {
         message={getMessage()}
         greet={greet as unknown as ServerFunction<typeof greet>}
       />
+      <ServerForm />
     </div>
   );
 };

examples/11_form/src/components/ServerForm.tsx

@@ -0,0 +1,27 @@
+async function requestUsername(formData: FormData) {
+  'use server';
+  const username = formData.get('username');
+  console.log(`username: ${username}`);
+}
+
+// FIXME make this example more realistic
+export const ServerForm = () => {
+  return (
+    <>
+      <form action={requestUsername}>
+        <input type="text" name="username" />
+        <button type="submit">Request</button>
+      </form>
+      <form
+        action={async (formData: FormData) => {
+          'use server';
+          const hobby = formData.get('hobby');
+          console.log(`hobby: ${hobby}`);
+        }}
+      >
+        <input type="text" name="hobby" />
+        <button type="submit">Request</button>
+      </form>
+    </>
+  );
+};

examples/23_actions/src/components2/App.tsx

@@ -2,6 +2,7 @@ import { Balancer } from 'react-wrap-balancer';
 
 import Counter from './Counter';
 import { greet, getCounter, increment } from './funcs';
+import ButtonServer from './ButtonServer';
 
 type ServerFunction<T> = T extends (...args: infer A) => infer R
   ? (...args: A) => Promise<R>
@@ -19,6 +20,8 @@ const App = ({ name }: { name: string }) => {
         increment={increment as unknown as ServerFunction<typeof increment>}
       />
       <Balancer>My Awesome Title</Balancer>
+      <ButtonServer name="Button1" />
+      <ButtonServer name="Button2" />
     </div>
   );
 };

examples/23_actions/src/components2/ButtonClient.tsx

@@ -0,0 +1,9 @@
+'use client';
+
+export default function ButtonClient({
+  onClick,
+}: {
+  onClick: () => Promise<void>;
+}) {
+  return <button onClick={() => onClick()}>Click me!</button>;
+}

examples/23_actions/src/components2/ButtonServer.tsx

@@ -0,0 +1,18 @@
+import ButtonClient from './ButtonClient';
+
+let counter = 0;
+
+const ButtonServer = ({ name }: { name: string }) => {
+  const now = Date.now();
+  async function handleClick() {
+    'use server';
+    console.log('Button clicked!', name, now, ++counter);
+  }
+  return (
+    <div>
+      {name} <ButtonClient onClick={handleClick} />
+    </div>
+  );
+};
+
+export default ButtonServer;

examples/23_actions/src/entries.tsx

@@ -1,8 +1,7 @@
-import { lazy } from 'react';
 import { defineEntries } from 'waku/server';
 import { Slot } from 'waku/client';
 
-const App = lazy(() => import('./components2/App'));
+import App from './components2/App';
 
 export default defineEntries(
   // renderEntries

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

@@ -18,6 +18,54 @@ const hash = async (code: string): Promise<string> => {
     .slice(0, 9);
 };
 
+const isServerAction = (
+  node:
+    | swc.FunctionDeclaration
+    | swc.FunctionExpression
+    | swc.ArrowFunctionExpression,
+): boolean =>
+  node.body?.type === 'BlockStatement' &&
+  node.body.stmts.some(
+    (s) =>
+      s.type === 'ExpressionStatement' &&
+      s.expression.type === 'StringLiteral' &&
+      s.expression.value === 'use server',
+  );
+
+const containsServerAction = (mod: swc.Module): boolean => {
+  const walk = (node: swc.Node): boolean => {
+    if (
+      node.type === 'FunctionDeclaration' ||
+      node.type === 'FunctionExpression' ||
+      node.type === 'ArrowFunctionExpression'
+    ) {
+      if (
+        isServerAction(
+          node as
+            | swc.FunctionDeclaration
+            | swc.FunctionExpression
+            | swc.ArrowFunctionExpression,
+        )
+      ) {
+        return true;
+      }
+    }
+    // FIXME do we need to walk the entire tree? feels inefficient
+    return Object.values(node).some((value) =>
+      (Array.isArray(value) ? value : [value]).some((v) => {
+        if (typeof v?.type === 'string') {
+          return walk(v);
+        }
+        if (typeof v?.expression?.type === 'string') {
+          return walk(v.expression);
+        }
+        return false;
+      }),
+    );
+  };
+  return walk(mod);
+};
+
 export function rscAnalyzePlugin(
   opts:
     | {
@@ -54,6 +102,15 @@ export function rscAnalyzePlugin(
             }
           }
         }
+        if (
+          !opts.isClient &&
+          !opts.clientFileSet.has(id) &&
+          !opts.serverFileSet.has(id) &&
+          code.includes('use server') &&
+          containsServerAction(mod)
+        ) {
+          opts.serverFileSet.add(id);
+        }
       }
       // Avoid walking after the client boundary
       if (!opts.isClient && opts.clientFileSet.has(id)) {

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

@@ -68,6 +68,136 @@ export ${name === 'default' ? name : `const ${name} =`} createServerReference('$
   }
 };
 
+const createIdentifier = (value: string): swc.Identifier => ({
+  type: 'Identifier',
+  value,
+  optional: false,
+  span: { start: 0, end: 0, ctxt: 0 },
+});
+
+const createStringLiteral = (value: string): swc.StringLiteral => ({
+  type: 'StringLiteral',
+  value,
+  span: { start: 0, end: 0, ctxt: 0 },
+});
+
+const serverActionsInitCode = swc.parseSync(`
+import { registerServerReference as __waku_registerServerReference__ } from 'react-server-dom-webpack/server';
+export const __waku_serverActions__ = new Map();
+let __waku_actionIndex__ = 0;
+function __waku_registerServerAction__(fn, actionId) {
+  const actionName = 'action' + __waku_actionIndex__++;
+  __waku_registerServerReference__(fn, actionId, actionName);
+  // FIXME this can cause memory leaks
+  __waku_serverActions__.set(actionName, fn);
+  return fn;
+}
+`).body;
+
+type FunctionWithBlockBody = (
+  | swc.FunctionDeclaration
+  | swc.FunctionExpression
+  | swc.ArrowFunctionExpression
+) & { body: swc.BlockStatement };
+
+const isServerAction = (node: swc.Node): node is FunctionWithBlockBody =>
+  (node.type === 'FunctionDeclaration' ||
+    node.type === 'FunctionExpression' ||
+    node.type === 'ArrowFunctionExpression') &&
+  (node as { body?: { type: string } }).body?.type === 'BlockStatement' &&
+  (node as FunctionWithBlockBody).body.stmts.some(
+    (s) =>
+      s.type === 'ExpressionStatement' &&
+      s.expression.type === 'StringLiteral' &&
+      s.expression.value === 'use server',
+  );
+
+const transformServerActions = (
+  mod: swc.Module,
+  getActionId: () => string,
+): swc.Module | void => {
+  let hasServerActions = false;
+  const registerServerAction: {
+    (fn: swc.FunctionDeclaration): swc.ExpressionStatement;
+    (
+      fn: swc.FunctionExpression | swc.ArrowFunctionExpression,
+    ): swc.CallExpression;
+  } = (fn): any => {
+    hasServerActions = true;
+    const exp: swc.CallExpression = {
+      type: 'CallExpression',
+      callee: createIdentifier('__waku_registerServerAction__'),
+      arguments: [
+        { expression: fn.type === 'FunctionDeclaration' ? fn.identifier : fn },
+        { expression: createStringLiteral(getActionId()) },
+      ],
+      span: { start: 0, end: 0, ctxt: 0 },
+    };
+    if (fn.type !== 'FunctionDeclaration') {
+      return exp;
+    }
+    return {
+      type: 'ExpressionStatement',
+      expression: exp,
+      span: { start: 0, end: 0, ctxt: 0 },
+    };
+  };
+  const handleStatements = (stmts: swc.Statement[] | swc.ModuleItem[]) => {
+    for (let i = 0; i < stmts.length; ++i) {
+      const stmt = stmts[i]!;
+      if (isServerAction(stmt)) {
+        const registerStmt = registerServerAction(stmt);
+        stmts.splice(++i, 0, registerStmt);
+      }
+    }
+  };
+  const handleExpression = (exp: swc.Expression) => {
+    if (isServerAction(exp)) {
+      const callExp = registerServerAction(Object.assign({}, exp));
+      Object.keys(exp).forEach((key) => {
+        delete exp[key as keyof typeof exp];
+      });
+      Object.assign(exp, callExp);
+    }
+  };
+  const walk = (node: swc.Node) => {
+    // FIXME do we need to walk the entire tree? feels inefficient
+    Object.values(node).forEach((value) => {
+      (Array.isArray(value) ? value : [value]).forEach((v) => {
+        if (typeof v?.type === 'string') {
+          walk(v);
+        } else if (typeof v?.expression?.type === 'string') {
+          walk(v.expression);
+        }
+      });
+    });
+    if (node.type === 'Module') {
+      const { body } = node as swc.Module;
+      handleStatements(body);
+    } else if (node.type === 'BlockStatement') {
+      const { stmts } = node as swc.BlockStatement;
+      handleStatements(stmts);
+    } else if (
+      node.type === 'FunctionExpression' ||
+      node.type === 'ArrowFunctionExpression'
+    ) {
+      handleExpression(
+        node as swc.FunctionExpression | swc.ArrowFunctionExpression,
+      );
+    }
+  };
+  walk(mod);
+  if (!hasServerActions) {
+    return;
+  }
+  const lastImportIndex = mod.body.findIndex(
+    (node) =>
+      node.type !== 'ExpressionStatement' && node.type !== 'ImportDeclaration',
+  );
+  mod.body.splice(lastImportIndex, 0, ...serverActionsInitCode);
+  return mod;
+};
+
 const transformServer = (
   code: string,
   id: string,
@@ -119,6 +249,14 @@ if (typeof ${name} === 'function') {
     }
     return newCode;
   }
+  // transform server actions in server components
+  const newMod =
+    code.includes('use server') &&
+    transformServerActions(mod, () => getServerId(id));
+  if (newMod) {
+    const newCode = swc.printSync(newMod).code;
+    return newCode;
+  }
 };
 
 export function rscTransformPlugin(

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

@@ -195,7 +195,7 @@ export async function renderRsc(
       }
       mod = await loadModule(fileId.slice('@id/'.length));
     }
-    const fn = mod[name] || mod;
+    const fn = mod.__waku_serverActions__?.get(name) || mod[name] || mod;
     return renderWithContextWithAction(context, fn, args);
   }