From 6d515d0913e526824262232df3e4758be026c4c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20Burzy=C5=84ski?= <mateuszburzynski@gmail.com>
Date: Wed, 10 Jan 2024 17:40:23 +0100
Subject: [PATCH] Throw errors from actor snapshots to trigger error boundaries
 in React

---
 packages/xstate-react/src/useActor.ts    | 17 +++++++++++++++--
 packages/xstate-react/src/useActorRef.ts | 22 ++++++++++++++++++----
 packages/xstate-react/src/useSelector.ts | 18 +++++++++++++++---
 3 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/packages/xstate-react/src/useActor.ts b/packages/xstate-react/src/useActor.ts
index d122555a25e..d85a616d20c 100644
--- a/packages/xstate-react/src/useActor.ts
+++ b/packages/xstate-react/src/useActor.ts
@@ -1,7 +1,13 @@
 import isDevelopment from '#is-development';
 import { useCallback, useEffect } from 'react';
 import { useSyncExternalStore } from 'use-sync-external-store/shim';
-import { Actor, ActorOptions, AnyActorLogic, SnapshotFrom } from 'xstate';
+import {
+  Actor,
+  ActorOptions,
+  AnyActorLogic,
+  Snapshot,
+  SnapshotFrom
+} from 'xstate';
 import { stopRootWithRehydration } from './stopRootWithRehydration.ts';
 import { useIdleActorRef } from './useActorRef.ts';
 
@@ -28,7 +34,10 @@ export function useActor<TLogic extends AnyActorLogic>(
 
   const subscribe = useCallback(
     (handleStoreChange) => {
-      const { unsubscribe } = actorRef.subscribe(handleStoreChange);
+      const { unsubscribe } = actorRef.subscribe(
+        handleStoreChange,
+        handleStoreChange
+      );
       return unsubscribe;
     },
     [actorRef]
@@ -40,6 +49,10 @@ export function useActor<TLogic extends AnyActorLogic>(
     getSnapshot
   );
 
+  if ((actorSnapshot as Snapshot<any>).status === 'error') {
+    throw (actorSnapshot as Snapshot<any>).error;
+  }
+
   useEffect(() => {
     actorRef.start();
 
diff --git a/packages/xstate-react/src/useActorRef.ts b/packages/xstate-react/src/useActorRef.ts
index e491829f4d1..d6c58607cd1 100644
--- a/packages/xstate-react/src/useActorRef.ts
+++ b/packages/xstate-react/src/useActorRef.ts
@@ -6,6 +6,7 @@ import {
   AnyActorLogic,
   AnyStateMachine,
   Observer,
+  Snapshot,
   SnapshotFrom,
   createActor,
   toObserver
@@ -42,6 +43,8 @@ export function useIdleActorRef<TLogic extends AnyActorLogic>(
   return actorRef;
 }
 
+const UNIQUE = {};
+
 export function useActorRef<TLogic extends AnyActorLogic>(
   machine: TLogic,
   options: ActorOptions<TLogic> = {},
@@ -50,12 +53,23 @@ export function useActorRef<TLogic extends AnyActorLogic>(
     | ((value: SnapshotFrom<TLogic>) => void)
 ): Actor<TLogic> {
   const actorRef = useIdleActorRef(machine, options);
+  const [reactError, setReactError] = useState(() => {
+    const initialSnapshot: Snapshot<any> = actorRef.getSnapshot();
+    return initialSnapshot.status === 'error' ? initialSnapshot.error : UNIQUE;
+  });
+
+  if (reactError !== UNIQUE) {
+    throw reactError;
+  }
 
   useEffect(() => {
-    if (!observerOrListener) {
-      return;
-    }
-    let sub = actorRef.subscribe(toObserver(observerOrListener));
+    const observer = toObserver(observerOrListener);
+    const errorListener = observer.error;
+    observer.error = (error) => {
+      setReactError(error);
+      errorListener?.(error);
+    };
+    let sub = actorRef.subscribe(observer);
     return () => {
       sub.unsubscribe();
     };
diff --git a/packages/xstate-react/src/useSelector.ts b/packages/xstate-react/src/useSelector.ts
index 8509856515f..5a0cde792d1 100644
--- a/packages/xstate-react/src/useSelector.ts
+++ b/packages/xstate-react/src/useSelector.ts
@@ -1,6 +1,6 @@
 import { useCallback } from 'react';
 import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector';
-import { ActorRef, SnapshotFrom } from 'xstate';
+import { ActorRef, Snapshot, SnapshotFrom } from 'xstate';
 
 function defaultCompare<T>(a: T, b: T) {
   return a === b;
@@ -13,19 +13,31 @@ export function useSelector<TActor extends ActorRef<any, any>, T>(
 ): T {
   const subscribe = useCallback(
     (handleStoreChange) => {
-      const { unsubscribe } = actor.subscribe(handleStoreChange);
+      const { unsubscribe } = actor.subscribe(
+        handleStoreChange,
+        handleStoreChange
+      );
       return unsubscribe;
     },
     [actor]
   );
 
   const boundGetSnapshot = useCallback(() => actor.getSnapshot(), [actor]);
+  const boundSelector: typeof selector = useCallback(
+    (snapshot: Snapshot<any>) => {
+      if (snapshot.status === 'error') {
+        throw snapshot.error;
+      }
+      return selector(snapshot as never);
+    },
+    [selector]
+  );
 
   const selectedSnapshot = useSyncExternalStoreWithSelector(
     subscribe,
     boundGetSnapshot,
     boundGetSnapshot,
-    selector,
+    boundSelector,
     compare
   );