From e201278a0b36aed294b51e30bda8dc28a5b8578f Mon Sep 17 00:00:00 2001
From: Cristian Velasquez Ramos <seeve.c@gmail.com>
Date: Tue, 16 Jan 2024 11:33:35 -0500
Subject: [PATCH] refactor system subscribe to make use of inspect. add
 registration events to inspect, and have return a subscription.

---
 .changeset/tasty-baboons-dance.md  |  12 +--
 packages/core/src/system.ts        | 143 ++++++++++-------------------
 packages/core/test/inspect.test.ts | 129 ++++++++++++++++++++++++++
 packages/core/test/system.test.ts  | 114 -----------------------
 4 files changed, 185 insertions(+), 213 deletions(-)

diff --git a/.changeset/tasty-baboons-dance.md b/.changeset/tasty-baboons-dance.md
index b8c2643132a..a79930e29ce 100644
--- a/.changeset/tasty-baboons-dance.md
+++ b/.changeset/tasty-baboons-dance.md
@@ -2,28 +2,28 @@
 'xstate': minor
 ---
 
-added support for `actor.system.subscribe` to subscribe to registration and unregistration events within an actor's system.
-`actor.system.subscribe` returns a `Subscription` object you can `.unsubscribe` to at any time.
+added registration and unregistration events to `actor.system.inspect`.
+`actor.system.inspect` now returns a `Subscription` object you can `.unsubscribe` to at any time.
 
 ex:
 
 ```js
 // observer object
-const subscription = actor.system.subscribe({
+const subscription = actor.system.inspect({
   next: (event) => console.log(event),
   error: (err) => console.error(err),
   complete: () => console.log('done')
 });
 
 // observer parameters
-const subscription = actor.system.subscribe(
+const subscription = actor.system.inspect(
   (event) => console.log(event),
   (err) => console.error(err),
   () => console.log('done')
 );
 
-// callback function
-const subscription = actor.system.subscribe((event) => console.log(event));
+// single listener
+const subscription = actor.system.inspect((event) => console.log(event));
 
 // unsubscribe
 subscription.unsubscribe();
diff --git a/packages/core/src/system.ts b/packages/core/src/system.ts
index 9f6fda68ddd..d05b2609b38 100644
--- a/packages/core/src/system.ts
+++ b/packages/core/src/system.ts
@@ -1,14 +1,12 @@
-import { ProcessingStatus } from './createActor.ts';
-import { reportUnhandledError } from './reportUnhandledError.ts';
 import {
   AnyEventObject,
   ActorSystemInfo,
   AnyActorRef,
   Observer,
   Snapshot,
-  Subscribable,
   HomomorphicOmit,
-  EventObject
+  EventObject,
+  Subscription
 } from './types.ts';
 import { toObserver } from './utils.ts';
 
@@ -47,13 +45,7 @@ function createScheduledEventId(
   return `${actorRef.sessionId}.${id}` as ScheduledEventId;
 }
 
-export interface ActorSystem<T extends ActorSystemInfo>
-  extends Subscribable<
-    RegistrationEvent<
-      T['actors'][keyof T['actors']],
-      keyof T['actors'] & string
-    >
-  > {
+export interface ActorSystem<T extends ActorSystemInfo> {
   /**
    * @internal
    */
@@ -66,18 +58,19 @@ export interface ActorSystem<T extends ActorSystemInfo>
    * @internal
    */
   _unregister: (actorRef: AnyActorRef) => void;
-  /**
-   * @internal
-   */
-  _sendRegistrationEvent: (
-    event: RegistrationEvent<T['actors'][keyof T['actors']]>
-  ) => void;
   /**
    * @internal
    */
   _set: <K extends keyof T['actors']>(key: K, actorRef: T['actors'][K]) => void;
   get: <K extends keyof T['actors']>(key: K) => T['actors'][K] | undefined;
-  inspect: (observer: Observer<InspectionEvent>) => void;
+  inspect: {
+    (observer: Observer<InspectionEvent>): Subscription;
+    (
+      next: (value: InspectionEvent) => void,
+      error?: (error: any) => void,
+      complete?: () => void
+    ): Subscription;
+  };
   /**
    * @internal
    */
@@ -118,10 +111,7 @@ export function createSystem<T extends ActorSystemInfo>(
   const children = new Map<string, AnyActorRef>();
   const keyedActors = new Map<keyof T['actors'], AnyActorRef | undefined>();
   const reverseKeyedActors = new WeakMap<AnyActorRef, keyof T['actors']>();
-  const inspectionObservers = new Set<Observer<InspectionEvent>>();
-  const registrationObservers = new Set<
-    Observer<RegistrationEvent<T['actors'][keyof T['actors']]>>
-  >();
+  const observers = new Set<Observer<InspectionEvent>>();
   const timerMap: { [id: ScheduledEventId]: number } = {};
   const clock = options.clock;
 
@@ -185,10 +175,10 @@ export function createSystem<T extends ActorSystemInfo>(
       children.set(sessionId, actorRef);
       const systemId = reverseKeyedActors.get(actorRef);
       if (systemId !== undefined) {
-        system._sendRegistrationEvent({
-          type: `@xstate.actor.register`,
+        system._sendInspectionEvent({
+          type: `@xstate.register`,
           systemId: systemId as string,
-          actorRef: actorRef as T['actors'][keyof T['actors']]
+          actorRef
         });
       }
       return sessionId;
@@ -200,20 +190,31 @@ export function createSystem<T extends ActorSystemInfo>(
       if (systemId !== undefined) {
         keyedActors.delete(systemId);
         reverseKeyedActors.delete(actorRef);
-        system._sendRegistrationEvent({
-          type: `@xstate.actor.unregister`,
+        system._sendInspectionEvent({
+          type: `@xstate.unregister`,
           systemId: systemId as string,
-          actorRef: actorRef as T['actors'][keyof T['actors']]
-        } as const);
+          actorRef
+        });
       }
     },
     get: (systemId) => {
       return keyedActors.get(systemId) as T['actors'][any];
     },
-    subscribe: (
+    _set: (systemId, actorRef) => {
+      const existing = keyedActors.get(systemId);
+      if (existing && existing !== actorRef) {
+        throw new Error(
+          `Actor with system ID '${systemId as string}' already exists.`
+        );
+      }
+
+      keyedActors.set(systemId, actorRef);
+      reverseKeyedActors.set(actorRef, systemId);
+    },
+    inspect: (
       nextListenerOrObserver:
-        | ((event: RegistrationEvent<T['actors'][keyof T['actors']]>) => void)
-        | Observer<RegistrationEvent<T['actors'][keyof T['actors']]>>,
+        | ((event: InspectionEvent) => void)
+        | Observer<InspectionEvent>,
       errorListener?: (error: any) => void,
       completeListener?: () => void
     ) => {
@@ -223,53 +224,20 @@ export function createSystem<T extends ActorSystemInfo>(
         completeListener
       );
 
-      if (rootActor._processingStatus !== ProcessingStatus.Stopped) {
-        registrationObservers.add(observer);
-      } else {
-        const snapshot = rootActor.getSnapshot();
-        switch (snapshot.status) {
-          case 'done':
-            try {
-              observer.complete?.();
-            } catch (err) {
-              reportUnhandledError(err);
-            }
-            break;
-          // can this error?
-        }
-      }
+      observers.add(observer);
 
       return {
         unsubscribe: () => {
-          registrationObservers.delete(observer);
+          observers.delete(observer);
         }
       };
     },
-    _set: (systemId, actorRef) => {
-      const existing = keyedActors.get(systemId);
-      if (existing && existing !== actorRef) {
-        throw new Error(
-          `Actor with system ID '${systemId as string}' already exists.`
-        );
-      }
-
-      keyedActors.set(systemId, actorRef);
-      reverseKeyedActors.set(actorRef, systemId);
-    },
-    inspect: (observer) => {
-      inspectionObservers.add(observer);
-    },
     _sendInspectionEvent: (event) => {
       const resolvedInspectionEvent: InspectionEvent = {
         ...event,
         rootId: rootActor.sessionId
       };
-      inspectionObservers.forEach(
-        (observer) => observer.next?.(resolvedInspectionEvent)
-      );
-    },
-    _sendRegistrationEvent: (event) => {
-      registrationObservers.forEach((observer) => observer.next?.(event));
+      observers.forEach((observer) => observer.next?.(resolvedInspectionEvent));
     },
     _relay: (source, target, event) => {
       system._sendInspectionEvent({
@@ -307,6 +275,7 @@ export interface BaseInspectionEventProperties {
    * - For snapshot events, this is the `actorRef` of the snapshot.
    * - For event events, this is the target `actorRef` (recipient of event).
    * - For actor events, this is the `actorRef` of the registered actor.
+   * - For registration events, this is the `actorRef` of the registered actor.
    */
   actorRef: AnyActorRef;
 }
@@ -330,32 +299,20 @@ export interface InspectedActorEvent extends BaseInspectionEventProperties {
   type: '@xstate.actor';
 }
 
-export type InspectionEvent =
-  | InspectedSnapshotEvent
-  | InspectedEventEvent
-  | InspectedActorEvent;
-
-export interface RegisteredActorEvent<
-  TActorRef extends AnyActorRef,
-  TSystemId extends string = string
-> {
-  type: `@xstate.actor.register`;
-  systemId: TSystemId;
-  actorRef: TActorRef;
+export interface InspectedRegisterEvent extends BaseInspectionEventProperties {
+  type: `@xstate.register`;
+  systemId: string;
 }
 
-export interface UnregisteredActorEvent<
-  TActorRef extends AnyActorRef,
-  TSystemId extends string = string
-> {
-  type: `@xstate.actor.unregister`;
-  systemId: TSystemId;
-  actorRef: TActorRef;
+export interface InspectedUnregisterEvent
+  extends BaseInspectionEventProperties {
+  type: `@xstate.unregister`;
+  systemId: string;
 }
 
-export type RegistrationEvent<
-  TActorRef extends AnyActorRef = AnyActorRef,
-  TSystemId extends string = string
-> =
-  | RegisteredActorEvent<TActorRef, TSystemId>
-  | UnregisteredActorEvent<TActorRef, TSystemId>;
+export type InspectionEvent =
+  | InspectedSnapshotEvent
+  | InspectedEventEvent
+  | InspectedActorEvent
+  | InspectedRegisterEvent
+  | InspectedUnregisterEvent;
diff --git a/packages/core/test/inspect.test.ts b/packages/core/test/inspect.test.ts
index b0a22a09934..b66c1b03473 100644
--- a/packages/core/test/inspect.test.ts
+++ b/packages/core/test/inspect.test.ts
@@ -404,4 +404,133 @@ describe('inspect', () => {
       ]
     `);
   });
