@bastani/atomic 0.9.11-alpha.4 → 0.9.11-alpha.6

package/dist/builtin/workflows/src/engine/run-tool-execution-tracker.ts

@@ -0,0 +1,152 @@
+export interface AdmittedToolFailure {
+  readonly admissionOrder: number;
+  readonly error: unknown;
+  readonly nodeId?: string;
+}
+
+export type AdmittedToolExecutionAdmission =
+  | { readonly accepted: true; bindNode(nodeId: string): void }
+  | { readonly accepted: false; readonly error: Error; bindNode(nodeId: string): void };
+
+export interface AdmittedToolExecutionTracker {
+  track<T>(execution: Promise<T>): AdmittedToolExecutionAdmission;
+  observeFailure(error: unknown, nodeId: string): void;
+  closeAndDrain(): Promise<void>;
+  abandonNonFailed(): readonly string[];
+  firstFailure(): AdmittedToolFailure | undefined;
+  uniqueFailureFor(error: unknown): AdmittedToolFailure | undefined;
+}
+
+export interface AdmittedToolExecutionTrackerOptions {
+  readonly onFailureObserved?: (failure: AdmittedToolFailure) => void;
+  readonly onFailureDuringDrain?: (failure: AdmittedToolFailure) => void;
+}
+
+/** Tracks each logical ctx.tool promise without changing the promise returned to authors. */
+export function createAdmittedToolExecutionTracker(
+  options: AdmittedToolExecutionTrackerOptions = {},
+): AdmittedToolExecutionTracker {
+  const inFlight = new Set<Promise<void>>();
+  const admissions: Array<{
+    readonly admissionOrder: number;
+    nodeId?: string;
+    failure?: AdmittedToolFailure;
+    settled: boolean;
+    observed?: Promise<void>;
+  }> = [];
+  const failures: AdmittedToolFailure[] = [];
+  let nextAdmissionOrder = 0;
+  let state: "OPEN" | "DRAINING" | "CLOSED" = "OPEN";
+  let closing: Promise<void> | undefined;
+  const abandonment = Promise.withResolvers<void>();
+  let abandoned = false;
+  let drainFailureReported = false;
+
+  const recordFailure = (
+    admission: (typeof admissions)[number],
+    error: unknown,
+    nodeId?: string,
+  ): void => {
+    if (admission.failure !== undefined) return;
+    const failure = { admissionOrder: admission.admissionOrder, error, ...(nodeId === undefined ? {} : { nodeId }) };
+    admission.failure = failure;
+    failures.push(failure);
+    options.onFailureObserved?.(failure);
+  };
+
+  const reportDrainFailure = (): void => {
+    const failure = failures[0];
+    if (state !== "DRAINING" || failure === undefined || drainFailureReported) return;
+    drainFailureReported = true;
+    options.onFailureDuringDrain?.(failure);
+  };
+
+  const drainRecordedFailures = (): Promise<void[]> => Promise.all(admissions.flatMap((admission) =>
+    admission.failure !== undefined && admission.observed !== undefined ? [admission.observed] : []
+  ));
+
+  const closeAtFixedPoint = async (): Promise<void> => {
+    while (!abandoned) {
+      const current = [...inFlight];
+      if (current.length > 0) {
+        const drained = await Promise.race([
+          Promise.all(current).then(() => true),
+          abandonment.promise.then(() => false),
+        ]);
+        if (!drained) break;
+      }
+      // Let author settlement continuations admit follow-up tools before the
+      // fixed-point check. The tracker observer is attached before return.
+      await Promise.resolve();
+      if (inFlight.size === 0) {
+        state = "CLOSED";
+        return;
+      }
+    }
+    await drainRecordedFailures();
+  };
+
+  return {
+    track<T>(execution: Promise<T>): AdmittedToolExecutionAdmission {
+      if (state === "CLOSED") {
+        const error = new Error("atomic-workflows: ctx.tool admission is closed for this run");
+        void execution.catch(() => undefined);
+        return { accepted: false, error, bindNode(): void {} };
+      }
+      const admission: (typeof admissions)[number] = {
+        admissionOrder: ++nextAdmissionOrder,
+        nodeId: undefined,
+        failure: undefined,
+        settled: false,
+      };
+      admissions.push(admission);
+      const observed = execution.then(
+        () => undefined,
+        (error: unknown) => {
+          recordFailure(admission, error, admission.nodeId);
+          reportDrainFailure();
+        },
+      );
+      admission.observed = observed;
+      inFlight.add(observed);
+      void observed.then(() => {
+        admission.settled = true;
+        inFlight.delete(observed);
+      });
+      return { accepted: true, bindNode(id: string): void { admission.nodeId = id; } };
+    },
+    observeFailure(error: unknown, nodeId: string): void {
+      const admission = admissions.find((candidate) => candidate.nodeId === nodeId);
+      if (admission !== undefined) {
+        recordFailure(admission, error, nodeId);
+        reportDrainFailure();
+      }
+    },
+    closeAndDrain(): Promise<void> {
+      if (closing !== undefined) return closing;
+      state = "DRAINING";
+      reportDrainFailure();
+      closing = closeAtFixedPoint();
+      return closing;
+    },
+    abandonNonFailed(): readonly string[] {
+      if (!abandoned) {
+        abandoned = true;
+        state = "CLOSED";
+        abandonment.resolve();
+      }
+      return admissions.flatMap((admission) =>
+        !admission.settled && admission.failure === undefined && admission.nodeId !== undefined
+          ? [admission.nodeId]
+          : []
+      );
+    },
+    firstFailure(): AdmittedToolFailure | undefined {
+      return failures[0];
+    },
+    uniqueFailureFor(error: unknown): AdmittedToolFailure | undefined {
+      const matching = failures.filter((failure) => Object.is(failure.error, error));
+      return matching.length === 1 ? matching[0] : undefined;
+    },
+  };
+}

