nexus-ai-pro 2.0.0 → 2.2.0

package/dist-cjs/graph/graph.js

@@ -37,8 +37,12 @@ exports.CompiledGraph = exports.StateGraph = void 0;
 exports.createGraph = createGraph;
 const node_crypto_1 = require("node:crypto");
 const graph_js_1 = require("../types/graph.js");
+const tenant_js_1 = require("../store/tenant.js");
+const signals_js_1 = require("../utils/signals.js");
 const checkpointer_js_1 = require("./checkpointer.js");
 const checkpoint_migration_js_1 = require("./checkpoint-migration.js");
+const durability_js_1 = require("./durability.js");
+const event_stream_js_1 = require("./event-stream.js");
 const errors_js_1 = require("./errors.js");
 const DEFAULT_MAX_STEPS = 25;
 const DEFAULT_MAX_CONCURRENCY = 16;
@@ -50,6 +54,7 @@ const DEFAULT_RETRY = {
     jitter: true,
 };
 const SUBGRAPH = Symbol('nexus.graph.subgraph');
+const FORWARD_EVENT = Symbol('nexus.graph.forwardEvent');
 const FOLLOWING_UPDATES = Symbol('nexus.graph.followingUpdates');
 class StateGraph {
     channels;
@@ -161,6 +166,7 @@ class CompiledGraph {
     now;
     store;
     nodeCache;
+    writers = new Set();
     constructor(channels, nodes, edges, options, scope = {}) {
         this.channels = channels;
         this.nodes = nodes;
@@ -170,6 +176,9 @@ class CompiledGraph {
         this.now = options.now ?? (() => new Date());
         this.store = options.checkpointer === false ? undefined : (options.checkpointer ?? new checkpointer_js_1.MemoryGraphCheckpointer());
     }
+    async flush() {
+        await Promise.all([...this.writers].map((writer) => writer.flush()));
+    }
     async invoke(input = {}, runOptions = {}) {
         const threadId = resolveThreadId(runOptions.threadId);
         let last;
@@ -181,11 +190,42 @@ class CompiledGraph {
         const threadId = resolveThreadId(runOptions.threadId);
         return this.run(threadId, runOptions, async () => {
             this.assertInput(input);
+            if (runOptions.tenantId !== undefined) {
+                const existing = await this.checkpointer()?.get(threadId);
+                if (existing && !ownedBy(existing, runOptions.tenantId)) {
+                    throw new errors_js_1.GraphValidationError(`Thread "${threadId}" cannot be used by this tenant; choose another id`);
+                }
+            }
             const state = this.seedState(input);
             const next = this.entryNodes();
             return { threadId, step: 0, state, next, status: 'running', createdAt: this.now().toISOString() };
         });
     }
+    events(input = {}, options = {}) {
+        const threadId = resolveThreadId(options.threadId);
+        return this.eventStream(threadId, options, (runOptions) => this.stream(input, runOptions));
+    }
+    resumeEvents(threadId, value, options = {}) {
+        return this.eventStream(threadId, options, (runOptions) => this.resume(threadId, value, runOptions));
+    }
+    continueEvents(threadId, options = {}) {
+        return this.eventStream(threadId, options, (runOptions) => this.continue(threadId, runOptions));
+    }
+    eventStream(threadId, options, run) {
+        return new event_stream_js_1.GraphEventStream((onEvent, signal) => {
+            const linked = (0, signals_js_1.linkSignals)([options.signal, signal]);
+            return releasing(run({
+                ...options,
+                threadId,
+                signal: linked.signal,
+                ...(options.subgraphs ? { subgraphEvents: true } : {}),
+                onEvent: (event) => {
+                    options.onEvent?.(event);
+                    onEvent(event);
+                },
+            }), linked.dispose);
+        }, (last) => this.toResult(threadId, last), (0, event_stream_js_1.streamSettings)(options));
+    }
     resume(threadId, value, runOptions = {}) {
         return this.run(threadId, runOptions, async () => {
             const checkpoint = await this.requireCheckpoint(threadId);
@@ -235,7 +275,7 @@ class CompiledGraph {
     continue(threadId, runOptions = {}) {
         return this.run(threadId, runOptions, async () => {
             const checkpoint = await this.requireCheckpoint(threadId);
-            return { ...checkpoint, status: 'running', error: undefined };
+            return { ...checkpoint, status: 'running', error: undefined, drained: undefined };
         });
     }
     resumeFrom(threadId, step, runOptions = {}) {
@@ -246,19 +286,25 @@ class CompiledGraph {
             return { ...(0, checkpoint_migration_js_1.migrateCheckpoint)(stored), status: 'running' };
         });
     }
-    async state(threadId, step) {
+    async state(threadId, step, options = {}) {
         const stored = await this.checkpointer()?.get(threadId, step);
-        return stored ? (0, checkpoint_migration_js_1.migrateCheckpoint)(stored) : undefined;
+        if (!stored || !ownedBy(stored, options.tenantId))
+            return undefined;
+        return (0, checkpoint_migration_js_1.migrateCheckpoint)(stored);
     }
-    async history(threadId, limit) {
+    async history(threadId, limit, options = {}) {
         const stored = (await this.checkpointer()?.history(threadId, limit)) ?? [];
+        if (stored[0] && !ownedBy(stored[0], options.tenantId))
+            return [];
         return stored.map((item) => (0, checkpoint_migration_js_1.migrateCheckpoint)(item));
     }
     asNode() {
         const node = async (context) => {
+            const forward = context[FORWARD_EVENT];
             const runOptions = {
                 threadId: `${context.threadId}:${context.taskId}`,
                 signal: context.signal,
+                ...(forward ? { onEvent: forward, subgraphEvents: true } : {}),
             };
             const paused = await this.state(runOptions.threadId);
             let result;
@@ -301,13 +347,18 @@ class CompiledGraph {
     describe() {
         const nodes = [...this.nodes.entries()].map(([id, definition]) => {
             const subgraph = definition.fn[SUBGRAPH];
-            const { ends, defer, retry, timeoutMs, cache } = definition.options;
+            const { ends, defer, timeoutMs, cache, idempotent } = definition.options;
+            const maxAttempts = this.retryPolicy(definition).maxAttempts;
+            const timeout = this.timeoutOf(definition);
             return {
                 id,
                 ...(ends?.length ? { ends: [...ends] } : {}),
                 ...(defer ? { defer } : {}),
-                ...((retry?.maxAttempts ?? this.options.retry?.maxAttempts ?? 1) > 1 ? { retry: true } : {}),
+                ...(maxAttempts > 1 ? { retry: true, maxAttempts } : {}),
                 ...(timeoutMs !== undefined ? { timeoutMs } : {}),
+                ...(timeout.runMs !== undefined || timeout.idleMs !== undefined ? { timeout } : {}),
+                ...(this.errorHandler(definition) ? { onError: true } : {}),
+                ...(idempotent ? { idempotent: true } : {}),
                 ...(cache ? { cache: true } : {}),
                 ...(subgraph ? { subgraph: subgraph.describe() } : {}),
             };
@@ -342,10 +393,13 @@ class CompiledGraph {
             dynamic: [...dynamic],
             ...(this.scope.input ? { input: [...this.scope.input] } : {}),
             ...(this.scope.output ? { output: [...this.scope.output] } : {}),
+            durability: this.options.durability ?? 'sync',
+            checkpointer: this.options.checkpointer === false ? 'none' : this.options.checkpointer === undefined ? 'memory' : 'custom',
+            maxConcurrency: this.options.maxConcurrency ?? DEFAULT_MAX_CONCURRENCY,
         };
     }
     async updateState(threadId, update, options = {}) {
-        const checkpoint = await this.requireCheckpoint(threadId);
+        const checkpoint = await this.requireCheckpoint(threadId, options.tenantId);
         const state = this.reduce(checkpoint.state, [update]);
         const createdAt = this.now().toISOString();
         if (!options.asNode) {
@@ -380,8 +434,9 @@ class CompiledGraph {
         if (!store)
             throw new errors_js_1.GraphValidationError('fork() needs a checkpointer to copy history from');
         const history = (await store.history(threadId, Number.MAX_SAFE_INTEGER)).map((item) => (0, checkpoint_migration_js_1.migrateCheckpoint)(item));
-        if (history.length === 0)
+        if (history.length === 0 || !ownedBy(history[0], options.tenantId)) {
             throw new errors_js_1.GraphThreadNotFoundError(threadId);
+        }
         const cutoff = options.step ?? history[0].step;
         const lineage = history.filter((item) => item.step <= cutoff).reverse();
         if (lineage.at(-1)?.step !== cutoff) {
@@ -432,22 +487,69 @@ class CompiledGraph {
         }
     }
     async *runSteps(threadId, runOptions, seed) {
+        const mode = runOptions.durability ?? this.options.durability ?? 'sync';
+        const writer = new durability_js_1.CheckpointWriter(async (written) => {
+            await this.save(written);
+            notify(runOptions, { type: 'checkpoint', step: written.step, status: written.status, next: written.next });
+        }, mode, this.options.maxPendingWrites);
+        this.writers.add(writer);
+        let failed = false;
+        try {
+            yield* this.steps(threadId, runOptions, seed, (written, durable) => writer.persist(written, durable));
+        }
+        catch (error) {
+            failed = true;
+            throw error;
+        }
+        finally {
+            this.writers.delete(writer);
+            if (failed)
+                await writer.flush().catch(() => undefined);
+            else
+                await writer.flush();
+        }
+    }
+    async *steps(threadId, runOptions, seed, persist) {
         const maxSteps = runOptions.maxSteps ?? this.options.maxSteps ?? DEFAULT_MAX_STEPS;
         const stopBefore = runOptions.interruptBefore ?? this.options.interruptBefore ?? [];
         const stopAfter = runOptions.interruptAfter ?? this.options.interruptAfter ?? [];
-        const persist = async (written) => {
-            await this.save(written);
-            notify(runOptions, { type: 'checkpoint', step: written.step, status: written.status, next: written.next });
-        };
         let checkpoint = await seed();
+        const owner = tenantOf(checkpoint);
+        if (runOptions.tenantId !== undefined && owner !== undefined && owner !== runOptions.tenantId) {
+            throw new errors_js_1.GraphThreadNotFoundError(threadId);
+        }
+        const tenantId = runOptions.tenantId ?? owner;
+        if (tenantId !== undefined) {
+            runOptions = { ...runOptions, tenantId, metadata: { ...runOptions.metadata, tenantId } };
+            if (owner === undefined)
+                checkpoint = { ...checkpoint, metadata: { ...checkpoint.metadata, tenantId } };
+        }
         let passBreakpoint = checkpoint.breakpoint?.when === 'before';
         if (checkpoint.breakpoint)
             checkpoint = { ...checkpoint, breakpoint: undefined };
+        if (checkpoint.next.length === 0) {
+            if (checkpoint.status === 'running')
+                checkpoint = { ...checkpoint, status: 'completed' };
+            await persist(checkpoint, true);
+            yield this.toEvent(checkpoint, []);
+            return;
+        }
         await persist(checkpoint);
         while (checkpoint.next.length > 0) {
+            if (runOptions.control?.draining) {
+                const at = this.now().toISOString();
+                checkpoint = {
+                    ...checkpoint,
+                    status: 'interrupted',
+                    drained: { ...(runOptions.control.reason ? { reason: runOptions.control.reason } : {}), at },
+                    createdAt: at,
+                };
+                await persist(checkpoint, true);
+                throw new errors_js_1.GraphDrainedError(threadId, checkpoint.step, runOptions.control.reason);
+            }
             if (runOptions.signal?.aborted) {
                 checkpoint = { ...checkpoint, status: 'interrupted', createdAt: this.now().toISOString() };
-                await persist(checkpoint);
+                await persist(checkpoint, true);
                 yield this.toEvent(checkpoint, []);
                 return;
             }
@@ -472,7 +574,7 @@ class CompiledGraph {
                     breakpoint: { when: 'before', nodes: breakBefore },
                     createdAt: this.now().toISOString(),
                 };
-                await persist(checkpoint);
+                await persist(checkpoint, true);
                 yield this.toEvent(checkpoint, []);
                 return;
             }
@@ -492,6 +594,7 @@ class CompiledGraph {
             const gotos = { ...(checkpoint.gotos ?? {}) };
             let parentCommand;
             let failure;
+            const recovered = [];
             for (const [index, task] of pendingTasks.entries()) {
                 const outcome = outcomes[index];
                 if (!outcome)
@@ -503,6 +606,8 @@ class CompiledGraph {
                     if (outcome.goto?.length)
                         gotos[task.id] = outcome.goto;
                     parentCommand ??= outcome.parent;
+                    if (outcome.recovered)
+                        recovered.push(outcome.recovered);
                     finished.push(task.id);
                 }
                 else if (outcome.kind === 'interrupt') {
@@ -518,6 +623,7 @@ class CompiledGraph {
             const ranNodes = distinctNodes(allTasks.filter((task) => completed.includes(task.id)));
             const ranTasks = pendingTasks;
             const routeFields = Object.keys(gotos).length > 0 ? { gotos } : {};
+            const recoveryFields = recovered.length > 0 ? { recovered } : {};
             if (failure) {
                 const failed = {
                     ...checkpoint,
@@ -525,6 +631,7 @@ class CompiledGraph {
                     ...pausedFields(allTasks, unfinished),
                     completed,
                     ...routeFields,
+                    ...recoveryFields,
                     status: 'failed',
                     error: {
                         name: failure.error instanceof Error ? failure.error.name : 'Error',
@@ -532,7 +639,7 @@ class CompiledGraph {
                     },
                     createdAt: this.now().toISOString(),
                 };
-                await persist(failed);
+                await persist(failed, true);
                 throw failure.error instanceof errors_js_1.GraphNodeError
                     ? failure.error
                     : new errors_js_1.GraphNodeError(failure.task.node, step, failure.error);
@@ -545,11 +652,12 @@ class CompiledGraph {
                     ...pausedFields(allTasks, unfinished),
                     completed,
                     ...routeFields,
+                    ...recoveryFields,
                     status: asking ? 'awaiting_input' : 'interrupted',
                     interrupts,
                     createdAt: this.now().toISOString(),
                 };
-                await persist(checkpoint);
+                await persist(checkpoint, true);
                 yield this.toEvent(checkpoint, ranNodes, ranTasks, attempts);
                 return;
             }
@@ -561,10 +669,11 @@ class CompiledGraph {
                     next: [],
                     status: 'completed',
                     resolved: checkpoint.resolved,
+                    ...recoveryFields,
                     createdAt: this.now().toISOString(),
                     ...(runOptions.metadata ? { metadata: runOptions.metadata } : {}),
                 };
-                await persist(checkpoint);
+                await persist(checkpoint, true);
                 yield { ...this.toEvent(checkpoint, ranNodes, ranTasks, attempts), parentCommand };
                 return;
             }
@@ -583,10 +692,11 @@ class CompiledGraph {
                 status: breakAfter.length > 0 ? 'interrupted' : next.length > 0 ? 'running' : 'completed',
                 ...(breakAfter.length > 0 ? { breakpoint: { when: 'after', nodes: breakAfter } } : {}),
                 resolved: checkpoint.resolved,
+                ...recoveryFields,
                 createdAt: this.now().toISOString(),
                 ...(runOptions.metadata ? { metadata: runOptions.metadata } : {}),
             };
-            await persist(checkpoint);
+            await persist(checkpoint, recovered.length > 0 || next.length === 0 || breakAfter.length > 0);
             yield this.toEvent(checkpoint, ranNodes, ranTasks, attempts);
             if (breakAfter.length > 0)
                 return;
@@ -596,8 +706,8 @@ class CompiledGraph {
         const definition = this.nodes.get(task.node);
         if (!definition)
             throw new errors_js_1.GraphValidationError(`Node "${task.node}" disappeared before it could run`);
-        const policy = { ...DEFAULT_RETRY, ...this.options.retry, ...definition.options.retry };
-        const timeoutMs = definition.options.timeoutMs;
+        const policy = this.retryPolicy(definition);
+        const limits = { ...this.timeoutOf(definition), ...task.timeout };
         const cachePolicy = definition.options.cache;
         const cache = cachePolicy ? await this.loadNodeCache() : undefined;
         const cacheKey = cachePolicy && cache ? cache.key(cachePolicy, task.node, checkpoint.state, task.input) : undefined;
@@ -619,22 +729,18 @@ class CompiledGraph {
             }
         }
         let attempts = 0;
+        const firstAttemptAt = this.now().toISOString();
         while (true) {
             attempts += 1;
             notify(runOptions, { type: 'task_start', step, taskId: task.id, node: task.node, attempt: attempts });
-            const timeout = timeoutMs === undefined ? undefined : AbortSignal.timeout(timeoutMs);
-            const sources = [stepSignal, ...(runOptions.signal ? [runOptions.signal] : []), ...(timeout ? [timeout] : [])];
-            const signal = sources.length === 1 ? sources[0] : AbortSignal.any(sources);
+            const timer = attemptTimer(task.node, limits);
+            const linked = (0, signals_js_1.linkSignals)([stepSignal, runOptions.signal, timer?.signal]);
+            const signal = linked.signal;
             try {
-                const running = this.runNode(definition.fn, task, step, threadId, checkpoint, runOptions, signal, attempts);
-                if (timeout)
+                const running = this.runNode(definition.fn, task, step, threadId, checkpoint, runOptions, signal, attempts, timer?.touch);
+                if (timer)
                     running.catch(() => undefined);
-                const output = timeout
-                    ? await Promise.race([
-                        running,
-                        rejectWhenAborted(timeout, () => new errors_js_1.GraphNodeTimeoutError(task.node, timeoutMs)),
-                    ])
-                    : await running;
+                const output = timer ? await Promise.race([running, timer.expired]) : await running;
                 const parent = output.command?.graph === graph_js_1.Command.PARENT ? output.command : undefined;
                 const goto = parent ? undefined : toRouteRecords(output.command?.goto);
                 const updates = parent ? [] : output.updates;
@@ -652,6 +758,7 @@ class CompiledGraph {
                 return { kind: 'ok', updates, goto, parent, attempts };
             }
             catch (error) {
+                timer?.stop();
                 if (error instanceof errors_js_1.GraphInterrupt) {
                     return {
                         kind: 'interrupt',
@@ -669,10 +776,24 @@ class CompiledGraph {
                 }
                 if (runOptions.signal?.aborted || stepSignal.aborted)
                     return { kind: 'aborted', attempts };
-                const failure = timeout?.aborted ? new errors_js_1.GraphNodeTimeoutError(task.node, timeoutMs) : error;
+                const failure = timer?.error() ?? error;
                 const retryable = policy.retryOn ? policy.retryOn(failure, attempts) : isRetryable(failure);
                 if (attempts >= Math.max(1, policy.maxAttempts) || !retryable) {
-                    return { kind: 'failed', error: failure, attempts };
+                    return this.recover(definition, task, step, checkpoint, runOptions, {
+                        node: task.node,
+                        taskId: task.id,
+                        step,
+                        attempts,
+                        firstAttemptAt,
+                        failedAt: this.now().toISOString(),
+                        error: failure,
+                        retryExhausted: retryable,
+                        ...(failure instanceof errors_js_1.GraphNodeTimeoutError
+                            ? { timeout: { type: failure.kind, limitMs: failure.timeoutMs } }
+                            : {}),
+                        ...(task.input === undefined ? {} : { input: task.input }),
+                        threadId,
+                    });
                 }
                 notify(runOptions, {
                     type: 'task_retry',
@@ -684,9 +805,81 @@ class CompiledGraph {
                 });
                 await delay(backoffDelay(policy, attempts), runOptions.signal);
             }
+            finally {
+                timer?.stop();
+                linked.dispose();
+            }
         }
     }
-    async runNode(fn, task, step, threadId, checkpoint, runOptions, signal, attempt) {
+    async recover(definition, task, step, checkpoint, runOptions, failure) {
+        const handler = this.errorHandler(definition);
+        const event = {
+            type: 'task_failed',
+            step,
+            taskId: task.id,
+            node: task.node,
+            attempts: failure.attempts,
+            error: describe(failure.error),
+        };
+        if (!handler || failure.error instanceof errors_js_1.GraphValidationError) {
+            notify(runOptions, { ...event, recovered: false });
+            return { kind: 'failed', error: failure.error, attempts: failure.attempts };
+        }
+        let decided;
+        try {
+            decided = await handler(failure, { state: Object.freeze({ ...checkpoint.state }) });
+        }
+        catch (error) {
+            notify(runOptions, { ...event, recovered: false });
+            return { kind: 'failed', error, attempts: failure.attempts };
+        }
+        const command = decided instanceof graph_js_1.Command ? decided : undefined;
+        const parent = command?.graph === graph_js_1.Command.PARENT ? command : undefined;
+        const goto = parent ? undefined : toRouteRecords(command?.goto);
+        const own = command ? command.update : decided;
+        const updates = parent || !own ? [] : [own];
+        notify(runOptions, { ...event, recovered: true });
+        return {
+            kind: 'ok',
+            updates,
+            goto,
+            parent,
+            attempts: failure.attempts,
+            recovered: {
+                node: failure.node,
+                taskId: failure.taskId,
+                step,
+                attempts: failure.attempts,
+                error: {
+                    name: failure.error instanceof Error ? failure.error.name : 'Error',
+                    message: describe(failure.error),
+                },
+                failedAt: failure.failedAt,
+                retryExhausted: failure.retryExhausted,
+                ...(failure.timeout ? { timeout: failure.timeout } : {}),
+                ...(goto?.length ? { goto: goto.map(routeName) } : {}),
+            },
+        };
+    }
+    retryPolicy(definition) {
+        return {
+            ...DEFAULT_RETRY,
+            ...this.options.retry,
+            ...this.options.nodeDefaults?.retry,
+            ...definition.options.retry,
+        };
+    }
+    timeoutOf(definition) {
+        const defaults = this.options.nodeDefaults?.timeout;
+        const own = definition.options;
+        const runMs = own.timeout?.runMs ?? own.timeoutMs ?? defaults?.runMs;
+        const idleMs = own.timeout?.idleMs ?? defaults?.idleMs;
+        return { ...(runMs === undefined ? {} : { runMs }), ...(idleMs === undefined ? {} : { idleMs }) };
+    }
+    errorHandler(definition) {
+        return definition.options.onError ?? this.options.nodeDefaults?.onError;
+    }
+    async runNode(fn, task, step, threadId, checkpoint, runOptions, signal, attempt, touch = noop) {
         const node = task.node;
         let interruptIndex = 0;
         const context = {
@@ -694,11 +887,18 @@ class CompiledGraph {
             node,
             step,
             threadId,
+            ...(runOptions.tenantId !== undefined ? { tenantId: runOptions.tenantId } : {}),
             taskId: task.id,
             input: task.input,
             attempt,
             signal,
-            ...(this.options.store ? { store: this.options.store } : {}),
+            ...(this.options.store
+                ? {
+                    store: runOptions.tenantId !== undefined
+                        ? (0, tenant_js_1.tenantStore)(this.options.store, runOptions.tenantId)
+                        : this.options.store,
+                }
+                : {}),
             interrupt: (request) => {
                 const index = interruptIndex++;
                 const key = (0, errors_js_1.interruptKey)({ node, taskId: task.id, step, index });
@@ -710,8 +910,21 @@ class CompiledGraph {
                 }
                 throw new errors_js_1.GraphInterrupt(request, node, step, index, task.id);
             },
-            emit: (data) => notify(runOptions, { type: 'custom', step, taskId: task.id, node, data }),
+            emit: (data) => {
+                touch();
+                notify(runOptions, { type: 'custom', step, taskId: task.id, node, data });
+            },
+            heartbeat: touch,
+            message: (chunk) => {
+                touch();
+                notify(runOptions, { type: 'message', step, taskId: task.id, node, chunk });
+            },
+            tool: (event) => {
+                touch();
+                notify(runOptions, { type: 'tool', step, taskId: task.id, node, tool: event });
+            },
             report: (progress) => {
+                touch();
                 if (!runOptions.onProgress)
                     return;
                 try {
@@ -721,6 +934,11 @@ class CompiledGraph {
                 }
             },
         };
+        if (runOptions.subgraphEvents) {
+            Object.defineProperty(context, FORWARD_EVENT, {
+                value: (event) => notify(runOptions, { ...event, namespace: [node, ...(event.namespace ?? [])] }),
+            });
+        }
         const result = await fn(context);
         const interrupted = interruptIndex > 0;
         if (result instanceof graph_js_1.Command) {
@@ -780,7 +998,12 @@ class CompiledGraph {
                 if (!this.nodes.has(target.node)) {
                     throw new errors_js_1.GraphValidationError(`Send from "${from}" targets unknown node "${target.node}"`);
                 }
-                tasks.push({ id: `${target.node}#${step + 1}.${sends++}`, node: target.node, input: target.input });
+                tasks.push({
+                    id: `${target.node}#${step + 1}.${sends++}`,
+                    node: target.node,
+                    input: target.input,
+                    ...(target.options?.timeout ? { timeout: target.options.timeout } : {}),
+                });
                 return;
             }
             if (target === graph_js_1.END)
@@ -807,7 +1030,9 @@ class CompiledGraph {
             }
         }
         for (const { from, target } of routes) {
-            add(typeof target === 'string' ? target : new graph_js_1.Send(target.node, target.input), from);
+            add(typeof target === 'string'
+                ? target
+                : new graph_js_1.Send(target.node, target.input, target.timeout ? { timeout: target.timeout } : undefined), from);
         }
         return tasks;
     }
@@ -823,9 +1048,9 @@ class CompiledGraph {
         }
         return this.reduce(state, [input]);
     }
-    async requireCheckpoint(threadId) {
+    async requireCheckpoint(threadId, tenantId) {
         const checkpoint = await this.checkpointer()?.get(threadId);
-        if (!checkpoint)
+        if (!checkpoint || !ownedBy(checkpoint, tenantId))
             throw new errors_js_1.GraphThreadNotFoundError(threadId);
         return (0, checkpoint_migration_js_1.migrateCheckpoint)(checkpoint);
     }
@@ -875,6 +1100,7 @@ class CompiledGraph {
                 ? { interrupt: checkpoint.interrupts[0], interrupts: checkpoint.interrupts }
                 : {}),
             ...(checkpoint.breakpoint ? { breakpoint: checkpoint.breakpoint } : {}),
+            ...(checkpoint.recovered?.length ? { recovered: checkpoint.recovered } : {}),
         };
     }
     toResult(threadId, last) {
@@ -971,21 +1197,77 @@ function toRouteRecords(goto) {
     if (goto === undefined)
         return undefined;
     return (Array.isArray(goto) ? goto : [goto]).map((target) => target instanceof graph_js_1.Send
-        ? { node: target.node, ...(target.input === undefined ? {} : { input: target.input }) }
+        ? {
+            node: target.node,
+            ...(target.input === undefined ? {} : { input: target.input }),
+            ...(target.options?.timeout ? { timeout: target.options.timeout } : {}),
+        }
         : target);
 }
 function routeName(route) {
     return typeof route === 'string' ? route : route.node;
 }
-function rejectWhenAborted(signal, error) {
-    return new Promise((_, reject) => {
-        if (signal.aborted) {
-            reject(error());
-            return;
-        }
-        signal.addEventListener('abort', () => reject(error()), { once: true });
+function noop() { }
+function attemptTimer(node, limits) {
+    const { runMs, idleMs } = limits;
+    if (runMs === undefined && idleMs === undefined)
+        return undefined;
+    const controller = new AbortController();
+    let reached;
+    let reject = noop;
+    const expired = new Promise((_, fail) => {
+        reject = fail;
     });
+    expired.catch(noop);
+    const fire = (error) => {
+        if (reached)
+            return;
+        reached = error;
+        controller.abort(error);
+        reject(error);
+    };
+    const runTimer = runMs === undefined ? undefined : setTimeout(() => fire(new errors_js_1.GraphNodeTimeoutError(node, runMs)), runMs);
+    let idleTimer;
+    const armIdle = () => {
+        if (idleMs === undefined || reached)
+            return;
+        if (idleTimer)
+            clearTimeout(idleTimer);
+        idleTimer = setTimeout(() => fire(new errors_js_1.GraphNodeTimeoutError(node, idleMs, 'idle')), idleMs);
+    };
+    armIdle();
+    return {
+        signal: controller.signal,
+        expired,
+        touch: armIdle,
+        error: () => reached,
+        stop: () => {
+            if (runTimer)
+                clearTimeout(runTimer);
+            if (idleTimer)
+                clearTimeout(idleTimer);
+            idleTimer = undefined;
+        },
+    };
 }
 function isRetryable(error) {
     return !(error instanceof errors_js_1.GraphValidationError);
 }
+function tenantOf(checkpoint) {
+    const tenantId = checkpoint.metadata?.tenantId;
+    return typeof tenantId === 'string' ? tenantId : undefined;
+}
+function ownedBy(checkpoint, tenantId) {
+    if (tenantId === undefined)
+        return true;
+    const owner = tenantOf(checkpoint);
+    return owner === undefined || owner === tenantId;
+}
+async function* releasing(source, release) {
+    try {
+        yield* source;
+    }
+    finally {
+        release();
+    }
+}

package/dist-cjs/graph/index.d.ts

@@ -2,6 +2,8 @@ export { CompiledGraph, StateGraph, createGraph } from './graph.js';
 export { type CheckpointDraft, type GraphCheckpointV2, migrateCheckpoint, toCheckpoint, } from './checkpoint-migration.js';
 export { appendList, appendSet, counter, lastValue, mergeObject, reducerChannel, } from './channels.js';
 export { MemoryGraphCheckpointer, OperationStoreCheckpointer, type MemoryGraphCheckpointerOptions, type OperationStoreCheckpointerOptions, } from './checkpointer.js';
-export { GraphError, GraphInterrupt, GraphNodeError, GraphNodeTimeoutError, GraphNotInterruptedError, GraphStepLimitError, GraphThreadNotFoundError, GraphValidationError, interruptKey, } from './errors.js';
+export { GraphDrainedError, GraphError, GraphInterrupt, GraphNodeError, GraphNodeTimeoutError, GraphNotInterruptedError, GraphStepLimitError, GraphThreadNotFoundError, GraphValidationError, interruptKey, } from './errors.js';
+export { RunControl } from './run-control.js';
+export { GraphEventStream, GraphStreamOverflowError } from './event-stream.js';
 export { Command, END, Send, START } from '../types/graph.js';
-export type { Channel, ChannelSchema, CommandTarget, CompileOptions, EdgeRouter, GraphBreakpoint, GraphCheckpoint, GraphCheckpointV1, GraphCheckpointer, GraphDescription, GraphEvent, GraphInput, GraphProgress, GraphResult, GraphRunOptions, GraphRouteTarget, GraphStatus, GraphStepEvent, GraphTask, InterruptRequest, NodeCacheEntry, NodeCachePolicy, NodeContext, NodeFn, NodeOptions, NodeResult, PendingInterrupt, RetryPolicy, StateOf, StateUpdate, StoredGraphCheckpoint, } from '../types/graph.js';
+export type { Channel, ChannelSchema, CommandTarget, CompileOptions, DurabilityMode, EdgeRouter, GraphBreakpoint, GraphCheckpoint, GraphCheckpointV1, GraphCheckpointer, GraphDescription, GraphEvent, GraphEventsOptions, GraphInput, GraphProgress, GraphResult, GraphRunOptions, GraphRouteTarget, GraphStatus, GraphStepEvent, GraphMessageChunk, GraphStreamEvent, GraphStreamOverflow, GraphStreamProjection, GraphStreamStats, GraphToolEvent, GraphTask, InterruptRequest, NodeCacheEntry, NodeCachePolicy, NodeContext, NodeErrorHandler, NodeFailure, NodeFn, NodeOptions, NodeResult, NodeTimeout, PendingInterrupt, RecoveredFailure, RetryPolicy, RunControlLike, SendOptions, StateOf, StateUpdate, StoredGraphCheckpoint, } from '../types/graph.js';