@aztec/pxe 4.0.0-nightly.20260108 → 4.0.0-nightly.20260111

package/src/job_coordinator/job_coordinator.ts

@@ -0,0 +1,149 @@
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { 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 = createLogger('pxe:job_coordinator');
+
+  /** The underlying KV store */
+  kvStore: AztecAsyncKVStore;
+
+  #currentJobId: string | undefined;
+  #stores: Map<string, StagedStore> = new Map();
+
+  constructor(kvStore: AztecAsyncKVStore) {
+    this.kvStore = kvStore;
+  }
+
+  /**
+   * 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;
+  }
+}

package/src/pxe.ts

@@ -60,9 +60,11 @@ import {
 } from './contract_function_simulator/contract_function_simulator.js';
 import { readCurrentClassId } from './contract_function_simulator/oracle/private_execution.js';
 import { ProxiedContractStoreFactory } from './contract_function_simulator/proxied_contract_data_source.js';
+import { ProxiedNodeFactory } from './contract_function_simulator/proxied_node.js';
 import { PXEDebugUtils } from './debug/pxe_debug_utils.js';
 import { enrichPublicSimulationError, enrichSimulationError } from './error_enriching.js';
 import { PrivateEventFilterValidator } from './events/private_event_filter_validator.js';
+import { JobCoordinator } from './job_coordinator/job_coordinator.js';
 import {
   PrivateKernelExecutionProver,
   type PrivateKernelExecutionProverConfig,
@@ -107,6 +109,7 @@ export class PXE {
     private protocolContractsProvider: ProtocolContractsProvider,
     private log: Logger,
     private jobQueue: SerialQueue,
+    private jobCoordinator: JobCoordinator,
     public debug: PXEDebugUtils,
   ) {}
 
@@ -153,6 +156,8 @@ export class PXE {
       loggerOrSuffix,
     );
 
+    const jobCoordinator = new JobCoordinator(store);
+
     const debugUtils = new PXEDebugUtils(contractStore, noteStore);
 
     const jobQueue = new SerialQueue();
@@ -176,6 +181,7 @@ export class PXE {
       protocolContractsProvider,
       log,
       jobQueue,
+      jobCoordinator,
       debugUtils,
     );
 
@@ -199,7 +205,7 @@ export class PXE {
       this.noteStore,
       this.keyStore,
       this.addressStore,
-      this.node,
+      ProxiedNodeFactory.create(this.node),
       this.anchorBlockStore,
       this.senderTaggingStore,
       this.recipientTaggingStore,
@@ -230,7 +236,7 @@ export class PXE {
    *
    * Useful for tasks that cannot run concurrently, such as contract function simulation.
    */
-  #putInJobQueue<T>(fn: () => Promise<T>): Promise<T> {
+  #putInJobQueue<T>(fn: (jobId: string) => Promise<T>): Promise<T> {
     // TODO(#12636): relax the conditions under which we forbid concurrency.
     if (this.jobQueue.length() != 0) {
       this.log.warn(
@@ -238,7 +244,22 @@ export class PXE {
       );
     }
 
-    return this.jobQueue.put(fn);
+    return this.jobQueue.put(async () => {
+      const jobId = this.jobCoordinator.beginJob();
+      this.log.verbose(`Beginning job ${jobId}`);
+
+      try {
+        const result = await fn(jobId);
+        this.log.verbose(`Committing job ${jobId}`);
+
+        await this.jobCoordinator.commitJob(jobId);
+        return result;
+      } catch (err) {
+        this.log.verbose(`Aborting job ${jobId}`);
+        await this.jobCoordinator.abortJob(jobId);
+        throw err;
+      }
+    });
   }
 
   async #registerProtocolContracts() {
@@ -271,7 +292,8 @@ export class PXE {
   async #executePrivate(
     contractFunctionSimulator: ContractFunctionSimulator,
     txRequest: TxExecutionRequest,
-    scopes?: AztecAddress[],
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ): Promise<PrivateExecutionResult> {
     const { origin: contractAddress, functionSelector } = txRequest;
 
@@ -288,6 +310,7 @@ export class PXE {
         // contract entrypoint
         undefined, // senderForTags
         scopes,
+        jobId,
       );
       this.log.debug(`Private simulation completed for ${contractAddress.toString()}:${functionSelector}`);
       return result;
@@ -306,17 +329,19 @@ export class PXE {
    * @param authWitnesses - Authentication witnesses required for the function call.
    * @param scopes - Optional array of account addresses whose notes can be accessed in this call. Defaults to all
    * accounts if not specified.
+   * @param jobId - The job ID for staged writes.
    * @returns The simulation result containing the outputs of the utility function.
    */
   async #simulateUtility(
     contractFunctionSimulator: ContractFunctionSimulator,
     call: FunctionCall,
-    authWitnesses?: AuthWitness[],
-    scopes?: AztecAddress[],
+    authWitnesses: AuthWitness[] | undefined,
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ) {
     try {
       const anchorBlockHeader = await this.anchorBlockStore.getBlockHeader();
-      return contractFunctionSimulator.runUtility(call, authWitnesses ?? [], anchorBlockHeader, scopes);
+      return contractFunctionSimulator.runUtility(call, authWitnesses ?? [], anchorBlockHeader, scopes, jobId);
     } catch (err) {
       if (err instanceof SimulationError) {
         await enrichSimulationError(err, this.contractStore, this.log);
@@ -665,14 +690,14 @@ export class PXE {
     let privateExecutionResult: PrivateExecutionResult;
     // We disable proving concurrently mostly out of caution, since it accesses some of our stores. Proving is so
     // computationally demanding that it'd be rare for someone to try to do it concurrently regardless.
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
       try {
         const syncTimer = new Timer();
         await this.blockStateSynchronizer.sync();
         const syncTime = syncTimer.ms();
         const contractFunctionSimulator = this.#getSimulatorForTx();
-        privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest);
+        privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, undefined, jobId);
 
         const {
           publicInputs,
@@ -749,7 +774,7 @@ export class PXE {
     skipProofGeneration: boolean = true,
   ): Promise<TxProfileResult> {
     // We disable concurrent profiles for consistency with simulateTx.
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
       try {
         const txInfo = {
@@ -769,7 +794,12 @@ export class PXE {
         const syncTime = syncTimer.ms();
 
         const contractFunctionSimulator = this.#getSimulatorForTx();
-        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest);
+        const privateExecutionResult = await this.#executePrivate(
+          contractFunctionSimulator,
+          txRequest,
+          undefined,
+          jobId,
+        );
 
         const { executionSteps, timings: { proving } = {} } = await this.#prove(
           txRequest,
@@ -849,7 +879,7 @@ export class PXE {
     // We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another simulation is currently modifying).
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       try {
         const totalTimer = new Timer();
         const txInfo = {
@@ -875,7 +905,7 @@ export class PXE {
         const skipKernels = overrides?.contracts !== undefined && Object.keys(overrides.contracts ?? {}).length > 0;
 
         // Execution of private functions only; no proving, and no kernel logic.
-        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, scopes);
+        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, scopes, jobId);
 
         let publicInputs: PrivateKernelTailCircuitPublicInputs | undefined;
         let executionSteps: PrivateExecutionStep[] = [];
@@ -990,7 +1020,7 @@ export class PXE {
     // We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another simulation is currently modifying).
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       try {
         const totalTimer = new Timer();
         const syncTimer = new Timer();
@@ -1000,10 +1030,16 @@ export class PXE {
         const contractFunctionSimulator = this.#getSimulatorForTx();
 
         await this.contractStore.syncPrivateState(call.to, call.selector, privateSyncCall =>
-          this.#simulateUtility(contractFunctionSimulator, privateSyncCall),
+          this.#simulateUtility(contractFunctionSimulator, privateSyncCall, [], undefined, jobId),
         );
 
-        const executionResult = await this.#simulateUtility(contractFunctionSimulator, call, authwits ?? [], scopes);
+        const executionResult = await this.#simulateUtility(
+          contractFunctionSimulator,
+          call,
+          authwits ?? [],
+          scopes,
+          jobId,
+        );
         const functionTime = functionTimer.ms();
 
         const totalTime = totalTimer.ms();
@@ -1045,14 +1081,15 @@ export class PXE {
    * @returns - The packed events with block and tx metadata.
    */
   public getPrivateEvents(eventSelector: EventSelector, filter: PrivateEventFilter): Promise<PackedPrivateEvent[]> {
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       await this.blockStateSynchronizer.sync();
       const contractFunctionSimulator = this.#getSimulatorForTx();
 
       await this.contractStore.syncPrivateState(
         filter.contractAddress,
         null,
-        async privateSyncCall => await this.#simulateUtility(contractFunctionSimulator, privateSyncCall),
+        async privateSyncCall =>
+          await this.#simulateUtility(contractFunctionSimulator, privateSyncCall, [], undefined, jobId),
       );
 
       const sanitizedFilter = await new PrivateEventFilterValidator(this.anchorBlockStore).validate(filter);

package/src/storage/private_event_store/private_event_store.ts

@@ -19,6 +19,7 @@ export type PrivateEventStoreFilter = {
 };
 
 type PrivateEventEntry = {
+  randomness: Fr; // Note that this value is currently not being returned on queries and is therefore temporarily unused
   msgContent: Buffer;
   eventCommitmentIndex: number;
   l2BlockNumber: number;
@@ -74,6 +75,7 @@ export class PrivateEventStore {
    */
   storePrivateEventLog(
     eventSelector: EventSelector,
+    randomness: Fr,
     msgContent: Fr[],
     eventCommitmentIndex: number,
     metadata: PrivateEventMetadata,
@@ -93,6 +95,7 @@ export class PrivateEventStore {
       this.logger.verbose('storing private event log', { contractAddress, scope, msgContent, l2BlockNumber });
 
       await this.#eventLogs.set(eventCommitmentIndex, {
+        randomness,
         msgContent: serializeToBuffer(msgContent),
         l2BlockNumber,
         l2BlockHash: l2BlockHash.toBuffer(),

package/src/storage/tagging_store/sender_tagging_store.ts

@@ -3,7 +3,7 @@ import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
 import type { DirectionalAppTaggingSecret, PreTag } from '@aztec/stdlib/logs';
 import { TxHash } from '@aztec/stdlib/tx';
 
-import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../../tagging/index.js';
+import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../../tagging/constants.js';
 
 /**
  * Data provider of tagging data used when syncing the sender tagging indexes. The recipient counterpart of this class

package/src/tagging/constants.ts

@@ -0,0 +1,10 @@
+// This window has to be as large as the largest expected number of logs emitted in a tx for a given directional app
+// tagging secret. If we get more tag indexes consumed than this window, an error is thrown in `PXE::proveTx` function.
+// This is set to a larger value than MAX_PRIVATE_LOGS_PER_TX (currently 64) because there could be more than
+// MAX_PRIVATE_LOGS_PER_TX indexes consumed in case the logs are squashed. This happens when the log contains a note
+// and the note is nullified in the same tx.
+//
+// Having a large window significantly slowed down `e2e_l1_with_wall_time` test as there we perform sync for more than
+// 1000 secrets. For this reason we set it to a relatively low value of 20. 20 should be sufficient for all the use
+// cases.
+export const UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN = 20;

package/src/tagging/index.ts

@@ -11,17 +11,7 @@
 
 export { loadPrivateLogsForSenderRecipientPair } from './recipient_sync/load_private_logs_for_sender_recipient_pair.js';
 export { syncSenderTaggingIndexes } from './sender_sync/sync_sender_tagging_indexes.js';
-
-// This window has to be as large as the largest expected number of logs emitted in a tx for a given directional app
-// tagging secret. If we get more tag indexes consumed than this window, an error is thrown in `PXE::proveTx` function.
-// This is set to a larger value than MAX_PRIVATE_LOGS_PER_TX (currently 64) because there could be more than
-// MAX_PRIVATE_LOGS_PER_TX indexes consumed in case the logs are squashed. This happens when the log contains a note
-// and the note is nullified in the same tx.
-//
-// Having a large window significantly slowed down `e2e_l1_with_wall_time` test as there we perform sync for more than
-// 1000 secrets. For this reason we set it to a relatively low value of 20. 20 should be sufficient for all the use
-// cases.
-export const UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN = 20;
+export { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from './constants.js';
 
 // Re-export tagging-related types from stdlib
 export { DirectionalAppTaggingSecret, Tag, SiloedTag } from '@aztec/stdlib/logs';

package/src/tagging/recipient_sync/load_private_logs_for_sender_recipient_pair.ts

@@ -4,7 +4,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import type { DirectionalAppTaggingSecret, TxScopedL2Log } from '@aztec/stdlib/logs';
 
 import type { RecipientTaggingStore } from '../../storage/tagging_store/recipient_tagging_store.js';
-import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../index.js';
+import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../constants.js';
 import { findHighestIndexes } from './utils/find_highest_indexes.js';
 import { loadLogsForRange } from './utils/load_logs_for_range.js';
 

package/src/tagging/sender_sync/sync_sender_tagging_indexes.ts

@@ -3,7 +3,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/server';
 import type { DirectionalAppTaggingSecret } from '@aztec/stdlib/logs';
 
 import type { SenderTaggingStore } from '../../storage/tagging_store/sender_tagging_store.js';
-import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../index.js';
+import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../constants.js';
 import { getStatusChangeOfPending } from './utils/get_status_change_of_pending.js';
 import { loadAndStoreNewTaggingIndexes } from './utils/load_and_store_new_tagging_indexes.js';