feat(rsc): expose callable argument metadata

packages/plugin-rsc/e2e/use-cache-callable.test.ts

@@ -13,6 +13,23 @@ test.describe('build', () => {
 })
 
 function defineTests(f: Fixture) {
+  test('argument admission and capture identity', async ({ page }) => {
+    using _errors = expectNoPageError(page)
+    await page.goto(f.url('/argument-admission'))
+
+    await expect(page.getByTestId('fixed-extra-admission')).toHaveText(
+      'excluded',
+    )
+    await expect(page.getByTestId('fixed-declared-admission')).toHaveText(
+      'included',
+    )
+    await expect(page.getByTestId('fixed-implementation-arguments')).toHaveText(
+      'preserved',
+    )
+    await expect(page.getByTestId('rest-admission')).toHaveText('included')
+    await expect(page.getByTestId('capture-admission')).toHaveText('included')
+  })
+
   test('inline directive', async ({ page }) => {
     using _errors = expectNoPageError(page)
     await page.goto(f.url())

packages/plugin-rsc/examples/use-cache-callable/callable-cache-plugin.ts

@@ -4,6 +4,7 @@ import {
   transformDirectiveProxyExport,
   transformHoistInlineDirective,
   transformWrapExport,
+  type FunctionParameters,
 } from '@vitejs/plugin-rsc/transforms'
 import { parseAstAsync, type Plugin } from 'vite'
 
@@ -29,9 +30,19 @@ export function callableCachePlugin(): Plugin {
       const environmentName = this.environment.name
 
       if (environmentName === 'rsc') {
-        const runtime = (value: string, name: string) =>
+        const runtime = (
+          value: string,
+          name: string,
+          meta: {
+            parameters?: FunctionParameters
+            boundArgumentCount?: number
+          },
+        ) =>
           `$$ReactServer.registerServerReference(` +
-          `$$cacheWrapper(${value}),` +
+          `$$cacheWrapper(${value}, ${JSON.stringify({
+            parameters: meta.parameters,
+            boundArgumentCount: meta.boundArgumentCount ?? 0,
+          })}),` +
           `${JSON.stringify(reference.referenceKey)},` +
           `${JSON.stringify(name)})`
         const result = hasDirective(ast.body, directive)
@@ -51,6 +62,10 @@ export function callableCachePlugin(): Plugin {
               rejectNonAsyncFunction: true,
               hoistRuntime: true,
               runtime,
+              encode: (value) =>
+                `$$Encryption.encryptActionBoundArgs(${value})`,
+              decode: (value) =>
+                `await $$Encryption.decryptActionBoundArgs(${value})`,
             })
         if (!result.output.hasChanged()) {
           manager.serverReferences.deleteClaim(pluginName, id)
@@ -63,6 +78,7 @@ export function callableCachePlugin(): Plugin {
         })
         result.output.prepend(
           `import $$cacheWrapper from "/src/framework/use-cache-runtime";\n` +
+            `import * as $$Encryption from "@vitejs/plugin-rsc/utils/encryption-runtime";\n` +
             `import * as $$ReactServer from "@vitejs/plugin-rsc/react/rsc/server";\n`,
         )
         return {

packages/plugin-rsc/examples/use-cache-callable/src/features/argument-admission/server.tsx

@@ -0,0 +1,60 @@
+async function fixed(declared = 'default') {
+  'use cache'
+  return `${declared}:${String(arguments[1])}:${crypto.randomUUID()}`
+}
+
+async function rest(declared: string, ...remaining: string[]) {
+  'use cache'
+  return `${declared}:${remaining.join(':')}:${crypto.randomUUID()}`
+}
+
+function withCapture(captured: string) {
+  return async function cached(declared: string) {
+    'use cache'
+    return `${captured}:${declared}:${crypto.randomUUID()}`
+  }
+}
+
+export async function ArgumentAdmission() {
+  const callFixed = fixed as (...args: string[]) => Promise<string>
+  const fixedFirst = await callFixed('same', 'first extra')
+  const fixedSecond = await callFixed('same', 'second extra')
+  const fixedDifferent = await callFixed('different', 'second extra')
+
+  const restFirst = await rest('same', 'first extra')
+  const restSecond = await rest('same', 'second extra')
+
+  const capturedAlpha = withCapture('alpha')
+  const capturedAlphaAgain = withCapture('alpha')
+  const capturedBeta = withCapture('beta')
+  const captureFirst = await capturedAlpha('same')
+  const captureSecond = await capturedAlphaAgain('same')
+  const captureDifferent = await capturedBeta('same')
+
+  return (
+    <dl>
+      <dt>Undeclared fixed-signature extras</dt>
+      <dd data-testid="fixed-extra-admission">
+        {fixedFirst === fixedSecond ? 'excluded' : 'included'}
+      </dd>
+      <dt>Declared fixed-signature arguments</dt>
+      <dd data-testid="fixed-declared-admission">
+        {fixedSecond !== fixedDifferent ? 'included' : 'excluded'}
+      </dd>
+      <dt>Complete implementation arguments</dt>
+      <dd data-testid="fixed-implementation-arguments">
+        {fixedFirst.startsWith('same:first extra:') ? 'preserved' : 'truncated'}
+      </dd>
+      <dt>Rest arguments</dt>
+      <dd data-testid="rest-admission">
+        {restFirst !== restSecond ? 'included' : 'excluded'}
+      </dd>
+      <dt>Decoded captures</dt>
+      <dd data-testid="capture-admission">
+        {captureFirst === captureSecond && captureSecond !== captureDifferent
+          ? 'included'
+          : 'excluded'}
+      </dd>
+    </dl>
+  )
+}

packages/plugin-rsc/examples/use-cache-callable/src/framework/use-cache-runtime.tsx

@@ -1,3 +1,4 @@
+import { decryptActionBoundArgs } from '@vitejs/plugin-rsc/rsc'
 import {
   createClientTemporaryReferenceSet,
   createFromReadableStream,
@@ -19,7 +20,15 @@ let cachedFnCacheEntries = new WeakMap<
   Record<string, Promise<StreamCacher>>
 >()
 
-export default function cacheWrapper(fn: (...args: any[]) => Promise<unknown>) {
+type CacheFunctionMetadata = {
+  parameters?: { count: number; hasRest: boolean }
+  boundArgumentCount: number
+}
+
+export default function cacheWrapper(
+  fn: (...args: any[]) => Promise<unknown>,
+  metadata: CacheFunctionMetadata,
+) {
   if (cachedFnMap.has(fn)) {
     return cachedFnMap.get(fn)!
   }
@@ -31,17 +40,36 @@ export default function cacheWrapper(fn: (...args: any[]) => Promise<unknown>) {
       cachedFnCacheEntries.set(cachedFn, cacheEntries)
     }
 
-    // Serialize arguments to a cache key via `encodeReply` from `react-server-dom/client`.
+    const boundArguments = args.slice(0, metadata.boundArgumentCount)
+    const invocationArguments = args.slice(metadata.boundArgumentCount)
+    const decodedBoundArguments =
+      metadata.boundArgumentCount === 0
+        ? []
+        : ((await decryptActionBoundArgs(boundArguments[0])) as unknown[])
+    const admittedInvocationArguments =
+      !metadata.parameters || metadata.parameters.hasRest
+        ? invocationArguments
+        : invocationArguments.slice(0, metadata.parameters.count)
+
+    // Serialize admitted arguments to a cache key via `encodeReply` from `react-server-dom/client`.
     // NOTE: using `renderToReadableStream` here for arguments serialization would end up
     // serializing react elements (e.g. children props), which causes
     // those arguments to be included as a cache key and it doesn't achieve
     // "use cache static shell + dynamic children props" pattern.
     // cf. https://nextjs.org/docs/app/api-reference/directives/use-cache#non-serializable-arguments
     const clientTemporaryReferences = createClientTemporaryReferenceSet()
+    const encodedCacheKey = await encodeReply(
+      [...decodedBoundArguments, ...admittedInvocationArguments],
+      {
+        temporaryReferences: clientTemporaryReferences,
+      },
+    )
+    const serializedCacheKey = await replyToCacheKey(encodedCacheKey)
+
+    // Keep the complete generated argument list for implementation invocation.
     const encodedArguments = await encodeReply(args, {
       temporaryReferences: clientTemporaryReferences,
     })
-    const serializedCacheKey = await replyToCacheKey(encodedArguments)
 
     // cache `fn` result as stream
     // (cache value is promise so that it dedupes concurrent async calls)

packages/plugin-rsc/examples/use-cache-callable/src/root.tsx

@@ -1,8 +1,16 @@
+import { ArgumentAdmission } from './features/argument-admission/server'
 import { FileDirectiveFromClientServer } from './features/file-directive-from-client/server'
 import { FileDirectiveFromServer } from './features/file-directive-from-server/server'
 import { InlineDirective } from './features/inline-directive/server'
 
 const routes = [
+  {
+    path: '/argument-admission',
+    title: 'Argument admission',
+    description:
+      'Source parameters and decoded closure captures determine cache identity.',
+    Component: ArgumentAdmission,
+  },
   {
     path: '/inline-directive',
     title: 'Inline directive',

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

@@ -493,11 +493,11 @@ export async function kv() {
       }),
     ).toMatchInlineSnapshot(`
       "
-      export const none = /* #__PURE__ */ $$register($$hoist_0_none, "<id>", "$$hoist_0_none", {"directiveMatch":["use cache",null]});
+      export const none = /* #__PURE__ */ $$register($$hoist_0_none, "<id>", "$$hoist_0_none", {"directiveMatch":["use cache",null],"parameters":{"count":0,"hasRest":false},"boundArgumentCount":0});
 
-      export const fs = /* #__PURE__ */ $$register($$hoist_1_fs, "<id>", "$$hoist_1_fs", {"directiveMatch":["use cache: fs",": fs"]});
+      export const fs = /* #__PURE__ */ $$register($$hoist_1_fs, "<id>", "$$hoist_1_fs", {"directiveMatch":["use cache: fs",": fs"],"parameters":{"count":0,"hasRest":false},"boundArgumentCount":0});
 
-      export const kv = /* #__PURE__ */ $$register($$hoist_2_kv, "<id>", "$$hoist_2_kv", {"directiveMatch":["use cache: kv",": kv"]});
+      export const kv = /* #__PURE__ */ $$register($$hoist_2_kv, "<id>", "$$hoist_2_kv", {"directiveMatch":["use cache: kv",": kv"],"parameters":{"count":0,"hasRest":false},"boundArgumentCount":0});
 
       ;async function $$hoist_0_none() {
         "use cache";
@@ -520,6 +520,69 @@ export async function kv() {
     `)
   })
 
+  it('reports source parameters and generated bound argument count', async () => {
+    const input = `
+async function noCapture(first = 1, { second }) {
+  "use server";
+}
+
+function outer(firstCapture, secondCapture) {
+  async function fixed(first = 1, { second }) {
+    "use server";
+    return [firstCapture, secondCapture, first, second];
+  }
+  async function rest(first, ...remaining) {
+    "use server";
+    return [firstCapture, first, remaining];
+  }
+}
+`
+    const ast = await parseAstAsync(input)
+
+    function collect(encoded: boolean) {
+      const metadata: unknown[] = []
+      transformHoistInlineDirective(input, ast, {
+        directive: 'use server',
+        runtime: (value, _name, meta) => {
+          metadata.push(meta)
+          return value
+        },
+        encode: encoded ? (value) => `encode(${value})` : undefined,
+        decode: encoded ? (value) => `decode(${value})` : undefined,
+      })
+      return metadata.map(({ directiveMatch: _, ...meta }: any) => meta)
+    }
+
+    expect(collect(false)).toEqual([
+      {
+        parameters: { count: 2, hasRest: false },
+        boundArgumentCount: 0,
+      },
+      {
+        parameters: { count: 2, hasRest: false },
+        boundArgumentCount: 2,
+      },
+      {
+        parameters: { count: 2, hasRest: true },
+        boundArgumentCount: 1,
+      },
+    ])
+    expect(collect(true)).toEqual([
+      {
+        parameters: { count: 2, hasRest: false },
+        boundArgumentCount: 0,
+      },
+      {
+        parameters: { count: 2, hasRest: false },
+        boundArgumentCount: 1,
+      },
+      {
+        parameters: { count: 2, hasRest: true },
+        boundArgumentCount: 1,
+      },
+    ])
+  })
+
   it('no ending new line', async () => {
     const input = `\
 export async function test() {

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

@@ -9,8 +9,16 @@ import type {
 import { walk } from 'estree-walker'
 import MagicString from 'magic-string'
 import type { ESTree } from 'vite'
+import type { FunctionParameters } from './module-export-scan'
 import { buildScopeTree, type ScopeTree } from './scope'
 
+export type InlineDirectiveMeta = {
+  directiveMatch: RegExpMatchArray
+  parameters: FunctionParameters
+  /** Number of generated leading arguments bound before source arguments. */
+  boundArgumentCount: number
+}
+
 /**
  * Turns an inline directive function into a module-level registered function.
  * Conceptually:
@@ -90,11 +98,7 @@ export function transformHoistInlineDirective(
     rejectNonAsyncFunction,
     ...options
   }: {
-    runtime: (
-      value: string,
-      name: string,
-      meta: { directiveMatch: RegExpMatchArray },
-    ) => string
+    runtime: (value: string, name: string, meta: InlineDirectiveMeta) => string
     directive: string | RegExp
     rejectNonAsyncFunction?: boolean
     encode?: (value: string) => string
@@ -166,6 +170,14 @@ export function transformHoistInlineDirective(
         // function. At the original call site, registration below binds the
         // corresponding values in the same order.
         const bindVars = getBindVars(node, scopeTree)
+        const parameters: FunctionParameters = {
+          count: node.params.length,
+          hasRest: node.params.some(
+            (parameter) => parameter.type === 'RestElement',
+          ),
+        }
+        const boundArgumentCount =
+          bindVars.length > 0 ? (options.encode ? 1 : bindVars.length) : 0
         let newParams = [
           ...bindVars.map((b) => b.root),
           ...node.params.map((n) => input.slice(n.start, n.end)),
@@ -206,7 +218,7 @@ export function transformHoistInlineDirective(
         const runtimeCode = `/* #__PURE__ */ ${runtime(
           implementationName,
           newName,
-          { directiveMatch: match },
+          { directiveMatch: match, parameters, boundArgumentCount },
         )}`
         if (options.hoistRuntime) {
           runtimeHoists.push(