@osolmaz/pi-workflows 0.9.1 → 0.10.0

package/src/workflows/definition.ts

@@ -16,27 +16,137 @@ import type {
   NotifyNodeDefinition,
   ShellActionNodeDefinition,
   WorkflowDefinition,
+  WorkflowExitMap,
+  WorkflowIncludeDefinition,
+  WorkflowIncludedResult,
+  WorkflowIncludeMap,
+  WorkflowInputOf,
+  WorkflowNodeContext,
+  WorkflowNodeDefinition,
+  WorkflowTypedEdge,
+  WorkflowValueParser,
 } from "./types.js";
 
 const WORKFLOW_DEFINITION_BRAND = Symbol.for("pi-workflows.definition");
 
-export function defineWorkflow<TWorkflow extends WorkflowDefinition>(
-  definition: TWorkflow,
-): TWorkflow {
-  assertValidWorkflowDefinitionShape(definition);
-  if (isWorkflowDefinition(definition)) {
-    return definition;
+type WorkflowDefinitionInput<
+  TInput,
+  TNodes extends Record<string, WorkflowNodeDefinition>,
+  TIncludes extends WorkflowIncludeMap,
+  TExits extends WorkflowExitMap,
+> = Omit<
+  WorkflowDefinition<TInput, TExits, TIncludes>,
+  "input" | "nodes" | "includes" | "exits" | "edges"
+> & {
+  input?: WorkflowValueParser<TInput>;
+  nodes: TNodes;
+  includes?: TIncludes;
+  exits?: TExits;
+  edges: WorkflowTypedEdge<TNodes, TIncludes>[];
+};
+
+export function defineWorkflow<
+  TInput = unknown,
+  const TNodes extends Record<string, WorkflowNodeDefinition> = Record<
+    string,
+    WorkflowNodeDefinition
+  >,
+  const TIncludes extends WorkflowIncludeMap = Record<never, never>,
+  const TExits extends WorkflowExitMap = Record<never, never>,
+>(
+  definition: WorkflowDefinitionInput<TInput, TNodes, TIncludes, TExits>,
+): WorkflowDefinition<TInput, TExits, TIncludes> & {
+  nodes: TNodes;
+  includes?: TIncludes;
+  exits?: TExits;
+} {
+  assertValidWorkflowDefinitionShape(definition as WorkflowDefinition);
+  const typed = definition as WorkflowDefinition<TInput, TExits, TIncludes> & {
+    nodes: TNodes;
+    includes?: TIncludes;
+    exits?: TExits;
+  };
+  if (isWorkflowDefinition(typed)) {
+    return typed;
   }
-  Object.defineProperty(definition, WORKFLOW_DEFINITION_BRAND, {
+  Object.defineProperty(typed, WORKFLOW_DEFINITION_BRAND, {
     value: true,
     enumerable: false,
     configurable: false,
     writable: false,
   });
+  return typed;
+}
+
+export function includeWorkflow<TWorkflow extends WorkflowDefinition<any, any, any>>(
+  workflow: TWorkflow,
+  options?: {
+    input?: (
+      context: WorkflowNodeContext,
+    ) => Promise<WorkflowInputOf<TWorkflow>> | WorkflowInputOf<TWorkflow>;
+  },
+): WorkflowIncludeDefinition<TWorkflow>;
+export function includeWorkflow<TWorkflow extends WorkflowDefinition<any, any, any>>(
+  definition: WorkflowIncludeDefinition<TWorkflow>,
+): WorkflowIncludeDefinition<TWorkflow>;
+export function includeWorkflow<TWorkflow extends WorkflowDefinition<any, any, any>>(
+  workflowOrDefinition: TWorkflow | WorkflowIncludeDefinition<TWorkflow>,
+  options: {
+    input?: (
+      context: WorkflowNodeContext,
+    ) => Promise<WorkflowInputOf<TWorkflow>> | WorkflowInputOf<TWorkflow>;
+  } = {},
+): WorkflowIncludeDefinition<TWorkflow> {
+  const definition = isWorkflowDefinition(workflowOrDefinition)
+    ? { workflow: workflowOrDefinition, ...options }
+    : workflowOrDefinition;
+  if (typeof definition.workflow !== "string" && !isWorkflowDefinition(definition.workflow)) {
+    throw new Error("Included workflow must be a defined workflow or a workflow reference");
+  }
+  if (definition.input !== undefined && typeof definition.input !== "function") {
+    throw new Error("Included workflow input must be a function");
+  }
+  if (definition.contract !== undefined && !isWorkflowDefinition(definition.contract)) {
+    throw new Error("Included workflow contract must be defined with defineWorkflow");
+  }
   return definition;
 }
 
-export function isWorkflowDefinition(value: unknown): value is WorkflowDefinition {
+export function includedResult<TWorkflow extends WorkflowDefinition<any, any, any>>(
+  workflow: TWorkflow,
+  value: unknown,
+): WorkflowIncludedResult<TWorkflow> {
+  if (value === null || typeof value !== "object" || Array.isArray(value)) {
+    throw new Error(`Included ${workflow.name} result must be an object`);
+  }
+  const exit = (value as { exit?: unknown }).exit;
+  if (typeof exit !== "string" || !Object.hasOwn(workflow.exits ?? {}, exit)) {
+    throw new Error(`Included ${workflow.name} result has unknown exit ${JSON.stringify(exit)}`);
+  }
+  if (!Object.hasOwn(value, "output")) {
+    throw new Error(`Included ${workflow.name} result requires output`);
+  }
+  return value as WorkflowIncludedResult<TWorkflow>;
+}
+
+/** Preserve exact workflow types while checking duplicate registry names. */
+export function defineWorkflowRegistry<
+  const TRegistry extends Record<string, WorkflowDefinition<any, any, any>>,
+>(registry: TRegistry): Readonly<TRegistry> {
+  const names = new Set<string>();
+  for (const [key, workflow] of Object.entries(registry)) {
+    if (!isWorkflowDefinition(workflow)) {
+      throw new Error(`Workflow registry entry ${key} is not defined with defineWorkflow`);
+    }
+    if (names.has(workflow.name)) {
+      throw new Error(`Workflow registry has duplicate workflow name: ${workflow.name}`);
+    }
+    names.add(workflow.name);
+  }
+  return Object.freeze({ ...registry });
+}
+
+export function isWorkflowDefinition(value: unknown): value is WorkflowDefinition<any, any, any> {
   return (
     value != null &&
     typeof value === "object" &&

package/src/workflows/engine.ts

@@ -1,6 +1,11 @@
-import { randomUUID } from "node:crypto";
+import { createHash, randomUUID } from "node:crypto";
 import { isDeepStrictEqual } from "node:util";
 import { resolveArtifacts } from "./artifacts.js";
+import {
+  compileWorkflowDefinition,
+  compositionMetadata,
+  isCompiledWorkflow,
+} from "./composition.js";
 import {
   CancelledError,
   errorMessage,
@@ -14,7 +19,13 @@ import {
 import { resolveNext, resolveNextForOutcome, validateWorkflowDefinition } from "./graph.js";
 import { extractJsonValue } from "./json.js";
 import { runShellAction, shellResultFromError } from "./shell.js";
-import { RUN_STATE_SCHEMA, WorkflowRunStore, createRunId, readRunBundle } from "./store.js";
+import {
+  RUN_STATE_SCHEMA,
+  WorkflowRunStore,
+  createDefinitionSnapshot,
+  createRunId,
+  readRunBundle,
+} from "./store.js";
 import type {
   AgentNodeDefinition,
   AgentStepExecutor,
@@ -215,10 +226,12 @@ export class WorkflowEngine {
     input: unknown,
     options: { workflowSource?: WorkflowSource; runId?: string } = {},
   ): Promise<WorkflowRunResult> {
+    workflow = isCompiledWorkflow(workflow) ? workflow : compileWorkflowDefinition(workflow);
     validateWorkflowDefinition(workflow);
     // Fail before any bundle exists so bad input cannot leave a partial run
     // on disk or silently change shape when state.json round-trips.
-    const normalizedInput = input === undefined ? null : input;
+    const suppliedInput = input === undefined ? null : input;
+    const normalizedInput = workflow.input ? await workflow.input(suppliedInput) : suppliedInput;
     assertJsonSerializable(normalizedInput, "Workflow run input");
     if (options.runId !== undefined && !/^[A-Za-z0-9][A-Za-z0-9._-]{0,199}$/.test(options.runId)) {
       throw new Error(`Invalid workflow run id: ${JSON.stringify(options.runId)}`);
@@ -270,6 +283,7 @@ export class WorkflowEngine {
     runId: string,
     options: { workflowSource?: WorkflowSource; force?: boolean } = {},
   ): Promise<WorkflowRunResult> {
+    workflow = isCompiledWorkflow(workflow) ? workflow : compileWorkflowDefinition(workflow);
     validateWorkflowDefinition(workflow);
     // Reset before any await: a park or cancel landing during preparation
     // must survive, or a host drain would hang while the run executes.
@@ -279,7 +293,7 @@ export class WorkflowEngine {
     const bundle = await this.store.prepareRunResume(runId);
     const { runDir } = bundle;
     const state = bundle.state;
-    const sourceMismatch = workflowSourceMismatch(state, options.workflowSource);
+    const sourceMismatch = workflowIdentityMismatch(state, workflow, options.workflowSource);
     if (sourceMismatch && options.force !== true) {
       throw new WorkflowSourceChangedError(runId);
     }
@@ -330,7 +344,7 @@ export class WorkflowEngine {
         state,
         runDir,
         point.nodeId,
-        state.steps.length,
+        countExecutableSteps(workflow, state.steps),
         point.lastOutput,
       );
     } catch (error) {
@@ -354,6 +368,7 @@ export class WorkflowEngine {
     input: unknown,
     options: { workflowSource?: WorkflowSource; runId?: string; force?: boolean } = {},
   ): Promise<WorkflowRunResult> {
+    workflow = isCompiledWorkflow(workflow) ? workflow : compileWorkflowDefinition(workflow);
     validateWorkflowDefinition(workflow);
     this.cancelled = false;
     this.paused = false;
@@ -367,12 +382,13 @@ export class WorkflowEngine {
         `Cannot continue workflow run ${parentRunId} with status ${parent.state.status}`,
       );
     }
-    const sourceMismatch = workflowSourceMismatch(parent.state, options.workflowSource);
+    const sourceMismatch = workflowIdentityMismatch(parent.state, workflow, options.workflowSource);
     if (sourceMismatch && options.force !== true) {
       throw new WorkflowSourceChangedError(parentRunId);
     }
 
-    const normalizedInput = input === undefined ? null : input;
+    const suppliedInput = input === undefined ? null : input;
+    const normalizedInput = workflow.input ? await workflow.input(suppliedInput) : suppliedInput;
     assertJsonSerializable(normalizedInput, "Workflow run input");
     if (options.runId !== undefined && !/^[A-Za-z0-9][A-Za-z0-9._-]{0,199}$/.test(options.runId)) {
       throw new Error(`Invalid workflow run id: ${JSON.stringify(options.runId)}`);
@@ -428,7 +444,7 @@ export class WorkflowEngine {
         state,
         runDir,
         point.nodeId,
-        state.steps.length,
+        countExecutableSteps(workflow, state.steps),
         point.lastOutput,
       );
     } catch (error) {
@@ -546,6 +562,7 @@ export class WorkflowEngine {
     runId: string | undefined,
   ): Promise<WorkflowRunState> {
     const now = new Date().toISOString();
+    const composition = compositionMetadata(workflow);
     return {
       schema: RUN_STATE_SCHEMA,
       traceSeq: 0,
@@ -553,6 +570,10 @@ export class WorkflowEngine {
       workflowName: workflow.name,
       ...(await this.resolveTitleBounded(workflow, input)),
       ...(workflowSource !== undefined ? { workflowSource } : {}),
+      ...(composition?.sources.length ? { workflowSources: composition.sources } : {}),
+      ...(composition?.snapshot.mounts.length
+        ? { definitionDigest: definitionDigest(workflow) }
+        : {}),
       startedAt: now,
       updatedAt: now,
       status: "running",
@@ -573,17 +594,24 @@ export class WorkflowEngine {
     initialLastOutput?: unknown,
   ): Promise<void> {
     const maxSteps = workflow.maxSteps ?? this.maxSteps;
+    const composition = compositionMetadata(workflow);
     let currentNodeId: string | null = startNodeId;
     let executedSteps = executedStepsBase;
     let lastOutput: unknown = initialLastOutput;
 
     while (currentNodeId !== null) {
       await this.holdWhilePaused(state, runDir);
-      executedSteps += 1;
-      if (executedSteps > maxSteps) {
-        throw new Error(
-          `Workflow exceeded maxSteps=${maxSteps}; aborting to avoid an unbounded loop`,
-        );
+      const isTransition =
+        composition?.entries[currentNodeId] !== undefined ||
+        composition?.exits[currentNodeId] !== undefined;
+      if (!isTransition) {
+        executedSteps += 1;
+        if (executedSteps > maxSteps) {
+          throw new Error(
+            `Workflow exceeded maxSteps=${maxSteps}; aborting to avoid an unbounded loop`,
+          );
+        }
+        assertInvocationStepLimit(composition, currentNodeId, state.steps);
       }
 
       const node = workflow.nodes[currentNodeId];
@@ -597,7 +625,7 @@ export class WorkflowEngine {
         // as in-flight, and resume reruns it with a fresh attempt.
         throw new RunParkedError();
       }
-      this.recordAttempt(state, attempt);
+      this.recordAttempt(workflow, state, attempt);
       // The terminal node event carries the output, receipt, and conversation
       // linkage so the trace alone is sufficient to reconstruct the run.
       await this.persist(runDir, state, {
@@ -622,6 +650,35 @@ export class WorkflowEngine {
         continue;
       }
 
+      const entered = composition?.entries[attempt.result.nodeId];
+      if (entered !== undefined) {
+        const value = attempt.result.output as { invocation?: number } | undefined;
+        await this.persist(runDir, state, {
+          scope: "run",
+          type: "include_entered",
+          payload: {
+            mountPath: entered.mountPath.split("/"),
+            workflowName: entered.workflowName,
+            invocation: value?.invocation ?? 1,
+          },
+        });
+      }
+      const exited = composition?.exits[attempt.result.nodeId];
+      if (exited !== undefined) {
+        const entrySteps = state.steps.filter((step) => step.nodeId === exited.mountPath);
+        await this.persist(runDir, state, {
+          scope: "run",
+          type: "include_exited",
+          payload: {
+            mountPath: exited.mountPath.split("/"),
+            workflowName: composition?.scopes[exited.mountPath]?.workflowName ?? exited.mountName,
+            invocation: entrySteps.length,
+            exit: exited.exitName,
+            output: attempt.result.output ?? null,
+          },
+        });
+      }
+
       lastOutput = attempt.result.output;
       if (node.nodeType === "checkpoint") {
         await this.finishRun(runDir, state, "waiting", {
@@ -678,13 +735,13 @@ export class WorkflowEngine {
     state: WorkflowRunState,
     attempt: NodeAttempt,
   ): string | null {
+    if (attempt.result.outcome === "cancelled" || this.cancelled) {
+      throw new CancelledError();
+    }
     const next = resolveNextForOutcome(workflow.edges, attempt.result.nodeId, attempt.result);
     if (next !== null) {
       return next;
     }
-    if (attempt.result.outcome === "cancelled" || this.cancelled) {
-      throw new CancelledError();
-    }
     if (attempt.result.outcome === "timed_out") {
       state.status = "timed_out";
     }
@@ -693,7 +750,11 @@ export class WorkflowEngine {
       : new Error(attempt.result.error ?? `Workflow node failed: ${attempt.result.nodeId}`);
   }
 
-  private recordAttempt(state: WorkflowRunState, attempt: NodeAttempt): void {
+  private recordAttempt(
+    workflow: WorkflowDefinition,
+    state: WorkflowRunState,
+    attempt: NodeAttempt,
+  ): void {
     state.results[attempt.result.nodeId] = attempt.result;
     if (attempt.result.outcome === "ok") {
       state.outputs[attempt.result.nodeId] = attempt.result.output;
@@ -702,6 +763,11 @@ export class WorkflowEngine {
       // must not survive next to a non-ok latest result.
       delete state.outputs[attempt.result.nodeId];
     }
+    const exit = compositionMetadata(workflow)?.exits[attempt.result.nodeId];
+    if (exit !== undefined && attempt.result.outcome === "ok") {
+      state.outputs[exit.mountPath] = attempt.result.output;
+      state.results[exit.mountPath] = { ...attempt.result, nodeId: exit.mountPath };
+    }
     const step: WorkflowStepRecord = {
       attemptId: attempt.result.attemptId,
       nodeId: attempt.result.nodeId,
@@ -1234,6 +1300,54 @@ function abortRejection(signal: AbortSignal): Promise<never> {
   });
 }
 
+function countExecutableSteps(workflow: WorkflowDefinition, steps: WorkflowStepRecord[]): number {
+  const metadata = compositionMetadata(workflow);
+  if (metadata === undefined) return steps.length;
+  return steps.filter(
+    (step) =>
+      metadata.entries[step.nodeId] === undefined && metadata.exits[step.nodeId] === undefined,
+  ).length;
+}
+
+function assertInvocationStepLimit(
+  metadata: ReturnType<typeof compositionMetadata>,
+  nodeId: string,
+  steps: WorkflowStepRecord[],
+): void {
+  if (metadata === undefined) return;
+  const scopes = Object.values(metadata.scopes)
+    .filter(
+      (candidate) =>
+        candidate.path !== "" &&
+        candidate.maxSteps !== undefined &&
+        nodeId.startsWith(`${candidate.path}/`),
+    )
+    .sort((a, b) => a.path.length - b.path.length);
+  for (const scope of scopes) {
+    let entryIndex = -1;
+    for (let index = steps.length - 1; index >= 0; index -= 1) {
+      if (steps[index]?.nodeId === scope.path) {
+        entryIndex = index;
+        break;
+      }
+    }
+    if (entryIndex < 0) throw new Error(`Workflow include entry is missing: ${scope.path}`);
+    const attempts = steps
+      .slice(entryIndex + 1)
+      .filter(
+        (step) =>
+          step.nodeId.startsWith(`${scope.path}/`) &&
+          metadata.entries[step.nodeId] === undefined &&
+          metadata.exits[step.nodeId] === undefined,
+      ).length;
+    if (attempts >= (scope.maxSteps as number)) {
+      throw new Error(
+        `Included workflow ${scope.workflowName} at ${scope.path} exceeded maxSteps=${scope.maxSteps}`,
+      );
+    }
+  }
+}
+
 /**
  * Outputs are persisted to the run bundle, so they must be JSON-serializable.
  * Failing here turns a bad callback return value into a normal node failure
@@ -1255,6 +1369,25 @@ function workflowSourceMismatch(
   );
 }
 
+function workflowIdentityMismatch(
+  state: WorkflowRunState,
+  workflow: WorkflowDefinition,
+  source: WorkflowSource | undefined,
+): boolean {
+  if (workflowSourceMismatch(state, source)) return true;
+  const metadata = compositionMetadata(workflow);
+  const currentSources = metadata?.sources ?? [];
+  if (!isDeepStrictEqual(state.workflowSources ?? [], currentSources)) return true;
+  const currentDigest = metadata?.snapshot.mounts.length ? definitionDigest(workflow) : undefined;
+  return state.definitionDigest !== currentDigest;
+}
+
+function definitionDigest(workflow: WorkflowDefinition): string {
+  return `sha256:${createHash("sha256")
+    .update(JSON.stringify(createDefinitionSnapshot(workflow)))
+    .digest("hex")}`;
+}
+
 function assertJsonSerializable(value: unknown, what: string): void {
   let encoded: string | undefined;
   try {

package/src/workflows/graph.ts

@@ -1,3 +1,4 @@
+import { compileWorkflowDefinition, isCompiledWorkflow } from "./composition.js";
 import { assertValidWorkflowDefinitionShape } from "./schema.js";
 import type { WorkflowDefinition, WorkflowEdge, WorkflowNodeResult } from "./types.js";
 
@@ -6,17 +7,18 @@ import type { WorkflowDefinition, WorkflowEdge, WorkflowNodeResult } from "./typ
  * at most one outgoing edge per node.
  */
 export function validateWorkflowDefinition(workflow: WorkflowDefinition): void {
-  assertValidWorkflowDefinitionShape(workflow);
-  if (!Object.hasOwn(workflow.nodes, workflow.startAt)) {
-    throw new Error(`Workflow start node is missing: ${workflow.startAt}`);
+  const executable = isCompiledWorkflow(workflow) ? workflow : compileWorkflowDefinition(workflow);
+  assertValidWorkflowDefinitionShape(executable, { compiled: true });
+  if (!Object.hasOwn(executable.nodes, executable.startAt)) {
+    throw new Error(`Workflow start node is missing: ${executable.startAt}`);
   }
 
   const outgoingEdges = new Set<string>();
-  for (const edge of workflow.edges) {
-    validateWorkflowEdge(workflow, edge, outgoingEdges);
+  for (const edge of executable.edges) {
+    validateWorkflowEdge(executable, edge, outgoingEdges);
   }
 
-  assertAllNodesReachable(workflow);
+  assertAllNodesReachable(executable);
 }
 
 /** Reject nodes that no path from `startAt` can ever reach. */

package/src/workflows/index.ts

@@ -4,11 +4,21 @@ export {
   checkpoint,
   compute,
   defineWorkflow,
+  defineWorkflowRegistry,
+  includeWorkflow,
+  includedResult,
   isWorkflowDefinition,
   notify,
   shell,
 } from "./definition.js";
 export { decision, decisionEdge, type DecisionDefinition } from "./decision.js";
+export {
+  compileWorkflowDefinition,
+  compositionMetadata,
+  isCompiledWorkflow,
+  type CompileWorkflowOptions,
+  type WorkflowCompositionMetadata,
+} from "./composition.js";
 export { WorkflowEngine, appendStepContract } from "./engine.js";
 export { CancelledError, TimeoutError } from "./errors.js";
 export {
@@ -96,9 +106,19 @@ export type {
   ShellActionResult,
   WorkflowActionContext,
   WorkflowActionReceipt,
+  WorkflowCompositionSnapshot,
   WorkflowDefinition,
   WorkflowDefinitionSnapshot,
   WorkflowEdge,
+  WorkflowExitDefinition,
+  WorkflowExitMap,
+  WorkflowExitOutputs,
+  WorkflowIncludeDefinition,
+  WorkflowIncludedResult,
+  WorkflowInputOf,
+  WorkflowMountedSource,
+  WorkflowMountSnapshot,
+  WorkflowValueParser,
   WorkflowEngineOptions,
   WorkflowNodeCommon,
   WorkflowNodeContext,