@langchain/langgraph 0.2.8 → 0.2.10-rc.0

package/dist/pregel/index.cjs

@@ -1,9 +1,13 @@
 "use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.Pregel = exports.Channel = void 0;
 /* eslint-disable no-param-reassign */
 const runnables_1 = require("@langchain/core/runnables");
 const langgraph_checkpoint_1 = require("@langchain/langgraph-checkpoint");
+const double_ended_queue_1 = __importDefault(require("double-ended-queue"));
 const base_js_1 = require("../channels/base.cjs");
 const read_js_1 = require("./read.cjs");
 const validate_js_1 = require("./validate.cjs");
@@ -13,11 +17,12 @@ const write_js_1 = require("./write.cjs");
 const constants_js_1 = require("../constants.cjs");
 const errors_js_1 = require("../errors.cjs");
 const algo_js_1 = require("./algo.cjs");
-const utils_js_1 = require("./utils.cjs");
+const index_js_1 = require("./utils/index.cjs");
 const loop_js_1 = require("./loop.cjs");
 const retry_js_1 = require("./retry.cjs");
 const base_js_2 = require("../managed/base.cjs");
-const utils_js_2 = require("../utils.cjs");
+const utils_js_1 = require("../utils.cjs");
+const config_js_1 = require("./utils/config.cjs");
 function isString(value) {
     return typeof value === "string";
 }
