@langchain/langgraph 0.2.8 → 0.2.9

package/dist/pregel/index.js

@@ -1,20 +1,22 @@
 /* eslint-disable no-param-reassign */
-import { Runnable, RunnableSequence, _coerceToRunnable, ensureConfig, getCallbackManagerForConfig, patchConfig, } from "@langchain/core/runnables";
+import { Runnable, RunnableSequence, _coerceToRunnable, getCallbackManagerForConfig, mergeConfigs, patchConfig, } from "@langchain/core/runnables";
 import { compareChannelVersions, copyCheckpoint, emptyCheckpoint, uuid5, } from "@langchain/langgraph-checkpoint";
+import Deque from "double-ended-queue";
 import { createCheckpoint, emptyChannels, isBaseChannel, } from "../channels/base.js";
 import { PregelNode } from "./read.js";
 import { validateGraph, validateKeys } from "./validate.js";
 import { readChannels } from "./io.js";
 import { printStepCheckpoint, printStepTasks, printStepWrites, tasksWithWrites, } from "./debug.js";
 import { ChannelWrite, PASSTHROUGH } from "./write.js";
-import { CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STORE, ERROR, INTERRUPT, } from "../constants.js";
+import { CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STORE, ERROR, INTERRUPT, CHECKPOINT_NAMESPACE_SEPARATOR, CHECKPOINT_NAMESPACE_END, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, } from "../constants.js";
 import { GraphRecursionError, GraphValueError, InvalidUpdateError, isGraphInterrupt, } from "../errors.js";
 import { _prepareNextTasks, _localRead, _applyWrites, } from "./algo.js";
-import { _coerceToDict, getNewChannelVersions } from "./utils.js";
-import { PregelLoop } from "./loop.js";
+import { _coerceToDict, getNewChannelVersions, patchCheckpointMap, } from "./utils/index.js";
+import { PregelLoop, StreamProtocol } from "./loop.js";
 import { executeTasksWithRetry } from "./retry.js";
 import { ChannelKeyPlaceholder, isConfiguredManagedValue, ManagedValueMapping, NoopManagedValue, } from "../managed/base.js";
-import { patchConfigurable } from "../utils.js";
+import { gatherIterator, patchConfigurable } from "../utils.js";
+import { ensureLangGraphConfig } from "./utils/config.js";
 function isString(value) {
     return typeof value === "string";
 }
@@ -72,6 +74,16 @@ export class Channel {
         return new ChannelWrite(channelWriteEntries);
     }
 }
