feat(rsc): Detect single RSA functions (not just entire files)

packages/vite/src/plugins/__tests__/vite-plugin-rsc-analyze.test.ts

@@ -0,0 +1,138 @@
+import type { TransformPluginContext } from 'rollup'
+import { beforeEach, describe, expect, it } from 'vitest'
+
+import { rscAnalyzePlugin } from '../vite-plugin-rsc-analyze.js'
+
+const foundFiles: Array<string> = []
+
+function callback(id: string) {
+  foundFiles.push(id)
+}
+
+function getPluginTransform() {
+  const plugin = rscAnalyzePlugin(callback, callback)
+
+  if (typeof plugin.transform !== 'function') {
+    throw new Error('Plugin does not have a transform function')
+  }
+
+  // Calling `bind` to please TS
+  // See https://stackoverflow.com/a/70463512/88106
+  // Typecasting because we're only going to call transform, and we don't need
+  // anything provided by the context.
+  return plugin.transform.bind({} as TransformPluginContext)
+}
+
+const pluginTransform = getPluginTransform()
+
+beforeEach(() => {
+  foundFiles.length = 0
+})
+
+describe('vite-plugin-rsc-analyze', () => {
+  it('finds "use server" action inlined as an arrow function', async () => {
+    const code = `
+      import { jsx, jsxs } from "react/jsx-runtime";
+      import fs from "node:fs";
+      import "./ServerDelayForm.css";
+      const ServerDelayForm = () => {
+        let delay = 0;
+        if (fs.existsSync("settings.json")) {
+          delay = JSON.parse(fs.readFileSync("settings.json", "utf8")).delay || 0;
+        }
+        return /* @__PURE__ */ jsx("div", { className: "server-delay-form", children: /* @__PURE__ */ jsxs("form", { action: async (formData) => {
+          "use server";
+          await fs.promises.writeFile("settings.json", \`{ "delay": \${formData.get("delay")} }
+      \`);
+        }, children: [
+          /* @__PURE__ */ jsxs("label", { htmlFor: "delay", children: [
+            /* @__PURE__ */ jsxs("div", { children: [
+              "Delay (",
+              delay,
+              "ms)"
+            ] }),
+            /* @__PURE__ */ jsx("input", { type: "number", id: "delay", name: "delay" })
+          ] }),
+          /* @__PURE__ */ jsx("button", { type: "submit", children: "Set" })
+        ] }) });
+      };
+      export default ServerDelayForm;
+      `
+
+    pluginTransform(code, 'test.tsx')
+
+    expect(foundFiles).toHaveLength(1)
+    expect(foundFiles[0]).toEqual('test.tsx')
+  })
+
+  it('finds "use server" action inlined as a named function', async () => {
+    const code = `
+      import { jsx, jsxs } from "react/jsx-runtime";
+      import fs from "node:fs";
+      import "./ServerDelayForm.css";
+      const ServerDelayForm = () => {
+        let delay = 0;
+        if (fs.existsSync("settings.json")) {
+          delay = JSON.parse(fs.readFileSync("settings.json", "utf8")).delay || 0;
+        }
+        return /* @__PURE__ */ jsx("div", { className: "server-delay-form", children: /* @__PURE__ */ jsxs("form", { action: async function formAction(formData) {
+          "use server";
+          await fs.promises.writeFile("settings.json", \`{ "delay": \${formData.get("delay")} }
+      \`);
+        }, children: [
+          /* @__PURE__ */ jsxs("label", { htmlFor: "delay", children: [
+            /* @__PURE__ */ jsxs("div", { children: [
+              "Delay (",
+              delay,
+              "ms)"
+            ] }),
+            /* @__PURE__ */ jsx("input", { type: "number", id: "delay", name: "delay" })
+          ] }),
+          /* @__PURE__ */ jsx("button", { type: "submit", children: "Set" })
+        ] }) });
+      };
+      export default ServerDelayForm;
+      `
+
+    pluginTransform(code, 'test.tsx')
+
+    expect(foundFiles).toHaveLength(1)
+    expect(foundFiles[0]).toEqual('test.tsx')
+  })
+
+  it('finds "use server" action as a named function', async () => {
+    const code = `
+      import { jsx, jsxs } from "react/jsx-runtime";
+      import fs from "node:fs";
+      import "./ServerDelayForm.css";
+      async function formAction(formData) {
+        "use server";
+        await fs.promises.writeFile("settings.json", \`{ "delay": \${formData.get("delay")} }
+      \`);
+      }
+      const ServerDelayForm = () => {
+        let delay = 0;
+        if (fs.existsSync("settings.json")) {
+          delay = JSON.parse(fs.readFileSync("settings.json", "utf8")).delay || 0;
+        }
+        return /* @__PURE__ */ jsx("div", { className: "server-delay-form", children: /* @__PURE__ */ jsxs("form", { action: formAction, children: [
+          /* @__PURE__ */ jsxs("label", { htmlFor: "delay", children: [
+            /* @__PURE__ */ jsxs("div", { children: [
+              "Delay (",
+              delay,
+              "ms)"
+            ] }),
+            /* @__PURE__ */ jsx("input", { type: "number", id: "delay", name: "delay" })
+          ] }),
+          /* @__PURE__ */ jsx("button", { type: "submit", children: "Set" })
+        ] }) });
+      };
+      export default ServerDelayForm;
+      `
+
+    pluginTransform(code, 'test.tsx')
+
+    expect(foundFiles).toHaveLength(1)
+    expect(foundFiles[0]).toEqual('test.tsx')
+  })
+})

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

