diff --git a/src/entry-options.ts b/src/entry-options.ts
index 3972fd25..3f28af6c 100644
--- a/src/entry-options.ts
+++ b/src/entry-options.ts
@@ -2,6 +2,6 @@ import type { EntryNodeKey } from './tree-map'
 import type { _ObjectFlat } from './utils'
 
 /**
- * Key used to identify a query. Always an array.
+ * Key used to identify a query or a mutation. Always an array.
  */
 export type EntryKey = Array<EntryNodeKey | _ObjectFlat>
diff --git a/src/use-mutation.test-d.ts b/src/use-mutation.test-d.ts
index 7f57c5f1..52c19417 100644
--- a/src/use-mutation.test-d.ts
+++ b/src/use-mutation.test-d.ts
@@ -2,6 +2,21 @@ import { expectTypeOf, it } from 'vitest'
 import { useMutation } from './use-mutation'
 
 it('types the parameters for the key', () => {
+  const { mutate } = useMutation({
+    mutation: (_one: string) => Promise.resolve({ name: 'foo' }),
+    key(result, one) {
+      expectTypeOf(one).toBeString()
+      expectTypeOf(result).toEqualTypeOf<{ name: string }>()
+      return ['foo']
+    },
+  })
+
+  mutate('one')
+  // @ts-expect-error: missing arg
+  mutate()
+})
+
+it('types the parameters for the keys', () => {
   const { mutate } = useMutation({
     mutation: (_one: string) => Promise.resolve({ name: 'foo' }),
     keys(result, one) {
diff --git a/src/use-mutation.ts b/src/use-mutation.ts
index 2c833187..d801c226 100644
--- a/src/use-mutation.ts
+++ b/src/use-mutation.ts
@@ -1,10 +1,14 @@
 import { computed, shallowRef } from 'vue'
-import type { ComputedRef, MaybeRefOrGetter, ShallowRef } from 'vue'
+import type { ComputedRef, ShallowRef } from 'vue'
 import { useQueryCache } from './query-store'
 import type { EntryKey } from './entry-options'
 import type { ErrorDefault } from './types-extension'
 import { type _Awaitable, noop } from './utils'
 
+type _MutationKey<TVars, TResult> =
+  | EntryKey
+  | ((data: TResult, vars: TVars) => EntryKey)
+
 // TODO: move to a plugin
 /**
  * The keys to invalidate when a mutation succeeds.
@@ -69,7 +73,7 @@ export interface UseMutationOptions<
    */
   mutation: (vars: TVars, context: NoInfer<TContext>) => Promise<TResult>
 
-  key?: MaybeRefOrGetter<EntryKey>
+  key?: _MutationKey<TVars, TResult>
 
   // TODO: move this to a plugin that calls invalidateEntry()
   /**