+// 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);
+}
 export class Pregel extends Runnable {
     static lc_name() {
         return "LangGraph";
@@ -163,6 +175,12 @@ export class Pregel extends 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,
@@ -186,11 +204,18 @@ export class Pregel extends 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() {
         validateGraph({
             nodes: this.nodes,
@@ -222,50 +247,181 @@ export class Pregel extends Runnable {
             return Object.keys(this.channels);
         }
     }
-    /**
-     * Get the current state of the graph.
-     */
-    async getState(config) {
-        if (!this.checkpointer) {
-            throw new 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}${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 : emptyCheckpoint();
-        const channels = 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 = _prepareNextTasks(checkpoint, this.nodes, channels, managed, saved !== undefined ? saved.config : config, false, { step: saved ? (saved.metadata?.step ?? -1) + 1 : -1 });
+        const channels = emptyChannels(this.channels, saved.checkpoint);
+        const nextTasks = Object.values(_prepareNextTasks(saved.checkpoint, this.nodes, channels, managed, saved.config, false, { step: (saved.metadata?.step ?? -1) + 1 }));
+        const subgraphs = await 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}${CHECKPOINT_NAMESPACE_END}${task.id}`;
+            if (parentNamespace) {
+                taskNs = `${parentNamespace}${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: {
+                        [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: readChannels(channels, this.streamChannelsAsIs),
             next: nextTasks.map((task) => task.name),
-            tasks: tasksWithWrites(nextTasks, saved?.pendingWrites ?? []),
-            metadata: saved?.metadata,
-            config: saved ? saved.config : config,
-            createdAt: saved?.checkpoint.ts,
-            parentConfig: saved?.parentConfig,
+            tasks: tasksWithWrites(nextTasks, saved?.pendingWrites ?? [], taskStates),
+            metadata: saved.metadata,
+            config: 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?.[CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
+        if (!checkpointer) {
+            throw new GraphValueError("No checkpointer set");
+        }
+        const checkpointNamespace = config.configurable?.checkpoint_ns ?? "";
+        if (checkpointNamespace !== "" &&
+            config.configurable?.[CONFIG_KEY_CHECKPOINTER] === undefined) {
+            // remove task_ids from checkpoint_ns
+            const recastCheckpointNamespace = checkpointNamespace
+                .split(CHECKPOINT_NAMESPACE_SEPARATOR)
+                .map((part) => part.split(CHECKPOINT_NAMESPACE_END)[0])
+                .join(CHECKPOINT_NAMESPACE_SEPARATOR);
+            for (const [name, subgraph] of this.getSubgraphs(recastCheckpointNamespace, true)) {
+                if (name === recastCheckpointNamespace) {
+                    return await subgraph.getState(patchConfigurable(config, {
+                        [CONFIG_KEY_CHECKPOINTER]: checkpointer,
+                    }), { subgraphs: options?.subgraphs });
+                }
+            }
+            throw new Error(`Subgraph with namespace "${recastCheckpointNamespace}" not found.`);
+        }
+        const mergedConfig = 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 GraphValueError("No checkpointer set");
+        const checkpointer = config.configurable?.[CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
+        if (!checkpointer) {
+            throw new Error("No checkpointer set");
+        }
+        const checkpointNamespace = config.configurable?.checkpoint_ns ?? "";
+        if (checkpointNamespace !== "" &&
+            config.configurable?.[CONFIG_KEY_CHECKPOINTER] === undefined) {
+            const recastCheckpointNamespace = checkpointNamespace
+                .split(CHECKPOINT_NAMESPACE_SEPARATOR)
+                .map((part) => part.split(CHECKPOINT_NAMESPACE_END)[0])
+                .join(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(patchConfigurable(config, {
+                        [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 = emptyChannels(this.channels, saved.checkpoint);
-            const nextTasks = _prepareNextTasks(saved.checkpoint, this.nodes, channels, managed, saved.config, false, { step: -1 });
-            yield {
-                values: readChannels(channels, this.streamChannelsAsIs),
-                next: nextTasks.map((task) => task.name),
-                tasks: tasksWithWrites(nextTasks, saved.pendingWrites ?? []),
-                metadata: saved.metadata,
-                config: saved.config,
-                createdAt: saved.checkpoint.ts,
-                parentConfig: saved.parentConfig,
-            };
+        const mergedConfig = 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,
+            });
         }
     }
     /**
@@ -273,34 +429,59 @@ export class Pregel extends 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?.[CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
+        if (!checkpointer) {
             throw new 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?.[CONFIG_KEY_CHECKPOINTER] === undefined) {
+            // remove task_ids from checkpoint_ns
+            const recastCheckpointNamespace = checkpointNamespace
+                .split(CHECKPOINT_NAMESPACE_SEPARATOR)
+                .map((part) => {
+                return part.split(CHECKPOINT_NAMESPACE_END)[0];
+            })
+                .join(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(patchConfigurable(inputConfig, {
+                    [CONFIG_KEY_CHECKPOINTER]: checkpointer,
+                }), values, asNode);
+            }
+            throw new Error(`Subgraph "${recastCheckpointNamespace}" not found`);
+        }
+        // get last checkpoint
+        const config = this.config
+            ? mergeConfigs(this.config, inputConfig)
+            : inputConfig;
+        const saved = await checkpointer.getTuple(config);
+        const checkpoint = saved !== undefined
             ? copyCheckpoint(saved.checkpoint)
             : 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 = patchConfigurable(config, {
+            checkpoint_ns: config.configurable?.checkpoint_ns ?? "",
+        });
+        if (saved) {
+            checkpointConfig = 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, createCheckpoint(checkpoint, undefined, step), {
+            const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
                 source: "update",
                 step,
                 writes: {},
+                parents: saved?.metadata?.parents ?? {},
             }, {});
+            return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
         }
         const nonNullVersion = Object.values(checkpoint.versions_seen)
             .map((seenVersions) => {
@@ -315,7 +496,6 @@ export class Pregel extends Runnable {
             }
         }
         else if (asNode === undefined) {
-            // TODO: Double check
             const lastSeenByNode = Object.entries(checkpoint.versions_seen)
                 .map(([n, seen]) => {
                 return Object.values(seen).map((v) => {
@@ -372,18 +552,21 @@ export class Pregel extends 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?
-        _applyWrites(checkpoint, channels, [task], this.checkpointer.getNextVersion.bind(this.checkpointer));
+        _applyWrites(checkpoint, channels, [task], checkpointer.getNextVersion.bind(this.checkpointer));
         const newVersions = getNewChannelVersions(checkpointPreviousVersions, checkpoint.channel_versions);
-        return await this.checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, channels, step + 1), {
+        const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, channels, step + 1), {
             source: "update",
             step: step + 1,
             writes: { [asNode]: values },
+            parents: saved?.metadata?.parents ?? {},
         }, newVersions);
+        return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
     }
     _defaults(config) {
         const { debug, streamMode, inputKeys, outputKeys, interruptAfter, interruptBefore, ...rest } = config;
@@ -411,14 +594,13 @@ export class Pregel extends Runnable {
         else {
             defaultStreamMode = this.streamMode;
         }
-        let defaultCheckpointer;
-        if (config.configurable !== undefined &&
-            config.configurable[CONFIG_KEY_READ] !== undefined) {
+        // if being called as a node in another graph, always use values mode
+        if (config.configurable?.[CONFIG_KEY_TASK_ID] !== undefined) {
             defaultStreamMode = ["values"];
         }
+        let defaultCheckpointer;
         if (config !== undefined &&
-            config.configurable?.[CONFIG_KEY_CHECKPOINTER] !== undefined &&
-            (defaultInterruptAfter.length > 0 || defaultInterruptBefore.length > 0)) {
+            config.configurable?.[CONFIG_KEY_CHECKPOINTER] !== undefined) {
             defaultCheckpointer = config.configurable[CONFIG_KEY_CHECKPOINTER];
         }
         else {
@@ -497,7 +679,8 @@ export class Pregel extends Runnable {
         };
     }
     async *_streamIterator(input, options) {
-        const inputConfig = ensureConfig(options);
+        const streamSubgraphs = options?.subgraphs;
+        const inputConfig = ensureLangGraphConfig(this.config, options);
         if (inputConfig.recursionLimit === undefined ||
             inputConfig.recursionLimit < 1) {
             throw new Error(`Passed "recursionLimit" must be at least 1.`);
@@ -513,6 +696,30 @@ export class Pregel extends Runnable {
         const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer,] = this._defaults(inputConfig);
         const { channelSpecs, managed } = await this.prepareSpecs(config);
         let loop;
+        const stream = new Deque();
+        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 PregelLoop.initialize({
                 input,
@@ -524,7 +731,14 @@ export class Pregel extends Runnable {
                 outputKeys,
                 streamKeys: this.streamChannelsAsIs,
                 store: this.store,
+                stream: new StreamProtocol((chunk) => stream.push(chunk), new Set(streamMode)),
             });
+            if (options?.subgraphs) {
+                loop.config.configurable = {
+                    ...loop.config.configurable,
+                    [CONFIG_KEY_STREAM]: loop.stream,
+                };
+            }
             while (await loop.tick({
                 inputKeys: this.inputChannels,
                 interruptAfter,
@@ -534,27 +748,14 @@ export class Pregel extends Runnable {
                 if (debug) {
                     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) {
-                    printStepTasks(loop.step, loop.tasks);
+                    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 = executeTasksWithRetry(loop.tasks.filter((task) => task.writes.length === 0), {
+                const taskStream = executeTasksWithRetry(Object.values(loop.tasks).filter((task) => task.writes.length === 0), {
                     stepTimeout: this.stepTimeout,
                     signal: config.signal,
                     retryPolicy: this.retryPolicy,
@@ -563,7 +764,12 @@ export class Pregel extends Runnable {
                 for await (const { task, error } of taskStream) {
                     if (error !== undefined) {
                         if (isGraphInterrupt(error)) {
-                            loop.putWrites(task.id, error.interrupts.map((interrupt) => [INTERRUPT, interrupt]));
+                            if (loop.isNested) {
+                                throw error;
+                            }
+                            if (error.interrupts.length) {
+                                loop.putWrites(task.id, error.interrupts.map((interrupt) => [INTERRUPT, interrupt]));
+                            }
                         }
                         else {
                             loop.putWrites(task.id, [
@@ -574,42 +780,18 @@ export class Pregel extends 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 && !isGraphInterrupt(error)) {
                         throw error;
                     }
                 }
                 if (debug) {
-                    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;
-                    }
+                    printStepWrites(loop.step, Object.values(loop.tasks)
+                        .map((task) => task.writes)
+                        .flat(), this.streamChannelsList);
                 }
             }
+            yield* emitCurrentLoopOutputs();
             if (loop.status === "out_of_steps") {
                 throw new GraphRecursionError([
                     `Recursion limit of ${config.recursionLimit} reached`,
@@ -618,7 +800,7 @@ export class Pregel extends Runnable {
                 ].join(" "));
             }
             await Promise.all(loop?.checkpointerPromises ?? []);
-            await runManager?.handleChainEnd(readChannels(loop.channels, outputKeys));
+            await runManager?.handleChainEnd(loop.output);
         }
         catch (e) {
             await runManager?.handleChainError(e);
@@ -653,7 +835,7 @@ export class Pregel extends Runnable {
     async invoke(input, options) {
         const streamMode = options?.streamMode ?? "values";
         const config = {
-            ...ensureConfig(options),
+            ...options,
             outputKeys: options?.outputKeys ?? this.outputChannels,
             streamMode,
         };

package/dist/pregel/io.cjs

@@ -79,12 +79,14 @@ function* mapOutputValues(outputChannels, pendingWrites, channels
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 ) {
     if (Array.isArray(outputChannels)) {
-        if (pendingWrites.find(([chan, _]) => outputChannels.includes(chan))) {
+        if (pendingWrites === true ||
+            pendingWrites.find(([chan, _]) => outputChannels.includes(chan))) {
             yield readChannels(channels, outputChannels);
         }
     }
     else {
-        if (pendingWrites.some(([chan, _]) => chan === outputChannels)) {
+        if (pendingWrites === true ||
+            pendingWrites.some(([chan, _]) => chan === outputChannels)) {
             // eslint-disable-next-line @typescript-eslint/no-explicit-any
             yield readChannel(channels, outputChannels);
         }
@@ -94,30 +96,32 @@ exports.mapOutputValues = mapOutputValues;
 /**
  * Map pending writes (a sequence of tuples (channel, value)) to output chunk.
  */
-function* mapOutputUpdates(outputChannels, tasks
+function* mapOutputUpdates(outputChannels, tasks, cached
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 ) {
-    const outputTasks = tasks.filter((task) => task.config === undefined || !task.config.tags?.includes(constants_js_1.TAG_HIDDEN));
+    const outputTasks = tasks.filter(([task]) => {
+        return task.config === undefined || !task.config.tags?.includes(constants_js_1.TAG_HIDDEN);
+    });
     if (!outputTasks.length) {
         return;
     }
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     let updated;
     if (!Array.isArray(outputChannels)) {
-        updated = outputTasks.flatMap((task) => task.writes
+        updated = outputTasks.flatMap(([task]) => task.writes
             .filter(([chan, _]) => chan === outputChannels)
             // eslint-disable-next-line @typescript-eslint/no-explicit-any
             .map(([_, value]) => [task.name, value]));
     }
     else {
         updated = outputTasks
-            .filter((task) => task.writes.some(([chan]) => outputChannels.includes(chan)))
-            .map((task) => [
+            .filter(([task]) => task.writes.some(([chan]) => outputChannels.includes(chan)))
+            .map(([task]) => [
             task.name,
             Object.fromEntries(task.writes.filter(([chan]) => outputChannels.includes(chan))),
         ]);
     }
-    const grouped = Object.fromEntries(outputTasks.map((t) => [t.name, []])
+    const grouped = Object.fromEntries(outputTasks.map(([t]) => [t.name, []])
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     );
     for (const [node, value] of updated) {
@@ -133,6 +137,9 @@ function* mapOutputUpdates(outputChannels, tasks
             grouped[node] = value[0];
         }
     }
+    if (cached) {
+        grouped["__metadata__"] = { cached };
+    }
     yield grouped;
 }
 exports.mapOutputUpdates = mapOutputUpdates;

package/dist/pregel/io.d.ts

@@ -10,9 +10,9 @@ export declare function mapInput<C extends PropertyKey>(inputChannels: C | Array
 /**
  * Map pending writes (a sequence of tuples (channel, value)) to output chunk.
  */
-export declare function mapOutputValues<C extends PropertyKey>(outputChannels: C | Array<C>, pendingWrites: readonly PendingWrite<C>[], channels: Record<C, BaseChannel>): Generator<Record<string, any>, any>;
+export declare function mapOutputValues<C extends PropertyKey>(outputChannels: C | Array<C>, pendingWrites: readonly PendingWrite<C>[] | true, channels: Record<C, BaseChannel>): Generator<Record<string, any>, any>;
 /**
  * Map pending writes (a sequence of tuples (channel, value)) to output chunk.
  */
-export declare function mapOutputUpdates<N extends PropertyKey, C extends PropertyKey>(outputChannels: C | Array<C>, tasks: readonly PregelExecutableTask<N, C>[]): Generator<Record<N, Record<string, any> | Record<string, any>[]>>;
+export declare function mapOutputUpdates<N extends PropertyKey, C extends PropertyKey>(outputChannels: C | Array<C>, tasks: readonly [PregelExecutableTask<N, C>, PendingWrite<C>[]][], cached?: boolean): Generator<Record<N, Record<string, any> | any>>;
 export declare function single<T>(iter: IterableIterator<T>): T | null;