@langchain/langgraph
Version:
216 lines (215 loc) • 6.89 kB
JavaScript
const require_constants = require("../constants.cjs");
const require_errors = require("../errors.cjs");
const require_index = require("./utils/index.cjs");
let _langchain_core_callbacks_base = require("@langchain/core/callbacks/base");
//#region src/pregel/timeout.ts
/**
* Tracks the live progress state of a single timed node attempt.
*
* The scope guards observable-progress channels (writes, child-task
* scheduling, the custom stream writer, callbacks) so that:
*
* 1. `idleTimeout` is refreshed whenever the node makes progress, and
* 2. once the attempt is `close()`d after a timeout fires, late writes/calls
* from the still-running background task are dropped (Python parity:
* buffered writes from the failed attempt must not leak into the
* checkpoint).
*
* @internal
*/
var TimedAttemptScope = class {
active = true;
lastProgress = Date.now();
refreshOn;
constructor(refreshOn) {
this.refreshOn = refreshOn;
}
/** Record progress now. Always honored (used by `runtime.heartbeat()`). */
touch() {
this.lastProgress = Date.now();
}
/**
* Record progress for an automatic signal (write/call/stream/callback).
* No-op when `refreshOn === "heartbeat"`, where only explicit heartbeats
* count as progress.
*/
autoTouch() {
if (this.refreshOn === "auto") this.lastProgress = Date.now();
}
close() {
this.active = false;
}
};
/**
* Callback handler that refreshes a {@link TimedAttemptScope} on any LangChain
* callback event emitted under the node's run. Because it is attached via
* `config.callbacks`, it only observes events from runs descended from this
* node's attempt, not from sibling nodes.
*
* @internal
*/
var IdleProgressCallbackHandler = class extends _langchain_core_callbacks_base.BaseCallbackHandler {
name = "IdleProgressCallbackHandler";
awaitHandlers = false;
#scope;
constructor(scope) {
super();
this.#scope = scope;
}
#touch = () => {
this.#scope.autoTouch();
};
handleLLMStart = this.#touch;
handleChatModelStart = this.#touch;
handleLLMNewToken = this.#touch;
handleLLMEnd = this.#touch;
handleLLMError = this.#touch;
handleChainStart = this.#touch;
handleChainEnd = this.#touch;
handleChainError = this.#touch;
handleToolStart = this.#touch;
handleToolEnd = this.#touch;
handleToolError = this.#touch;
handleText = this.#touch;
handleRetrieverStart = this.#touch;
handleRetrieverEnd = this.#touch;
handleRetrieverError = this.#touch;
handleCustomEvent = this.#touch;
};
/**
* Wrap the node attempt config so observable-progress signals refresh the idle
* clock and are dropped once the scope is closed. Also injects
* {@link LangGraphRunnableConfig.heartbeat}.
*/
function wrapConfig(config, scope, policy, taskName) {
const configurable = config.configurable ?? {};
const patch = {};
const send = configurable[require_constants.CONFIG_KEY_SEND];
if (typeof send === "function") patch[require_constants.CONFIG_KEY_SEND] = (writes) => {
if (!scope.active) return void 0;
if (writes && writes.length) scope.autoTouch();
return send(writes);
};
const callFn = configurable[require_constants.CONFIG_KEY_CALL];
if (typeof callFn === "function") patch[require_constants.CONFIG_KEY_CALL] = (...args) => {
if (!scope.active) throw new Error(`Node "${taskName}" attempt was cancelled after its timeout fired`);
scope.autoTouch();
return callFn(...args);
};
const wrapped = { ...Object.keys(patch).length > 0 ? require_index.patchConfigurable(config, patch) : config };
wrapped.heartbeat = () => {
if (policy.idleTimeout !== void 0) scope.touch();
};
if (typeof wrapped.writer === "function") {
const writer = wrapped.writer;
wrapped.writer = ((chunk) => {
if (!scope.active) return void 0;
scope.autoTouch();
return writer(chunk);
});
}
if ((policy.refreshOn ?? "auto") === "auto" && policy.idleTimeout !== void 0) {
const handler = new IdleProgressCallbackHandler(scope);
const cb = wrapped.callbacks;
if (cb === void 0) wrapped.callbacks = [handler];
else if (Array.isArray(cb)) wrapped.callbacks = [...cb, handler];
else {
const copied = cb.copy();
copied.addHandler(handler, true);
wrapped.callbacks = copied;
}
}
return wrapped;
}
/**
* Run a single node attempt under a {@link TimeoutPolicy}.
*
* Races the node invocation against per-attempt run/idle watchdogs. On
* successful completion (or node error), returns/rethrows normally. When a
* watchdog fires first, the scope is closed, the task's buffered writes are
* dropped, the attempt's {@link AbortSignal} is aborted, and a
* {@link NodeTimeoutError} is thrown.
*
* @internal
*/
async function runAttemptWithTimeout(task, config, policy, invoke) {
const scope = new TimedAttemptScope(policy.refreshOn ?? "auto");
const timeoutController = new AbortController();
const { signal: composedSignal, dispose } = require_index.combineAbortSignals(config.signal, timeoutController.signal);
const scopedConfig = wrapConfig({
...config,
signal: composedSignal
}, scope, policy, String(task.name));
const start = Date.now();
const nodeOutcome = invoke(scopedConfig).then((value) => ({
type: "ok",
value
}), (error) => ({
type: "err",
error
}));
let runTimer;
let idleTimer;
const clearTimers = () => {
if (runTimer !== void 0) clearTimeout(runTimer);
if (idleTimer !== void 0) clearTimeout(idleTimer);
};
const watchdog = new Promise((resolve) => {
if (policy.runTimeout !== void 0) runTimer = setTimeout(() => resolve({
type: "timeout",
kind: "run"
}), policy.runTimeout);
if (policy.idleTimeout !== void 0) {
const idleMs = policy.idleTimeout;
const checkIdle = () => {
const remaining = scope.lastProgress + idleMs - Date.now();
if (remaining <= 0) resolve({
type: "timeout",
kind: "idle"
});
else idleTimer = setTimeout(checkIdle, remaining);
};
idleTimer = setTimeout(checkIdle, idleMs);
}
});
let outcome;
try {
outcome = await Promise.race([nodeOutcome, watchdog]);
} finally {
clearTimers();
}
if (outcome.type !== "timeout") {
const now = Date.now();
if (policy.runTimeout !== void 0 && now - start >= policy.runTimeout) outcome = {
type: "timeout",
kind: "run"
};
else if (policy.idleTimeout !== void 0 && now - scope.lastProgress >= policy.idleTimeout) outcome = {
type: "timeout",
kind: "idle"
};
}
if (outcome.type === "ok") {
dispose?.();
return outcome.value;
}
if (outcome.type === "err") {
dispose?.();
throw outcome.error;
}
const elapsed = Date.now() - start;
scope.close();
task.writes.splice(0, task.writes.length);
timeoutController.abort();
dispose?.();
throw new require_errors.NodeTimeoutError({
node: String(task.name),
elapsed,
kind: outcome.kind,
runTimeout: policy.runTimeout,
idleTimeout: policy.idleTimeout
});
}
//#endregion
exports.runAttemptWithTimeout = runAttemptWithTimeout;
//# sourceMappingURL=timeout.cjs.map