Fix HMR issue after patching the client module (#43819)

Co-authored-by: Tim Neutkens <tim@timneutkens.nl>

packages/next/client/app-index.tsx

@@ -29,7 +29,10 @@ __webpack_require__.u = (chunkId: any) => {
 // Ignore the module ID transform in client.
 // eslint-disable-next-line no-undef
 // @ts-expect-error TODO: fix type
-self.__next_require__ = __webpack_require__
+self.__next_require__ = (id: string) => {
+  const modId = id.replace(/\?.+$/, '')
+  return __webpack_require__(modId)
+}
 
 // eslint-disable-next-line no-undef
 ;(self as any).__next_chunk_load__ = (chunk: string) => {

packages/next/server/app-render.tsx

@@ -1598,11 +1598,35 @@ export async function renderToHTMLOrFlight(
         ).slice(1),
       ]
 
+      const serverComponentManifestWithHMR = dev
+        ? new Proxy(serverComponentManifest, {
+            get: (target, prop) => {
+              if (
+                typeof prop === 'string' &&
+                !prop.startsWith('_') &&
+                target[prop]
+              ) {
+                // Attach TS (timestamp) query param to IDs to get rid of flight client's module cache on HMR.
+                const namedExports: any = {}
+                const ts = Date.now()
+                for (let key in target[prop]) {
+                  namedExports[key] = {
+                    ...target[prop][key],
+                    id: `${target[prop][key].id}?ts=${ts}`,
+                  }
+                }
+                return namedExports
+              }
+              return target[prop]
+            },
+          })
+        : serverComponentManifest
+
       // For app dir, use the bundled version of Fizz renderer (renderToReadableStream)
       // which contains the subset React.
       const readable = ComponentMod.renderToReadableStream(
         flightData,
-        serverComponentManifest,
+        serverComponentManifestWithHMR,
         {
           context: serverContexts,
           onError: flightDataRendererErrorHandler,

test/e2e/app-dir/index.test.ts

@@ -441,7 +441,7 @@ describe('app dir', () => {
         }
       })
 
-      it.skip('should HMR correctly when changing the component type', async () => {
+      it('should HMR correctly when changing the component type', async () => {
         const filePath = 'app/dashboard/page/page.jsx'
         const origContent = await next.readFile(filePath)
 
@@ -466,13 +466,12 @@ describe('app dir', () => {
           )
 
           // Change to client component
-          await new Promise((resolve) => setTimeout(resolve, 1000))
           await next.patchFile(
             filePath,
             origContent
               .replace("// 'use client'", "'use client'")
               .replace(
-                'hello dashboard/page in server component!',
+                'hello dashboard/page!',
                 'hello dashboard/page in client component!'
               )
           )
@@ -482,18 +481,31 @@ describe('app dir', () => {
           )
 
           // Change back to server component
+          await next.patchFile(
+            filePath,
+            origContent.replace(
+              'hello dashboard/page!',
+              'hello dashboard/page in server component2!'
+            )
+          )
+          await check(
+            () => browser.elementByCss('p').text(),
+            /in server component2/
+          )
+
+          // Change to client component again
           await next.patchFile(
             filePath,
             origContent
-              .replace("'use client'", "// 'use client'")
+              .replace("// 'use client'", "'use client'")
               .replace(
-                'hello dashboard/page in client component!',
-                'hello dashboard/page in server component!'
+                'hello dashboard/page!',
+                'hello dashboard/page in client component2!'
               )
           )
           await check(
             () => browser.elementByCss('p').text(),
-            /in server component/
+            /in client component2/
           )
         } finally {
           await next.patchFile(filePath, origContent)