@kiberon-labs/behave-graph-suspendable 1.0.1

package/README.md

@@ -0,0 +1,13 @@
+# Behave-Graph Suspendable
+
+This is an extension to behaviour graphs to implement a suspendable engine. While behaviour graphs are great, they have issues with serializing them during an existing execution. This implements a suspendable engine, which relies on nodes implementing the suspendable interface. The engine can be suspended mid execution, serialized and then unsuspended.
+
+This is a work in progress as there are a number of edge cases that need to be accounted for, namely serialzing nodes that don't implement the required interface as well as hydration issues.
+
+To use it you can do the following :
+
+```ts
+import 
+
+```
+

package/package.json

@@ -0,0 +1,41 @@
+{
+  "name": "@kiberon-labs/behave-graph-suspendable",
+  "description": "An auxiliary engine for the behaviour graph that can be suspended mid execution.",
+  "version": "1.0.1",
+  "type": "module",
+  "types": "./dist/index.d.ts",
+  "main": "./dist/index.js",
+  "source": "./src/index.ts",
+  "author": "kiberonlabsdev",
+  "scripts": {
+    "build": "tsdown",
+    "check": "oxfmt --check",
+    "format": "oxfmt",
+    "test": "vitest --passWithNoTests",
+    "dev": "tsdown --watch src",
+    "lint": "oxlint",
+    "typecheck": "tsc"
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/kiberon-labs/behave-graph",
+    "directory": "packages/scene"
+  },
+  "bugs": "https://github.com/kiberon-labs/behave-graph/issues",
+  "devDependencies": {
+    "@kiberon-labs/behave-graph": "workspace:*",
+    "@types/node": "^24.10.1",
+    "oxfmt": "^0.14.0",
+    "oxlint": "^1.29.0",
+    "tsdown": "^0.16.5",
+    "vitest": "^4.0.10"
+  },
+  "dependencies": {},
+  "license": "ISC",
+  "peerDependencies": {
+    "@kiberon-labs/behave-graph": "workspace:*"
+  },
+  "publishConfig": {
+    "access": "public"
+  }
+}

package/src/engine.ts

