[codex] Fix wrapped use server exports

packages/vinext/src/index.ts

@@ -64,6 +64,7 @@ import { asyncHooksStubPlugin } from "./plugins/async-hooks-stub.js";
 import { clientReferenceDedupPlugin } from "./plugins/client-reference-dedup.js";
 import { createInstrumentationClientTransformPlugin } from "./plugins/instrumentation-client.js";
 import { createOptimizeImportsPlugin } from "./plugins/optimize-imports.js";
+import { fixUseServerWrappedExportsPlugin } from "./plugins/fix-use-server-wrapped-exports.js";
 import { fixUseServerClosureCollisionPlugin } from "./plugins/fix-use-server-closure-collision.js";
 import { createOgInlineFetchAssetsPlugin, ogAssetsPlugin } from "./plugins/og-assets.js";
 import { createServerExternalsManifestPlugin } from "./plugins/server-externals-manifest.js";
@@ -1050,6 +1051,9 @@ export default function vinext(options: VinextOptions = {}): PluginOption[] {
     // Transform CJS require()/module.exports to ESM before other plugins
     // analyze imports (RSC directive scanning, shim resolution, etc.)
     commonjs(),
+    // Fix file-level "use server" exports that wrap an inline async action
+    // (e.g. next-safe-action) before plugin-rsc validates export shapes.
+    fixUseServerWrappedExportsPlugin,
     // Fix 'use server' closure variable collision with local declarations.
     // See packages/vinext/src/plugins/fix-use-server-closure-collision.ts for details.
     fixUseServerClosureCollisionPlugin,

packages/vinext/src/plugins/fix-use-server-wrapped-exports.ts

@@ -0,0 +1,235 @@
+import type { Plugin } from "vite";
+import { parseAst } from "vite";
+import MagicString from "magic-string";
+
+/**
+ * Fix file-level "use server" exports that wrap an inline async action.
+ *
+ * Libraries like next-safe-action expose server actions as wrapped call
+ * expressions:
+ *
+ *   "use server";
+ *   export const action = actionClient.action(async () => { ... });
+ *
+ * @vitejs/plugin-rsc's client proxy transform validates `export const`
+ * declarations syntactically and rejects anything whose initializer is not a
+ * direct async arrow/function expression. That makes wrapped server actions
+ * fail in vinext even though Next.js accepts them.
+ *
+ * Fix: before plugin-rsc runs, rewrite only those exports to:
+ *
+ *   const action = actionClient.action(async () => { ... });
+ *   export { action };
+ *
+ * The upstream proxy transform accepts export specifiers without re-validating
+ * the initializer shape, while the server-side transform still registers the
+ * correct local binding.
+ */
+export const fixUseServerWrappedExportsPlugin: Plugin = {
+  name: "vinext:fix-use-server-wrapped-exports",
+  enforce: "pre" as const,
+  transform(code: string, id: string) {
+    if (!code.includes("use server")) return null;
+    if (!/\.(js|jsx|ts|tsx|mjs|cjs)$/.test(id.split("?")[0])) return null;
+
+    // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+    let ast: any;
+    try {
+      ast = parseAst(code);
+    } catch {
+      return null;
+    }
+
+    if (!hasTopLevelUseServerDirective(ast.body)) return null;
+
+    const s = new MagicString(code);
+    const topLevelBindings = collectTopLevelBindings(ast.body);
+    let changed = false;
+
+    for (const node of ast.body) {
+      if (
+        node.type === "ExportNamedDeclaration" &&
+        node.declaration?.type === "VariableDeclaration" &&
+        shouldRewriteNamedDeclaration(node.declaration)
+      ) {
+        const exportNames = new Set<string>();
+        for (const decl of node.declaration.declarations) {
+          collectPatternNames(decl.id, exportNames);
+        }
+        if (exportNames.size === 0) continue;
+
+        s.overwrite(node.start, node.declaration.start, "");
+        s.appendLeft(node.end, `\nexport { ${[...exportNames].join(", ")} };`);
+        changed = true;
+        continue;
+      }
+
+      if (node.type === "ExportDefaultDeclaration" && shouldRewriteValue(node.declaration)) {
+        const localName = createUniqueName("__vinext_server_default__", topLevelBindings);
+        s.overwrite(node.start, node.declaration.start, `const ${localName} = `);
+        s.appendLeft(node.end, `\nexport default ${localName};`);
+        changed = true;
+      }
+    }
+
+    if (!changed) return null;
+    return {
+      code: s.toString(),
+      map: s.generateMap({ hires: "boundary" }),
+    };
+  },
+};
+
+function shouldRewriteNamedDeclaration(
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+  declaration: any,
+): boolean {
+  return declaration.declarations.some(
+    // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+    (decl: any) => decl.init && shouldRewriteValue(decl.init),
+  );
+}
+
+function shouldRewriteValue(
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+  node: any,
+): boolean {
+  return !isDirectAsyncFunctionNode(node) && containsInlineAsyncFunction(node);
+}
+
+function isDirectAsyncFunctionNode(
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+  node: any,
+): boolean {
+  return (
+    (node?.type === "FunctionDeclaration" ||
+      node?.type === "FunctionExpression" ||
+      node?.type === "ArrowFunctionExpression") &&
+    node.async === true
+  );
+}
+
+function containsInlineAsyncFunction(
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+  node: any,
+): boolean {
+  if (!node || typeof node !== "object") return false;
+
+  if (
+    (node.type === "FunctionExpression" || node.type === "ArrowFunctionExpression") &&
+    node.async === true
+  ) {
+    return true;
+  }
+
+  for (const key of Object.keys(node)) {
+    if (key === "type" || key === "start" || key === "end" || key === "loc" || key === "parent") {
+      continue;
+    }
+    const child = node[key];
+    if (Array.isArray(child)) {
+      for (const item of child) {
+        if (containsInlineAsyncFunction(item)) return true;
+      }
+    } else if (containsInlineAsyncFunction(child)) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+function hasTopLevelUseServerDirective(
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+  body: any[],
+): boolean {
+  for (const stmt of body) {
+    if (
+      stmt.type === "ExpressionStatement" &&
+      stmt.expression?.type === "Literal" &&
+      typeof stmt.expression.value === "string"
+    ) {
+      if (stmt.expression.value === "use server") return true;
+      continue;
+    }
+    break;
+  }
+  return false;
+}
+
+function createUniqueName(base: string, names: Set<string>): string {
+  let candidate = base;
+  let suffix = 0;
+  while (names.has(candidate)) {
+    suffix++;
+    candidate = `${base}${suffix}`;
+  }
+  names.add(candidate);
+  return candidate;
+}
+
+function collectTopLevelBindings(
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+  body: any[],
+): Set<string> {
+  const names = new Set<string>();
+
+  for (const node of body) {
+    if (node.type === "ImportDeclaration") {
+      for (const specifier of node.specifiers ?? []) {
+        if (specifier.local?.name) names.add(specifier.local.name);
+      }
+      continue;
+    }
+
+    if (node.type === "FunctionDeclaration" && node.id?.name) {
+      names.add(node.id.name);
+      continue;
+    }
+
+    if (node.type === "ClassDeclaration" && node.id?.name) {
+      names.add(node.id.name);
+      continue;
+    }
+
+    if (node.type === "VariableDeclaration") {
+      for (const decl of node.declarations) collectPatternNames(decl.id, names);
+      continue;
+    }
+
+    if (node.type === "ExportNamedDeclaration" && node.declaration) {
+      if (node.declaration.type === "FunctionDeclaration" && node.declaration.id?.name) {
+        names.add(node.declaration.id.name);
+      } else if (node.declaration.type === "ClassDeclaration" && node.declaration.id?.name) {
+        names.add(node.declaration.id.name);
+      } else if (node.declaration.type === "VariableDeclaration") {
+        for (const decl of node.declaration.declarations) collectPatternNames(decl.id, names);
+      }
+      continue;
+    }
+  }
+
+  return names;
+}
+
+function collectPatternNames(
+  // oxlint-disable-next-line @typescript-eslint/no-explicit-any
+  pattern: any,
+  names: Set<string>,
+) {
+  if (!pattern) return;
+
+  if (pattern.type === "Identifier") {
+    names.add(pattern.name);
+  } else if (pattern.type === "ObjectPattern") {
+    for (const prop of pattern.properties) {
+      collectPatternNames(prop.value ?? prop.argument, names);
+    }
+  } else if (pattern.type === "ArrayPattern") {
+    for (const elem of pattern.elements) {
+      collectPatternNames(elem, names);
+    }
+  } else if (pattern.type === "RestElement" || pattern.type === "AssignmentPattern") {
+    collectPatternNames(pattern.left ?? pattern.argument, names);
+  }
+}

tests/use-server-wrapped-exports.test.ts

@@ -0,0 +1,114 @@
+import { describe, expect, it } from "vite-plus/test";
+import { parseAst } from "vite-plus";
+import {
+  transformDirectiveProxyExport,
+  transformServerActionServer,
+} from "@vitejs/plugin-rsc/transforms";
+import vinext from "../packages/vinext/src/index.js";
+import type { Plugin } from "vite-plus";
+
+function unwrapHook(hook: unknown): Function {
+  if (typeof hook === "function") return hook;
+  const handler = (hook as { handler?: Function } | undefined)?.handler;
+  if (!handler) throw new Error("expected plugin hook handler");
+  return handler;
+}
+
+function getWrappedExportsPlugin(): Plugin | undefined {
+  const plugins = (vinext() as Plugin[]).flat(Infinity) as Plugin[];
+  return plugins.find((plugin) => plugin?.name === "vinext:fix-use-server-wrapped-exports");
+}
+
+async function runPlugin(source: string, id = "/app/actions.ts"): Promise<string> {
+  const plugin = getWrappedExportsPlugin();
+  let code = source;
+
+  if (!plugin?.transform) return code;
+
+  const transform = unwrapHook(plugin.transform);
+  const result = await transform.call(plugin, code, id);
+  if (result != null) {
+    code = typeof result === "string" ? result : result.code;
+  }
+  return code;
+}
+
+const WRAPPED_EXPORT_SOURCE = `
+"use server";
+
+import { actionClient } from "./lib/safe-action";
+
+export const testAction = actionClient.action(async () => {
+  return { message: "Hello, world!" };
+});
+`.trimStart();
+
+describe("vinext:fix-use-server-wrapped-exports", () => {
+  it("plugin is present in the vinext() plugin array", () => {
+    expect(getWrappedExportsPlugin()).toBeDefined();
+  });
+
+  it("reproduces the upstream bug: strict proxy transform rejects wrapped async exports", () => {
+    const ast = parseAst(WRAPPED_EXPORT_SOURCE);
+
+    expect(() =>
+      transformDirectiveProxyExport(ast as Parameters<typeof transformDirectiveProxyExport>[0], {
+        code: WRAPPED_EXPORT_SOURCE,
+        runtime: (name: string) => `createRef(${JSON.stringify(name)})`,
+        directive: "use server",
+        rejectNonAsyncFunction: true,
+      }),
+    ).toThrowError(/unsupported non async function/);
+  });
+
+  it("fix: rewrites wrapped async exports into local bindings plus export specifiers", async () => {
+    const output = await runPlugin(WRAPPED_EXPORT_SOURCE);
+
+    expect(output).toContain("const testAction = actionClient.action(async () => {");
+    expect(output).toContain("export { testAction };");
+
+    const ast = parseAst(output);
+    const result = transformDirectiveProxyExport(
+      ast as Parameters<typeof transformDirectiveProxyExport>[0],
+      {
+        code: output,
+        runtime: (name: string) => `createRef(${JSON.stringify(name)})`,
+        directive: "use server",
+        rejectNonAsyncFunction: true,
+      },
+    );
+
+    if (!result) throw new Error("expected proxy transform result");
+    expect(result.output.toString()).toContain(
+      'export const testAction = /* #__PURE__ */ createRef("testAction");',
+    );
+  });
+
+  it("preserves the server transform for wrapped async exports", async () => {
+    const output = await runPlugin(WRAPPED_EXPORT_SOURCE);
+    const ast = parseAst(output);
+
+    const result = transformServerActionServer(
+      output,
+      ast as Parameters<typeof transformServerActionServer>[1],
+      {
+        runtime: (_value: string, name: string) => `register(${JSON.stringify(name)})`,
+        rejectNonAsyncFunction: true,
+      },
+    );
+
+    expect(result.output.toString()).toContain('register("testAction")');
+  });
+
+  it("does not rewrite direct async function exports", async () => {
+    const source = `
+"use server";
+
+export const testAction = async () => {
+  return { message: "Hello, world!" };
+};
+`.trimStart();
+
+    expect(await runPlugin(source)).toBe(source);
+  });
+});