nexus-ai-pro 2.0.0 → 2.2.0

package/dist/graph/graph.js

@@ -1,8 +1,12 @@
 import { randomBytes } from 'node:crypto';
 import { Command, END, Send, START } from '../types/graph.js';
+import { tenantStore } from '../store/tenant.js';
+import { linkSignals } from '../utils/signals.js';
 import { MemoryGraphCheckpointer } from './checkpointer.js';
 import { migrateCheckpoint, toCheckpoint } from './checkpoint-migration.js';
-import { GraphInterrupt, GraphNodeError, GraphNodeTimeoutError, GraphNotInterruptedError, GraphStepLimitError, GraphThreadNotFoundError, GraphValidationError, interruptKey, } from './errors.js';
+import { CheckpointWriter } from './durability.js';
+import { GraphEventStream, streamSettings } from './event-stream.js';
+import { GraphDrainedError, GraphInterrupt, GraphNodeError, GraphNodeTimeoutError, GraphNotInterruptedError, GraphStepLimitError, GraphThreadNotFoundError, GraphValidationError, interruptKey, } from './errors.js';
 const DEFAULT_MAX_STEPS = 25;
 const DEFAULT_MAX_CONCURRENCY = 16;
 const DEFAULT_RETRY = {
@@ -13,6 +17,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');
 export class StateGraph {
     channels;
@@ -123,6 +128,7 @@ export class CompiledGraph {
     now;
     store;
     nodeCache;
+    writers = new Set();
     constructor(channels, nodes, edges, options, scope = {}) {
         this.channels = channels;
         this.nodes = nodes;
@@ -132,6 +138,9 @@ export class CompiledGraph {
         this.now = options.now ?? (() => new Date());
         this.store = options.checkpointer === false ? undefined : (options.checkpointer ?? new MemoryGraphCheckpointer());
     }
+    async flush() {
+        await Promise.all([...this.writers].map((writer) => writer.flush()));
+    }
     async invoke(input = {}, runOptions = {}) {
         const threadId = resolveThreadId(runOptions.threadId);
         let last;
@@ -143,11 +152,42 @@ export 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 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 GraphEventStream((onEvent, signal) => {
+            const linked = 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), streamSettings(options));
+    }
     resume(threadId, value, runOptions = {}) {
         return this.run(threadId, runOptions, async () => {
             const checkpoint = await this.requireCheckpoint(threadId);
@@ -197,7 +237,7 @@ export 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 = {}) {
@@ -208,19 +248,25 @@ export class CompiledGraph {
             return { ...migrateCheckpoint(stored), status: 'running' };
         });
     }
-    async state(threadId, step) {
+    async state(threadId, step, options = {}) {
         const stored = await this.checkpointer()?.get(threadId, step);
-        return stored ? migrateCheckpoint(stored) : undefined;
+        if (!stored || !ownedBy(stored, options.tenantId))
+            return undefined;
+        return 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) => 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;
@@ -263,13 +309,18 @@ export 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() } : {}),
             };
@@ -304,10 +355,13 @@ export 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) {
@@ -342,8 +396,9 @@ export class CompiledGraph {
         if (!store)
             throw new GraphValidationError('fork() needs a checkpointer to copy history from');
         const history = (await store.history(threadId, Number.MAX_SAFE_INTEGER)).map((item) => migrateCheckpoint(item));
-        if (history.length === 0)
+        if (history.length === 0 || !ownedBy(history[0], options.tenantId)) {
             throw new GraphThreadNotFoundError(threadId);
+        }
         const cutoff = options.step ?? history[0].step;
         const lineage = history.filter((item) => item.step <= cutoff).reverse();
         if (lineage.at(-1)?.step !== cutoff) {
@@ -394,22 +449,69 @@ export class CompiledGraph {
         }
     }
     async *runSteps(threadId, runOptions, seed) {
+        const mode = runOptions.durability ?? this.options.durability ?? 'sync';
+        const writer = new 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 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 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;
             }
@@ -434,7 +536,7 @@ export class CompiledGraph {
                     breakpoint: { when: 'before', nodes: breakBefore },
                     createdAt: this.now().toISOString(),
                 };
-                await persist(checkpoint);
+                await persist(checkpoint, true);
                 yield this.toEvent(checkpoint, []);
                 return;
             }
@@ -454,6 +556,7 @@ export 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)
@@ -465,6 +568,8 @@ export 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') {
@@ -480,6 +585,7 @@ export 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,
@@ -487,6 +593,7 @@ export class CompiledGraph {
                     ...pausedFields(allTasks, unfinished),
                     completed,
                     ...routeFields,
+                    ...recoveryFields,
                     status: 'failed',
                     error: {
                         name: failure.error instanceof Error ? failure.error.name : 'Error',
@@ -494,7 +601,7 @@ export class CompiledGraph {
                     },
                     createdAt: this.now().toISOString(),
                 };
-                await persist(failed);
+                await persist(failed, true);
                 throw failure.error instanceof GraphNodeError
                     ? failure.error
                     : new GraphNodeError(failure.task.node, step, failure.error);
@@ -507,11 +614,12 @@ export 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;
             }
@@ -523,10 +631,11 @@ export 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;
             }
@@ -545,10 +654,11 @@ export 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;
@@ -558,8 +668,8 @@ export class CompiledGraph {
         const definition = this.nodes.get(task.node);
         if (!definition)
             throw new 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;
@@ -581,22 +691,18 @@ export 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 = 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 GraphNodeTimeoutError(task.node, timeoutMs)),
-                    ])
-                    : await running;
+                const output = timer ? await Promise.race([running, timer.expired]) : await running;
                 const parent = output.command?.graph === Command.PARENT ? output.command : undefined;
                 const goto = parent ? undefined : toRouteRecords(output.command?.goto);
                 const updates = parent ? [] : output.updates;
