@langchain/langgraph 0.3.4 → 0.3.6

package/dist/pregel/index.cjs

@@ -679,7 +679,7 @@ class Pregel extends PartialRunnable {
         return {
             values: (0, io_js_1.readChannels)(channels, this.streamChannelsAsIs),
             next: nextList,
-            tasks: (0, debug_js_1.tasksWithWrites)(nextTasks, saved?.pendingWrites ?? [], taskStates),
+            tasks: (0, debug_js_1.tasksWithWrites)(nextTasks, saved?.pendingWrites ?? [], taskStates, this.streamChannelsAsIs),
             metadata,
             config: (0, index_js_1.patchCheckpointMap)(saved.config, saved.metadata),
             createdAt: saved.checkpoint.ts,
@@ -862,7 +862,7 @@ class Pregel extends PartialRunnable {
                     // tasks for this checkpoint
                     const nextTasks = (0, algo_js_1._prepareNextTasks)(checkpoint, saved.pendingWrites || [], this.nodes, channels, managed, saved.config, true, {
                         step: (saved.metadata?.step ?? -1) + 1,
-                        checkpointer: this.checkpointer || undefined,
+                        checkpointer,
                         store: this.store,
                     });
                     // apply null writes
@@ -870,13 +870,13 @@ class Pregel extends PartialRunnable {
                         .filter((w) => w[0] === constants_js_1.NULL_TASK_ID)
                         .map((w) => w.slice(1));
                     if (nullWrites.length > 0) {
-                        (0, algo_js_1._applyWrites)(saved.checkpoint, channels, [
+                        (0, algo_js_1._applyWrites)(checkpoint, channels, [
                             {
                                 name: constants_js_1.INPUT,
                                 writes: nullWrites,
                                 triggers: [],
                             },
-                        ], undefined, this.triggerToNodes);
+                        ], checkpointer.getNextVersion.bind(checkpointer), this.triggerToNodes);
                     }
                     // apply writes from tasks that already ran
                     for (const [taskId, k, v] of saved.pendingWrites || []) {
@@ -889,29 +889,61 @@ class Pregel extends PartialRunnable {
                         nextTasks[taskId].writes.push([k, v]);
                     }
                     // clear all current tasks
-                    (0, algo_js_1._applyWrites)(checkpoint, channels, Object.values(nextTasks), undefined, this.triggerToNodes);
+                    (0, algo_js_1._applyWrites)(checkpoint, channels, Object.values(nextTasks), checkpointer.getNextVersion.bind(checkpointer), this.triggerToNodes);
                 }
                 // save checkpoint
-                const nextConfig = await checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, undefined, step), {
+                const nextConfig = await checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, channels, step), {
                     ...checkpointMetadata,
                     source: "update",
                     step: step + 1,
                     writes: {},
                     parents: saved?.metadata?.parents ?? {},
-                }, {});
+                }, (0, index_js_1.getNewChannelVersions)(checkpointPreviousVersions, checkpoint.channel_versions));
                 return (0, index_js_1.patchCheckpointMap)(nextConfig, saved ? saved.metadata : undefined);
             }
