feat(rsc): ability to register server function via plugins

packages/plugin-rsc/README.md

@@ -27,6 +27,7 @@ npm create vite@latest -- --template rsc
 
 - [`./examples/basic`](./examples/basic) - Comprehensive showcase of standard RSC features and the primary E2E test fixture.
 - [`./examples/use-cache`](./examples/use-cache) - Minimal cache feature inspired by Next.js's `"use cache"`, built with generic transform and RSC runtime APIs.
+- [`./examples/custom-server-function`](./examples/custom-server-function) - Third-party Server Function directive integration using server reference claims.
 - [`./examples/ssg`](./examples/ssg) - Static site generation with MDX and client components for interactivity.
 - [`./examples/ppr`](./examples/ppr) - Partial prerendering with a reusable static HTML shell and request-time RSC content.
 - [`./examples/no-ssr`](./examples/no-ssr) - RSC application without an SSR environment.

packages/plugin-rsc/e2e/custom-server-function.test.ts

@@ -0,0 +1,108 @@
+import { expect, test, type Page } from '@playwright/test'
+import { useFixture } from './fixture'
+import {
+  expectNoPageError,
+  expectNoReload,
+  testNoJs,
+  waitForHydration,
+} from './helper'
+
+test.describe('dev-custom-server-function', () => {
+  const f = useFixture({
+    root: 'examples/custom-server-function',
+    mode: 'dev',
+  })
+  defineTest(f)
+
+  test('moves an action between custom and built-in ownership', async ({
+    page,
+  }) => {
+    using _ = expectNoPageError(page)
+    await page.goto(f.url())
+    await waitForHydration(page)
+    await using _noReload = await expectNoReload(page)
+
+    const editor = f.createEditor('src/features/mixed-directives/actions.ts')
+    // Switch one export from the custom plugin to the built-in plugin. HMR
+    // must remove the custom claim while preserving the module's other claims.
+    editor.edit((source) =>
+      source
+        .replace(`'use custom-server'`, `'use server'`)
+        .replace(
+          `customLabel = 'Custom'`,
+          `customLabel = 'Custom changed to built-in'`,
+        )
+        .replace('customCount++', 'customCount += 10'),
+    )
+
+    await expect(
+      page.getByRole('button', { name: 'Custom changed to built-in: 0' }),
+    ).toBeVisible()
+    await page.getByRole('button', { name: 'Built-in: 0', exact: true }).click()
+    await expect(
+      page.getByRole('button', { name: 'Built-in: 1', exact: true }),
+    ).toBeVisible()
+    await page
+      .getByRole('button', { name: 'Custom changed to built-in: 0' })
+      .click()
+    await expect(
+      page.getByRole('button', { name: 'Custom changed to built-in: 10' }),
+    ).toBeVisible()
+
+    // Switch the export back and verify neither owner retained stale state.
+    editor.reset()
+    await expect(page.getByRole('button', { name: 'Custom: 0' })).toBeVisible()
+    await page.getByRole('button', { name: 'Built-in: 0', exact: true }).click()
+    await expect(
+      page.getByRole('button', { name: 'Built-in: 1', exact: true }),
+    ).toBeVisible()
+    await page.getByRole('button', { name: 'Custom: 0' }).click()
+    await expect(page.getByRole('button', { name: 'Custom: 1' })).toBeVisible()
+  })
+})
+
+test.describe('build-custom-server-function', () => {
+  const f = useFixture({
+    root: 'examples/custom-server-function',
+    mode: 'build',
+  })
+  defineTest(f)
+})
+
+function defineTest(f: ReturnType<typeof useFixture>) {
+  test('built-in and custom server functions', async ({ page }) => {
+    using _ = expectNoPageError(page)
+    await page.goto(f.url())
+    await waitForHydration(page)
+    await testActions(page)
+  })
+
+  testNoJs('progressive forms', async ({ page }) => {
+    await page.goto(f.url())
+    await testActions(page)
+  })
+
+  async function testActions(page: Page) {
+    // The first two actions are inline functions in one RSC-reachable module.
+    await page.getByRole('button', { name: 'Built-in: 0' }).click()
+    await expect(
+      page.getByRole('button', { name: 'Built-in: 1' }),
+    ).toBeVisible()
+
+    await page.getByRole('button', { name: 'Custom: 0' }).click()
+    await expect(page.getByRole('button', { name: 'Custom: 1' })).toBeVisible()
+
+    // This module is only imported by a Client Component, so its implementation
+    // reaches the RSC build through the aggregated server reference manifest.
+    await page.getByRole('button', { name: 'From client: 0' }).click()
+    await expect(
+      page.getByRole('button', { name: 'From client: 1' }),
+    ).toBeVisible()
+
+    await page.getByRole('button', { name: 'Reset' }).click()
+    await expect(
+      page.getByRole('button', { name: 'Built-in: 0' }),
+    ).toBeVisible()
+    await expect(page.getByRole('button', { name: 'Custom: 0' })).toBeVisible()
+  }
+}

