fix: edge case when rsa depends on another rsa (#785)

Co-authored-by: daishi <daishi@axlight.com>

e2e/fixtures/rsc-basic/modules/ai/package.json

@@ -0,0 +1,21 @@
+{
+  "name": "ai",
+  "type": "module",
+  "version": "1.0.0",
+  "description": "Vercel AI mockup",
+  "exports": {
+    "./rsc": {
+      "types": "./src/index.d.ts",
+      "react-server": "./src/server.js",
+      "import": "./src/client.js"
+    }
+  },
+  "devDependencies": {
+    "react-dom": "^18",
+    "react-server-dom-webpack": "18.3.0-canary-eb33bd747-20240312"
+  },
+  "peerDependencies": {
+    "react": "^18 || ^19"
+  },
+  "private": true
+}

e2e/fixtures/rsc-basic/modules/ai/src/client.js

@@ -0,0 +1,8 @@
+'use client';
+import { useActions } from './shared.js';
+
+export { useActions };
+
+export function createAI() {
+  throw new Error('You should not call createAI in the client side');
+}

e2e/fixtures/rsc-basic/modules/ai/src/index.d.ts

@@ -0,0 +1,7 @@
+import type { ReactNode } from 'react';
+
+declare function createAI(
+  actions: Record<string, (...args: any[]) => any>,
+): (props: { children: ReactNode }) => ReactNode;
+
+declare function useActions(): Record<string, any>;

e2e/fixtures/rsc-basic/modules/ai/src/server.js

@@ -0,0 +1,25 @@
+'use server';
+import { InternalProvider } from './shared.js';
+import { jsx } from 'react/jsx-runtime';
+
+async function innerAction({ action }, state, ...args) {
+  'use server';
+  return action(...args);
+}
+
+function wrapAction(action, options) {
+  return innerAction.bind(null, { action, options });
+}
+
+export function createAI(actions) {
+  const wrappedActions = {};
+  for (const name in actions) {
+    wrappedActions[name] = wrapAction(actions[name]);
+  }
+  return function AI(props) {
+    return jsx(InternalProvider, {
+      actions: wrappedActions,
+      children: props.children,
+    });
+  };
+}

e2e/fixtures/rsc-basic/modules/ai/src/shared.js

@@ -0,0 +1,19 @@
+'use client';
+import { createContext, useContext } from 'react';
+import { jsx } from 'react/jsx-runtime';
+
+const ActionContext = createContext(null);
+
+export function useActions() {
+  return useContext(ActionContext);
+}
+
+export function InternalProvider(props) {
+  return jsx('div', {
+    'data-testid': 'ai-internal-provider',
+    children: jsx(ActionContext.Provider, {
+      value: props.actions,
+      children: props.children,
+    }),
+  });
+}

e2e/fixtures/rsc-basic/package.json

@@ -9,6 +9,7 @@
     "start": "waku start"
   },
   "dependencies": {
+    "ai": "link:./modules/ai",
     "react": "19.0.0-rc-3da26163a3-20240704",
     "react-dom": "19.0.0-rc-3da26163a3-20240704",
     "react-server-dom-webpack": "19.0.0-rc-3da26163a3-20240704",

e2e/fixtures/rsc-basic/src/components/App.tsx

@@ -4,6 +4,8 @@
 import { ClientCounter } from './ClientCounter.js';
 import { ServerPing } from './ServerPing/index.js';
 import { ServerBox } from './Box.js';
+import { ServerProvider } from './ServerAction/Server.js';
+import { ClientActionsConsumer } from './ServerAction/Client.js';
 
 const App = ({ name }: { name: string }) => {
   return (
@@ -12,6 +14,9 @@ const App = ({ name }: { name: string }) => {
       <p data-testid="app-name">{name}</p>
       <ClientCounter />
       <ServerPing />
+      <ServerProvider>
+        <ClientActionsConsumer />
+      </ServerProvider>
     </ServerBox>
   );
 };

e2e/fixtures/rsc-basic/src/components/ServerAction/Client.tsx

@@ -0,0 +1,14 @@
+'use client';
+
+import { useActions } from 'ai/rsc';
+import { useEffect } from 'react';
+
+export const ClientActionsConsumer = () => {
+  const actions = useActions();
+  useEffect(() => {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-expect-error
+    globalThis.actions = actions;
+  }, [actions]);
+  return <div>globalThis.actions: {JSON.stringify(Object.keys(actions))}</div>;
+};

e2e/fixtures/rsc-basic/src/components/ServerAction/Server.tsx

@@ -0,0 +1,13 @@
+import type { ReactNode } from 'react';
+import { createAI } from 'ai/rsc';
+
+const AI = createAI({
+  foo: async () => {
+    'use server';
+    return 0;
+  },
+});
+
+export function ServerProvider({ children }: { children: ReactNode }) {
+  return <AI>{children}</AI>;
+}

e2e/rsc-basic.spec.ts

@@ -81,5 +81,18 @@ for (const { build, command } of commands) {
         page.getByTestId('server-ping').getByTestId('counter'),
       ).toHaveText('2');
     });