@@ -614,6 +720,7 @@ export class CompiledGraph {
                 return { kind: 'ok', updates, goto, parent, attempts };
             }
             catch (error) {
+                timer?.stop();
                 if (error instanceof GraphInterrupt) {
                     return {
                         kind: 'interrupt',
@@ -631,10 +738,24 @@ export class CompiledGraph {
                 }
                 if (runOptions.signal?.aborted || stepSignal.aborted)
                     return { kind: 'aborted', attempts };
-                const failure = timeout?.aborted ? new 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 GraphNodeTimeoutError
+                            ? { timeout: { type: failure.kind, limitMs: failure.timeoutMs } }
+                            : {}),
+                        ...(task.input === undefined ? {} : { input: task.input }),
+                        threadId,
+                    });
                 }
                 notify(runOptions, {
                     type: 'task_retry',
@@ -646,9 +767,81 @@ export class CompiledGraph {
                 });
                 await delay(backoffDelay(policy, attempts), runOptions.signal);
             }
+            finally {
+                timer?.stop();
+                linked.dispose();
+            }
+        }
+    }
+    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 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 Command ? decided : undefined;
+        const parent = command?.graph === 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,
+        };
     }
-    async runNode(fn, task, step, threadId, checkpoint, runOptions, signal, attempt) {
+    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 = {
@@ -656,11 +849,18 @@ export 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
+                        ? tenantStore(this.options.store, runOptions.tenantId)
+                        : this.options.store,
+                }
+                : {}),
             interrupt: (request) => {
                 const index = interruptIndex++;
                 const key = interruptKey({ node, taskId: task.id, step, index });
@@ -672,8 +872,21 @@ export class CompiledGraph {
                 }
                 throw new 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 {
@@ -683,6 +896,11 @@ export 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 Command) {
@@ -742,7 +960,12 @@ export class CompiledGraph {
                 if (!this.nodes.has(target.node)) {
                     throw new 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 === END)
@@ -769,7 +992,9 @@ export class CompiledGraph {
             }
         }
         for (const { from, target } of routes) {
-            add(typeof target === 'string' ? target : new Send(target.node, target.input), from);
+            add(typeof target === 'string'
+                ? target
+                : new Send(target.node, target.input, target.timeout ? { timeout: target.timeout } : undefined), from);
         }
         return tasks;
     }
@@ -785,9 +1010,9 @@ export 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 GraphThreadNotFoundError(threadId);
         return migrateCheckpoint(checkpoint);
     }
@@ -837,6 +1062,7 @@ export class CompiledGraph {
                 ? { interrupt: checkpoint.interrupts[0], interrupts: checkpoint.interrupts }
                 : {}),
             ...(checkpoint.breakpoint ? { breakpoint: checkpoint.breakpoint } : {}),
+            ...(checkpoint.recovered?.length ? { recovered: checkpoint.recovered } : {}),
         };
     }
     toResult(threadId, last) {
@@ -932,21 +1158,77 @@ function toRouteRecords(goto) {
     if (goto === undefined)
         return undefined;
     return (Array.isArray(goto) ? goto : [goto]).map((target) => target instanceof 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 GraphNodeTimeoutError(node, runMs)), runMs);
+    let idleTimer;
+    const armIdle = () => {
+        if (idleMs === undefined || reached)
+            return;
+        if (idleTimer)
+            clearTimeout(idleTimer);
+        idleTimer = setTimeout(() => fire(new 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 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/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';