packages/plugin-rsc/examples/browser-mode/vite.config.ts

@@ -216,7 +216,7 @@ function rscBrowserModePlugin(): Plugin[] {
             return `export default {}` // no-op during dev
           }
           let code = ''
-          for (const meta of Object.values(manager.serverReferenceMetaMap)) {
+          for (const meta of manager.serverReferences.metaMap.values()) {
             code += `${JSON.stringify(meta.referenceKey)}: () => import(${JSON.stringify(meta.importId)}),`
           }
           return `export default {${code}}`

packages/plugin-rsc/examples/custom-server-function/README.md

@@ -0,0 +1,29 @@
+# Custom Server Function
+
+This example demonstrates a third-party Vite plugin implementing a custom `"use custom-server"` directive alongside the built-in `"use server"` directive.
+
+## Background
+
+Server Function extensibility has two separate concerns:
+
+- The directive owner transforms its syntax and registers the function with the React runtime.
+- `@vitejs/plugin-rsc` owns bundler-level module identity, graph visibility, manifests, and reference resolution.
+
+The custom plugin in [`custom-server-function-plugin.ts`](./custom-server-function-plugin.ts) transforms `"use custom-server"` and reports its exports as server reference claims. The RSC plugin aggregates those claims with its built-in `"use server"` claim instead of requiring one transform to own the entire module. This keeps custom syntax and metadata policy outside the RSC plugin while preserving a single canonical reference identity for the bundler.
+
+The example separates two import graph shapes under [`src/features`](./src/features):
+
+- `mixed-directives` exports inline built-in and custom Server Functions from one RSC-reachable module.
+- `action-from-client` imports a module-level custom Server Function only from a Client Component. The custom plugin creates its client and SSR proxies, while the server reference manifest brings its implementation into the RSC build.
+
+This is a low-level integration example rather than a proposed high-level Server Function API.
+
+## Current E2E Coverage
+
+[`../../e2e/custom-server-function.test.ts`](../../e2e/custom-server-function.test.ts) verifies in both development and production build modes that:
+
+- a built-in Server Function and a custom Server Function can coexist in one module
+- each function reaches the server and updates the rendered result
+- a custom Server Function that is not statically imported by the RSC entry works through both the SSR and client proxy paths
+- built-in, inline custom, and client-imported custom functions work as progressively enhanced forms without JavaScript
+- in development, an export can move from the custom owner to the built-in owner and back without reloading, losing the built-in claim, or retaining a conflicting stale claim

packages/plugin-rsc/examples/custom-server-function/custom-server-function-plugin.ts

@@ -0,0 +1,103 @@
+import { getPluginApi, type RscPluginManager } from '@vitejs/plugin-rsc'
+import {
+  hasDirective,
+  transformDirectiveProxyExport,
+  transformHoistInlineDirective,
+  transformWrapExport,
+} from '@vitejs/plugin-rsc/transforms'
+import { parseAstAsync, type Plugin } from 'vite'
+
+const directive = 'use custom-server'
+const pluginName = 'example:custom-server-function'
+
+// This intentionally mirrors the built-in "use server" pipeline through the
+// public integration APIs, but uses a separate directive and plugin name and
+// omits export-all expansion and action encryption.
+// TODO: Give the custom directive observable runtime semantics so ownership
+// swaps can assert the active registration path in addition to claim cleanup.
+export function customServerFunctionPlugin(): Plugin {
+  let manager: RscPluginManager
+
+  return {
+    name: pluginName,
+    configResolved(config) {
+      manager = getPluginApi(config)!.manager
+    },
+    async transform(code, id) {
+      const environmentName = this.environment.name
+      if (!code.includes(directive)) {
+        manager.serverReferences.deleteClaim(pluginName, id)
+        return
+      }
+
+      const reference = manager.serverReferences.resolve(id, 'rsc')
+      const ast = (await parseAstAsync(code)) as unknown as Parameters<
+        typeof transformHoistInlineDirective
+      >[1]
+
+      if (environmentName === 'rsc') {
+        const runtime = (value: string, name: string) =>
+          `$$CustomReactServer.registerServerReference(${value}, ${JSON.stringify(reference.referenceKey)}, ${JSON.stringify(name)})`
+        const result = hasDirective(ast.body, directive)
+          ? transformWrapExport(code, ast, {
+              runtime,
+              rejectNonAsyncFunction: true,
+            })
+          : transformHoistInlineDirective(code, ast, {
+              directive,
+              runtime,
+              rejectNonAsyncFunction: true,
+            })
+        if (!result.output.hasChanged()) {
+          manager.serverReferences.deleteClaim(pluginName, id)
+          return
+        }
+
+        manager.serverReferences.replaceClaim(pluginName, id, {
+          ...reference,
+          exportNames: 'names' in result ? result.names : result.exportNames,
+        })
+        result.output.prepend(
+          `import * as $$CustomReactServer from "@vitejs/plugin-rsc/react/rsc/server";\n`,
+        )
+        return {
+          code: result.output.toString(),
+          map: result.output.generateMap({ hires: 'boundary' }),
+        }
+      }
+
+      const result = transformDirectiveProxyExport(ast, {
+        code,
+        directive,
+        rejectNonAsyncFunction: true,
+        runtime: (name) =>
+          `$$CustomReactClient.createServerReference(` +
+          `${JSON.stringify(reference.referenceKey + '#' + name)},` +
+          `$$CustomReactClient.callServer,` +
+          `undefined,` +
+          (this.environment.mode === 'dev'
+            ? `$$CustomReactClient.findSourceMapURL,`
+            : `undefined,`) +
+          `${JSON.stringify(name)})`,
+      })
+      if (!result?.output.hasChanged()) {
+        manager.serverReferences.deleteClaim(pluginName, id)
+        return
+      }
+
+      manager.serverReferences.replaceClaim(pluginName, id, {
+        ...reference,
+        exportNames: result.exportNames,
+      })
+      const runtimeEnvironment =
+        environmentName === 'client' ? 'browser' : 'ssr'
+      result.output.prepend(
+        `import * as $$CustomReactClient from "@vitejs/plugin-rsc/react/${runtimeEnvironment}";\n`,
+      )
+      return {
+        code: result.output.toString(),
+        map: result.output.generateMap({ hires: 'boundary' }),
+      }
+    },
+  }
+}

packages/plugin-rsc/examples/custom-server-function/package.json

@@ -0,0 +1,24 @@
+{
+  "name": "@vitejs/plugin-rsc-examples-custom-server-function",
+  "version": "0.0.0",
+  "private": true,
+  "license": "MIT",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "vite build",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "react": "^19.2.8",
+    "react-dom": "^19.2.8"
+  },
+  "devDependencies": {
+    "@types/react": "^19.2.17",
+    "@types/react-dom": "^19.2.3",
+    "@vitejs/plugin-react": "latest",
+    "@vitejs/plugin-rsc": "latest",
+    "rsc-html-stream": "^0.0.7",
+    "vite": "^8.1.4"
+  }
+}

