Skip to content

Commit 142fb07

Browse files
committed
feat(rsc): support custom server function directives
1 parent ac2325a commit 142fb07

5 files changed

Lines changed: 650 additions & 12 deletions

File tree

packages/plugin-rsc/src/plugin.ts

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,12 @@ import {
3939
withResolvedIdProxy,
4040
} from './plugins/resolved-id-proxy'
4141
import { scanBuildStripPlugin } from './plugins/scan'
42+
import {
43+
SERVER_FUNCTION_DIRECTIVE_MARKER,
44+
type ServerFunctionDirective,
45+
type ServerFunctionDirectiveContext,
46+
vitePluginServerFunctionDirectives,
47+
} from './plugins/server-function-directives'
4248
import {
4349
parseCssVirtual,
4450
toCssVirtual,
@@ -171,6 +177,8 @@ class RscPluginManager {
171177
}
172178

173179
export type RscPluginOptions = {
180+
/** @experimental */
181+
serverFunctionDirectives?: ServerFunctionDirective[]
174182
/**
175183
* shorthand for configuring `environments.(name).build.rollupOptions.input.index`
176184
*/
@@ -283,6 +291,8 @@ export type RscPluginOptions = {
283291
customClientEntry?: boolean
284292
}
285293

294+
export type { ServerFunctionDirective, ServerFunctionDirectiveContext }
295+
286296
export type PluginApi = {
287297
manager: RscPluginManager
288298
}
@@ -339,6 +349,30 @@ export function vitePluginRscMinimal(
339349
},
340350
...vitePluginRscCore(),
341351
...vitePluginUseClient(rscPluginOptions, manager),
352+
...(rscPluginOptions.serverFunctionDirectives?.length
353+
? [
354+
vitePluginServerFunctionDirectives({
355+
definitions: rscPluginOptions.serverFunctionDirectives,
356+
manager,
357+
serverEnvironmentName: rscPluginOptions.environment?.rsc ?? 'rsc',
358+
browserEnvironmentName:
359+
rscPluginOptions.environment?.browser ?? 'client',
360+
rscRuntime: resolvePackage(`${PKG_NAME}/react/rsc`),
361+
ssrRuntime: resolvePackage(`${PKG_NAME}/react/ssr`),
362+
browserRuntime: resolvePackage(`${PKG_NAME}/react/browser`),
363+
encryptionRuntime: resolvePackage(
364+
`${PKG_NAME}/utils/encryption-runtime`,
365+
),
366+
expandExportAll: (context, code, ast, id) =>
367+
transformExpandExportAll({
368+
code,
369+
ast,
370+
importer: id,
371+
...createTransformExpandExportAllContext(context),
372+
}),
373+
}),
374+
]
375+
: []),
342376
...vitePluginUseServer(rscPluginOptions, manager),
343377
...vitePluginDefineEncryptionKey(rscPluginOptions),
344378
{
@@ -1946,7 +1980,9 @@ function vitePluginUseServer(
19461980
// filter: { code: 'use server' },
19471981
async handler(code, id) {
19481982
if (!code.includes('use server')) {
1949-
delete manager.serverReferenceMetaMap[id]
1983+
if (!code.includes(SERVER_FUNCTION_DIRECTIVE_MARKER)) {
1984+
delete manager.serverReferenceMetaMap[id]
1985+
}
19501986
return
19511987
}
19521988
let ast = await parseAstAsync(code)
@@ -2017,11 +2053,15 @@ function vitePluginUseServer(
20172053
delete manager.serverReferenceMetaMap[id]
20182054
return
20192055
}
2056+
const customExportNames =
2057+
manager.serverReferenceMetaMap[id]?.exportNames ?? []
20202058
manager.serverReferenceMetaMap[id] = {
20212059
importId: id,
20222060
referenceKey: getNormalizedId(),
2023-
exportNames:
2024-
'names' in result ? result.names : result.exportNames,
2061+
exportNames: [
2062+
...customExportNames,
2063+
...('names' in result ? result.names : result.exportNames),
2064+
],
20252065
}
20262066
const importSource = resolvePackage(`${PKG_NAME}/react/rsc`)
20272067
output.prepend(
Lines changed: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
1+
import type { Rollup } from 'vite'
2+
import { describe, expect, it, vi } from 'vitest'
3+
import {
4+
SERVER_FUNCTION_DIRECTIVE_MARKER,
5+
vitePluginServerFunctionDirectives,
6+
type ServerFunctionDirective,
7+
} from './server-function-directives'
8+
9+
type EnvironmentName = 'rsc' | 'ssr' | 'client'
10+
11+
type Manager = Parameters<
12+
typeof vitePluginServerFunctionDirectives
13+
>[0]['manager']
14+
15+
function createHarness(
16+
definitions: ServerFunctionDirective[],
17+
command: 'build' | 'serve' = 'build',
18+
) {
19+
const manager: Manager = {
20+
config: { command },
21+
server: { environments: {} as Manager['server']['environments'] },
22+
toRelativeId: (id) => id,
23+
serverReferenceMetaMap: {},
24+
}
25+
const expandExportAll = vi.fn(async () => undefined)
26+
const plugin = vitePluginServerFunctionDirectives({
27+
definitions,
28+
manager,
29+
serverEnvironmentName: 'rsc',
30+
browserEnvironmentName: 'client',
31+
encryptionRuntime: '/encryption-runtime.js',
32+
rscRuntime: '/rsc-runtime.js',
33+
browserRuntime: '/browser-runtime.js',
34+
ssrRuntime: '/ssr-runtime.js',
35+
expandExportAll,
36+
})
37+
const transformHook = plugin.transform
38+
if (
39+
!transformHook ||
40+
typeof transformHook === 'function' ||
41+
!('handler' in transformHook)
42+
) {
43+
throw new Error('expected an object transform hook')
44+
}
45+
const transform = transformHook.handler
46+
47+
async function run(
48+
code: string,
49+
environment: EnvironmentName = 'rsc',
50+
id = '/src/example.ts',
51+
) {
52+
const context = {
53+
environment: {
54+
name: environment,
55+
mode: command === 'build' ? 'build' : 'dev',
56+
},
57+
} as Rollup.TransformPluginContext
58+
const result = await transform.call(context, code, id, { moduleType: 'js' })
59+
if (!result) return
60+
if (typeof result === 'string') return { code: result, map: null }
61+
return {
62+
code:
63+
typeof result.code === 'string' ? result.code : result.code?.toString(),
64+
map: result.map,
65+
}
66+
}
67+
68+
return { expandExportAll, manager, run }
69+
}
70+
71+
function cacheDirective(
72+
overrides: Partial<ServerFunctionDirective> = {},
73+
): ServerFunctionDirective {
74+
return {
75+
directive: /^use cache(?:: .+)?$/,
76+
test: (code) => code.includes('use cache'),
77+
rejectNonAsyncFunction: true,
78+
wrap: ({ value, directiveMatch, location }) =>
79+
`cache(${value}, ${JSON.stringify(directiveMatch[0])}, ${JSON.stringify(location)})`,
80+
...overrides,
81+
}
82+
}
83+
84+
describe('vitePluginServerFunctionDirectives', () => {
85+
it('hoists, wraps, and registers inline functions in RSC', async () => {
86+
const { manager, run } = createHarness([cacheDirective()])
87+
const result = await run(`
88+
export async function getData() {
89+
"use cache";
90+
return 1;
91+
}
92+
`)
93+
expect(result?.code).toContain(SERVER_FUNCTION_DIRECTIVE_MARKER)
94+
expect(result?.code).toContain('cache($$hoist_')
95+
expect(result?.code).toContain('$$ReactServer.registerServerReference')
96+
expect(result?.code).toContain('/rsc-runtime.js')
97+
expect(
98+
manager.serverReferenceMetaMap['/src/example.ts']?.exportNames,
99+
).toEqual([expect.stringMatching(/^\$\$hoist_/)])
100+
})
101+
102+
it('encrypts captured values without adding ciphertext to the wrapper', async () => {
103+
const wrap = vi.fn(({ value }: { value: string }) => `cache(${value})`)
104+
const { run } = createHarness([cacheDirective({ wrap })])
105+
const result = await run(`
106+
export async function outer(value) {
107+
return async function cached() {
108+
"use cache";
109+
return value;
110+
};
111+
}
112+
`)
113+
expect(result?.code).toContain('encryptActionBoundArgs([value])')
114+
expect(result?.code).toContain('decryptActionBoundArgs($$encoded)')
115+
expect(result?.code).toContain('/encryption-runtime.js')
116+
expect(wrap).toHaveBeenCalledWith(
117+
expect.objectContaining({ location: 'inline', hasBoundArgs: true }),
118+
)
119+
})
120+
121+
it('wraps module exports and records only selected references', async () => {
122+
const filterExport = vi.fn(
123+
({ name }: { name: string }) => name !== 'metadata',
124+
)
125+
const { expandExportAll, manager, run } = createHarness([
126+
cacheDirective({ filterExport }),
127+
])
128+
const result = await run(`
129+
"use cache";
130+
export async function getData() { return 1 }
131+
export const metadata = { title: "test" };
132+
`)
133+
expect(expandExportAll).toHaveBeenCalledOnce()
134+
expect(result?.code).toContain('cache(getData, "use cache", "module")')
135+
expect(result?.code).not.toContain('cache(metadata')
136+
expect(
137+
manager.serverReferenceMetaMap['/src/example.ts']?.exportNames,
138+
).toEqual(['getData'])
139+
expect(filterExport).toHaveBeenCalledWith(
140+
expect.objectContaining({ name: 'metadata', id: '/src/example.ts' }),
141+
)
142+
})
143+
144+
it.each([
145+
['client', '/browser-runtime.js'],
146+
['ssr', '/ssr-runtime.js'],
147+
] as const)('creates module proxies in %s', async (environment, runtime) => {
148+
const { run } = createHarness([cacheDirective()])
149+
const result = await run(
150+
`"use cache"; export async function getData() { return 1 }`,
151+
environment,
152+
)
153+
expect(result?.code).toContain(runtime)
154+
expect(result?.code).toContain('$$ReactClient.createServerReference')
155+
expect(result?.code).toContain('#getData')
156+
})
157+
158+
it('uses clientError for non-RSC module boundaries', async () => {
159+
const { run } = createHarness([
160+
cacheDirective({
161+
clientError: ({ id, environment }) => `${environment}:${id}`,
162+
}),
163+
])
164+
await expect(
165+
run(`"use cache"; export async function getData() {}`, 'client'),
166+
).rejects.toThrow('client:/src/example.ts')
167+
})
168+
169+
it('leaves non-server inline directives untouched', async () => {
170+
const { run } = createHarness([cacheDirective()])
171+
const code = `export async function getData() { "use cache" }`
172+
await expect(run(code, 'client')).resolves.toBeUndefined()
173+
await expect(run(code, 'ssr')).resolves.toBeUndefined()
174+
})
175+
176+
it('wraps inline directives inside use-server modules without owning metadata', async () => {
177+
const { manager, run } = createHarness([cacheDirective()])
178+
manager.serverReferenceMetaMap['/src/example.ts'] = {
179+
importId: '/src/example.ts',
180+
referenceKey: 'existing',
181+
exportNames: ['action'],
182+
}
183+
const result = await run(`
184+
"use server";
185+
export async function action() {
186+
async function cached() { "use cache"; return 1 }
187+
return cached();
188+
}
189+
`)
190+
expect(result?.code).toContain('cache($$hoist_')
191+
expect(result?.code).not.toContain('$$ReactServer.registerServerReference')
192+
expect(manager.serverReferenceMetaMap['/src/example.ts']).toEqual({
193+
importId: '/src/example.ts',
194+
referenceKey: 'existing',
195+
exportNames: ['action'],
196+
})
197+
})
198+
199+
it('rejects conflicting module-level custom and use-server directives', async () => {
200+
const { run } = createHarness([cacheDirective()])
201+
await expect(
202+
run(`"use cache"; "use server"; export async function action() {}`),
203+
).rejects.toThrow('cannot contain both')
204+
})
205+
206+
it('runs validation for inline and module directives', async () => {
207+
const validate = vi.fn()
208+
const { run } = createHarness([cacheDirective({ validate })])
209+
await run(`export async function getData() { "use cache: remote" }`)
210+
await run(`"use cache"; export async function getData() {}`)
211+
expect(validate).toHaveBeenCalledWith(
212+
expect.objectContaining({
213+
directive: 'use cache: remote',
214+
location: 'inline',
215+
}),
216+
)
217+
expect(validate).toHaveBeenCalledWith(
218+
expect.objectContaining({ directive: 'use cache', location: 'module' }),
219+
)
220+
})
221+
222+
it('rejects synchronous functions when configured', async () => {
223+
const { run } = createHarness([cacheDirective()])
224+
await expect(
225+
run(`export function getData() { "use cache" }`),
226+
).rejects.toThrow('non async function')
227+
})
228+
229+
it('respects source and id prefilters and clears stale metadata', async () => {
230+
const test = vi.fn(() => false)
231+
const filter = vi.fn(() => false)
232+
const { manager, run } = createHarness([
233+
cacheDirective({ test }),
234+
cacheDirective({ test: undefined, filter }),
235+
])
236+
manager.serverReferenceMetaMap['/src/example.ts'] = {
237+
importId: '/src/example.ts',
238+
referenceKey: 'stale',
239+
exportNames: ['stale'],
240+
}
241+
await expect(
242+
run(`export async function value() { "use cache" }`),
243+
).resolves.toBeUndefined()
244+
expect(test).toHaveBeenCalled()
245+
expect(filter).toHaveBeenCalled()
246+
expect(manager.serverReferenceMetaMap['/src/example.ts']).toBeUndefined()
247+
})
248+
249+
it('rejects overlapping module directive definitions', async () => {
250+
const { run } = createHarness([
251+
cacheDirective({ directive: /^use cache/ }),
252+
cacheDirective({ directive: 'use cache' }),
253+
])
254+
await expect(
255+
run(`"use cache"; export async function getData() {}`, 'client'),
256+
).rejects.toThrow('Multiple server function directives')
257+
})
258+
259+
it('returns source maps for server and proxy transforms', async () => {
260+
const { run } = createHarness([cacheDirective()])
261+
const server = await run(`export async function getData() { "use cache" }`)
262+
const client = await run(
263+
`"use cache"; export async function getData() {}`,
264+
'client',
265+
)
266+
expect(server?.map).toMatchObject({ version: 3 })
267+
expect(client?.map).toMatchObject({ version: 3 })
268+
})
269+
})

0 commit comments

Comments
 (0)