+
+  it('should allow subscribing to system registration inspection events', () => {
+    const aSystemId = 'a_child';
+    const bSystemId = 'b_child';
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          invoke: {
+            src: createMachine({}),
+            systemId: aSystemId
+          },
+          on: {
+            to_b: 'b'
+          }
+        },
+        b: {
+          invoke: {
+            src: createMachine({}),
+            systemId: bSystemId
+          },
+          on: {
+            to_a: 'a'
+          }
+        }
+      }
+    });
+
+    const actorRef = createActor(machine).start();
+
+    const events: string[] = [];
+
+    actorRef.system.inspect((event) => {
+      if (
+        event.type === '@xstate.register' ||
+        event.type === '@xstate.unregister'
+      ) {
+        events.push(`${event.type}.${event.systemId}`);
+      }
+    });
+
+    actorRef.send({ type: 'to_b' });
+    expect(events).toEqual([
+      `@xstate.unregister.${aSystemId}`,
+      `@xstate.register.${bSystemId}`
+    ]);
+    actorRef.send({ type: 'to_a' });
+    expect(events).toEqual([
+      `@xstate.unregister.${aSystemId}`,
+      `@xstate.register.${bSystemId}`,
+      `@xstate.unregister.${bSystemId}`,
+      `@xstate.register.${aSystemId}`
+    ]);
+  });
+
+  it('should allow unsubscribing from a system registration inspection subscription', () => {
+    const aSystemId = 'a_child';
+    const bSystemId = 'b_child';
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          invoke: {
+            src: createMachine({}),
+            systemId: aSystemId
+          },
+          on: {
+            to_b: 'b'
+          }
+        },
+        b: {
+          invoke: {
+            src: createMachine({}),
+            systemId: bSystemId
+          },
+          on: {
+            to_a: 'a'
+          }
+        }
+      }
+    });
+
+    const actorRef = createActor(machine).start();
+
+    const events: string[] = [];
+    const unsubscribedEvents: string[] = [];
+
+    const subscription = actorRef.system.inspect((event) => {
+      if (
+        event.type === '@xstate.register' ||
+        event.type === '@xstate.unregister'
+      ) {
+        events.push(`${event.type}.${event.systemId}`);
+      }
+    });
+
+    actorRef.system.inspect((event) => {
+      if (
+        event.type === '@xstate.register' ||
+        event.type === '@xstate.unregister'
+      ) {
+        unsubscribedEvents.push(`${event.type}.${event.systemId}`);
+      }
+    });
+
+    actorRef.send({ type: 'to_b' });
+    expect(events).toEqual([
+      `@xstate.unregister.${aSystemId}`,
+      `@xstate.register.${bSystemId}`
+    ]);
+    expect(unsubscribedEvents).toEqual([
+      `@xstate.unregister.${aSystemId}`,
+      `@xstate.register.${bSystemId}`
+    ]);
+
+    subscription.unsubscribe();
+    actorRef.send({ type: 'to_a' });
+
+    expect(events).toEqual([
+      `@xstate.unregister.${aSystemId}`,
+      `@xstate.register.${bSystemId}`
+    ]);
+    expect(unsubscribedEvents).toEqual([
+      `@xstate.unregister.${aSystemId}`,
+      `@xstate.register.${bSystemId}`,
+      `@xstate.unregister.${bSystemId}`,
+      `@xstate.register.${aSystemId}`
+    ]);
+  });
 });
