From c58b36dc35991e20ba5d3c6e212e075f9b27f37d Mon Sep 17 00:00:00 2001
From: David Khourshid <davidkpiano@gmail.com>
Date: Tue, 18 Jun 2024 14:04:04 +0200
Subject: [PATCH] [core] Inspect subscription (#4936)

* Make actor.system.inspect return a subscription

* Changeset
---
 .changeset/cold-planes-clap.md     | 24 ++++++++
 packages/core/src/system.ts        | 20 ++++++-
 packages/core/test/inspect.test.ts | 90 ++++++++++++++++++++++++++++++
 3 files changed, 131 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/cold-planes-clap.md

diff --git a/.changeset/cold-planes-clap.md b/.changeset/cold-planes-clap.md
new file mode 100644
index 0000000000..27ff74d03a
--- /dev/null
+++ b/.changeset/cold-planes-clap.md
@@ -0,0 +1,24 @@
+---
+'xstate': minor
+---
+
+Inspecting an actor system via `actor.system.inspect(ev => …)` now accepts a function or observer, and returns a subscription:
+
+```ts
+const actor = createActor(someMachine);
+
+const sub = actor.system.inspect((inspectionEvent) => {
+  console.log(inspectionEvent);
+});
+
+// Inspection events will be logged
+actor.start();
+actor.send({ type: 'anEvent' });
+
+// ...
+
+sub.unsubscribe();
+
+// Will no longer log inspection events
+actor.send({ type: 'someEvent' });
+```
diff --git a/packages/core/src/system.ts b/packages/core/src/system.ts
index 288d15b815..b96d584861 100644
--- a/packages/core/src/system.ts
+++ b/packages/core/src/system.ts
@@ -7,8 +7,10 @@ import {
   Snapshot,
   HomomorphicOmit,
   EventObject,
-  AnyTransitionDefinition
+  AnyTransitionDefinition,
+  Subscription
 } from './types.ts';
+import { toObserver } from './utils.ts';
 
 export interface ScheduledEvent {
   id: string;
@@ -63,7 +65,12 @@ export interface ActorSystem<T extends ActorSystemInfo> {
    */
   _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>
+      | ((inspectionEvent: InspectionEvent) => void)
+  ) => Subscription;
   /**
    * @internal
    */
@@ -206,8 +213,15 @@ export function createSystem<T extends ActorSystemInfo>(
       keyedActors.set(systemId, actorRef);
       reverseKeyedActors.set(actorRef, systemId);
     },
-    inspect: (observer) => {
+    inspect: (observerOrFn) => {
+      const observer = toObserver(observerOrFn);
       inspectionObservers.add(observer);
+
+      return {
+        unsubscribe() {
+          inspectionObservers.delete(observer);
+        }
+      };
     },
     _sendInspectionEvent: sendInspectionEvent as any,
     _relay: (source, target, event) => {
diff --git a/packages/core/test/inspect.test.ts b/packages/core/test/inspect.test.ts
index 96702f23a5..b87b6fee39 100644
--- a/packages/core/test/inspect.test.ts
+++ b/packages/core/test/inspect.test.ts
@@ -1032,4 +1032,94 @@ describe('inspect', () => {
     actor.start();
     actor.send({ type: 'any' });
   });
+
+  it('actor.system.inspect(…) can inspect actors', () => {
+    const actor = createActor(createMachine({}));
+    const events: InspectionEvent[] = [];
+
+    actor.system.inspect((ev) => {
+      events.push(ev);
+    });
+
+    actor.start();
+
+    expect(events).toContainEqual(
+      expect.objectContaining({
+        type: '@xstate.event'
+      })
+    );
+    expect(events).toContainEqual(
+      expect.objectContaining({
+        type: '@xstate.snapshot'
+      })
+    );
+  });
+
+  it('actor.system.inspect(…) can inspect actors (observer)', () => {
+    const actor = createActor(createMachine({}));
+    const events: InspectionEvent[] = [];
+
+    actor.system.inspect({
+      next: (ev) => {
+        events.push(ev);
+      }
+    });
+
+    actor.start();
+
+    expect(events).toContainEqual(
+      expect.objectContaining({
+        type: '@xstate.event'
+      })
+    );
+    expect(events).toContainEqual(
+      expect.objectContaining({
+        type: '@xstate.snapshot'
+      })
+    );
+  });
+
+  it('actor.system.inspect(…) can be unsubscribed', () => {
+    const actor = createActor(createMachine({}));
+    const events: InspectionEvent[] = [];
+
+    const sub = actor.system.inspect((ev) => {
+      events.push(ev);
+    });
+
+    actor.start();
+
+    expect(events.length).toEqual(2);
+
+    events.length = 0;
+
+    sub.unsubscribe();
+
+    actor.send({ type: 'someEvent' });
+
+    expect(events.length).toEqual(0);
+  });
+
+  it('actor.system.inspect(…) can be unsubscribed (observer)', () => {
+    const actor = createActor(createMachine({}));
+    const events: InspectionEvent[] = [];
+
+    const sub = actor.system.inspect({
+      next: (ev) => {
+        events.push(ev);
+      }
+    });
+
+    actor.start();
+
+    expect(events.length).toEqual(2);
+
+    events.length = 0;
+
+    sub.unsubscribe();
+
+    actor.send({ type: 'someEvent' });
+
+    expect(events.length).toEqual(0);
+  });
 });