dai-shi · dai-shi · Jan 28, 2025 · Jan 30, 2025 · Jan 30, 2025 · Jan 30, 2025
diff --git a/examples/54_jotai/package.json b/examples/54_jotai/package.json
@@ -0,0 +1,23 @@
+{
+  "name": "54_jotai",
+  "version": "0.1.0",
+  "type": "module",
+  "private": true,
+  "scripts": {
+    "dev": "waku dev",
+    "build": "waku build",
+    "start": "waku start"
+  },
+  "dependencies": {
+    "jotai": "https://pkg.csb.dev/pmndrs/jotai/commit/4bd935a7/jotai",
+    "react": "19.0.0",
+    "react-dom": "19.0.0",
+    "react-server-dom-webpack": "19.0.0",
+    "waku": "0.21.18"
+  },
+  "devDependencies": {
+    "@types/react": "19.0.8",
+    "@types/react-dom": "19.0.3",
+    "typescript": "5.7.3"
+  }
+}
diff --git a/examples/54_jotai/src/components/app.tsx b/examples/54_jotai/src/components/app.tsx
@@ -0,0 +1,43 @@
+import { atom } from 'jotai/vanilla';
+
+import { getStore, Provider } from '../lib/waku-jotai/server';
+import { Counter, countAtom } from './counter';
+
+// server-only atom
+const doubleCountAtom = atom(async (get) => {
+  await new Promise((r) => setTimeout(r, 1000));
+  return get(countAtom) * 2;
+});
+
+const MyApp = ({ name }: { name: string }) => {
+  const store = getStore();
+  const doubleCount = store.get(doubleCountAtom);
+  return (
+    <html>
+      <head>
+        <title>Waku</title>
+      </head>
+      <body>
+        <div
+          style={{ border: '3px red dashed', margin: '1em', padding: '1em' }}
+        >
+          <h1>Hello {name}!!</h1>
+          <h2>(doubleCount={doubleCount})</h2>
+          <h3>This is a server component.</h3>
+          <Counter />
+          <div>{new Date().toISOString()}</div>
+        </div>
+      </body>
+    </html>
+  );
+};
+
+const App = ({ name, rscParams }: { name: string; rscParams: unknown }) => {
+  return (
+    <Provider rscParams={rscParams}>
+      <MyApp name={name} />
+    </Provider>
+  );
+};
+
+export default App;
diff --git a/examples/54_jotai/src/components/counter.tsx b/examples/54_jotai/src/components/counter.tsx
@@ -0,0 +1,24 @@
+'use client';
+
+import { useTransition } from 'react';
+import { unstable_allowServer as allowServer } from 'waku/client';
+import { atom, useAtom } from 'jotai';
+
+export const countAtom = allowServer(atom(1));
+
+export const Counter = () => {
+  const [count, setCount] = useAtom(countAtom);
+  const [isPending, startTransition] = useTransition();
+  const inc = () => {
+    startTransition(() => {
+      setCount((c) => c + 1);
+    });
+  };
+  return (
+    <div style={{ border: '3px blue dashed', margin: '1em', padding: '1em' }}>
+      <p>Count: {count}</p>
+      <button onClick={inc}>Increment</button> {isPending ? 'Pending...' : ''}
+      <h3>This is a client component.</h3>
+    </div>
+  );
+};
diff --git a/examples/54_jotai/src/entries.tsx b/examples/54_jotai/src/entries.tsx
@@ -0,0 +1,60 @@
+import { unstable_defineEntries as defineEntries } from 'waku/minimal/server';
+import { Slot } from 'waku/minimal/client';
+import { unstable_createAsyncIterable as createAsyncIterable } from 'waku/server';
+
+import App from './components/app';
+
+export default defineEntries({
+  handleRequest: async (input, { renderRsc, renderHtml }) => {
+    if (input.type === 'component') {
+      return renderRsc({
+        App: <App name={input.rscPath || 'Waku'} rscParams={input.rscParams} />,
+      });
+    }
+    if (input.type === 'custom' && input.pathname === '/') {
+      return renderHtml(
+        { App: <App name="Waku" rscParams={undefined} /> },
+        <Slot id="App" />,
+        { rscPath: '' },
+      );
+    }
+  },
+  handleBuild: ({
+    // renderRsc,
+    // renderHtml,
+    // rscPath2pathname,
+    unstable_generatePrefetchCode,
+  }) =>
+    createAsyncIterable(async () => {
+      const moduleIds = new Set<string>();
+      const generateHtmlHead = () =>
+        `<script type="module" async>${unstable_generatePrefetchCode(
+          [''],
+          moduleIds,
+        )}</script>`;
+      const tasks = [
+        async () => ({
+          type: 'htmlHead' as const,
+          pathSpec: [],
+          head: generateHtmlHead(),
+        }),
+        // async () => ({
+        //   type: 'file' as const,
+        //   pathname: rscPath2pathname(''),
+        //   body: await renderRsc(
+        //     { App: <App name="Waku" /> },
+        //     { moduleIdCallback: (id) => moduleIds.add(id) },
+        //   ),
+        // }),
+        // async () => ({
+        //   type: 'file' as const,
+        //   pathname: '/',
+        //   body: renderHtml({ App: <App name="Waku" /> }, <Slot id="App" />, {
+        //     rscPath: '',
+        //     htmlHead: generateHtmlHead(),
+        //   }).then(({ body }) => body),
+        // }),
+      ];
+      return tasks;
+    }),
+});
diff --git a/examples/54_jotai/src/lib/waku-jotai/client.tsx b/examples/54_jotai/src/lib/waku-jotai/client.tsx
@@ -0,0 +1,58 @@
+'use client';
+
+import { useEffect, useRef } from 'react';
+import { useRefetch } from 'waku/minimal/client';
+import { atom, useStore } from 'jotai';
+import type { Atom } from 'jotai';
+
+export const SyncAtoms = ({
+  atomsPromise,
+}: {
+  atomsPromise: Promise<Map<Atom<unknown>, string>>;
+}) => {
+  const store = useStore();
+  const refetch = useRefetch();
+  const prevAtomValues = useRef<Map<Atom<unknown>, unknown>>(new Map());
+  useEffect(() => {
+    const controller = new AbortController();
+    // eslint-disable-next-line @typescript-eslint/no-floating-promises
+    atomsPromise.then((atoms) => {
+      if (controller.signal.aborted) {
+        return;
+      }
+      const atomValuesAtom = atom(
+        (get) =>
+          new Map<Atom<unknown>, unknown>(
+            Array.from(atoms).map(([a]) => [a, get(a)]),
+          ),
+      );
+      const callback = (atomValues: Map<Atom<unknown>, unknown>) => {
+        prevAtomValues.current = atomValues;
+        const rscParams = new Map(
+          Array.from(atomValues).map(([a, value]) => [atoms.get(a)!, value]),
+        );
+        // TODO rscPath==='' is hardcoded
+        refetch('', rscParams);
+      };
+      const unsub = store.sub(atomValuesAtom, () => {
+        callback(store.get(atomValuesAtom));
+      });
+      const atomValues = store.get(atomValuesAtom);
+      // HACK check if atom values have already been changed
+      if (
+        Array.from(atomValues).some(([a, value]) =>
+          prevAtomValues.current.has(a)
+            ? prevAtomValues.current.get(a) !== value
+            : 'init' in a && a.init !== value,
+        )
+      ) {
+        callback(atomValues);
+      }
+      controller.signal.addEventListener('abort', () => {
+        unsub();
+      });
+    });
+    return () => controller.abort();
+  }, [store, atomsPromise, refetch]);
+  return null;
+};
diff --git a/examples/54_jotai/src/lib/waku-jotai/server.tsx b/examples/54_jotai/src/lib/waku-jotai/server.tsx
@@ -0,0 +1,91 @@
+import { cache } from 'react';
+import type { ReactNode } from 'react';
+import type { Atom } from 'jotai/vanilla';
+import { INTERNAL_buildStoreRev1 as buildStore } from 'jotai/vanilla/internals';
+import type {
+  INTERNAL_AtomState as AtomState,
+  INTERNAL_AtomStateMap as AtomStateMap,
+} from 'jotai/vanilla/internals';
+
+import { SyncAtoms } from './client';
+
+const CLIENT_REFERENCE_TAG = Symbol.for('react.client.reference');
+
+type ClientReferenceId = string;
+
+const getClientReferenceId = (a: Atom<unknown>) => {
+  if ((a as any)['$$typeof'] === CLIENT_REFERENCE_TAG) {
+    const id: ClientReferenceId = (a as any)['$$id'];
+    return id;
+  }
+  return null;
+};
+
+export const getStore = cache(() => {
+  const clientAtoms = new Map<Atom<unknown>, ClientReferenceId>();
+  const clientAtomValues = new Map<ClientReferenceId, unknown>();
+  const atomStateMap = new Map<Atom<unknown>, AtomState>();
+  const patchedAtomStateMap: AtomStateMap = {
+    get: (a) => atomStateMap.get(a),
+    set: (a, s) => {
+      const id = getClientReferenceId(a);
+      if (id) {
+        clientAtoms.set(a, id);
+        if (clientAtomValues.has(id)) {
+          s.v = clientAtomValues.get(id) as never;
+        }
+      }
+      atomStateMap.set(a, s);
+    },
+  };
+  const store = buildStore(patchedAtomStateMap);
+  const getAtoms = () => clientAtoms;
+  const setAtomValues = (values: Iterable<[ClientReferenceId, unknown]>) => {
+    for (const [id, value] of values) {
+      clientAtomValues.set(id, value);
+    }
+  };
+  const waitForAtoms = async () => {
+    let size: number;
+    do {
+      size = atomStateMap.size;
+      await Promise.all(Array.from(atomStateMap.values()).map((s) => s.v));
+    } while (size !== atomStateMap.size);
+  };
+  return Object.assign(store, {
+    getAtoms,
+    setAtomValues,
+    waitForAtoms,
+  });
+});
+
+export const Provider = ({
+  children,
+  rscParams,
+}: {
+  children: ReactNode;
+  rscParams: unknown;
+}) => {
+  const atomValues = rscParams instanceof Map ? rscParams : new Map();
+  let resolveAtoms: (m: Map<Atom<unknown>, string>) => void;
+  const atomsPromise = new Promise<Map<Atom<unknown>, string>>((r) => {
+    resolveAtoms = r;
+  });
+  const store = getStore();
+  store.setAtomValues(atomValues);
+  setTimeout(() => {
+    store
+      .waitForAtoms()
+      .then(() => {
+        const atoms = store.getAtoms();
+        resolveAtoms(atoms);
+      })
+      .catch(() => {});
+  });
+  return (
+    <>
+      {children}
+      <SyncAtoms atomsPromise={atomsPromise} />
+    </>
+  );
+};
diff --git a/examples/54_jotai/src/main.tsx b/examples/54_jotai/src/main.tsx
@@ -0,0 +1,17 @@
+import { StrictMode } from 'react';
+import { createRoot, hydrateRoot } from 'react-dom/client';
+import { Root, Slot } from 'waku/minimal/client';
+
+const rootElement = (
+  <StrictMode>
+    <Root>
+      <Slot id="App" />
+    </Root>
+  </StrictMode>
+);
+
+if ((globalThis as any).__WAKU_HYDRATE__) {
+  hydrateRoot(document, rootElement);
+} else {
+  createRoot(document as any).render(rootElement);
+}
diff --git a/examples/54_jotai/tsconfig.json b/examples/54_jotai/tsconfig.json
@@ -0,0 +1,15 @@
+{
+  "compilerOptions": {
+    "strict": true,
+    "target": "esnext",
+    "downlevelIteration": true,
+    "esModuleInterop": true,
+    "module": "esnext",
+    "moduleResolution": "bundler",
+    "skipLibCheck": true,
+    "noUncheckedIndexedAccess": true,
+    "exactOptionalPropertyTypes": true,
+    "types": ["react/experimental"],
+    "jsx": "react-jsx"
+  }
+}
diff --git a/packages/waku/src/lib/plugins/vite-plugin-rsc-hmr.ts b/packages/waku/src/lib/plugins/vite-plugin-rsc-hmr.ts
@@ -62,6 +62,10 @@ if (import.meta.hot && !globalThis.__WAKU_HMR_CONFIGURED__) {
       queueMicrotask(() => style.parentElement?.removeChild(style));
     }
   });
+  import.meta.hot.on('vite:invalidate', () => {
+    // FIXME is there a better solution?
+    location.reload();
+  });
 }
 `;