@@ -0,0 +1,237 @@
+import {
+  Assert,
+  Engine,
+  EventEmitter,
+  isAsyncNode,
+  isEventNode,
+  Link,
+  isFlowNode,
+  type FiberListenerInner,
+  type GraphInstance,
+  type INode,
+  type IRegistry
+} from '@kiberon-labs/behave-graph';
+import { SuspendableFiber } from './fiber';
+import { isSuspendable, type SerializedSuspension } from './types';
+
+type StatefulNode = INode & {
+  getState(): unknown;
+  setState(value: unknown): void;
+};
+
+/** Concrete nodes carry getState/setState (from the Node base class) even
+ * though INode does not declare them. */
+const hasNodeState = (node: INode): node is StatefulNode => {
+  const candidate = node as Partial<StatefulNode>;
+  return (
+    typeof candidate.getState === 'function' &&
+    typeof candidate.setState === 'function'
+  );
+};
+
+export class SuspendableEngine extends Engine {
+  protected override fiberQueue: SuspendableFiber[] = [];
+  public readonly onSuspension = new EventEmitter<SerializedSuspension>();
+
+  constructor(graph: GraphInstance, registry: IRegistry) {
+    super(graph, registry);
+  }
+
+  //Note this is a weak point in the system to completely override the existing logic
+  override commitToNewFiber(
+    node: INode,
+    outputFlowSocketName: string,
+    fiberCompletedListener: FiberListenerInner = undefined
+  ) {
+    try {
+      Assert.mustBeTrue(isEventNode(node) || isAsyncNode(node));
+      const outputSocket = node.outputs.find(
+        (socket) => socket.name === outputFlowSocketName
+      );
+      if (outputSocket === undefined) {
+        throw new Error(`no socket with the name ${outputFlowSocketName}`);
+      }
+      if (outputSocket.links.length > 1) {
+        throw new Error(
+          'invalid for an output flow socket to have multiple downstream links:' +
+            `${node.description.typeName}.${outputSocket.name} has ${outputSocket.links.length} downlinks`
+        );
+      }
+      if (outputSocket.links.length === 1) {
+        const fiber = new SuspendableFiber(
+          this,
+          outputSocket.links[0]!,
+          fiberCompletedListener
+        );
+        this.onNodeCommit.emit({ node, socket: outputFlowSocketName });
+
+        this.fiberQueue.push(fiber);
+      }
+    } catch (error) {
+      this.onNodeExecutionError.emit({ node, error });
+      throw error;
+    }
+  }
+
+  commitContinuedFiber(node: INode) {
+    const fiber = new SuspendableFiber(this, new Link(node.id, 'flow'));
+    this.fiberQueue.push(fiber);
+  }
+
+  suspend(): SerializedSuspension {
+    const serializedNodes: Record<string, any> = {};
+    Object.entries(this.nodes).forEach(([id, node]) => {
+      if (isSuspendable(node)) {
+        serializedNodes[id] = node.suspend();
+        return;
+      }
+      // Flow nodes that keep their cursor in node state (e.g. flow/forLoop,
+      // flow/sequence, flow/counter) are captured generically. Event/async
+      // node state can hold live listeners or timers, so those must opt in
+      // via ISuspendable instead.
+      if (isFlowNode(node) && hasNodeState(node)) {
+        const state = node.getState();
+        if (state !== undefined) {
+          serializedNodes[id] = structuredClone(state);
+        }
+      }
+    });
+
+    const serializedFiberQueue = this.fiberQueue.map((fiber) =>
+      fiber.serialize()
+    );
+
+    // Snapshot the downstream links of each node's output sockets so that
+    // connections made at runtime survive the suspension round-trip.
+    const sockets: SerializedSuspension['sockets'] = {};
+    Object.entries(this.nodes).forEach(([id, node]) => {
+      const linkedOutputs: SerializedSuspension['sockets'][string] = {};
+      node.outputs.forEach((socket) => {
+        if (socket.links.length > 0) {
+          linkedOutputs[socket.name] = socket.links.map((link) => ({
+            nodeId: link.nodeId,
+            socketName: link.socketName
+          }));
+        }
+      });
+      if (Object.keys(linkedOutputs).length > 0) {
+        sockets[id] = linkedOutputs;
+      }
+    });
+
+    // Snapshot non-flow output socket values: resumed flows may read an
+    // upstream output (e.g. a loop's `index`) before that node re-executes.
+    const socketValues: SerializedSuspension['socketValues'] = {};
+    Object.entries(this.nodes).forEach(([id, node]) => {
+      const values: Record<string, any> = {};
+      node.outputs.forEach((socket) => {
+        if (socket.valueTypeName === 'flow' || socket.value === undefined) {
+          return;
+        }
+        const valueType = this.registry.values[socket.valueTypeName];
+        values[socket.name] = valueType
+          ? valueType.serialize(socket.value)
+          : socket.value;
+      });
+      if (Object.keys(values).length > 0) {
+        socketValues[id] = values;
+      }
+    });
+
+    const variables: SerializedSuspension['variables'] = {};
+    Object.entries(this.graph.variables).forEach(([id, variable]) => {
+      const valueType = this.registry.values[variable.valueTypeName];
+      variables[id] = {
+        type: variable.valueTypeName,
+        value: valueType ? valueType.serialize(variable.get()) : variable.get()
+      };
+    });
+
+    return {
+      fiberQueue: serializedFiberQueue,
+      nodes: serializedNodes,
+      sockets,
+      socketValues,
+      variables
+    };
+  }
+
+  unsuspend(suspension: SerializedSuspension, continuanceData: any) {
+    Object.entries(suspension.variables).forEach(([id, serialized]) => {
+      const variable = this.graph.variables[id];
+      if (!variable) {
+        throw new Error(`Could not find missing variable ${id}`);
+      }
+      const valueType = this.registry.values[serialized.type];
+      variable.set(
+        valueType ? valueType.deserialize(serialized.value) : serialized.value
+      );
+    });
+
+    Object.entries(suspension.sockets).forEach(([nodeId, linkedOutputs]) => {
+      const node = this.nodes[nodeId];
+      if (!node) {
+        throw new Error(`Could not find missing node ${nodeId}`);
+      }
+      Object.entries(linkedOutputs).forEach(([socketName, links]) => {
+        const socket = node.outputs.find((s) => s.name === socketName);
+        if (!socket) {
+          throw new Error(
+            `Could not find missing socket ${socketName} on node ${nodeId}`
+          );
+        }
+        socket.links.splice(
+          0,
+          socket.links.length,
+          ...links.map((link) => new Link(link.nodeId, link.socketName))
+        );
+      });
+    });
+
+    Object.entries(suspension.socketValues ?? {}).forEach(
+      ([nodeId, values]) => {
+        const node = this.nodes[nodeId];
+        if (!node) {
+          throw new Error(`Could not find missing node ${nodeId}`);
+        }
+        Object.entries(values).forEach(([socketName, serialized]) => {
+          const socket = node.outputs.find((s) => s.name === socketName);
+          if (!socket) {
+            throw new Error(
+              `Could not find missing socket ${socketName} on node ${nodeId}`
+            );
+          }
+          const valueType = this.registry.values[socket.valueTypeName];
+          socket.value = valueType
+            ? valueType.deserialize(serialized)
+            : serialized;
+        });
+      }
+    );
+
+    Object.entries(suspension.nodes).forEach(([id, serialized]) => {
+      const node = this.nodes[id];
+      if (!node) {
+        throw new Error(`Could not find missing node ${id}`);
+      }
+
+      if (isSuspendable(node)) {
+        node.hydrate?.(serialized);
+      } else if (isFlowNode(node) && hasNodeState(node)) {
+        node.setState(serialized);
+      }
+    });
+
+    //Recreate the fiberqueue
+    this.fiberQueue = suspension.fiberQueue.map((x) =>
+      SuspendableFiber.rehydrate(this, x)
+    );
+
+    // Only fibers suspended at an async node need an explicit continuation;
+    // fibers suspended between steps resume through normal execution.
+    const first = this.fiberQueue[0];
+    if (first?.canContinue()) {
+      first.continue(continuanceData);
+    }
+  }
+}

