nexus-ai-pro 2.1.0 → 2.3.0

package/dist/graph/functional.d.ts

@@ -1,14 +1,19 @@
 import type { OperationLifecycleLike } from '../types/lifecycle.js';
+import type { Principal } from '../types/principal.js';
 import { GraphError } from './errors.js';
-import type { GraphCheckpoint, GraphCheckpointer, GraphDescription, GraphEvent, GraphStatus, InterruptRequest, PendingInterrupt, RetryPolicy } from '../types/graph.js';
+import type { GraphCheckpoint, GraphCheckpointer, GraphDescription, GraphEvent, GraphStatus, InterruptRequest, PendingInterrupt, RetryPolicy, RunControlLike } from '../types/graph.js';
 import type { Store } from '../types/store.js';
 /** What a workflow checkpoint holds: the input, every completed step's result, and the output. */
 export interface WorkflowState<I = unknown, O = unknown> {
     /** The input the workflow was started with. */
     input: I;
-    /** Each completed step's result, keyed by step name — `name#2` for the second call with a name. */
+    /**
+     * Each completed step's result, keyed by step name — `name#2` for the second call with a name. A
+     * step whose `onError` supplied its result after its retries ran out also records `recovered`.
+     */
     steps: Record<string, {
         value?: unknown;
+        recovered?: StepRecovery;
     }>;
     /** The workflow's return value, once it has completed. */
     output?: O;
@@ -23,21 +28,70 @@ export interface StepContext {
     readonly attempt: number;
     /** Aborted when the run is cancelled or the step exceeds its timeout. */
     readonly signal: AbortSignal;
+    /**
+     * Shows the step is still making progress, which refreshes its idle timeout. For long work that
+     * sends nothing else, such as a large upload or a slow tool call.
+     */
+    heartbeat(): void;
+}
+/** Limits on one attempt of a step: its whole run, and the time it may go without progress. */
+export interface StepTimeout {
+    /** Fails the attempt when it runs longer than this, in milliseconds. */
+    runMs?: number;
+    /**
+     * Fails the attempt when this long passes without `context.heartbeat()`, in milliseconds, which
+     * tells a hung step from a slow one that is working.
+     */
+    idleMs?: number;
+}
+/** What a step's `onError` receives once its retries have run out. */
+export interface StepFailure {
+    /** The step's name. */
+    name: string;
+    /** The key its result is recorded under. */
+    key: string;
+    /** Attempts made. */
+    attempts: number;
+    /** The last attempt's error. */
+    error: unknown;
+}
+/** How a step's result was recovered, as recorded in the workflow state. */
+export interface StepRecovery {
+    /** The error that exhausted the step's retries. */
+    error: {
+        name: string;
+        message: string;
+    };
+    /** Attempts made before `onError` decided. */
+    attempts: number;
 }
 /** How one step runs. */
-export interface StepOptions {
+export interface StepOptions<T = unknown> {
     /** Retries a failing attempt, overriding the workflow's default policy. */
     retry?: RetryPolicy;
     /** Aborts the step's signal and fails the attempt when it runs longer than this, in milliseconds. */
     timeoutMs?: number;
+    /** Run and idle limits on each attempt. `runMs` is `timeoutMs` by another name; `timeoutMs` wins when both are set. */
+    timeout?: StepTimeout;
+    /**
+     * Decides what happens once the step's retries run out. What it returns becomes the step's
+     * result, recorded with `recovered`, so a resumed run does not run the step again — a refund after
+     * a charge that kept failing, say. Throwing fails the step with that error. A cancelled run, a
+     * question, or a drain is never passed to it.
+     */
+    onError?: (failure: StepFailure) => T | Promise<T>;
 }
 /** What a workflow function receives besides its input. */
 export interface WorkflowContext {
     /** The thread this run belongs to. */
     readonly threadId: string;
+    /** The tenant the run is for, when it has one. */
+    readonly tenantId?: string;
+    /** Who the run is for, when it was started with a principal: the caller the server authenticated. */
+    readonly principal?: Readonly<Principal>;
     /** Aborted when the run is cancelled. */
     readonly signal: AbortSignal;
-    /** Long-term memory, when the workflow was created with a store. */
+    /** Long-term memory, when the workflow was created with a store. In a tenant's run, the tenant's view of it. */
     readonly store?: Store;
     /**
      * Runs a piece of work once per thread and records its result.
@@ -46,7 +100,7 @@ export interface WorkflowContext {
      * running. Name steps so the same call has the same name every time; a name used more than once
      * is numbered in call order. The result must survive JSON serialization.
      */
-    step<T>(name: string, run: (context: StepContext) => Promise<T> | T, options?: StepOptions): Promise<T>;
+    step<T>(name: string, run: (context: StepContext) => Promise<T> | T, options?: StepOptions<T>): Promise<T>;
     /**
      * Suspends the workflow until a person answers.
      *
@@ -73,6 +127,11 @@ export interface WorkflowOptions {
     store?: Store;
     /** Default retry policy for every step. A step's own policy wins. Defaults to one attempt. */
     retry?: RetryPolicy;
+    /**
+     * Defaults for every step: its retries, its run and idle timeouts, and its `onError`. A step's own
+     * options win, field by field.
+     */
+    stepDefaults?: Pick<StepOptions, 'retry' | 'timeout' | 'onError'>;
     /** Steps that may run at once, for a workflow that starts many with `Promise.all`. Defaults to 16. */
     maxConcurrency?: number;
     /**
@@ -93,6 +152,25 @@ export interface WorkflowRunOptions {
     metadata?: Record<string, unknown>;
     /** Receives each step starting, retrying, and finishing, each checkpoint, and custom events. */
     onEvent?: (event: GraphEvent) => void;
+    /**
+     * Lets the caller stop the run between steps: pass a `RunControl` and call `drain()`. Steps in
+     * flight finish and are recorded, no new step starts, and the run ends with `GraphDrainedError`;
+     * `continue()` carries it on, on this worker or another. The agent server drains a workflow it
+     * serves this way.
+     */
+    control?: RunControlLike;
+    /**
+     * The tenant the run is for. Recorded on every checkpoint as `metadata.tenantId`, read as
+     * `context.tenantId`, and scoping `context.store`; another tenant's thread is not found. A run
+     * without one carries on the thread's.
+     */
+    tenantId?: string;
+    /**
+     * Who the run is for, read as `context.principal`. Its `tenantId` is the run's tenant when
+     * `tenantId` is not given, and a run whose two tenants differ is refused; its `userId` is recorded
+     * on every checkpoint as `metadata.userId`.
+     */
+    principal?: Principal;
 }
 /** One event of a streamed workflow run. */
 export type WorkflowEvent<O = unknown> = {
@@ -152,17 +230,21 @@ export interface WorkflowResult<O = unknown> {
     /** Every pending question. */
     interrupts?: PendingInterrupt[];
 }
-/** Raised when a step outlives its `timeoutMs`. */
+/** Raised when a step outlives its run timeout, or goes longer than its idle timeout without progress. */
 export declare class WorkflowStepTimeoutError extends GraphError {
     /** The step. */
     readonly stepName: string;
     /** The limit it exceeded, in milliseconds. */
     readonly timeoutMs: number;
+    /** Which limit: `run` for the whole attempt, `idle` for the time without progress. */
+    readonly kind: 'run' | 'idle';
     constructor(
     /** The step. */
     stepName: string, 
     /** The limit it exceeded, in milliseconds. */
-    timeoutMs: number);
+    timeoutMs: number, 
+    /** Which limit: `run` for the whole attempt, `idle` for the time without progress. */
+    kind?: 'run' | 'idle');
 }
 /**
  * A durable workflow: a function whose steps are checkpointed, so it survives interrupts, restarts,

package/dist/graph/functional.js

@@ -1,17 +1,28 @@
 import { MemoryGraphCheckpointer } from './checkpointer.js';
 import { withLifecycle } from '../lifecycle/run.js';
 import { migrateCheckpoint, toCheckpoint } from './checkpoint-migration.js';
-import { GraphError, GraphNotInterruptedError, GraphThreadNotFoundError } from './errors.js';
+import { tenantStore } from '../store/tenant.js';
+import { withPrincipal } from '../utils/tenant.js';
+import { GraphDrainedError, GraphError, GraphNotInterruptedError, GraphThreadNotFoundError } from './errors.js';
 export class WorkflowStepTimeoutError extends GraphError {
     stepName;
     timeoutMs;
-    constructor(stepName, timeoutMs) {
-        super(`Step "${stepName}" exceeded its ${timeoutMs} ms timeout`, 'WORKFLOW_STEP_TIMEOUT');
+    kind;
+    constructor(stepName, timeoutMs, kind = 'run') {
+        super(kind === 'idle'
+            ? `Step "${stepName}" made no progress for ${timeoutMs} ms`
+            : `Step "${stepName}" exceeded its ${timeoutMs} ms timeout`, 'WORKFLOW_STEP_TIMEOUT');
         this.stepName = stepName;
         this.timeoutMs = timeoutMs;
+        this.kind = kind;
         this.name = 'WorkflowStepTimeoutError';
     }
 }
+class WorkflowDraining extends Error {
+    constructor() {
+        super('The workflow is draining');
+    }
+}
 class WorkflowSuspended extends Error {
     pending;
     constructor(pending) {
@@ -40,10 +51,17 @@ export class Workflow {
     }
     stream(input, runOptions = {}) {
         const threadId = runOptions.threadId ?? generatedId();
-        return this.execute(threadId, runOptions, async () => ({
-            ...this.fresh(threadId, input),
-            step: (await this.latest(threadId))?.step ?? 0,
-        }));
+        return this.execute(threadId, runOptions, async () => {
+            const latest = await this.latest(threadId);
+            const tenantId = runOptions.tenantId ?? runOptions.principal?.tenantId;
+            if (latest && tenantId !== undefined) {
+                const owner = latest.metadata?.tenantId;
+                if (typeof owner === 'string' && owner !== tenantId) {
+                    throw new GraphError(`Thread "${threadId}" cannot be used by this tenant; choose another id`, 'GRAPH_VALIDATION_ERROR');
+                }
+            }
+            return { ...this.fresh(threadId, input), step: latest?.step ?? 0 };
+        });
     }
     resume(threadId, value, runOptions = {}) {
         return this.execute(threadId, runOptions, async () => {
@@ -132,16 +150,40 @@ export class Workflow {
     }
     execute(threadId, runOptions, start) {
         const queue = new EventQueue();
-        void withLifecycle(this.options.lifecycle, { family: 'graph', operation: `workflow.${this.name}`, metadata: { threadId, graph: this.name } }, runOptions.signal ? { signal: runOptions.signal } : {}, () => this.run(threadId, runOptions, start, queue)).then(() => queue.close(), (error) => queue.fail(error));
+        let options;
+        try {
+            options = withPrincipal(runOptions);
+        }
+        catch (error) {
+            queue.fail(error);
+            return queue;
+        }
+        const metadata = {
+            threadId,
+            graph: this.name,
+            ...(options.tenantId !== undefined ? { tenantId: options.tenantId } : {}),
+            ...(options.principal?.userId !== undefined ? { userId: options.principal.userId } : {}),
+        };
+        void withLifecycle(this.options.lifecycle, { family: 'graph', operation: `workflow.${this.name}`, metadata }, options.signal ? { signal: options.signal } : {}, () => this.run(threadId, options, start, queue)).then(() => queue.close(), (error) => queue.fail(error));
         return queue;
     }
     async run(threadId, runOptions, start, queue) {
         const initial = await start();
+        const owner = typeof initial.metadata?.tenantId === 'string' ? initial.metadata.tenantId : undefined;
+        if (runOptions.tenantId !== undefined && owner !== undefined && owner !== runOptions.tenantId) {
+            throw new GraphThreadNotFoundError(threadId);
+        }
+        const tenantId = runOptions.tenantId ?? owner;
         const checkpointer = this.checkpointer();
         const state = structuredClone(initial.state);
         state.steps ??= {};
         const resolved = { ...initial.resolved };
-        const metadata = { ...initial.metadata, graph: this.name, ...runOptions.metadata };
+        const metadata = {
+            ...initial.metadata,
+            graph: this.name,
+            ...runOptions.metadata,
+            ...(tenantId === undefined ? {} : { tenantId }),
+        };
         let stepNumber = initial.step;
         let writes = Promise.resolve();
         const persist = (checkpoint) => {
@@ -171,11 +213,17 @@ export class Workflow {
         const inFlight = new Set();
         const suspended = [];
         let interruptIndex = 0;
+        let drained = false;
         const limiter = new Limiter(this.options.maxConcurrency ?? 16);
+        const defaults = this.options.stepDefaults ?? {};
         const context = {
             threadId,
+            ...(tenantId === undefined ? {} : { tenantId }),
+            ...(runOptions.principal ? { principal: runOptions.principal } : {}),
             signal: controller.signal,
-            ...(this.options.store ? { store: this.options.store } : {}),
+            ...(this.options.store
+                ? { store: tenantId === undefined ? this.options.store : tenantStore(this.options.store, tenantId) }
+                : {}),
             step: (name, run, options = {}) => {
                 const count = (counts.get(name) ?? 0) + 1;
                 counts.set(name, count);
@@ -185,10 +233,22 @@ export class Workflow {
                     queue.push({ type: 'step', step: stepNumber, name, key, cached: true, state: snapshot(), status: 'running' });
                     return Promise.resolve(recorded.value);
                 }
+                if (runOptions.control?.draining) {
+                    drained = true;
+                    return Promise.reject(new WorkflowDraining());
+                }
+                const merged = {
+                    ...defaults,
+                    ...options,
+                    timeout: { ...defaults.timeout, ...options.timeout },
+                };
                 const work = limiter
-                    .run(() => this.runStep(name, key, run, { ...this.options.retry, ...options.retry }, options, controller.signal, runOptions, stepNumber))
-                    .then(async (value) => {
-                    state.steps[key] = value === undefined ? {} : { value: structuredClone(value) };
+                    .run(() => this.runStep(name, key, run, { ...this.options.retry, ...defaults.retry, ...options.retry }, merged, controller.signal, runOptions, stepNumber))
+                    .then(async ({ value, recovered }) => {
+                    state.steps[key] = {
+                        ...(value === undefined ? {} : { value: structuredClone(value) }),
+                        ...(recovered ? { recovered } : {}),
+                    };
                     await persist({ state: state, next: [this.name], status: 'running' });
                     queue.push({
                         type: 'step',
@@ -262,6 +322,18 @@ export class Workflow {
             queue.push({ type: 'interrupt', step: stepNumber, state: snapshot(), status: 'interrupted' });
             return;
         }
+        if (drained && outcome.kind === 'failed') {
+            const reason = runOptions.control?.reason;
+            await persist({
+                state: state,
+                next: [this.name],
+                status: 'interrupted',
+                resolved,
+                drained: { ...(reason ? { reason } : {}), at: this.now().toISOString() },
+            });
+            await writes;
+            throw new GraphDrainedError(threadId, stepNumber, reason);
+        }
         if (outcome.kind === 'failed') {
             const error = outcome.error;
             await persist({
@@ -284,41 +356,87 @@ export class Workflow {
     }
     async runStep(name, key, run, retry, options, signal, runOptions, step) {
         const policy = { ...DEFAULT_RETRY, ...retry };
+        const runMs = options.timeoutMs ?? options.timeout?.runMs;
+        const idleMs = options.timeout?.idleMs;
         for (let attempt = 1;; attempt += 1) {
             notify(runOptions, { type: 'task_start', step, taskId: key, node: name, attempt });
             const stepController = new AbortController();
             const forward = () => stepController.abort(signal.reason);
             signal.addEventListener('abort', forward, { once: true });
-            let timer;
-            let timedOut = false;
+            let runTimer;
+            let idleTimer;
+            let timedOut;
+            let expire = () => undefined;
+            const limits = runMs === undefined && idleMs === undefined
+                ? undefined
+                : new Promise((_, reject) => {
+                    expire = (error) => {
+                        if (timedOut)
+                            return;
+                        timedOut = error;
+                        reject(error);
+                        stepController.abort();
+                    };
+                });
+            limits?.catch(() => undefined);
+            if (runMs !== undefined)
+                runTimer = setTimeout(() => expire(new WorkflowStepTimeoutError(name, runMs)), runMs);
+            const armIdle = () => {
+                if (idleMs === undefined || timedOut)
+                    return;
+                if (idleTimer)
+                    clearTimeout(idleTimer);
+                idleTimer = setTimeout(() => expire(new WorkflowStepTimeoutError(name, idleMs, 'idle')), idleMs);
+            };
+            armIdle();
             try {
-                const work = Promise.resolve(run({ name, key, attempt, signal: stepController.signal }));
-                const value = options.timeoutMs === undefined
-                    ? await work
-                    : await Promise.race([
-                        work,
-                        new Promise((_, reject) => {
-                            timer = setTimeout(() => {
-                                timedOut = true;
-                                reject(new WorkflowStepTimeoutError(name, options.timeoutMs));
-                                stepController.abort();
-                            }, options.timeoutMs);
-                        }),
-                    ]);
+                const work = Promise.resolve(run({ name, key, attempt, signal: stepController.signal, heartbeat: armIdle }));
+                const value = limits ? await Promise.race([work, limits]) : await work;
                 if (timedOut)
-                    throw new WorkflowStepTimeoutError(name, options.timeoutMs);
+                    throw timedOut;
                 notify(runOptions, { type: 'task_end', step, taskId: key, node: name, update: value });
-                return value;
+                return { value };
             }
             catch (error) {
                 if (signal.aborted || error instanceof WorkflowSuspended)
                     throw error;
-                const failure = timedOut ? new WorkflowStepTimeoutError(name, options.timeoutMs) : error;
+                const failure = timedOut ?? error;
                 const retryable = policy.retryOn
                     ? policy.retryOn(failure, attempt)
-                    : !(failure instanceof GraphError) || timedOut;
-                if (attempt >= Math.max(1, policy.maxAttempts) || !retryable)
-                    throw failure;
+                    : !(failure instanceof GraphError) || timedOut !== undefined;
+                if (attempt >= Math.max(1, policy.maxAttempts) || !retryable) {
+                    if (!options.onError)
+                        throw failure;
+                    const failed = (recovered) => notify(runOptions, {
+                        type: 'task_failed',
+                        step,
+                        taskId: key,
+                        node: name,
+                        attempts: attempt,
+                        error: failure instanceof Error ? failure.message : String(failure),
+                        recovered,
+                    });
+                    let value;
+                    try {
+                        value = await options.onError({ name, key, attempts: attempt, error: failure });
+                    }
+                    catch (decision) {
+                        failed(false);
+                        throw decision;
+                    }
+                    failed(true);
+                    notify(runOptions, { type: 'task_end', step, taskId: key, node: name, update: value });
+                    return {
+                        value,
+                        recovered: {
+                            error: {
+                                name: failure instanceof Error ? failure.name : 'Error',
+                                message: failure instanceof Error ? failure.message : String(failure),
+                            },
+                            attempts: attempt,
+                        },
+                    };
+                }
                 notify(runOptions, {
                     type: 'task_retry',
                     step,
@@ -330,8 +448,10 @@ export class Workflow {
                 await sleep(backoff(policy, attempt), signal);
             }
             finally {
-                if (timer)
-                    clearTimeout(timer);
+                if (runTimer)
+                    clearTimeout(runTimer);
+                if (idleTimer)
+                    clearTimeout(idleTimer);
                 signal.removeEventListener('abort', forward);
             }
         }

package/dist/graph/graph.d.ts

@@ -111,10 +111,17 @@ export declare class CompiledGraph<S extends ChannelSchema, I extends keyof S =
      * once the rewound checkpoint is written.
      */
     resumeFrom(threadId: string, step: number, runOptions?: GraphRunOptions): AsyncIterable<GraphStepEvent<S>>;
-    /** Latest checkpoint, or the one at `step`. */
-    state(threadId: string, step?: number): Promise<GraphCheckpoint<S> | undefined>;
-    /** Checkpoints for a thread, newest first. */
-    history(threadId: string, limit?: number): Promise<Array<GraphCheckpoint<S>>>;
+    /**
+     * Latest checkpoint, or the one at `step`. With a `tenantId`, a thread another tenant owns reads
+     * as absent.
+     */
+    state(threadId: string, step?: number, options?: {
+        tenantId?: string;
+    }): Promise<GraphCheckpoint<S> | undefined>;
+    /** Checkpoints for a thread, newest first. With a `tenantId`, another tenant's thread has none. */
+    history(threadId: string, limit?: number, options?: {
+        tenantId?: string;
+    }): Promise<Array<GraphCheckpoint<S>>>;
     /**
      * Wraps this graph as a node in another graph.
      *
@@ -143,17 +150,20 @@ export declare class CompiledGraph<S extends ChannelSchema, I extends keyof S =
      */
     updateState(threadId: string, update: StateUpdate<S>, options?: {
         asNode?: string;
+        tenantId?: string;
     }): Promise<GraphCheckpoint<S>>;
     /**
      * Copies a thread's history up to `step` into a new thread, and returns the new thread's id.
      *
      * Rewinding with `resumeFrom()` replaces the original timeline. A fork keeps it: both threads stay
      * readable and runnable, so two answers to the same question can be compared side by side. Every
-     * copied checkpoint records `metadata.forkedFrom`.
+     * copied checkpoint records `metadata.forkedFrom`, and keeps the tenant of the thread it came from.
+     * With a `tenantId`, another tenant's thread is not found.
      */
     fork(threadId: string, options?: {
         step?: number;
         threadId?: string;
+        tenantId?: string;
     }): Promise<string>;
     /**
      * Runs the supersteps, as one operation of the compiled lifecycle when there is one. The run is