@aztec/pxe 0.0.1-commit.9b94fc1 → 0.0.1-commit.9badcec54

package/src/entrypoints/client/lazy/utils.ts

@@ -1,5 +1,4 @@
 import { BBLazyPrivateKernelProver } from '@aztec/bb-prover/client/lazy';
-import { randomBytes } from '@aztec/foundation/crypto';
 import { createLogger } from '@aztec/foundation/log';
 import { createStore } from '@aztec/kv-store/indexeddb';
 import { LazyProtocolContractsProvider } from '@aztec/protocol-contracts/providers/lazy';
@@ -8,7 +7,8 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
 import type { PXEConfig } from '../../../config/index.js';
 import { PXE } from '../../../pxe.js';
-import type { PXECreationOptions } from '../../pxe_creation_options.js';
+import { PXE_DATA_SCHEMA_VERSION } from '../../../storage/metadata.js';
+import { type PXECreationOptions, isPrivateKernelProver } from '../../pxe_creation_options.js';
 
 /**
  * Create and start an PXE instance with the given AztecNode.
@@ -16,7 +16,7 @@ import type { PXECreationOptions } from '../../pxe_creation_options.js';
  *
  * @param aztecNode - The AztecNode instance to be used by the server.
  * @param config - The PXE Config to use
- * @param
+ * @param options - (Optional) Optional information for creating an PXE.
  * @returns A Promise that resolves to the started PXE instance.
  */
 export async function createPXE(
@@ -24,12 +24,7 @@ export async function createPXE(
   config: PXEConfig,
   options: PXECreationOptions = { loggers: {} },
 ) {
-  const logSuffix =
-    typeof options.useLogSuffix === 'boolean'
-      ? options.useLogSuffix
-        ? randomBytes(3).toString('hex')
-        : undefined
-      : options.useLogSuffix;
+  const actor = options.loggerActorLabel;
 
   const l1Contracts = await aztecNode.getL1ContractAddresses();
   const configWithContracts = {
@@ -39,20 +34,31 @@ export async function createPXE(
 
   const loggers = options.loggers ?? {};
 
-  const storeLogger = loggers.store ? loggers.store : createLogger('pxe:data:idb' + (logSuffix ? `:${logSuffix}` : ''));
+  const storeLogger = loggers.store ?? createLogger('pxe:data:idb', { actor });
 
-  const store = options.store ?? (await createStore('pxe_data', configWithContracts, storeLogger));
+  const store =
+    options.store ?? (await createStore('pxe_data', configWithContracts, PXE_DATA_SCHEMA_VERSION, storeLogger));
 
-  const simulator = new WASMSimulator();
-  const proverLogger = loggers.prover
-    ? loggers.prover
-    : createLogger('pxe:bb:wasm:bundle' + (logSuffix ? `:${logSuffix}` : ''));
-
-  const prover = options.prover ?? new BBLazyPrivateKernelProver(simulator, proverLogger);
+  const simulator = options.simulator ?? new WASMSimulator();
+  const proverLogger = loggers.prover ?? createLogger('pxe:bb:wasm:bundle', { actor });
 
+  let prover;
+  if (isPrivateKernelProver(options.proverOrOptions)) {
+    prover = options.proverOrOptions;
+  } else {
+    prover = new BBLazyPrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });
+  }
   const protocolContractsProvider = new LazyProtocolContractsProvider();
 
-  const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));
-  const pxe = await PXE.create(aztecNode, store, prover, simulator, protocolContractsProvider, config, pxeLogger);
+  const pxeLogger = loggers.pxe ?? createLogger('pxe:service', { actor });
+  const pxe = await PXE.create({
+    node: aztecNode,
+    store,
+    proofCreator: prover,
+    simulator,
+    protocolContractsProvider,
+    config,
+    loggerOrSuffix: pxeLogger,
+  });
   return pxe;
 }

package/src/entrypoints/pxe_creation_options.ts

@@ -1,10 +1,21 @@
+import type { BBPrivateKernelProverOptions } from '@aztec/bb-prover/client';
 import type { Logger } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore } from '@aztec/kv-store';
