@aztec/prover-client 0.0.1-commit.96bb3f7 → 0.0.1-commit.993d240

package/src/orchestrator/top-tree-proving-state.ts

@@ -0,0 +1,220 @@
+import type { BatchedBlob, BatchedBlobAccumulator, FinalBlobBatchingChallenges } from '@aztec/blob-lib';
+import type { NESTED_RECURSIVE_PROOF_LENGTH, NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH } from '@aztec/constants';
+import { EpochNumber } from '@aztec/foundation/branded-types';
+import { type TreeNodeLocation, UnbalancedTreeStore } from '@aztec/foundation/trees';
+import type { PublicInputsAndRecursiveProof } from '@aztec/stdlib/interfaces/server';
+import type { Proof } from '@aztec/stdlib/proofs';
+import {
+  CheckpointMergeRollupPrivateInputs,
+  CheckpointPaddingRollupPrivateInputs,
+  CheckpointRollupPublicInputs,
+  RootRollupPrivateInputs,
+  type RootRollupPublicInputs,
+} from '@aztec/stdlib/rollup';
+
+import { toProofData } from './block-building-helpers.js';
+import type { ProofState } from './block-proving-state.js';
+
+enum TOP_TREE_LIFECYCLE {
+  CREATED,
+  RESOLVED,
+  REJECTED,
+}
+
+/**
+ * Lean top-tree-only state. Owns the merge tree of checkpoint root proofs, the
+ * single-checkpoint padding proof slot, the final root rollup proof, and the blob
+ * accumulator endpoints needed to finalise the epoch's batched blob proof.
+ *
+ * Constructed with `totalNumCheckpoints` and `finalBlobBatchingChallenges` upfront —
+ * by the time the top tree starts, all checkpoints are known and the challenges are
+ * derivable from their blob fields.
+ */
+export class TopTreeProvingState {
+  private checkpointProofs: UnbalancedTreeStore<
+    ProofState<CheckpointRollupPublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
+  >;
+  private checkpointPaddingProof:
+    | ProofState<CheckpointRollupPublicInputs, typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH>
+    | undefined;
+  private rootRollupProof: ProofState<RootRollupPublicInputs, typeof NESTED_RECURSIVE_PROOF_LENGTH> | undefined;
+  private endBlobAccumulator: BatchedBlobAccumulator | undefined;
+  private finalBatchedBlob: BatchedBlob | undefined;
+  private lifecycle = TOP_TREE_LIFECYCLE.CREATED;
+
+  constructor(
+    public readonly epochNumber: EpochNumber,
+    public readonly totalNumCheckpoints: number,
+    public readonly finalBlobBatchingChallenges: FinalBlobBatchingChallenges,
+    public readonly startBlobAccumulator: BatchedBlobAccumulator,
+    private readonly completionCallback: () => void,
+    private readonly rejectionCallback: (reason: string) => void,
+  ) {
+    if (totalNumCheckpoints < 1) {
+      throw new Error(`TopTreeProvingState requires at least one checkpoint; got ${totalNumCheckpoints}.`);
+    }
+    this.checkpointProofs = new UnbalancedTreeStore(totalNumCheckpoints);
+  }
+
+  // --- checkpoint root rollup ---
+
+  public setCheckpointRootRollupProof(
+    checkpointIndex: number,
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
+      typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
+    >,
+  ): TreeNodeLocation {
+    return this.checkpointProofs.setLeaf(checkpointIndex, { provingOutput });
+  }
+
+  // --- checkpoint merge rollup ---
+
+  public tryStartProvingCheckpointMerge(location: TreeNodeLocation) {
+    if (this.checkpointProofs.getNode(location)?.isProving) {
+      return false;
+    }
+    this.checkpointProofs.setNode(location, { isProving: true });
+    return true;
+  }
+
+  public setCheckpointMergeRollupProof(
+    location: TreeNodeLocation,
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
+      typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
+    >,
+  ) {
+    this.checkpointProofs.setNode(location, { provingOutput });
+  }
+
+  public isReadyForCheckpointMerge(location: TreeNodeLocation) {
+    return !!this.checkpointProofs.getSibling(location)?.provingOutput;
+  }
+
+  public getParentLocation(location: TreeNodeLocation) {
+    return this.checkpointProofs.getParentLocation(location);
+  }
+
+  public getCheckpointMergeRollupInputs(mergeLocation: TreeNodeLocation) {
+    const [left, right] = this.checkpointProofs.getChildren(mergeLocation).map(c => c?.provingOutput);
+    if (!left || !right) {
+      throw new Error('At least one child is not ready for the checkpoint merge rollup.');
+    }
+    return new CheckpointMergeRollupPrivateInputs([toProofData(left), toProofData(right)]);
+  }
+
+  // --- padding (single-checkpoint case) ---
+
+  public tryStartProvingPaddingCheckpoint() {
+    if (this.checkpointPaddingProof?.isProving) {
+      return false;
+    }
+    this.checkpointPaddingProof = { isProving: true };
+    return true;
+  }
+
+  public setCheckpointPaddingProof(
+    provingOutput: PublicInputsAndRecursiveProof<
+      CheckpointRollupPublicInputs,
+      typeof NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH
+    >,
+  ) {
+    this.checkpointPaddingProof = { provingOutput };
+  }
+
+  public getPaddingCheckpointInputs() {
+    return new CheckpointPaddingRollupPrivateInputs();
+  }
+
+  // --- root rollup ---
+
+  public tryStartProvingRootRollup() {
+    if (this.rootRollupProof?.isProving) {
+      return false;
+    }
+    this.rootRollupProof = { isProving: true };
+    return true;
+  }
+
+  public setRootRollupProof(provingOutput: PublicInputsAndRecursiveProof<RootRollupPublicInputs>) {
+    this.rootRollupProof = { provingOutput };
+  }
+
+  public isReadyForRootRollup() {
+    const childProofs = this.#getChildProofsForRoot();
+    return childProofs.every(p => !!p);
+  }
+
+  public getRootRollupInputs() {
+    const [left, right] = this.#getChildProofsForRoot();
+    if (!left || !right) {
+      throw new Error('At least one child is not ready for the root rollup.');
+    }
+    return RootRollupPrivateInputs.from({
+      previousRollups: [toProofData(left), toProofData(right)],
+    });
+  }
+
+  // --- blob accumulator finalisation ---
+
+  /**
+   * Sets the end-of-epoch blob accumulator, computed by the top-tree orchestrator
+   * from the surviving checkpoints' blob fields. Required before `finalizeBatchedBlob`.
+   */
+  public setEndBlobAccumulator(accumulator: BatchedBlobAccumulator) {
+    this.endBlobAccumulator = accumulator;
+  }
+
+  public async finalizeBatchedBlob() {
+    if (!this.endBlobAccumulator) {
+      throw new Error('End blob accumulator not set; call setEndBlobAccumulator before finalize.');
+    }
+    this.finalBatchedBlob = await this.endBlobAccumulator.finalize(true /* verifyProof */);
+  }
+
+  public getEpochProofResult(): { proof: Proof; publicInputs: RootRollupPublicInputs; batchedBlobInputs: BatchedBlob } {
+    const provingOutput = this.rootRollupProof?.provingOutput;
+    if (!provingOutput || !this.finalBatchedBlob) {
+      throw new Error('Top-tree proof not ready; root rollup or batched blob missing.');
+    }
+    return {
+      proof: provingOutput.proof.binaryProof,
+      publicInputs: provingOutput.inputs,
+      batchedBlobInputs: this.finalBatchedBlob,
+    };
+  }
+
+  // --- lifecycle ---
+
+  public verifyState() {
+    return this.lifecycle === TOP_TREE_LIFECYCLE.CREATED;
+  }
+
+  public resolve() {
+    if (!this.verifyState()) {
+      return;
+    }
+    this.lifecycle = TOP_TREE_LIFECYCLE.RESOLVED;
+    this.completionCallback();
+  }
+
+  public reject(reason: string) {
+    if (!this.verifyState()) {
+      return;
+    }
+    this.lifecycle = TOP_TREE_LIFECYCLE.REJECTED;
+    this.rejectionCallback(reason);
+  }
+
+  public cancel() {
+    this.reject('Top-tree proving cancelled');
+  }
+
+  #getChildProofsForRoot() {
+    const rootLocation = { level: 0, index: 0 };
+    return this.totalNumCheckpoints === 1
+      ? [this.checkpointProofs.getNode(rootLocation)?.provingOutput, this.checkpointPaddingProof?.provingOutput]
+      : this.checkpointProofs.getChildren(rootLocation).map(c => c?.provingOutput);
+  }
+}

