ai-sdk-graph 0.1.0

package/README.md

@@ -0,0 +1,213 @@
+# ai-sdk-graph
+
+A TypeScript library for building stateful, resumable workflows with human-in-the-loop support.
+
+## Features
+
+- **Human in the Loop** — Suspend execution to wait for user input, approvals, or external data, then resume seamlessly
+- **Type-safe** — Full TypeScript support with generic state management
+- **Resumable** — Checkpoint and resume execution from any point
+- **Composable** — Nested subgraphs with state mapping
+- **Parallel Execution** — Fork workflows with multiple edges from a single node
+- **Conditional Routing** — Dynamic edges based on state
+- **Visualization** — Generate Mermaid diagrams of your workflows
+
+## Installation
+
+```bash
+npm i ai-sdk-graph
+```
+
+## Quick Start
+
+```typescript
+import { graph } from 'ai-sdk-graph'
+
+const workflow = graph<{ value: number }>()
+  .node('validate', ({ update }) => {
+    update({ value: 10 })
+  })
+  .node('transform', ({ update, state }) => {
+    update({ value: state().value * 2 })
+  })
+  .edge('START', 'validate')
+  .edge('validate', 'transform')
+  .edge('transform', 'END')
+
+// Execute the workflow
+const stream = workflow.execute('run-1', { value: 0 })
+```
+
+## Core Concepts
+
+### Nodes
+
+Nodes are execution units that receive a context object with:
+
+- `state()` — Read the current state
+- `update(changes)` — Update state with partial object or function
+- `suspense(data?)` — Pause execution for human-in-the-loop
+- `writer` — Stream writer for UI integration
+
+```typescript
+graph<{ count: number }>()
+  .node('increment', ({ state, update }) => {
+    update({ count: state().count + 1 })
+  })
+```
+
+### Edges
+
+Connect nodes with static or dynamic routing:
+
+```typescript
+// Static edge
+.edge('START', 'validate')
+
+// Dynamic edge based on state
+.edge('router', (state) => state.isValid ? 'process' : 'reject')
+```
+
+### State Management
+
+State is type-safe and immutable. Updates can be partial objects or functions:
+
+```typescript
+// Partial update
+update({ status: 'complete' })
+
+// Functional update
+update((state) => ({ count: state.count + 1 }))
+```
+
+### Human in the Loop
+
+Suspend execution to wait for user input, approvals, or external data:
+
+```typescript
+const workflow = graph<{ approved: boolean }>()
+  .node('review', ({ state, suspense }) => {
+    if (!state().approved) {
+      suspense({ reason: 'Waiting for approval' })
+    }
+  })
+  .edge('START', 'review')
+  .edge('review', 'END')
+
+// First execution suspends
+workflow.execute('run-1', { approved: false })
+
+// Resume with updated state after user approves
+workflow.execute('run-1', (existing) => ({ ...existing, approved: true }))
+```
+
+## Parallel Execution
+
+Multiple edges from the same node execute targets in parallel:
+
+```typescript
+graph<{ results: string[] }>()
+  .node('fork', () => {})
+  .node('taskA', ({ update }) => update({ results: ['A'] }))
+  .node('taskB', ({ update }) => update({ results: ['B'] }))
+  .node('join', () => {})
+  .edge('START', 'fork')
+  .edge('fork', 'taskA')  // Both taskA and taskB
+  .edge('fork', 'taskB')  // execute in parallel
+  .edge('taskA', 'join')
+  .edge('taskB', 'join')
+  .edge('join', 'END')
+```
+
+## Subgraphs
+
+Compose workflows with nested graphs:
+
+```typescript
+const childGraph = graph<{ value: number }>()
+  .node('double', ({ update, state }) => {
+    update({ value: state().value * 2 })
+  })
+  .edge('START', 'double')
+  .edge('double', 'END')
+
+const parentGraph = graph<{ input: number; result: number }>()
+  .graph('process', childGraph, {
+    input: (parentState) => ({ value: parentState.input }),
+    output: (childState) => ({ result: childState.value })
+  })
+  .edge('START', 'process')
+  .edge('process', 'END')
+```
+
+## Storage
+
+By default, graphs use in-memory storage. For production, use Redis:
+
+```typescript
+import { graph } from 'ai-sdk-graph'
+import { RedisStorage } from 'ai-sdk-graph/storage'
+import Redis from 'ioredis'
+
+const redis = new Redis()
+const storage = new RedisStorage(redis)
+
+const workflow = graph<State>(storage)
+  // ... define nodes and edges
+```
+
+## Visualization
+
+Generate Mermaid diagrams of your workflows:
+
+```typescript
+const diagram = workflow.toMermaid()
+// or with direction
+const diagram = workflow.toMermaid({ direction: 'LR' })
+```
+
+Output:
+```mermaid
+flowchart TB
+    START([START])
+    validate[validate]
+    transform[transform]
+    END([END])
+    START --> validate
+    validate --> transform
+    transform --> END
+```
+
+## API Reference
+
+### `graph<State>(storage?)`
+
+Create a new graph with optional storage backend.
+
+### `.node(id, handler)`
+
+Add a node with an execution handler.
+
+### `.edge(from, to)`
+
+Add a static edge between nodes.
+
+### `.edge(from, (state) => nodeId)`
+
+Add a dynamic edge that routes based on state.
+
+### `.graph(id, subgraph, options)`
+
+Add a nested subgraph with state mapping.
+
+### `.execute(runId, initialState)`
+
+Execute the graph and return a readable stream.
+
+### `.toMermaid(options?)`
+
+Generate a Mermaid flowchart diagram.
+
+## License
+
+MIT

