fix(rsc): preserve custom directive references

Author: james-elicx

packages/plugin-rsc/src/plugins/server-function-directives.test.ts

@@ -1,6 +1,7 @@
 import type { Rollup } from 'vite'
 import { describe, expect, it, vi } from 'vitest'
 import {
+  SERVER_FUNCTION_DIRECTIVE_MARKER,
   vitePluginServerFunctionDirectives,
   type ServerFunctionDirective,
 } from './server-function-directives'
@@ -181,7 +182,7 @@ export const metadata = { title: "test" };
     expect(expandExportAll).toHaveBeenCalledOnce()
     expect(result?.code).toMatchInlineSnapshot(`
       "/* __vite_rsc_server_function_directives__ */
-
+      import * as $$ReactServer from "/rsc-runtime.js";
 
       /* "use cache" */;
       async function getData() { return 1 }
@@ -199,17 +200,38 @@ export const metadata = { title: "test" };
     )
   })
 
-  it('creates module proxies in client', async () => {
+  it('ignores jsxDEV source metadata while rejecting user this expressions', async () => {
     const { run } = createHarness([cacheDirective()])
+    await expect(
+      run(`
+export async function Page() {
+  "use cache";
+  return _jsxDEV("p", { children: "test" }, void 0, false, { fileName: "page.tsx" }, this);
+}
+`),
+    ).resolves.toBeDefined()
+    await expect(
+      run(`export async function getData() { "use cache"; return this.value }`),
+    ).rejects.toThrow('"use cache" functions cannot use "this"')
+  })
+
+  it('creates module proxies in client', async () => {
+    const { manager, run } = createHarness([cacheDirective()])
     const result = await run(
       `"use cache"; export async function getData() { return 1 }`,
       'client',
     )
     expect(result?.code).toMatchInlineSnapshot(`
-      "import * as $$ReactClient from "/browser-runtime.js";
+      "/* __vite_rsc_server_function_directives__ */
+      import * as $$ReactClient from "/browser-runtime.js";
        export const getData = /* #__PURE__ */ $$ReactClient.createServerReference("53eb073e2100#getData",$$ReactClient.callServer,undefined,undefined,"getData");
       "
     `)
+    expect(manager.serverReferenceMetaMap['/src/example.ts']).toEqual({
+      importId: '/src/example.ts',
+      referenceKey: '53eb073e2100',
+      exportNames: ['getData'],
+    })
   })
 
   it('creates module proxies in SSR', async () => {
@@ -219,7 +241,8 @@ export const metadata = { title: "test" };
       'ssr',
     )
     expect(result?.code).toMatchInlineSnapshot(`
-      "import * as $$ReactClient from "/ssr-runtime.js";
+      "/* __vite_rsc_server_function_directives__ */
+      import * as $$ReactClient from "/ssr-runtime.js";
        export const getData = /* #__PURE__ */ $$ReactClient.createServerReference("53eb073e2100#getData",$$ReactClient.callServer,undefined,undefined,"getData");
       "
     `)
@@ -386,6 +409,18 @@ export class CacheClass {
     expect(manager.serverReferenceMetaMap['/src/example.ts']).toBeUndefined()
   })
 
+  it('preserves metadata when transformed output is processed again', async () => {
+    const { manager, run } = createHarness([cacheDirective()])
+    const result = await run(
+      `"use cache"; export async function getData() { return 1 }`,
+    )
+    expect(result?.code).toContain(SERVER_FUNCTION_DIRECTIVE_MARKER)
+    const metadata = manager.serverReferenceMetaMap['/src/example.ts']
+
+    await expect(run(result!.code!)).resolves.toBeUndefined()
+    expect(manager.serverReferenceMetaMap['/src/example.ts']).toEqual(metadata)
+  })
+
   it('rejects overlapping module directive definitions', async () => {
     const { run } = createHarness([
       cacheDirective({ directive: /^use cache/ }),

packages/plugin-rsc/src/plugins/server-function-directives.ts

@@ -167,6 +167,8 @@ export function vitePluginServerFunctionDirectives(options: Options): Plugin {
     name: 'rsc:server-function-directives',
     transform: {
       async handler(code, id) {
+        if (code.includes(SERVER_FUNCTION_DIRECTIVE_MARKER)) return
+
         const active = definitions.filter(
           (definition) =>
             (definition.test?.(code) ?? code.includes('use ')) &&
@@ -244,8 +246,13 @@ export function vitePluginServerFunctionDirectives(options: Options): Plugin {
               `$$ReactClient.createServerReference(${JSON.stringify(normalizedId + '#' + name)},$$ReactClient.callServer,undefined,${this.environment.mode === 'dev' ? '$$ReactClient.findSourceMapURL' : 'undefined'},${JSON.stringify(name)})`,
           })
           if (!result?.output.hasChanged()) return
+          manager.serverReferenceMetaMap[id] = {
+            importId: id,
+            referenceKey: normalizedId,
+            exportNames: result.exportNames,
+          }
           result.output.prepend(
-            `import * as $$ReactClient from ${JSON.stringify(this.environment.name === options.browserEnvironmentName ? options.browserRuntime : options.ssrRuntime)};\n`,
+            `${SERVER_FUNCTION_DIRECTIVE_MARKER}\nimport * as $$ReactClient from ${JSON.stringify(this.environment.name === options.browserEnvironmentName ? options.browserRuntime : options.ssrRuntime)};\n`,
           )
           return {
             code: result.output.toString(),
@@ -309,6 +316,7 @@ export function vitePluginServerFunctionDirectives(options: Options): Plugin {
             moduleDirective,
             moduleRuntime: (value, name, meta) => {
               if (!moduleMatch) return value
+              needsReactRuntime = true
               return `$$ReactServer.registerServerReference(${definition.wrap({ value, name, id, directiveMatch: moduleMatch, location: 'module', hasBoundArgs: false, parameters: meta.parameters, runtime: getRuntime(), meta })}, ${JSON.stringify(normalizedId)}, ${JSON.stringify(name)})`
             },
             inlineRuntime: (value, name, meta) => {

packages/plugin-rsc/src/transforms/hoist.ts

@@ -237,7 +237,7 @@ export function transformHoistInlineDirective(
 
 function validateForbiddenExpressions(body: BlockStatement, directive: string) {
   walk(body, {
-    enter(node) {
+    enter(node, parent) {
       if (
         node !== body &&
         (node.type === 'FunctionDeclaration' ||
@@ -246,8 +246,14 @@ function validateForbiddenExpressions(body: BlockStatement, directive: string) {
         this.skip()
         return
       }
+      const isJsxDevSourceThis =
+        node.type === 'ThisExpression' &&
+        parent?.type === 'CallExpression' &&
+        parent.arguments.at(-1) === node &&
+        parent.callee.type === 'Identifier' &&
+        /(?:^|_)jsxDEV$/.test(parent.callee.name)
       const expression =
-        node.type === 'ThisExpression'
+        node.type === 'ThisExpression' && !isJsxDevSourceThis
           ? 'this'
           : node.type === 'Super'
             ? 'super'