@@ -13,26 +13,106 @@ export function rscAnalyzePlugin(
       const ext = path.extname(id)
 
       if (['.ts', '.tsx', '.js', '.jsx'].includes(ext)) {
+        // TODO (RSC): In a larger codebase, see if we'd be any faster by doing
+        // a simple code.includes('use client') || code.includes('use server')
+        // check first before parsing the code
+
         const mod = swc.parseSync(code, {
           syntax: ext === '.ts' || ext === '.tsx' ? 'typescript' : 'ecmascript',
           tsx: ext === '.tsx',
         })
 
+        let directiveFound = false
+
+        // The `item`s in mod.body are the top-level statements in the file
         for (const item of mod.body) {
           if (
             item.type === 'ExpressionStatement' &&
             item.expression.type === 'StringLiteral'
           ) {
             if (item.expression.value === 'use client') {
               clientEntryCallback(id)
+              directiveFound = true
             } else if (item.expression.value === 'use server') {
               serverEntryCallback(id)
+              directiveFound = true
             }
           }
         }
+
+        if (
+          !directiveFound &&
+          code.includes('use server') &&
+          containsServerAction(mod)
+        ) {
+          serverEntryCallback(id)
+        }
       }
 
       return code
     },
   }
 }
+
+function isServerAction(
+  node:
+    | swc.FunctionDeclaration
+    | swc.FunctionExpression
+    | swc.ArrowFunctionExpression,
+): boolean {
+  return (
+    node.body?.type === 'BlockStatement' &&
+    node.body.stmts.some(
+      (s) =>
+        s.type === 'ExpressionStatement' &&
+        s.expression.type === 'StringLiteral' &&
+        s.expression.value === 'use server',
+    )
+  )
+}
+
+function isFunctionDeclaration(
+  node: swc.Node,
+): node is swc.FunctionDeclaration {
+  return node.type === 'FunctionDeclaration'
+}
+
+function isFunctionExpression(node: swc.Node): node is swc.FunctionExpression {
+  return node.type === 'FunctionExpression'
+}
+
+function isArrowFunctionExpression(
+  node: swc.Node,
+): node is swc.ArrowFunctionExpression {
+  return node.type === 'ArrowFunctionExpression'
+}
+
+function containsServerAction(mod: swc.Module) {
+  function walk(node: swc.Node): boolean {
+    if (
+      isFunctionDeclaration(node) ||
+      isFunctionExpression(node) ||
+      isArrowFunctionExpression(node)
+    ) {
+      if (isServerAction(node)) {
+        return true
+      }
+    }
+
+    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)
+}

packages/vite/src/rsc/rscWorker.ts

@@ -428,8 +428,6 @@ async function handleRsa(input: RenderInput): Promise<PipeableStream> {
     throw new Error('Unexpected input')
   }
 
-  const config = await getViteConfig()
-
   const [fileName, actionName] = input.rsfId.split('#')
   console.log('Server Action fileName', fileName, 'actionName', actionName)
   const module = await loadServerFile(fileName)
@@ -450,6 +448,12 @@ async function handleRsa(input: RenderInput): Promise<PipeableStream> {
     input.args[0] = formData
   }
 
-  const data = await (module[actionName] || module)(...input.args)
+  const method = module[actionName] || module
+  console.log('rscWorker.ts method', method)
+  console.log('rscWorker.ts args', ...input.args)
+
+  const data = await method(...input.args)
+  const config = await getViteConfig()
+
   return renderToPipeableStream(data, getBundlerConfig(config))
 }