package/dist/graph.d.ts

@@ -0,0 +1,69 @@
+import type { GraphSDK } from './types';
+import { createUIMessageStream } from 'ai';
+export declare class SuspenseError extends Error {
+    readonly data?: unknown;
+    constructor(data?: unknown);
+}
+type Writer = Parameters<Parameters<typeof createUIMessageStream>[0]['execute']>[0]['writer'];
+export declare class Graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'> {
+    private readonly nodeRegistry;
+    private readonly edgeRegistry;
+    private readonly subgraphRegistry;
+    private readonly storage;
+    private readonly emitter;
+    private readonly stateManager;
+    constructor(storage?: GraphSDK.GraphStorage<State, NodeKeys>);
+    node<NewKey extends string>(id: NewKey, execute: ({ state, writer, suspense, update }: {
+        state: () => Readonly<State>;
+        writer: Writer;
+        suspense: (data?: unknown) => never;
+        update: (update: GraphSDK.StateUpdate<State>) => 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>;
+    get nodes(): ReadonlyMap<NodeKeys, GraphSDK.Node<State, NodeKeys>>;
+    get edges(): ReadonlyMap<NodeKeys, GraphSDK.Edge<State, NodeKeys>[]>;
+    get subgraphs(): ReadonlyMap<NodeKeys, {
+        subgraph: Graph<any, any>;
+        options: GraphSDK.SubgraphOptions<State, any>;
+    }>;
+    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;
+    private addEdgeToRegistry;
+    private createExecutionContext;
+    private runExecutionLoop;
+    private runExecutionLoopInternal;
+    private executeWithStrategy;
+    private resumeWithStrategy;
+    private executeNodesWithStrategy;
+    private handleBatchResultWithStrategy;
+    private hasSuspendedNodes;
+    private hasNodesToExecute;
+    private restoreCheckpoint;
+    private isValidCheckpoint;
+    private hasNodeIds;
+    private hasAtLeastOneNode;
+    private restoreFromCheckpoint;
+    private createFreshExecution;
+    private persistCheckpoint;
+    private executeBatch;
+    private executeSingleNode;
+    private createNodeExecutionParams;
+    private executeSubgraphNode;
+    private generateSubgraphRunId;
+    private createSuspenseFunction;
+    private resolveNodeIds;
+    private computeNextNodes;
+    private findSuccessors;
+    private resolveEdgeTarget;
+    private deduplicateNodes;
+    private excludeTerminalNodes;
+}
+export declare function graph<State extends Record<string, unknown>, NodeKeys extends string = 'START' | 'END'>(storage?: GraphSDK.GraphStorage<State, NodeKeys>): Graph<State, NodeKeys>;
+export {};

package/dist/index.d.ts

@@ -0,0 +1,3 @@
+export * from './graph';
+export * from './storage';
+export * from './types';

package/dist/index.js

@@ -0,0 +1,379 @@
+// src/storage.ts
+import Redis from "ioredis";
+
+class InMemoryStorage {
+  store = new Map;
+  async save(runId, checkpoint) {
+    this.store.set(runId, checkpoint);
+  }
+  async load(runId) {
+    return this.store.get(runId) ?? null;
+  }
+  async delete(runId) {
+    this.store.delete(runId);
+  }
+}
+
+class RedisStorage {
+  redis;
+  constructor(redisUrl) {
+    this.redis = new Redis(redisUrl);
+  }
+  async save(runId, checkpoint) {
+    await this.redis.set(runId, JSON.stringify(checkpoint));
+  }
+  async load(runId) {
+    const data = await this.redis.get(runId);
+    return data ? JSON.parse(data) : null;
+  }
+  async delete(runId) {
+    await this.redis.del(runId);
+  }
+}
+
+// src/graph.ts
+import { createUIMessageStream } from "ai";
+var BUILT_IN_NODES = {
+  START: "START",
+  END: "END"
+};
+
+class SuspenseError extends Error {
+  data;
+  constructor(data) {
+    super("Suspense");
+    this.name = "SuspenseError";
+    this.data = data;
+  }
+}
+
+class Graph {
+  nodeRegistry = new Map;
+  edgeRegistry = new Map;
+  subgraphRegistry = new Map;
+  storage;
+  emitter = new NodeEventEmitter;
+  stateManager = new StateManager;
+  constructor(storage = new InMemoryStorage) {
+    this.storage = storage;
+    this.registerBuiltInNodes();
+  }
+  node(id, execute) {
+    const node = { id, execute };
+    this.nodeRegistry.set(node.id, node);
+    return this;
+  }
+  edge(from, to) {
+    const edge = { from, to };
+    this.addEdgeToRegistry(edge);
+    return this;
+  }
+  graph(id, subgraph, options) {
+    this.subgraphRegistry.set(id, { subgraph, options });
+    const node = {
+      id,
+      execute: async () => {}
+    };
+    this.nodeRegistry.set(node.id, node);
+    return this;
+  }
+  get nodes() {
+    return this.nodeRegistry;
+  }
+  get edges() {
+    return this.edgeRegistry;
+  }
+  get subgraphs() {
+    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));
+    });
+  }
+  execute(runId, initialState) {
+    return createUIMessageStream({
+      execute: async ({ writer }) => {
+        const context = await this.createExecutionContext(runId, initialState, writer);
+        await this.runExecutionLoop(context);
+      }
+    });
+  }
+  async executeInternal(runId, initialState, writer) {
+    const context = await this.createExecutionContext(runId, initialState, writer);
+    await this.runExecutionLoopInternal(context);
+    return context.state;
+  }
+  registerBuiltInNodes() {
+    this.node(BUILT_IN_NODES.START, () => {});
+    this.node(BUILT_IN_NODES.END, () => {});
+  }
+  addEdgeToRegistry(edge) {
+    const existingEdges = this.edgeRegistry.get(edge.from) ?? [];
+    existingEdges.push(edge);
+    this.edgeRegistry.set(edge.from, existingEdges);
+  }
+  async createExecutionContext(runId, initialState, writer) {
+    const restored = await this.restoreCheckpoint(runId, initialState);
+    return { runId, writer, ...restored };
+  }
+  async runExecutionLoop(context) {
+    await this.executeWithStrategy(context, "return");
+  }
+  async runExecutionLoopInternal(context) {
+    await this.executeWithStrategy(context, "throw");
+  }
+  async executeWithStrategy(context, strategy) {
+    if (this.hasSuspendedNodes(context)) {
+      const shouldContinue = await this.resumeWithStrategy(context, strategy);
+      if (!shouldContinue)
+        return;
+    }
+    await this.executeNodesWithStrategy(context, strategy);
+  }
+  async resumeWithStrategy(context, strategy) {
+    const suspenses = await this.executeBatch(context.suspendedNodes, context);
+    return this.handleBatchResultWithStrategy(context, suspenses, strategy);
+  }
+  async executeNodesWithStrategy(context, strategy) {
+    while (this.hasNodesToExecute(context)) {
+      const suspenses = await this.executeBatch(context.currentNodes, context);
+      const shouldContinue = await this.handleBatchResultWithStrategy(context, suspenses, strategy);
+      if (!shouldContinue)
+        return;
+    }
+    await this.storage.delete(context.runId);
+  }
+  async handleBatchResultWithStrategy(context, suspenses, strategy) {
+    if (suspenses.length > 0) {
+      context.suspendedNodes = suspenses.map((s) => s.node);
+      await this.persistCheckpoint(context);
+      if (strategy === "throw") {
+        throw new SuspenseError({ type: "subgraph-suspended" });
+      }
+      return false;
+    }
+    context.currentNodes = this.computeNextNodes(context.currentNodes, context.state);
+    context.suspendedNodes = [];
+    await this.persistCheckpoint(context);
+    return true;
+  }
+  hasSuspendedNodes(context) {
+    return context.suspendedNodes.length > 0;
+  }
+  hasNodesToExecute(context) {
+    return context.currentNodes.length > 0;
+  }
+  async restoreCheckpoint(runId, initialState) {
+    const checkpoint = await this.storage.load(runId);
+    if (this.isValidCheckpoint(checkpoint)) {
+      return this.restoreFromCheckpoint(checkpoint, initialState);
+    }
+    return this.createFreshExecution(initialState);
+  }
+  isValidCheckpoint(checkpoint) {
+    return this.hasNodeIds(checkpoint) && this.hasAtLeastOneNode(checkpoint);
+  }
+  hasNodeIds(checkpoint) {
+    return checkpoint?.nodeIds != null;
+  }
+  hasAtLeastOneNode(checkpoint) {
+    return checkpoint.nodeIds.length > 0;
+  }
+  restoreFromCheckpoint(checkpoint, initialState) {
+    const state = this.stateManager.resolve(initialState, checkpoint.state);
+    return {
+      state,
+      currentNodes: this.resolveNodeIds(checkpoint.nodeIds),
+      suspendedNodes: this.resolveNodeIds(checkpoint.suspendedNodes)
+    };
+  }
+  createFreshExecution(initialState) {
+    const state = this.stateManager.resolve(initialState, undefined);
+    const startNode = this.nodeRegistry.get(BUILT_IN_NODES.START);
+    return {
+      state,
+      currentNodes: [startNode],
+      suspendedNodes: []
+    };
+  }
+  async persistCheckpoint(context) {
+    await this.storage.save(context.runId, {
+      state: context.state,
+      nodeIds: context.currentNodes.map((n) => n.id),
+      suspendedNodes: context.suspendedNodes.map((n) => n.id)
+    });
+  }
+  async executeBatch(nodes, context) {
+    const results = await Promise.all(nodes.map((node) => this.executeSingleNode(node, context)));
+    return results.filter((r) => r !== null);
+  }
+  async executeSingleNode(node, context) {
+    const subgraphEntry = this.subgraphRegistry.get(node.id);
+    if (subgraphEntry) {
+      return this.executeSubgraphNode(node, context, subgraphEntry);
+    }
+    this.emitter.emitStart(context.writer, node.id);
+    try {
+      await node.execute(this.createNodeExecutionParams(context));
+      this.emitter.emitEnd(context.writer, node.id);
+      return null;
+    } catch (error) {
+      if (error instanceof SuspenseError) {
+        this.emitter.emitSuspense(context.writer, node.id, error.data);
+        return { node, error };
+      }
+      throw error;
+    }
+  }
+  createNodeExecutionParams(context) {
+    return {
+      state: () => context.state,
+      writer: context.writer,
+      suspense: this.createSuspenseFunction(),
+      update: (update) => {
+        context.state = this.stateManager.apply(context.state, update);
+      }
+    };
+  }
+  async executeSubgraphNode(node, context, entry) {
+    const { subgraph, options } = entry;
+    this.emitter.emitStart(context.writer, node.id);
+    const subgraphRunId = this.generateSubgraphRunId(context.runId, node.id);
+    try {
+      const childFinalState = await subgraph.executeInternal(subgraphRunId, options.input(context.state), context.writer);
+      const parentUpdate = options.output(childFinalState, context.state);
+      context.state = this.stateManager.apply(context.state, parentUpdate);
+      this.emitter.emitEnd(context.writer, node.id);
+      return null;
+    } catch (error) {
+      if (error instanceof SuspenseError) {
+        this.emitter.emitSuspense(context.writer, node.id, error.data);
+        return { node, error };
+      }
+      throw error;
+    }
+  }
+  generateSubgraphRunId(parentRunId, nodeId) {
+    return `${parentRunId}:subgraph:${nodeId}`;
+  }
+  createSuspenseFunction() {
+    return (data) => {
+      throw new SuspenseError(data);
+    };
+  }
+  resolveNodeIds(nodeIds) {
+    return nodeIds.map((id) => this.nodeRegistry.get(id)).filter((node) => node != null);
+  }
+  computeNextNodes(currentNodes, state) {
+    const successors = currentNodes.flatMap((node) => this.findSuccessors(node.id, state));
+    const uniqueSuccessors = this.deduplicateNodes(successors);
+    return this.excludeTerminalNodes(uniqueSuccessors);
+  }
+  findSuccessors(nodeId, state) {
+    const outgoingEdges = this.edgeRegistry.get(nodeId) ?? [];
+    return outgoingEdges.map((edge) => this.resolveEdgeTarget(edge, state)).filter((node) => node != null);
+  }
+  resolveEdgeTarget(edge, state) {
+    const targetId = typeof edge.to === "function" ? edge.to(state) : edge.to;
+    return this.nodeRegistry.get(targetId);
+  }
+  deduplicateNodes(nodes) {
+    return [...new Set(nodes)];
+  }
+  excludeTerminalNodes(nodes) {
+    return nodes.filter((node) => node.id !== BUILT_IN_NODES.END);
+  }
+}
+function graph(storage) {
+  return new Graph(storage);
+}
+
+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) {
+    return {
+      ...state,
+      ...typeof update === "function" ? update(state) : update
+    };
+  }
+  resolve(initialState, existingState) {
+    if (this.isStateFactory(initialState)) {
+      return initialState(existingState);
+    }
+    return existingState ?? initialState;
+  }
+  isStateFactory(initialState) {
+    return typeof initialState === "function";
+  }
+}
+export {
+  graph,
+  SuspenseError,
+  RedisStorage,
+  InMemoryStorage,
+  Graph
+};

