From fedc95969b5d8135bf69dc252dbfd56b42a2fd96 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Fri, 5 Jan 2024 12:04:28 +0000
Subject: [PATCH] Fix provider prop types

---
 .changeset/long-bees-wash.md                  |  5 +++++
 packages/jotai-x/src/createAtomStore.spec.tsx | 12 ++++++++++-
 packages/jotai-x/src/createAtomStore.ts       | 20 +++++++++++++------
 3 files changed, 30 insertions(+), 7 deletions(-)
 create mode 100644 .changeset/long-bees-wash.md

diff --git a/.changeset/long-bees-wash.md b/.changeset/long-bees-wash.md
new file mode 100644
index 0000000..c7f5ee0
--- /dev/null
+++ b/.changeset/long-bees-wash.md
@@ -0,0 +1,5 @@
+---
+'jotai-x': patch
+---
+
+Fix: Provider prop types expect atoms instead of values for stores created with custom atoms
diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index b519f32..5fce47c 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -483,7 +483,7 @@ describe('createAtomStore', () => {
       isCustomAtom: true,
     });
 
-    const { customStore } = createAtomStore(
+    const { customStore, useCustomStore, CustomProvider } = createAtomStore(
       {
         x: createCustomAtom(1),
       },
@@ -496,6 +496,16 @@ describe('createAtomStore', () => {
       const myAtom = customStore.atom.x as CustomAtom<number>;
       expect(myAtom.isCustomAtom).toBe(true);
     });
+
+    it('accepts initial values', () => {
+      const { result } = renderHook(() => useCustomStore().get.x(), {
+        wrapper: ({ children }) => (
+          <CustomProvider x={2}>{children}</CustomProvider>
+        ),
+      });
+
+      expect(result.current).toBe(2);
+    });
   });
 
   describe('arbitrary atom accessors', () => {
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 4df6c02..7ea8217 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -41,6 +41,12 @@ type StoreAtomsWithoutExtend<T> = {
   [K in keyof T]: T[K] extends Atom<any> ? T[K] : SimpleWritableAtom<T[K]>;
 };
 
+type ValueTypesForAtoms<T> = {
+  [K in keyof T]: T[K] extends Atom<infer V> ? V : never;
+};
+
+type StoreInitialValues<T> = ValueTypesForAtoms<StoreAtomsWithoutExtend<T>>;
+
 type StoreAtoms<T, E> = StoreAtomsWithoutExtend<T> & E;
 
 type FilterWritableAtoms<T> = {
@@ -108,7 +114,9 @@ export type AtomStoreApi<
 > = {
   name: N;
 } & {
-  [key in keyof Record<NameProvider<N>, object>]: FC<ProviderProps<T>>;
+  [key in keyof Record<NameProvider<N>, object>]: FC<
+    ProviderProps<StoreInitialValues<T>>
+  >;
 } & {
   [key in keyof Record<NameStore<N>, object>]: StoreApi<T, E, N>;
 } & {
@@ -185,6 +193,7 @@ export const createAtomStore = <
   type MyStoreAtomsWithoutExtend = StoreAtomsWithoutExtend<T>;
   type MyWritableStoreAtomsWithoutExtend =
     FilterWritableAtoms<MyStoreAtomsWithoutExtend>;
+  type MyStoreInitialValues = StoreInitialValues<T>;
 
   const providerIndex = getProviderIndex(name) as NameProvider<N>;
   const useStoreIndex = getUseStoreIndex(name) as UseNameStore<N>;
@@ -278,11 +287,10 @@ export const createAtomStore = <
     }
   }
 
-  const Provider: FC<ProviderProps<T>> = createAtomProvider(
-    name,
-    writableAtomsWithoutExtend,
-    { effect }
-  );
+  const Provider: FC<ProviderProps<MyStoreInitialValues>> = createAtomProvider<
+    MyStoreInitialValues,
+    N
+  >(name, writableAtomsWithoutExtend, { effect });
 
   const storeApi: StoreApi<T, E, N> = {
     atom: atoms,