package/dist/builtin/workflows/src/engine/run-tool-node-lifecycle.ts

@@ -0,0 +1,87 @@
+import { createToolPrimitive, type CreateToolPrimitiveInput } from "../durable/tool-primitive.js";
+import type { DurableWorkflowBackend } from "../durable/backend.js";
+import type { WorkflowToolPrimitive } from "../shared/types.js";
+import { unknownErrorMessage } from "../runs/foreground/executor-abort.js";
+import type { Store } from "../shared/store.js";
+import type { RunSnapshot, ToolNodeSnapshot } from "../shared/store-types.js";
+import type { GraphFrontierTracker } from "./graph-inference.js";
+import { durableRunTopology } from "./run-durable-topology.js";
+import { createAdmittedToolExecutionTracker, type AdmittedToolExecutionTracker } from "./run-tool-execution-tracker.js";
+import type { RunTerminalEventArbiter } from "./run-terminal-event.js";
+
+type ToolNodeLifecycle = Pick<
+  CreateToolPrimitiveInput,
+  "onNodeStart" | "onNodeRunning" | "onNodeEnd" | "onNodeSettle" | "runTopology"
+>;
+
+export function createToolNodeLifecycle(input: {
+  readonly store: Store;
+  readonly tracker: GraphFrontierTracker;
+  readonly run: RunSnapshot;
+  readonly sourceToReplayedNodeIds: Map<string, string>;
+}): ToolNodeLifecycle {
+  const { store, tracker, run, sourceToReplayedNodeIds } = input;
+  return {
+    onNodeStart: (node) => {
+      const inferredParents = tracker.onSpawn(node.id, node.name);
+      const sourceParents = node.replayed === true && node.topologyState !== "unavailable"
+        ? node.parentIds
+        : undefined;
+      const restored = sourceParents?.map((sourceId) => sourceToReplayedNodeIds.get(sourceId));
+      const parentIds = restored !== undefined && restored.every((id): id is string => id !== undefined)
+        ? restored
+        : inferredParents;
+      tracker.replaceParents(node.id, parentIds);
+      (node as ToolNodeSnapshot & { parentIds: readonly string[] }).parentIds = Object.freeze([...parentIds]);
+      sourceToReplayedNodeIds.set(node.id, node.id);
+      store.recordToolNodeStart(run.id, node);
+    },
+    onNodeRunning: (nodeId, startedAt) => { store.recordToolNodeRunning(run.id, nodeId, startedAt); },
+    onNodeEnd: (nodeId, update) => { store.recordToolNodeEnd(run.id, nodeId, update); },
+    onNodeSettle: (nodeId) => { tracker.onSettle(nodeId); },
+    runTopology: durableRunTopology(run),
+  };
+}
+
+/** Wire durable tool execution to terminal-event arbitration and graph state. */
+export function createTrackedToolPrimitive(input: {
+  readonly workflowId: string;
+  readonly backend: DurableWorkflowBackend;
+  readonly nextCheckpointId: () => string;
+  readonly controller: AbortController;
+  readonly terminalEvents: RunTerminalEventArbiter;
+  readonly store: Store;
+  readonly tracker: GraphFrontierTracker;
+  readonly run: RunSnapshot;
+  readonly sourceToReplayedNodeIds: Map<string, string>;
+}): { readonly tool: WorkflowToolPrimitive; readonly admittedTools: AdmittedToolExecutionTracker } {
+  const admittedTools = createAdmittedToolExecutionTracker({
+    onFailureObserved: ({ error, nodeId }) => { input.terminalEvents.selectFailure(error, nodeId); },
+    onFailureDuringDrain: (failure) => {
+      if (input.terminalEvents.winner()?.kind === "failure") input.controller.abort(failure.error);
+    },
+  });
+  const lifecycle = createToolNodeLifecycle(input);
+  const tool = createToolPrimitive({
+    workflowId: input.workflowId,
+    backend: input.backend,
+    onFailureObserved: admittedTools.observeFailure,
+    nextCheckpointId: input.nextCheckpointId,
+    throwIfCancelled: () => {
+      if (input.controller.signal.aborted) throw new Error("atomic-workflows: workflow cancelled");
+    },
+    signal: input.controller.signal,
+    trackExecution: admittedTools.track,
+    ...lifecycle,
+  });
+  input.controller.signal.addEventListener("abort", () => {
+    if (input.terminalEvents.winner()?.kind !== "failure") return;
+    const endedAt = Date.now();
+    const error = unknownErrorMessage(input.controller.signal.reason);
+    for (const nodeId of admittedTools.abandonNonFailed()) {
+      lifecycle.onNodeEnd?.(nodeId, { status: "cancelled", endedAt, error });
+      lifecycle.onNodeSettle?.(nodeId);
+    }
+  }, { once: true });
+  return { tool, admittedTools };
+}