Allow standard schemas to validate endpoint values (#4864)

* add argument validation

* add result, error, and meta schemas

* declare error shape once

* add type tests

* pull types from spec instead of inlining

* assign argSchema result

* move common option to common options section

* tweak test

* add metaSchema type test

* add test for inference from schemas

* resultSchema -> responseSchema

* add argSchema inference test

* errorSchema -> errorResponse

* add inference of raw result type from schema (or specifying from a third type parameter)

* add api level and endpoint level handlers for schema failures

* export NamedSchemaError

* added skipSchemaValidation flag

* persist raw result type after building

Author: EskiMojo14

packages/toolkit/package.json

@@ -99,6 +99,7 @@
     "tsup": "^8.2.3",
     "tsx": "^4.19.0",
     "typescript": "^5.8.2",
+    "valibot": "^1.0.0",
     "vite-tsconfig-paths": "^4.3.1",
     "vitest": "^1.6.0",
     "yargs": "^15.3.1"
@@ -124,6 +125,8 @@
     "react"
   ],
   "dependencies": {
+    "@standard-schema/spec": "^1.0.0",
+    "@standard-schema/utils": "^0.3.0",
     "immer": "^10.0.3",
     "redux": "^5.0.1",
     "redux-thunk": "^3.1.0",

packages/toolkit/src/query/core/apiState.ts

@@ -179,7 +179,7 @@ export type MutationKeys<Definitions extends EndpointDefinitions> = {
 }[keyof Definitions]
 
 type BaseQuerySubState<
-  D extends BaseEndpointDefinition<any, any, any>,
+  D extends BaseEndpointDefinition<any, any, any, any>,
   DataType = ResultTypeFrom<D>,
 > = {
   /**
@@ -217,7 +217,7 @@ type BaseQuerySubState<
 }
 
 export type QuerySubState<
-  D extends BaseEndpointDefinition<any, any, any>,
+  D extends BaseEndpointDefinition<any, any, any, any>,
   DataType = ResultTypeFrom<D>,
 > = Id<
   | ({ status: QueryStatus.fulfilled } & WithRequiredProp<
@@ -244,15 +244,17 @@ export type QuerySubState<
 export type InfiniteQueryDirection = 'forward' | 'backward'
 
 export type InfiniteQuerySubState<
-  D extends BaseEndpointDefinition<any, any, any>,
+  D extends BaseEndpointDefinition<any, any, any, any>,
 > =
   D extends InfiniteQueryDefinition<any, any, any, any, any>
     ? QuerySubState<D, InfiniteData<ResultTypeFrom<D>, PageParamFrom<D>>> & {
         direction?: InfiniteQueryDirection
       }
     : never
 
-type BaseMutationSubState<D extends BaseEndpointDefinition<any, any, any>> = {
+type BaseMutationSubState<
+  D extends BaseEndpointDefinition<any, any, any, any>,
+> = {
   requestId: string
   data?: ResultTypeFrom<D>
   error?:
@@ -265,8 +267,12 @@ type BaseMutationSubState<D extends BaseEndpointDefinition<any, any, any>> = {
   fulfilledTimeStamp?: number
 }
 
-export type MutationSubState<D extends BaseEndpointDefinition<any, any, any>> =
-  | (({ status: QueryStatus.fulfilled } & WithRequiredProp<
+export type MutationSubState<
+  D extends BaseEndpointDefinition<any, any, any, any>,
+> =
+  | (({
+      status: QueryStatus.fulfilled
+    } & WithRequiredProp<
       BaseMutationSubState<D>,
       'data' | 'fulfilledTimeStamp'
     >) & { error: undefined })

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -1,5 +1,9 @@
 import type { ThunkDispatch, UnknownAction } from '@reduxjs/toolkit'
-import type { BaseQueryFn, BaseQueryMeta } from '../../baseQueryTypes'
+import type {
+  BaseQueryFn,
+  BaseQueryMeta,
+  BaseQueryResult,
+} from '../../baseQueryTypes'
 import type { BaseEndpointDefinition } from '../../endpointDefinitions'
 import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { QueryCacheKey, RootState } from '../apiState'
@@ -32,7 +36,8 @@ export interface QueryBaseLifecycleApi<
     { type: DefinitionType.query } & BaseEndpointDefinition<
       QueryArg,
       BaseQuery,
-      ResultType
+      ResultType,
+      BaseQueryResult<BaseQuery>
     >
   >
   /**
@@ -55,7 +60,8 @@ export type MutationBaseLifecycleApi<
     { type: DefinitionType.mutation } & BaseEndpointDefinition<
       QueryArg,
       BaseQuery,
-      ResultType
+      ResultType,
+      BaseQueryResult<BaseQuery>
     >
   >
 }

packages/toolkit/src/query/core/buildThunks.ts

@@ -28,6 +28,8 @@ import type {
   QueryDefinition,
   ResultDescription,
   ResultTypeFrom,
+  SchemaFailureHandler,
+  SchemaFailureInfo,
 } from '../endpointDefinitions'
 import {
   calculateProvidedBy,
@@ -65,6 +67,7 @@ import {
   isRejectedWithValue,
   SHOULD_AUTOBATCH,
 } from './rtkImports'
+import { parseWithSchema, NamedSchemaError } from '../standardSchema'
 
 export type BuildThunksApiEndpointQuery<
   Definition extends QueryDefinition<any, any, any, any, any>,
@@ -329,6 +332,8 @@ export function buildThunks<
   api,
   assertTagType,
   selectors,
+  onSchemaFailure,
+  skipSchemaValidation: globalSkipSchemaValidation,
 }: {
   baseQuery: BaseQuery
   reducerPath: ReducerPath
@@ -337,6 +342,8 @@ export function buildThunks<
   api: Api<BaseQuery, Definitions, ReducerPath, any>
   assertTagType: AssertTagTypes
   selectors: AllSelectors
+  onSchemaFailure: SchemaFailureHandler | undefined
+  skipSchemaValidation: boolean | undefined
 }) {
   type State = RootState<any, string, ReducerPath>
 
@@ -491,10 +498,14 @@ export function buildThunks<
     },
   ) => {
     const endpointDefinition = endpointDefinitions[arg.endpointName]
+    const { metaSchema, skipSchemaValidation = globalSkipSchemaValidation } =
+      endpointDefinition
 
     try {
-      let transformResponse: TransformCallback =
-        getTransformCallbackForEndpoint(endpointDefinition, 'transformResponse')
+      let transformResponse = getTransformCallbackForEndpoint(
+        endpointDefinition,
+        'transformResponse',
+      )
 
       const baseQueryApi = {
         signal,
@@ -551,7 +562,16 @@ export function buildThunks<
         finalQueryArg: unknown,
       ): Promise<QueryReturnValue> {
         let result: QueryReturnValue
-        const { extraOptions } = endpointDefinition
+        const { extraOptions, argSchema, rawResponseSchema, responseSchema } =
+          endpointDefinition
+
+        if (argSchema && !skipSchemaValidation) {
+          finalQueryArg = await parseWithSchema(
+            argSchema,
+            finalQueryArg,
+            'argSchema',
+          )
+        }
 
         if (forceQueryFn) {
           // upsertQueryData relies on this to pass in the user-provided value
@@ -606,13 +626,34 @@ export function buildThunks<
 
         if (result.error) throw new HandledError(result.error, result.meta)
 
-        const transformedResponse = await transformResponse(
-          result.data,
+        let { data } = result
+
+        if (rawResponseSchema && !skipSchemaValidation) {
+          data = await parseWithSchema(
+            rawResponseSchema,
+            result.data,
+            'rawResponseSchema',
+          )
+        }
+
+        let transformedResponse = await transformResponse(
+          data,
           result.meta,
           finalQueryArg,
         )
 
-        return { ...result, data: transformedResponse }
+        if (responseSchema && !skipSchemaValidation) {
+          transformedResponse = await parseWithSchema(
+            responseSchema,
+            transformedResponse,
+            'responseSchema',
+          )
+        }
+
+        return {
+          ...result,
+          data: transformedResponse,
+        }
       }
 
       if (
@@ -698,6 +739,14 @@ export function buildThunks<
         finalQueryReturnValue = await executeRequest(arg.originalArgs)
       }
 
+      if (metaSchema && !skipSchemaValidation && finalQueryReturnValue.meta) {
+        finalQueryReturnValue.meta = await parseWithSchema(
+          metaSchema,
+          finalQueryReturnValue.meta,
+          'metaSchema',
+        )
+      }
+
       // console.log('Final result: ', transformedData)
       return fulfillWithValue(
         finalQueryReturnValue.data,
@@ -707,40 +756,78 @@ export function buildThunks<
         }),
       )
     } catch (error) {
-      let catchedError = error
-      if (catchedError instanceof HandledError) {
-        let transformErrorResponse: TransformCallback =
-          getTransformCallbackForEndpoint(
+      try {
+        let caughtError = error
+        if (caughtError instanceof HandledError) {
+          let transformErrorResponse = getTransformCallbackForEndpoint(
             endpointDefinition,
             'transformErrorResponse',
           )
+          const { rawErrorResponseSchema, errorResponseSchema } =
+            endpointDefinition
+
+          let { value, meta } = caughtError
+
+          if (rawErrorResponseSchema && !skipSchemaValidation) {
+            value = await parseWithSchema(
+              rawErrorResponseSchema,
+              value,
+              'rawErrorResponseSchema',
+            )
+          }
+
+          if (metaSchema && !skipSchemaValidation) {
+            meta = await parseWithSchema(metaSchema, meta, 'metaSchema')
+          }
 
-        try {
-          return rejectWithValue(
-            await transformErrorResponse(
-              catchedError.value,
-              catchedError.meta,
+          try {
+            let transformedErrorResponse = await transformErrorResponse(
+              value,
+              meta,
               arg.originalArgs,
-            ),
-            addShouldAutoBatch({ baseQueryMeta: catchedError.meta }),
-          )
-        } catch (e) {
-          catchedError = e
+            )
+            if (errorResponseSchema && !skipSchemaValidation) {
+              transformedErrorResponse = await parseWithSchema(
+                errorResponseSchema,
+                transformedErrorResponse,
+                'errorResponseSchema',
+              )
+            }
+
+            return rejectWithValue(
+              transformedErrorResponse,
+              addShouldAutoBatch({ baseQueryMeta: meta }),
+            )
+          } catch (e) {
+            caughtError = e
+          }
         }
-      }
-      if (
-        typeof process !== 'undefined' &&
-        process.env.NODE_ENV !== 'production'
-      ) {
-        console.error(
-          `An unhandled error occurred processing a request for the endpoint "${arg.endpointName}".
+        if (
+          typeof process !== 'undefined' &&
+          process.env.NODE_ENV !== 'production'
+        ) {
+          console.error(
+            `An unhandled error occurred processing a request for the endpoint "${arg.endpointName}".
 In the case of an unhandled error, no tags will be "provided" or "invalidated".`,
-          catchedError,
-        )
-      } else {
-        console.error(catchedError)
+            caughtError,
+          )
+        } else {
+          console.error(caughtError)
+        }
+        throw caughtError
+      } catch (error) {
+        if (error instanceof NamedSchemaError) {
+          const info: SchemaFailureInfo = {
+            endpoint: arg.endpointName,
+            arg: arg.originalArgs,
+            type: arg.type,
+            queryCacheKey: arg.type === 'query' ? arg.queryCacheKey : undefined,
+          }
+          endpointDefinition.onSchemaFailure?.(error, info)
+          onSchemaFailure?.(error, info)
+        }
+        throw error
       }
-      throw catchedError
     }
   }
 

packages/toolkit/src/query/core/module.ts

@@ -516,6 +516,8 @@ export const coreModule = ({
       refetchOnFocus,
       refetchOnReconnect,
       invalidationBehavior,
+      onSchemaFailure,
+      skipSchemaValidation,
     },
     context,
   ) {
@@ -582,6 +584,8 @@ export const coreModule = ({
       serializeQueryArgs,
       assertTagType,
       selectors,
+      onSchemaFailure,
+      skipSchemaValidation,
     })
 
     const { reducer, actions: sliceActions } = buildSlice({

packages/toolkit/src/query/createApi.ts

@@ -6,6 +6,7 @@ import { defaultSerializeQueryArgs } from './defaultSerializeQueryArgs'
 import type {
   EndpointBuilder,
   EndpointDefinitions,
+  SchemaFailureHandler,
 } from './endpointDefinitions'
 import {
   DefinitionType,
@@ -212,6 +213,9 @@ export interface CreateApiOptions<
         NoInfer<TagTypes>,
         NoInfer<ReducerPath>
       >
+
+  onSchemaFailure?: SchemaFailureHandler
+  skipSchemaValidation?: boolean
 }
 
 export type CreateApi<Modules extends ModuleName> = {