ai-workflow-logic 0.1.3 → 0.1.5

package/src/checkpoint-sidecar.ts

@@ -0,0 +1,419 @@
+import fs from 'node:fs';
+import path from 'node:path';
+import { createHash, randomUUID } from 'node:crypto';
+import type { AIWorkflowFlowRunCheckpoint, FlowClosedValue } from 'ai-workflow-contract';
+import type {
+  DefinitionStepDigest,
+  ImmutableStepRecord,
+  StepMerkleNode,
+  StepMerkleTree,
+} from 'flow-step-space-contract';
+import {
+  buildStepMerkleTree,
+  createImmutableStepRecord,
+  diffStepMerkleTrees,
+} from 'flow-step-space-logic';
+import { normalizeFlowRunCheckpoint } from './run-checkpoint.js';
+
+export interface FlowRunCheckpointHead {
+  readonly schemaVersion: 'depa.flow-run-checkpoint-head/v1';
+  readonly instanceId: string;
+  readonly runId: string;
+  readonly definition: AIWorkflowFlowRunCheckpoint['definition'];
+  readonly version: number;
+  readonly logicalDigest: DefinitionStepDigest;
+  readonly runStepStateRootDigest: DefinitionStepDigest;
+}
+
+export interface FlowRunCheckpointSidecarReceipt {
+  readonly schemaVersion: 'depa.flow-run-checkpoint-sidecar-receipt/v1';
+  readonly version: number;
+  readonly previousRootDigest: DefinitionStepDigest | null;
+  readonly rootDigest: DefinitionStepDigest;
+  readonly logicalDigest: DefinitionStepDigest;
+  readonly addedRecordDigests: readonly DefinitionStepDigest[];
+  readonly addedTreeNodeDigests: readonly DefinitionStepDigest[];
+}
+
+type EncodedCheckpoint = {
+  readonly head: FlowRunCheckpointHead;
+  readonly records: readonly ImmutableStepRecord[];
+  readonly tree: StepMerkleTree;
+  readonly receipt: FlowRunCheckpointSidecarReceipt;
+};
+
+type Marker =
+  | { readonly kind: 'object' }
+  | { readonly kind: 'array'; readonly length: number }
+  | { readonly kind: 'value'; readonly value: FlowClosedValue };
+
+const NO_FOLLOW = fs.constants.O_NOFOLLOW ?? 0;
+
+function compareCodeUnits(left: string, right: string): number {
+  return left < right ? -1 : left > right ? 1 : 0;
+}
+
+function canonicalJson(value: unknown): string {
+  if (value === null || typeof value !== 'object') return JSON.stringify(value);
+  if (Array.isArray(value)) return `[${value.map(canonicalJson).join(',')}]`;
+  const record = value as Readonly<Record<string, unknown>>;
+  return `{${Object.keys(record).sort(compareCodeUnits).map((key) => `${JSON.stringify(key)}:${canonicalJson(record[key])}`).join(',')}}`;
+}
+
+function digest(value: unknown): DefinitionStepDigest {
+  return `sha256:${createHash('sha256').update(canonicalJson(value)).digest('hex')}`;
+}
+
+function digestFileName(value: DefinitionStepDigest): string {
+  if (!/^sha256:[0-9a-f]{64}$/.test(value)) throw new Error(`Invalid sidecar digest ${value}`);
+  return `${value.slice('sha256:'.length)}.json`;
+}
+
+function flatten(value: FlowClosedValue): ImmutableStepRecord[] {
+  const records: ImmutableStepRecord[] = [];
+  const visit = (current: FlowClosedValue, segments: readonly string[]): void => {
+    const key = JSON.stringify(segments);
+    if (Array.isArray(current)) {
+      records.push(createImmutableStepRecord(key, { kind: 'array', length: current.length }));
+      current.forEach((item, index) => visit(item, [...segments, String(index)]));
+      return;
+    }
+    if (current !== null && typeof current === 'object') {
+      records.push(createImmutableStepRecord(key, { kind: 'object' }));
+      const object = current as Readonly<Record<string, FlowClosedValue>>;
+      Object.keys(object).sort(compareCodeUnits).forEach((name) => visit(object[name], [...segments, name]));
+      return;
+    }
+    records.push(createImmutableStepRecord(key, { kind: 'value', value: current }));
+  };
+  visit(value, []);
+  return records;
+}
+
+function parseMarker(record: ImmutableStepRecord): Marker {
+  const value = record.value;
+  if (value === null || typeof value !== 'object' || Array.isArray(value)) throw new Error('Sidecar marker must be a closed object.');
+  const marker = value as Readonly<Record<string, FlowClosedValue>>;
+  const keys = Object.keys(marker).sort(compareCodeUnits);
+  if (marker.kind === 'object' && canonicalJson(keys) === canonicalJson(['kind'])) return { kind: 'object' };
+  if (marker.kind === 'array' && canonicalJson(keys) === canonicalJson(['kind', 'length'])
+    && typeof marker.length === 'number' && Number.isSafeInteger(marker.length) && marker.length >= 0) {
+    return { kind: 'array', length: marker.length };
+  }
+  if (marker.kind === 'value' && canonicalJson(keys) === canonicalJson(['kind', 'value'])) {
+    return { kind: 'value', value: marker.value };
+  }
+  throw new Error(`Invalid sidecar marker for ${record.key}`);
+}
+
+function decodePath(key: string): readonly string[] {
+  const value = JSON.parse(key) as unknown;
+  if (!Array.isArray(value) || value.some((segment) => typeof segment !== 'string')) {
+    throw new Error(`Invalid sidecar record key ${key}`);
+  }
+  return value;
+}
+
+function reconstruct(records: readonly ImmutableStepRecord[]): AIWorkflowFlowRunCheckpoint {
+  const entries = records.map((record) => ({ record, path: decodePath(record.key), marker: parseMarker(record) }))
+    .sort((left, right) => left.path.length - right.path.length || compareCodeUnits(left.record.key, right.record.key));
+  if (entries.length === 0 || entries[0].path.length !== 0) throw new Error('Sidecar root record is missing.');
+  const materialize = (marker: Marker): unknown => marker.kind === 'object'
+    ? {}
+    : marker.kind === 'array'
+      ? new Array(marker.length)
+      : marker.value;
+  let root = materialize(entries[0].marker);
+  for (const entry of entries.slice(1)) {
+    if (entry.path.length === 0) throw new Error('Sidecar root record is duplicated.');
+    let parent: unknown = root;
+    for (const segment of entry.path.slice(0, -1)) {
+      if (parent === null || typeof parent !== 'object') throw new Error(`Sidecar parent ${segment} is not a container.`);
+      parent = (parent as Record<string, unknown>)[segment];
+    }
+    if (parent === null || typeof parent !== 'object') throw new Error('Sidecar record parent is not a container.');
+    const segment = entry.path.at(-1)!;
+    if (Array.isArray(parent)) {
+      if (!/^(0|[1-9][0-9]*)$/.test(segment) || Number(segment) >= parent.length) throw new Error('Sidecar array index is invalid.');
+      parent[Number(segment)] = materialize(entry.marker);
+    } else {
+      if (Object.prototype.hasOwnProperty.call(parent, segment)) throw new Error('Sidecar object key is duplicated.');
+      (parent as Record<string, unknown>)[segment] = materialize(entry.marker);
+    }
+  }
+  const normalized = normalizeFlowRunCheckpoint(root);
+  return normalized;
+}
+
+function exactKeys(value: unknown, expected: readonly string[], label: string): Record<string, unknown> {
+  if (value === null || typeof value !== 'object' || Array.isArray(value)) throw new Error(`${label} must be an object.`);
+  const keys = Object.keys(value).sort(compareCodeUnits);
+  const wanted = [...expected].sort(compareCodeUnits);
+  if (canonicalJson(keys) !== canonicalJson(wanted)) throw new Error(`${label} fields are invalid.`);
+  return value as Record<string, unknown>;
+}
+
+function normalizeHead(value: unknown): FlowRunCheckpointHead {
+  const record = exactKeys(value, [
+    'schemaVersion', 'instanceId', 'runId', 'definition', 'version', 'logicalDigest', 'runStepStateRootDigest',
+  ], 'checkpoint head');
+  if (record.schemaVersion !== 'depa.flow-run-checkpoint-head/v1'
+    || typeof record.instanceId !== 'string' || !record.instanceId
+    || typeof record.runId !== 'string' || !record.runId
+    || !Number.isSafeInteger(record.version) || (record.version as number) < 0
+    || typeof record.logicalDigest !== 'string' || !/^sha256:[0-9a-f]{64}$/.test(record.logicalDigest)
+    || typeof record.runStepStateRootDigest !== 'string' || !/^sha256:[0-9a-f]{64}$/.test(record.runStepStateRootDigest)) {
+    throw new Error('Checkpoint head scalar fields are invalid.');
+  }
+  const definition = exactKeys(record.definition, ['revision', 'digest', 'provenance'], 'checkpoint definition');
+  if (typeof definition.revision !== 'string' || !definition.revision
+    || typeof definition.digest !== 'string' || !/^sha256:[0-9a-f]{64}$/.test(definition.digest)) {
+    throw new Error('Checkpoint definition identity is invalid.');
+  }
+  return Object.freeze(record as unknown as FlowRunCheckpointHead);
+}
+
+function readJson(file: string): unknown {
+  const stat = fs.lstatSync(file);
+  if (!stat.isFile() || stat.isSymbolicLink()) throw new Error(`Sidecar path ${file} must be a regular file.`);
+  let descriptor: number | undefined;
+  try {
+    descriptor = fs.openSync(file, fs.constants.O_RDONLY | NO_FOLLOW);
+    return JSON.parse(fs.readFileSync(descriptor, 'utf8')) as unknown;
+  } finally {
+    if (descriptor !== undefined) fs.closeSync(descriptor);
+  }
+}
+
+function readNode(treeDir: string, nodeDigest: DefinitionStepDigest): StepMerkleNode {
+  const file = path.join(treeDir, digestFileName(nodeDigest));
+  const decoded = readJson(file);
+  const kind = decoded && typeof decoded === 'object' && !Array.isArray(decoded)
+    ? (decoded as { readonly kind?: unknown }).kind
+    : undefined;
+  const value = exactKeys(
+    decoded,
+    kind === 'leaf'
+      ? ['kind', 'key', 'recordDigest', 'digest']
+      : ['kind', 'firstKey', 'lastKey', 'leftDigest', 'rightDigest', 'digest'],
+    'sidecar tree node',
+  ) as unknown as StepMerkleNode;
+  const expected = value.kind === 'leaf'
+    ? digest({ kind: 'leaf', key: value.key, recordDigest: value.recordDigest })
+    : value.kind === 'branch'
+      ? digest({
+          kind: 'branch', firstKey: value.firstKey, lastKey: value.lastKey,
+          leftDigest: value.leftDigest, rightDigest: value.rightDigest,
+        })
+      : undefined;
+  if (expected !== nodeDigest || value.digest !== nodeDigest) throw new Error(`Sidecar tree node ${nodeDigest} digest mismatch.`);
+  return value;
+}
+
+function readRecord(recordDir: string, recordDigest: DefinitionStepDigest): ImmutableStepRecord {
+  const value = exactKeys(
+    readJson(path.join(recordDir, digestFileName(recordDigest))),
+    ['schemaVersion', 'key', 'value', 'digest'],
+    'sidecar record',
+  ) as unknown as ImmutableStepRecord;
+  const expected = createImmutableStepRecord(value.key, value.value);
+  if (expected.digest !== recordDigest || value.digest !== recordDigest || canonicalJson(expected) !== canonicalJson(value)) {
+    throw new Error(`Sidecar record ${recordDigest} digest mismatch.`);
+  }
+  return expected;
+}
+
+function collectRecords(rootDir: string, rootDigest: DefinitionStepDigest): ImmutableStepRecord[] {
+  const treeDir = path.join(rootDir, 'trees');
+  const recordDir = path.join(rootDir, 'records');
+  const records = new Map<string, ImmutableStepRecord>();
+  const visited = new Set<string>();
+  const visit = (nodeDigest: DefinitionStepDigest): void => {
+    if (visited.has(nodeDigest)) return;
+    visited.add(nodeDigest);
+    const node = readNode(treeDir, nodeDigest);
+    if (node.kind === 'leaf') {
+      const record = readRecord(recordDir, node.recordDigest);
+      if (record.key !== node.key) throw new Error(`Sidecar leaf ${nodeDigest} key mismatch.`);
+      if (records.has(record.key)) throw new Error(`Sidecar record key ${record.key} is duplicated.`);
+      records.set(record.key, record);
+      return;
+    }
+    visit(node.leftDigest);
+    visit(node.rightDigest);
+  };
+  visit(rootDigest);
+  return [...records.values()].sort((left, right) => compareCodeUnits(left.key, right.key));
+}
+
+function writeImmutable(file: string, value: unknown): void {
+  fs.mkdirSync(path.dirname(file), { recursive: true });
+  if (fs.existsSync(file)) {
+    if (canonicalJson(readJson(file)) !== canonicalJson(value)) throw new Error(`Immutable sidecar collision at ${file}`);
+    return;
+  }
+  const staging = `${file}.candidate-${randomUUID()}`;
+  let descriptor: number | undefined;
+  try {
+    descriptor = fs.openSync(staging, fs.constants.O_CREAT | fs.constants.O_EXCL | fs.constants.O_WRONLY | NO_FOLLOW, 0o600);
+    fs.writeFileSync(descriptor, canonicalJson(value), 'utf8');
+    fs.fsyncSync(descriptor);
+    fs.closeSync(descriptor);
+    descriptor = undefined;
+    if (canonicalJson(readJson(staging)) !== canonicalJson(value)) throw new Error(`Sidecar candidate ${staging} readback mismatch.`);
+    try {
+      fs.renameSync(staging, file);
+    } catch (error) {
+      if (!fs.existsSync(file) || canonicalJson(readJson(file)) !== canonicalJson(value)) throw error;
+      fs.unlinkSync(staging);
+    }
+  } finally {
+    if (descriptor !== undefined) fs.closeSync(descriptor);
+    if (fs.existsSync(staging)) fs.unlinkSync(staging);
+  }
+}
+
+export function encodeCheckpointSidecars(
+  checkpointValue: AIWorkflowFlowRunCheckpoint,
+  previousTree: StepMerkleTree | undefined,
+): EncodedCheckpoint {
+  const checkpoint = normalizeFlowRunCheckpoint(checkpointValue);
+  const records = flatten(checkpoint as unknown as FlowClosedValue);
+  const tree = buildStepMerkleTree(records);
+  const treeDiff = diffStepMerkleTrees(previousTree, tree);
+  const previousRecords = new Set<string>();
+  if (previousTree) {
+    Object.values(previousTree.nodes).forEach((node) => {
+      if (node.kind === 'leaf') previousRecords.add(node.recordDigest);
+    });
+  }
+  const addedRecordDigests = records.map((record) => record.digest).filter((value) => !previousRecords.has(value)).sort(compareCodeUnits);
+  const logicalDigest = digest(checkpoint);
+  const head = Object.freeze({
+    schemaVersion: 'depa.flow-run-checkpoint-head/v1' as const,
+    instanceId: checkpoint.instanceId,
+    runId: checkpoint.runId,
+    definition: checkpoint.definition,
+    version: checkpoint.version,
+    logicalDigest,
+    runStepStateRootDigest: tree.rootDigest,
+  });
+  return {
+    head,
+    records,
+    tree,
+    receipt: Object.freeze({
+      schemaVersion: 'depa.flow-run-checkpoint-sidecar-receipt/v1' as const,
+      version: checkpoint.version,
+      previousRootDigest: previousTree?.rootDigest ?? null,
+      rootDigest: tree.rootDigest,
+      logicalDigest,
+      addedRecordDigests: Object.freeze(addedRecordDigests),
+      addedTreeNodeDigests: treeDiff.addedNodeDigests,
+    }),
+  };
+}
+
+export function persistCheckpointSidecars(runDir: string, encoded: EncodedCheckpoint): void {
+  const root = path.join(runDir, 'step-space');
+  const addedRecords = new Set(encoded.receipt.addedRecordDigests);
+  for (const record of encoded.records.filter((value) => addedRecords.has(value.digest))) {
+    writeImmutable(path.join(root, 'records', digestFileName(record.digest)), record);
+  }
+  const addedNodes = new Set(encoded.receipt.addedTreeNodeDigests);
+  for (const node of Object.values(encoded.tree.nodes).filter((value) => addedNodes.has(value.digest))) {
+    writeImmutable(path.join(root, 'trees', digestFileName(node.digest)), node);
+  }
+  writeImmutable(path.join(root, 'receipts', `${encoded.head.version}.json`), encoded.receipt);
+}
+
+function assertReceipt(root: string, head: FlowRunCheckpointHead): void {
+  const value = exactKeys(
+    readJson(path.join(root, 'receipts', `${head.version}.json`)),
+    [
+      'schemaVersion', 'version', 'previousRootDigest', 'rootDigest', 'logicalDigest',
+      'addedRecordDigests', 'addedTreeNodeDigests',
+    ],
+    'sidecar receipt',
+  );
+  const digestArray = (field: 'addedRecordDigests' | 'addedTreeNodeDigests'): readonly string[] => {
+    const entries = value[field];
+    if (!Array.isArray(entries)
+      || Reflect.ownKeys(entries).some((key) => typeof key !== 'string'
+        || (key !== 'length' && (!/^(0|[1-9][0-9]*)$/.test(key) || Number(key) >= entries.length)))
+      || entries.some((entry) => typeof entry !== 'string' || !/^sha256:[0-9a-f]{64}$/.test(entry))) {
+      throw new Error(`Sidecar receipt ${field} is invalid.`);
+    }
+    const sorted = [...entries].sort(compareCodeUnits);
+    if (new Set(sorted).size !== sorted.length || canonicalJson(entries) !== canonicalJson(sorted)) {
+      throw new Error(`Sidecar receipt ${field} must be sorted and unique.`);
+    }
+    return entries;
+  };
+  digestArray('addedRecordDigests');
+  digestArray('addedTreeNodeDigests');
+  if (value.schemaVersion !== 'depa.flow-run-checkpoint-sidecar-receipt/v1'
+    || value.version !== head.version
+    || value.rootDigest !== head.runStepStateRootDigest
+    || value.logicalDigest !== head.logicalDigest
+    || (value.previousRootDigest !== null
+      && (typeof value.previousRootDigest !== 'string' || !/^sha256:[0-9a-f]{64}$/.test(value.previousRootDigest)))) {
+    throw new Error('Sidecar receipt does not bind the accepted head.');
+  }
+}
+
+export function readCheckpointHeadAndLogical(file: string): {
+  readonly head: FlowRunCheckpointHead;
+  readonly checkpoint: AIWorkflowFlowRunCheckpoint;
+  readonly tree: StepMerkleTree;
+} {
+  const head = normalizeHead(readJson(file));
+  const root = path.join(path.dirname(file), 'step-space');
+  assertReceipt(root, head);
+  const records = collectRecords(root, head.runStepStateRootDigest);
+  const checkpoint = reconstruct(records);
+  if (digest(checkpoint) !== head.logicalDigest
+    || checkpoint.instanceId !== head.instanceId
+    || checkpoint.runId !== head.runId
+    || checkpoint.version !== head.version
+    || canonicalJson(checkpoint.definition) !== canonicalJson(head.definition)) {
+    throw new Error('Checkpoint head and reconstructed logical value differ.');
+  }
+  const nodes: Record<string, StepMerkleNode> = {};
+  const leafByKey: Record<string, DefinitionStepDigest> = {};
+  const keyOrder: string[] = [];
+  const visit = (nodeDigest: DefinitionStepDigest): void => {
+    if (nodes[nodeDigest]) return;
+    const node = readNode(path.join(root, 'trees'), nodeDigest);
+    nodes[nodeDigest] = node;
+    if (node.kind === 'leaf') {
+      leafByKey[node.key] = nodeDigest;
+      keyOrder.push(node.key);
+    } else {
+      visit(node.leftDigest);
+      visit(node.rightDigest);
+    }
+  };
+  visit(head.runStepStateRootDigest);
+  keyOrder.sort(compareCodeUnits);
+  return {
+    head,
+    checkpoint,
+    tree: Object.freeze({
+      schemaVersion: 'depa.step-merkle-tree/v1',
+      rootDigest: head.runStepStateRootDigest,
+      keyOrder: Object.freeze(keyOrder),
+      leafByKey: Object.freeze(leafByKey),
+      nodes: Object.freeze(nodes),
+    }),
+  };
+}
+
+export function isCheckpointHead(value: unknown): boolean {
+  return !!value && typeof value === 'object' && !Array.isArray(value)
+    && (value as { readonly schemaVersion?: unknown }).schemaVersion === 'depa.flow-run-checkpoint-head/v1';
+}
+
+export function checkpointHeadJson(head: FlowRunCheckpointHead): string {
+  return canonicalJson(head);
+}