@@ -76,6 +81,16 @@ class Channel {
     }
 }
 exports.Channel = Channel;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function isPregel(x) {
+    return ("inputChannels" in x &&
+        x.inputChannels !== undefined &&
+        "outputChannels" &&
+        x.outputChannels !== undefined);
+}
+function isRunnableSequence(x) {
+    return "steps" in x && Array.isArray(x.steps);
+}
 class Pregel extends runnables_1.Runnable {
     static lc_name() {
         return "LangGraph";
@@ -167,6 +182,12 @@ class Pregel extends runnables_1.Runnable {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "config", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "store", {
             enumerable: true,
             configurable: true,
@@ -190,11 +211,18 @@ class Pregel extends runnables_1.Runnable {
         this.debug = fields.debug ?? this.debug;
         this.checkpointer = fields.checkpointer;
         this.retryPolicy = fields.retryPolicy;
+        this.config = fields.config;
         this.store = fields.store;
         if (this.autoValidate) {
             this.validate();
         }
     }
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore Remove ignore when we remove support for 0.2 versions of core
+    withConfig(config) {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        return new this.constructor({ ...this, config });
+    }
     validate() {
         (0, validate_js_1.validateGraph)({
             nodes: this.nodes,
@@ -226,50 +254,181 @@ class Pregel extends runnables_1.Runnable {
             return Object.keys(this.channels);
         }
     }
-    /**
-     * Get the current state of the graph.
-     */
-    async getState(config) {
-        if (!this.checkpointer) {
-            throw new errors_js_1.GraphValueError("No checkpointer set");
+    *getSubgraphs(namespace, recurse
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    ) {
+        for (const [name, node] of Object.entries(this.nodes)) {
+            // filter by prefix
+            if (namespace !== undefined) {
+                if (!namespace.startsWith(name)) {
+                    continue;
+                }
+            }
+            // find the subgraph if any
+            let graph;
+            const candidates = [node.bound];
+            for (const candidate of candidates) {
+                if (isPregel(candidate)) {
+                    graph = candidate;
+                    break;
+                }
+                else if (isRunnableSequence(candidate)) {
+                    candidates.push(...candidate.steps);
+                }
+            }
+            // if found, yield recursively
+            if (graph !== undefined) {
+                if (name === namespace) {
+                    yield [name, graph];
+                    return;
+                }
+                if (namespace === undefined) {
+                    yield [name, graph];
+                }
+                if (recurse) {
+                    let newNamespace = namespace;
+                    if (namespace !== undefined) {
+                        newNamespace = namespace.slice(name.length + 1);
+                    }
+                    for (const [subgraphName, subgraph] of graph.getSubgraphs(newNamespace, recurse)) {
+                        yield [
+                            `${name}${constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR}${subgraphName}`,
+                            subgraph,
+                        ];
+                    }
+                }
+            }
+        }
+    }
+    async _prepareStateSnapshot({ config, saved, subgraphCheckpointer, }) {
+        if (saved === undefined) {
+            return {
+                values: {},
+                next: [],
+                config,
+                tasks: [],
+            };
         }
-        const saved = await this.checkpointer.getTuple(config);
-        const checkpoint = saved ? saved.checkpoint : (0, langgraph_checkpoint_1.emptyCheckpoint)();
-        const channels = (0, base_js_1.emptyChannels)(this.channels, checkpoint);
         // Pass `skipManaged: true` as managed values should not be returned in get state calls.
         const { managed } = await this.prepareSpecs(config, { skipManaged: true });
-        const nextTasks = (0, algo_js_1._prepareNextTasks)(checkpoint, this.nodes, channels, managed, saved !== undefined ? saved.config : config, false, { step: saved ? (saved.metadata?.step ?? -1) + 1 : -1 });
+        const channels = (0, base_js_1.emptyChannels)(this.channels, saved.checkpoint);
+        const nextTasks = Object.values((0, algo_js_1._prepareNextTasks)(saved.checkpoint, this.nodes, channels, managed, saved.config, false, { step: (saved.metadata?.step ?? -1) + 1 }));
+        const subgraphs = await (0, utils_js_1.gatherIterator)(this.getSubgraphs());
+        const parentNamespace = saved.config.configurable?.checkpoint_ns ?? "";
+        const taskStates = {};
+        for (const task of nextTasks) {
+            const matchingSubgraph = subgraphs.find(([name]) => name === task.name);
+            if (!matchingSubgraph) {
+                continue;
+            }
+            // assemble checkpoint_ns for this task
+            let taskNs = `${task.name}${constants_js_1.CHECKPOINT_NAMESPACE_END}${task.id}`;
+            if (parentNamespace) {
+                taskNs = `${parentNamespace}${constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR}${taskNs}`;
+            }
+            if (subgraphCheckpointer === undefined) {
+                // set config as signal that subgraph checkpoints exist
+                const config = {
+                    configurable: {
+                        thread_id: saved.config.configurable?.thread_id,
+                        checkpoint_ns: taskNs,
+                    },
+                };
+                taskStates[task.id] = config;
+            }
+            else {
+                // get the state of the subgraph
+                const config = {
+                    configurable: {
+                        [constants_js_1.CONFIG_KEY_CHECKPOINTER]: subgraphCheckpointer,
+                        thread_id: saved.config.configurable?.thread_id,
+                        checkpoint_ns: taskNs,
+                    },
+                };
+                taskStates[task.id] = await matchingSubgraph[1].getState(config, {
+                    subgraphs: true,
+                });
+            }
+        }
+        // assemble the state snapshot
         return {
             values: (0, io_js_1.readChannels)(channels, this.streamChannelsAsIs),
             next: nextTasks.map((task) => task.name),
-            tasks: (0, debug_js_1.tasksWithWrites)(nextTasks, saved?.pendingWrites ?? []),
-            metadata: saved?.metadata,
-            config: saved ? saved.config : config,
-            createdAt: saved?.checkpoint.ts,
-            parentConfig: saved?.parentConfig,
+            tasks: (0, debug_js_1.tasksWithWrites)(nextTasks, saved?.pendingWrites ?? [], taskStates),
+            metadata: saved.metadata,
+            config: (0, index_js_1.patchCheckpointMap)(saved.config, saved.metadata),
+            createdAt: saved.checkpoint.ts,
+            parentConfig: saved.parentConfig,
         };
     }
+    /**
+     * Get the current state of the graph.
+     */
+    async getState(config, options) {
+        const checkpointer = config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
+        if (!checkpointer) {
+            throw new errors_js_1.GraphValueError("No checkpointer set");
+        }
+        const checkpointNamespace = config.configurable?.checkpoint_ns ?? "";
+        if (checkpointNamespace !== "" &&
+            config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] === undefined) {
+            // remove task_ids from checkpoint_ns
+            const recastCheckpointNamespace = checkpointNamespace
+                .split(constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR)
+                .map((part) => part.split(constants_js_1.CHECKPOINT_NAMESPACE_END)[0])
+                .join(constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR);
+            for (const [name, subgraph] of this.getSubgraphs(recastCheckpointNamespace, true)) {
+                if (name === recastCheckpointNamespace) {
+                    return await subgraph.getState((0, utils_js_1.patchConfigurable)(config, {
+                        [constants_js_1.CONFIG_KEY_CHECKPOINTER]: checkpointer,
+                    }), { subgraphs: options?.subgraphs });
+                }
+            }
+            throw new Error(`Subgraph with namespace "${recastCheckpointNamespace}" not found.`);
+        }
+        const mergedConfig = (0, runnables_1.mergeConfigs)(this.config, config);
+        const saved = await checkpointer.getTuple(config);
+        const snapshot = await this._prepareStateSnapshot({
+            config: mergedConfig,
+            saved,
+            subgraphCheckpointer: options?.subgraphs ? checkpointer : undefined,
+        });
+        return snapshot;
+    }
     /**
      * Get the history of the state of the graph.
      */
     async *getStateHistory(config, options) {
-        if (!this.checkpointer) {
-            throw new errors_js_1.GraphValueError("No checkpointer set");
+        const checkpointer = config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
+        if (!checkpointer) {
+            throw new Error("No checkpointer set");
+        }
+        const checkpointNamespace = config.configurable?.checkpoint_ns ?? "";
+        if (checkpointNamespace !== "" &&
+            config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] === undefined) {
+            const recastCheckpointNamespace = checkpointNamespace
+                .split(constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR)
+                .map((part) => part.split(constants_js_1.CHECKPOINT_NAMESPACE_END)[0])
+                .join(constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR);
+            // find the subgraph with the matching name
+            for (const [name, pregel] of this.getSubgraphs(recastCheckpointNamespace, true)) {
+                if (name === recastCheckpointNamespace) {
+                    yield* pregel.getStateHistory((0, utils_js_1.patchConfigurable)(config, {
+                        [constants_js_1.CONFIG_KEY_CHECKPOINTER]: checkpointer,
+                    }), options);
+                    return;
+                }
+            }
+            throw new Error(`Subgraph with namespace "${recastCheckpointNamespace}" not found.`);
         }
-        // Pass `skipManaged: true` as managed values should not be returned in get state calls.
-        const { managed } = await this.prepareSpecs(config, { skipManaged: true });
-        for await (const saved of this.checkpointer.list(config, options)) {
-            const channels = (0, base_js_1.emptyChannels)(this.channels, saved.checkpoint);
-            const nextTasks = (0, algo_js_1._prepareNextTasks)(saved.checkpoint, this.nodes, channels, managed, saved.config, false, { step: -1 });
-            yield {
-                values: (0, io_js_1.readChannels)(channels, this.streamChannelsAsIs),
-                next: nextTasks.map((task) => task.name),
-                tasks: (0, debug_js_1.tasksWithWrites)(nextTasks, saved.pendingWrites ?? []),
-                metadata: saved.metadata,
-                config: saved.config,
-                createdAt: saved.checkpoint.ts,
-                parentConfig: saved.parentConfig,
-            };
+        const mergedConfig = (0, runnables_1.mergeConfigs)(this.config, config, {
+            configurable: { checkpoint_ns: checkpointNamespace },
+        });
+        for await (const checkpointTuple of checkpointer.list(mergedConfig, options)) {
+            yield this._prepareStateSnapshot({
+                config: checkpointTuple.config,
+                saved: checkpointTuple,
+            });
         }
     }
     /**
@@ -277,34 +436,59 @@ class Pregel extends runnables_1.Runnable {
      * node `as_node`. If `as_node` is not provided, it will be set to the last node
      * that updated the state, if not ambiguous.
      */
-    async updateState(config, values, asNode) {
-        if (!this.checkpointer) {
+    async updateState(inputConfig, values, asNode) {
+        const checkpointer = inputConfig.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
+        if (!checkpointer) {
             throw new errors_js_1.GraphValueError("No checkpointer set");
         }
-        // Get the latest checkpoint
-        const saved = await this.checkpointer.getTuple(config);
-        const checkpoint = saved
+        // delegate to subgraph
+        const checkpointNamespace = inputConfig.configurable?.checkpoint_ns ?? "";
+        if (checkpointNamespace !== "" &&
+            inputConfig.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] === undefined) {
+            // remove task_ids from checkpoint_ns
+            const recastCheckpointNamespace = checkpointNamespace
+                .split(constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR)
+                .map((part) => {
+                return part.split(constants_js_1.CHECKPOINT_NAMESPACE_END)[0];
+            })
+                .join(constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR);
+            // find the subgraph with the matching name
+            // eslint-disable-next-line no-unreachable-loop
+            for (const [, pregel] of this.getSubgraphs(recastCheckpointNamespace, true)) {
+                return await pregel.updateState((0, utils_js_1.patchConfigurable)(inputConfig, {
+                    [constants_js_1.CONFIG_KEY_CHECKPOINTER]: checkpointer,
+                }), values, asNode);
+            }
+            throw new Error(`Subgraph "${recastCheckpointNamespace}" not found`);
+        }
+        // get last checkpoint
+        const config = this.config
+            ? (0, runnables_1.mergeConfigs)(this.config, inputConfig)
+            : inputConfig;
+        const saved = await checkpointer.getTuple(config);
+        const checkpoint = saved !== undefined
             ? (0, langgraph_checkpoint_1.copyCheckpoint)(saved.checkpoint)
             : (0, langgraph_checkpoint_1.emptyCheckpoint)();
-        const checkpointPreviousVersions = saved?.checkpoint.channel_versions ?? {};
+        const checkpointPreviousVersions = {
+            ...saved?.checkpoint.channel_versions,
+        };
         const step = saved?.metadata?.step ?? -1;
         // merge configurable fields with previous checkpoint config
-        const checkpointConfig = {
-            ...config,
-            configurable: {
-                ...config.configurable,
-                // TODO: add proper support for updating nested subgraph state
-                checkpoint_ns: "",
-                ...saved?.config.configurable,
-            },
-        };
+        let checkpointConfig = (0, utils_js_1.patchConfigurable)(config, {
+            checkpoint_ns: config.configurable?.checkpoint_ns ?? "",
+        });
+        if (saved) {
+            checkpointConfig = (0, utils_js_1.patchConfigurable)(config, saved.config.configurable);
+        }
         // Find last node that updated the state, if not provided
         if (values == null && asNode === undefined) {
-            return await this.checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, undefined, step), {
+            const nextConfig = await checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, undefined, step), {
                 source: "update",
                 step,
                 writes: {},
+                parents: saved?.metadata?.parents ?? {},
             }, {});
+            return (0, index_js_1.patchCheckpointMap)(nextConfig, saved ? saved.metadata : undefined);
         }
         const nonNullVersion = Object.values(checkpoint.versions_seen)
             .map((seenVersions) => {
@@ -319,7 +503,6 @@ class Pregel extends runnables_1.Runnable {
             }
         }
         else if (asNode === undefined) {
-            // TODO: Double check
             const lastSeenByNode = Object.entries(checkpoint.versions_seen)
                 .map(([n, seen]) => {
                 return Object.values(seen).map((v) => {
@@ -367,7 +550,10 @@ class Pregel extends runnables_1.Runnable {
             id: (0, langgraph_checkpoint_1.uuid5)(constants_js_1.INTERRUPT, checkpoint.id),
         };
         // execute task
-        await task.proc.invoke(task.input, (0, runnables_1.patchConfig)(config, {
+        await task.proc.invoke(task.input, (0, runnables_1.patchConfig)({
+            ...config,
+            store: config?.store ?? this.store,
+        }, {
             runName: config.runName ?? `${this.getName()}UpdateState`,
             configurable: {
                 [constants_js_1.CONFIG_KEY_SEND]: (items) => task.writes.push(...items),
@@ -376,18 +562,21 @@ class Pregel extends runnables_1.Runnable {
                 task, select_, fresh_),
             },
         }));
+        // save task writes
         if (saved !== undefined) {
-            await this.checkpointer.putWrites(checkpointConfig, task.writes, task.id);
+            await checkpointer.putWrites(checkpointConfig, task.writes, task.id);
         }
         // apply to checkpoint
         // TODO: Why does keyof StrRecord allow number and symbol?
-        (0, algo_js_1._applyWrites)(checkpoint, channels, [task], this.checkpointer.getNextVersion.bind(this.checkpointer));
-        const newVersions = (0, utils_js_1.getNewChannelVersions)(checkpointPreviousVersions, checkpoint.channel_versions);
-        return await this.checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, channels, step + 1), {
+        (0, algo_js_1._applyWrites)(checkpoint, channels, [task], checkpointer.getNextVersion.bind(this.checkpointer));
+        const newVersions = (0, index_js_1.getNewChannelVersions)(checkpointPreviousVersions, checkpoint.channel_versions);
+        const nextConfig = await checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, channels, step + 1), {
             source: "update",
             step: step + 1,
             writes: { [asNode]: values },
+            parents: saved?.metadata?.parents ?? {},
         }, newVersions);
+        return (0, index_js_1.patchCheckpointMap)(nextConfig, saved ? saved.metadata : undefined);
     }
     _defaults(config) {
         const { debug, streamMode, inputKeys, outputKeys, interruptAfter, interruptBefore, ...rest } = config;
@@ -415,19 +604,19 @@ class Pregel extends runnables_1.Runnable {
         else {
             defaultStreamMode = this.streamMode;
         }
-        let defaultCheckpointer;
-        if (config.configurable !== undefined &&
-            config.configurable[constants_js_1.CONFIG_KEY_READ] !== undefined) {
+        // if being called as a node in another graph, always use values mode
+        if (config.configurable?.[constants_js_1.CONFIG_KEY_TASK_ID] !== undefined) {
             defaultStreamMode = ["values"];
         }
+        let defaultCheckpointer;
         if (config !== undefined &&
-            config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] !== undefined &&
-            (defaultInterruptAfter.length > 0 || defaultInterruptBefore.length > 0)) {
+            config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] !== undefined) {
             defaultCheckpointer = config.configurable[constants_js_1.CONFIG_KEY_CHECKPOINTER];
         }
         else {
             defaultCheckpointer = this.checkpointer;
         }
+        const defaultStore = config.store ?? this.store;
         return [
             defaultDebug,
             defaultStreamMode,
@@ -437,6 +626,7 @@ class Pregel extends runnables_1.Runnable {
             defaultInterruptBefore,
             defaultInterruptAfter,
             defaultCheckpointer,
+            defaultStore,
         ];
     }
     /**
@@ -458,9 +648,10 @@ class Pregel extends runnables_1.Runnable {
         return super.stream(input, options);
     }
     async prepareSpecs(config, options) {
-        const configForManaged = (0, utils_js_2.patchConfigurable)(config, {
-            [constants_js_1.CONFIG_KEY_STORE]: this.store,
-        });
+        const configForManaged = {
+            ...config,
+            store: this.store,
+        };
         const channelSpecs = {};
         const managedSpecs = {};
         for (const [name, spec] of Object.entries(this.channels)) {
@@ -501,7 +692,8 @@ class Pregel extends runnables_1.Runnable {
         };
     }
     async *_streamIterator(input, options) {
-        const inputConfig = (0, runnables_1.ensureConfig)(options);
+        const streamSubgraphs = options?.subgraphs;
+        const inputConfig = (0, config_js_1.ensureLangGraphConfig)(this.config, options);
         if (inputConfig.recursionLimit === undefined ||
             inputConfig.recursionLimit < 1) {
             throw new Error(`Passed "recursionLimit" must be at least 1.`);
@@ -511,12 +703,36 @@ class Pregel extends runnables_1.Runnable {
             throw new Error(`Checkpointer requires one or more of the following "configurable" keys: "thread_id", "checkpoint_ns", "checkpoint_id"`);
         }
         const callbackManager = await (0, runnables_1.getCallbackManagerForConfig)(inputConfig);
-        const runManager = await callbackManager?.handleChainStart(this.toJSON(), (0, utils_js_1._coerceToDict)(input, "input"), inputConfig.runId, undefined, undefined, undefined, inputConfig?.runName ?? this.getName());
+        const runManager = await callbackManager?.handleChainStart(this.toJSON(), (0, index_js_1._coerceToDict)(input, "input"), inputConfig.runId, undefined, undefined, undefined, inputConfig?.runName ?? this.getName());
         delete inputConfig.runId;
         // assign defaults
-        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer,] = this._defaults(inputConfig);
+        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store,] = this._defaults(inputConfig);
         const { channelSpecs, managed } = await this.prepareSpecs(config);
         let loop;
+        const stream = new double_ended_queue_1.default();
+        function* emitCurrentLoopOutputs() {
+            while (loop !== undefined && stream.length > 0) {
+                const nextItem = stream.shift();
+                if (nextItem === undefined) {
+                    throw new Error("Data structure error.");
+                }
+                const [namespace, mode, payload] = nextItem;
+                if (streamMode.includes(mode)) {
+                    if (streamSubgraphs && streamMode.length > 1) {
+                        yield [namespace, mode, payload];
+                    }
+                    else if (streamMode.length > 1) {
+                        yield [mode, payload];
+                    }
+                    else if (streamSubgraphs) {
+                        yield [namespace, payload];
+                    }
+                    else {
+                        yield payload;
+                    }
+                }
+            }
+        }
         try {
             loop = await loop_js_1.PregelLoop.initialize({
                 input,
@@ -527,8 +743,15 @@ class Pregel extends runnables_1.Runnable {
                 managed,
                 outputKeys,
                 streamKeys: this.streamChannelsAsIs,
-                store: this.store,
+                store,
+                stream: new loop_js_1.StreamProtocol((chunk) => stream.push(chunk), new Set(streamMode)),
             });
+            if (options?.subgraphs) {
+                loop.config.configurable = {
+                    ...loop.config.configurable,
+                    [constants_js_1.CONFIG_KEY_STREAM]: loop.stream,
+                };
+            }
             while (await loop.tick({
                 inputKeys: this.inputChannels,
                 interruptAfter,
@@ -538,27 +761,14 @@ class Pregel extends runnables_1.Runnable {
                 if (debug) {
                     (0, debug_js_1.printStepCheckpoint)(loop.checkpointMetadata.step, loop.channels, this.streamChannelsList);
                 }
-                while (loop.stream.length > 0) {
-                    const nextItem = loop.stream.shift();
-                    if (nextItem === undefined) {
-                        throw new Error("Data structure error.");
-                    }
-                    if (streamMode.includes(nextItem[0])) {
-                        if (streamMode.length === 1) {
-                            yield nextItem[1];
-                        }
-                        else {
-                            yield nextItem;
-                        }
-                    }
-                }
+                yield* emitCurrentLoopOutputs();
                 if (debug) {
-                    (0, debug_js_1.printStepTasks)(loop.step, loop.tasks);
+                    (0, debug_js_1.printStepTasks)(loop.step, Object.values(loop.tasks));
                 }
                 // execute tasks, and wait for one to fail or all to finish.
                 // each task is independent from all other concurrent tasks
                 // yield updates/debug output as each task finishes
-                const taskStream = (0, retry_js_1.executeTasksWithRetry)(loop.tasks.filter((task) => task.writes.length === 0), {
+                const taskStream = (0, retry_js_1.executeTasksWithRetry)(Object.values(loop.tasks).filter((task) => task.writes.length === 0), {
                     stepTimeout: this.stepTimeout,
                     signal: config.signal,
                     retryPolicy: this.retryPolicy,
@@ -567,7 +777,12 @@ class Pregel extends runnables_1.Runnable {
                 for await (const { task, error } of taskStream) {
                     if (error !== undefined) {
                         if ((0, errors_js_1.isGraphInterrupt)(error)) {
-                            loop.putWrites(task.id, error.interrupts.map((interrupt) => [constants_js_1.INTERRUPT, interrupt]));
+                            if (loop.isNested) {
+                                throw error;
+                            }
+                            if (error.interrupts.length) {
+                                loop.putWrites(task.id, error.interrupts.map((interrupt) => [constants_js_1.INTERRUPT, interrupt]));
+                            }
                         }
                         else {
                             loop.putWrites(task.id, [
@@ -578,42 +793,18 @@ class Pregel extends runnables_1.Runnable {
                     else {
                         loop.putWrites(task.id, task.writes);
                     }
-                    while (loop.stream.length > 0) {
-                        const nextItem = loop.stream.shift();
-                        if (nextItem === undefined) {
-                            throw new Error("Data structure error.");
-                        }
-                        if (streamMode.includes(nextItem[0])) {
-                            if (streamMode.length === 1) {
-                                yield nextItem[1];
-                            }
-                            else {
-                                yield nextItem;
-                            }
-                        }
-                    }
+                    yield* emitCurrentLoopOutputs();
                     if (error !== undefined && !(0, errors_js_1.isGraphInterrupt)(error)) {
                         throw error;
                     }
                 }
                 if (debug) {
-                    (0, debug_js_1.printStepWrites)(loop.step, loop.tasks.map((task) => task.writes).flat(), this.streamChannelsList);
-                }
-            }
-            while (loop.stream.length > 0) {
-                const nextItem = loop.stream.shift();
-                if (nextItem === undefined) {
-                    throw new Error("Data structure error.");
-                }
-                if (streamMode.includes(nextItem[0])) {
-                    if (streamMode.length === 1) {
-                        yield nextItem[1];
-                    }
-                    else {
-                        yield nextItem;
-                    }
+                    (0, debug_js_1.printStepWrites)(loop.step, Object.values(loop.tasks)
+                        .map((task) => task.writes)
+                        .flat(), this.streamChannelsList);
                 }
             }
+            yield* emitCurrentLoopOutputs();
             if (loop.status === "out_of_steps") {
                 throw new errors_js_1.GraphRecursionError([
                     `Recursion limit of ${config.recursionLimit} reached`,
@@ -622,7 +813,7 @@ class Pregel extends runnables_1.Runnable {
                 ].join(" "));
             }
             await Promise.all(loop?.checkpointerPromises ?? []);
-            await runManager?.handleChainEnd((0, io_js_1.readChannels)(loop.channels, outputKeys));
+            await runManager?.handleChainEnd(loop.output);
         }
         catch (e) {
             await runManager?.handleChainError(e);
@@ -631,7 +822,7 @@ class Pregel extends runnables_1.Runnable {
         finally {
             // Call `.stop()` again incase it was not called in the loop, e.g due to an error.
             if (loop) {
-                loop.store?.stop();
+                await loop.store?.stop();
             }
             await Promise.all([
                 loop?.checkpointerPromises ?? [],
@@ -657,7 +848,7 @@ class Pregel extends runnables_1.Runnable {
     async invoke(input, options) {
         const streamMode = options?.streamMode ?? "values";
         const config = {
-            ...(0, runnables_1.ensureConfig)(options),
+            ...options,
             outputKeys: options?.outputKeys ?? this.outputChannels,
             streamMode,
         };