packages/plugin-rsc/examples/custom-server-function/src/features/action-from-client/action.ts

@@ -0,0 +1,8 @@
+'use custom-server'
+
+// @ts-ignore -- virtualized by @vitejs/plugin-rsc
+import 'server-only'
+
+export async function incrementFromClient(previous: number) {
+  return previous + 1
+}

packages/plugin-rsc/examples/custom-server-function/src/features/action-from-client/client.tsx

@@ -0,0 +1,13 @@
+'use client'
+
+import { useActionState } from 'react'
+import { incrementFromClient } from './action.ts'
+
+export function ActionFromClient() {
+  const [count, action] = useActionState(incrementFromClient, 0)
+  return (
+    <form action={action}>
+      <button>From client: {count}</button>
+    </form>
+  )
+}

packages/plugin-rsc/examples/custom-server-function/src/features/mixed-directives/actions.ts

@@ -0,0 +1,24 @@
+let builtinCount = 0
+let customCount = 0
+
+export const customLabel = 'Custom'
+
+export function getCounts() {
+  return { builtinCount, customCount }
+}
+
+export async function incrementBuiltin() {
+  'use server'
+  builtinCount++
+}
+
+export async function incrementCustom() {
+  'use custom-server'
+  customCount++
+}
+
+export async function resetCounts() {
+  'use server'
+  builtinCount = 0
+  customCount = 0
+}

packages/plugin-rsc/examples/custom-server-function/src/features/mixed-directives/server.tsx

@@ -0,0 +1,26 @@
+import {
+  customLabel,
+  getCounts,
+  incrementBuiltin,
+  incrementCustom,
+  resetCounts,
+} from './actions.ts'
+
+export function MixedDirectives() {
+  const { builtinCount, customCount } = getCounts()
+  return (
+    <>
+      <form action={incrementBuiltin}>
+        <button>Built-in: {builtinCount}</button>
+      </form>
+      <form action={incrementCustom}>
+        <button>
+          {customLabel}: {customCount}
+        </button>
+      </form>
+      <form action={resetCounts}>
+        <button>Reset</button>
+      </form>
+    </>
+  )
+}