Improved Fast Refresh support (#4436)

* Initial groundwork for Fast Refresh changes

* Remove `persistedRef`

* Expand on the custom rehydration logic in effect to cover for reconnecting effects

* Revert test changes from "Initial groundwork for Fast Refresh changes"

This partially reverts commit af40d5b.

* Allow rehydrating inline actors

* use `??`

* Changeset

---------

Co-authored-by: David Khourshid <davidkpiano@gmail.com>

.changeset/good-masks-play.md

@@ -0,0 +1,5 @@
+---
+'@xstate/react': minor
+---
+
+Fast refresh now works as expected for most use-cases.

packages/core/src/State.ts

@@ -299,7 +299,8 @@ export function getPersistedState<
     TTag,
     TOutput,
     TResolvedTypesMeta
-  >
+  >,
+  options?: unknown
 ): Snapshot<unknown> {
   const { configuration, tags, machine, children, context, ...jsonValues } =
     state;
@@ -308,11 +309,15 @@ export function getPersistedState<
 
   for (const id in children) {
     const child = children[id] as any;
-    if (isDevelopment && typeof child.src !== 'string') {
+    if (
+      isDevelopment &&
+      typeof child.src !== 'string' &&
+      (!options || !('__unsafeAllowInlineActors' in (options as object)))
+    ) {
       throw new Error('An inline child actor cannot be persisted.');
     }
     childrenJson[id as keyof typeof childrenJson] = {
-      state: child.getPersistedState(),
+      state: child.getPersistedState(options),
       src: child.src,
       systemId: child._systemId
     };

packages/core/src/StateMachine.ts

@@ -45,7 +45,8 @@ import type {
   Equals,
   TODO,
   SnapshotFrom,
-  Snapshot
+  Snapshot,
+  AnyActorLogic
 } from './types.ts';
 import { isErrorActorEvent, resolveReferencedActor } from './utils.ts';
 import { $$ACTOR_TYPE, createActor } from './interpreter.ts';
@@ -180,7 +181,6 @@ export class StateMachine<
     this.getInitialState = this.getInitialState.bind(this);
     this.restoreState = this.restoreState.bind(this);
     this.start = this.start.bind(this);
-    this.getPersistedState = this.getPersistedState.bind(this);
 
     this.root = new StateNode(config, {
       _key: this.id,
@@ -534,9 +534,10 @@ export class StateMachine<
       TTag,
       TOutput,
       TResolvedTypesMeta
-    >
+    >,
+    options?: unknown
   ) {
-    return getPersistedState(state);
+    return getPersistedState(state, options);
   }
 
   public createState(
@@ -587,7 +588,11 @@ export class StateMachine<
     const children: Record<string, AnyActorRef> = {};
     const snapshotChildren: Record<
       string,
-      { src: string; state: Snapshot<unknown>; systemId?: string }
+      {
+        src: string | AnyActorLogic;
+        state: Snapshot<unknown>;
+        systemId?: string;
+      }
     > = (snapshot as any).children;
 
     Object.keys(snapshotChildren).forEach((actorId) => {
@@ -596,7 +601,8 @@ export class StateMachine<
       const childState = actorData.state;
       const src = actorData.src;
 
-      const logic = src ? resolveReferencedActor(this, src)?.src : undefined;
+      const logic =
+        typeof src === 'string' ? resolveReferencedActor(this, src)?.src : src;
 
       if (!logic) {
         return;

packages/core/src/actions/spawn.ts

@@ -59,7 +59,7 @@ function resolveSpawn(
     const configuredInput = input || referenced.input;
     actorRef = createActor(referenced.src, {
       id: resolvedId,
-      src: typeof src === 'string' ? src : undefined,
+      src,
       parent: actorScope?.self,
       systemId,
       input:

packages/core/src/interpreter.ts

@@ -129,7 +129,7 @@ export class Actor<TLogic extends AnyActorLogic>
   public system: ActorSystem<any>;
   private _doneEvent?: DoneActorEvent;
 
-  public src?: string;
+  public src: string | AnyActorLogic;
 
   /**
    * Creates a new actor instance for the given logic with the provided options, if any.
@@ -158,7 +158,7 @@ export class Actor<TLogic extends AnyActorLogic>
     this.clock = clock;
     this._parent = parent;
     this.options = resolvedOptions;
-    this.src = resolvedOptions.src;
+    this.src = resolvedOptions.src ?? logic;
     this.ref = this;
     this._actorScope = {
       self: this,
@@ -186,19 +186,19 @@ export class Actor<TLogic extends AnyActorLogic>
       type: '@xstate.actor',
       actorRef: this
     });
-    this._initState();
+    this._initState(options?.state);
 
     if (systemId && (this._state as any).status === 'active') {
       this._systemId = systemId;
       this.system._set(systemId, this);
     }
   }
 
-  private _initState() {
-    this._state = this.options.state
+  private _initState(persistedState?: Snapshot<unknown>) {
+    this._state = persistedState
       ? this.logic.restoreState
-        ? this.logic.restoreState(this.options.state, this._actorScope)
-        : this.options.state
+        ? this.logic.restoreState(persistedState, this._actorScope)
+        : persistedState
       : this.logic.getInitialState(this._actorScope, this.options?.input);
   }
 
@@ -217,7 +217,6 @@ export class Actor<TLogic extends AnyActorLogic>
     }
 
     for (const observer of this.observers) {
-      // TODO: should observers be notified in case of the error?
       try {
         observer.next?.(snapshot);
       } catch (err) {
@@ -357,6 +356,7 @@ export class Actor<TLogic extends AnyActorLogic>
     }
     this._processingStatus = ProcessingStatus.Running;
 
+    // TODO: this isn't correct when rehydrating
     const initEvent = createInitEvent(this.options.input);
 
     this.system._sendInspectionEvent({
@@ -598,8 +598,9 @@ export class Actor<TLogic extends AnyActorLogic>
     };
   }
 
-  public getPersistedState(): Snapshot<unknown> {
-    return this.logic.getPersistedState(this._state);
+  public getPersistedState(): Snapshot<unknown>;
+  public getPersistedState(options?: unknown): Snapshot<unknown> {
+    return this.logic.getPersistedState(this._state, options);
   }
 
   public [symbolObservable](): InteropSubscribable<SnapshotFrom<TLogic>> {

packages/core/src/spawn.ts

@@ -109,12 +109,11 @@ export function createSpawner(
       }
       return actorRef;
     } else {
-      // TODO: this should also receive `src`
       const actorRef = createActor(src, {
         id: options.id,
         parent: actorScope.self,
         input: options.input,
-        src: undefined,
+        src,
         systemId
       });
 

packages/core/src/types.ts

@@ -1702,7 +1702,7 @@ export interface ActorOptions<TLogic extends AnyActorLogic> {
   /**
    * The source definition.
    */
-  src?: string;
+  src?: string | AnyActorLogic;
 
   inspect?:
     | Observer<InspectionEvent>
@@ -1790,7 +1790,7 @@ export interface ActorRef<
   system?: ActorSystem<any>;
   /** @internal */
   _processingStatus: ProcessingStatus;
-  src?: string;
+  src: string | AnyActorLogic;
 }
 
 export type AnyActorRef = ActorRef<any, any>;
@@ -1981,7 +1981,7 @@ export interface ActorLogic<
   /**
    * @returns Persisted state
    */
-  getPersistedState: (state: TSnapshot) => Snapshot<unknown>;
+  getPersistedState: (state: TSnapshot, options?: unknown) => Snapshot<unknown>;
 }
 
 export type AnyActorLogic = ActorLogic<

packages/xstate-react/src/stopRootWithRehydration.ts

@@ -0,0 +1,35 @@
+import { AnyActorRef, Snapshot } from 'xstate';
+
+const forEachActor = (
+  actorRef: AnyActorRef,
+  callback: (ref: AnyActorRef) => void
+) => {
+  callback(actorRef);
+  const children = actorRef.getSnapshot().children;
+  if (children) {
+    Object.values(children).forEach((child) => {
+      forEachActor(child as AnyActorRef, callback);
+    });
+  }
+};
+
+export function stopRootWithRehydration(actorRef: AnyActorRef) {
+  // persist state here in a custom way allows us to persist inline actors and to preserve actor references
+  // we do it to avoid setState in useEffect when the effect gets "reconnected"
+  // this currently only happens in Strict Effects but it simulates the Offscreen aka Activity API
+  // it also just allows us to end up with a somewhat more predictable behavior for the users
+  const persistedSnapshots: Array<[AnyActorRef, Snapshot<unknown>]> = [];
+  forEachActor(actorRef, (ref) => {
+    persistedSnapshots.push([ref, ref.getSnapshot()]);
+    // muting observers allow us to avoid `useSelector` from being notified about the stopped state
+    // React reconnects its subscribers (from the useSyncExternalStore) on its own
+    // and userland subscibers should basically always do the same anyway
+    // as each subscription should have its own cleanup logic and that should be called each such reconnect
+    (ref as any).observers = new Set();
+  });
+  actorRef.stop();
+  persistedSnapshots.forEach(([ref, snapshot]) => {
+    (ref as any)._processingStatus = 0;
+    (ref as any)._state = snapshot;
+  });
+}

packages/xstate-react/src/useActor.ts

@@ -7,7 +7,8 @@ import {
   ActorOptions,
   SnapshotFrom
 } from 'xstate';
-import { useIdleInterpreter } from './useActorRef.ts';
+import { useIdleActor } from './useActorRef.ts';
+import { stopRootWithRehydration } from './stopRootWithRehydration.ts';
 
 export function useActor<TLogic extends AnyActorLogic>(
   logic: TLogic,
@@ -24,7 +25,7 @@ export function useActor<TLogic extends AnyActorLogic>(
     );
   }
 
-  const actorRef = useIdleInterpreter(logic, options as any);
+  const actorRef = useIdleActor(logic, options as any);
 
   const getSnapshot = useCallback(() => {
     return actorRef.getSnapshot();
@@ -48,9 +49,7 @@ export function useActor<TLogic extends AnyActorLogic>(
     actorRef.start();
 
     return () => {
-      actorRef.stop();
-      (actorRef as any)._processingStatus = 0;
-      (actorRef as any)._initState();
+      stopRootWithRehydration(actorRef);
     };
   }, [actorRef]);
 

packages/xstate-react/src/useActorRef.ts

@@ -1,4 +1,3 @@
-import isDevelopment from '#is-development';
 import { useEffect, useState } from 'react';
 import {
   AnyActorLogic,
@@ -15,35 +14,37 @@ import {
   SnapshotFrom,
   TODO
 } from 'xstate';
-import useConstant from './useConstant.ts';
 import useIsomorphicLayoutEffect from 'use-isomorphic-layout-effect';
+import { stopRootWithRehydration } from './stopRootWithRehydration';
 
-export function useIdleInterpreter(
-  machine: AnyActorLogic,
+export function useIdleActor(
+  logic: AnyActorLogic,
   options: Partial<ActorOptions<AnyActorLogic>>
 ): AnyActor {
-  if (isDevelopment) {
-    const [initialMachine] = useState(machine);
+  let [[currentConfig, actorRef], setCurrent] = useState(() => {
+    const actorRef = createActor(logic, options);
+    return [logic.config, actorRef];
+  });
 
-    if (machine.config !== initialMachine.config) {
-      console.warn(
-        `Actor logic has changed between renders. This is not supported and may lead to invalid snapshots.`
-      );
-    }
+  if (logic.config !== currentConfig) {
+    const newActorRef = createActor(logic, {
+      ...options,
+      state: (actorRef.getPersistedState as any)({
+        __unsafeAllowInlineActors: true
+      })
+    });
+    setCurrent([logic.config, newActorRef]);
+    actorRef = newActorRef;
   }
 
-  const actorRef = useConstant(() => {
-    return createActor(machine as AnyStateMachine, options);
-  });
-
   // TODO: consider using `useAsapEffect` that would do this in `useInsertionEffect` is that's available
   useIsomorphicLayoutEffect(() => {
     (actorRef.logic as AnyStateMachine).implementations = (
-      machine as AnyStateMachine
+      logic as AnyStateMachine
     ).implementations;
   });
 
-  return actorRef as any;
+  return actorRef;
 }
 
 export function useActorRef<TLogic extends AnyActorLogic>(
@@ -53,7 +54,7 @@ export function useActorRef<TLogic extends AnyActorLogic>(
     | Observer<SnapshotFrom<TLogic>>
     | ((value: SnapshotFrom<TLogic>) => void)
 ): ActorRefFrom<TLogic> {
-  const actorRef = useIdleInterpreter(machine, options);
+  const actorRef = useIdleActor(machine, options);
 
   useEffect(() => {
     if (!observerOrListener) {
@@ -69,11 +70,9 @@ export function useActorRef<TLogic extends AnyActorLogic>(
     actorRef.start();
 
     return () => {
-      actorRef.stop();
-      (actorRef as any)._processingStatus = 0;
-      (actorRef as any)._initState();
+      stopRootWithRehydration(actorRef);
     };
-  }, []);
+  }, [actorRef]);
 
   return actorRef as any;
 }