Add support for client components

package.json

@@ -4,6 +4,7 @@
   "license": "MIT",
   "dependencies": {
     "@koa/router": "^12.0.1",
+    "esbuild": "^0.20.2",
     "koa": "^2.15.1",
     "react": "^18.3.0-canary-a870b2d54-20240314",
     "react-dom": "^18.3.0-canary-a870b2d54-20240314",

src/app/favorite.jsx

@@ -0,0 +1,28 @@
+'use client';
+
+import React from 'react';
+
+export default function Favorite({ name }) {
+  const [isFavorite, setIsFavorite] = React.useState(null);
+
+  React.useEffect(() => {
+    const isFavorite = localStorage.getItem(`${name}:favorite`) === 'true';
+
+    setIsFavorite(isFavorite);
+  }, [name]);
+
+  const onClick = () => {
+    setIsFavorite(!isFavorite);
+    localStorage.setItem(`${name}:favorite`, String(!isFavorite));
+  };
+
+  return (
+    <button disabled={isFavorite == null} onClick={onClick}>
+      {isFavorite == null
+        ? '…'
+        : isFavorite
+        ? 'Remove Favorite'
+        : 'Add Favorite'}
+    </button>
+  );
+}

src/app/pokemon.jsx

@@ -1,4 +1,5 @@
 import React from 'react';
+import Favorite from './favorite.jsx';
 
 export default async function Pokemon({ query }) {
   const data = await fetch(
@@ -11,6 +12,9 @@ export default async function Pokemon({ query }) {
       <h1>
         {query.name} ({data.types.map((type) => type.type.name).join(', ')})
       </h1>
+      <div>
+        <Favorite name={query.name} />
+      </div>
       <img src={data.sprites.front_default} />
     </div>
   );

src/client.js

@@ -24,6 +24,8 @@ function parseJSX(key, value) {
     return Symbol.for('react.element');
   } else if (typeof value === 'string' && value.startsWith('$$')) {
     return value.slice(1);
+  } else if (value?.$$typeof === '$RE_M') {
+    return window.__CLIENT_MODULES__[value.filename];
   } else {
     return value;
   }

src/loader.js

@@ -0,0 +1,27 @@
+import { relative } from 'path';
+import { fileURLToPath } from 'url';
+import { renderToString } from 'react-dom/server';
+
+export async function load(url, context, defaultLoad) {
+  const result = await defaultLoad(url, context, defaultLoad);
+
+  if (result.format === 'module' && !url.includes('?original')) {
+    const code = result.source.toString();
+
+    if (/['"]use client['"]/.test(code)) {
+      const source = `
+        export default {
+          $$typeof: Symbol.for('react.client.reference'),
+          name: 'default',
+          filename: ${JSON.stringify(
+            relative(import.meta.dirname, fileURLToPath(url))
+          )},
+        }
+      `;
+
+      return { source, format: 'module' };
+    }
+  }
+
+  return result;
+}

src/server.jsx

@@ -1,22 +1,31 @@
 import Router from '@koa/router';
+import { transform } from 'esbuild';
 import Koa from 'koa';
-import { createReadStream } from 'node:fs';
-import { stat } from 'node:fs/promises';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { readFile, readdir, stat } from 'node:fs/promises';
+import { register } from 'node:module';
 import { join } from 'node:path';
 import { renderToString } from 'react-dom/server';
 
+register('./loader.js', import.meta.url);
+
 const html = String.raw;
 
 const PORT = 5172;
 
 const app = new Koa();
 const router = new Router();
 
-router.get('/client.js', (ctx) => {
-  const stream = createReadStream(join(import.meta.dirname, ctx.path));
+router.get('/(.*).(js|jsx)', async (ctx) => {
+  const content = await readFile(join(import.meta.dirname, ctx.path), 'utf8');
+  const transformed = await transform(content, {
+    loader: 'jsx',
+    format: 'esm',
+    target: 'es2020',
+  });
 
   ctx.type = 'text/javascript';
-  ctx.body = stream;
+  ctx.body = transformed.code;
 });
 
 router.get('/:path*', async (ctx) => {
@@ -56,9 +65,28 @@ async function respond(ctx, jsx) {
     return;
   }
 
+  const files = await readdir(join(import.meta.dirname, 'app'));
+  const imports = await Promise.all(
+    files.map((file) => import(join(import.meta.dirname, 'app', file)))
+  );
+
+  const modules = imports.reduce((acc, mod, i) => {
+    if (mod.default?.$$typeof === Symbol.for('react.client.reference')) {
+      acc += `
+        import $${i} from './${mod.default.filename}';
+
+        window.__CLIENT_MODULES__['${mod.default.filename}'] = $${i};
+      `;
+    }
+
+    return acc;
+  }, `window.__CLIENT_MODULES__ = {};`);
+
   const clientJSXString = JSON.stringify(clientJSX, stringifyJSX);
   const clientHtml = html`
-    ${renderToString(clientJSX)}
+    ${renderToString(
+      await storage.run({ ssr: true }, () => renderJSXToClientJSX(jsx))
+    )}
 
     <script>
       window.__INITIAL_CLIENT_JSX_STRING__ = ${JSON.stringify(
@@ -73,6 +101,9 @@ async function respond(ctx, jsx) {
         }
       }
     </script>
+    <script type="module">
+      ${modules};
+    </script>
     <script type="module" src="/client.js"></script>
   `;
 
@@ -83,13 +114,17 @@ async function respond(ctx, jsx) {
 function stringifyJSX(key, value) {
   if (value === Symbol.for('react.element')) {
     return '$RE';
+  } else if (value === Symbol.for('react.client.reference')) {
+    return '$RE_M';
   } else if (typeof value === 'string' && value.startsWith('$')) {
     return '$' + value;
   } else {
     return value;
   }
 }
 
+const storage = new AsyncLocalStorage();
+
 async function renderJSXToClientJSX(jsx) {
   if (
     typeof jsx === 'string' ||
@@ -113,6 +148,16 @@ async function renderJSXToClientJSX(jsx) {
         const returnedJsx = await Component(props);
 
         return renderJSXToClientJSX(returnedJsx);
+      } else if (jsx.type.$$typeof === Symbol.for('react.client.reference')) {
+        const ssr = storage.getStore()?.ssr;
+
+        if (!ssr) {
+          return jsx;
+        }
+
+        const m = await import(`./${jsx.type.filename}?original`);
+
+        return { ...jsx, type: m[jsx.type.name] };
       } else {
         throw new Error('Not implemented.');
       }