-            if (values == null && asNode === constants_js_1.COPY) {
+            if (asNode === constants_js_1.COPY) {
                 if (updates.length > 1) {
                     throw new errors_js_1.InvalidUpdateError(`Cannot copy checkpoint with multiple updates`);
                 }
-                const nextConfig = await checkpointer.put(saved?.parentConfig ?? checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, undefined, step), {
+                if (saved == null) {
+                    throw new errors_js_1.InvalidUpdateError(`Cannot copy a non-existent checkpoint`);
+                }
+                const isCopyWithUpdates = (values) => {
+                    if (!Array.isArray(values))
+                        return false;
+                    if (values.length === 0)
+                        return false;
+                    return values.every((v) => Array.isArray(v) && v.length === 2);
+                };
+                const nextCheckpoint = (0, base_js_1.createCheckpoint)(checkpoint, undefined, step);
+                const nextConfig = await checkpointer.put(saved.parentConfig ??
+                    (0, utils_js_1.patchConfigurable)(saved.config, { checkpoint_id: undefined }), nextCheckpoint, {
                     source: "fork",
                     step: step + 1,
                     writes: {},
-                    parents: saved?.metadata?.parents ?? {},
+                    parents: saved.metadata?.parents ?? {},
                 }, {});
-                return (0, index_js_1.patchCheckpointMap)(nextConfig, saved ? saved.metadata : undefined);
+                // We want to both clone a checkpoint and update state in one go.
+                // Reuse the same task ID if possible.
+                if (isCopyWithUpdates(values)) {
+                    // figure out the task IDs for the next update checkpoint
+                    const nextTasks = (0, algo_js_1._prepareNextTasks)(nextCheckpoint, saved.pendingWrites, this.nodes, channels, managed, nextConfig, false, { step: step + 2 });
+                    const tasksGroupBy = Object.values(nextTasks).reduce((acc, { name, id }) => {
+                        acc[name] ??= [];
+                        acc[name].push({ id });
+                        return acc;
+                    }, {});
+                    const userGroupBy = values.reduce((acc, item) => {
+                        const [values, asNode] = item;
+                        acc[asNode] ??= [];
+                        const targetIdx = acc[asNode].length;
+                        const taskId = tasksGroupBy[asNode]?.[targetIdx]?.id;
+                        acc[asNode].push({ values, asNode, taskId });
+                        return acc;
+                    }, {});
+                    return updateSuperStep((0, index_js_1.patchCheckpointMap)(nextConfig, saved.metadata), Object.values(userGroupBy).flat());
+                }
+                return (0, index_js_1.patchCheckpointMap)(nextConfig, saved.metadata);
             }
             if (asNode === constants_js_1.INPUT) {
                 if (updates.length > 1) {
@@ -983,7 +1015,7 @@ class Pregel extends PartialRunnable {
             const validUpdates = [];
             if (updates.length === 1) {
                 // eslint-disable-next-line prefer-const
-                let { values, asNode } = updates[0];
+                let { values, asNode, taskId } = updates[0];
                 if (asNode === undefined && Object.keys(this.nodes).length === 1) {
                     // if only one node, use it
                     [asNode] = Object.keys(this.nodes);
@@ -1019,18 +1051,18 @@ class Pregel extends PartialRunnable {
                 if (asNode === undefined) {
                     throw new errors_js_1.InvalidUpdateError(`Ambiguous update, specify "asNode"`);
                 }
-                validUpdates.push({ values, asNode });
+                validUpdates.push({ values, asNode, taskId });
             }
             else {
-                for (const { asNode, values } of updates) {
+                for (const { asNode, values, taskId } of updates) {
                     if (asNode == null) {
                         throw new errors_js_1.InvalidUpdateError(`"asNode" is required when applying multiple updates`);
                     }
-                    validUpdates.push({ values, asNode });
+                    validUpdates.push({ values, asNode, taskId });
                 }
             }
             const tasks = [];
-            for (const { asNode, values } of validUpdates) {
+            for (const { asNode, values, taskId } of validUpdates) {
                 if (this.nodes[asNode] === undefined) {
                     throw new errors_js_1.InvalidUpdateError(`Node "${asNode.toString()}" does not exist`);
                 }
@@ -1050,7 +1082,7 @@ class Pregel extends PartialRunnable {
                         : writers[0],
                     writes: [],
                     triggers: [constants_js_1.INTERRUPT],
-                    id: (0, langgraph_checkpoint_1.uuid5)(constants_js_1.INTERRUPT, checkpoint.id),
+                    id: taskId ?? (0, langgraph_checkpoint_1.uuid5)(constants_js_1.INTERRUPT, checkpoint.id),
                     writers: [],
                 });
             }
@@ -1139,6 +1171,8 @@ class Pregel extends PartialRunnable {
      * - checkpointer
      * - store
      * - whether stream mode is single
+     * - node cache
+     * - whether checkpoint during is enabled
      * @internal
      */
     _defaults(config) {
@@ -1167,13 +1201,16 @@ class Pregel extends PartialRunnable {
             streamModeSingle = typeof streamMode === "string";
         }
         else {
-            defaultStreamMode = this.streamMode;
+            // if being called as a node in another graph, default to values mode
+            // but don't overwrite `streamMode`if provided
+            if (config.configurable?.[constants_js_1.CONFIG_KEY_TASK_ID] !== undefined) {
+                defaultStreamMode = ["values"];
+            }
+            else {
+                defaultStreamMode = this.streamMode;
+            }
             streamModeSingle = true;
         }
-        // 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 (this.checkpointer === false) {
             defaultCheckpointer = undefined;
@@ -1182,11 +1219,17 @@ class Pregel extends PartialRunnable {
             config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINTER] !== undefined) {
             defaultCheckpointer = config.configurable[constants_js_1.CONFIG_KEY_CHECKPOINTER];
         }
+        else if (this.checkpointer === true) {
+            throw new Error("checkpointer: true cannot be used for root graphs.");
+        }
         else {
             defaultCheckpointer = this.checkpointer;
         }
         const defaultStore = config.store ?? this.store;
         const defaultCache = config.cache ?? this.cache;
+        const defaultCheckpointDuring = config.checkpointDuring ??
+            config?.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINT_DURING] ??
+            true;
         return [
             defaultDebug,
             defaultStreamMode,
@@ -1199,6 +1242,7 @@ class Pregel extends PartialRunnable {
             defaultStore,
             streamModeSingle,
             defaultCache,
+            defaultCheckpointDuring,
         ];
     }
     /**
@@ -1345,11 +1389,20 @@ class Pregel extends PartialRunnable {
         const validInput = await this._validateInput(input);
         const { runId, ...restConfig } = inputConfig;
         // assign defaults
-        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store, streamModeSingle, cache,] = this._defaults(restConfig);
+        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store, streamModeSingle, cache, checkpointDuring,] = this._defaults(restConfig);
         config.configurable = await this._validateConfigurable(config.configurable);
         const stream = new stream_js_1.IterableReadableWritableStream({
             modes: new Set(streamMode),
         });
+        // set up subgraph checkpointing
+        if (this.checkpointer === true) {
+            config.configurable ??= {};
+            const ns = config.configurable[constants_js_1.CONFIG_KEY_CHECKPOINT_NS] ?? "";
+            config.configurable[constants_js_1.CONFIG_KEY_CHECKPOINT_NS] = ns
+                .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);
+        }
         // set up messages stream mode
         if (streamMode.includes("messages")) {
             const messageStreamer = new messages_js_1.StreamMessagesHandler((chunk) => stream.push(chunk));
@@ -1409,6 +1462,7 @@ class Pregel extends PartialRunnable {
                     manager: runManager,
                     debug: this.debug,
                     triggerToNodes: this.triggerToNodes,
+                    checkpointDuring,
                 });
                 const runner = new runner_js_1.PregelRunner({
                     loop,

package/dist/pregel/index.d.ts

@@ -222,7 +222,7 @@ export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channel
      * When provided, saves a checkpoint of the graph state at every superstep.
      * When false or undefined, checkpointing is disabled, and the graph will not be able to save or restore state.
      */
-    checkpointer?: BaseCheckpointSaver | false;
+    checkpointer?: BaseCheckpointSaver | boolean;
     /** Optional retry policy for handling failures in node execution */
     retryPolicy?: RetryPolicy;
     /** The default configuration for graph execution, can be overridden on a per-invocation basis */
@@ -412,6 +412,8 @@ export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channel
      * - checkpointer
      * - store
      * - whether stream mode is single
+     * - node cache
+     * - whether checkpoint during is enabled
      * @internal
      */
     _defaults(config: PregelOptions<Nodes, Channels>): [
@@ -428,9 +430,12 @@ export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channel
         All | string[],
         // interrupt after
         BaseCheckpointSaver | undefined,
+        // checkpointer
         BaseStore | undefined,
         boolean,
-        BaseCache | undefined
+        // stream mode single
+        BaseCache | undefined,
+        boolean
     ];
     /**
      * Streams the execution of the graph, emitting state updates as they occur.

package/dist/pregel/index.js

@@ -2,7 +2,7 @@
 import { _coerceToRunnable, getCallbackManagerForConfig, mergeConfigs, patchConfig, Runnable, RunnableSequence, } from "@langchain/core/runnables";
 import { compareChannelVersions, copyCheckpoint, emptyCheckpoint, SCHEDULED, uuid5, } from "@langchain/langgraph-checkpoint";
 import { createCheckpoint, emptyChannels, isBaseChannel, } from "../channels/base.js";
-import { CHECKPOINT_NAMESPACE_END, CHECKPOINT_NAMESPACE_SEPARATOR, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_NODE_FINISHED, CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, COPY, END, ERROR, INPUT, INTERRUPT, isInterrupted, NULL_TASK_ID, PUSH, } from "../constants.js";
+import { CHECKPOINT_NAMESPACE_END, CHECKPOINT_NAMESPACE_SEPARATOR, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_NODE_FINISHED, CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, COPY, END, ERROR, INPUT, INTERRUPT, isInterrupted, NULL_TASK_ID, PUSH, CONFIG_KEY_CHECKPOINT_DURING, CONFIG_KEY_CHECKPOINT_NS, } from "../constants.js";
 import { GraphRecursionError, GraphValueError, InvalidUpdateError, } from "../errors.js";
 import { ChannelKeyPlaceholder, isConfiguredManagedValue, ManagedValueMapping, NoopManagedValue, } from "../managed/base.js";
 import { gatherIterator, patchConfigurable } from "../utils.js";
@@ -675,7 +675,7 @@ export class Pregel extends PartialRunnable {
         return {
             values: readChannels(channels, this.streamChannelsAsIs),
             next: nextList,
-            tasks: tasksWithWrites(nextTasks, saved?.pendingWrites ?? [], taskStates),
+            tasks: tasksWithWrites(nextTasks, saved?.pendingWrites ?? [], taskStates, this.streamChannelsAsIs),
             metadata,
             config: patchCheckpointMap(saved.config, saved.metadata),
             createdAt: saved.checkpoint.ts,
@@ -858,7 +858,7 @@ export class Pregel extends PartialRunnable {
                     // tasks for this checkpoint
                     const nextTasks = _prepareNextTasks(checkpoint, saved.pendingWrites || [], this.nodes, channels, managed, saved.config, true, {
                         step: (saved.metadata?.step ?? -1) + 1,
-                        checkpointer: this.checkpointer || undefined,
+                        checkpointer,
                         store: this.store,
                     });
                     // apply null writes
@@ -866,13 +866,13 @@ export class Pregel extends PartialRunnable {
                         .filter((w) => w[0] === NULL_TASK_ID)
                         .map((w) => w.slice(1));
                     if (nullWrites.length > 0) {
-                        _applyWrites(saved.checkpoint, channels, [
+                        _applyWrites(checkpoint, channels, [
                             {
                                 name: INPUT,
                                 writes: nullWrites,
                                 triggers: [],
                             },
-                        ], undefined, this.triggerToNodes);
+                        ], checkpointer.getNextVersion.bind(checkpointer), this.triggerToNodes);
                     }
                     // apply writes from tasks that already ran
                     for (const [taskId, k, v] of saved.pendingWrites || []) {
@@ -885,29 +885,61 @@ export class Pregel extends PartialRunnable {
                         nextTasks[taskId].writes.push([k, v]);
                     }
                     // clear all current tasks
-                    _applyWrites(checkpoint, channels, Object.values(nextTasks), undefined, this.triggerToNodes);
+                    _applyWrites(checkpoint, channels, Object.values(nextTasks), checkpointer.getNextVersion.bind(checkpointer), this.triggerToNodes);
                 }
                 // save checkpoint
-                const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
+                const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, channels, step), {
                     ...checkpointMetadata,
                     source: "update",
                     step: step + 1,
                     writes: {},
                     parents: saved?.metadata?.parents ?? {},
-                }, {});
+                }, getNewChannelVersions(checkpointPreviousVersions, checkpoint.channel_versions));
                 return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
             }
-            if (values == null && asNode === COPY) {
+            if (asNode === COPY) {
                 if (updates.length > 1) {
                     throw new InvalidUpdateError(`Cannot copy checkpoint with multiple updates`);
                 }
-                const nextConfig = await checkpointer.put(saved?.parentConfig ?? checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
+                if (saved == null) {
+                    throw new InvalidUpdateError(`Cannot copy a non-existent checkpoint`);
+                }
+                const isCopyWithUpdates = (values) => {
+                    if (!Array.isArray(values))
+                        return false;
+                    if (values.length === 0)
+                        return false;
+                    return values.every((v) => Array.isArray(v) && v.length === 2);
+                };
+                const nextCheckpoint = createCheckpoint(checkpoint, undefined, step);
+                const nextConfig = await checkpointer.put(saved.parentConfig ??
+                    patchConfigurable(saved.config, { checkpoint_id: undefined }), nextCheckpoint, {
                     source: "fork",
                     step: step + 1,
                     writes: {},
-                    parents: saved?.metadata?.parents ?? {},
+                    parents: saved.metadata?.parents ?? {},
                 }, {});
-                return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
+                // We want to both clone a checkpoint and update state in one go.
+                // Reuse the same task ID if possible.
+                if (isCopyWithUpdates(values)) {
+                    // figure out the task IDs for the next update checkpoint
+                    const nextTasks = _prepareNextTasks(nextCheckpoint, saved.pendingWrites, this.nodes, channels, managed, nextConfig, false, { step: step + 2 });
+                    const tasksGroupBy = Object.values(nextTasks).reduce((acc, { name, id }) => {
+                        acc[name] ??= [];
+                        acc[name].push({ id });
+                        return acc;
+                    }, {});
+                    const userGroupBy = values.reduce((acc, item) => {
+                        const [values, asNode] = item;
+                        acc[asNode] ??= [];
+                        const targetIdx = acc[asNode].length;
+                        const taskId = tasksGroupBy[asNode]?.[targetIdx]?.id;
+                        acc[asNode].push({ values, asNode, taskId });
+                        return acc;
+                    }, {});
+                    return updateSuperStep(patchCheckpointMap(nextConfig, saved.metadata), Object.values(userGroupBy).flat());
+                }
+                return patchCheckpointMap(nextConfig, saved.metadata);
             }
             if (asNode === INPUT) {
                 if (updates.length > 1) {
@@ -979,7 +1011,7 @@ export class Pregel extends PartialRunnable {
             const validUpdates = [];
             if (updates.length === 1) {
                 // eslint-disable-next-line prefer-const
-                let { values, asNode } = updates[0];
+                let { values, asNode, taskId } = updates[0];
                 if (asNode === undefined && Object.keys(this.nodes).length === 1) {
                     // if only one node, use it
                     [asNode] = Object.keys(this.nodes);
@@ -1015,18 +1047,18 @@ export class Pregel extends PartialRunnable {
                 if (asNode === undefined) {
                     throw new InvalidUpdateError(`Ambiguous update, specify "asNode"`);
                 }
-                validUpdates.push({ values, asNode });
+                validUpdates.push({ values, asNode, taskId });
             }
             else {
-                for (const { asNode, values } of updates) {
+                for (const { asNode, values, taskId } of updates) {
                     if (asNode == null) {
                         throw new InvalidUpdateError(`"asNode" is required when applying multiple updates`);
                     }
-                    validUpdates.push({ values, asNode });
+                    validUpdates.push({ values, asNode, taskId });
                 }
             }
             const tasks = [];
-            for (const { asNode, values } of validUpdates) {
+            for (const { asNode, values, taskId } of validUpdates) {
                 if (this.nodes[asNode] === undefined) {
                     throw new InvalidUpdateError(`Node "${asNode.toString()}" does not exist`);
                 }
@@ -1046,7 +1078,7 @@ export class Pregel extends PartialRunnable {
                         : writers[0],
                     writes: [],
                     triggers: [INTERRUPT],
-                    id: uuid5(INTERRUPT, checkpoint.id),
+                    id: taskId ?? uuid5(INTERRUPT, checkpoint.id),
                     writers: [],
                 });
             }
@@ -1135,6 +1167,8 @@ export class Pregel extends PartialRunnable {
      * - checkpointer
      * - store
      * - whether stream mode is single
+     * - node cache
+     * - whether checkpoint during is enabled
      * @internal
      */
     _defaults(config) {
@@ -1163,13 +1197,16 @@ export class Pregel extends PartialRunnable {
             streamModeSingle = typeof streamMode === "string";
         }
         else {
-            defaultStreamMode = this.streamMode;
+            // if being called as a node in another graph, default to values mode
+            // but don't overwrite `streamMode`if provided
+            if (config.configurable?.[CONFIG_KEY_TASK_ID] !== undefined) {
+                defaultStreamMode = ["values"];
+            }
+            else {
+                defaultStreamMode = this.streamMode;
+            }
             streamModeSingle = true;
         }
-        // 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 (this.checkpointer === false) {
             defaultCheckpointer = undefined;
@@ -1178,11 +1215,17 @@ export class Pregel extends PartialRunnable {
             config.configurable?.[CONFIG_KEY_CHECKPOINTER] !== undefined) {
             defaultCheckpointer = config.configurable[CONFIG_KEY_CHECKPOINTER];
         }
+        else if (this.checkpointer === true) {
+            throw new Error("checkpointer: true cannot be used for root graphs.");
+        }
         else {
             defaultCheckpointer = this.checkpointer;
         }
         const defaultStore = config.store ?? this.store;
         const defaultCache = config.cache ?? this.cache;
+        const defaultCheckpointDuring = config.checkpointDuring ??
+            config?.configurable?.[CONFIG_KEY_CHECKPOINT_DURING] ??
+            true;
         return [
             defaultDebug,
             defaultStreamMode,
@@ -1195,6 +1238,7 @@ export class Pregel extends PartialRunnable {
             defaultStore,
             streamModeSingle,
             defaultCache,
+            defaultCheckpointDuring,
         ];
     }
     /**
@@ -1341,11 +1385,20 @@ export class Pregel extends PartialRunnable {
         const validInput = await this._validateInput(input);
         const { runId, ...restConfig } = inputConfig;
         // assign defaults
-        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store, streamModeSingle, cache,] = this._defaults(restConfig);
+        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store, streamModeSingle, cache, checkpointDuring,] = this._defaults(restConfig);
         config.configurable = await this._validateConfigurable(config.configurable);
         const stream = new IterableReadableWritableStream({
             modes: new Set(streamMode),
         });
+        // set up subgraph checkpointing
+        if (this.checkpointer === true) {
+            config.configurable ??= {};
+            const ns = config.configurable[CONFIG_KEY_CHECKPOINT_NS] ?? "";
+            config.configurable[CONFIG_KEY_CHECKPOINT_NS] = ns
+                .split(CHECKPOINT_NAMESPACE_SEPARATOR)
+                .map((part) => part.split(CHECKPOINT_NAMESPACE_END)[0])
+                .join(CHECKPOINT_NAMESPACE_SEPARATOR);
+        }
         // set up messages stream mode
         if (streamMode.includes("messages")) {
             const messageStreamer = new StreamMessagesHandler((chunk) => stream.push(chunk));
@@ -1405,6 +1458,7 @@ export class Pregel extends PartialRunnable {
                     manager: runManager,
                     debug: this.debug,
                     triggerToNodes: this.triggerToNodes,
+                    checkpointDuring,
                 });
                 const runner = new PregelRunner({
                     loop,