package/src/prover-client/prover-client.ts

@@ -1,7 +1,9 @@
 import { type ACVMConfig, type BBConfig, BBNativeRollupProver, TestCircuitProver } from '@aztec/bb-prover';
+import type { EpochNumber } from '@aztec/foundation/branded-types';
 import { times } from '@aztec/foundation/collection';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { EthAddress } from '@aztec/foundation/eth-address';
-import { createLogger } from '@aztec/foundation/log';
+import { type Logger, createLogger } from '@aztec/foundation/log';
 import { NativeACVMSimulator } from '@aztec/simulator/server';
 import {
   type ActualProverConfig,
@@ -15,41 +17,166 @@ import {
   type ServerCircuitProver,
   tryStop,
 } from '@aztec/stdlib/interfaces/server';
+import type { CheckpointConstantData } from '@aztec/stdlib/rollup';
+import type { BlockHeader } from '@aztec/stdlib/tx';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
 import type { ProverClientConfig } from '../config.js';
+import { CheckpointSubTreeOrchestrator } from '../orchestrator/checkpoint-sub-tree-orchestrator.js';
+import { EpochProvingContext } from '../orchestrator/epoch-proving-context.js';
 import { ProvingOrchestrator } from '../orchestrator/orchestrator.js';
+import { TopTreeOrchestrator } from '../orchestrator/top-tree-orchestrator.js';
 import { BrokerCircuitProverFacade } from '../proving_broker/broker_prover_facade.js';
 import { InlineProofStore, type ProofStore, createProofStore } from '../proving_broker/proof_store/index.js';
 import { ProvingAgent } from '../proving_broker/proving_agent.js';
 import { ServerEpochProver } from './server-epoch-prover.js';
 
+/**
+ * The factory surface that `EpochProvingJob` (in `prover-node`) depends on. Implemented
+ * by `ProverClient`. Defined here rather than in stdlib because the return types
+ * (`CheckpointSubTreeOrchestrator`, `TopTreeOrchestrator`) are concrete classes from
+ * this package.
+ *
+ * A single `BrokerCircuitProverFacade` is owned by `ProverClient` and shared across
+ * every orchestrator (every sub-tree and every top-tree across every concurrent epoch
+ * job). The broker delivers each completed-job notification exactly once (drained on
+ * the first `getCompletedJobs` poll), so multiple facades polling the same broker
+ * race and lose notifications
+ *
+ * The facade's job map cleans up entries on resolve/reject, and the prover-node
+ * keeps `ProverClient` alive for its whole lifetime
+ */
+export interface EpochProverFactory {
+  getProverId(): EthAddress;
+  /**
+   * Constructs a per-epoch shared context for the caching of e.g. chonk verifier results
+   */
+  createEpochProvingContext(epochNumber: EpochNumber): EpochProvingContext;
+  /**
+   * Constructs and starts a `CheckpointSubTreeOrchestrator` for a single checkpoint.
+   */
+  createCheckpointSubTreeOrchestrator(
+    epochContext: EpochProvingContext,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator>;
+  createTopTreeOrchestrator(): TopTreeOrchestrator;
+}
+
 /** Manages proving of epochs by orchestrating the proving of individual blocks relying on a pool of prover agents. */
-export class ProverClient implements EpochProverManager {
+export class ProverClient implements EpochProverManager, EpochProverFactory {
   private running = false;
   private agents: ProvingAgent[] = [];
-
-  private proofStore: ProofStore;
-  private failedProofStore: ProofStore | undefined;
+  /**
+   * The single broker facade shared by every orchestrator created from this client.
+   * Constructed lazily on `start()` and torn down on `stop()` — see the comment on
+   * `EpochProverFactory` for why a single shared facade is required.
+   */
+  private facade: BrokerCircuitProverFacade | undefined;
 
   private constructor(
     private config: ProverClientConfig,
     private worldState: ForkMerkleTreeOperations & ReadonlyWorldStateAccess,
     private orchestratorClient: ProvingJobProducer,
+    private proofStore: ProofStore,
+    private failedProofStore: ProofStore | undefined,
     private agentClient?: ProvingJobConsumer,
     private telemetry: TelemetryClient = getTelemetryClient(),
-    private log = createLogger('prover-client:tx-prover'),
-  ) {
-    this.proofStore = new InlineProofStore();
-    this.failedProofStore = this.config.failedProofStore ? createProofStore(this.config.failedProofStore) : undefined;
+    private log: Logger = createLogger('prover-client:tx-prover'),
+  ) {}
+
+  /**
+   * Lazy-init the shared facade. The broker delivers each completed-job notification
+   * exactly once (drained on the first `getCompletedJobs` poll), so we cannot start
+   * a shared facade alongside the per-call facades that `createEpochProver` builds —
+   * they would race for notifications and one side would silently drop them. Starting
+   * the shared facade only on first use of one of the new factory methods keeps the
+   * legacy `createEpochProver` path race-free.
+   */
+  private getFacade(): BrokerCircuitProverFacade {
+    if (!this.running) {
+      throw new Error('ProverClient is not running; call start() before constructing orchestrators.');
+    }
+    if (!this.facade) {
+      this.facade = new BrokerCircuitProverFacade(
+        this.orchestratorClient,
+        this.proofStore,
+        this.failedProofStore,
+        undefined,
+        this.log.getBindings(),
+      );
+      this.facade.start();
+    }
+    return this.facade;
   }
 
+  /**
+   * Legacy single-class epoch prover. Each call constructs its own
+   * `BrokerCircuitProverFacade`; the new factory methods (`createCheckpointSubTreeOrchestrator`,
+   * `createTopTreeOrchestrator`, `createEpochProvingContext`) share a single facade
+   * owned by `ProverClient`. Both APIs coexist while the prover-node migrates onto
+   * the new pair.
+   */
   public createEpochProver(): EpochProver {
-    const facade = new BrokerCircuitProverFacade(this.orchestratorClient, this.proofStore, this.failedProofStore);
-    const orchestrator = new ProvingOrchestrator(this.worldState, facade, this.config.proverId, this.telemetry);
+    const bindings = this.log.getBindings();
+    const facade = new BrokerCircuitProverFacade(
+      this.orchestratorClient,
+      this.proofStore,
+      this.failedProofStore,
+      undefined,
+      bindings,
+    );
+    const orchestrator = new ProvingOrchestrator(
+      this.worldState,
+      facade,
+      this.config.proverId,
+      this.config.cancelJobsOnStop,
+      this.config.enqueueConcurrency,
+      this.telemetry,
+      bindings,
+    );
     return new ServerEpochProver(facade, orchestrator);
   }
 
+  public createEpochProvingContext(epochNumber: EpochNumber): EpochProvingContext {
+    return new EpochProvingContext(this.getFacade(), epochNumber, this.log.getBindings());
+  }
+
+  public createCheckpointSubTreeOrchestrator(
+    epochContext: EpochProvingContext,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator> {
+    return CheckpointSubTreeOrchestrator.start(
+      this.worldState,
+      this.getFacade(),
+      this.config.proverId,
+      epochContext,
+      this.config.cancelJobsOnStop,
+      this.config.enqueueConcurrency,
+      checkpointConstants,
+      l1ToL2Messages,
+      totalNumBlocks,
+      headerOfLastBlockInPreviousCheckpoint,
+      this.telemetry,
+      this.log.getBindings(),
+    );
+  }
+
+  public createTopTreeOrchestrator(): TopTreeOrchestrator {
+    return new TopTreeOrchestrator(
+      this.getFacade(),
+      this.config.proverId,
+      this.config.enqueueConcurrency,
+      this.telemetry,
+      this.log.getBindings(),
+    );
+  }
+
   public getProverId(): EthAddress {
     return this.config.proverId;
   }
@@ -89,6 +216,14 @@ export class ProverClient implements EpochProverManager {
     }
     this.running = false;
     await this.stopAgents();
+    if (this.facade) {
+      try {
+        await this.facade.stop();
+      } catch (err) {
+        this.log.error('Error stopping shared broker facade', err);
+      }
+      this.facade = undefined;
+    }
     await tryStop(this.orchestratorClient);
   }
 
@@ -104,7 +239,9 @@ export class ProverClient implements EpochProverManager {
     broker: ProvingJobBroker,
     telemetry: TelemetryClient = getTelemetryClient(),
   ) {
-    const prover = new ProverClient(config, worldState, broker, broker, telemetry);
+    const proofStore = await createProofStore(config.proofStore);
+    const failedProofStore = config.failedProofStore ? await createProofStore(config.failedProofStore) : undefined;
+    const prover = new ProverClient(config, worldState, broker, proofStore, failedProofStore, broker, telemetry);
     await prover.start();
     return prover;
   }
@@ -128,9 +265,11 @@ export class ProverClient implements EpochProverManager {
 
     const proofStore = new InlineProofStore();
     const prover = await buildServerCircuitProver(this.config, this.telemetry);
+    const bindings = this.log.getBindings();
     this.agents = times(
       this.config.proverAgentCount,
-      () => new ProvingAgent(this.agentClient!, proofStore, prover, [], this.config.proverAgentPollIntervalMs),
+      () =>
+        new ProvingAgent(this.agentClient!, proofStore, prover, [], this.config.proverAgentPollIntervalMs, bindings),
     );
 
     await Promise.all(this.agents.map(agent => agent.start()));
@@ -142,15 +281,16 @@ export class ProverClient implements EpochProverManager {
 }
 
 export function buildServerCircuitProver(
-  config: ActualProverConfig & ACVMConfig & BBConfig,
+  config: Omit<ActualProverConfig, 'enqueueConcurrency'> & ACVMConfig & BBConfig,
   telemetry: TelemetryClient,
 ): Promise<ServerCircuitProver> {
   if (config.realProofs) {
     return BBNativeRollupProver.new(config, telemetry);
   }
 
+  const logger = createLogger('prover-client:acvm-native');
   const simulator = config.acvmBinaryPath
-    ? new NativeACVMSimulator(config.acvmWorkingDirectory, config.acvmBinaryPath)
+    ? new NativeACVMSimulator(config.acvmWorkingDirectory, config.acvmBinaryPath, undefined, logger)
     : undefined;
 
   return Promise.resolve(new TestCircuitProver(simulator, config, telemetry));

package/src/proving_broker/broker_prover_facade.ts

@@ -4,9 +4,10 @@ import type {
   NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH,
   RECURSIVE_PROOF_LENGTH,
 } from '@aztec/constants';
+import { asyncPool } from '@aztec/foundation/async-pool';
 import { EpochNumber } from '@aztec/foundation/branded-types';
-import { sha256 } from '@aztec/foundation/crypto/sha256';
-import { createLogger } from '@aztec/foundation/log';
+import { chunk } from '@aztec/foundation/collection';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { type PromiseWithResolvers, RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
 import { truncate } from '@aztec/foundation/string';
 import type { AvmCircuitInputs } from '@aztec/stdlib/avm';
@@ -46,6 +47,8 @@ import type {
   TxRollupPublicInputs,
 } from '@aztec/stdlib/rollup';
 
+import { createHash } from 'node:crypto';
+
 import { InlineProofStore, type ProofStore } from './proof_store/index.js';
 
 // Perform a snapshot sync every 30 seconds
@@ -68,14 +71,17 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
   private runningPromise?: RunningPromise;
   private timeOfLastSnapshotSync = Date.now();
   private jobsToRetrieve: Set<ProvingJobId> = new Set();
+  private log: Logger;
 
   constructor(
     private broker: ProvingJobProducer,
     private proofStore: ProofStore = new InlineProofStore(),
     private failedProofStore?: ProofStore,
     private pollIntervalMs = 1000,
-    private log = createLogger('prover-client:broker-circuit-prover-facade'),
-  ) {}
+    bindings?: LoggerBindings,
+  ) {
+    this.log = createLogger('prover-client:broker-circuit-prover-facade', bindings);
+  }
 
   /**
    * This is a critical section. This function can not be async since it writes
@@ -222,17 +228,11 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
       // We collect all returned notifications and return them
       const allCompleted = new Set<ProvingJobId>();
       try {
-        let numRequests = 0;
-        while (ids.length > 0) {
-          const slice = ids.splice(0, SNAPSHOT_SYNC_CHECK_MAX_REQUEST_SIZE);
-          const completed = await this.broker.getCompletedJobs(slice);
+        const batches = ids.length > 0 ? chunk(ids, SNAPSHOT_SYNC_CHECK_MAX_REQUEST_SIZE) : [[]];
+        await asyncPool(1, batches, async batch => {
+          const completed = await this.broker.getCompletedJobs(batch);
           completed.forEach(id => allCompleted.add(id));
-          ++numRequests;
-        }
-        if (numRequests === 0) {
-          const final = await this.broker.getCompletedJobs([]);
-          final.forEach(id => allCompleted.add(id));
-        }
+        });
       } catch (err) {
         this.log.error(`Error thrown when requesting completed job notifications from the broker`, err);
       }
@@ -348,12 +348,8 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
       .map(id => this.jobs.get(id)!)
       .filter(x => x !== undefined);
     const totalJobsToRetrieve = toBeRetrieved.length;
-    let totalJobsRetrieved = 0;
-    while (toBeRetrieved.length > 0) {
-      const slice = toBeRetrieved.splice(0, MAX_CONCURRENT_JOB_SETTLED_REQUESTS);
-      const results = await Promise.all(slice.map(job => processJob(job!)));
-      totalJobsRetrieved += results.filter(x => x).length;
-    }
+    const results = await asyncPool(MAX_CONCURRENT_JOB_SETTLED_REQUESTS, toBeRetrieved, job => processJob(job));
+    const totalJobsRetrieved = results.filter(x => x).length;
     if (totalJobsToRetrieve > 0) {
       this.log.verbose(
         `Successfully retrieved ${totalJobsRetrieved} of ${totalJobsToRetrieve} jobs that should be ready, total ready jobs is now: ${this.jobsToRetrieve.size}`,
@@ -656,8 +652,12 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
     );
   }
 
-  private generateId(type: ProvingRequestType, inputs: { toBuffer(): Buffer }, epochNumber = EpochNumber.ZERO) {
-    const inputsHash = sha256(inputs.toBuffer());
-    return makeProvingJobId(epochNumber, type, inputsHash.toString('hex'));
+  private generateId(
+    type: ProvingRequestType,
+    inputs: { toBuffer(): Buffer },
+    epochNumber = EpochNumber.ZERO,
+  ): ProvingJobId {
+    const inputsHash = createHash('sha256').update(inputs.toBuffer()).digest('hex');
+    return makeProvingJobId(epochNumber, type, inputsHash);
   }
 }

package/src/proving_broker/config.ts

@@ -4,10 +4,11 @@ import {
   booleanConfigHelper,
   getDefaultConfig,
   numberConfigHelper,
+  optionalNumberConfigHelper,
 } from '@aztec/foundation/config';
 import { pickConfigMappings } from '@aztec/foundation/config';
-import { type DataStoreConfig, dataConfigMappings } from '@aztec/kv-store/config';
 import { type ChainConfig, chainConfigMappings } from '@aztec/stdlib/config';
+import { type DataStoreConfig, dataConfigMappings } from '@aztec/stdlib/kv-store';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
 
 import { z } from 'zod';
@@ -31,6 +32,8 @@ export const ProverBrokerConfig = z.object({
   proverBrokerBatchIntervalMs: z.number().int().nonnegative(),
   /** The maximum number of epochs to keep results for */
   proverBrokerMaxEpochsToKeepResultsFor: z.number().int().nonnegative(),
+  /** Enable debug replay mode for replaying proving jobs from stored inputs */
+  proverBrokerDebugReplayEnabled: z.boolean(),
 });
 
 export type ProverBrokerConfig = z.infer<typeof ProverBrokerConfig> &
@@ -71,9 +74,14 @@ export const proverBrokerConfigMappings: ConfigMappingsType<ProverBrokerConfig>
   },
   proverBrokerStoreMapSizeKb: {
     env: 'PROVER_BROKER_STORE_MAP_SIZE_KB',
-    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    ...optionalNumberConfigHelper(),
     description: "The size of the prover broker's database. Will override the dataStoreMapSizeKb if set.",
   },
+  proverBrokerDebugReplayEnabled: {
+    env: 'PROVER_BROKER_DEBUG_REPLAY_ENABLED',
+    description: 'Enable debug replay mode for replaying proving jobs from stored inputs',
+    ...booleanConfigHelper(false),
+  },
   ...dataConfigMappings,
   ...l1ReaderConfigMappings,
   ...pickConfigMappings(chainConfigMappings, ['rollupVersion']),
@@ -100,6 +108,10 @@ export const ProverAgentConfig = z.object({
   proverTestDelayFactor: z.number(),
   /** The delay (ms) to inject during fake proof verification */
   proverTestVerificationDelayMs: z.number().optional(),
+  /** Whether to abort pending proving jobs when the orchestrator is cancelled */
+  cancelJobsOnStop: z.boolean(),
+  /** Where to store proving results. Must be accessible to both prover node and agents. If not set will inline-encode the parameters */
+  proofStore: z.string().optional(),
 });
 
 export type ProverAgentConfig = z.infer<typeof ProverAgentConfig>;
@@ -153,4 +165,15 @@ export const proverAgentConfigMappings: ConfigMappingsType<ProverAgentConfig> =
     description: 'The delay (ms) to inject during fake proof verification',
     ...numberConfigHelper(10),
   },
+  cancelJobsOnStop: {
+    env: 'PROVER_CANCEL_JOBS_ON_STOP',
+    description:
+      'Whether to abort pending proving jobs when the orchestrator is cancelled. ' +
+      'When false (default), jobs remain in the broker queue and can be reused on restart/reorg.',
+    ...booleanConfigHelper(false),
+  },
+  proofStore: {
+    env: 'PROVER_PROOF_STORE',
+    description: 'Optional proof input store for the prover',
+  },
 };

package/src/proving_broker/index.ts

@@ -7,3 +7,4 @@ export * from './proving_broker_database/persisted.js';
 export * from './proof_store/index.js';
 export * from './factory.js';
 export * from './config.js';
+export { BrokerCircuitProverFacade } from './broker_prover_facade.js';

package/src/proving_broker/proof_store/factory.ts

@@ -1,42 +1,20 @@
 import { createLogger } from '@aztec/foundation/log';
+import { createFileStore } from '@aztec/stdlib/file-store';
 
-import { GoogleCloudStorageProofStore } from './gcs_proof_store.js';
+import { FileStoreProofStore } from './file_store_proof_store.js';
 import { InlineProofStore } from './inline_proof_store.js';
 import type { ProofStore } from './proof_store.js';
 
-export function createProofStore(config: string | undefined, logger = createLogger('prover-client:proof-store')) {
-  if (config === undefined) {
+export async function createProofStore(
+  config: string | undefined,
+  logger = createLogger('prover-client:proof-store'),
+): Promise<ProofStore> {
+  if (!config) {
     logger.info('Creating inline proof store');
     return new InlineProofStore();
-  } else if (config.startsWith('gs://')) {
-    try {
-      const url = new URL(config);
-      const bucket = url.host;
-      const path = url.pathname.replace(/^\/+/, '');
-      logger.info(`Creating google cloud proof store at ${bucket}`, { bucket, path });
-      return new GoogleCloudStorageProofStore(bucket, path);
-    } catch {
-      throw new Error(
-        `Invalid google cloud proof store definition: '${config}'. Supported values are 'gs://bucket-name/path/to/store'.`,
-      );
-    }
-  } else {
-    throw new Error(`Unknown proof store config: '${config}'. Supported values are 'gs://bucket-name/path/to/store'.`);
   }
-}
 
-export function createProofStoreForUri(
-  uri: string,
-  logger = createLogger('prover-client:proof-store'),
-): Pick<ProofStore, 'getProofInput' | 'getProofOutput'> {
-  if (uri.startsWith('data://')) {
-    return createProofStore(undefined, logger);
-  } else if (uri.startsWith('gs://')) {
-    const url = new URL(uri);
-    const basePath = url.pathname.replace(/^\/+/, '').split('/').slice(0, -3);
-    url.pathname = basePath.join('/');
-    return createProofStore(uri, logger);
-  } else {
-    throw new Error(`Unknown proof store config: '${uri}'. Supported protocols are 'data://' and 'gs://'.`);
-  }
+  const fileStore = await createFileStore(config, logger);
+  logger.info(`Creating file store proof store at ${config}`);
+  return new FileStoreProofStore(fileStore);
 }