fix: React context loss in a monorepo setup

e2e/fixtures/monorepo/packages/context-library/package.json

@@ -0,0 +1,16 @@
+{
+  "name": "context-library",
+  "type": "module",
+  "version": "0.0.0",
+  "description": "A library with React Context",
+  "exports": {
+    "./entry-point": {
+      "types": "./src/index.d.ts",
+      "import": "./src/index.js"
+    }
+  },
+  "private": true,
+  "peerDependencies": {
+    "react": ">=18.0.0"
+  }
+}

e2e/fixtures/monorepo/packages/context-library/src/context-consumer.js

@@ -0,0 +1,27 @@
+'use client';
+
+// eslint-disable-next-line import/no-unresolved
+import { createElement, useContext, useEffect, useState } from 'react';
+
+// Do not add '.js' extension to reproduce the issue
+// https://github.com/dai-shi/waku/pull/1162
+import { Context } from './context-provider';
+
+export const ContextConsumer = () => {
+  const [mounted, setMounted] = useState(false);
+  useEffect(() => {
+    setMounted(true);
+  }, []);
+  const value = useContext(Context);
+  return createElement(
+    'div',
+    null,
+    mounted &&
+      createElement(
+        'div',
+        { 'data-testid': 'context-consumer-mounted' },
+        'true',
+      ),
+    createElement('div', { 'data-testid': 'context-consumer-value' }, value),
+  );
+};

e2e/fixtures/monorepo/packages/context-library/src/context-provider.js

@@ -0,0 +1,10 @@
+'use client';
+
+// eslint-disable-next-line import/no-unresolved
+import { createContext, createElement } from 'react';
+
+export const Context = createContext('original');
+
+export const ContextProvider = ({ children }) => {
+  return createElement(Context.Provider, { value: 'provider value' }, children);
+};

e2e/fixtures/monorepo/packages/context-library/src/index.d.ts

@@ -0,0 +1,6 @@
+import type { ReactElement, ReactNode } from 'react';
+
+export declare function ContextProvider(propx: {
+  children: ReactNode;
+}): ReactElement;
+export declare function ContextConsumer(): ReactElement;

e2e/fixtures/monorepo/packages/context-library/src/index.js

@@ -0,0 +1,2 @@
+export { ContextProvider } from './context-provider';
+export { ContextConsumer } from './context-consumer';

e2e/fixtures/monorepo/packages/waku-project/package.json

@@ -9,6 +9,7 @@
     "start": "waku start"
   },
   "dependencies": {
+    "context-library": "0.0.0",
     "dummy-library": "0.0.0",
     "react": "19.0.0",
     "react-dom": "19.0.0",

e2e/fixtures/monorepo/packages/waku-project/src/pages/index.tsx

@@ -2,6 +2,9 @@ import { Link } from 'waku';
 // @ts-expect-error no types
 // eslint-disable-next-line import/no-unresolved
 import { Hello } from 'dummy-library/entry-point';
+// @ts-expect-error no types
+// eslint-disable-next-line import/no-unresolved
+import { ContextProvider, ContextConsumer } from 'context-library/entry-point';
 
 import { Counter } from '../components/counter';
 
@@ -20,6 +23,9 @@ export default async function HomePage() {
       <Link to="/about" className="mt-4 inline-block underline">
         About page
       </Link>
+      <ContextProvider>
+        <ContextConsumer />
+      </ContextProvider>
     </div>
   );
 }

e2e/monorepo.spec.ts

@@ -23,6 +23,11 @@ for (const mode of ['DEV', 'PRD'] as const) {
       test('renders the home page', async ({ page }) => {
         await page.goto(`http://localhost:${port}`);
         await expect(page.getByTestId('header')).toHaveText('Waku');
+        // it should show context value from provider correctly
+        await page.waitForSelector('[data-testid="context-consumer-mounted"]');
+        await expect(page.getByTestId('context-consumer-value')).toHaveText(
+          'provider value',
+        );
       });
     });
   }

packages/waku/src/lib/middleware/dev-server-impl.ts

@@ -73,16 +73,15 @@ const createStreamPair = (): [Writable, Promise<ReadableStream | null>] => {
   return [writable, promise];
 };
 
