Skip to content

Commit 1aabe31

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

2 files changed

Lines changed: 326 additions & 3 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: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
1+
import { exactRegex } from '@rolldown/pluginutils'
2+
import type { Program } from 'estree'
3+
import type { Plugin, ResolvedConfig, Rollup, ViteDevServer } from 'vite'
4+
import { parseAstAsync } from 'vite'
5+
import {
6+
hasDirective,
7+
transformDirectiveProxyExport,
8+
transformServerActionServer,
9+
type TransformWrapExportFilter,
10+
} from '../transforms'
11+
import { hashString } from './utils'
12+
import { normalizeViteImportAnalysisUrl } from './vite-utils'
13+
14+
export const SERVER_FUNCTION_DIRECTIVE_MARKER =
15+
'/* __vite_rsc_server_function_directives__ */'
16+
17+
export type ServerFunctionDirectiveContext = {
18+
value: string
19+
name: string
20+
id: string
21+
directiveMatch: RegExpMatchArray
22+
location: 'inline' | 'module'
23+
hasBoundArgs: boolean
24+
meta?: Parameters<TransformWrapExportFilter>[1]
25+
}
26+
27+
export type ServerFunctionDirective = {
28+
directive: string | RegExp
29+
test?: (code: string) => boolean
30+
filter?: (id: string) => boolean
31+
validate?: (context: {
32+
id: string
33+
directive: string
34+
location: 'inline' | 'module'
35+
}) => void
36+
rejectNonAsyncFunction?: boolean
37+
wrap: (context: ServerFunctionDirectiveContext) => string
38+
filterExport?: (context: {
39+
name: string
40+
id: string
41+
meta: Parameters<TransformWrapExportFilter>[1]
42+
}) => boolean
43+
clientError?: (context: { id: string; environment: string }) => string
44+
}
45+
46+
type Options = {
47+
definitions: ServerFunctionDirective[]
48+
manager: {
49+
config: Pick<ResolvedConfig, 'command'>
50+
server: Pick<ViteDevServer, 'environments'>
51+
toRelativeId: (id: string) => string
52+
serverReferenceMetaMap: Record<
53+
string,
54+
{ importId: string; referenceKey: string; exportNames: string[] }
55+
>
56+
}
57+
serverEnvironmentName: string
58+
browserEnvironmentName: string
59+
encryptionRuntime: string
60+
rscRuntime: string
61+
browserRuntime: string
62+
ssrRuntime: string
63+
expandExportAll: (
64+
context: Rollup.TransformPluginContext,
65+
code: string,
66+
ast: Program,
67+
id: string,
68+
) => Promise<{ code: string } | undefined>
69+
}
70+
71+
function matchDirective(value: string, directive: string | RegExp) {
72+
const pattern =
73+
typeof directive === 'string'
74+
? exactRegex(directive)
75+
: new RegExp(directive.source, directive.flags)
76+
pattern.lastIndex = 0
77+
return value.match(pattern) ?? undefined
78+
}
79+
80+
function findModuleDirective(ast: Program, directive: string | RegExp) {
81+
const statement = ast.body.find(
82+
(node) =>
83+
node.type === 'ExpressionStatement' &&
84+
node.expression.type === 'Literal' &&
85+
typeof node.expression.value === 'string' &&
86+
matchDirective(node.expression.value, directive),
87+
)
88+
return statement?.type === 'ExpressionStatement' &&
89+
statement.expression.type === 'Literal'
90+
? statement.expression
91+
: undefined
92+
}
93+
94+
export function vitePluginServerFunctionDirectives(options: Options): Plugin {
95+
const { definitions, manager } = options
96+
return {
97+
name: 'rsc:server-function-directives',
98+
transform: {
99+
async handler(code, id) {
100+
const active = definitions.filter(
101+
(definition) =>
102+
(definition.test?.(code) ?? code.includes('use ')) &&
103+
(!definition.filter || definition.filter(id)),
104+
)
105+
const isServer = this.environment.name === options.serverEnvironmentName
106+
if (active.length === 0) {
107+
if (isServer) delete manager.serverReferenceMetaMap[id]
108+
return
109+
}
110+
111+
let ast = await parseAstAsync(code)
112+
const useServerBoundary = hasDirective(ast.body, 'use server')
113+
if (!isServer && useServerBoundary) return
114+
115+
let normalizedId: string | undefined
116+
const getNormalizedId = () =>
117+
(normalizedId ??=
118+
manager.config.command === 'build'
119+
? hashString(manager.toRelativeId(id))
120+
: normalizeViteImportAnalysisUrl(
121+
manager.server.environments[options.serverEnvironmentName]!,
122+
id,
123+
))
124+
125+
if (isServer) {
126+
const exportNames = new Set<string>()
127+
let needsReactRuntime = false
128+
let needsEncryptionRuntime = false
129+
let outputMap
130+
131+
for (const definition of active) {
132+
let moduleDirective = findModuleDirective(ast, definition.directive)
133+
if (moduleDirective && useServerBoundary) {
134+
throw Object.assign(
135+
new Error(
136+
`A module cannot contain both ${JSON.stringify(moduleDirective.value)} and "use server" directives.`,
137+
),
138+
{ pos: moduleDirective.start },
139+
)
140+
}
141+
if (moduleDirective) {
142+
const expanded = await options.expandExportAll(
143+
this,
144+
code,
145+
ast,
146+
id,
147+
)
148+
if (expanded) {
149+
code = expanded.code
150+
ast = await parseAstAsync(code)
151+
moduleDirective = findModuleDirective(ast, definition.directive)
152+
}
153+
}
154+
const moduleMatch = moduleDirective
155+
? matchDirective(
156+
moduleDirective.value as string,
157+
definition.directive,
158+
)
159+
: undefined
160+
if (moduleMatch) {
161+
definition.validate?.({
162+
id,
163+
directive: moduleMatch[0],
164+
location: 'module',
165+
})
166+
}
167+
168+
const result = transformServerActionServer(code, ast, {
169+
runtime: (value, name) =>
170+
`$$ReactServer.registerServerReference(${value}, ${JSON.stringify(getNormalizedId())}, ${JSON.stringify(name)})`,
171+
directive: definition.directive,
172+
moduleDirective,
173+
moduleRuntime: (value, name, meta) =>
174+
`$$ReactServer.registerServerReference(${definition.wrap({ value, name, id, directiveMatch: moduleMatch!, location: 'module', hasBoundArgs: false, meta })}, ${JSON.stringify(getNormalizedId())}, ${JSON.stringify(name)})`,
175+
inlineRuntime: (value, name, meta) => {
176+
definition.validate?.({
177+
id,
178+
directive: meta.directiveMatch[0],
179+
location: 'inline',
180+
})
181+
const wrapped = definition.wrap({
182+
value,
183+
name,
184+
id,
185+
directiveMatch: meta.directiveMatch,
186+
location: 'inline',
187+
hasBoundArgs: meta.hasBoundArgs,
188+
})
189+
if (useServerBoundary) return wrapped
190+
needsReactRuntime = true
191+
if (meta.hasBoundArgs) {
192+
needsEncryptionRuntime = true
193+
return `$$ReactServer.registerServerReference((($$wrapped) => async ($$encoded, ...$$args) => $$wrapped(...await __vite_rsc_encryption_runtime.decryptActionBoundArgs($$encoded), ...$$args))(${wrapped}), ${JSON.stringify(getNormalizedId())}, ${JSON.stringify(name)})`
194+
}
195+
return `$$ReactServer.registerServerReference(${wrapped}, ${JSON.stringify(getNormalizedId())}, ${JSON.stringify(name)})`
196+
},
197+
filter: (name, meta) =>
198+
definition.filterExport?.({ name, id, meta }) ?? true,
199+
rejectNonAsyncFunction: definition.rejectNonAsyncFunction,
200+
encode: (value) => {
201+
needsEncryptionRuntime = true
202+
return `__vite_rsc_encryption_runtime.encryptActionBoundArgs(${value})`
203+
},
204+
stableName: true,
205+
})
206+
if (!result.output.hasChanged()) continue
207+
for (const name of 'names' in result
208+
? result.names
209+
: result.exportNames) {
210+
exportNames.add(name)
211+
}
212+
outputMap = result.output.generateMap({
213+
hires: 'boundary',
214+
source: id,
215+
})
216+
code = result.output.toString()
217+
ast = await parseAstAsync(code)
218+
}
219+
220+
if (!useServerBoundary) {
221+
if (exportNames.size === 0)
222+
delete manager.serverReferenceMetaMap[id]
223+
else {
224+
manager.serverReferenceMetaMap[id] = {
225+
importId: id,
226+
referenceKey: getNormalizedId(),
227+
exportNames: [...exportNames],
228+
}
229+
}
230+
}
231+
const imports = [
232+
needsReactRuntime &&
233+
`import * as $$ReactServer from ${JSON.stringify(options.rscRuntime)};`,
234+
needsEncryptionRuntime &&
235+
`import * as __vite_rsc_encryption_runtime from ${JSON.stringify(options.encryptionRuntime)};`,
236+
].filter(Boolean)
237+
return {
238+
code: `${SERVER_FUNCTION_DIRECTIVE_MARKER}\n${imports.join('\n')}\n${code}`,
239+
map: outputMap,
240+
}
241+
}
242+
243+
const matches = active.flatMap((definition) => {
244+
const moduleDirective = findModuleDirective(ast, definition.directive)
245+
return moduleDirective ? [[definition, moduleDirective] as const] : []
246+
})
247+
if (matches.length === 0) return
248+
if (matches.length > 1) {
249+
throw Object.assign(
250+
new Error('Multiple server function directives match this module.'),
251+
{ pos: matches[1]![1].start },
252+
)
253+
}
254+
const [definition, moduleDirective] = matches[0]!
255+
if (definition.clientError) {
256+
throw Object.assign(
257+
new Error(
258+
definition.clientError({
259+
id,
260+
environment: this.environment.name,
261+
}),
262+
),
263+
{ pos: moduleDirective.start },
264+
)
265+
}
266+
const result = transformDirectiveProxyExport(ast, {
267+
code,
268+
directive: moduleDirective.value as string,
269+
runtime: (name) =>
270+
`$$ReactClient.createServerReference(${JSON.stringify(getNormalizedId() + '#' + name)},$$ReactClient.callServer,undefined,${this.environment.mode === 'dev' ? '$$ReactClient.findSourceMapURL' : 'undefined'},${JSON.stringify(name)})`,
271+
})
272+
if (!result?.output.hasChanged()) return
273+
result.output.prepend(
274+
`import * as $$ReactClient from ${JSON.stringify(this.environment.name === options.browserEnvironmentName ? options.browserRuntime : options.ssrRuntime)};\n`,
275+
)
276+
return {
277+
code: result.output.toString(),
278+
map: result.output.generateMap({ hires: 'boundary', source: id }),
279+
}
280+
},
281+
},
282+
}
283+
}

0 commit comments

Comments
 (0)