diff --git a/examples/zustand/store.ts b/examples/zustand/store.ts
index 6de5f4e9..da7cf051 100644
--- a/examples/zustand/store.ts
+++ b/examples/zustand/store.ts
@@ -1,6 +1,6 @@
 export { useStore }
 
-import { create, server } from 'vike-react-zustand'
+import { create, server, withPageContext } from 'vike-react-zustand'
 
 interface Store {
   counter: number
@@ -8,11 +8,11 @@ interface Store {
   serverEnv: string
 }
 
-const useStore = create<Store>(
+const useStore = create<Store>()(
   (
     set,
     get
-  /* TODO
+    /* TODO
   pageContext
   */
   ) => ({
diff --git a/packages/vike-react-zustand/src/index.tsx b/packages/vike-react-zustand/src/index.tsx
index 9008fa8a..77eb67d7 100644
--- a/packages/vike-react-zustand/src/index.tsx
+++ b/packages/vike-react-zustand/src/index.tsx
@@ -3,9 +3,42 @@ export { create, server }
 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'
 
-function create(createStore: any): any {
-  setCreateStore(() => {
+type Create = {
+  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>
+  ): UseBoundStore<Mutate<ZustandStoreApi<T>, Mos>>
+  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>
+  ) => UseBoundStore<Mutate<ZustandStoreApi<T>, Mos>>
+  /**
+   * @deprecated Use `useStore` hook to bind store
+   */
+  <S extends ZustandStoreApi<unknown>>(store: S): UseBoundStore<S>
+}
+
+type Get<T, K, F> = K extends keyof T ? T[K] : F
+export type StateCreator<
+  T,
+  Mis extends [StoreMutatorIdentifier, unknown][] = [],
+  Mos extends [StoreMutatorIdentifier, unknown][] = [],
+  U = T
+> = ((
+  setState: Get<Mutate<ZustandStoreApi<T>, Mis>, 'setState', never>,
+  getState: Get<Mutate<ZustandStoreApi<T>, Mis>, 'getState', never>,
+  store: Mutate<ZustandStoreApi<T>, Mis>
+) => U) & {
+  $$storeMutators?: Mos
+}
+
+const create: Create = ((createState: any) => {
+  return createState ? createImpl(createState) : createImpl
+}) as any
+
+function createImpl(createStore: any): any {
+  // @ts-ignore
+  setCreateStore((pageContext: any) => {
     return create_(createStore)
   })
 
@@ -26,3 +59,12 @@ function server<T extends Record<string, any>>(fn: () => T) {
   }
   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)
+  // }
+}