ai-sdk-graph 0.1.2 → 0.1.4

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Can Temizyurek
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/dist/graph.d.ts

@@ -13,10 +13,8 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     private readonly emitter;
     private readonly stateManager;
     private readonly onFinish;
-    constructor(options?: {
-        storage?: GraphSDK.GraphStorage<State, NodeKeys>;
-        onFinish?: (state: State) => Promise<void> | void;
-    });
+    private readonly onStart;
+    constructor(options?: GraphSDK.GraphOptions<State, NodeKeys>);
     node<NewKey extends string>(id: NewKey, execute: ({ state, writer, suspense, update }: {
         state: () => Readonly<State>;
         writer: Writer;
@@ -34,8 +32,6 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     toMermaid(options?: {
         direction?: 'TB' | 'LR';
     }): string;
-    private generateMermaid;
-    private extractPossibleTargets;
     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: Writer): Promise<State>;
     private registerBuiltInNodes;
@@ -69,8 +65,5 @@ export declare class Graph<State extends Record<string, unknown>, NodeKeys exten
     private deduplicateNodes;
     private excludeTerminalNodes;
 }
-export declare function graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'>(options?: {
-    storage?: GraphSDK.GraphStorage<State, NodeKeys>;
-    onFinish?: (state: State) => Promise<void> | void;
-}): Graph<State, NodeKeys>;
+export declare function graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'>(options?: GraphSDK.GraphOptions<State, NodeKeys>): Graph<State, NodeKeys>;
 export {};

package/dist/index.js

