ai-sdk-graph 0.1.3 → 0.2.0

package/dist/index.js

@@ -98,8 +98,12 @@ class Graph {
     let context;
     return createUIMessageStream({
       execute: async ({ writer }) => {
-        context = await this.createExecutionContext(runId, initialState, writer);
-        await this.onStart({ state: context.state, 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 () => {
@@ -110,7 +114,7 @@ class Graph {
     });
   }
   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;
   }
@@ -124,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, writer);
+    return { context: { ...context, runId, writer }, firstTime };
   }
   async runExecutionLoop(context) {
     await this.executeWithStrategy(context, "return");
@@ -174,12 +178,12 @@ class Graph {
   hasNodesToExecute(context) {
     return context.currentNodes.length > 0;
   }
-  async restoreCheckpoint(runId, initialState) {
+  async restoreCheckpoint(runId, initialState, writer) {
     const checkpoint = await this.storage.load(runId);
     if (this.isValidCheckpoint(checkpoint)) {
-      return this.restoreFromCheckpoint(checkpoint, initialState);
+      return { context: this.restoreFromCheckpoint(checkpoint, initialState, writer), firstTime: false };
     }
-    return this.createFreshExecution(initialState);
+    return { context: this.createFreshExecution(initialState, writer), firstTime: true };
   }
   isValidCheckpoint(checkpoint) {
     return this.hasNodeIds(checkpoint) && this.hasAtLeastOneNode(checkpoint);
@@ -190,16 +194,16 @@ class Graph {
   hasAtLeastOneNode(checkpoint) {
     return checkpoint.nodeIds.length > 0;
   }
-  restoreFromCheckpoint(checkpoint, initialState) {
-    const state = this.stateManager.resolve(initialState, checkpoint.state);
+  restoreFromCheckpoint(checkpoint, initialState, writer) {
+    const state = this.stateManager.resolve(initialState, checkpoint.state, writer);
     return {
       state,
       currentNodes: this.resolveNodeIds(checkpoint.nodeIds),
       suspendedNodes: this.resolveNodeIds(checkpoint.suspendedNodes)
     };
   }
-  createFreshExecution(initialState) {
-    const state = this.stateManager.resolve(initialState, undefined);
+  createFreshExecution(initialState, writer) {
+    const state = this.stateManager.resolve(initialState, undefined, writer);
     const startNode = this.nodeRegistry.get(BUILT_IN_NODES.START);
     return {
       state,
@@ -242,7 +246,7 @@ class Graph {
       writer: context.writer,
       suspense: this.createSuspenseFunction(),
       update: (update) => {
-        context.state = this.stateManager.apply(context.state, update);
+        context.state = this.stateManager.apply(context.state, update, context.writer);
       }
     };
   }
@@ -253,7 +257,7 @@ class Graph {
     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);
+      context.state = this.stateManager.apply(context.state, parentUpdate, context.writer);
       this.emitter.emitEnd(context.writer, node.id);
       return null;
     } catch (error) {
@@ -312,17 +316,23 @@ class NodeEventEmitter {
 }
 
 class StateManager {
-  apply(state, update) {
-    return {
+  apply(state, update, writer) {
+    const newState = {
       ...state,
       ...typeof update === "function" ? update(state) : update
     };
+    writer.write({ type: "data-state", data: newState });
+    return newState;
   }
-  resolve(initialState, existingState) {
+  resolve(initialState, existingState, writer) {
     if (this.isStateFactory(initialState)) {
-      return initialState(existingState);
+      const newState2 = initialState(existingState);
+      writer.write({ type: "data-state", data: newState2 });
+      return newState2;
     }
-    return existingState ?? initialState;
+    const newState = existingState ?? initialState;
+    writer.write({ type: "data-state", data: newState });
+    return newState;
   }
   isStateFactory(initialState) {
     return typeof initialState === "function";

package/package.json

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