fix(query-core, vue-query): fix type inference in setQueryData when used in functions with explicit return types (#9030)

* fix(query-core): fix setQueryData type inference in typed contexts

* test(query-core): add test case for setQueryData type inference in functions

* fix(vue-query): fix setQueryData type inference in typed contexts

* test(vue-query): add test case for setQueryData type inference in functions

Author: custardcream98

packages/query-core/src/__tests__/queryClient.test-d.tsx

@@ -113,6 +113,19 @@ describe('setQueryData', () => {
 
     expectTypeOf(data).toEqualTypeOf<string | undefined>()
   })
+
+  it('should preserve updater parameter type inference when used in functions with explicit return types', () => {
+    const queryKey = ['key'] as DataTag<Array<string>, number>
+    const queryClient = new QueryClient()
+
+    // Simulate usage inside a function with explicit return type
+    // The outer function returns 'unknown' but this shouldn't affect the updater's type inference
+    ;(() =>
+      queryClient.setQueryData(queryKey, (data) => {
+        expectTypeOf(data).toEqualTypeOf<number | undefined>()
+        return data
+      })) satisfies () => unknown
+  })
 })
 
 describe('getQueryState', () => {
@@ -590,3 +603,28 @@ describe('resetQueries', () => {
     })
   })
 })
+type SuccessCallback = () => unknown
+it('should infer types correctly with expression body arrow functions', () => {
+  const queryKey = ['key'] as DataTag<Array<string>, number>
+  const queryClient = new QueryClient()
+
+  // @ts-expect-error
+  const callbackTest: SuccessCallback = () =>
+    queryClient.setQueryData(queryKey, (data) => {
+      expectTypeOf(data).toEqualTypeOf<number | undefined>()
+      return data
+    })
+})
+
+it('should infer types correctly with block body arrow functions', () => {
+  const queryKey = ['key'] as DataTag<Array<string>, number>
+  const queryClient = new QueryClient()
+
+  // @ts-expect-error
+  const callbackTest2: SuccessCallback = () => {
+    queryClient.setQueryData(queryKey, (data) => {
+      expectTypeOf(data).toEqualTypeOf<number | undefined>()
+      return data
+    })
+  }
+})

packages/query-core/src/queryClient.ts

@@ -185,7 +185,7 @@ export class QueryClient {
       NoInfer<TInferredQueryFnData> | undefined
     >,
     options?: SetDataOptions,
-  ): TInferredQueryFnData | undefined {
+  ): NoInfer<TInferredQueryFnData> | undefined {
     const defaultedOptions = this.defaultQueryOptions<
       any,
       any,

packages/vue-query/src/__tests__/queryClient.test-d.ts

@@ -96,6 +96,19 @@ describe('setQueryData', () => {
 
     expectTypeOf(data).toEqualTypeOf<string | undefined>()
   })
+
+  it('should preserve updater parameter type inference when used in functions with explicit return types', () => {
+    const queryKey = ['key'] as DataTag<Array<string>, number>
+    const queryClient = new QueryClient()
+
+    // Simulate usage inside a function with explicit return type
+    // The outer function returns 'unknown' but this shouldn't affect the updater's type inference
+    ;(() =>
+      queryClient.setQueryData(queryKey, (data) => {
+        expectTypeOf(data).toEqualTypeOf<number | undefined>()
+        return data
+      })) satisfies () => unknown
+  })
 })
 
 describe('fetchInfiniteQuery', () => {

packages/vue-query/src/queryClient.ts

@@ -114,17 +114,17 @@ export class QueryClient extends QC {
       NoInfer<TInferredQueryFnData> | undefined
     >,
     options?: MaybeRefDeep<SetDataOptions>,
-  ): TInferredQueryFnData | undefined
+  ): NoInfer<TInferredQueryFnData> | undefined
   setQueryData<TQueryFnData, TData = NoUnknown<TQueryFnData>>(
     queryKey: MaybeRefDeep<QueryKey>,
     updater: Updater<NoInfer<TData> | undefined, NoInfer<TData> | undefined>,
     options?: MaybeRefDeep<SetDataOptions>,
-  ): TData | undefined
+  ): NoInfer<TData> | undefined
   setQueryData<TData>(
     queryKey: MaybeRefDeep<QueryKey>,
     updater: Updater<TData | undefined, TData | undefined>,
     options: MaybeRefDeep<SetDataOptions> = {},
-  ): TData | undefined {
+  ): NoInfer<TData> | undefined {
     return super.setQueryData(
       cloneDeepUnref(queryKey),
       updater,