@langchain/langgraph
Version:
941 lines (940 loc) • 41.7 kB
JavaScript
import { CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_READ, CONFIG_KEY_REPLAY_STATE, CONFIG_KEY_RESUME_MAP, CONFIG_KEY_RESUMING, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_STREAM, ERROR, INPUT, INTERRUPT, NULL_TASK_ID, PUSH, RESUME, _isOverwriteValue, _isSend, isCommand } from "../constants.js";
import { EmptyInputError, GraphInterrupt, isGraphInterrupt } from "../errors.js";
import { channelsFromCheckpoint, createCheckpoint, deltaChannelsToSnapshot, exitDeltaTaskId, isDeltaChannel } from "../channels/base.js";
import { gatherIterator, gatherIteratorSync, prefixGenerator } from "../utils.js";
import { isXXH3 } from "../hash.js";
import { mapCommand, mapInput, mapOutputUpdates, mapOutputValues, readChannels } from "./io.js";
import { getNewChannelVersions, patchConfigurable } from "./utils/index.js";
import { _applyWrites, _prepareNextTasks, _prepareNodeErrorHandlerTask, _prepareSingleTask, increment, sanitizeUntrackedValuesInSend, shouldInterrupt } from "./algo.js";
import { mapDebugCheckpoint, mapDebugTaskResults, mapDebugTasks, printStepTasks } from "./debug.js";
import { ReplayState } from "./replay.js";
import { createDuplexStream } from "./stream.js";
import { AsyncBatchedStore, BaseCache, WRITES_IDX_MAP, copyCheckpoint, emptyCheckpoint } from "@langchain/langgraph-checkpoint";
import { v4 } from "@langchain/core/utils/uuid";
import { BaseMessage } from "@langchain/core/messages";
//#region src/pregel/loop.ts
const INPUT_DONE = Symbol.for("INPUT_DONE");
const INPUT_RESUMING = Symbol.for("INPUT_RESUMING");
const DEFAULT_LOOP_LIMIT = 25;
/**
* Recursively assign a stable UUID to any {@link BaseMessage} (in a value, an
* array, or an object's values) that is missing an `id`. Used so DeltaChannel
* writes — replayed on every read — reconstruct identical message identities.
*/
function ensureMessageIds(value) {
if (value == null || typeof value !== "object") return;
if (BaseMessage.isInstance(value)) {
const msg = value;
if (msg.id == null) {
msg.id = v4();
if (msg.lc_kwargs != null) msg.lc_kwargs.id = msg.id;
}
return;
}
if (Array.isArray(value)) {
for (const item of value) ensureMessageIds(item);
return;
}
}
/**
* Split a serialized checkpoint namespace into its path segments.
*
* Checkpoint namespaces are stored as a single string whose nested levels are
* joined by {@link CHECKPOINT_NAMESPACE_SEPARATOR} (e.g. `"parent|child"`).
* The root namespace — represented as `undefined` or the empty string — maps
* to an empty array.
*
* @param ns - The serialized checkpoint namespace, or `undefined`.
* @returns The namespace as an array of path segments (`[]` for the root).
*/
function checkpointNamespaceFromNs(ns) {
if (ns === void 0 || ns === "") return [];
return ns.split("|");
}
/**
* Find the most deeply nested namespace recorded in a checkpoint map.
*
* The checkpoint map ({@link CONFIG_KEY_CHECKPOINT_MAP}) associates every
* namespace seen on a thread with its checkpoint id. Because nested namespaces
* are built by appending segments to their parent, a deeper namespace always
* yields a longer key — so the longest non-empty key is the deepest one.
*
* Used by the loop's `#interruptStreamNamespace()` during subgraph
* time-travel: interrupt events must be emitted against the active (deepest)
* subgraph namespace rather than the root graph.
*
* @param map - The checkpoint map (namespace -> checkpoint id), or `undefined`.
* @returns The deepest namespace as path segments, or `[]` when the map is
* absent, empty, or only contains the root namespace.
*/
function deepestCheckpointMapNamespace(map) {
if (!map) return [];
let deepest = "";
for (const key of Object.keys(map)) if (key !== "" && key.length > deepest.length) deepest = key;
return checkpointNamespaceFromNs(deepest);
}
var AsyncBatchedCache = class extends BaseCache {
cache;
queue = Promise.resolve();
constructor(cache) {
super();
this.cache = cache;
}
async get(keys) {
return this.enqueueOperation("get", keys);
}
async set(pairs) {
return this.enqueueOperation("set", pairs);
}
async clear(namespaces) {
return this.enqueueOperation("clear", namespaces);
}
async stop() {
await this.queue;
}
enqueueOperation(type, ...args) {
const newPromise = this.queue.then(() => {
return this.cache[type](...args);
});
this.queue = newPromise.then(() => void 0, () => void 0);
return newPromise;
}
};
var PregelLoop = class PregelLoop {
input;
output;
config;
checkpointer;
checkpointerGetNextVersion;
channels;
checkpoint;
checkpointIdSaved;
/**
* Exit-mode accumulator of DeltaChannel writes across the whole run, as
* `[step, taskId, channel, value]`. `undefined` outside "exit" durability.
*/
_exitDeltaWrites;
/**
* DeltaChannels that saw an Overwrite since the last checkpoint. These
* channels are force-snapshotted at the next checkpoint so reconstruction
* starts from the post-overwrite value and never has to replay across the
* reset (the live `update` discards every sibling write in the overwriting
* super-step). Cleared once the channel snapshots.
*/
_deltaChannelsWithOverwrite = /* @__PURE__ */ new Set();
/** Whether a real checkpoint was loaded from the saver at initialization. */
_hasPersistedParent = false;
/** The checkpointConfig as captured at initialization (anchor for exit writes). */
_initialCheckpointConfig;
checkpointConfig;
checkpointMetadata;
checkpointNamespace;
checkpointPendingWrites = [];
checkpointPreviousVersions;
step;
stop;
durability;
outputKeys;
streamKeys;
nodes;
skipDoneTasks;
prevCheckpointConfig;
updatedChannels;
status = "pending";
/**
* Run-scoped control surface for cooperative draining. Populated from the
* run config. When `control.drainRequested` is true, the loop stops at the
* next superstep boundary instead of dispatching more tasks.
*/
control;
tasks = {};
stream;
checkpointerPromises = /* @__PURE__ */ new Set();
isNested;
/** True when an explicit checkpoint_id targets the latest saved checkpoint. */
resumeAtHead;
_checkpointerChainedPromise = Promise.resolve();
/**
* Track a checkpointer promise, removing it from the set on success.
* Failed promises are kept so that Promise.all() in the finally block
* of _streamIterator can surface the error.
*
* @internal
*/
_trackCheckpointerPromise(promise) {
const tracked = promise.then((value) => {
this.checkpointerPromises.delete(tracked);
return value;
}, (error) => {
throw error;
});
this.checkpointerPromises.add(tracked);
}
store;
cache;
manager;
interruptAfter;
interruptBefore;
toInterrupt = [];
debug = false;
triggerToNodes;
get isResuming() {
let hasChannelVersions = false;
if ("__start__" in this.checkpoint.channel_versions) hasChannelVersions = true;
else for (const chan in this.checkpoint.channel_versions) if (Object.prototype.hasOwnProperty.call(this.checkpoint.channel_versions, chan)) {
hasChannelVersions = true;
break;
}
const configIsResuming = this.config.configurable?.["__pregel_resuming"] !== void 0 && this.config.configurable?.["__pregel_resuming"];
const inputIsNullOrUndefined = this.input === null || this.input === void 0;
const inputIsCommandResuming = isCommand(this.input) && this.input.resume != null;
const inputIsResuming = this.input === INPUT_RESUMING;
const runIdMatchesPrevious = !this.isNested && this.config.metadata?.run_id !== void 0 && this.checkpointMetadata?.run_id !== void 0 && this.config.metadata.run_id === this.checkpointMetadata?.run_id;
return hasChannelVersions && (configIsResuming || inputIsNullOrUndefined || inputIsCommandResuming || inputIsResuming || runIdMatchesPrevious);
}
get isReplaying() {
return !this.skipDoneTasks;
}
constructor(params) {
this.input = params.input;
this.checkpointer = params.checkpointer;
if (this.checkpointer !== void 0) this.checkpointerGetNextVersion = this.checkpointer.getNextVersion.bind(this.checkpointer);
else this.checkpointerGetNextVersion = increment;
this.checkpoint = params.checkpoint;
this.checkpointMetadata = params.checkpointMetadata;
this.checkpointPreviousVersions = params.checkpointPreviousVersions;
this.channels = params.channels;
this.checkpointPendingWrites = params.checkpointPendingWrites;
this.step = params.step;
this.stop = params.stop;
this.config = params.config;
this.checkpointConfig = params.checkpointConfig;
this.isNested = params.isNested;
this.resumeAtHead = params.resumeAtHead;
this.manager = params.manager;
this.outputKeys = params.outputKeys;
this.streamKeys = params.streamKeys;
this.nodes = params.nodes;
this.skipDoneTasks = params.skipDoneTasks;
this.store = params.store;
this.cache = params.cache ? new AsyncBatchedCache(params.cache) : void 0;
this.stream = params.stream;
this.checkpointNamespace = params.checkpointNamespace;
this.prevCheckpointConfig = params.prevCheckpointConfig;
this.interruptAfter = params.interruptAfter;
this.interruptBefore = params.interruptBefore;
this.durability = params.durability;
this.debug = params.debug;
this.triggerToNodes = params.triggerToNodes;
this.control = this.config.control;
this._exitDeltaWrites = this.durability === "exit" && this.checkpointer != null ? [] : void 0;
this._hasPersistedParent = params.hasPersistedParent ?? false;
this._initialCheckpointConfig = params.checkpointConfig;
this.checkpointIdSaved = params.checkpoint.id;
}
static async initialize(params) {
let { config, stream } = params;
if (stream !== void 0 && config.configurable?.["__pregel_stream"] !== void 0) stream = createDuplexStream(stream, config.configurable[CONFIG_KEY_STREAM]);
const skipDoneTasks = config.configurable ? !("checkpoint_id" in config.configurable) : true;
const scratchpad = config.configurable?.[CONFIG_KEY_SCRATCHPAD];
if (config.configurable && scratchpad) {
if (scratchpad.subgraphCounter > 0) config = patchConfigurable(config, { [CONFIG_KEY_CHECKPOINT_NS]: [config.configurable[CONFIG_KEY_CHECKPOINT_NS], scratchpad.subgraphCounter.toString()].join("|") });
scratchpad.subgraphCounter += 1;
}
const requestedCheckpointId = config.configurable?.checkpoint_id;
const isNested = CONFIG_KEY_READ in (config.configurable ?? {});
if (!isNested && config.configurable?.checkpoint_ns !== void 0 && config.configurable?.checkpoint_ns !== "") config = patchConfigurable(config, {
checkpoint_ns: "",
checkpoint_id: void 0
});
let checkpointConfig = config;
if (config.configurable?.checkpoint_id === void 0 && config.configurable?.["checkpoint_map"] !== void 0 && config.configurable?.["checkpoint_map"]?.[config.configurable?.checkpoint_ns]) checkpointConfig = patchConfigurable(config, { checkpoint_id: config.configurable[CONFIG_KEY_CHECKPOINT_MAP][config.configurable?.checkpoint_ns] });
const checkpointNamespace = checkpointNamespaceFromNs(config.configurable?.checkpoint_ns);
let saved;
if (!params.checkpointer) saved = void 0;
else if (checkpointConfig.configurable?.["checkpoint_id"]) saved = await params.checkpointer.getTuple(checkpointConfig);
else if (config.configurable?.["__pregel_replay_state"]) {
saved = await config.configurable[CONFIG_KEY_REPLAY_STATE].getCheckpoint(config.configurable?.["checkpoint_ns"] ?? "", params.checkpointer, checkpointConfig);
if (config.configurable) delete config.configurable[CONFIG_KEY_RESUMING];
} else saved = await params.checkpointer.getTuple(checkpointConfig);
const hasPersistedParent = saved !== void 0;
if (!saved) saved = {
config,
checkpoint: emptyCheckpoint(),
metadata: {
source: "input",
step: -2,
parents: {}
},
pendingWrites: []
};
checkpointConfig = {
...config,
...saved.config,
configurable: {
checkpoint_ns: "",
...config.configurable,
...saved.config.configurable
}
};
const prevCheckpointConfig = saved.parentConfig;
const checkpoint = copyCheckpoint(saved.checkpoint);
const checkpointMetadata = { ...saved.metadata };
let checkpointPendingWrites = saved.pendingWrites ?? [];
const currentCheckpointNamespace = config.configurable?.checkpoint_ns;
const checkpointMap = config.configurable?.[CONFIG_KEY_CHECKPOINT_MAP];
if (typeof currentCheckpointNamespace === "string" && currentCheckpointNamespace !== "" && typeof checkpointMap === "object" && checkpointMap !== null && currentCheckpointNamespace in checkpointMap && checkpointPendingWrites.length > 0) checkpointPendingWrites = checkpointPendingWrites.filter(([, channel]) => channel !== RESUME);
let resumeAtHead = false;
const threadId = checkpointConfig.configurable?.thread_id;
const checkpointNs = checkpointConfig.configurable?.checkpoint_ns ?? "";
if (params.checkpointer && requestedCheckpointId && typeof threadId === "string") resumeAtHead = (await params.checkpointer.getTuple({ configurable: {
thread_id: threadId,
checkpoint_ns: checkpointNs
} }))?.config.configurable?.checkpoint_id === requestedCheckpointId && checkpointMetadata.source !== "update" && checkpointMetadata.source !== "fork";
const channels = await channelsFromCheckpoint(params.channelSpecs, checkpoint, {
saver: params.checkpointer,
config: checkpointConfig
});
const step = (checkpointMetadata.step ?? 0) + 1;
const stop = step + (config.recursionLimit ?? DEFAULT_LOOP_LIMIT) + 1;
const checkpointPreviousVersions = { ...checkpoint.channel_versions };
const store = params.store ? new AsyncBatchedStore(params.store) : void 0;
if (store) await store.start();
return new PregelLoop({
input: params.input,
config,
checkpointer: params.checkpointer,
checkpoint,
checkpointMetadata,
checkpointConfig,
prevCheckpointConfig,
checkpointNamespace,
channels,
isNested,
resumeAtHead,
manager: params.manager,
skipDoneTasks,
step,
stop,
checkpointPreviousVersions,
checkpointPendingWrites,
outputKeys: params.outputKeys ?? [],
streamKeys: params.streamKeys ?? [],
nodes: params.nodes,
stream,
store,
cache: params.cache,
interruptAfter: params.interruptAfter,
interruptBefore: params.interruptBefore,
durability: params.durability,
debug: params.debug,
triggerToNodes: params.triggerToNodes,
hasPersistedParent
});
}
_checkpointerPutAfterPrevious(input) {
this._checkpointerChainedPromise = this._checkpointerChainedPromise.then(() => {
return this.checkpointer?.put(input.config, input.checkpoint, input.metadata, input.newVersions);
});
this._trackCheckpointerPromise(this._checkpointerChainedPromise);
}
/**
* Put writes for a task, to be read by the next tick.
* @param taskId
* @param writes
*/
putWrites(taskId, writes) {
let writesCopy = writes;
if (writesCopy.length === 0) return;
if (writesCopy.every(([key]) => key in WRITES_IDX_MAP)) writesCopy = Array.from(new Map(writesCopy.map((w) => [w[0], w])).values());
let hasUntrackedChannels = false;
for (const key in this.channels) if (Object.prototype.hasOwnProperty.call(this.channels, key)) {
if (this.channels[key].lc_graph_name === "UntrackedValue") {
hasUntrackedChannels = true;
break;
}
}
let writesToSave = writesCopy;
if (hasUntrackedChannels) writesToSave = writesCopy.filter(([c]) => {
const channel = this.channels[c];
return !channel || channel.lc_graph_name !== "UntrackedValue";
}).map(([c, v]) => {
if (c === "__pregel_tasks" && _isSend(v)) return [c, sanitizeUntrackedValuesInSend(v, this.channels)];
return [c, v];
});
this.checkpointPendingWrites = this.checkpointPendingWrites.filter((w) => w[0] !== taskId);
for (const [c, v] of writesToSave) this.checkpointPendingWrites.push([
taskId,
c,
v
]);
for (const [c, v] of writesToSave) {
const channel = this.channels[c];
if (channel != null && isDeltaChannel(channel)) ensureMessageIds(v);
}
const config = patchConfigurable(this.checkpointConfig, {
[CONFIG_KEY_CHECKPOINT_NS]: this.config.configurable?.checkpoint_ns ?? "",
[CONFIG_KEY_CHECKPOINT_ID]: this.checkpoint.id
});
if (this.durability !== "exit" && this.checkpointer != null) this._trackCheckpointerPromise(this.checkpointer.putWrites(config, writesToSave, taskId));
if (this.tasks) this._outputWrites(taskId, writesCopy);
if (!writes.length || !this.cache || !this.tasks) return;
const task = this.tasks[taskId];
if (task == null || task.cache_key == null) return;
if (writes[0][0] === "__error__" || writes[0][0] === "__interrupt__") return;
this.cache.set([{
key: [task.cache_key.ns, task.cache_key.key],
value: task.writes,
ttl: task.cache_key.ttl
}]);
}
_outputWrites(taskId, writes, cached = false) {
const task = this.tasks[taskId];
if (task !== void 0) {
if (task.config !== void 0 && (task.config.tags ?? []).includes("langsmith:hidden")) return;
if (writes.length > 0) {
if (writes[0][0] === "__interrupt__") {
if (task.path?.[0] === "__pregel_push" && task.path?.[task.path.length - 1] === true) return;
const interruptWrites = writes.filter((w) => w[0] === INTERRUPT).flatMap((w) => w[1]);
this._emit([["updates", { [INTERRUPT]: interruptWrites }], ["values", { [INTERRUPT]: interruptWrites }]]);
} else if (writes[0][0] !== "__error__") this._emit(gatherIteratorSync(prefixGenerator(mapOutputUpdates(this.outputKeys, [[task, writes]], cached), "updates")));
}
if (!cached) this._emit(gatherIteratorSync(prefixGenerator(mapDebugTaskResults([[task, writes]], this.streamKeys), "tasks")));
}
}
async _matchCachedWrites() {
if (!this.cache) return [];
const matched = [];
const serializeKey = ([ns, key]) => {
return `ns:${ns.join(",")}|key:${key}`;
};
const keys = [];
const keyMap = {};
for (const task of Object.values(this.tasks)) if (task.cache_key != null && !task.writes.length) {
keys.push([task.cache_key.ns, task.cache_key.key]);
keyMap[serializeKey([task.cache_key.ns, task.cache_key.key])] = task;
}
if (keys.length === 0) return [];
const cache = await this.cache.get(keys);
for (const { key, value } of cache) {
const task = keyMap[serializeKey(key)];
if (task != null) {
task.writes.push(...value);
matched.push({
task,
result: value
});
}
}
return matched;
}
/**
* Execute a single iteration of the Pregel loop.
* Returns true if more iterations are needed.
* @param params - The input keys to use for the tick.
* @returns True if more iterations are needed, false otherwise.
*/
async tick(params) {
if (this.store && !this.store.isRunning) await this.store?.start();
const { inputKeys = [] } = params;
if (this.status !== "pending") throw new Error(`Cannot tick when status is no longer "pending". Current status: "${this.status}"`);
if (![INPUT_DONE, INPUT_RESUMING].includes(this.input)) await this._first(inputKeys);
else if (this.toInterrupt.length > 0) {
this.status = "interrupt_before";
throw new GraphInterrupt();
} else if (Object.values(this.tasks).every((task) => task.writes.length > 0)) {
const finishTaskList = Object.values(this.tasks);
const writes = finishTaskList.flatMap((t) => t.writes);
this.updatedChannels = _applyWrites(this.checkpoint, this.channels, finishTaskList, this.checkpointerGetNextVersion, this.triggerToNodes);
for (const [ch, v] of writes) {
const channel = this.channels[ch];
if (channel != null && isDeltaChannel(channel) && _isOverwriteValue(v)) this._deltaChannelsWithOverwrite.add(ch);
}
const valuesOutput = await gatherIterator(prefixGenerator(mapOutputValues(this.outputKeys, writes, this.channels), "values"));
if (this._exitDeltaWrites !== void 0) for (const [tid, ch, v] of this.checkpointPendingWrites) {
const channel = this.channels[ch];
if (channel != null && isDeltaChannel(channel)) this._exitDeltaWrites.push([
this.step,
tid,
ch,
v
]);
}
this.checkpointPendingWrites = [];
await this._putCheckpoint({ source: "loop" });
this._emitValuesWithCheckpointMeta(valuesOutput);
if (shouldInterrupt(this.checkpoint, this.interruptAfter, finishTaskList)) {
this.status = "interrupt_after";
throw new GraphInterrupt();
}
if (this.config.configurable?.["__pregel_resuming"] !== void 0) delete this.config.configurable?.[CONFIG_KEY_RESUMING];
} else return false;
if (this.step > this.stop) {
this.status = "out_of_steps";
return false;
}
this.tasks = _prepareNextTasks(this.checkpoint, this.checkpointPendingWrites, this.nodes, this.channels, this.config, true, {
step: this.step,
checkpointer: this.checkpointer,
isResuming: this.isResuming,
manager: this.manager,
store: this.store,
stream: this.stream,
triggerToNodes: this.triggerToNodes,
updatedChannels: this.updatedChannels
});
let taskList = Object.values(this.tasks);
if (this.checkpointer && (this.stream.modes.has("checkpoints") || this.stream.modes.has("debug"))) this._emit(await gatherIterator(prefixGenerator(mapDebugCheckpoint(this.checkpointConfig, this.channels, this.streamKeys, this.checkpointMetadata, taskList, this.checkpointPendingWrites, this.prevCheckpointConfig, this.outputKeys), "checkpoints")));
if (taskList.length === 0) {
this.status = "done";
return false;
}
if (this.control != null && this.control.drainRequested) {
this.status = "draining";
return false;
}
if (this.skipDoneTasks && this.checkpointPendingWrites.length > 0) {
for (const [tid, k, v] of this.checkpointPendingWrites) {
if (k === "__error__" || k === "__error_source_node__" || k === "__interrupt__" || k === "__resume__") continue;
const task = taskList.find((t) => t.id === tid);
if (task) task.writes.push([k, v]);
}
this._resumeErrorHandlersIfApplicable();
taskList = Object.values(this.tasks);
for (const task of taskList) if (task.writes.length > 0) this._outputWrites(task.id, task.writes, true);
}
if (taskList.every((task) => task.writes.length > 0)) return this.tick({ inputKeys });
if (shouldInterrupt(this.checkpoint, this.interruptBefore, taskList)) {
this.status = "interrupt_before";
throw new GraphInterrupt();
}
if (this.stream.modes.has("tasks") || this.stream.modes.has("debug")) {
const debugOutput = await gatherIterator(prefixGenerator(mapDebugTasks(taskList), "tasks"));
this._emit(debugOutput);
}
return true;
}
async finishAndHandleError(error) {
if (this.durability === "exit" && (!this.isNested || typeof error !== "undefined" || this.checkpointNamespace.every((part) => !part.includes(":")))) {
await this._putExitDeltaWrites();
this._putCheckpoint(this.checkpointMetadata);
this._flushPendingWrites();
}
const suppress = this._suppressInterrupt(error);
if (suppress || error === void 0) this.output = readChannels(this.channels, this.outputKeys);
if (suppress) {
if (this.tasks !== void 0 && this.checkpointPendingWrites.length > 0 && Object.values(this.tasks).some((task) => task.writes.length > 0)) {
this.updatedChannels = _applyWrites(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion, this.triggerToNodes);
this._emitValuesWithCheckpointMeta(gatherIteratorSync(prefixGenerator(mapOutputValues(this.outputKeys, Object.values(this.tasks).flatMap((t) => t.writes), this.channels), "values")));
}
if (isGraphInterrupt(error) && !error.interrupts.length) this._emit([["updates", { [INTERRUPT]: [] }], ["values", { [INTERRUPT]: [] }]], this.#interruptStreamNamespace());
}
return suppress;
}
async acceptPush(task, writeIdx, call) {
if (this.interruptAfter?.length > 0 && shouldInterrupt(this.checkpoint, this.interruptAfter, [task])) {
this.toInterrupt.push(task);
return;
}
const pushed = _prepareSingleTask([
PUSH,
task.path ?? [],
writeIdx,
task.id,
call
], this.checkpoint, this.checkpointPendingWrites, this.nodes, this.channels, task.config ?? {}, true, {
step: this.step,
checkpointer: this.checkpointer,
manager: this.manager,
store: this.store,
stream: this.stream
});
if (!pushed) return;
if (this.interruptBefore?.length > 0 && shouldInterrupt(this.checkpoint, this.interruptBefore, [pushed])) {
this.toInterrupt.push(pushed);
return;
}
if (this.stream.modes.has("tasks") || this.stream.modes.has("debug")) this._emit(gatherIteratorSync(prefixGenerator(mapDebugTasks([pushed]), "tasks")));
if (this.debug) printStepTasks(this.step, [pushed]);
this.tasks[pushed.id] = pushed;
if (this.skipDoneTasks) this._matchWrites({ [pushed.id]: pushed });
const tasks = await this._matchCachedWrites();
for (const { task } of tasks) this._outputWrites(task.id, task.writes, true);
return pushed;
}
/**
* Returns the name of the error handler node registered for `nodeName`, or
* `undefined` if none is configured.
*/
getErrorHandlerNode(nodeName) {
return this.nodes[nodeName]?.errorHandlerNode;
}
/**
* Whether `nodeName` is itself an auto-generated error handler node.
*/
isErrorHandlerNode(nodeName) {
return this.nodes[nodeName]?.isErrorHandler === true;
}
/**
* Schedule a node-level error handler task for a task that failed after its
* retry policy was exhausted. Prepares the handler task (injecting a
* {@link NodeError}), registers it so the runner executes it within the
* current step, and returns it (or `undefined` if no handler applies).
*
* The failure provenance (`ERROR` + `ERROR_SOURCE_NODE`) is checkpointed by
* the runner via {@link PregelLoop#putWrites} so handlers observe the same
* context after a resume.
*/
scheduleErrorHandler(failedTask, error) {
const handlerNode = this.getErrorHandlerNode(String(failedTask.name));
if (!handlerNode) return void 0;
const handlerTask = _prepareNodeErrorHandlerTask(failedTask, handlerNode, error, this.checkpoint, this.checkpointPendingWrites, this.nodes, this.channels, failedTask.config ?? this.config, {
step: this.step,
checkpointer: this.checkpointer,
manager: this.manager,
store: this.store,
stream: this.stream
});
if (handlerTask === void 0) return void 0;
this.tasks[handlerTask.id] = handlerTask;
this._emit(gatherIteratorSync(prefixGenerator(mapDebugTasks([handlerTask]), "tasks")));
if (this.debug) printStepTasks(this.step, [handlerTask]);
return handlerTask;
}
/**
* On resume, re-schedule error handlers for tasks that failed in a prior run
* but had not finished being handled. Scans pending writes for
* `ERROR_SOURCE_NODE` markers (paired with `ERROR`), marks the originating
* task as done (so the runner won't re-run it), and prepares a fresh handler
* task so the runner picks it up.
*/
_resumeErrorHandlersIfApplicable() {
const failed = /* @__PURE__ */ new Map();
for (const [tid, chan] of this.checkpointPendingWrites) {
if (chan !== "__error_source_node__") continue;
const errorWrite = this.checkpointPendingWrites.find(([t, c]) => t === tid && c === "__error__");
if (errorWrite === void 0) continue;
const value = errorWrite[2];
const error = new Error(value?.message ?? String(value));
if (value?.name) error.name = value.name;
failed.set(tid, error);
}
for (const [tid, error] of failed) {
const task = this.tasks[tid];
if (task === void 0) continue;
if (!this.getErrorHandlerNode(String(task.name))) continue;
if (task.writes.length === 0) task.writes.push([ERROR, {
message: error.message,
name: error.name
}]);
this.scheduleErrorHandler(task, error);
}
}
_suppressInterrupt(e) {
return isGraphInterrupt(e) && !this.isNested;
}
async _first(inputKeys) {
const { configurable } = this.config;
const scratchpad = configurable?.[CONFIG_KEY_SCRATCHPAD];
if (scratchpad && scratchpad.nullResume !== void 0) this.putWrites(NULL_TASK_ID, [[RESUME, scratchpad.nullResume]]);
if (isCommand(this.input)) {
const hasResume = this.input.resume != null;
if (this.input.resume != null && typeof this.input.resume === "object" && Object.keys(this.input.resume).every(isXXH3)) {
this.config.configurable ??= {};
this.config.configurable[CONFIG_KEY_RESUME_MAP] = this.input.resume;
}
if (hasResume && this.checkpointer == null) throw new Error("Cannot use Command(resume=...) without checkpointer");
const writes = {};
for (const [tid, key, value] of mapCommand(this.input, this.checkpointPendingWrites)) {
writes[tid] ??= [];
writes[tid].push([key, value]);
}
if (Object.keys(writes).length === 0) throw new EmptyInputError("Received empty Command input");
for (const [tid, ws] of Object.entries(writes)) this.putWrites(tid, ws);
}
const nullWrites = (this.checkpointPendingWrites ?? []).filter((w) => w[0] === NULL_TASK_ID).map((w) => w.slice(1));
if (nullWrites.length > 0) _applyWrites(this.checkpoint, this.channels, [{
name: INPUT,
writes: nullWrites,
triggers: []
}], this.checkpointerGetNextVersion, this.triggerToNodes);
const inputIsCommand = isCommand(this.input);
const isCommandUpdateOrGoto = inputIsCommand && nullWrites.length > 0;
const isTimeTraveling = this.isReplaying && (this.isNested && configurable?.["checkpoint_ns"] !== void 0 && configurable?.["checkpoint_ns"] !== "" && configurable?.["checkpoint_map"] !== void 0 && configurable["checkpoint_ns"] in configurable["checkpoint_map"] || !(inputIsCommand && this.input.resume != null || configurable?.["__pregel_resuming"] === true || this.resumeAtHead));
if (isTimeTraveling) this.checkpointPendingWrites = this.checkpointPendingWrites.filter((w) => w[1] !== RESUME);
const cachedIsResuming = this.isResuming;
if (cachedIsResuming || isCommandUpdateOrGoto) {
const interruptSeen = { ...this.checkpoint.versions_seen[INTERRUPT] };
for (const channelName in this.channels) {
if (!Object.prototype.hasOwnProperty.call(this.channels, channelName)) continue;
if (this.checkpoint.channel_versions[channelName] !== void 0) interruptSeen[channelName] = this.checkpoint.channel_versions[channelName];
}
this.checkpoint.versions_seen[INTERRUPT] = interruptSeen;
if (isTimeTraveling && this.checkpointMetadata.source !== "update" && this.checkpointMetadata.source !== "fork") {
this.checkpointPendingWrites = this.checkpointPendingWrites.filter((w) => w[1] !== INTERRUPT);
await this._putCheckpoint({ source: "fork" });
}
const valuesOutput = await gatherIterator(prefixGenerator(mapOutputValues(this.outputKeys, true, this.channels), "values"));
if (cachedIsResuming) this.input = INPUT_RESUMING;
else if (isCommandUpdateOrGoto) {
await this._putCheckpoint({ source: "input" });
this.input = INPUT_DONE;
}
this._emitValuesWithCheckpointMeta(valuesOutput);
} else {
const inputWrites = await gatherIterator(mapInput(inputKeys, this.input));
if (inputWrites.length > 0) {
const discardTasks = _prepareNextTasks(this.checkpoint, this.checkpointPendingWrites, this.nodes, this.channels, this.config, true, { step: this.step });
this.updatedChannels = _applyWrites(this.checkpoint, this.channels, Object.values(discardTasks).concat([{
name: INPUT,
writes: inputWrites,
triggers: []
}]), this.checkpointerGetNextVersion, this.triggerToNodes);
const deltaInput = inputWrites.filter(([c]) => {
const channel = this.channels[c];
return channel != null && isDeltaChannel(channel);
});
for (const [c, v] of deltaInput) if (_isOverwriteValue(v)) this._deltaChannelsWithOverwrite.add(c);
if (deltaInput.length > 0) {
if (this._exitDeltaWrites !== void 0) for (const [c, v] of deltaInput) this._exitDeltaWrites.push([
this.step,
NULL_TASK_ID,
c,
v
]);
else if (this.checkpointer != null) this.putWrites(NULL_TASK_ID, deltaInput);
}
await this._putCheckpoint({ source: "input" });
this.input = INPUT_DONE;
} else if (!("__pregel_resuming" in (this.config.configurable ?? {}))) throw new EmptyInputError(`Received no input writes for ${JSON.stringify(inputKeys, null, 2)}`);
else this.input = INPUT_DONE;
}
if (!this.isNested) {
let replayState;
if (isTimeTraveling) {
let replayCheckpointId = this.checkpoint.id;
if ((this.checkpointMetadata.source === "update" || this.checkpointMetadata.source === "fork") && this.prevCheckpointConfig) replayCheckpointId = this.prevCheckpointConfig.configurable?.["checkpoint_id"] ?? replayCheckpointId;
replayState = new ReplayState(replayCheckpointId);
}
this.config = patchConfigurable(this.config, {
[CONFIG_KEY_RESUMING]: this.isResuming,
[CONFIG_KEY_REPLAY_STATE]: replayState
});
}
}
#interruptStreamNamespace() {
const ns = this.checkpointNamespace;
if (!(ns.length === 0 || ns.length === 1 && ns[0] === "") || this.config.configurable?.["__pregel_stream"] === void 0) return ns;
const deepest = deepestCheckpointMapNamespace(this.config.configurable?.[CONFIG_KEY_CHECKPOINT_MAP]);
return deepest.length > 0 ? deepest : ns;
}
_emit(values, namespace = this.checkpointNamespace) {
for (const [mode, payload] of values) {
if (this.stream.modes.has(mode)) this.stream.push([
namespace,
mode,
payload
]);
if ((mode === "checkpoints" || mode === "tasks") && this.stream.modes.has("debug")) {
const step = mode === "checkpoints" ? this.step - 1 : this.step;
const timestamp = (/* @__PURE__ */ new Date()).toISOString();
const type = (() => {
if (mode === "checkpoints") return "checkpoint";
else if (typeof payload === "object" && payload != null && "result" in payload) return "task_result";
else return "task";
})();
this.stream.push([
namespace,
"debug",
{
step,
type,
timestamp,
payload
}
]);
}
}
}
/**
* Build a {@link StreamChunkMeta} describing the currently active checkpoint.
* Emitted as a separate ``[namespace, "checkpoints", envelope]`` chunk before
* the paired ``values`` chunk. Returns `undefined` if no checkpoint metadata
* is available yet.
*/
_currentCheckpointMeta() {
if (!this.checkpointMetadata || !this.checkpoint?.id) return void 0;
const parent_id = this.prevCheckpointConfig?.configurable?.checkpoint_id;
return { checkpoint: {
id: this.checkpoint.id,
...parent_id ? { parent_id } : {},
step: this.checkpointMetadata.step,
source: this.checkpointMetadata.source
} };
}
/**
* Emit stream entries. When checkpoint meta is available, push a lightweight
* ``[namespace, "checkpoints", envelope]`` chunk before each ``values`` chunk.
*/
_emitValuesWithCheckpointMeta(entries) {
const meta = this._currentCheckpointMeta();
for (const [mode, payload] of entries) {
if (mode === "values" && meta?.checkpoint != null && !this.stream.modes.has("checkpoints")) this.stream.push([
this.checkpointNamespace,
"checkpoints",
meta.checkpoint
]);
if (this.stream.modes.has(mode)) this.stream.push([
this.checkpointNamespace,
mode,
payload
]);
}
}
_putCheckpoint(inputMetadata) {
const exiting = this.checkpointMetadata === inputMetadata;
const doCheckpoint = this.checkpointer != null && (this.durability !== "exit" || exiting);
const storeCheckpoint = (checkpoint) => {
this.prevCheckpointConfig = this.checkpointConfig?.configurable?.checkpoint_id ? this.checkpointConfig : void 0;
this.checkpointConfig = patchConfigurable(this.checkpointConfig, { [CONFIG_KEY_CHECKPOINT_NS]: this.config.configurable?.checkpoint_ns ?? "" });
const channelVersions = { ...this.checkpoint.channel_versions };
const newVersions = getNewChannelVersions(this.checkpointPreviousVersions, channelVersions);
this.checkpointPreviousVersions = channelVersions;
this._checkpointerPutAfterPrevious({
config: { ...this.checkpointConfig },
checkpoint: copyCheckpoint(checkpoint),
metadata: { ...this.checkpointMetadata },
newVersions
});
this.checkpointConfig = {
...this.checkpointConfig,
configurable: {
...this.checkpointConfig.configurable,
checkpoint_id: this.checkpoint.id
}
};
};
let newCounters;
if (!exiting) {
const prevCounters = this.checkpointMetadata.counters_since_delta_snapshot ?? {};
newCounters = {};
const updated = this.updatedChannels ?? /* @__PURE__ */ new Set();
for (const chName in this.channels) {
if (!Object.prototype.hasOwnProperty.call(this.channels, chName)) continue;
if (!isDeltaChannel(this.channels[chName])) continue;
const [u, s] = prevCounters[chName] ?? [0, 0];
newCounters[chName] = [updated.has(chName) ? u + 1 : u, s + 1];
}
this.checkpointMetadata = {
...inputMetadata,
step: this.step,
parents: this.config.configurable?.["checkpoint_map"] ?? {}
};
} else newCounters = { ...this.checkpointMetadata.counters_since_delta_snapshot ?? {} };
const channelsToSnapshot = doCheckpoint ? deltaChannelsToSnapshot(this.channels, newCounters) : /* @__PURE__ */ new Set();
if (doCheckpoint) for (const ch of this._deltaChannelsWithOverwrite) channelsToSnapshot.add(ch);
this.checkpoint = createCheckpoint(this.checkpoint, doCheckpoint ? this.channels : void 0, this.step, {
id: exiting ? this.checkpoint.id : void 0,
channelsToSnapshot,
updatedChannels: this.updatedChannels,
getNextVersion: doCheckpoint ? (current) => this.checkpointerGetNextVersion(current) : void 0
});
for (const k of channelsToSnapshot) {
newCounters[k] = [0, 0];
this._deltaChannelsWithOverwrite.delete(k);
}
const nonZero = {};
for (const k in newCounters) {
if (!Object.prototype.hasOwnProperty.call(newCounters, k)) continue;
const [u, s] = newCounters[k];
if (u !== 0 || s !== 0) nonZero[k] = [u, s];
}
if (Object.keys(nonZero).length > 0) this.checkpointMetadata.counters_since_delta_snapshot = nonZero;
else delete this.checkpointMetadata.counters_since_delta_snapshot;
if (doCheckpoint) storeCheckpoint(this.checkpoint);
if (!exiting) this.step += 1;
}
/**
* Stage the exit-mode accumulator of DeltaChannel writes so the final
* checkpoint can be reconstructed. In "exit" durability per-step writes are
* not persisted, so delta writes are accumulated across the run and anchored
* here — under the saved parent, or a freshly-created stub when this is a
* first run with no persisted parent. Channels that will snapshot in the
* final checkpoint are excluded (their full value lives in `channel_values`).
*
* Must run BEFORE the final `_putCheckpoint` so the stub branch can adjust
* `checkpointConfig` to anchor the final checkpoint on the stub.
*/
async _putExitDeltaWrites() {
if (this._exitDeltaWrites === void 0 || this._exitDeltaWrites.length === 0 || this.checkpointer == null || this._initialCheckpointConfig === void 0) return;
const counters = this.checkpointMetadata.counters_since_delta_snapshot ?? {};
const channelsToSnapshot = deltaChannelsToSnapshot(this.channels, counters);
for (const ch of this._deltaChannelsWithOverwrite) channelsToSnapshot.add(ch);
const pending = this._exitDeltaWrites.filter(([, , ch]) => !channelsToSnapshot.has(ch));
if (pending.length === 0) return;
let anchorConfig;
if (this._hasPersistedParent) anchorConfig = this._initialCheckpointConfig;
else {
const stubCp = emptyCheckpoint();
stubCp.id = this.checkpointIdSaved ?? stubCp.id;
stubCp.ts = (/* @__PURE__ */ new Date()).toISOString();
const stubPutConfig = patchConfigurable(this._initialCheckpointConfig, { [CONFIG_KEY_CHECKPOINT_ID]: void 0 });
anchorConfig = patchConfigurable(this._initialCheckpointConfig, { [CONFIG_KEY_CHECKPOINT_ID]: stubCp.id });
this._trackCheckpointerPromise(this.checkpointer.put(stubPutConfig, stubCp, {
source: "loop",
step: -2,
parents: {}
}, {}));
this.checkpointConfig = anchorConfig;
}
const anchorWriteConfig = patchConfigurable(anchorConfig, {
[CONFIG_KEY_CHECKPOINT_NS]: this.config.configurable?.checkpoint_ns ?? "",
[CONFIG_KEY_CHECKPOINT_ID]: anchorConfig.configurable?.[CONFIG_KEY_CHECKPOINT_ID]
});
const grouped = /* @__PURE__ */ new Map();
const order = [];
for (const [step, tid, ch, v] of pending) {
const key = `${step}\u0000${tid}`;
let group = grouped.get(key);
if (group === void 0) {
group = [];
grouped.set(key, group);
order.push({
key,
step,
tid
});
}
group.push([ch, v]);
}
for (const { key, step, tid } of order) {
const synthTid = exitDeltaTaskId(step, tid);
this._trackCheckpointerPromise(this.checkpointer.putWrites(anchorWriteConfig, grouped.get(key), synthTid));
}
}
_flushPendingWrites() {
if (this.checkpointer == null) return;
if (this.checkpointPendingWrites.length === 0) return;
const config = patchConfigurable(this.checkpointConfig, {
[CONFIG_KEY_CHECKPOINT_NS]: this.config.configurable?.checkpoint_ns ?? "",
[CONFIG_KEY_CHECKPOINT_ID]: this.checkpoint.id
});
const byTask = {};
for (const [tid, key, value] of this.checkpointPendingWrites) {
byTask[tid] ??= [];
byTask[tid].push([key, value]);
}
for (const [tid, ws] of Object.entries(byTask)) this._trackCheckpointerPromise(this.checkpointer.putWrites(config, ws, tid));
}
_matchWrites(tasks) {
for (const [tid, k, v] of this.checkpointPendingWrites) {
if (k === "__error__" || k === "__interrupt__" || k === "__resume__") continue;
const task = Object.values(tasks).find((t) => t.id === tid);
if (task) task.writes.push([k, v]);
}
for (const task of Object.values(tasks)) if (task.writes.length > 0) this._outputWrites(task.id, task.writes, true);
}
};
//#endregion
export { PregelLoop };
//# sourceMappingURL=loop.js.map