@@ -55,10 +55,12 @@ class Graph {
   emitter = new NodeEventEmitter;
   stateManager = new StateManager;
   onFinish;
+  onStart;
   constructor(options = {}) {
     this.storage = options.storage ?? new InMemoryStorage;
     this.registerBuiltInNodes();
-    this.onFinish = options.onFinish ?? ((state) => {});
+    this.onFinish = options.onFinish ?? (() => {});
+    this.onStart = options.onStart ?? (() => {});
   }
   node(id, execute) {
     const node = { id, execute };
@@ -89,78 +91,30 @@ class Graph {
     return this.subgraphRegistry;
   }
   toMermaid(options) {
-    const direction = options?.direction ?? "TB";
-    return this.generateMermaid(direction, "");
-  }
-  generateMermaid(direction, prefix) {
-    const lines = [];
-    const indent = prefix ? "        " : "    ";
-    if (!prefix) {
-      lines.push(`flowchart ${direction}`);
-    }
-    for (const [nodeId] of this.nodeRegistry) {
-      const prefixedId = prefix ? `${prefix}_${nodeId}` : nodeId;
-      const subgraphEntry = this.subgraphRegistry.get(nodeId);
-      if (subgraphEntry) {
-        lines.push(`${indent}subgraph ${prefixedId}[${nodeId}]`);
-        lines.push(`${indent}    direction ${direction}`);
-        const subgraphContent = subgraphEntry.subgraph.generateMermaid(direction, prefixedId);
-        const subgraphLines = subgraphContent.split(`
-`);
-        lines.push(...subgraphLines.map((line) => `${indent}${line}`));
-        lines.push(`${indent}end`);
-      } else if (nodeId === "START" || nodeId === "END") {
-        lines.push(`${indent}${prefixedId}([${nodeId}])`);
-      } else {
-        lines.push(`${indent}${prefixedId}[${nodeId}]`);
-      }
-    }
-    for (const [fromId, edges] of this.edgeRegistry) {
-      for (const edge of edges) {
-        const prefixedFrom = prefix ? `${prefix}_${fromId}` : fromId;
-        if (typeof edge.to === "function") {
-          const possibleTargets = this.extractPossibleTargets(edge.to);
-          for (const targetId of possibleTargets) {
-            const prefixedTo = prefix ? `${prefix}_${targetId}` : targetId;
-            lines.push(`${indent}${prefixedFrom} -.-> ${prefixedTo}`);
-          }
-        } else {
-          const prefixedTo = prefix ? `${prefix}_${edge.to}` : edge.to;
-          lines.push(`${indent}${prefixedFrom} --> ${prefixedTo}`);
-        }
-      }
-    }
-    return lines.join(`
-`);
-  }
-  extractPossibleTargets(edgeFn) {
-    const fnString = edgeFn.toString();
-    const nodeIds = Array.from(this.nodeRegistry.keys());
-    return nodeIds.filter((nodeId) => {
-      const patterns = [
-        `'${nodeId}'`,
-        `"${nodeId}"`,
-        `\`${nodeId}\``
-      ];
-      return patterns.some((pattern) => fnString.includes(pattern));
-    });
+    const generator = new MermaidGenerator(this.nodeRegistry, this.edgeRegistry, this.subgraphRegistry);
+    return generator.generate(options);
   }
   execute(runId, initialState) {
     let context;
     return createUIMessageStream({
       execute: async ({ writer }) => {
-        context = await this.createExecutionContext(runId, initialState, writer);
+        const result = await this.createExecutionContext(runId, initialState, 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(context.state);
+          await this.onFinish({ state: context.state });
         }
       }
     });
   }
   async executeInternal(runId, initialState, writer) {
-    const context = await this.createExecutionContext(runId, initialState, writer);
+    const { context } = await this.createExecutionContext(runId, initialState, writer);
     await this.runExecutionLoopInternal(context);
     return context.state;
   }
@@ -174,8 +128,8 @@ class Graph {
     this.edgeRegistry.set(edge.from, existingEdges);
   }
   async createExecutionContext(runId, initialState, writer) {
-    const restored = await this.restoreCheckpoint(runId, initialState);
-    return { runId, writer, ...restored };
+    const { context, firstTime } = await this.restoreCheckpoint(runId, initialState);
+    return { context: { ...context, runId, writer }, firstTime };
   }
   async runExecutionLoop(context) {
     await this.executeWithStrategy(context, "return");
@@ -227,9 +181,9 @@ class Graph {
   async restoreCheckpoint(runId, initialState) {
     const checkpoint = await this.storage.load(runId);
     if (this.isValidCheckpoint(checkpoint)) {
-      return this.restoreFromCheckpoint(checkpoint, initialState);
+      return { context: this.restoreFromCheckpoint(checkpoint, initialState), firstTime: false };
     }
-    return this.createFreshExecution(initialState);
+    return { context: this.createFreshExecution(initialState), firstTime: true };
   }
   isValidCheckpoint(checkpoint) {
     return this.hasNodeIds(checkpoint) && this.hasAtLeastOneNode(checkpoint);
@@ -346,10 +300,7 @@ class Graph {
   }
 }
 function graph(options = {}) {
-  return new Graph({
-    storage: options.storage,
-    onFinish: options.onFinish
-  });
+  return new Graph(options);
 }
 
 class NodeEventEmitter {
@@ -381,6 +332,75 @@ class StateManager {
     return typeof initialState === "function";
   }
 }
+
+class MermaidGenerator {
+  nodeRegistry;
+  edgeRegistry;
+  subgraphRegistry;
+  constructor(nodeRegistry, edgeRegistry, subgraphRegistry) {
+    this.nodeRegistry = nodeRegistry;
+    this.edgeRegistry = edgeRegistry;
+    this.subgraphRegistry = subgraphRegistry;
+  }
+  generate(options) {
+    const direction = options?.direction ?? "TB";
+    return this.generateMermaid(direction, "");
+  }
+  generateMermaid(direction, prefix) {
+    const lines = [];
+    const indent = prefix ? "        " : "    ";
+    if (!prefix) {
+      lines.push(`flowchart ${direction}`);
+    }
+    for (const [nodeId] of this.nodeRegistry) {
+      const prefixedId = prefix ? `${prefix}_${nodeId}` : nodeId;
+      const subgraphEntry = this.subgraphRegistry.get(nodeId);
+      if (subgraphEntry) {
+        lines.push(`${indent}subgraph ${prefixedId}[${nodeId}]`);
+        lines.push(`${indent}    direction ${direction}`);
+        const subgraphGenerator = new MermaidGenerator(subgraphEntry.subgraph.nodes, subgraphEntry.subgraph.edges, subgraphEntry.subgraph.subgraphs);
+        const subgraphContent = subgraphGenerator.generateMermaid(direction, prefixedId);
+        const subgraphLines = subgraphContent.split(`
+`);
+        lines.push(...subgraphLines.map((line) => `${indent}${line}`));
+        lines.push(`${indent}end`);
+      } else if (nodeId === "START" || nodeId === "END") {
+        lines.push(`${indent}${prefixedId}([${nodeId}])`);
+      } else {
+        lines.push(`${indent}${prefixedId}[${nodeId}]`);
+      }
+    }
+    for (const [fromId, edges] of this.edgeRegistry) {
+      for (const edge of edges) {
+        const prefixedFrom = prefix ? `${prefix}_${fromId}` : fromId;
+        if (typeof edge.to === "function") {
+          const possibleTargets = this.extractPossibleTargets(edge.to);
+          for (const targetId of possibleTargets) {
+            const prefixedTo = prefix ? `${prefix}_${targetId}` : targetId;
+            lines.push(`${indent}${prefixedFrom} -.-> ${prefixedTo}`);
+          }
+        } else {
+          const prefixedTo = prefix ? `${prefix}_${edge.to}` : edge.to;
+          lines.push(`${indent}${prefixedFrom} --> ${prefixedTo}`);
+        }
+      }
+    }
+    return lines.join(`
+`);
+  }
+  extractPossibleTargets(edgeFn) {
+    const fnString = edgeFn.toString();
+    const nodeIds = Array.from(this.nodeRegistry.keys());
+    return nodeIds.filter((nodeId) => {
+      const patterns = [
+        `'${nodeId}'`,
+        `"${nodeId}"`,
+        `\`${nodeId}\``
+      ];
+      return patterns.some((pattern) => fnString.includes(pattern));
+    });
+  }
+}
 export {
   graph,
   SuspenseError,

package/dist/types.d.ts

@@ -39,4 +39,14 @@ export declare namespace GraphSDK {
         suspendedNodes: Node<State, NodeKeys>[];
         writer: UIMessageStreamWriter;
     }
+    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;
+    }
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-sdk-graph",
-  "version": "0.1.2",
+  "version": "0.1.4",
   "description": "Graph-based workflows for the AI SDK",
   "type": "module",
   "main": "dist/index.js",
@@ -25,7 +25,8 @@
   },
   "devDependencies": {
     "@changesets/cli": "^2.29.8",
-    "@types/bun": "latest"
+    "@types/bun": "latest",
+    "prettier": "^3.8.0"
   },
   "peerDependencies": {
     "typescript": "^5"