+
+    test('server action', async ({ page }) => {
+      await page.goto(`http://localhost:${port}/`);
+      await expect(page.getByTestId('app-name')).toHaveText('Waku');
+      await expect(page.getByTestId('ai-internal-provider')).toHaveText(
+        'globalThis.actions: ["foo"]',
+      );
+      const result = await page.evaluate(() => {
+        // @ts-expect-error no types
+        return globalThis.actions.foo();
+      });
+      expect(result).toBe(0);
+    });
   });
 }

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

@@ -309,6 +309,9 @@ export const devServer: Middleware = (options) => {
   (globalThis as any).__WAKU_PRIVATE_ENV__ = options.env || {};
   const configPromise = resolveConfig(options.config);
 
+  (globalThis as any).__WAKU_HACK_IMPORT__ = async (id: string) =>
+    loadServerFileRsc(id);
+
   const {
     vitePromise,
     loadServerFileMain,

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

@@ -21,7 +21,7 @@ globalThis.__WAKU_${type}_CHUNK_LOAD__ ||= (
   if (!globalThis.__WAKU_${type}_MODULE_LOADING__.has(id)) {
     globalThis.__WAKU_${type}_MODULE_LOADING__.set(
       id,
-      customImport(id).then((m) => {
+      (customImport ? customImport(id) : import(id)).then((m) => {
         globalThis.__WAKU_${type}_MODULE_CACHE__.set(id, m);
       })
     );

packages/waku/src/lib/renderers/rsc-renderer.ts

@@ -92,7 +92,7 @@ export async function renderRsc(
     ),
   ]);
 
-  const bundlerConfig = new Proxy(
+  const clientBundlerConfig = new Proxy(
     {},
     {
       get(_target, encodedId: string) {
@@ -104,6 +104,22 @@ export async function renderRsc(
     },
   );
 
+  const serverBundlerConfig = new Proxy(
+    {},
+    {
+      get(_target, encodedId: string) {
+        const [fileId, name] = encodedId.split('#') as [string, string];
+        const id = filePathToFileURL(fileId);
+        if (fileId.startsWith('@id/assets/')) {
+          const id = '.' + fileId.slice('@id'.length);
+          return { id, chunks: [id], name, async: true };
+        } else {
+          return { id, chunks: [id], name, async: true };
+        }
+      },
+    },
+  );
+
   const renderWithContext = async (
     context: Record<string, unknown> | undefined,
     input: string,
@@ -128,7 +144,7 @@ export async function renderRsc(
       if (Object.keys(elements).some((key) => key.startsWith('_'))) {
         throw new Error('"_" prefix is reserved');
       }
-      return renderToReadableStream(elements, bundlerConfig, {
+      return renderToReadableStream(elements, clientBundlerConfig, {
         onError,
       });
     });
@@ -168,7 +184,7 @@ export async function renderRsc(
       }
       return renderToReadableStream(
         { ...elements, _value: actionValue },
-        bundlerConfig,
+        clientBundlerConfig,
         {
           onError,
         },
@@ -189,9 +205,9 @@ export async function renderRsc(
     ) {
       // XXX This doesn't support streaming unlike busboy
       const formData = parseFormData(bodyStr, contentType);
-      args = await decodeReply(formData);
+      args = await decodeReply(formData, serverBundlerConfig);
     } else if (bodyStr) {
-      args = await decodeReply(bodyStr);
+      args = await decodeReply(bodyStr, serverBundlerConfig);
     }
     const [fileId, name] = rsfId.split('#') as [string, string];
     let mod: any;