package/dist/storage.d.ts

@@ -0,0 +1,14 @@
+import type { GraphSDK } from './types';
+export declare class InMemoryStorage<State extends Record<string, unknown>, NodeKeys extends string> implements GraphSDK.GraphStorage<State, NodeKeys> {
+    private store;
+    save(runId: string, checkpoint: GraphSDK.Checkpoint<State, NodeKeys>): Promise<void>;
+    load(runId: string): Promise<GraphSDK.Checkpoint<State, NodeKeys> | null>;
+    delete(runId: string): Promise<void>;
+}
+export declare class RedisStorage<State extends Record<string, unknown>, NodeKeys extends string> implements GraphSDK.GraphStorage<State, NodeKeys> {
+    private redis;
+    constructor(redisUrl: string);
+    save(runId: string, checkpoint: GraphSDK.Checkpoint<State, NodeKeys>): Promise<void>;
+    load(runId: string): Promise<GraphSDK.Checkpoint<State, NodeKeys> | null>;
+    delete(runId: string): Promise<void>;
+}

package/dist/types.d.ts

@@ -0,0 +1,42 @@
+import type { UIMessageStreamWriter } from 'ai';
+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>;
+    }
+    interface Graph<State extends Record<string, unknown>, NodeKeys extends string> {
+        nodes: Map<NodeKeys, Node<State, NodeKeys>>;
+        edges: Map<NodeKeys, Edge<State, NodeKeys>[]>;
+    }
+    interface Node<State extends Record<string, unknown>, NodeKeys extends string> {
+        id: NodeKeys;
+        execute: ({ state, writer, suspense, update }: {
+            state: () => Readonly<State>;
+            writer: UIMessageStreamWriter;
+            suspense: (data?: unknown) => never;
+            update: (update: StateUpdate<State>) => void;
+        }) => Promise<void> | void;
+    }
+    interface Edge<State extends Record<string, unknown>, NodeKeys extends string> {
+        from: NodeKeys;
+        to: NodeKeys | ((state: State) => NodeKeys);
+    }
+    interface Checkpoint<State extends Record<string, unknown>, NodeKeys extends string> {
+        state: State;
+        nodeIds: NodeKeys[];
+        suspendedNodes: NodeKeys[];
+    }
+    interface GraphStorage<State extends Record<string, unknown>, NodeKeys extends string> {
+        save(runId: string, checkpoint: Checkpoint<State, NodeKeys>): Promise<void>;
+        load(runId: string): Promise<Checkpoint<State, NodeKeys> | null>;
+        delete(runId: string): Promise<void>;
+    }
+    interface ExecutionContext<State extends Record<string, unknown>, NodeKeys extends string> {
+        runId: string;
+        state: State;
+        currentNodes: Node<State, NodeKeys>[];
+        suspendedNodes: Node<State, NodeKeys>[];
+        writer: UIMessageStreamWriter;
+    }
+}

package/package.json

@@ -0,0 +1,33 @@
+{
+  "name": "ai-sdk-graph",
+  "version": "0.1.0",
+  "description": "Graph-based workflows for the AI SDK",
+  "type": "module",
+  "main": "dist/index.js",
+  "module": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js",
+      "require": "./dist/index.js"
+    }
+  },
+  "files": [
+    "dist"
+  ],
+  "scripts": {
+    "test": "bun test",
+    "build": "bun build src/index.ts --outdir dist --target node --packages=external && tsc -p tsconfig.build.json"
+  },
+  "devDependencies": {
+    "@types/bun": "latest"
+  },
+  "peerDependencies": {
+    "typescript": "^5"
+  },
+  "dependencies": {
+    "ai": "^6.0.37",
+    "ioredis": "^5.9.2"
+  }
+}