@cellaflow/sdk 0.7.0 → 0.7.1

package/src/execution.ts

@@ -0,0 +1,224 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+import { CellaflowClient } from "./client.js";
+import { CacheStatus } from "./cellaflow/v1/idempotency_pb.js";
+import { LeaseHeartbeat } from "./lease.js";
+
+const DEFAULT_TTL_MS = 30_000;
+const DEFAULT_HEARTBEAT_INTERVAL_MS = 5_000;
+
+/** Raised when the requested execution lease cannot be acquired. */
+export class LeaseNotAcquired extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "LeaseNotAcquired";
+  }
+}
+
+/** Raised by `LeaseHandle.check()` once the lease has been lost. */
+export class LeaseLostError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "LeaseLostError";
+  }
+}
+
+/**
+ * Handle to an active execution lease.
+ *
+ * Node cannot interrupt a running task the way Python's asyncio can cancel one,
+ * so losing a lease here is **cooperative**: the work keeps running until it
+ * checks. Two ways to check, and long or irreversible work should use one:
+ *
+ * - `check()` throws {@link LeaseLostError} once the lease is gone.
+ * - `signal` aborts, so it can be handed to `fetch` or any AbortSignal-aware API.
+ */
+export class LeaseHandle {
+  readonly fencingToken: number;
+  private readonly controller = new AbortController();
+  private lostDetail?: string;
+
+  constructor(fencingToken: number) {
+    this.fencingToken = fencingToken;
+  }
+
+  /** Aborts when the lease is lost. Pass to `fetch`, streams, or your own loops. */
+  get signal(): AbortSignal {
+    return this.controller.signal;
+  }
+
+  get isLost(): boolean {
+    return this.controller.signal.aborted;
+  }
+
+  /** @internal */
+  markLost(detail: string): void {
+    if (this.controller.signal.aborted) return;
+    this.lostDetail = detail;
+    this.controller.abort(new LeaseLostError(`Execution lease was lost: ${detail}`));
+  }
+
+  /**
+   * Throws if the lease has been lost.
+   *
+   * Call this before anything irreversible inside a long block. Nothing else
+   * stops the work: losing a lease cannot preempt a running function in Node.
+   */
+  check(): void {
+    if (this.isLost) {
+      throw new LeaseLostError(`Execution lease was lost: ${this.lostDetail}`);
+    }
+  }
+}
+
+const currentLeaseStore = new AsyncLocalStorage<LeaseHandle>();
+
+/**
+ * Returns the lease held by the enclosing block.
+ *
+ * Lets code inside an {@link executionLease} block, or inside a
+ * {@link taskLease} function which has no handle to receive, reach the fencing
+ * token to pass downstream.
+ */
+export function currentLease(): LeaseHandle {
+  const handle = currentLeaseStore.getStore();
+  if (!handle) {
+    throw new Error(
+      "No execution lease is active. currentLease() is only valid inside an " +
+        "executionLease(...) block or a taskLease(...) function.",
+    );
+  }
+  return handle;
+}
+
+export interface ExecutionLeaseOptions {
+  /** Identifies this worker. Sent as `agentId`. */
+  workerId: string;
+  target?: string;
+  secure?: boolean;
+  ttlMs?: number;
+  heartbeatIntervalMs?: number;
+  /** Notified when the lease is lost. A notification, not a substitute for checking. */
+  onLeaseLost?: (detail: string) => void;
+  /**
+   * How long to wait for a lease another worker is holding. Defaults to 0, which
+   * fails immediately with {@link LeaseNotAcquired}.
+   */
+  waitMs?: number;
+}
+
+/**
+ * Runs `fn` holding a distributed lock on `key`, renewed by a heartbeat.
+ *
+ * Only one worker runs the block at a time. If this process dies, the lease
+ * stops being renewed and another worker takes it once the TTL elapses, which is
+ * the property a plain database lock cannot offer for a holder that hangs
+ * without dying.
+ *
+ * ```ts
+ * await executionLease("task:abc-123", { workerId: "worker-1" }, async (lease) => {
+ *   lease.check();
+ *   await doTheWork({ signal: lease.signal });
+ * });
+ * ```
+ *
+ * Unlike the Python `async_execution_lease`, losing the lease does **not**
+ * interrupt `fn`: Node has no task cancellation. The handle exposes `check()`
+ * and `signal` so the work can abort itself, and anything irreversible should
+ * check first.
+ */
+export async function executionLease<T>(
+  key: string,
+  options: ExecutionLeaseOptions,
+  fn: (lease: LeaseHandle) => Promise<T>,
+): Promise<T> {
+  const {
+    workerId,
+    target = "localhost:50051",
+    secure = false,
+    ttlMs = DEFAULT_TTL_MS,
+    heartbeatIntervalMs = DEFAULT_HEARTBEAT_INTERVAL_MS,
+    onLeaseLost,
+    waitMs = 0,
+  } = options;
+
+  const client = new CellaflowClient({ target, secure });
+  let hb: LeaseHeartbeat | undefined;
+  let fencingToken = 0;
+
+  try {
+    const resp = await client.checkIdempotencyCache(workerId, key, waitMs, ttlMs);
+
+    if (resp.status === CacheStatus.HIT) {
+      throw new LeaseNotAcquired(
+        `'${key}' is already committed: the operation it names has completed. An ` +
+          "execution lease locks work still to be done, so a committed key means " +
+          "this work is finished, not that the lock is busy.",
+      );
+    }
+    if (resp.status !== CacheStatus.ACQUIRED) {
+      const holder = resp.currentHolderId ? ` (held by ${resp.currentHolderId})` : "";
+      throw new LeaseNotAcquired(
+        `Could not acquire execution lease '${key}'${holder}. Another worker is ` +
+          "running it. Raise waitMs to wait for them, or treat this as the " +
+          "signal that the work is already in hand.",
+      );
+    }
+
+    fencingToken = Number(resp.fencingToken ?? 0n);
+    const handle = new LeaseHandle(fencingToken);
+
+    hb = new LeaseHeartbeat({
+      client,
+      agentId: workerId,
+      idempotencyKey: key,
+      fencingToken,
+      heartbeatIntervalMs: Number(resp.heartbeatIntervalMs ?? BigInt(heartbeatIntervalMs)),
+      leaseTtlMs: ttlMs,
+      onLeaseLost: (detail) => {
+        handle.markLost(detail);
+        try {
+          onLeaseLost?.(detail);
+        } catch {
+          // A caller's callback must not mask the loss itself.
+        }
+      },
+    });
+    hb.start();
+
+    return await currentLeaseStore.run(handle, () => fn(handle));
+  } finally {
+    await hb?.stop();
+    if (fencingToken > 0) {
+      await client
+        .releaseLease(workerId, key, fencingToken, "BLOCK_EXIT")
+        .catch(() => {
+          // The lease expires on its own; a failed release is not worth masking
+          // whatever the block was doing.
+        });
+    }
+    client.close();
+  }
+}
+
+/**
+ * Wraps a function so every call runs under an execution lease.
+ *
+ * The key is derived per call, so a task id argument becomes the lock:
+ *
+ * ```ts
+ * const processOrder = taskLease(
+ *   async (orderId: string) => { currentLease().check(); await ship(orderId); },
+ *   { workerId: "worker-1", key: (orderId) => `order:${orderId}` },
+ * );
+ * ```
+ */
+export function taskLease<A extends unknown[], R>(
+  fn: (...args: A) => Promise<R>,
+  options: ExecutionLeaseOptions & { key: string | ((...args: A) => string) },
+): (...args: A) => Promise<R> {
+  const { key, ...leaseOptions } = options;
+  return async function leasedTask(...args: A): Promise<R> {
+    const resolved = typeof key === "function" ? key(...args) : key;
+    return executionLease(resolved, leaseOptions, () => fn(...args));
+  };
+}

