ai-sdk-graph 0.5.0 → 0.6.0

package/dist/compiled-graph.d.ts

@@ -9,16 +9,20 @@ export declare class CompiledGraph<State extends Record<string, unknown>, NodeKe
     private readonly edgeRegistry;
     private readonly subgraphRegistry;
     private readonly storage;
-    private readonly emitter;
     private readonly stateManager;
-    private readonly onFinish;
-    private readonly onStart;
+    private readonly graphMiddleware;
+    private readonly nodeMiddleware;
+    private readonly stateMiddleware;
+    private readonly eventMiddleware;
     constructor(nodeRegistry: ReadonlyMap<NodeKeys, GraphSDK.Node<State, NodeKeys>>, edgeRegistry: ReadonlyMap<NodeKeys, GraphSDK.Edge<State, NodeKeys>[]>, subgraphRegistry: ReadonlyMap<NodeKeys, {
         subgraph: Graph<any, any>;
         options: GraphSDK.SubgraphOptions<State, any>;
-    }>, options?: GraphSDK.CompileOptions<State, NodeKeys>);
-    execute(runId: string, initialState: State | ((state: State | undefined) => State)): ReadableStream<import("ai").InferUIMessageChunk<import("ai").UIMessage<unknown, import("ai").UIDataTypes, import("ai").UITools>>>;
-    executeInternal(runId: string, initialState: State, writer: GraphSDK.Writer): Promise<State>;
+    }>, options?: GraphSDK.CompileOptions<State, NodeKeys>, graphMiddleware?: GraphSDK.GraphMiddleware<State, NodeKeys>[], nodeMiddleware?: GraphSDK.NodeMiddleware<State, NodeKeys>[], stateMiddleware?: GraphSDK.StateMiddleware<State, NodeKeys>[], eventMiddleware?: GraphSDK.EventMiddleware<State, NodeKeys>[]);
+    stream(runId: string, initialState: State | ((state: State | undefined) => State)): ReadableStream<import("ai").InferUIMessageChunk<import("ai").UIMessage<unknown, import("ai").UIDataTypes, import("ai").UITools>>>;
+    execute(runId: string, initialState: State | ((state: State | undefined) => State), options?: {
+        onEvent?: (event: GraphSDK.GraphEvent<State, NodeKeys>) => void;
+    }): Promise<State>;
+    executeInternal(runId: string, initialState: State, writer: GraphSDK.Writer, emit: (event: GraphSDK.GraphEvent<State, NodeKeys>) => void): Promise<State>;
     private createExecutionContext;
     private runExecutionLoop;
     private runExecutionLoopInternal;

package/dist/graph.d.ts

@@ -5,12 +5,13 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     private readonly nodeRegistry;
     private readonly edgeRegistry;
     private readonly subgraphRegistry;
+    private readonly middlewares;
     constructor();
     node<NewKey extends string>(id: NewKey, execute: ({ state, writer, suspense, update }: {
         state: () => Readonly<State>;
         writer: GraphSDK.Writer;
         suspense: (data?: unknown) => never;
-        update: (update: GraphSDK.StateUpdate<State>) => void;
+        update: (update: GraphSDK.StateUpdate<State>) => Promise<void>;
     }) => Promise<void> | void): Graph<State, NodeKeys | NewKey>;
     edge(from: NodeKeys, to: NodeKeys | ((state: State) => NodeKeys)): Graph<State, NodeKeys>;
     graph<NewKey extends string, ChildState extends Record<string, unknown>, ChildNodeKeys extends string = 'START' | 'END'>(id: NewKey, subgraph: Graph<ChildState, ChildNodeKeys>, options: GraphSDK.SubgraphOptions<State, ChildState>): Graph<State, NodeKeys | NewKey>;
@@ -20,6 +21,7 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
         subgraph: Graph<any, any>;
         options: GraphSDK.SubgraphOptions<State, any>;
     }>;
