fix(lit-query): propagate mutation state generics

.changeset/lit-mutation-state-generics.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/lit-query': patch
+---
+
+fix(types): propagate generic type parameters to the `useMutationState` select callback

packages/lit-query/src/__tests__/mutation-state-types.test.ts

@@ -0,0 +1,79 @@
+import { QueryClient } from '@tanstack/query-core'
+import { LitElement } from 'lit'
+import { describe, expectTypeOf, it } from 'vitest'
+import { useMutationState } from '../useMutationState.js'
+import { generateElementName } from './utils.js'
+import type {
+  Mutation,
+  MutationState,
+  MutationStatus,
+} from '@tanstack/query-core'
+
+describe('useMutationState types', () => {
+  it('defaults to mutation state with the default error type', () => {
+    class Host extends LitElement {}
+    customElements.define(generateElementName(), Host)
+
+    const result = useMutationState(
+      new Host(),
+      { filters: { status: 'pending' } },
+      new QueryClient(),
+    )
+
+    expectTypeOf(result()).toEqualTypeOf<
+      Array<MutationState<unknown, Error, unknown, unknown>>
+    >()
+  })
+
+  it('infers the selected status', () => {
+    class Host extends LitElement {}
+    customElements.define(generateElementName(), Host)
+
+    const result = useMutationState(
+      new Host(),
+      {
+        filters: { status: 'pending' },
+        select: (mutation) => mutation.state.status,
+      },
+      new QueryClient(),
+    )
+
+    expectTypeOf(result()).toEqualTypeOf<Array<MutationStatus>>()
+  })
+
+  it('allows selecting a default error message', () => {
+    class Host extends LitElement {}
+    customElements.define(generateElementName(), Host)
+
+    const result = useMutationState(
+      new Host(),
+      { select: (mutation) => mutation.state.error?.message },
+      new QueryClient(),
+    )
+
+    expectTypeOf(result()).toEqualTypeOf<Array<string | undefined>>()
+  })
+
+  it('propagates typed mutation state generics to the select callback', () => {
+    type MyData = { data: Array<string> }
+    type MyError = { code: number; message: string }
+    type MyVariables = { id: number }
+
+    class Host extends LitElement {}
+    customElements.define(generateElementName(), Host)
+
+    useMutationState<MutationState<MyData, MyError, MyVariables>>(
+      new Host(),
+      {
+        filters: { mutationKey: ['key'] },
+        select: (mutation) => {
+          expectTypeOf(mutation).toEqualTypeOf<
+            Mutation<MyData, MyError, MyVariables, unknown>
+          >()
+          return mutation.state
+        },
+      },
+      new QueryClient(),
+    )
+  })
+})

packages/lit-query/src/useMutationState.ts

@@ -10,14 +10,28 @@ import type {
 } from '@tanstack/query-core'
 import type { ReactiveControllerHost } from 'lit'
 
-/**
- * Options accepted by `useMutationState`.
- */
-export type MutationStateOptions<TResult> = {
+/** Recovers mutation generics when the selected result is a mutation state. */
+type MutationTypeFromResult<TResult> = [TResult] extends [
+  MutationState<
+    infer TData,
+    infer TError,
+    infer TVariables,
+    infer TOnMutateResult
+  >,
+]
+  ? Mutation<TData, TError, TVariables, TOnMutateResult>
+  : Mutation
+
+/** Options accepted by `useMutationState`. */
+export type MutationStateOptions<
+  TResult,
+  TMutation extends Mutation<any, any, any, any> =
+    MutationTypeFromResult<TResult>,
+> = {
   /** Filters used to select mutations from the mutation cache. */
   filters?: Accessor<MutationFilters>
   /** Maps each matching mutation to the value returned by the accessor. */
-  select?: (mutation: Mutation) => TResult
+  select?: (mutation: TMutation) => TResult
 }
 
 /**
@@ -31,13 +45,16 @@ export type MutationStateAccessor<TResult> = ValueAccessor<Array<TResult>> & {
   destroy: () => void
 }
 
-class MutationStateController<TResult> extends BaseController<Array<TResult>> {
+class MutationStateController<
+  TResult,
+  TMutation extends Mutation<any, any, any, any>,
+> extends BaseController<Array<TResult>> {
   private queryClient: QueryClient | undefined
   private unsubscribe: (() => void) | undefined
 
   constructor(
     host: ReactiveControllerHost,
-    private readonly options: MutationStateOptions<TResult>,
+    private readonly options: MutationStateOptions<TResult, TMutation>,
     queryClient?: QueryClient,
   ) {
     super(host, [], queryClient)
@@ -144,7 +161,7 @@ class MutationStateController<TResult> extends BaseController<Array<TResult>> {
 
     return mutations.map((mutation) => {
       if (select) {
-        return select(mutation)
+        return select(mutation as TMutation)
       }
 
       return mutation.state as TResult
@@ -184,13 +201,19 @@ class MutationStateController<TResult> extends BaseController<Array<TResult>> {
  * ```
  */
 export function useMutationState<
-  TResult = MutationState<unknown, unknown, unknown, unknown>,
+  TResult = MutationState,
+  TMutation extends Mutation<any, any, any, any> =
+    MutationTypeFromResult<TResult>,
 >(
   host: ReactiveControllerHost,
-  options: MutationStateOptions<TResult> = {},
+  options: MutationStateOptions<TResult, TMutation> = {},
   queryClient?: QueryClient,
 ): MutationStateAccessor<TResult> {
-  const controller = new MutationStateController(host, options, queryClient)
+  const controller = new MutationStateController<TResult, TMutation>(
+    host,
+    options,
+    queryClient,
+  )
   return Object.assign(
     createValueAccessor(() => controller.current),
     {