Expose per-state schemas on machine state nodes

.changeset/proud-states-schemas.md

@@ -0,0 +1,24 @@
+---
+'xstate': patch
+---
+
+Per-state schemas declared in `setup({ states })` are now available on the compiled machine's state nodes via `machine.states.X.schemas`, so tooling (e.g. per-state snapshot validators) can be derived from the machine alone. Previously they were only accessible on the setup return value.
+
+```ts
+import { setup, types } from 'xstate';
+
+const machine = setup({
+  states: {
+    running: {
+      schemas: { context: types<{ startedAt: number }>() }
+    }
+  }
+}).createMachine({
+  initial: 'running',
+  states: {
+    running: {}
+  }
+});
+
+machine.states.running.schemas?.context; // the schema declared in setup
+```

packages/core/src/StateNode.ts

@@ -1,5 +1,6 @@
 import isDevelopment from '#is-development';
 import { NULL_EVENT, STATE_DELIMITER } from './constants.ts';
+import type { SetupStateSchemas } from './schema.types.ts';
 import { createInvokeTimeoutEventId } from './eventUtils.ts';
 import { memo } from './memo.ts';
 import {
@@ -134,6 +135,8 @@ export class StateNode<
 
   public description?: string;
 
+  public schemas: SetupStateSchemas | undefined;
+
   public tags: string[] = [];
   public transitions!: Map<string, AnyTransitionDefinition[]>;
   public always?: Array<AnyTransitionDefinition>;
@@ -165,6 +168,7 @@ export class StateNode<
           ? 'history'
           : 'atomic');
     this.description = this.config.description;
+    this.schemas = this.config.schemas;
 
     validateStateNodeConfig(this);
 

packages/core/src/schema.types.ts

@@ -4,6 +4,12 @@ export interface StandardSchemaV1<Input = unknown, Output = Input> {
   readonly '~standard': StandardSchemaV1.Props<Input, Output>;
 }
 