-const hotUpdateCallbackSet = new Set<(payload: HotUpdatePayload) => void>();
-const registerHotUpdateCallback = (fn: (payload: HotUpdatePayload) => void) =>
-  hotUpdateCallbackSet.add(fn);
-const hotUpdateCallback = (payload: HotUpdatePayload) =>
-  hotUpdateCallbackSet.forEach((fn) => fn(payload));
-
 const createMainViteServer = (
   env: Record<string, string>,
   configPromise: ReturnType<typeof resolveConfig>,
+  hotUpdateCallbackSet: Set<(payload: HotUpdatePayload) => void>,
+  resolvedMap: Map<string, string>,
 ) => {
+  const registerHotUpdateCallback = (fn: (payload: HotUpdatePayload) => void) =>
+    hotUpdateCallbackSet.add(fn);
+
   const vitePromise = configPromise.then(async (config) => {
     const vite = await createViteServer(
       extendViteConfig(
@@ -158,6 +157,15 @@ const createMainViteServer = (
       // HACK `external: ['waku']` doesn't do the same
       return import(/* @vite-ignore */ filePathToFileURL(file));
     }
+    {
+      let id = file;
+      while (resolvedMap.has(id)) {
+        id = resolvedMap.get(id)!;
+      }
+      if (!id.startsWith('/')) {
+        return vite.ssrLoadModule(id);
+      }
+    }
     if (file.includes('/node_modules/')) {
       // HACK node_modules should be externalized
       return import(/* @vite-ignore */ filePathToFileURL(file));
@@ -226,7 +234,11 @@ const createMainViteServer = (
 const createRscViteServer = (
   env: Record<string, string>,
   configPromise: ReturnType<typeof resolveConfig>,
+  hotUpdateCallbackSet: Set<(payload: HotUpdatePayload) => void>,
+  resolvedMap: Map<string, string>,
 ) => {
+  const hotUpdateCallback = (payload: HotUpdatePayload) =>
+    hotUpdateCallbackSet.forEach((fn) => fn(payload));
   const dummyServer = new Server(); // FIXME we hope to avoid this hack
 
   const vitePromise = configPromise.then(async (config) => {
@@ -246,7 +258,11 @@ const createRscViteServer = (
               hotUpdateCallback,
             }),
             rscManagedPlugin(config),
-            rscTransformPlugin({ isClient: false, isBuild: false }),
+            rscTransformPlugin({
+              isClient: false,
+              isBuild: false,
+              resolvedMap,
+            }),
             rscDelegatePlugin(hotUpdateCallback),
           ],
           optimizeDeps: {
@@ -346,15 +362,23 @@ export const devServer: Middleware = (options) => {
   (globalThis as any).__WAKU_CLIENT_IMPORT__ = (id: string) =>
     loadServerModuleMain(id);
 
+  const hotUpdateCallbackSet = new Set<(payload: HotUpdatePayload) => void>();
+  const resolvedMap = new Map<string, string>();
+
   const {
     vitePromise,
     loadServerModuleMain,
     transformIndexHtml,
     willBeHandled,
-  } = createMainViteServer(env, configPromise);
+  } = createMainViteServer(
+    env,
+    configPromise,
+    hotUpdateCallbackSet,
+    resolvedMap,
+  );
 
   const { loadServerModuleRsc, loadEntriesDev, resolveClientEntry } =
-    createRscViteServer(env, configPromise);
+    createRscViteServer(env, configPromise, hotUpdateCallbackSet, resolvedMap);
 
   let initialModules: ClonableModuleNode[];
 

packages/waku/src/lib/plugins/vite-plugin-rsc-analyze.ts

@@ -82,6 +82,7 @@ export function rscAnalyzePlugin(
   const rscTransform = rscTransformPlugin({
     isClient: false,
     isBuild: false,
+    resolvedMap: new Map(),
   }).transform;
   return {
     name: 'rsc-analyze-plugin',

packages/waku/src/lib/plugins/vite-plugin-rsc-transform.ts

@@ -2,12 +2,7 @@ import type { Plugin } from 'vite';
 import * as swc from '@swc/core';
 
 import { EXTENSIONS } from '../constants.js';
-import {
-  extname,
-  joinPath,
-  fileURLToFilePath,
-  decodeFilePathFromAbsolute,
-} from '../utils/path.js';
+import { extname, joinPath } from '../utils/path.js';
 import { parseOpts } from '../utils/swc.js';
 
 const collectExportNames = (mod: swc.Module) => {
@@ -620,6 +615,7 @@ export function rscTransformPlugin(
     | {
         isClient: false;
         isBuild: false;
+        resolvedMap: Map<string, string>;
       }
     | {
         isClient: false;
@@ -628,25 +624,12 @@ export function rscTransformPlugin(
         serverEntryFiles: Record<string, string>;
       },
 ): Plugin {
-  let rootDir: string;
-  const resolvedMap = new Map<string, string>();
   const getClientId = (id: string): string => {
     if (opts.isClient) {
       throw new Error('getClientId is only for server');
     }
     if (!opts.isBuild) {
-      // HACK this logic is too heuristic
-      if (id.startsWith(rootDir) && !id.includes('?')) {
-        return id;
-      }
-      const origId = resolvedMap.get(id);
-      if (origId) {
-        if (origId.startsWith('/@fs/') && !origId.includes('?')) {
-          return origId;
-        }
-        return getClientId(origId);
-      }
-      return id;
+      return id.split('?')[0]!;
     }
     for (const [k, v] of Object.entries(opts.clientEntryFiles)) {
       if (v === id) {
@@ -657,18 +640,7 @@ export function rscTransformPlugin(
   };
   const getServerId = (id: string): string => {
     if (!opts.isBuild) {
-      // HACK this logic is too heuristic
-      if (id.startsWith(rootDir) && !id.includes('?')) {
-        return id;
-      }
-      const origId = resolvedMap.get(id);
-      if (origId) {
-        if (origId.startsWith('/@fs/') && !origId.includes('?')) {
-          return origId;
-        }
-        return getServerId(origId);
-      }
-      return id;
+      return id.split('?')[0]!;
     }
     for (const [k, v] of Object.entries(opts.serverEntryFiles)) {
       if (v === id) {
@@ -677,16 +649,9 @@ export function rscTransformPlugin(
     }
     throw new Error('server id not found: ' + id);
   };
-  const wakuDist = joinPath(
-    decodeFilePathFromAbsolute(fileURLToFilePath(import.meta.url)),
-    '../../..',
-  );
   return {
     name: 'rsc-transform-plugin',
     enforce: 'pre', // required for `resolveId`
-    configResolved(config) {
-      rootDir = config.root;
-    },
     async resolveId(id, importer, options) {
       if (opts.isBuild) {
         return;
@@ -699,17 +664,17 @@ export function rscTransformPlugin(
         return (await this.resolve(id.slice('/@fs'.length), importer, options))
           ?.id;
       }
-      const resolved = await this.resolve(id, importer, options);
-      let srcId =
-        importer && (id.startsWith('./') || id.startsWith('../'))
-          ? joinPath(importer.split('?')[0]!, '..', id)
-          : id;
-      if (srcId.startsWith('waku/')) {
-        srcId = wakuDist + srcId.slice('waku'.length) + '.js';
-      }
-      if (resolved && resolved.id !== srcId) {
-        if (!resolvedMap.has(resolved.id)) {
-          resolvedMap.set(resolved.id, srcId);
+      if ('resolvedMap' in opts) {
+        const resolved = await this.resolve(id, importer, options);
+        const srcId =
+          importer && (id.startsWith('./') || id.startsWith('../'))
+            ? joinPath(importer.split('?')[0]!, '..', id)
+            : id;
+        const dstId = resolved && resolved.id.split('?')[0]!;
+        if (dstId && dstId !== srcId) {
+          if (!opts.resolvedMap.has(dstId)) {
+            opts.resolvedMap.set(dstId, srcId);
+          }
         }
       }
     },

packages/waku/tests/vite-plugin-rsc-transform-internals.test.ts

@@ -6,6 +6,7 @@ describe('internal transform function for server environment', () => {
   const { transform } = rscTransformPlugin({
     isClient: false,
     isBuild: false,
+    resolvedMap: new Map(),
   }) as {
     transform(
       code: string,