fix(rsc): handle added/removed "use client" during dev (#750)

Author: hi-ogawa

packages/plugin-rsc/e2e/basic.test.ts

@@ -431,6 +431,48 @@ function defineTest(f: Fixture) {
       await page.reload()
       await expect(page.getByText('ok (test-shared)')).toBeVisible()
     })
+
+    test('hmr switch server to client', async ({ page }) => {
+      await page.goto(f.url())
+      await waitForHydration(page)
+      await using _ = await expectNoReload(page)
+
+      await expect(page.getByTestId('test-hmr-switch-server')).toContainText(
+        '(useState: false)',
+      )
+      const editor = f.createEditor('src/routes/hmr-switch/server.tsx')
+      editor.edit((s) => `"use client";\n` + s)
+      await expect(page.getByTestId('test-hmr-switch-server')).toContainText(
+        '(useState: true)',
+      )
+
+      await page.waitForTimeout(100)
+      editor.reset()
+      await expect(page.getByTestId('test-hmr-switch-server')).toContainText(
+        '(useState: false)',
+      )
+    })
+
+    test('hmr switch client to server', async ({ page }) => {
+      await page.goto(f.url())
+      await waitForHydration(page)
+      await using _ = await expectNoReload(page)
+
+      await expect(page.getByTestId('test-hmr-switch-client')).toContainText(
+        '(useState: true)',
+      )
+      const editor = f.createEditor('src/routes/hmr-switch/client.tsx')
+      editor.edit((s) => s.replace(`'use client'`, ''))
+      await expect(page.getByTestId('test-hmr-switch-client')).toContainText(
+        '(useState: false)',
+      )
+
+      await page.waitForTimeout(100)
+      editor.reset()
+      await expect(page.getByTestId('test-hmr-switch-client')).toContainText(
+        '(useState: true)',
+      )
+    })
   })
 
   test('css @js', async ({ page }) => {

packages/plugin-rsc/examples/basic/src/routes/hmr-switch/client.tsx

@@ -0,0 +1,11 @@
+'use client'
+
+import React from 'react'
+
+export function TestHmrSwitchClient() {
+  return (
+    <div data-testid="test-hmr-switch-client">
+      test-hmr-switch-client (useState: {String(!!React.useState)})
+    </div>
+  )
+}

packages/plugin-rsc/examples/basic/src/routes/hmr-switch/server.tsx

@@ -0,0 +1,9 @@
+import React from 'react'
+
+export function TestHmrSwitchServer() {
+  return (
+    <div data-testid="test-hmr-switch-server">
+      test-hmr-switch-server (useState: {String(!!React.useState)})
+    </div>
+  )
+}

packages/plugin-rsc/examples/basic/src/routes/root.tsx

@@ -38,6 +38,8 @@ import { TestHmrSharedAtomic } from './hmr-shared/atomic/server'
 import { TestCssQueries } from './css-queries/server'
 import { TestImportMetaGlob } from './import-meta-glob/server'
 import { TestAssetsServer } from './assets/server'
+import { TestHmrSwitchServer } from './hmr-switch/server'
+import { TestHmrSwitchClient } from './hmr-switch/client'
 
 export function Root(props: { url: URL }) {
   return (
@@ -65,6 +67,8 @@ export function Root(props: { url: URL }) {
         <TestHmrSharedServer />
         <TestHmrSharedClient />
         <TestHmrSharedAtomic />
+        <TestHmrSwitchServer />
+        <TestHmrSwitchClient />
         <TestTemporaryReference />
         <TestServerActionError />
         <TestReplayConsoleLogs url={props.url} />

packages/plugin-rsc/src/plugin.ts

@@ -453,6 +453,22 @@ export default function vitePluginRsc(
         const ids = ctx.modules.map((mod) => mod.id).filter((v) => v !== null)
         if (ids.length === 0) return
 
+        // handle client -> server switch (i.e. "use client" removal)
+        // by eagerly transforming new module on "rsc" environment.
+        if (this.environment.name === 'rsc') {
+          for (const mod of ctx.modules) {
+            if (
+              mod.type === 'js' &&
+              mod.id &&
+              mod.id in manager.clientReferenceMetaMap
+            ) {
+              try {
+                await this.environment.transformRequest(mod.url)
+              } catch {}
+            }
+          }
+        }
+
         // a shared component/module will have `isInsideClientBoundary = false` on `rsc` environment
         // and `isInsideClientBoundary = true` on `client` environment,
         // which means both server hmr and client hmr will be triggered.
@@ -1009,10 +1025,16 @@ function vitePluginUseClient(
       name: 'rsc:use-client',
       async transform(code, id) {
         if (this.environment.name !== serverEnvironmentName) return
-        if (!code.includes('use client')) return
+        if (!code.includes('use client')) {
+          delete manager.clientReferenceMetaMap[id]
+          return
+        }
 
         const ast = await parseAstAsync(code)
-        if (!hasDirective(ast.body, 'use client')) return
+        if (!hasDirective(ast.body, 'use client')) {
+          delete manager.clientReferenceMetaMap[id]
+          return
+        }
 
         let importId: string
         let referenceKey: string