add rehydrate and store options

src/hydrate-atoms/HydrationBoundary.tsx

@@ -1,16 +1,24 @@
 'use client';
 
 import type { PropsWithChildren } from 'react';
-import type { AnyWritableAtom, InferAtomTuples } from './types.js';
+import type {
+  AnyWritableAtom,
+  HydrateAtomOptions,
+  InferAtomTuples,
+} from './types.js';
 import { useHydrateAtoms } from './use-hydrate-atoms.js';
 
 export function HydrationBoundary<
   T extends (readonly [AnyWritableAtom, ...unknown[]])[],
 >({
   children,
   hydrateAtoms,
-}: PropsWithChildren<{ hydrateAtoms: InferAtomTuples<T> }>) {
-  useHydrateAtoms(hydrateAtoms);
+  options,
+}: PropsWithChildren<{
+  hydrateAtoms: InferAtomTuples<T>;
+  options?: HydrateAtomOptions | undefined;
+}>) {
+  useHydrateAtoms(hydrateAtoms, options);
 
   return <>{children}</>;
 }

src/hydrate-atoms/types.ts

@@ -1,4 +1,5 @@
 import type { WritableAtom } from 'jotai';
+import type { Store } from '../shared.js';
 
 export type AnyWritableAtom = WritableAtom<unknown, never[], unknown>;
 
@@ -9,3 +10,8 @@ export type InferAtomTuples<T> = {
       : T[K]
     : never;
 };
+
+export type HydrateAtomOptions = {
+  store: Store;
+  rehydrateKey?: string | undefined;
+};

src/hydrate-atoms/use-hydrate-atoms.ts

@@ -1,13 +1,19 @@
 import { atom, useStore } from 'jotai';
-import type { AnyWritableAtom, InferAtomTuples } from './types.js';
+import type {
+  AnyWritableAtom,
+  HydrateAtomOptions,
+  InferAtomTuples,
+} from './types.js';
 import { useCallback, useEffect, useRef } from 'react';
 
 export function useHydrateAtoms<
   T extends (readonly [AnyWritableAtom, ...unknown[]])[],
->(hydrateAtoms: InferAtomTuples<T>) {
+>(hydrateAtoms: InferAtomTuples<T>, options?: HydrateAtomOptions) {
   const isHydratedRef = useRef(false);
-  const lastHydrateAtoms = useRef(hydrateAtoms);
-  const store = useStore();
+  const lastRehydrateKey = useRef(options?.rehydrateKey);
+  const store = useStore(
+    options?.store != null ? { store: options?.store } : undefined,
+  );
 
   const hydrate = useCallback(() => {
     store.set(isHydratingPrimitiveAtom, true);
@@ -23,11 +29,11 @@ export function useHydrateAtoms<
   }
 
   useEffect(() => {
-    if (lastHydrateAtoms.current !== hydrateAtoms) {
-      lastHydrateAtoms.current = hydrateAtoms;
+    if (lastRehydrateKey.current !== options?.rehydrateKey) {
+      lastRehydrateKey.current = options?.rehydrateKey;
       hydrate();
     }
-  }, [hydrate, hydrateAtoms]);
+  }, [hydrate, options?.rehydrateKey]);
 }
 
 const isHydratingPrimitiveAtom = atom(false);

src/shared.ts

@@ -0,0 +1,3 @@
+import type { createStore } from 'jotai';
+
+export type Store = ReturnType<typeof createStore>;

src/ssr-provider/SSRProvider.tsx

@@ -1,10 +1,9 @@
 'use client';
 
-import { Provider, type createStore } from 'jotai';
+import { Provider } from 'jotai';
 import type { PropsWithChildren } from 'react';
 import { useCreateStore } from './use-create-store.js';
-
-type Store = ReturnType<typeof createStore>;
+import type { Store } from '../shared.js';
 
 function SSRProviderWithGivenStore({
   store,