@aztec/prover-client 0.0.1-commit.d3ec352c → 0.0.1-commit.d58ff9d0

package/src/proving_broker/proving_agent.ts

@@ -1,5 +1,5 @@
 import { AbortError } from '@aztec/foundation/error';
-import { createLogger } from '@aztec/foundation/log';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { RunningPromise } from '@aztec/foundation/running-promise';
 import { truncate } from '@aztec/foundation/string';
 import { ProvingError } from '@aztec/stdlib/errors';
@@ -13,27 +13,17 @@ import type {
   ServerCircuitProver,
 } from '@aztec/stdlib/interfaces/server';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
-import {
-  type TelemetryClient,
-  type Traceable,
-  type Tracer,
-  getTelemetryClient,
-  trackSpan,
-} from '@aztec/telemetry-client';
 
 import type { ProofStore } from './proof_store/index.js';
-import { ProvingAgentInstrumentation } from './proving_agent_instrumentation.js';
 import { ProvingJobController, ProvingJobControllerStatus } from './proving_job_controller.js';
 
 /**
  * A helper class that encapsulates a circuit prover and connects it to a job source.
  */
-export class ProvingAgent implements Traceable {
+export class ProvingAgent {
   private currentJobController?: ProvingJobController;
   private runningPromise: RunningPromise;
-  private instrumentation: ProvingAgentInstrumentation;
-
-  public readonly tracer: Tracer;
+  private log: Logger;
 
   constructor(
     /** The source of proving jobs */
@@ -46,12 +36,9 @@ export class ProvingAgent implements Traceable {
     private proofAllowList: Array<ProvingRequestType> = [],
     /** How long to wait between jobs */
     private pollIntervalMs = 1000,
-    /** A telemetry client through which to emit metrics */
-    client: TelemetryClient = getTelemetryClient(),
-    private log = createLogger('prover-client:proving-agent'),
+    bindings?: LoggerBindings,
   ) {
-    this.tracer = client.getTracer('ProvingAgent');
-    this.instrumentation = new ProvingAgentInstrumentation(client);
+    this.log = createLogger('prover-client:proving-agent', bindings);
     this.runningPromise = new RunningPromise(this.work.bind(this), this.log, this.pollIntervalMs);
   }
 
@@ -85,7 +72,6 @@ export class ProvingAgent implements Traceable {
     return this.runningPromise.isRunning() ? { status: 'running' } : { status: 'stopped' };
   }
 
-  @trackSpan('ProvingAgent.safeWork')
   private async work() {
     // every tick we need to take one of the following actions:
     // 1. send a hearbeat to the broker that we're working on some job
@@ -175,6 +161,7 @@ export class ProvingAgent implements Traceable {
         // no need to await this here. The controller will stay alive (in DONE state) until the result is send to the broker
         void this.runningPromise.trigger();
       },
+      this.log.getBindings(),
     );
 
     if (abortedProofJobId) {

package/src/proving_broker/proving_broker.ts

@@ -7,6 +7,7 @@ import {
   type GetProvingJobResponse,
   type ProofUri,
   type ProvingJob,
+  type ProvingJobBrokerDebug,
   type ProvingJobConsumer,
   type ProvingJobFilter,
   type ProvingJobId,
@@ -16,13 +17,7 @@ import {
   tryStop,
 } from '@aztec/stdlib/interfaces/server';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
-import {
-  type TelemetryClient,
-  type Traceable,
-  type Tracer,
-  getTelemetryClient,
-  trackSpan,
-} from '@aztec/telemetry-client';
+import { type TelemetryClient, type Traceable, type Tracer, getTelemetryClient } from '@aztec/telemetry-client';
 
 import assert from 'assert';
 
@@ -42,7 +37,7 @@ type EnqueuedProvingJob = Pick<ProvingJob, 'id' | 'epochNumber'>;
  * A broker that manages proof requests and distributes them to workers based on their priority.
  * It takes a backend that is responsible for storing and retrieving proof requests and results.
  */
-export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Traceable {
+export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, ProvingJobBrokerDebug, Traceable {
   private queues: ProvingQueues = {
     [ProvingRequestType.PUBLIC_VM]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
     [ProvingRequestType.PUBLIC_CHONK_VERIFIER]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
@@ -120,6 +115,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
   private started = false;
 
+  private debugReplayEnabled: boolean;
+
   public constructor(
     private database: ProvingBrokerDatabase,
     {
@@ -127,6 +124,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       proverBrokerPollIntervalMs,
       proverBrokerJobMaxRetries,
       proverBrokerMaxEpochsToKeepResultsFor,
+      proverBrokerDebugReplayEnabled,
     }: Required<
       Pick<
         ProverBrokerConfig,
@@ -134,6 +132,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
         | 'proverBrokerPollIntervalMs'
         | 'proverBrokerJobMaxRetries'
         | 'proverBrokerMaxEpochsToKeepResultsFor'
+        | 'proverBrokerDebugReplayEnabled'
       >
     > = defaultProverBrokerConfig,
     client: TelemetryClient = getTelemetryClient(),
@@ -145,6 +144,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     this.jobTimeoutMs = proverBrokerJobTimeoutMs!;
     this.maxRetries = proverBrokerJobMaxRetries!;
     this.maxEpochsToKeepResultsFor = proverBrokerMaxEpochsToKeepResultsFor!;
+    this.debugReplayEnabled = proverBrokerDebugReplayEnabled ?? false;
   }
 
   private measureQueueDepth: MonitorCallback = (type: ProvingRequestType) => {
@@ -247,6 +247,29 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return Promise.resolve(this.#reportProvingJobProgress(id, startedAt, filter));
   }
 
+  public async replayProvingJob(
+    jobId: ProvingJobId,
+    type: ProvingRequestType,
+    epochNumber: EpochNumber,
+    inputsUri: ProofUri,
+  ): Promise<ProvingJobStatus> {
+    if (!this.debugReplayEnabled) {
+      throw new Error('Debug replay not enabled. Set PROVER_BROKER_DEBUG_REPLAY_ENABLED=true');
+    }
+
+    this.logger.info(`Replaying proving job`, { provingJobId: jobId, epochNumber, inputsUri });
+
+    // Clear existing state and enqueue
+    this.cleanUpProvingJobState([jobId]);
+
+    const job: ProvingJob = { id: jobId, type, epochNumber, inputsUri };
+    this.jobsCache.set(jobId, job);
+    await this.database.addProvingJob(job);
+    this.enqueueJobInternal(job);
+
+    return { status: 'in-queue' };
+  }
+
   async #enqueueProvingJob(job: ProvingJob): Promise<ProvingJobStatus> {
     // We return the job status at the start of this call
     const jobStatus = this.#getProvingJobStatus(job.id);
@@ -291,18 +314,25 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     // notify listeners of the cancellation
     if (!this.resultsCache.has(id)) {
       this.logger.info(`Cancelling job id=${id}`, { provingJobId: id });
-      await this.#reportProvingJobError(id, 'Aborted', false);
+      await this.#reportProvingJobError(id, 'Aborted', false, undefined, true);
     }
   }
 
   private cleanUpProvingJobState(ids: ProvingJobId[]) {
+    const idsToClean = new Set(ids);
     for (const id of ids) {
       this.jobsCache.delete(id);
+      const deferred = this.promises.get(id);
+      if (deferred) {
+        deferred.resolve({ status: 'rejected', reason: 'Proving job cleaned up' });
+      }
       this.promises.delete(id);
       this.resultsCache.delete(id);
       this.inProgress.delete(id);
       this.retries.delete(id);
+      this.enqueuedAt.delete(id);
     }
+    this.completedJobNotifications = this.completedJobNotifications.filter(id => !idsToClean.has(id));
   }
 
   #getProvingJobStatus(id: ProvingJobId): ProvingJobStatus {
@@ -354,6 +384,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
           const enqueuedAt = this.enqueuedAt.get(job.id);
           if (enqueuedAt) {
             this.instrumentation.recordJobWait(job.type, enqueuedAt);
+            // we can clear this flag now.
+            this.enqueuedAt.delete(job.id);
           }
 
           return { job, time };
@@ -369,6 +401,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     err: string,
     retry = false,
     filter?: ProvingJobFilter,
+    aborted = false,
   ): Promise<GetProvingJobResponse | undefined> {
     const info = this.inProgress.get(id);
     const item = this.jobsCache.get(id);
@@ -429,7 +462,11 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     this.promises.get(id)!.resolve(result);
     this.completedJobNotifications.push(id);
 
-    this.instrumentation.incRejectedJobs(item.type);
+    if (aborted) {
+      this.instrumentation.incAbortedJobs(item.type);
+    } else {
+      this.instrumentation.incRejectedJobs(item.type);
+    }
     if (info) {
       const duration = this.msTimeSource() - info.startedAt;
       this.instrumentation.recordJobDuration(item.type, duration);
@@ -562,23 +599,22 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return this.#getProvingJob(filter);
   }
 
-  @trackSpan('ProvingBroker.cleanupPass')
   private async cleanupPass() {
-    this.cleanupStaleJobs();
     this.reEnqueueExpiredJobs();
     const oldestEpochToKeep = this.oldestEpochToKeep();
     if (oldestEpochToKeep > 0) {
+      this.cleanupJobsOlderThanEpoch(EpochNumber(oldestEpochToKeep));
       await this.database.deleteAllProvingJobsOlderThanEpoch(EpochNumber(oldestEpochToKeep));
       this.logger.trace(`Deleted all epochs older than ${oldestEpochToKeep}`);
     }
   }
 
-  private cleanupStaleJobs() {
+  private cleanupJobsOlderThanEpoch(epochNumber: EpochNumber) {
     const jobIds = Array.from(this.jobsCache.keys());
     const jobsToClean: ProvingJobId[] = [];
     for (const id of jobIds) {
       const job = this.jobsCache.get(id)!;
-      if (this.isJobStale(job)) {
+      if (job.epochNumber < epochNumber) {
         jobsToClean.push(id);
       }
     }
@@ -602,10 +638,26 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       const now = this.msTimeSource();
       const msSinceLastUpdate = now - metadata.lastUpdatedAt;
       if (msSinceLastUpdate >= this.jobTimeoutMs) {
-        this.logger.warn(`Proving job id=${id} timed out. Adding it back to the queue.`, { provingJobId: id });
         this.inProgress.delete(id);
-        this.enqueueJobInternal(item);
         this.instrumentation.incTimedOutJobs(item.type);
+
+        const retries = this.retries.get(id) ?? 0;
+        if (retries + 1 < this.maxRetries && !this.isJobStale(item)) {
+          this.logger.warn(`Proving job id=${id} timed out. Re-enqueueing (retry ${retries + 1}/${this.maxRetries}).`, {
+            provingJobId: id,
+          });
+          this.retries.set(id, retries + 1);
+          this.enqueueJobInternal(item);
+        } else {
+          this.logger.error(`Proving job id=${id} timed out after ${retries + 1} attempts. Marking as failed.`, {
+            provingJobId: id,
+          });
+          const result: ProvingJobSettledResult = { status: 'rejected', reason: 'Timed out' };
+          this.resultsCache.set(id, result);
+          this.promises.get(id)?.resolve(result);
+          this.completedJobNotifications.push(id);
+          this.instrumentation.incRejectedJobs(item.type);
+        }
       }
     }
   }

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -11,7 +11,14 @@ import {
   ProvingJobSettledResult,
   getEpochFromProvingJobId,
 } from '@aztec/stdlib/interfaces/server';
-import { Attributes, LmdbMetrics, type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
+import {
+  Attributes,
+  LmdbMetrics,
+  type TelemetryClient,
+  type Tracer,
+  getTelemetryClient,
+  trackSpan,
+} from '@aztec/telemetry-client';
 
 import { mkdir, readdir } from 'fs/promises';
 import { join } from 'path';
@@ -78,6 +85,8 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
 
   private batchQueue: BatchQueue<ProvingJob | [ProvingJobId, ProvingJobSettledResult], number>;
 
+  public readonly tracer: Tracer;
+
   private constructor(
     private epochs: Map<number, SingleEpochDatabase>,
     private config: ProverBrokerConfig,
@@ -92,6 +101,8 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       () => this.estimateSize(),
     );
 
+    this.tracer = client.getTracer('KVBrokerDatabase');
+
     this.batchQueue = new BatchQueue(
       (items, key) => this.commitWrites(items, key),
       config.proverBrokerBatchSize,
@@ -143,7 +154,7 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       const db = await openVersionedStoreAt(
         fullDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
-        config.l1Contracts.rollupAddress,
+        config.rollupAddress,
         config.dataStoreMapSizeKb,
       );
       const epochDb = new SingleEpochDatabase(db);
@@ -165,6 +176,9 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     }
   }
 
+  @trackSpan('KVBrokerDatabase.deleteAllProvingJobsOlderThanEpoch', epochNumber => ({
+    [Attributes.EPOCH_NUMBER]: epochNumber,
+  }))
   async deleteAllProvingJobsOlderThanEpoch(epochNumber: EpochNumber): Promise<void> {
     const oldEpochs = Array.from(this.epochs.keys()).filter(e => e < Number(epochNumber));
     for (const old of oldEpochs) {
@@ -208,7 +222,7 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       const db = await openVersionedStoreAt(
         newEpochDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
-        this.config.l1Contracts.rollupAddress,
+        this.config.rollupAddress,
         this.config.dataStoreMapSizeKb,
       );
       epochDb = new SingleEpochDatabase(db);

package/src/proving_broker/proving_broker_instrumentation.ts

@@ -8,7 +8,7 @@ import {
   type ObservableResult,
   type TelemetryClient,
   type UpDownCounter,
-  ValueType,
+  createUpDownCounterWithDefault,
 } from '@aztec/telemetry-client';
 
 export type MonitorCallback = (proofType: ProvingRequestType) => number;
@@ -18,6 +18,7 @@ export class ProvingBrokerInstrumentation {
   private activeJobs: ObservableGauge;
   private resolvedJobs: UpDownCounter;
   private rejectedJobs: UpDownCounter;
+  private abortedJobs: UpDownCounter;
   private timedOutJobs: UpDownCounter;
   private cachedJobs: UpDownCounter;
   private totalJobs: UpDownCounter;
@@ -28,49 +29,30 @@ export class ProvingBrokerInstrumentation {
   constructor(client: TelemetryClient, name = 'ProvingBroker') {
     const meter = client.getMeter(name);
 
-    this.queueSize = meter.createObservableGauge(Metrics.PROVING_QUEUE_SIZE, {
-      valueType: ValueType.INT,
-    });
+    this.queueSize = meter.createObservableGauge(Metrics.PROVING_QUEUE_SIZE);
 
-    this.activeJobs = meter.createObservableGauge(Metrics.PROVING_QUEUE_ACTIVE_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.activeJobs = meter.createObservableGauge(Metrics.PROVING_QUEUE_ACTIVE_JOBS);
 
-    this.resolvedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_RESOLVED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    const provingJobTypes = Object.values(ProvingRequestType).filter(v => typeof v === 'string');
+    const provingJobAttrs = { [Attributes.PROVING_JOB_TYPE]: provingJobTypes };
 
-    this.rejectedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_REJECTED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.resolvedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RESOLVED_JOBS, provingJobAttrs);
 
-    this.retriedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_RETRIED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.rejectedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_REJECTED_JOBS, provingJobAttrs);
 
-    this.timedOutJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.abortedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_ABORTED_JOBS, provingJobAttrs);
 
-    this.cachedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_CACHED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.retriedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RETRIED_JOBS, provingJobAttrs);
 
-    this.totalJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_TOTAL_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.timedOutJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, provingJobAttrs);
 
-    this.jobWait = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_WAIT, {
-      description: 'Records how long a job sits in the queue',
-      unit: 'ms',
-      valueType: ValueType.INT,
-    });
+    this.cachedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_CACHED_JOBS, provingJobAttrs);
 
-    this.jobDuration = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_DURATION, {
-      description: 'Records how long a job takes to complete',
-      unit: 'ms',
-      valueType: ValueType.INT,
-    });
+    this.totalJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TOTAL_JOBS, provingJobAttrs);
+
+    this.jobWait = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_WAIT);
+
+    this.jobDuration = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_DURATION);
   }
 
   monitorQueueDepth(fn: MonitorCallback) {
@@ -93,6 +75,12 @@ export class ProvingBrokerInstrumentation {
     });
   }
 
+  incAbortedJobs(proofType: ProvingRequestType) {
+    this.abortedJobs.add(1, {
+      [Attributes.PROVING_JOB_TYPE]: ProvingRequestType[proofType],
+    });
+  }
+
   incRetriedJobs(proofType: ProvingRequestType) {
     this.retriedJobs.add(1, {
       [Attributes.PROVING_JOB_TYPE]: ProvingRequestType[proofType],

package/src/proving_broker/proving_job_controller.ts

@@ -1,7 +1,7 @@
 import { EpochNumber } from '@aztec/foundation/branded-types';
-import { randomBytes } from '@aztec/foundation/crypto';
+import { randomBytes } from '@aztec/foundation/crypto/random';
 import { AbortError } from '@aztec/foundation/error';
-import { createLogger } from '@aztec/foundation/log';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import type {
   ProvingJobId,
   ProvingJobInputs,
@@ -21,6 +21,7 @@ export class ProvingJobController {
   private promise?: Promise<void>;
   private abortController = new AbortController();
   private result?: ProvingJobResultsMap[ProvingRequestType] | Error;
+  private log: Logger;
 
   constructor(
     private jobId: ProvingJobId,
@@ -29,8 +30,13 @@ export class ProvingJobController {
     private startedAt: number,
     private circuitProver: ServerCircuitProver,
     private onComplete: () => void,
-    private log = createLogger('prover-client:proving-agent:job-controller-' + randomBytes(4).toString('hex')),
-  ) {}
+    bindings?: LoggerBindings,
+  ) {
+    this.log = createLogger('prover-client:proving-agent:job-controller', {
+      instanceId: randomBytes(4).toString('hex'),
+      ...bindings,
+    });
+  }
 
   public start(): void {
     if (this.status !== ProvingJobControllerStatus.IDLE) {
@@ -125,8 +131,7 @@ export class ProvingJobController {
     const signal = this.abortController.signal;
     switch (type) {
       case ProvingRequestType.PUBLIC_VM: {
-        // TODO(#14234)[Unconditional PIs validation]: Remove argument "undefined".
-        return await this.circuitProver.getAvmProof(inputs, undefined, signal, this.epochNumber);
+        return await this.circuitProver.getAvmProof(inputs, signal, this.epochNumber);
       }
 
       case ProvingRequestType.PUBLIC_CHONK_VERIFIER: {

package/src/proving_broker/rpc.ts

@@ -1,9 +1,11 @@
+import { EpochNumberSchema } from '@aztec/foundation/branded-types';
 import { createSafeJsonRpcClient } from '@aztec/foundation/json-rpc/client';
 import {
   type GetProvingJobResponse,
   ProofUri,
   ProvingJob,
   type ProvingJobBroker,
+  type ProvingJobBrokerDebug,
   type ProvingJobConsumer,
   ProvingJobId,
   type ProvingJobProducer,
@@ -16,6 +18,15 @@ import { makeTracedFetch } from '@aztec/telemetry-client';
 
 import { z } from 'zod';
 
+/** Indefinite backoff for broker communication: 1, 1, 1, 2, 4, 4, 4, ... seconds. */
+export function* proverBrokerBackoff() {
+  const v = [1, 1, 1, 2, 4];
+  let i = 0;
+  while (true) {
+    yield v[Math.min(i++, v.length - 1)];
+  }
+}
+
 const ProvingJobFilterSchema = z.object({
   allowList: z.array(z.nativeEnum(ProvingRequestType)),
 });
@@ -26,26 +37,29 @@ const GetProvingJobResponse = z.object({
 });
 
 export const ProvingJobProducerSchema: ApiSchemaFor<ProvingJobProducer> = {
-  enqueueProvingJob: z.function().args(ProvingJob).returns(ProvingJobStatus),
-  getProvingJobStatus: z.function().args(ProvingJobId).returns(ProvingJobStatus),
-  cancelProvingJob: z.function().args(ProvingJobId).returns(z.void()),
-  getCompletedJobs: z.function().args(z.array(ProvingJobId)).returns(z.array(ProvingJobId)),
+  enqueueProvingJob: z.function({ input: z.tuple([ProvingJob]), output: ProvingJobStatus }),
+  getProvingJobStatus: z.function({ input: z.tuple([ProvingJobId]), output: ProvingJobStatus }),
+  cancelProvingJob: z.function({ input: z.tuple([ProvingJobId]), output: z.void() }),
+  getCompletedJobs: z.function({ input: z.tuple([z.array(ProvingJobId)]), output: z.array(ProvingJobId) }),
 };
 
 export const ProvingJobConsumerSchema: ApiSchemaFor<ProvingJobConsumer> = {
-  getProvingJob: z.function().args(optional(ProvingJobFilterSchema)).returns(GetProvingJobResponse.optional()),
-  reportProvingJobError: z
-    .function()
-    .args(ProvingJobId, z.string(), optional(z.boolean()), optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
-  reportProvingJobProgress: z
-    .function()
-    .args(ProvingJobId, z.number(), optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
-  reportProvingJobSuccess: z
-    .function()
-    .args(ProvingJobId, ProofUri, optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
+  getProvingJob: z.function({
+    input: z.tuple([optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobError: z.function({
+    input: z.tuple([ProvingJobId, z.string(), optional(z.boolean()), optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobProgress: z.function({
+    input: z.tuple([ProvingJobId, z.number(), optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobSuccess: z.function({
+    input: z.tuple([ProvingJobId, ProofUri, optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
 };
 
 export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
@@ -53,10 +67,22 @@ export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
   ...ProvingJobProducerSchema,
 };
 
+export const ProvingJobBrokerDebugSchema: ApiSchemaFor<ProvingJobBrokerDebug> = {
+  replayProvingJob: z.function({
+    input: z.tuple([ProvingJobId, z.nativeEnum(ProvingRequestType), EpochNumberSchema, ProofUri]),
+    output: ProvingJobStatus,
+  }),
+};
+
+export const ProvingJobBrokerSchemaWithDebug: ApiSchemaFor<ProvingJobBroker & ProvingJobBrokerDebug> = {
+  ...ProvingJobBrokerSchema,
+  ...ProvingJobBrokerDebugSchema,
+};
+
 export function createProvingJobBrokerClient(
   url: string,
   versions: Partial<ComponentsVersions>,
-  fetch = makeTracedFetch([1, 2, 3], false),
+  fetch = makeTracedFetch(proverBrokerBackoff, false),
 ): ProvingJobBroker {
   return createSafeJsonRpcClient(url, ProvingJobBrokerSchema, {
     namespaceMethods: 'proverBroker',
@@ -68,7 +94,7 @@ export function createProvingJobBrokerClient(
 export function createProvingJobProducerClient(
   url: string,
   versions: Partial<ComponentsVersions>,
-  fetch = makeTracedFetch([1, 2, 3], false),
+  fetch = makeTracedFetch(proverBrokerBackoff, false),
 ): ProvingJobProducer {
   return createSafeJsonRpcClient(url, ProvingJobProducerSchema, {
     namespaceMethods: 'provingJobProducer',
@@ -80,7 +106,7 @@ export function createProvingJobProducerClient(
 export function createProvingJobConsumerClient(
   url: string,
   versions: Partial<ComponentsVersions>,
-  fetch = makeTracedFetch([1, 2, 3], false),
+  fetch = makeTracedFetch(proverBrokerBackoff, false),
 ): ProvingJobConsumer {
   return createSafeJsonRpcClient(url, ProvingJobConsumerSchema, {
     namespaceMethods: 'provingJobConsumer',

package/src/test/epoch_settlement.ts

@@ -0,0 +1,82 @@
+import type { RollupCheatCodes } from '@aztec/ethereum/test';
+import type { CheckpointNumber, EpochNumber, SlotNumber } from '@aztec/foundation/branded-types';
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import type { Logger } from '@aztec/foundation/log';
+import type { L2BlockSource } from '@aztec/stdlib/block';
+import { computeEpochOutHash } from '@aztec/stdlib/messaging';
+
+/** Arguments for {@link settleEpochOutbox}. */
+export type SettleEpochOutboxArgs = {
+  /** Cheat codes used to write the computed out hash directly into the L1 Outbox storage. */
+  rollupCheatCodes: RollupCheatCodes;
+  /** Source of checkpointed L2 blocks for the epoch being settled. */
+  l2BlockSource: L2BlockSource;
+  /** Epoch to settle. */
+  epoch: EpochNumber;
+  /**
+   * When set, only checkpoints up to and including this number are covered. Used to settle a partial
+   * epoch (the AZIP-14 Outbox keeps one root per `numCheckpointsInEpoch`, so a prefix of an epoch can
+   * be settled and consumed before the epoch completes). When omitted, every checkpointed block in the
+   * epoch is covered.
+   */
+  maxCheckpoint?: CheckpointNumber;
+  log: Logger;
+};
+
+/**
+ * Computes the epoch out hash over the checkpointed blocks of an epoch (optionally only up to
+ * `maxCheckpoint`) and writes it into the L1 Outbox via cheat codes. This is the synthetic,
+ * no-prover equivalent of an epoch root proof landing on L1: it makes the L2-to-L1 messages in the
+ * covered checkpoints consumable. It does NOT advance the rollup's proven tip — callers
+ * `markAsProven` separately so a single prove call can settle multiple epochs before marking proven.
+ *
+ * Used by the test-only proving drivers: the `AutomineSequencer`'s auto-settle loop and the standalone
+ * `EpochTestSettler`. Lives here (rather than in either consumer) so both share one implementation.
+ *
+ * @returns The last checkpoint number covered by the settled range, or `undefined` if the epoch has
+ * no checkpointed blocks (within the `maxCheckpoint` bound).
+ */
+export async function settleEpochOutbox({
+  rollupCheatCodes,
+  l2BlockSource,
+  epoch,
+  maxCheckpoint,
+  log,
+}: SettleEpochOutboxArgs): Promise<CheckpointNumber | undefined> {
+  let blocks = await l2BlockSource.getBlocks({ epoch, onlyCheckpointed: true });
+  if (maxCheckpoint !== undefined) {
+    blocks = blocks.filter(block => block.checkpointNumber <= maxCheckpoint);
+  }
+  if (blocks.length === 0) {
+    return undefined;
+  }
+
+  log.info(
+    `Settling epoch ${epoch}${maxCheckpoint !== undefined ? ` up to checkpoint ${maxCheckpoint}` : ''} with blocks ${blocks[0]?.header.getBlockNumber()} to ${blocks.at(-1)?.header.getBlockNumber()}`,
+    { epoch, maxCheckpoint, blocks: blocks.map(block => block.toBlockInfo()) },
+  );
+
+  const messagesInEpoch: Fr[][][][] = [];
+  // Undefined (not SlotNumber.ZERO) so a first checkpointed block at slot 0 still opens checkpoint 0.
+  let previousSlotNumber: SlotNumber | undefined;
+  let checkpointIndex = -1;
+
+  for (const block of blocks) {
+    const slotNumber = block.header.globalVariables.slotNumber;
+    if (slotNumber !== previousSlotNumber) {
+      checkpointIndex++;
+      messagesInEpoch[checkpointIndex] = [];
+      previousSlotNumber = slotNumber;
+    }
+    messagesInEpoch[checkpointIndex].push(block.body.txEffects.map(txEffect => txEffect.l2ToL1Msgs));
+  }
+
+  const outHash = computeEpochOutHash(messagesInEpoch);
+  if (!outHash.isZero()) {
+    await rollupCheatCodes.insertOutbox(epoch, messagesInEpoch.length, outHash.toBigInt());
+  } else {
+    log.info(`No L2 to L1 messages in epoch ${epoch}`, { epoch });
+  }
+
+  return blocks.at(-1)?.checkpointNumber;
+}

package/src/test/index.ts

@@ -0,0 +1 @@
+export { settleEpochOutbox, type SettleEpochOutboxArgs } from './epoch_settlement.js';