package/src/fiber.ts

@@ -0,0 +1,162 @@
+import {
+  Assert,
+  Fiber,
+  isAsyncNode,
+  Link,
+  type FiberListenerInner,
+  type IAsyncNode
+} from '@kiberon-labs/behave-graph';
+import type { SuspendableEngine } from './engine';
+import { isSuspendable, type IAsyncSuspendable } from './types';
+
+type SerializedLink = {
+  nodeId: string;
+  socketName: string;
+};
+
+export type SerializedSuspendedFiber = {
+  curr?: SerializedLink;
+  link?: SerializedLink;
+  queue: string[];
+};
+
+export class SuspendableFiber extends Fiber {
+  protected currentLink?: Link;
+  declare public readonly engine: SuspendableEngine;
+
+  constructor(
+    engine: SuspendableEngine,
+    nextEval: Link | null,
+    fiberCompletedListener: FiberListenerInner = undefined
+  ) {
+    super(engine, nextEval, fiberCompletedListener);
+  }
+
+  static rehydrate(
+    engine: SuspendableEngine,
+    serialized: SerializedSuspendedFiber
+  ): SuspendableFiber {
+    const nextlink = serialized.link
+      ? new Link(serialized.link.nodeId, serialized.link.socketName)
+      : null;
+    const fiber = new SuspendableFiber(engine, nextlink);
+    fiber.currentLink = serialized.curr
+      ? new Link(serialized.curr.nodeId, serialized.curr.socketName)
+      : undefined;
+
+    //Rehydrate the listener stack
+    serialized.queue.forEach((nodeId) => {
+      fiber.fiberCompletedListenerStack.push({
+        nodeId,
+        cb: () => engine.commitContinuedFiber(engine.nodes[nodeId]!)
+      });
+    });
+
+    return fiber;
+  }
+
+  serialize(): SerializedSuspendedFiber {
+    return {
+      queue: this.fiberCompletedListenerStack
+        .map((l) => l.nodeId!)
+        .filter((id) => !!id),
+      curr: this.currentLink
+        ? {
+            nodeId: this.currentLink.nodeId,
+            socketName: this.currentLink.socketName
+          }
+        : undefined,
+      link: this.nextEval
+        ? {
+            nodeId: this.nextEval.nodeId,
+            socketName: this.nextEval.socketName
+          }
+        : undefined
+    };
+  }
+
+  clear() {
+    this.fiberCompletedListenerStack.splice(
+      0,
+      this.fiberCompletedListenerStack.length
+    );
+    this.currentLink = undefined;
+    this.nextEval = null;
+  }
+
+  /** Whether this fiber was suspended at an async node that must be resumed
+   * via `continue()`. Fibers suspended between steps have no current link and
+   * simply resume through normal engine execution. */
+  canContinue(): boolean {
+    return !!this.currentLink;
+  }
+
+  /**
+   * Track the link while executing an async suspendable node. Such nodes can
+   * park the fiber on a pending promise; if the engine is suspended while
+   * parked, `serialize()` must record where execution stopped so `continue()`
+   * can resume the node with continuance data. The link is cleared once the
+   * node completes normally.
+   */
+  override async executeStep() {
+    const link = this.nextEval;
+    if (link) {
+      const node = this.nodes[link.nodeId];
+      if (node && isAsyncNode(node) && isSuspendable(node)) {
+        this.currentLink = link;
+      }
+    }
+
+    await super.executeStep();
+
+    if (link && this.currentLink === link) {
+      this.currentLink = undefined;
+    }
+  }
+
+  continue(continuanceData: any) {
+    Assert.mustBeTrue(
+      !!this.currentLink,
+      'Cannot continue fiber if no current link'
+    );
+
+    const nodeID = this.currentLink?.nodeId;
+    const node = this.nodes[nodeID!] as IAsyncSuspendable;
+
+    if (!node) {
+      throw new Error(`Missing node for continuance: ${nodeID}`);
+    }
+    const socketName = this.currentLink?.socketName;
+    Assert.mustBeTrue(
+      !!socketName,
+      'Cannot continue fiber if no current socket name'
+    );
+
+    this.engine.onNodeExecutionStart.emit(node!);
+    return node.unsuspend(continuanceData, this.engine, socketName!, () => {
+      //Only on successful unsuspend do we clear the current link
+      //This is to allow a node to immediately re-suspend if needed
+      this.currentLink = undefined;
+      this.engine.onNodeExecutionEnd.emit(node);
+
+      //remove from the list of pending async nodes (it will not be present
+      //on a freshly rehydrated engine, where the node was never triggered)
+      const index = this.engine.asyncNodes.indexOf(
+        node as unknown as IAsyncNode
+      );
+      if (index >= 0) {
+        this.engine.asyncNodes.splice(index, 1);
+      }
+      this.executionSteps++;
+
+      //There might be pending callbacks on the queue
+      while (this.fiberCompletedListenerStack.length > 0) {
+        const awaitingCallback = this.fiberCompletedListenerStack.pop();
+        if (awaitingCallback === undefined) {
+          throw new Error('awaitingCallback is empty');
+        }
+        awaitingCallback.cb();
+      }
+    });
+  }
+}