+import type { CircuitSimulator } from '@aztec/simulator/client';
 import type { PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
 
 export type PXECreationOptions = {
   loggers?: { store?: Logger; pxe?: Logger; prover?: Logger };
-  useLogSuffix?: boolean | string;
-  prover?: PrivateKernelProver;
+  /** Actor label to include in log output (e.g., 'pxe-0', 'pxe-test'). */
+  loggerActorLabel?: string;
+  proverOrOptions?: PrivateKernelProver | BBPrivateKernelProverOptions;
   store?: AztecAsyncKVStore;
+  simulator?: CircuitSimulator;
 };
+
+/** Checks if the given value implements the PrivateKernelProver interface via duck-typing. */
+export function isPrivateKernelProver(value: unknown): value is PrivateKernelProver {
+  return (
+    typeof value === 'object' && value !== null && typeof (value as PrivateKernelProver).createChonkProof === 'function'
+  );
+}

package/src/entrypoints/server/index.ts

@@ -1,8 +1,11 @@
+export * from '../../notes_filter.js';
 export * from '../../pxe.js';
 export * from '../../config/index.js';
 export * from '../../error_enriching.js';
 export * from '../../storage/index.js';
 export * from './utils.js';
-export { PXEOracleInterface } from '../../contract_function_simulator/pxe_oracle_interface.js';
-export { ORACLE_VERSION } from '../../oracle_version.js';
+export { NoteService } from '../../notes/note_service.js';
+export { ORACLE_VERSION_MAJOR, ORACLE_VERSION_MINOR } from '../../oracle_version.js';
 export { type PXECreationOptions } from '../pxe_creation_options.js';
+export { JobCoordinator } from '../../job_coordinator/job_coordinator.js';
+export { ContractSyncService } from '../../contract_sync/contract_sync_service.js';

package/src/entrypoints/server/utils.ts

@@ -1,20 +1,15 @@
 import { BBBundlePrivateKernelProver } from '@aztec/bb-prover/client/bundle';
-import { randomBytes } from '@aztec/foundation/crypto';
-import { type Logger, createLogger } from '@aztec/foundation/log';
+import { createLogger } from '@aztec/foundation/log';
+import { createStore } from '@aztec/kv-store/lmdb-v2';
 import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import {
-  type CircuitSimulator,
-  MemoryCircuitRecorder,
-  SimulatorRecorderWrapper,
-  WASMSimulator,
-} from '@aztec/simulator/client';
+import { MemoryCircuitRecorder, SimulatorRecorderWrapper, WASMSimulator } from '@aztec/simulator/client';
 import { FileCircuitRecorder } from '@aztec/simulator/testing';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
 import type { PXEConfig } from '../../config/index.js';
 import { PXE } from '../../pxe.js';
 import { PXE_DATA_SCHEMA_VERSION } from '../../storage/index.js';
-import type { PXECreationOptions } from '../pxe_creation_options.js';
+import { type PXECreationOptions, isPrivateKernelProver } from '../pxe_creation_options.js';
 
 type PXEConfigWithoutDefaults = Omit<PXEConfig, 'l1Contracts' | 'l1ChainId' | 'l2BlockBatchSize' | 'rollupVersion'>;
 
@@ -23,17 +18,13 @@ export async function createPXE(
   config: PXEConfigWithoutDefaults,
   options: PXECreationOptions = { loggers: {} },
 ) {
+  const actor = options.loggerActorLabel;
+  const recorderLogger = createLogger('simulator:acvm:recording', { actor });
   const recorder = process.env.CIRCUIT_RECORD_DIR
-    ? new FileCircuitRecorder(process.env.CIRCUIT_RECORD_DIR)
-    : new MemoryCircuitRecorder();
-  const simulator = new SimulatorRecorderWrapper(new WASMSimulator(), recorder);
-
-  const logSuffix =
-    typeof options.useLogSuffix === 'boolean'
-      ? options.useLogSuffix
-        ? randomBytes(3).toString('hex')
-        : undefined
-      : options.useLogSuffix;
+    ? new FileCircuitRecorder(process.env.CIRCUIT_RECORD_DIR, recorderLogger)
+    : new MemoryCircuitRecorder(recorderLogger);
+  const simulatorLogger = createLogger('wasm-simulator', { actor });
+  const simulator = new SimulatorRecorderWrapper(new WASMSimulator(simulatorLogger), recorder);
   const loggers = options.loggers ?? {};
 
   const { l1ChainId, l1ContractAddresses: l1Contracts, rollupVersion } = await aztecNode.getNodeInfo();
@@ -46,34 +37,34 @@ export async function createPXE(
   };
 
   if (!options.store) {
-    // TODO once https://github.com/AztecProtocol/aztec-packages/issues/13656 is fixed, we can remove this and always
-    // import the lmdb-v2 version
-    const { createStore } = await import('@aztec/kv-store/lmdb-v2');
-    const storeLogger = loggers.store
-      ? loggers.store
-      : createLogger('pxe:data:lmdb' + (logSuffix ? `:${logSuffix}` : ''));
-    options.store = await createStore('pxe_data', PXE_DATA_SCHEMA_VERSION, configWithContracts, storeLogger);
+    const storeLogger = loggers.store ?? createLogger('pxe:data:lmdb', { actor });
+    options.store = await createStore(
+      'pxe_data',
+      PXE_DATA_SCHEMA_VERSION,
+      configWithContracts,
+      storeLogger.getBindings(),
+    );
+  }
+  const proverLogger = loggers.prover ?? createLogger('pxe:bb:native', { actor });
+
+  let prover;
+  if (isPrivateKernelProver(options.proverOrOptions)) {
+    prover = options.proverOrOptions;
+  } else {
+    prover = new BBBundlePrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });
   }
