feat(query-core): add mutation context to all mutation callbacks

Author: joshuaellis

packages/query-core/src/mutation.ts

@@ -175,11 +175,11 @@ export class Mutation<
       this.#dispatch({ type: 'continue' })
     }
 
-    const mutationFnContext: MutationFunctionContext = {
+    const mutationFnContext = {
       client: this.#client,
       meta: this.options.meta,
       mutationKey: this.options.mutationKey,
-    }
+    } satisfies MutationFunctionContext
 
     this.#retryer = createRetryer({
       fn: () => {
@@ -239,7 +239,12 @@ export class Mutation<
         this as Mutation<unknown, unknown, unknown, unknown>,
       )
 
-      await this.options.onSuccess?.(data, variables, this.state.scope!)
+      await this.options.onSuccess?.(
+        data,
+        variables,
+        this.state.scope,
+        mutationFnContext,
+      )
 
       // Notify cache callback
       await this.#mutationCache.config.onSettled?.(
@@ -250,7 +255,13 @@ export class Mutation<
         this as Mutation<unknown, unknown, unknown, unknown>,
       )
 
-      await this.options.onSettled?.(data, null, variables, this.state.scope)
+      await this.options.onSettled?.(
+        data,
+        null,
+        variables,
+        this.state.scope,
+        mutationFnContext,
+      )
 
       this.#dispatch({ type: 'success', data })
       return data
@@ -268,6 +279,7 @@ export class Mutation<
           error as TError,
           variables,
           this.state.scope,
+          mutationFnContext,
         )
 
         // Notify cache callback
@@ -284,6 +296,7 @@ export class Mutation<
           error as TError,
           variables,
           this.state.scope,
+          mutationFnContext,
         )
         throw error
       } finally {

packages/query-core/src/mutationObserver.ts

@@ -6,6 +6,7 @@ import type { QueryClient } from './queryClient'
 import type {
   DefaultError,
   MutateOptions,
+  MutationFunctionContext,
   MutationObserverOptions,
   MutationObserverResult,
 } from './types'
@@ -148,16 +149,34 @@ export class MutationObserver<
         const variables = this.#currentResult.variables!
         const scope = this.#currentResult.scope
 
+        const context = {
+          client: this.#client,
+          meta: this.options.meta,
+          mutationKey: this.options.mutationKey,
+        } satisfies MutationFunctionContext
+
         if (action?.type === 'success') {
-          this.#mutateOptions.onSuccess?.(action.data, variables, scope!)
-          this.#mutateOptions.onSettled?.(action.data, null, variables, scope)
+          this.#mutateOptions.onSuccess?.(
+            action.data,
+            variables,
+            scope,
+            context,
+          )
+          this.#mutateOptions.onSettled?.(
+            action.data,
+            null,
+            variables,
+            scope,
+            context,
+          )
         } else if (action?.type === 'error') {
-          this.#mutateOptions.onError?.(action.error, variables, scope)
+          this.#mutateOptions.onError?.(action.error, variables, scope, context)
           this.#mutateOptions.onSettled?.(
             undefined,
             action.error,
             variables,
             scope,
+            context,
           )
         }
       }

packages/query-core/src/types.ts

@@ -1118,17 +1118,20 @@ export interface MutationOptions<
     data: TData,
     variables: TVariables,
     scope: TScope | undefined,
+    context: MutationFunctionContext,
   ) => Promise<unknown> | unknown
   onError?: (
     error: TError,
     variables: TVariables,
     scope: TScope | undefined,
+    context: MutationFunctionContext,
   ) => Promise<unknown> | unknown
   onSettled?: (
     data: TData | undefined,
     error: TError | null,
     variables: TVariables,
     scope: TScope | undefined,
+    context: MutationFunctionContext,
   ) => Promise<unknown> | unknown
   retry?: RetryValue<TError>
   retryDelay?: RetryDelayValue<TError>
@@ -1158,17 +1161,20 @@ export interface MutateOptions<
     data: TData,
     variables: TVariables,
     scope: TScope | undefined,
+    context: MutationFunctionContext,
   ) => void
   onError?: (
     error: TError,
     variables: TVariables,
     scope: TScope | undefined,
+    context: MutationFunctionContext,
   ) => void
   onSettled?: (
     data: TData | undefined,
     error: TError | null,
     variables: TVariables,
     scope: TScope | undefined,
+    context: MutationFunctionContext,
   ) => void
 }
 

packages/react-query/src/__tests__/mutationOptions.test-d.tsx

@@ -4,6 +4,7 @@ import { useIsMutating, useMutation, useMutationState } from '..'
 import { mutationOptions } from '../mutationOptions'
 import type {
   DefaultError,
+  MutationFunctionContext,
   MutationState,
   WithRequired,
 } from '@tanstack/query-core'
@@ -62,7 +63,29 @@ describe('mutationOptions', () => {
         return { name: 'scope' }
       },
       onSuccess: (_data, _variables, scope) => {
-        expectTypeOf(scope).toEqualTypeOf<{ name: string }>()
+        expectTypeOf(scope).toEqualTypeOf<{ name: string } | undefined>()
+      },
+    })
+  })
+
+  it('should infer context type correctly', () => {
+    mutationOptions<number>({
+      mutationFn: (_variables, context) => {
+        expectTypeOf(context).toEqualTypeOf<MutationFunctionContext>()
+        return Promise.resolve(5)
+      },
+      mutationKey: ['key'],
+      onMutate: (_variables, context) => {
+        expectTypeOf(context).toEqualTypeOf<MutationFunctionContext>()
+      },
+      onSuccess: (_data, _variables, _scope, context) => {
+        expectTypeOf(context).toEqualTypeOf<MutationFunctionContext>()
+      },
+      onError: (_error, _variables, _scope, context) => {
+        expectTypeOf(context).toEqualTypeOf<MutationFunctionContext>()
+      },
+      onSettled: (_data, _error, _variables, _scope, context) => {
+        expectTypeOf(context).toEqualTypeOf<MutationFunctionContext>()
       },
     })
   })