diff --git a/packages/core/test/system.test.ts b/packages/core/test/system.test.ts
index 6393a38a9d8..dd95672d37e 100644
--- a/packages/core/test/system.test.ts
+++ b/packages/core/test/system.test.ts
@@ -535,118 +535,4 @@ describe('system', () => {
 
     expect(spy).toHaveBeenCalledTimes(1);
   });
-
-  it('should allow subscribing to system registration events', () => {
-    const aSystemId = 'a_child';
-    const bSystemId = 'b_child';
-    const machine = createMachine({
-      initial: 'a',
-      states: {
-        a: {
-          invoke: {
-            src: createMachine({}),
-            systemId: aSystemId
-          },
-          on: {
-            to_b: 'b'
-          }
-        },
-        b: {
-          invoke: {
-            src: createMachine({}),
-            systemId: bSystemId
-          },
-          on: {
-            to_a: 'a'
-          }
-        }
-      }
-    });
-
-    const actorRef = createActor(machine).start();
-
-    const events: string[] = [];
-
-    actorRef.system.subscribe((event) => {
-      events.push(`${event.type}.${event.systemId}`);
-    });
-
-    actorRef.send({ type: 'to_b' });
-    expect(events).toEqual([
-      `@xstate.actor.unregister.${aSystemId}`,
-      `@xstate.actor.register.${bSystemId}`
-    ]);
-    actorRef.send({ type: 'to_a' });
-    expect(events).toEqual([
-      `@xstate.actor.unregister.${aSystemId}`,
-      `@xstate.actor.register.${bSystemId}`,
-      `@xstate.actor.unregister.${bSystemId}`,
-      `@xstate.actor.register.${aSystemId}`
-    ]);
-  });
-
-  it('should allow unsubscribing from a system registration subscription', () => {
-    const aSystemId = 'a_child';
-    const bSystemId = 'b_child';
-    const machine = createMachine({
-      initial: 'a',
-      states: {
-        a: {
-          invoke: {
-            src: createMachine({}),
-            systemId: aSystemId
-          },
-          on: {
-            to_b: 'b'
-          }
-        },
-        b: {
-          invoke: {
-            src: createMachine({}),
-            systemId: bSystemId
-          },
-          on: {
-            to_a: 'a'
-          }
-        }
-      }
-    });
-
-    const actorRef = createActor(machine).start();
-
-    const events: string[] = [];
-    const unsubscribedEvents: string[] = [];
-
-    const subscription = actorRef.system.subscribe((event) => {
-      events.push(`${event.type}.${event.systemId}`);
-    });
-
-    actorRef.system.subscribe((event) => {
-      unsubscribedEvents.push(`${event.type}.${event.systemId}`);
-    });
-
-    actorRef.send({ type: 'to_b' });
-    expect(events).toEqual([
-      `@xstate.actor.unregister.${aSystemId}`,
-      `@xstate.actor.register.${bSystemId}`
-    ]);
-    expect(unsubscribedEvents).toEqual([
-      `@xstate.actor.unregister.${aSystemId}`,
-      `@xstate.actor.register.${bSystemId}`
-    ]);
-
-    subscription.unsubscribe();
-    actorRef.send({ type: 'to_a' });
-
-    expect(events).toEqual([
-      `@xstate.actor.unregister.${aSystemId}`,
-      `@xstate.actor.register.${bSystemId}`
-    ]);
-    expect(unsubscribedEvents).toEqual([
-      `@xstate.actor.unregister.${aSystemId}`,
-      `@xstate.actor.register.${bSystemId}`,
-      `@xstate.actor.unregister.${bSystemId}`,
-      `@xstate.actor.register.${aSystemId}`
-    ]);
-  });
 });