package/src/idempotency.ts

@@ -0,0 +1,134 @@
+import canonicalize from "canonicalize";
+import { createHash } from "node:crypto";
+
+/**
+ * How widely a derived idempotency key deduplicates.
+ *
+ * The numeric values match the Python SDK and the engine's enum. They are part
+ * of the wire contract, not an implementation detail.
+ */
+export enum IdempotencyScope {
+  UNSPECIFIED = 0,
+  /** Shared across all agents in the session. The default. */
+  SESSION_WIDE = 1,
+  /** Isolated to the executing agent. */
+  AGENT_PRIVATE = 2,
+  /** Isolated to the specific superstep / node. */
+  STEP_LOCAL = 3,
+  /**
+   * Shared across *sessions* within a declared coordination domain.
+   *
+   * The only scope whose key omits `sessionId`, so agents running different
+   * workflows in different sessions deduplicate one shared side effect.
+   */
+  SHARED = 4,
+}
+
+/**
+ * Hashes inputs with RFC 8785 Canonical JSON and SHA-256, returning the first
+ * 16 bytes hex-encoded.
+ *
+ * Byte-identical to the Python SDK's `_hash_inputs`, which is what allows a
+ * TypeScript agent and a Python agent to converge on one key under
+ * {@link IdempotencyScope.SHARED}. Verified against Python reference hashes
+ * covering unicode, escapes, nested objects, key reordering and exponent-form
+ * numbers. Changing the canonicalisation or the digest length silently stops
+ * cross-language deduplication, which fails by repeating the side effect rather
+ * than by raising.
+ */
+export function hashInputs(args: unknown[], kwargs: Record<string, unknown>): string {
+  const canon = canonicalize({ args, kwargs });
+  if (canon === undefined) {
+    throw new TypeError(
+      "Tool arguments could not be canonicalised. Arguments contributing to an " +
+        "idempotency key must be JSON-representable: no functions, symbols, " +
+        "BigInt, undefined, or circular references.",
+    );
+  }
+  return createHash("sha256").update(canon, "utf8").digest("hex").slice(0, 32);
+}
+
+export interface DeriveKeyOptions {
+  sessionId: string;
+  workflowVersion: string;
+  stepSequence: number;
+  agentId: string;
+  toolName: string;
+  scope: IdempotencyScope;
+  coordinationId?: string;
+  /** Positional arguments the tool was called with. */
+  args?: unknown[];
+  /** Named arguments, when the tool takes a single options object. */
+  kwargs?: Record<string, unknown>;
+  /**
+   * Restricts the hash to the named keys of `kwargs`.
+   *
+   * Heterogeneous agents converge on one side effect precisely when they
+   * disagree about everything else, so hashing everything they pass is the one
+   * thing guaranteed to keep them apart.
+   *
+   * JavaScript has no named arguments, so a tool using `sharedOn` must take a
+   * single options object; the names are read from it.
+   */
+  sharedOn?: readonly string[];
+}
+
+/** Derives the canonical idempotency key for a step or tool execution. */
+export function deriveIdempotencyKey(opts: DeriveKeyOptions): string {
+  const {
+    sessionId,
+    workflowVersion,
+    stepSequence,
+    agentId,
+    toolName,
+    scope,
+    coordinationId,
+    args = [],
+    kwargs = {},
+    sharedOn,
+  } = opts;
+
+  let inputsHash: string;
+  if (sharedOn !== undefined) {
+    const selected: Record<string, unknown> = {};
+    for (const k of sharedOn) {
+      if (Object.prototype.hasOwnProperty.call(kwargs, k)) selected[k] = kwargs[k];
+    }
+    inputsHash = hashInputs([], selected);
+  } else {
+    inputsHash = hashInputs(args, kwargs);
+  }
+
+  if (scope === IdempotencyScope.SHARED) {
+    // The only scope that omits sessionId, so agents in different sessions
+    // converge on one key. It also omits workflowVersion, because the whole
+    // point is that *different* workflows share the operation and they will
+    // not be on the same version.
+    //
+    // coordinationId is what keeps this from being too wide. Without it two
+    // unrelated callers of sendEmail({to: X}) would deduplicate, suppressing
+    // one of them silently. It is required, and the caller must choose it.
+    if (!coordinationId) {
+      throw new Error(
+        "IdempotencyScope.SHARED requires a coordinationId naming the work being " +
+          "shared: a ticket, task, or tenant id. Pass it when opening the session, " +
+          'e.g. durableTools(config, { coordinationId: "ticket-4417" }, fn). It has ' +
+          "no default: a shared one would deduplicate unrelated callers that happen " +
+          "to make the same call.",
+      );
+    }
+    return `shared:${coordinationId}:${toolName}:${inputsHash}`;
+  }
+
+  let seqPart = "session_wide";
+  let agentPart = "session_wide";
+
+  if (scope === IdempotencyScope.AGENT_PRIVATE) {
+    agentPart = agentId;
+  } else if (scope === IdempotencyScope.STEP_LOCAL) {
+    seqPart = String(stepSequence);
+    agentPart = agentId;
+  }
+
+  return `${sessionId}:${workflowVersion}:${seqPart}:${agentPart}:${toolName}:${inputsHash}`;
+}

