@osolmaz/pi-workflows 0.9.0 → 0.10.0

package/src/workflows/composition.ts

@@ -0,0 +1,627 @@
+import { createHash } from "node:crypto";
+import { readFileSync } from "node:fs";
+import { fileURLToPath } from "node:url";
+import type {
+  ActionNodeDefinition,
+  AgentNodeDefinition,
+  CheckpointNodeDefinition,
+  ComputeNodeDefinition,
+  FunctionActionNodeDefinition,
+  NotifyNodeDefinition,
+  ShellActionNodeDefinition,
+  WorkflowActionContext,
+  WorkflowCompositionSnapshot,
+  WorkflowDefinition,
+  WorkflowEdge,
+  WorkflowExitDefinition,
+  WorkflowIncludeDefinition,
+  WorkflowMountedSource,
+  WorkflowMountSnapshot,
+  WorkflowNodeContext,
+  WorkflowNodeDefinition,
+  WorkflowNodeResult,
+  WorkflowRunState,
+  WorkflowSource,
+  WorkflowStepRecord,
+} from "./types.js";
+
+const COMPILED_WORKFLOW_BRAND = Symbol.for("pi-workflows.compiled-definition");
+const COMPOSITION_METADATA = Symbol.for("pi-workflows.composition-metadata");
+
+export type WorkflowScopeMetadata = {
+  path: string;
+  workflowName: string;
+  maxSteps?: number;
+  authoredNodes: Record<string, string>;
+  childMounts: Record<string, string>;
+};
+
+export type WorkflowEntryMetadata = {
+  mountPath: string;
+  mountName: string;
+  parentPath: string;
+  workflowName: string;
+  input: WorkflowIncludeDefinition["input"];
+  normalizeInput?: WorkflowDefinition["input"];
+};
+
+export type WorkflowExitMetadata = {
+  nodeId: string;
+  mountPath: string;
+  mountName: string;
+  parentPath: string;
+  exitName: string;
+  fromNode: string;
+  validate?: WorkflowExitDefinition["validate"];
+};
+
+export type WorkflowCompositionMetadata = {
+  snapshot: WorkflowCompositionSnapshot;
+  scopes: Record<string, WorkflowScopeMetadata>;
+  entries: Record<string, WorkflowEntryMetadata>;
+  exits: Record<string, WorkflowExitMetadata>;
+  sources: WorkflowMountedSource[];
+};
+
+export type WorkflowCompositionSourceMap = Map<WorkflowDefinition<any, any, any>, WorkflowSource>;
+
+export type CompileWorkflowOptions = {
+  rootSource?: WorkflowSource;
+  sourceMap?: WorkflowCompositionSourceMap;
+};
+
+export function isCompiledWorkflow(workflow: WorkflowDefinition<any, any, any>): boolean {
+  return (workflow as Record<PropertyKey, unknown>)[COMPILED_WORKFLOW_BRAND] === true;
+}
+
+export function compositionMetadata(
+  workflow: WorkflowDefinition<any, any, any>,
+): WorkflowCompositionMetadata | undefined {
+  return (workflow as Record<PropertyKey, unknown>)[COMPOSITION_METADATA] as
+    | WorkflowCompositionMetadata
+    | undefined;
+}
+
+/**
+ * Compile direct includes into one qualified graph. Dynamic string references
+ * must be resolved to direct definitions before this function is called.
+ */
+export function compileWorkflowDefinition<TWorkflow extends WorkflowDefinition<any, any, any>>(
+  workflow: TWorkflow,
+  options: CompileWorkflowOptions = {},
+): TWorkflow {
+  if (isCompiledWorkflow(workflow)) return workflow;
+
+  validateAuthoredWorkflow(workflow);
+  const nodes: Record<string, WorkflowNodeDefinition> = {};
+  const edges: WorkflowEdge[] = [];
+  const mounts: WorkflowMountSnapshot[] = [];
+  const scopes: Record<string, WorkflowScopeMetadata> = {};
+  const entries: Record<string, WorkflowEntryMetadata> = {};
+  const exits: Record<string, WorkflowExitMetadata> = {};
+  const sources: WorkflowMountedSource[] = [];
+  const activeDefinitions: WorkflowDefinition<any, any, any>[] = [];
+
+  const mountWorkflow = (
+    current: WorkflowDefinition<any, any, any>,
+    scopePath: string,
+    parentPath: string | null,
+    mountName: string | null,
+  ): void => {
+    const cycleAt = activeDefinitions.indexOf(current);
+    if (cycleAt >= 0) {
+      const chain = [...activeDefinitions.slice(cycleAt).map((item) => item.name), current.name];
+      throw new Error(`Workflow include cycle: ${chain.join(" -> ")}`);
+    }
+    activeDefinitions.push(current);
+    validateAuthoredWorkflow(current, scopePath !== "");
+
+    const authoredNodes = Object.fromEntries(
+      Object.keys(current.nodes).map((nodeId) => [nodeId, qualify(scopePath, nodeId)]),
+    );
+    const childMounts = Object.fromEntries(
+      Object.keys(current.includes ?? {}).map((name) => [name, qualify(scopePath, name)]),
+    );
+    scopes[scopePath] = {
+      path: scopePath,
+      workflowName: current.name,
+      ...(current.maxSteps !== undefined ? { maxSteps: current.maxSteps } : {}),
+      authoredNodes,
+      childMounts,
+    };
+
+    if (scopePath !== "" && parentPath !== null && mountName !== null) {
+      const source = options.sourceMap?.get(current) ?? sourceFromDefinition(current);
+      if (source !== undefined) {
+        sources.push({
+          mountPath: scopePath.split("/"),
+          workflowName: current.name,
+          source,
+        });
+      }
+    }
+
+    for (const [localNodeId, node] of Object.entries(current.nodes)) {
+      const qualifiedNodeId = qualify(scopePath, localNodeId);
+      nodes[qualifiedNodeId] = wrapNode(node, scopePath, scopes, entries, exits);
+    }
+
+    for (const [name, include] of Object.entries(current.includes ?? {}) as [
+      string,
+      WorkflowIncludeDefinition,
+    ][]) {
+      const child = typeof include.workflow === "string" ? include.contract : include.workflow;
+      if (child === undefined) {
+        throw new Error(
+          `Workflow include ${qualify(scopePath, name)} is unresolved: ${include.workflow}`,
+        );
+      }
+      const mountPath = qualify(scopePath, name);
+      const childExits = child.exits ?? {};
+      if (Object.keys(childExits).length === 0) {
+        throw new Error(`Included workflow ${child.name} must declare at least one exit`);
+      }
+      const entryNode = mountPath;
+      entries[entryNode] = {
+        mountPath,
+        mountName: name,
+        parentPath: scopePath,
+        workflowName: child.name,
+        input: include.input,
+        normalizeInput: child.input,
+      };
+      nodes[entryNode] = {
+        nodeType: "compute",
+        statusDetail: `entering ${child.name}`,
+        run: async (context) => {
+          const parentContext = projectWorkflowContext(context, scopePath, scopes, entries, exits);
+          const mapped = include.input ? await include.input(parentContext) : parentContext.input;
+          const normalized = child.input ? await child.input(mapped) : mapped;
+          return {
+            schema: "pi-workflows.include-entry.v1",
+            invocation: countCompletedEntries(context.state, entryNode) + 1,
+            input: normalized === undefined ? null : normalized,
+          };
+        },
+      };
+      edges.push({ from: entryNode, to: qualify(mountPath, child.startAt) });
+
+      const exitNodes: Record<string, string> = {};
+      for (const [exitName, exit] of Object.entries(childExits) as [
+        string,
+        WorkflowExitDefinition,
+      ][]) {
+        const exitNodeId = qualify(mountPath, `__piw_exit_${exitName}`);
+        const fromNode = qualify(mountPath, exit.from);
+        exitNodes[exitName] = exitNodeId;
+        exits[exitNodeId] = {
+          nodeId: exitNodeId,
+          mountPath,
+          mountName: name,
+          parentPath: scopePath,
+          exitName,
+          fromNode,
+          validate: exit.validate,
+        };
+        nodes[exitNodeId] = {
+          nodeType: "compute",
+          statusDetail: `leaving ${child.name} through ${exitName}`,
+          run: async (context) => {
+            const childContext = projectWorkflowContext(context, mountPath, scopes, entries, exits);
+            const raw = childContext.outputs[exit.from];
+            const output = exit.validate ? await exit.validate(raw) : raw;
+            return { exit: exitName, output: output === undefined ? null : output };
+          },
+        };
+        edges.push({ from: fromNode, to: exitNodeId });
+      }
+
+      mounts.push({
+        mountPath: mountPath.split("/"),
+        workflowName: child.name,
+        entryNode,
+        exits: exitNodes,
+        ...(child.maxSteps !== undefined ? { maxSteps: child.maxSteps } : {}),
+      });
+      mountWorkflow(child, mountPath, scopePath, name);
+    }
+
+    for (const edge of current.edges) {
+      const from = resolveAuthoredFrom(scopePath, current, edge.from, exits);
+      if ("to" in edge) {
+        edges.push({ from, to: resolveAuthoredTo(scopePath, current, edge.to) });
+      } else {
+        edges.push({
+          from,
+          switch: {
+            on: edge.switch.on,
+            cases: Object.fromEntries(
+              Object.entries(edge.switch.cases).map(([key, target]) => [
+                key,
+                resolveAuthoredTo(scopePath, current, target),
+              ]),
+            ),
+          },
+        });
+      }
+    }
+
+    activeDefinitions.pop();
+  };
+
+  mountWorkflow(workflow, "", null, null);
+  const metadata: WorkflowCompositionMetadata = {
+    snapshot: { mounts: mounts.sort(compareMounts) },
+    scopes,
+    entries,
+    exits,
+    sources: sources.sort(compareSources),
+  };
+  const compiled = {
+    ...workflow,
+    nodes,
+    edges,
+  } as TWorkflow;
+  delete (compiled as { includes?: unknown }).includes;
+  delete (compiled as { exits?: unknown }).exits;
+  for (const symbol of Object.getOwnPropertySymbols(workflow)) {
+    const descriptor = Object.getOwnPropertyDescriptor(workflow, symbol);
+    if (descriptor !== undefined) Object.defineProperty(compiled, symbol, descriptor);
+  }
+  Object.defineProperty(compiled, COMPILED_WORKFLOW_BRAND, {
+    value: true,
+    enumerable: false,
+  });
+  Object.defineProperty(compiled, COMPOSITION_METADATA, {
+    value: metadata,
+    enumerable: false,
+  });
+  return compiled;
+}
+
+function validateAuthoredWorkflow(
+  workflow: WorkflowDefinition<any, any, any>,
+  included = false,
+): void {
+  if (!Object.hasOwn(workflow.nodes, workflow.startAt)) {
+    throw new Error(`Workflow start node is missing: ${workflow.startAt}`);
+  }
+  const includes = workflow.includes ?? {};
+  const exits = workflow.exits ?? {};
+  const outgoing = new Set<string>();
+  for (const edge of workflow.edges) {
+    if (outgoing.has(edge.from)) {
+      throw new Error(`Workflow node must not declare multiple outgoing edges: ${edge.from}`);
+    }
+    outgoing.add(edge.from);
+    assertAuthoredFrom(workflow, edge.from);
+    const targets = "to" in edge ? [edge.to] : Object.values(edge.switch.cases);
+    for (const target of targets) assertAuthoredTo(workflow, target);
+  }
+  const exitNodes = new Set<string>();
+  for (const [exitName, exit] of Object.entries(exits) as [string, WorkflowExitDefinition][]) {
+    if (!Object.hasOwn(workflow.nodes, exit.from)) {
+      throw new Error(`Workflow exit ${exitName} references unknown node: ${exit.from}`);
+    }
+    if (outgoing.has(exit.from)) {
+      throw new Error(`Workflow exit ${exitName} must reference a terminal node: ${exit.from}`);
+    }
+    if (exitNodes.has(exit.from)) {
+      throw new Error(`Workflow terminal node has more than one exit: ${exit.from}`);
+    }
+    exitNodes.add(exit.from);
+  }
+  if (included) {
+    const undeclared = Object.keys(workflow.nodes).filter(
+      (nodeId) => !outgoing.has(nodeId) && !exitNodes.has(nodeId),
+    );
+    if (undeclared.length > 0) {
+      throw new Error(
+        `Included workflow ${workflow.name} has terminal nodes without named exits: ${undeclared.join(", ")}`,
+      );
+    }
+  }
+
+  const reachable = new Set<string>([workflow.startAt]);
+  const queue = [workflow.startAt];
+  while (queue.length > 0) {
+    const from = queue.shift() as string;
+    const edge = workflow.edges.find((candidate) => candidate.from === from);
+    if (edge !== undefined) {
+      const targets = "to" in edge ? [edge.to] : Object.values(edge.switch.cases);
+      for (const target of targets) addReachable(target, reachable, queue);
+    }
+    if (Object.hasOwn(includes, from)) {
+      const include = includes[from];
+      const child =
+        include && typeof include.workflow !== "string" ? include.workflow : include?.contract;
+      for (const exitName of Object.keys(child?.exits ?? {})) {
+        addReachable(`${from}.${exitName}`, reachable, queue);
+      }
+    }
+  }
+  const unreachable = [...Object.keys(workflow.nodes), ...Object.keys(includes)].filter(
+    (id) => !reachable.has(id),
+  );
+  if (unreachable.length > 0) {
+    throw new Error(`Workflow has unreachable nodes: ${unreachable.join(", ")}`);
+  }
+}
+
+function assertAuthoredFrom(workflow: WorkflowDefinition<any, any, any>, from: string): void {
+  if (Object.hasOwn(workflow.nodes, from)) return;
+  const parsed = parseExitReference(from);
+  const include = parsed ? workflow.includes?.[parsed.mount] : undefined;
+  const child =
+    include && typeof include.workflow !== "string" ? include.workflow : include?.contract;
+  if (parsed && child?.exits && Object.hasOwn(child.exits, parsed.exit)) return;
+  throw new Error(`Workflow edge references unknown from-node or include exit: ${from}`);
+}
+
+function assertAuthoredTo(workflow: WorkflowDefinition<any, any, any>, to: string): void {
+  if (Object.hasOwn(workflow.nodes, to) || Object.hasOwn(workflow.includes ?? {}, to)) return;
+  throw new Error(`Workflow edge references unknown to-node: ${to}`);
+}
+
+function resolveAuthoredFrom(
+  scopePath: string,
+  workflow: WorkflowDefinition<any, any, any>,
+  from: string,
+  exitMetadata: Record<string, WorkflowExitMetadata>,
+): string {
+  if (Object.hasOwn(workflow.nodes, from)) return qualify(scopePath, from);
+  const parsed = parseExitReference(from);
+  if (!parsed) throw new Error(`Invalid include exit reference: ${from}`);
+  const mountPath = qualify(scopePath, parsed.mount);
+  const match = Object.values(exitMetadata).find(
+    (exit) => exit.mountPath === mountPath && exit.exitName === parsed.exit,
+  );
+  if (!match) throw new Error(`Unknown include exit reference: ${from}`);
+  return match.nodeId;
+}
+
+function resolveAuthoredTo(
+  scopePath: string,
+  workflow: WorkflowDefinition<any, any, any>,
+  target: string,
+): string {
+  if (Object.hasOwn(workflow.nodes, target) || Object.hasOwn(workflow.includes ?? {}, target)) {
+    return qualify(scopePath, target);
+  }
+  throw new Error(`Unknown workflow target: ${target}`);
+}
+
+function parseExitReference(value: string): { mount: string; exit: string } | undefined {
+  const dot = value.indexOf(".");
+  if (dot <= 0 || dot === value.length - 1 || value.indexOf(".", dot + 1) >= 0) return undefined;
+  return { mount: value.slice(0, dot), exit: value.slice(dot + 1) };
+}
+
+function addReachable(value: string, reachable: Set<string>, queue: string[]): void {
+  if (reachable.has(value)) return;
+  reachable.add(value);
+  queue.push(value);
+}
+
+function qualify(scopePath: string, localId: string): string {
+  return scopePath === "" ? localId : `${scopePath}/${localId}`;
+}
+
+function compareMounts(a: WorkflowMountSnapshot, b: WorkflowMountSnapshot): number {
+  return a.mountPath.join("/").localeCompare(b.mountPath.join("/"));
+}
+
+function compareSources(a: WorkflowMountedSource, b: WorkflowMountedSource): number {
+  return a.mountPath.join("/").localeCompare(b.mountPath.join("/"));
+}
+
+function sourceFromDefinition(
+  workflow: WorkflowDefinition<any, any, any>,
+): WorkflowSource | undefined {
+  if (workflow.source === undefined) return undefined;
+  let filePath: string;
+  try {
+    filePath = fileURLToPath(workflow.source);
+  } catch {
+    throw new Error(`Workflow ${workflow.name} source must be a file URL: ${workflow.source}`);
+  }
+  const hash = createHash("sha256").update(readFileSync(filePath)).digest("hex");
+  return { kind: "file", path: filePath, hash };
+}
+
+function wrapNode(
+  node: WorkflowNodeDefinition,
+  scopePath: string,
+  scopes: Record<string, WorkflowScopeMetadata>,
+  entries: Record<string, WorkflowEntryMetadata>,
+  exits: Record<string, WorkflowExitMetadata>,
+): WorkflowNodeDefinition {
+  if (scopePath === "" && Object.keys(scopes[scopePath]?.childMounts ?? {}).length === 0) {
+    return node;
+  }
+  const project = (context: WorkflowNodeContext) =>
+    projectWorkflowContext(context, scopePath, scopes, entries, exits);
+  const timeoutResolver = node.timeoutMs;
+  const timeoutMs =
+    typeof timeoutResolver === "function"
+      ? (context: WorkflowNodeContext) => timeoutResolver(project(context))
+      : timeoutResolver;
+  const common = {
+    ...(timeoutMs !== undefined ? { timeoutMs } : {}),
+    ...(node.statusDetail !== undefined ? { statusDetail: node.statusDetail } : {}),
+  };
+  switch (node.nodeType) {
+    case "agent":
+      return {
+        ...common,
+        nodeType: "agent",
+        prompt: (context) => (node as AgentNodeDefinition).prompt(project(context)),
+        ...(node.expectedOutput !== undefined ? { expectedOutput: node.expectedOutput } : {}),
+        ...(node.validate !== undefined
+          ? {
+              validate: (output: unknown, context: WorkflowNodeContext) =>
+                node.validate?.(output, project(context)),
+            }
+          : {}),
+      };
+    case "compute":
+      return {
+        ...common,
+        nodeType: "compute",
+        run: (context) => (node as ComputeNodeDefinition).run(project(context)),
+      };
+    case "notify":
+      return {
+        ...common,
+        nodeType: "notify",
+        message: (context) => (node as NotifyNodeDefinition).message(project(context)),
+        ...(node.kind !== undefined ? { kind: node.kind } : {}),
+      };
+    case "checkpoint":
+      return {
+        ...common,
+        nodeType: "checkpoint",
+        ...(node.summary !== undefined ? { summary: node.summary } : {}),
+        ...(node.run !== undefined
+          ? { run: (context: WorkflowNodeContext) => node.run?.(project(context)) }
+          : {}),
+      } as CheckpointNodeDefinition;
+    case "action":
+      if ("exec" in node) {
+        const shell = node as ShellActionNodeDefinition;
+        return {
+          ...common,
+          nodeType: "action",
+          exec: (context) => shell.exec(project(context)),
+          ...(shell.parse !== undefined
+            ? {
+                parse: (result, context) => shell.parse?.(result, project(context)),
+              }
+            : {}),
+          ...(shell.updates !== undefined
+            ? {
+                updates: {
+                  ...(shell.updates.streams !== undefined
+                    ? { streams: [...shell.updates.streams] }
+                    : {}),
+                  parseLine: (line, context) =>
+                    shell.updates?.parseLine(
+                      line,
+                      projectActionContext(context, scopePath, scopes, entries, exits),
+                    ),
+                },
+              }
+            : {}),
+        } as ShellActionNodeDefinition;
+      }
+      return {
+        ...common,
+        nodeType: "action",
+        run: (context) =>
+          (node as FunctionActionNodeDefinition).run(
+            projectActionContext(context, scopePath, scopes, entries, exits),
+          ),
+      } as ActionNodeDefinition;
+  }
+}
+
+function projectActionContext(
+  context: WorkflowActionContext,
+  scopePath: string,
+  scopes: Record<string, WorkflowScopeMetadata>,
+  entries: Record<string, WorkflowEntryMetadata>,
+  exits: Record<string, WorkflowExitMetadata>,
+): WorkflowActionContext {
+  return {
+    ...projectWorkflowContext(context, scopePath, scopes, entries, exits),
+    publishUpdate: context.publishUpdate,
+  };
+}
+
+export function projectWorkflowContext(
+  context: WorkflowNodeContext,
+  scopePath: string,
+  scopes: Record<string, WorkflowScopeMetadata>,
+  _entries: Record<string, WorkflowEntryMetadata>,
+  exits: Record<string, WorkflowExitMetadata>,
+): WorkflowNodeContext {
+  const scope = scopes[scopePath];
+  if (!scope) throw new Error(`Workflow scope is missing: ${scopePath || "root"}`);
+  const entryIndex = scopePath === "" ? -1 : latestStepIndex(context.state.steps, scopePath);
+  if (scopePath !== "" && entryIndex < 0) {
+    throw new Error(`Workflow include entry is missing: ${scopePath}`);
+  }
+  const entry =
+    scopePath === ""
+      ? undefined
+      : (context.state.steps[entryIndex]?.output as
+          | { schema?: string; input?: unknown }
+          | undefined);
+  const outputs: Record<string, unknown> = {};
+  const results: Record<string, WorkflowNodeResult> = {};
+  for (const [local, qualified] of Object.entries(scope.authoredNodes)) {
+    if (latestStepIndex(context.state.steps, qualified) <= entryIndex) continue;
+    if (Object.hasOwn(context.state.outputs, qualified))
+      outputs[local] = context.state.outputs[qualified];
+    const result = context.state.results[qualified];
+    if (result !== undefined) results[local] = result;
+  }
+  for (const [local, mountPath] of Object.entries(scope.childMounts)) {
+    const latestExit = Object.values(exits)
+      .filter((candidate) => candidate.mountPath === mountPath)
+      .map((candidate) => latestStepIndex(context.state.steps, candidate.nodeId))
+      .reduce((max, value) => Math.max(max, value), -1);
+    if (latestExit <= entryIndex) continue;
+    if (Object.hasOwn(context.state.outputs, mountPath))
+      outputs[local] = context.state.outputs[mountPath];
+    const result = context.state.results[mountPath];
+    if (result !== undefined) results[local] = result;
+  }
+  const localSteps = context.state.steps.slice(entryIndex + 1).flatMap((step) => {
+    if (Object.values(scope.authoredNodes).includes(step.nodeId)) {
+      return [{ ...step, nodeId: localNameFor(scope, step.nodeId) }];
+    }
+    const exit = exits[step.nodeId];
+    if (exit?.parentPath === scopePath) {
+      return [{ ...step, nodeId: exit.mountName }];
+    }
+    return [];
+  });
+  const currentNode = context.state.currentNode
+    ? localNameFor(scope, context.state.currentNode)
+    : undefined;
+  const localInput = scopePath === "" ? context.input : (entry?.input ?? null);
+  const state: WorkflowRunState = {
+    ...context.state,
+    input: localInput,
+    outputs,
+    results,
+    steps: localSteps,
+    ...(currentNode !== undefined ? { currentNode } : {}),
+  };
+  return {
+    input: localInput,
+    outputs,
+    results,
+    state,
+    signal: context.signal,
+  };
+}
+
+function localNameFor(scope: WorkflowScopeMetadata, qualified: string): string {
+  for (const [local, global] of Object.entries(scope.authoredNodes)) {
+    if (global === qualified) return local;
+  }
+  return qualified;
+}
+
+function latestStepIndex(steps: WorkflowStepRecord[], nodeId: string): number {
+  for (let index = steps.length - 1; index >= 0; index -= 1) {
+    if (steps[index]?.nodeId === nodeId) return index;
+  }
+  return -1;
+}
+
+function countCompletedEntries(state: WorkflowRunState, entryNode: string): number {
+  return state.steps.filter((step) => step.nodeId === entryNode && step.outcome === "ok").length;
+}