-  const proverLogger = loggers.prover
-    ? loggers.prover
-    : createLogger('pxe:bb:native' + (logSuffix ? `:${logSuffix}` : ''));
 
-  const prover = options.prover ?? createProver(simulator, proverLogger);
   const protocolContractsProvider = new BundledProtocolContractsProvider();
 
-  const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));
-  const pxe = await PXE.create(
-    aztecNode,
-    options.store,
-    prover,
+  const pxeLogger = loggers.pxe ?? createLogger('pxe:service', { actor });
+  const pxe = await PXE.create({
+    node: aztecNode,
+    store: options.store,
+    proofCreator: prover,
     simulator,
     protocolContractsProvider,
-    configWithContracts,
-    pxeLogger,
-  );
+    config: configWithContracts,
+    loggerOrSuffix: pxeLogger,
+  });
   return pxe;
 }
-
-function createProver(simulator: CircuitSimulator, logger?: Logger) {
-  return new BBBundlePrivateKernelProver(simulator, logger);
-}

package/src/error_enriching.ts

@@ -4,18 +4,14 @@ import { FunctionSelector } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { type SimulationError, isNoirCallStackUnresolved } from '@aztec/stdlib/errors';
 
-import type { ContractDataProvider } from './storage/contract_data_provider/contract_data_provider.js';
+import type { ContractStore } from './storage/contract_store/contract_store.js';
 
 /**
  * Adds contract and function names to a simulation error, if they
  * can be found in the PXE database
  * @param err - The error to enrich.
  */