package/src/index.ts

@@ -0,0 +1 @@
+export * from './engine';

package/src/types.ts

@@ -0,0 +1,67 @@
+import type { INode } from '@kiberon-labs/behave-graph';
+import type { SuspendableEngine } from './engine';
+import type { SerializedSuspendedFiber } from './fiber';
+
+export interface IAsyncSuspendable<T = any, State = any>
+  extends ISuspendable<State> {
+  /**
+   * Continue from suspension
+   */
+  unsuspend(
+    data: T,
+    engine: SuspendableEngine,
+    triggeringSocketName: string,
+    cb: () => void
+  ): Promise<void>;
+  triggered(
+    engine: SuspendableEngine,
+    triggeringSocketName: string,
+    finished: () => void
+  ): unknown;
+}
+
+export interface ISuspendable<T = any> extends INode {
+  /**
+   * Suspend and create a copy of your internal state
+   */
+  suspend(): T;
+  /**
+   * Rehydrate from suspended data
+   */
+  hydrate?(data: T): void;
+}
+
+export function isSuspendable(node: unknown): node is ISuspendable {
+  if (node == null || typeof node !== 'object') {
+    return false;
+  }
+  const candidate = node as Record<string, unknown>;
+  // ISuspendable only requires suspend(); unsuspend() belongs to
+  // IAsyncSuspendable and is checked at the point of continuation.
+  return typeof candidate.suspend === 'function';
+}
+
+export type SerializedSuspension = {
+  fiberQueue: SerializedSuspendedFiber[];
+  /**
+   * Serialized state of the nodes
+   */
+  nodes: Record<string, any>;
+  sockets: Record<
+    string,
+    Record<string, { nodeId: string; socketName: string }[]>
+  >;
+  /**
+   * Serialized output socket values per node. Resumed flows may read an
+   * upstream node's output (e.g. a loop's `index`) before that node is
+   * re-triggered, so the values present at suspension time must be restored.
+   */
+  socketValues: Record<string, Record<string, any>>;
+  variables: Record<
+    string,
+    {
+      type: string;
+      value: any;
+    }
+  >;
+};

package/tests/asyncNodeSuspension.test.ts