package/src/workflows/decision.ts

@@ -1,6 +1,6 @@
 import { agent } from "./definition.js";
 import { extractJsonValue } from "./json.js";
-import type { AgentNodeDefinition, WorkflowEdge, WorkflowNodeContext } from "./types.js";
+import type { AgentNodeDefinition, WorkflowNodeContext } from "./types.js";
 
 const DEFAULT_FIELD = "route";
 const SIMPLE_FIELD_PATTERN = /^[A-Za-z_][A-Za-z0-9_]*$/;
@@ -64,12 +64,19 @@ export function decision<TChoice extends string>(
  * Build the matching `switch` edge for a `decision` node. Typing `cases` as
  * `Record<TChoice, string>` makes a missing case a compile error.
  */
-export function decisionEdge<TChoice extends string>(args: {
-  from: string;
+export function decisionEdge<
+  TChoice extends string,
+  const TFrom extends string,
+  const TCases extends Record<TChoice, string>,
+>(args: {
+  from: TFrom;
   choices: readonly TChoice[];
   field?: string;
-  cases: Record<TChoice, string>;
-}): WorkflowEdge {
+  cases: TCases;
+}): {
+  from: TFrom;
+  switch: { on: string; cases: TCases };
+} {
   const field = normalizeField(args.field);
   assertValidChoices(args.choices);
   for (const choice of args.choices) {