+/** Schemas that can be declared for an individual state node. */
+export type SetupStateSchemas = {
+  context?: StandardSchemaV1;
+  input?: StandardSchemaV1;
+};
+
 /** A type-only Standard Schema produced by {@link types}. */
 export interface TypeSchema<T> extends StandardSchemaV1<T, T> {
   readonly '~standard': StandardSchemaV1.Props<T, T> & {

packages/core/src/setup.ts

@@ -1,4 +1,4 @@
-import { StandardSchemaV1 } from './schema.types.ts';
+import { SetupStateSchemas, StandardSchemaV1 } from './schema.types.ts';
 import { StateMachine } from './StateMachine.ts';
 import {
   createActor as createActorFromLogic,
@@ -171,10 +171,7 @@ type ValidateRegistryKeys<
           }
         : unknown);
 
-export type SetupStateSchemas = {
-  context?: StandardSchemaV1;
-  input?: StandardSchemaV1;
-};
+export type { SetupStateSchemas };
 
 export type SetupSchemas = {
   context?: StandardSchemaV1;
@@ -1942,6 +1939,7 @@ export function setup<
     createMachine(machineConfig) {
       const configSchemas = machineConfig.schemas;
       const mergedSchemas = mergeSchemas(configSchemas, schemas);
+      const mergedStates = mergeStateSchemas(machineConfig.states, states);
       const mergedActions = mergeMaps(actions, machineConfig.actions);
       const mergedActors = mergeMaps(actors, machineConfig.actors);
       const mergedGuards = mergeMaps(guards, machineConfig.guards);
@@ -1950,6 +1948,7 @@ export function setup<
       return new StateMachine({
         ...machineConfig,
         ...(mergedSchemas ? { schemas: mergedSchemas } : undefined),
+        ...(mergedStates ? { states: mergedStates } : undefined),
         ...(mergedActions ? { actions: mergedActions } : undefined),
         ...(mergedActors ? { actors: mergedActors } : undefined),
         ...(mergedGuards ? { guards: mergedGuards } : undefined),
@@ -2103,6 +2102,41 @@ function mergeSchemas(
   };
 }
 
+/**
+ * Setup schemas win over inline config schemas (same precedence as root
+ * `mergeSchemas`); setup states with no matching config state are skipped.
+ */
+function mergeStateSchemas(
+  configStates: Record<string, SetupStateSchema> | undefined,
+  setupStates: Record<string, SetupStateSchema> | undefined
+): Record<string, SetupStateSchema> | undefined {
+  if (!configStates || !setupStates) {
+    return configStates;
+  }
+
+  return Object.fromEntries(
+    Object.entries(configStates).map(([key, configState]) => {
+      const setupState = setupStates[key];
+
+      if (!setupState) {
+        return [key, configState];
+      }
+
+      const schemas = mergeMaps(configState.schemas, setupState.schemas);
+      const states = mergeStateSchemas(configState.states, setupState.states);
+
+      return [
+        key,
+        {
+          ...configState,
+          ...(schemas ? { schemas } : undefined),
+          ...(states ? { states } : undefined)
+        }
+      ];
+    })
+  );
+}
+
 function mergeSetupConfigs<
   TBase extends SetupConfig<any, any, any, any, any, any>,
   TExtension extends SetupConfig<any, any, any, any, any, any>

packages/core/src/types.v6.ts

@@ -1,4 +1,4 @@
-import { StandardSchemaV1 } from './schema.types.ts';
+import { SetupStateSchemas, StandardSchemaV1 } from './schema.types.ts';
 import { MachineSnapshot } from './State';
 import {
   Action,
@@ -274,7 +274,7 @@ export type Next_MachineConfig<
     DoNotInfer<TSystemRegistry>,
     DoNotInfer<InferOutput<TOutputSchema, unknown>>
   >,
-  'output'
+  'output' | 'schemas'
 > & {
   internalEvents?: readonly InternalEventDescriptorFor<
     InferEvents<TEventSchemaMap>
@@ -920,6 +920,7 @@ interface Next_ChoiceStateNodeConfig<
   TDelayMap extends Sources['delays']
 > {
   contextSchema?: StandardSchemaV1;
+  schemas?: SetupStateSchemas;
   type: 'choice';
   /** Function that resolves this choice state to a target. */
   choice: Next_ChoiceConfigFunction<
@@ -984,6 +985,7 @@ interface Next_RegularStateNodeConfig<
   TChildOutput = unknown
 > {
   contextSchema?: StandardSchemaV1;
+  schemas?: SetupStateSchemas;
   /** The initial state transition. */
   initial?:
     | string

packages/core/test/setup.test.ts

@@ -120,6 +120,47 @@ describe('setup', () => {
     expect(s.schemas.children.sibling).toBe(sibling);
   });
 
+  it('exposes per-state schemas on the machine state nodes', () => {
+    const rootContext = types<{ count: number }>();
+    const runningContext = types<{ startedAt: number }>();
+    const runningInput = types<{ timeout: number }>();
+    const retryingContext = types<{ attempt: number }>();
+
+    const machine = setup({
+      schemas: { context: rootContext },
+      states: {
+        running: {
+          schemas: { context: runningContext, input: runningInput },
+          states: {
+            retrying: {
+              schemas: { context: retryingContext }
+            }
+          }
+        }
+      }
+    }).createMachine({
+      context: { count: 0 },
+      initial: 'running',
+      states: {
+        running: {
+          initial: 'retrying',
+          states: {
+            retrying: {}
+          }
+        },
+        done: { type: 'final' }
+      }
+    });
+
+    expect(machine.schemas?.context).toBe(rootContext);
+    expect(machine.states.running.schemas?.context).toBe(runningContext);
+    expect(machine.states.running.schemas?.input).toBe(runningInput);
+    expect(machine.states.running.states.retrying.schemas?.context).toBe(
+      retryingContext
+    );
+    expect(machine.states.done.schemas).toBeUndefined();
+  });
+
   it('extends sources', () => {
     const calls: string[] = [];