+    use(middleware: GraphSDK.Middleware<State, NodeKeys>): Graph<State, NodeKeys>;
     compile(options?: GraphSDK.CompileOptions<State, NodeKeys>): CompiledGraph<State, NodeKeys>;
     toMermaid(options?: {
         direction?: 'TB' | 'LR';

package/dist/index.d.ts

@@ -3,3 +3,4 @@ export * from './compiled-graph';
 export * from './storage';
 export * from './types';
 export * from './utils';
+export * from './middleware';

package/dist/index.js

@@ -33,6 +33,21 @@ class RedisStorage {
 
 // src/compiled-graph.ts
 import { createUIMessageStream } from "ai";
+
+// src/middleware.ts
+function composeMiddleware(middlewares, action) {
+  return (ctx) => {
+    function dispatch(i) {
+      if (i === middlewares.length) {
+        return action(ctx);
+      }
+      return middlewares[i](ctx, () => dispatch(i + 1));
+    }
+    return dispatch(0);
+  };
+}
+
+// src/compiled-graph.ts
 var BUILT_IN_NODES = {
   START: "START",
   END: "END"
@@ -46,51 +61,108 @@ class SuspenseError extends Error {
     this.data = data;
   }
 }
+var nullWriter = {
+  write() {},
+  merge() {}
+};
+function composeEventMiddleware(middlewares, terminal) {
+  return (event) => {
+    function dispatch(i) {
+      if (i === middlewares.length) {
+        terminal(event);
+        return;
+      }
+      middlewares[i](event, () => dispatch(i + 1));
+    }
+    dispatch(0);
+  };
+}
 
 class CompiledGraph {
   nodeRegistry;
   edgeRegistry;
   subgraphRegistry;
   storage;
-  emitter = new NodeEventEmitter;
   stateManager = new StateManager;
-  onFinish;
-  onStart;
-  constructor(nodeRegistry, edgeRegistry, subgraphRegistry, options = {}) {
+  graphMiddleware;
+  nodeMiddleware;
+  stateMiddleware;
+  eventMiddleware;
+  constructor(nodeRegistry, edgeRegistry, subgraphRegistry, options = {}, graphMiddleware = [], nodeMiddleware = [], stateMiddleware = [], eventMiddleware = []) {
     this.nodeRegistry = nodeRegistry;
     this.edgeRegistry = edgeRegistry;
     this.subgraphRegistry = subgraphRegistry;
     this.storage = options.storage ?? new InMemoryStorage;
-    this.onFinish = options.onFinish ?? (() => {});
-    this.onStart = options.onStart ?? (() => {});
+    this.graphMiddleware = graphMiddleware;
+    this.nodeMiddleware = nodeMiddleware;
+    this.stateMiddleware = stateMiddleware;
+    this.eventMiddleware = eventMiddleware;
   }
-  execute(runId, initialState) {
+  stream(runId, initialState) {
     let context;
     return createUIMessageStream({
       execute: async ({ writer }) => {
-        const result = await this.createExecutionContext(runId, initialState, writer);
+        const emit = composeEventMiddleware(this.eventMiddleware, (event) => {
+          switch (event.type) {
+            case "state":
+              writer.write({ type: "data-state", data: event.state });
+              break;
+            case "node:start":
+              writer.write({ type: "data-node-start", data: event.nodeId });
+              break;
+            case "node:end":
+              writer.write({ type: "data-node-end", data: event.nodeId });
+              break;
+            case "node:suspense":
+              writer.write({ type: "data-node-suspense", data: { nodeId: event.nodeId, data: event.data } });
+              break;
+          }
+        });
+        const result = await this.createExecutionContext(runId, initialState, emit, writer);
         context = result.context;
         const firstTime = result.firstTime;
-        if (firstTime) {
-          await this.onStart({ state: context.state, writer });
-        }
-        await this.runExecutionLoop(context);
-      },
-      onFinish: async () => {
-        if (context) {
-          await this.onFinish({ state: context.state });
+        if (this.graphMiddleware.length > 0) {
+          const graphCtx = {
+            runId: context.runId,
+            state: () => context.state,
+            writer: context.writer,
+            isResume: !firstTime
+          };
+          await composeMiddleware(this.graphMiddleware, async () => {
+            await this.runExecutionLoop(context);
+          })(graphCtx);
+        } else {
+          await this.runExecutionLoop(context);
         }
       }
     });
   }
-  async executeInternal(runId, initialState, writer) {
-    const { context } = await this.createExecutionContext(runId, initialState, writer);
+  async execute(runId, initialState, options) {
+    const emit = composeEventMiddleware(this.eventMiddleware, options?.onEvent ?? (() => {}));
+    const { context, firstTime } = await this.createExecutionContext(runId, initialState, emit, nullWriter);
+    if (this.graphMiddleware.length > 0) {
+      const graphCtx = {
+        runId: context.runId,
+        state: () => context.state,
+        writer: context.writer,
+        isResume: !firstTime
+      };
+      await composeMiddleware(this.graphMiddleware, async () => {
+        await this.runExecutionLoop(context);
+      })(graphCtx);
+    } else {
+      await this.runExecutionLoop(context);
+    }
+    return context.state;
+  }
+  async executeInternal(runId, initialState, writer, emit) {
+    const { context } = await this.createExecutionContext(runId, initialState, emit, writer);
     await this.runExecutionLoopInternal(context);
     return context.state;
   }
-  async createExecutionContext(runId, initialState, writer) {
-    const { context, firstTime } = await this.restoreCheckpoint(runId, initialState, writer);
-    return { context: { ...context, runId, writer }, firstTime };
+  async createExecutionContext(runId, initialState, emit, writer) {
+    const { context, firstTime } = await this.restoreCheckpoint(runId, initialState, emit);
+    return { context: { ...context, runId, writer, emit }, firstTime };
   }
   async runExecutionLoop(context) {
     await this.executeWithStrategy(context, "return");
@@ -139,12 +211,12 @@ class CompiledGraph {
   hasNodesToExecute(context) {
     return context.currentNodes.length > 0;
   }
-  async restoreCheckpoint(runId, initialState, writer) {
+  async restoreCheckpoint(runId, initialState, emit) {
     const checkpoint = await this.storage.load(runId);
     if (this.isValidCheckpoint(checkpoint)) {
-      return { context: this.restoreFromCheckpoint(checkpoint, initialState, writer), firstTime: false };
+      return { context: this.restoreFromCheckpoint(checkpoint, initialState, emit), firstTime: false };
     }
-    return { context: this.createFreshExecution(initialState, writer), firstTime: true };
+    return { context: this.createFreshExecution(initialState, emit), firstTime: true };
   }
   isValidCheckpoint(checkpoint) {
     return this.hasNodeIds(checkpoint) && this.hasAtLeastOneNode(checkpoint);
@@ -155,16 +227,16 @@ class CompiledGraph {
   hasAtLeastOneNode(checkpoint) {
     return checkpoint.nodeIds.length > 0;
   }
-  restoreFromCheckpoint(checkpoint, initialState, writer) {
-    const state = this.stateManager.resolve(initialState, checkpoint.state, writer);
+  restoreFromCheckpoint(checkpoint, initialState, emit) {
+    const state = this.stateManager.resolve(initialState, checkpoint.state, emit);
     return {
       state,
       currentNodes: this.resolveNodeIds(checkpoint.nodeIds),
       suspendedNodes: this.resolveNodeIds(checkpoint.suspendedNodes)
     };
   }
-  createFreshExecution(initialState, writer) {
-    const state = this.stateManager.resolve(initialState, undefined, writer);
+  createFreshExecution(initialState, emit) {
+    const state = this.stateManager.resolve(initialState, undefined, emit);
     const startNode = this.nodeRegistry.get(BUILT_IN_NODES.START);
     return {
       state,
@@ -188,43 +260,110 @@ class CompiledGraph {
     if (subgraphEntry) {
       return this.executeSubgraphNode(node, context, subgraphEntry);
     }
-    this.emitter.emitStart(context.writer, node.id);
+    const isBuiltIn = node.id === BUILT_IN_NODES.START || node.id === BUILT_IN_NODES.END;
+    context.emit({ type: "node:start", nodeId: node.id });
     try {
-      await node.execute(this.createNodeExecutionParams(context));
-      this.emitter.emitEnd(context.writer, node.id);
+      const { params, pendingUpdates } = this.createNodeExecutionParams(context, node.id);
+      if (!isBuiltIn && this.nodeMiddleware.length > 0) {
+        const nodeCtx = {
+          runId: context.runId,
+          nodeId: node.id,
+          state: () => context.state,
+          writer: context.writer,
+          isSubgraph: false
+        };
+        await composeMiddleware(this.nodeMiddleware, async () => {
+          await node.execute(params);
+        })(nodeCtx);
+      } else {
+        await node.execute(params);
+      }
+      await Promise.all(pendingUpdates);
+      context.emit({ type: "node:end", nodeId: node.id });
       return null;
     } catch (error) {
       if (error instanceof SuspenseError) {
-        this.emitter.emitSuspense(context.writer, node.id, error.data);
+        context.emit({ type: "node:suspense", nodeId: node.id, data: error.data });
         return { node, error };
       }
       throw error;
     }
   }
-  createNodeExecutionParams(context) {
-    return {
+  createNodeExecutionParams(context, nodeId = null) {
+    const pendingUpdates = [];
+    const params = {
       state: () => context.state,
       writer: context.writer,
       suspense: this.createSuspenseFunction(),
       update: (update) => {
-        context.state = this.stateManager.apply(context.state, update, context.writer);
+        const p = (async () => {
+          if (this.stateMiddleware.length > 0) {
+            const resolvedUpdate = typeof update === "function" ? update(context.state) : update;
+            const stateCtx = {
+              runId: context.runId,
+              nodeId,
+              currentState: context.state,
+              update,
+              resolvedUpdate
+            };
+            const finalPartial = await composeMiddleware(this.stateMiddleware, async (ctx) => ctx.resolvedUpdate)(stateCtx);
+            context.state = { ...context.state, ...finalPartial };
+            context.emit({ type: "state", state: context.state });
+          } else {
+            context.state = this.stateManager.apply(context.state, update);
+            context.emit({ type: "state", state: context.state });
+          }
+        })();
+        pendingUpdates.push(p);
+        return p;
       }
     };
+    return { params, pendingUpdates };
   }
   async executeSubgraphNode(node, context, entry) {
     const { subgraph, options } = entry;
-    this.emitter.emitStart(context.writer, node.id);
+    context.emit({ type: "node:start", nodeId: node.id });
     const subgraphRunId = this.generateSubgraphRunId(context.runId, node.id);
     try {
-      const childRunner = new CompiledGraph(subgraph.nodes, subgraph.edges, subgraph.subgraphs, { storage: this.storage });
-      const childFinalState = await childRunner.executeInternal(subgraphRunId, options.input(context.state), context.writer);
-      const parentUpdate = options.output(childFinalState, context.state);
-      context.state = this.stateManager.apply(context.state, parentUpdate, context.writer);
-      this.emitter.emitEnd(context.writer, node.id);
+      const executeSubgraph = async () => {
+        const childRunner = new CompiledGraph(subgraph.nodes, subgraph.edges, subgraph.subgraphs, { storage: this.storage }, [], this.nodeMiddleware, this.stateMiddleware, this.eventMiddleware);
+        const childFinalState = await childRunner.executeInternal(subgraphRunId, options.input(context.state), context.writer, context.emit);
+        const parentUpdate = options.output(childFinalState, context.state);
+        if (this.stateMiddleware.length > 0) {
+          const resolvedUpdate = typeof parentUpdate === "function" ? parentUpdate(context.state) : parentUpdate;
+          const stateCtx = {
+            runId: context.runId,
+            nodeId: node.id,
+            currentState: context.state,
+            update: parentUpdate,
+            resolvedUpdate
+          };
+          const finalPartial = await composeMiddleware(this.stateMiddleware, async (ctx) => ctx.resolvedUpdate)(stateCtx);
+          context.state = { ...context.state, ...finalPartial };
+        } else {
+          context.state = this.stateManager.apply(context.state, parentUpdate);
+        }
+        context.emit({ type: "state", state: context.state });
+      };
+      if (this.nodeMiddleware.length > 0) {
+        const nodeCtx = {
+          runId: context.runId,
+          nodeId: node.id,
+          state: () => context.state,
+          writer: context.writer,
+          isSubgraph: true
+        };
+        await composeMiddleware(this.nodeMiddleware, async () => {
+          await executeSubgraph();
+        })(nodeCtx);
+      } else {
+        await executeSubgraph();
+      }
+      context.emit({ type: "node:end", nodeId: node.id });
       return null;
     } catch (error) {
       if (error instanceof SuspenseError) {
-        this.emitter.emitSuspense(context.writer, node.id, error.data);
+        context.emit({ type: "node:suspense", nodeId: node.id, data: error.data });
         return { node, error };
       }
       throw error;
@@ -262,35 +401,21 @@ class CompiledGraph {
   }
 }
 
-class NodeEventEmitter {
-  emitStart(writer, nodeId) {
-    writer.write({ type: "data-node-start", data: nodeId });
-  }
-  emitEnd(writer, nodeId) {
-    writer.write({ type: "data-node-end", data: nodeId });
-  }
-  emitSuspense(writer, nodeId, data) {
-    writer.write({ type: "data-node-suspense", data: { nodeId, data } });
-  }
-}
-
 class StateManager {
-  apply(state, update, writer) {
-    const newState = {
+  apply(state, update) {
+    return {
       ...state,
       ...typeof update === "function" ? update(state) : update
     };
-    writer.write({ type: "data-state", data: newState });
-    return newState;
   }
-  resolve(initialState, existingState, writer) {
+  resolve(initialState, existingState, emit) {
     if (this.isStateFactory(initialState)) {
       const newState2 = initialState(existingState);
-      writer.write({ type: "data-state", data: newState2 });
+      emit({ type: "state", state: newState2 });
       return newState2;
     }
     const newState = existingState ?? initialState;
-    writer.write({ type: "data-state", data: newState });
+    emit({ type: "state", state: newState });
     return newState;
   }
   isStateFactory(initialState) {
@@ -303,6 +428,7 @@ class Graph {
   nodeRegistry = new Map;
   edgeRegistry = new Map;
   subgraphRegistry = new Map;
+  middlewares = [];
   constructor() {
     this.registerBuiltInNodes();
   }
@@ -334,8 +460,26 @@ class Graph {
   get subgraphs() {
     return this.subgraphRegistry;
   }
+  use(middleware) {
+    this.middlewares.push(middleware);
+    return this;
+  }
   compile(options = {}) {
-    return new CompiledGraph(this.nodeRegistry, this.edgeRegistry, this.subgraphRegistry, options);
+    const graphMiddleware = [];
+    const nodeMiddleware = [];
+    const stateMiddleware = [];
+    const eventMiddleware = [];
+    for (const mw of this.middlewares) {
+      if (mw.graph)
+        graphMiddleware.push(mw.graph);
+      if (mw.node)
+        nodeMiddleware.push(mw.node);
+      if (mw.state)
+        stateMiddleware.push(mw.state);
+      if (mw.event)
+        eventMiddleware.push(mw.event);
+    }
+    return new CompiledGraph(this.nodeRegistry, this.edgeRegistry, this.subgraphRegistry, options, graphMiddleware, nodeMiddleware, stateMiddleware, eventMiddleware);
   }
   toMermaid(options) {
     const generator = new MermaidGenerator(this.nodeRegistry, this.edgeRegistry, this.subgraphRegistry);
@@ -468,6 +612,7 @@ export {
   isGraphDataPart,
   graph,
   consumeAndMergeStream,
+  composeMiddleware,
   SuspenseError,
   RedisStorage,
   InMemoryStorage,

package/dist/middleware.d.ts

@@ -0,0 +1 @@
+export declare function composeMiddleware<Ctx, R = void>(middlewares: Array<(ctx: Ctx, next: () => Promise<R>) => Promise<R>>, action: (ctx: Ctx) => Promise<R>): (ctx: Ctx) => Promise<R>;

package/dist/types.d.ts

@@ -3,7 +3,7 @@ export declare namespace GraphSDK {
     type StateUpdate<State> = Partial<State> | ((state: State) => Partial<State>);
     interface SubgraphOptions<ParentState extends Record<string, unknown>, ChildState extends Record<string, unknown>> {
         input: (parentState: ParentState) => ChildState;
-        output: (childState: ChildState, parentState: ParentState) => Partial<ParentState>;
+        output: (childState: ChildState, parentState: ParentState) => StateUpdate<ParentState>;
     }
     interface Graph<State extends Record<string, unknown>, NodeKeys extends string> {
         nodes: Map<NodeKeys, Node<State, NodeKeys>>;
@@ -15,7 +15,7 @@ export declare namespace GraphSDK {
             state: () => Readonly<State>;
             writer: UIMessageStreamWriter;
             suspense: (data?: unknown) => never;
-            update: (update: StateUpdate<State>) => void;
+            update: (update: StateUpdate<State>) => Promise<void>;
         }) => Promise<void> | void;
     }
     interface Edge<State extends Record<string, unknown>, NodeKeys extends string> {
@@ -37,27 +37,58 @@ export declare namespace GraphSDK {
         state: State;
         currentNodes: Node<State, NodeKeys>[];
         suspendedNodes: Node<State, NodeKeys>[];
-        writer: UIMessageStreamWriter;
+        writer: Writer;
+        emit: (event: GraphEvent<State, NodeKeys>) => void;
     }
     interface GraphOptions<State extends Record<string, unknown>, NodeKeys extends string> {
         storage?: GraphStorage<State, NodeKeys>;
-        onFinish?: (args: {
-            state: State;
-        }) => Promise<void> | void;
-        onStart?: (args: {
-            state: State;
-            writer: UIMessageStreamWriter;
-        }) => Promise<void> | void;
     }
     interface CompileOptions<State extends Record<string, unknown>, NodeKeys extends string> {
         storage?: GraphStorage<State, NodeKeys>;
-        onFinish?: (args: {
-            state: State;
-        }) => Promise<void> | void;
-        onStart?: (args: {
-            state: State;
-            writer: UIMessageStreamWriter;
-        }) => Promise<void> | void;
     }
     type Writer = Parameters<Parameters<typeof createUIMessageStream>[0]['execute']>[0]['writer'];
+    interface GraphMiddlewareContext<State extends Record<string, unknown>, NodeKeys extends string> {
+        runId: string;
+        state: () => Readonly<State>;
+        writer: Writer;
+        isResume: boolean;
+    }
+    interface NodeMiddlewareContext<State extends Record<string, unknown>, NodeKeys extends string> {
+        runId: string;
+        nodeId: NodeKeys;
+        state: () => Readonly<State>;
+        writer: Writer;
+        isSubgraph: boolean;
+    }
+    interface StateMiddlewareContext<State extends Record<string, unknown>, NodeKeys extends string> {
+        runId: string;
+        nodeId: NodeKeys | null;
+        currentState: Readonly<State>;
+        update: StateUpdate<State>;
+        resolvedUpdate: Partial<State>;
+    }
+    type GraphMiddleware<State extends Record<string, unknown>, NodeKeys extends string> = (ctx: GraphMiddlewareContext<State, NodeKeys>, next: () => Promise<void>) => Promise<void>;
+    type NodeMiddleware<State extends Record<string, unknown>, NodeKeys extends string> = (ctx: NodeMiddlewareContext<State, NodeKeys>, next: () => Promise<void>) => Promise<void>;
+    type StateMiddleware<State extends Record<string, unknown>, NodeKeys extends string> = (ctx: StateMiddlewareContext<State, NodeKeys>, next: () => Promise<Partial<State>>) => Promise<Partial<State>>;
+    type GraphEvent<State extends Record<string, unknown>, NodeKeys extends string> = {
+        type: 'state';
+        state: State;
+    } | {
+        type: 'node:start';
+        nodeId: NodeKeys;
+    } | {
+        type: 'node:end';
+        nodeId: NodeKeys;
+    } | {
+        type: 'node:suspense';
+        nodeId: NodeKeys;
+        data: unknown;
+    };
+    type EventMiddleware<State extends Record<string, unknown>, NodeKeys extends string> = (event: GraphEvent<State, NodeKeys>, next: () => void) => void;
+    interface Middleware<State extends Record<string, unknown>, NodeKeys extends string> {
+        graph?: GraphMiddleware<State, NodeKeys>;
+        node?: NodeMiddleware<State, NodeKeys>;
+        state?: StateMiddleware<State, NodeKeys>;
+        event?: EventMiddleware<State, NodeKeys>;
+    }
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-sdk-graph",
-  "version": "0.5.0",
+  "version": "0.6.0",
   "description": "Graph-based workflows for the AI SDK",
   "type": "module",
   "main": "dist/index.js",