From a2ac21df2849dbd34e76ae7cd9c99ab871fc6b2b Mon Sep 17 00:00:00 2001
From: Mark Erikson <mark@isquaredsoftware.com>
Date: Sun, 2 Mar 2025 23:06:59 -0500
Subject: [PATCH] Fix infinite query subscription/state hooks

---
 packages/toolkit/src/query/react/buildHooks.ts     |  4 ++--
 .../src/query/tests/infiniteQueries.test-d.ts      | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/packages/toolkit/src/query/react/buildHooks.ts b/packages/toolkit/src/query/react/buildHooks.ts
index 1f08eab47b..4c8ca51bef 100644
--- a/packages/toolkit/src/query/react/buildHooks.ts
+++ b/packages/toolkit/src/query/react/buildHooks.ts
@@ -940,7 +940,7 @@ export type UseInfiniteQuery<
 export type UseInfiniteQueryState<
   D extends InfiniteQueryDefinition<any, any, any, any, any>,
 > = <R extends Record<string, any> = UseInfiniteQueryStateDefaultResult<D>>(
-  arg: QueryArgFrom<D> | SkipToken,
+  arg: InfiniteQueryArgFrom<D> | SkipToken,
   options?: UseInfiniteQueryStateOptions<D, R>,
 ) => UseInfiniteQueryStateResult<D, R>
 
@@ -977,7 +977,7 @@ export type TypedUseInfiniteQueryState<
 export type UseInfiniteQuerySubscription<
   D extends InfiniteQueryDefinition<any, any, any, any, any>,
 > = (
-  arg: QueryArgFrom<D> | SkipToken,
+  arg: InfiniteQueryArgFrom<D> | SkipToken,
   options?: UseInfiniteQuerySubscriptionOptions<D>,
 ) => UseInfiniteQuerySubscriptionResult<D>
 
diff --git a/packages/toolkit/src/query/tests/infiniteQueries.test-d.ts b/packages/toolkit/src/query/tests/infiniteQueries.test-d.ts
index 132ebd12c8..2ada1aef09 100644
--- a/packages/toolkit/src/query/tests/infiniteQueries.test-d.ts
+++ b/packages/toolkit/src/query/tests/infiniteQueries.test-d.ts
@@ -75,6 +75,20 @@ describe('Infinite queries', () => {
 
     expectTypeOf(pokemonApi.useGetInfinitePokemonInfiniteQuery).toBeFunction()
 
+    expectTypeOf(pokemonApi.endpoints.getInfinitePokemon.useInfiniteQuery)
+      .parameter(0)
+      .toEqualTypeOf<string | typeof skipToken>()
+
+    expectTypeOf(pokemonApi.endpoints.getInfinitePokemon.useInfiniteQueryState)
+      .parameter(0)
+      .toEqualTypeOf<string | typeof skipToken>()
+
+    expectTypeOf(
+      pokemonApi.endpoints.getInfinitePokemon.useInfiniteQuerySubscription,
+    )
+      .parameter(0)
+      .toEqualTypeOf<string | typeof skipToken>()
+
     const slice = createSlice({
       name: 'pokemon',
       initialState: {} as { data: Pokemon[] },