fix withPageContext

examples/zustand/store.ts

@@ -1,30 +1,40 @@
 export { useStore }
 
-import { create, server, withPageContext } from 'vike-react-zustand'
+import { create, serverOnly, withPageContext } from 'vike-react-zustand'
 
 interface Store {
   counter: number
   setCounter: (value: number) => void
   serverEnv: string
+  url: string
 }
 
-const useStore = create<Store>()(
-  (
-    set,
-    get
-    /* TODO
-  pageContext
-  */
-  ) => ({
+const useStore = withPageContext((pageContext) =>
+  create<Store>()((set, get) => ({
     counter: Math.floor(10000 * Math.random()),
     setCounter(value) {
       set({ counter: value })
     },
+    url: pageContext.urlOriginal,
 
     // the callback only runs on the server,
     // the return value is passed to the client on the initial navigation
-    ...server(() => ({
+    ...serverOnly(() => ({
       serverEnv: process.env.SOME_ENV!
     }))
-  })
+  }))
 )
+
+// This works, too
+// const useStore = create<Store>()((set, get) => ({
+//   counter: Math.floor(10000 * Math.random()),
+//   setCounter(value) {
+//     set({ counter: value })
+//   },
+
+//   // the callback only runs on the server,
+//   // the return value is passed to the client on the initial navigation
+//   ...server(() => ({
+//     serverEnv: process.env.SOME_ENV!
+//   }))
+// }))

packages/vike-react-zustand/src/index.tsx

@@ -1,9 +1,10 @@
-export { create, server }
+export { create, serverOnly, withPageContext }
 
 import { useContext } from 'react'
 import { getContext, setCreateStore } from './renderer/context.js'
 import { create as create_ } from 'zustand'
 import type { StoreMutatorIdentifier, UseBoundStore, Mutate, StoreApi as ZustandStoreApi } from 'zustand'
+import type { PageContext } from 'vike/types'
 
 type Create = {
   <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
@@ -32,14 +33,25 @@ export type StateCreator<
   $$storeMutators?: Mos
 }
 
-const create: Create = ((createState: any) => {
-  return createState ? createImpl(createState) : createImpl
+const create: Create = ((storeCreatorFn: any) => {
+  return storeCreatorFn ? createImpl(storeCreatorFn) : createImpl
 }) as any
 
-function createImpl(createStore: any): any {
-  // @ts-ignore
+function createImpl(storeCreatorFn: any): any {
   setCreateStore((pageContext: any) => {
-    return create_(createStore)
+    // This is called only once per request
+    if (storeCreatorFn._withPageContext) {
+      // storeCreatorFn(pageContext) looks like this:
+      // (pageContext) =>
+      //   create((set, get) => ({
+      //     counter: 123
+      //   }))
+      // create calls createImpl a second time, and it returns useStore.
+      // but we need to pass the original storeCreatorFn(_storeCreatorFn) to create_
+      return create_()(storeCreatorFn(pageContext)._storeCreatorFn)
+    } else {
+      return create_()(storeCreatorFn)
+    }
   })
 
   function useStore(...args: any[]) {
@@ -50,21 +62,19 @@ function createImpl(createStore: any): any {
     return store(...args)
   }
 
+  useStore._storeCreatorFn = storeCreatorFn
   return useStore
 }
 
-function server<T extends Record<string, any>>(fn: () => T) {
+function serverOnly<T extends Record<string, any>>(fn: () => T) {
   if (typeof window === 'undefined') {
     return fn()
   }
   return {} as T
 }
-type StoreAndHook = ReturnType<typeof create>
-function withPageContext<S extends StoreAndHook>(storeCreatorCreatorFn: (pageContext: StoreAndHook) => S) {
-  //@ts-ignore
-  // createImpl._withPageContext_ = storeCreatorFn
-  // const storeCreatorFn = () => {
-  //   storeCreatorCreatorFn(pageContext)
-  //   return createImpl(storeCreatorFn)
-  // }
+
+function withPageContext<S extends ReturnType<typeof create_>>(storeCreatorFn: (pageContext: PageContext) => S): S {
+  const wrappedStoreCreatorFn = (pageContext: any) => storeCreatorFn(pageContext)
+  wrappedStoreCreatorFn._withPageContext = true
+  return createImpl(wrappedStoreCreatorFn)
 }

packages/vike-react-zustand/src/renderer/VikeReactZustandWrapper.tsx

@@ -13,7 +13,7 @@ export default function VikeReactZustandWrapper({ pageContext, children }: VikeR
   const context = getContext()
   assert(context)
   const createStore = getCreateStore()
-  const store = useMemo(() => createStore?.(), [createStore])
+  const store = useMemo(() => createStore?.(pageContext), [createStore])
   if (!store) {
     // Is that the best thing to do?
     return children

packages/vike-react-zustand/src/renderer/context.ts

@@ -1,18 +1,20 @@
+export { getContext, getCreateStore, setCreateStore }
+
 import React, { createContext } from 'react'
 import { getGlobalObject } from '../utils.js'
 import type { create } from 'zustand'
 
 type StoreAndHook = ReturnType<typeof create>
-type CreateStore = () => StoreAndHook & { __hydrated__?: true }
+type CreateStore = (pageContext: any) => StoreAndHook & { __hydrated__?: true }
 
 const globalObject = getGlobalObject('VikeReactZustandContext.ts', {
   createStore: undefined as CreateStore | undefined,
   context: createContext<StoreAndHook | undefined>(undefined)
 })
 
-export const getCreateStore = () => globalObject.createStore
-export const getContext = () => globalObject.context as unknown as React.Context<StoreAndHook | undefined>
+const getContext = () => globalObject.context as unknown as React.Context<StoreAndHook | undefined>
 
-export const setCreateStore = (createStore_: CreateStore) => {
+const getCreateStore = () => globalObject.createStore
+const setCreateStore = (createStore_: CreateStore) => {
   globalObject.createStore = createStore_
 }