-export async function enrichSimulationError(
-  err: SimulationError,
-  contractDataProvider: ContractDataProvider,
-  logger: Logger,
-) {
+export async function enrichSimulationError(err: SimulationError, contractStore: ContractStore, logger: Logger) {
   // Maps contract addresses to the set of function selectors that were in error.
   // Map and Set do reference equality for their keys instead of value equality, so we store the string
   // representation to get e.g. different contract address objects with the same address value to match.
@@ -34,7 +30,7 @@ export async function enrichSimulationError(
   await Promise.all(
     [...mentionedFunctions.entries()].map(async ([contractAddress, fnSelectors]) => {
       const parsedContractAddress = AztecAddress.fromString(contractAddress);
-      const contract = await contractDataProvider.getContract(parsedContractAddress);
+      const contract = await contractStore.getContract(parsedContractAddress);
       if (contract) {
         err.enrichWithContractName(parsedContractAddress, contract.name);
         // Map from function selector to function name. It uses a stringified key for the same reason as mentionedFunctions.
@@ -68,11 +64,7 @@ export async function enrichSimulationError(
   );
 }
 
-export async function enrichPublicSimulationError(
-  err: SimulationError,
-  contractDataProvider: ContractDataProvider,
-  logger: Logger,
-) {
+export async function enrichPublicSimulationError(err: SimulationError, contractStore: ContractStore, logger: Logger) {
   const callStack = err.getCallStack();
   const originalFailingFunction = callStack[callStack.length - 1];
 
@@ -80,7 +72,7 @@ export async function enrichPublicSimulationError(
     throw new Error(`Original failing function not found when enriching public simulation, missing callstack`);
   }
 
-  const artifact = await contractDataProvider.getPublicFunctionArtifact(originalFailingFunction.contractAddress);
+  const artifact = await contractStore.getPublicFunctionArtifact(originalFailingFunction.contractAddress);
   if (!artifact) {
     throw new Error(
       `Artifact not found when enriching public simulation error. Contract address: ${originalFailingFunction.contractAddress}.`,
@@ -92,7 +84,7 @@ export async function enrichPublicSimulationError(
     err.setOriginalMessage(err.getOriginalMessage() + `${assertionMessage}`);
   }
 
-  const debugInfo = await contractDataProvider.getPublicFunctionDebugMetadata(originalFailingFunction.contractAddress);
+  const debugInfo = await contractStore.getPublicFunctionDebugMetadata(originalFailingFunction.contractAddress);
 
   const noirCallStack = err.getNoirCallStack();
   if (debugInfo) {
@@ -110,6 +102,6 @@ export async function enrichPublicSimulationError(
         );
       }
     }
-    await enrichSimulationError(err, contractDataProvider, logger);
+    await enrichSimulationError(err, contractStore, logger);
   }
 }

package/src/events/event_service.ts

@@ -0,0 +1,96 @@
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import { createLogger } from '@aztec/foundation/log';
+import type { EventSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { computePrivateEventCommitment, siloNullifier } from '@aztec/stdlib/hash';
+import type { AztecNode } from '@aztec/stdlib/interfaces/server';
+import type { BlockHeader, TxHash } from '@aztec/stdlib/tx';
+
+import { PrivateEventStore } from '../storage/private_event_store/private_event_store.js';
+
+export class EventService {
+  constructor(
+    private readonly anchorBlockHeader: BlockHeader,
+    private readonly aztecNode: AztecNode,
+    private readonly privateEventStore: PrivateEventStore,
+    private readonly jobId: string,
+    private readonly log = createLogger('pxe:event_service'),
+  ) {}
+
+  public async validateAndStoreEvent(
+    contractAddress: AztecAddress,
+    selector: EventSelector,
+    randomness: Fr,
+    content: Fr[],
+    eventCommitment: Fr,
+    txHash: TxHash,
+    scope: AztecAddress,
+  ): Promise<void> {
+    // Defense-in-depth: the built-in private-event path derives this commitment from content before enqueueing, but
+    // unconstrained PXE-side code (e.g. a custom message handler) can reach this oracle with arbitrary
+    // (content, commitment) pairs. Without this check it could bind arbitrary content to a legitimate tx nullifier,
+    // causing PXE to surface fabricated event data.
+    const recomputedCommitment = await computePrivateEventCommitment(randomness, selector.toField(), content);
+    if (!recomputedCommitment.equals(eventCommitment)) {
+      this.log.warn(
+        `Skipping event whose content does not hash to the provided commitment. contract=${contractAddress}, selector=${selector}, eventCommitment=${eventCommitment}, txHash=${txHash}, recomputedCommitment=${recomputedCommitment}`,
+      );
+      return;
+    }
+
+    // While using 'latest' block number would be fine for private events since they cannot be accessed from Aztec.nr
+    // (and thus we're less concerned about being ahead of the synced block), we use the synced block number to
+    // maintain consistent behavior in the PXE. Additionally, events should never be ahead of the synced block here
+    // since `fetchTaggedLogs` only processes logs up to the synced block.
+    const [siloedEventCommitment, txEffect] = await Promise.all([
+      siloNullifier(contractAddress, eventCommitment),
+      this.aztecNode.getTxEffect(txHash),
+    ]);
+
+    const anchorBlockNumber = this.anchorBlockHeader.getBlockNumber();
+
+    if (!txEffect) {
+      // We error out instead of just logging a warning and skipping the event because this would indicate a bug. This
+      // is because the node has already served info about this tx either when obtaining the log (TxScopedL2Log contain
+      // tx info) or when getting metadata for the offchain message (before the message got passed to `process_log`).
+      throw new Error(`Could not find tx effect for tx hash ${txHash} when processing an event.`);
+    }
+
+    if (txEffect.l2BlockNumber > anchorBlockNumber) {
+      // We should never process a message from a tx past the anchor block. If we got here, a preprocessing step made
+      // a mistake.
+      throw new Error(
+        `Obtained a newer tx effect for ${txHash} for an event validation request than the anchor block ${anchorBlockNumber}. This is a bug as smart contracts should not issue event validation requests for events from blocks newer than the anchor block.`,
+      );
+    }
+
+    // Find the index of the event commitment in the nullifiers array to determine event ordering within the tx
+    const eventIndexInTx = txEffect.data.nullifiers.findIndex(n => n.equals(siloedEventCommitment));
+    if (eventIndexInTx === -1) {
+      // Unlike in NoteService, this might not be a bug since the commitment hasn't been verified yet in the message
+      // processing pipeline. A malformed or malicious message could trigger this condition. Because of this we don't
+      // error out and we just show a warning.
+      this.log.warn(
+        `Skipping event whose commitment is not present in its tx. siloedEventCommitment=${siloedEventCommitment}, contract=${contractAddress}, selector=${selector}, eventCommitment=${eventCommitment}, txHash=${txHash}`,
+      );
+      return;
+    }
+
+    return this.privateEventStore.storePrivateEventLog(
+      selector,
+      randomness,
+      content,
+      siloedEventCommitment,
+      {
+        contractAddress,
+        scope,
+        txHash,
+        l2BlockNumber: txEffect.l2BlockNumber,
+        l2BlockHash: txEffect.l2BlockHash,
+        txIndexInBlock: txEffect.txIndexInBlock,
+        eventIndexInTx,
+      },
+      this.jobId,
+    );
+  }
+}

package/src/events/index.ts

@@ -0,0 +1 @@
+export * from './private_event_filter_validator.js';

package/src/events/private_event_filter_validator.ts

@@ -0,0 +1,66 @@
+import type { PrivateEventFilter } from '@aztec/aztec.js/wallet';
+import { INITIAL_L2_BLOCK_NUM } from '@aztec/constants';
+import { BlockNumber } from '@aztec/foundation/branded-types';
+import { createLogger } from '@aztec/foundation/log';
+
+import type { PrivateEventStoreFilter } from '../storage/private_event_store/private_event_store.js';
+
+export class PrivateEventFilterValidator {
+  private readonly log = createLogger('pxe:private_event_filter_validator');
+
+  constructor(private readonly lastBlock: BlockNumber) {}
+
+  validate(filter: PrivateEventFilter): PrivateEventStoreFilter {
+    let { fromBlock, toBlock } = filter;
+
+    // Block range filters in Aztec Node are defined as closed-open intervals [fromBlock, toBlock), so
+    // we respect that convention here for consistency.
+    // We then default to [INITIAL_L2_BLOCK_NUM, latestKnownBlock + 1), ie: by default return events from
+    // the first block to the latest known block.
+    if (!fromBlock || !toBlock) {
+      fromBlock = fromBlock ?? BlockNumber(INITIAL_L2_BLOCK_NUM);
+      toBlock = toBlock ?? BlockNumber(this.lastBlock + 1);
+    }
+
+    if (filter.scopes.length === 0) {
+      throw new Error('At least one scope is required to get private events');
+    }
+
+    if (fromBlock < 1) {
+      throw new Error('fromBlock must be greater than or equal to 1');
+    }
+
+    if (toBlock < 1) {
+      throw new Error('toBlock must be greater than or equal to 1');
+    }
+
+    if (fromBlock >= toBlock) {
+      throw new Error('toBlock must be strictly greater than fromBlock');
+    }
+
+    // Cap the requested range to the synced block range. Without this, callers that pass a large
+    // toBlock (e.g. Number.MAX_SAFE_INTEGER as a "give me everything" idiom) would silently receive
+    // only the events that happen to be synced and believe they have complete coverage.
+    // We warn + cap rather than throw so callers don't need to query the last synced block before
+    // every request (which would also be unreliable, as the block can advance between the two calls).
+    const syncedUpperBound = BlockNumber(this.lastBlock + 1);
+    if (fromBlock >= syncedUpperBound) {
+      this.log.warn(
+        `Requested fromBlock ${fromBlock} is past last synced block ${this.lastBlock}; no events will be returned until PXE syncs further.`,
+      );
+    } else if (toBlock > syncedUpperBound) {
+      this.log.warn(
+        `Requested toBlock ${toBlock} exceeds last synced block ${this.lastBlock}; capping to ${syncedUpperBound}. Retry once PXE is further synced for complete coverage.`,
+      );
+      toBlock = syncedUpperBound;
+    }
+
+    return {
+      contractAddress: filter.contractAddress,
+      scopes: filter.scopes,
+      txHash: filter.txHash,
+      toBlock,
+      fromBlock,
+    };
+  }
+}

package/src/job_coordinator/job_coordinator.ts

@@ -0,0 +1,150 @@
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+
+/**
+ * Interface that stores must implement to support staged writes.
+ */
+export interface StagedStore {
+  /** Unique name identifying this store (used for tracking staged stores from JobCoordinator) */
+  readonly storeName: string;
+
+  /**
+   * Commits staged data to main storage.
+   * Should be called within a transaction for atomicity.
+   *
+   * @param jobId - The job identifier
+   */
+  commit(jobId: string): Promise<void>;
+
+  /**
+   * Discards staged data without committing.
+   * Called on abort.
+   *
+   * @param jobId - The job identifier
+   */
+  discardStaged(jobId: string): Promise<void>;
+}
+
+/**
+ * JobCoordinator manages job lifecycle and provides crash resilience for PXE operations.
+ *
+ * It uses a staged writes pattern:
+ * 1. When a job begins, a unique job ID is created
+ * 2. During the job, all writes go to staging (keyed by job ID)
+ * 3. On commit, staging is promoted to main storage
+ * 4. On abort, staged data is discarded
+ *
+ * Note: PXE should only rely on a single JobCoordinator instance, so it can eventually
+ * orchestrate concurrent jobs. Right now it doesn't make a difference because we're
+ * using a job queue with concurrency=1.
+ */
+export class JobCoordinator {
+  private readonly log: Logger;
+
+  /** The underlying KV store */
+  kvStore: AztecAsyncKVStore;
+
+  #currentJobId: string | undefined;
+  #stores: Map<string, StagedStore> = new Map();
+
+  constructor(kvStore: AztecAsyncKVStore, bindings?: LoggerBindings) {
+    this.kvStore = kvStore;
+    this.log = createLogger('pxe:job_coordinator', bindings);
+  }
+
+  /**
+   * Registers a staged store.
+   * Must be called during initialization for all stores that need staging support.
+   */
+  registerStore(store: StagedStore): void {
+    if (this.#stores.has(store.storeName)) {
+      throw new Error(`Store "${store.storeName}" is already registered`);
+    }
+    this.#stores.set(store.storeName, store);
+    this.log.debug(`Registered staged store: ${store.storeName}`);
+  }
+
+  /**
+   * Registers multiple staged stores.
+   */
+  registerStores(stores: StagedStore[]): void {
+    for (const store of stores) {
+      this.registerStore(store);
+    }
+  }
+
+  /**
+   * Begins a new job and returns a job ID for staged writes.
+   *
+   * @returns Job ID to pass to store operations
+   */
+  beginJob(): string {
+    if (this.#currentJobId) {
+      throw new Error(
+        `Cannot begin job: job ${this.#currentJobId} is already in progress. ` +
+          `This should not happen - ensure jobs are properly committed or aborted.`,
+      );
+    }
+
+    const jobId = randomBytes(8).toString('hex');
+    this.#currentJobId = jobId;
+
+    this.log.debug(`Started job ${jobId}`);
+    return jobId;
+  }
+
+  /**
+   * Commits a job by promoting all staged data to main storage.
+   *
+   * @param jobId - The job ID returned from beginJob
+   */
+  async commitJob(jobId: string): Promise<void> {
+    if (!this.#currentJobId || this.#currentJobId !== jobId) {
+      throw new Error(
+        `Cannot commit job ${jobId}: no matching job in progress. ` + `Current job: ${this.#currentJobId ?? 'none'}`,
+      );
+    }
+
+    this.log.debug(`Committing job ${jobId}`);
+
+    // Commit all stores atomically in a single transaction.
+    // Each store's commit is a no-op if it has no staged data (but that's up to each store to handle).
+    await this.kvStore.transactionAsync(async () => {
+      for (const store of this.#stores.values()) {
+        await store.commit(jobId);
+      }
+    });
+
+    this.#currentJobId = undefined;
+    this.log.debug(`Job ${jobId} committed successfully`);
+  }
+
+  /**
+   * Aborts a job by discarding all staged data.
+   *
+   * @param jobId - The job ID returned from beginJob
+   */
+  async abortJob(jobId: string): Promise<void> {
+    if (!this.#currentJobId || this.#currentJobId !== jobId) {
+      // Job may have already been aborted or never started properly
+      this.log.warn(`Abort called for job ${jobId} but current job is ${this.#currentJobId ?? 'none'}`);
+    }
+
+    this.log.debug(`Aborting job ${jobId}`);
+
+    for (const store of this.#stores.values()) {
+      await store.discardStaged(jobId);
+    }
+
+    this.#currentJobId = undefined;
+    this.log.debug(`Job ${jobId} aborted`);
+  }
+
+  /**
+   * Checks if there's a job currently in progress.
+   */
+  hasJobInProgress(): boolean {
+    return this.#currentJobId !== undefined;
+  }
+}