@@ -0,0 +1,124 @@
+import type { GraphJSON } from '@kiberon-labs/behave-graph';
+import { describe, expect, it } from 'vitest';
+import { makeRecorderDescription, WaitForSignal } from './fixtures/testNodes';
+import { makeTestEngine } from './testUtils';
+
+const COMPLETED = 999;
+
+/**
+ * onStart -> forLoop [0,1) -> waitForSignal (async, parks the fiber) ->
+ * recorder reading the delivered signal value. The orphan time/delay node is
+ * present only to assert async nodes are excluded from generic state capture.
+ */
+const waitGraph: GraphJSON = {
+  variables: [],
+  customEvents: [],
+  nodes: [
+    {
+      type: 'lifecycle/onStart',
+      id: '0',
+      flows: { flow: { nodeId: 'loop', socket: 'flow' } }
+    },
+    {
+      type: 'flow/forLoop',
+      id: 'loop',
+      parameters: {
+        startIndex: { value: 0 },
+        endIndex: { value: 1 }
+      },
+      flows: {
+        loopBody: { nodeId: 'wait', socket: 'flow' },
+        completed: { nodeId: 'done', socket: 'flow' }
+      }
+    },
+    {
+      type: 'test/waitForSignal',
+      id: 'wait',
+      flows: { flow: { nodeId: 'rec', socket: 'flow' } }
+    },
+    {
+      type: 'test/recorder',
+      id: 'rec',
+      parameters: { index: { link: { nodeId: 'wait', socket: 'value' } } }
+    },
+    {
+      type: 'test/recorder',
+      id: 'done',
+      parameters: { index: { value: COMPLETED } }
+    },
+    {
+      type: 'time/delay',
+      id: 'orphanDelay'
+    }
+  ]
+};
+
+const makeWaitEngine = (record: (index: number) => void) =>
+  makeTestEngine(waitGraph, {
+    'test/recorder': makeRecorderDescription(record),
+    [WaitForSignal.Description.typeName]: WaitForSignal.Description
+  });
+
+describe('async node suspension (vs flow nodes)', () => {
+  it('parks the fiber at the async node and resumes via unsuspend with continuance data', async () => {
+    const firstRun: number[] = [];
+    const first = makeWaitEngine((i) => firstRun.push(i));
+
+    first.start();
+    // do not await: the async node parks execution on a pending promise
+    void first.engine.executeAllSync();
+    const waitNode = first.graph.nodes['wait'] as WaitForSignal;
+    await waitNode.parked;
+
+    // nothing downstream of the wait has run
+    expect(firstRun).toEqual([]);
+
+    const suspension = first.engine.suspend();
+    first.engine.dispose();
+
+    // ASYNC nodes are captured through ISuspendable.suspend()...
+    expect(suspension.nodes['wait']).toEqual({ waiting: true });
+    // ...and the fiber records the async node as its resume point, unlike
+    // flow-node suspensions which carry no current link
+    expect(suspension.fiberQueue[0]!.curr).toEqual({
+      nodeId: 'wait',
+      socketName: 'flow'
+    });
+    // the enclosing loop's continuation is preserved alongside it
+    expect(suspension.fiberQueue[0]!.queue).toEqual(['loop']);
+
+    const persisted = JSON.parse(JSON.stringify(suspension));
+
+    // resume on a fresh engine, delivering the awaited data as continuance
+    const secondRun: number[] = [];
+    const second = makeWaitEngine((i) => secondRun.push(i));
+    second.engine.unsuspend(persisted, 42);
+    await second.engine.executeAllSync();
+
+    // the delivered data flows downstream of the wait node, then the
+    // enclosing loop finishes without restarting
+    expect(secondRun).toEqual([42, COMPLETED]);
+  });
+
+  it('excludes async and event node state from generic capture', async () => {
+    const firstRun: number[] = [];
+    const first = makeWaitEngine((i) => firstRun.push(i));
+
+    first.start();
+    void first.engine.executeAllSync();
+    await (first.graph.nodes['wait'] as WaitForSignal).parked;
+
+    const suspension = first.engine.suspend();
+    first.engine.dispose();
+
+    // flow nodes are captured generically from node state
+    expect(suspension.nodes['loop']).toEqual({ nextIndex: 1 });
+
+    // event nodes hold live listeners — never serialized
+    expect(suspension.nodes['0']).toBeUndefined();
+
+    // async nodes can hold timers, so they are only captured when they opt in
+    // via ISuspendable. time/delay has node state but no suspend() — excluded.
+    expect(suspension.nodes['orphanDelay']).toBeUndefined();
+  });
+});