package/src/index.ts

@@ -5,3 +5,21 @@ export * from "./cellaflow/v1/idempotency_pb.js";
 // Note: internal_pb is intentionally NOT re-exported here. It contains
 // engine-internal types (CacheRecord, LeaseRecord, etc.) that are not part
 // of the public SDK surface.
+
+// High-level API
+export { IdempotencyScope, deriveIdempotencyKey, hashInputs } from "./idempotency.js";
+export type { DeriveKeyOptions } from "./idempotency.js";
+export { tool, step, DivergentStepError } from "./tool.js";
+export type { ToolOptions } from "./tool.js";
+export { durableTools, toolSessionId } from "./durable.js";
+export type { DurableToolsOptions, DurableSession, ThreadRef } from "./durable.js";
+export {
+  executionLease,
+  taskLease,
+  currentLease,
+  LeaseHandle,
+  LeaseNotAcquired,
+  LeaseLostError,
+} from "./execution.js";
+export type { ExecutionLeaseOptions } from "./execution.js";
+export { WorkflowContext } from "./context.js";

package/src/lease.ts

@@ -0,0 +1,145 @@
+import type { CellaflowClient } from "./client.js";
+import { RenewFailureReason } from "./cellaflow/v1/idempotency_pb.js";
+
+/** Why the engine refused to renew, in words a caller can act on. */
+function renewFailureDetail(reason: RenewFailureReason | undefined): string {
+  switch (reason) {
+    case RenewFailureReason.EXPIRED:
+      return "the lease had already expired";
+    case RenewFailureReason.SUPERSEDED:
+      return "another worker took the lease";
+    case RenewFailureReason.NOT_FOUND:
+      return "the engine has no record of the lease";
+    case RenewFailureReason.COMPLETED:
+      return "the operation was already committed";
+    case RenewFailureReason.MAX_LIFETIME_EXCEEDED:
+      return "the lease hit its maximum lifetime and was reclaimed";
+    default:
+      return "the engine refused renewal";
+  }
+}
+
+export interface LeaseHeartbeatOptions {
+  client: CellaflowClient;
+  agentId: string;
+  idempotencyKey: string;
+  fencingToken: number;
+  /** How often to renew. The engine suggests this on acquisition. */
+  heartbeatIntervalMs: number;
+  /** How long each renewal extends the lease. Defaults to 3x the interval. */
+  extendMs?: number;
+  /**
+   * How long the lease survives without a confirmed renewal. Loss is declared
+   * against this, not against an error count.
+   */
+  leaseTtlMs?: number;
+  /** Called once when the lease is determined to be lost. */
+  onLeaseLost?: (detail: string) => void;
+}
+
+/**
+ * Renews a held lease on an interval until stopped.
+ *
+ * The engine requires a caller holding a lease to renew it, so this starts as
+ * soon as one is acquired and stops in a `finally`.
+ */
+export class LeaseHeartbeat {
+  private readonly opts: Required<Omit<LeaseHeartbeatOptions, "onLeaseLost">> &
+    Pick<LeaseHeartbeatOptions, "onLeaseLost">;
+  private timer?: NodeJS.Timeout;
+  private stopped = false;
+  private lastConfirmed = Date.now();
+  private inFlight?: Promise<void>;
+
+  /** Set once the lease is known to be lost. Read by `LeaseHandle.check()`. */
+  lost = false;
+  lostDetail?: string;
+
+  constructor(options: LeaseHeartbeatOptions) {
+    const intervalMs = options.heartbeatIntervalMs;
+    this.opts = {
+      ...options,
+      heartbeatIntervalMs: intervalMs,
+      extendMs: options.extendMs ?? intervalMs * 3,
+      leaseTtlMs: options.leaseTtlMs ?? intervalMs * 3,
+    };
+  }
+
+  start(): void {
+    if (this.timer) return;
+    this.timer = setInterval(() => {
+      // Never overlap renewals: a slow RPC would otherwise queue them and each
+      // would extend from a stale view of the lease.
+      if (this.inFlight) return;
+      this.inFlight = this.beat().finally(() => {
+        this.inFlight = undefined;
+      });
+    }, this.opts.heartbeatIntervalMs);
+    // Do not hold the event loop open on the heartbeat alone.
+    this.timer.unref?.();
+  }
+
+  async stop(): Promise<void> {
+    this.stopped = true;
+    if (this.timer) {
+      clearInterval(this.timer);
+      this.timer = undefined;
+    }
+    // Let an in-flight renewal settle so it cannot outlive the block that owns
+    // the lease and renew something the caller has already released.
+    if (this.inFlight) await this.inFlight.catch(() => {});
+  }
+
+  private notifyLost(detail: string): void {
+    if (this.lost) return;
+    this.lost = true;
+    this.lostDetail = detail;
+    void this.stop();
+    try {
+      this.opts.onLeaseLost?.(detail);
+    } catch {
+      // A caller's callback must not take down the heartbeat.
+    }
+  }
+
+  /**
+   * True once the lease can no longer be assumed held.
+   *
+   * A failed renewal is not itself proof of loss: the RPC failing means we could
+   * not confirm the lease, while the engine may still hold it until the TTL runs
+   * out. Declaring loss on an error count would abort work that still holds a
+   * perfectly valid lease, so loss is declared on elapsed time since the last
+   * *confirmed* renewal.
+   */
+  private ttlExhausted(): boolean {
+    return Date.now() - this.lastConfirmed >= this.opts.leaseTtlMs;
+  }
+
+  private async beat(): Promise<void> {
+    if (this.stopped || this.lost) return;
+    try {
+      const resp = await this.opts.client.renewLease(
+        this.opts.agentId,
+        this.opts.idempotencyKey,
+        this.opts.fencingToken,
+        this.opts.extendMs,
+        // Bound the RPC. Without a deadline a black-holed connection parks the
+        // renewal forever and the lease silently ages out while we wait on it.
+        Math.max(1, Math.ceil(this.opts.heartbeatIntervalMs / 1000)),
+      );
+      if (!resp.renewed) {
+        this.notifyLost(renewFailureDetail(resp.failureReason));
+        return;
+      }
+      this.lastConfirmed = Date.now();
+    } catch (err) {
+      if (this.ttlExhausted()) {
+        const msg = err instanceof Error ? err.message : String(err);
+        this.notifyLost(
+          `could not be renewed for ${this.opts.leaseTtlMs}ms (last error: ${msg})`,
+        );
+      }
+      // Otherwise: transient. The engine may still be holding it for us.
+    }
+  }
+}