@aztec/prover-client 0.0.1-commit.a072138 → 0.0.1-commit.a5db02d

package/src/proving_broker/proof_store/file_store_proof_store.ts

@@ -0,0 +1,78 @@
+import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
+import type { FileStore } from '@aztec/stdlib/file-store';
+import {
+  type ProofUri,
+  type ProvingJobId,
+  type ProvingJobInputs,
+  type ProvingJobInputsMap,
+  ProvingJobResult,
+  type ProvingJobResultsMap,
+  getProvingJobInputClassFor,
+} from '@aztec/stdlib/interfaces/server';
+import { ProvingRequestType } from '@aztec/stdlib/proofs';
+
+import type { ProofStore } from './proof_store.js';
+
+const INPUTS_PATH = 'inputs';
+const OUTPUTS_PATH = 'outputs';
+
+/**
+ * A proof store implementation backed by a generic FileStore.
+ * Supports any storage backend (GCS, S3, local filesystem) via the FileStore abstraction.
+ */
+export class FileStoreProofStore implements ProofStore {
+  constructor(private readonly fileStore: FileStore) {}
+
+  async saveProofInput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+  ): Promise<ProofUri> {
+    const path = `${INPUTS_PATH}/${ProvingRequestType[type]}/${id}`;
+    const uri = await this.fileStore.save(path, inputs.toBuffer());
+    return uri as ProofUri;
+  }
+
+  async saveProofOutput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    result: ProvingJobResultsMap[T],
+  ): Promise<ProofUri> {
+    const jobResult = { type, result } as ProvingJobResult;
+    const json = jsonStringify(jobResult);
+    const path = `${OUTPUTS_PATH}/${ProvingRequestType[type]}/${id}.json`;
+    const uri = await this.fileStore.save(path, Buffer.from(json, 'utf-8'));
+    return uri as ProofUri;
+  }
+
+  async getProofInput(uri: ProofUri): Promise<ProvingJobInputs> {
+    try {
+      const buffer = await this.fileStore.read(uri);
+      const type = this.extractTypeFromUri(uri);
+      const inputs = getProvingJobInputClassFor(type).fromBuffer(buffer);
+      return { inputs, type } as ProvingJobInputs;
+    } catch (err) {
+      throw new Error(`Error getting proof input at ${uri}: ${err}`);
+    }
+  }
+
+  async getProofOutput(uri: ProofUri): Promise<ProvingJobResult> {
+    try {
+      const buffer = await this.fileStore.read(uri);
+      return jsonParseWithSchema(buffer.toString('utf-8'), ProvingJobResult);
+    } catch (err) {
+      throw new Error(`Error getting proof output at ${uri}: ${err}`);
+    }
+  }
+
+  private extractTypeFromUri(uri: string): ProvingRequestType {
+    const url = new URL(uri);
+    const pathParts = url.pathname.split('/').filter(Boolean);
+    const typeString = pathParts.at(-2);
+    const type = typeString ? ProvingRequestType[typeString as keyof typeof ProvingRequestType] : undefined;
+    if (type === undefined) {
+      throw new Error(`Unrecognized proof type ${typeString} in URI ${uri}`);
+    }
+    return type;
+  }
+}

package/src/proving_broker/proof_store/index.ts

@@ -1,4 +1,4 @@
 export * from './proof_store.js';
 export * from './inline_proof_store.js';
 export * from './factory.js';
-export * from './gcs_proof_store.js';
+export * from './file_store_proof_store.js';

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,
@@ -36,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),
@@ -114,6 +115,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
   private started = false;
 
+  private debugReplayEnabled: boolean;
+
   public constructor(
     private database: ProvingBrokerDatabase,
     {
@@ -121,6 +124,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       proverBrokerPollIntervalMs,
       proverBrokerJobMaxRetries,
       proverBrokerMaxEpochsToKeepResultsFor,
+      proverBrokerDebugReplayEnabled,
     }: Required<
       Pick<
         ProverBrokerConfig,
@@ -128,6 +132,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
         | 'proverBrokerPollIntervalMs'
         | 'proverBrokerJobMaxRetries'
         | 'proverBrokerMaxEpochsToKeepResultsFor'
+        | 'proverBrokerDebugReplayEnabled'
       >
     > = defaultProverBrokerConfig,
     client: TelemetryClient = getTelemetryClient(),
@@ -139,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) => {
@@ -241,17 +247,65 @@ 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);
+    let jobStatus = this.#getProvingJobStatus(job.id);
     if (this.jobsCache.has(job.id)) {
       const existing = this.jobsCache.get(job.id);
       assert.deepStrictEqual(job, existing, 'Duplicate proving job ID');
-      this.logger.warn(`Cached proving job id=${job.id} epochNumber=${job.epochNumber}. Not enqueuing again`, {
-        provingJobId: job.id,
-      });
-      this.instrumentation.incCachedJobs(job.type);
-      return jobStatus;
+
+      if (this.resultsCache.get(job.id)?.status === 'aborted') {
+        // The producer is re-requesting a job it previously cancelled: revive it rather than
+        // returning the cached abort, clearing the aborted state in memory and in the database so the
+        // revival survives a restart.
+        //
+        // Concurrency model: `jobsCache` is the enqueue lock. Every path that puts a job on the queue
+        // populates `jobsCache` *synchronously, before its first await* (see the "New proving job"
+        // block below), so a second concurrent enqueue of the same id observes the entry at the top
+        // of this method and takes a cached, no-op branch instead of enqueuing a duplicate. The revive
+        // must keep holding that lock: we tear down the settled state and re-set `jobsCache` in a
+        // single synchronous span (no await in between), and only then await the database. Because a
+        // concurrent re-request can only interleave at that await — by which point `jobsCache` is
+        // populated again and the aborted result is gone — it falls into the cached branch and no-ops,
+        // so the job is enqueued exactly once. (`cleanUpProvingJobState` also drops the promise, which
+        // was resolved with the abort, so `enqueueJobInternal` below mints a fresh one for the retry.)
+        this.logger.info(`Reviving aborted proving job id=${job.id} epochNumber=${job.epochNumber}`, {
+          provingJobId: job.id,
+        });
+        this.cleanUpProvingJobState([job.id]);
+        this.jobsCache.set(job.id, job);
+        await this.database.deleteProvingJobResult(job.id);
+        jobStatus = this.#getProvingJobStatus(job.id);
+      } else {
+        this.logger.warn(`Cached proving job id=${job.id} epochNumber=${job.epochNumber}. Not enqueuing again`, {
+          provingJobId: job.id,
+        });
+        this.instrumentation.incCachedJobs(job.type);
+        return jobStatus;
+      }
     }
 
     if (this.isJobStale(job)) {
@@ -277,27 +331,53 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   }
 
   async #cancelProvingJob(id: ProvingJobId): Promise<void> {
-    if (!this.jobsCache.has(id)) {
+    const job = this.jobsCache.get(id);
+    if (!job) {
       this.logger.warn(`Can't cancel a job that doesn't exist id=${id}`, { provingJobId: id });
       return;
     }
 
-    // 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);
+    // Leave jobs that have already settled (completed or failed) alone: those results are terminal.
+    if (this.resultsCache.has(id)) {
+      return;
+    }
+
+    this.logger.info(`Cancelling job id=${id}`, { provingJobId: id });
+    this.inProgress.delete(id);
+
+    // Record the cancellation as its own settled state and persist it, so it survives a restart and
+    // notifies the current waiter. Unlike a completion or failure this is not terminal: re-enqueuing
+    // the same job id revives it (see #enqueueProvingJob), so the abort never permanently blocks the
+    // proof.
+    const result: ProvingJobSettledResult = { status: 'aborted' };
+    this.resultsCache.set(id, result);
+    this.promises.get(id)?.resolve(result);
+    this.completedJobNotifications.push(id);
+    this.instrumentation.incAbortedJobs(job.type);
+
+    try {
+      await this.database.setProvingJobAborted(id);
+    } catch (saveErr) {
+      this.logger.error(`Failed to save proving job aborted status id=${id}`, saveErr, { provingJobId: id });
+      throw saveErr;
     }
   }
 
   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 {
@@ -560,21 +640,21 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   }
 
   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);
       }
     }
@@ -598,10 +678,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/memory.ts

@@ -36,6 +36,16 @@ export class InMemoryBrokerDatabase implements ProvingBrokerDatabase {
     return Promise.resolve();
   }
 
+  setProvingJobAborted(id: ProvingJobId): Promise<void> {
+    this.results.set(id, { status: 'aborted' });
+    return Promise.resolve();
+  }
+
+  deleteProvingJobResult(id: ProvingJobId): Promise<void> {
+    this.results.delete(id);
+    return Promise.resolve();
+  }
+
   deleteProvingJobs(ids: ProvingJobId[]): Promise<void> {
     for (const id of ids) {
       this.jobs.delete(id);

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -41,7 +41,11 @@ class SingleEpochDatabase {
     return this.store.estimateSize();
   }
 
-  async batchWrite(jobs: ProvingJob[], results: Array<[ProvingJobId, ProvingJobSettledResult]>) {
+  async batchWrite(
+    jobs: ProvingJob[],
+    results: Array<[ProvingJobId, ProvingJobSettledResult]>,
+    deletedResults: ProvingJobId[] = [],
+  ) {
     await this.store.transactionAsync(async () => {
       for (const job of jobs) {
         await this.jobs.set(job.id, jsonStringify(job));
@@ -49,6 +53,11 @@ class SingleEpochDatabase {
       for (const [id, result] of results) {
         await this.jobResults.set(id, jsonStringify(result));
       }
+      // Deletes are applied after sets so that, if a set and a delete for the same id land in one
+      // batch (e.g. an abort immediately followed by a revive's result-delete), the delete wins.
+      for (const id of deletedResults) {
+        await this.jobResults.delete(id);
+      }
     });
   }
 
@@ -66,6 +75,11 @@ class SingleEpochDatabase {
     await this.jobResults.set(id, jsonStringify(result));
   }
 
+  async setProvingJobAborted(id: ProvingJobId): Promise<void> {
+    const result: ProvingJobSettledResult = { status: 'aborted' };
+    await this.jobResults.set(id, jsonStringify(result));
+  }
+
   async setProvingJobResult(id: ProvingJobId, value: ProofUri): Promise<void> {
     const result: ProvingJobSettledResult = { status: 'fulfilled', value };
     await this.jobResults.set(id, jsonStringify(result));
@@ -80,10 +94,19 @@ class SingleEpochDatabase {
   }
 }
 
+/**
+ * An item queued for a batched write. A {@link ProvingJob} adds/updates a job; a `[id, result]`
+ * tuple sets a job's settled result; a `[id, undefined]` tuple deletes a job's result. Deletes ride
+ * the same queue as writes so they stay FIFO-ordered against them: a revive's result-delete must
+ * land after the abort that preceded it, otherwise a still-queued abort could flush later and
+ * resurrect the cancelled state on disk.
+ */
+type BrokerWrite = ProvingJob | [ProvingJobId, ProvingJobSettledResult | undefined];
+
 export class KVBrokerDatabase implements ProvingBrokerDatabase {
   private metrics: LmdbMetrics;
 
-  private batchQueue: BatchQueue<ProvingJob | [ProvingJobId, ProvingJobSettledResult], number>;
+  private batchQueue: BatchQueue<BrokerWrite, number>;
 
   public readonly tracer: Tracer;
 
@@ -112,12 +135,16 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
   }
 
   // exposed for testing
-  public async commitWrites(items: Array<ProvingJob | [ProvingJobId, ProvingJobSettledResult]>, epochNumber: number) {
+  public async commitWrites(items: Array<BrokerWrite>, epochNumber: number) {
     const jobsToAdd = items.filter((item): item is ProvingJob => 'id' in item);
-    const resultsToAdd = items.filter((item): item is [ProvingJobId, ProvingJobSettledResult] => Array.isArray(item));
+    const resultOps = items.filter((item): item is [ProvingJobId, ProvingJobSettledResult | undefined] =>
+      Array.isArray(item),
+    );
+    const resultsToSet = resultOps.filter((op): op is [ProvingJobId, ProvingJobSettledResult] => op[1] !== undefined);
+    const resultsToDelete = resultOps.filter(op => op[1] === undefined).map(([id]) => id);
 
     const db = await this.getEpochDatabase(EpochNumber(epochNumber));
-    await db.batchWrite(jobsToAdd, resultsToAdd);
+    await db.batchWrite(jobsToAdd, resultsToSet, resultsToDelete);
   }
 
   private async estimateSize() {
@@ -154,7 +181,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);
@@ -207,6 +234,14 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     return this.batchQueue.put([id, { status: 'rejected', reason }], getEpochFromProvingJobId(id));
   }
 
+  setProvingJobAborted(id: ProvingJobId): Promise<void> {
+    return this.batchQueue.put([id, { status: 'aborted' }], getEpochFromProvingJobId(id));
+  }
+
+  deleteProvingJobResult(id: ProvingJobId): Promise<void> {
+    return this.batchQueue.put([id, undefined], getEpochFromProvingJobId(id));
+  }
+
   setProvingJobResult(id: ProvingJobId, value: ProofUri): Promise<void> {
     return this.batchQueue.put([id, { status: 'fulfilled', value }], getEpochFromProvingJobId(id));
   }
@@ -222,7 +257,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_database.ts

@@ -38,6 +38,22 @@ export interface ProvingBrokerDatabase {
    */
   setProvingJobError(id: ProvingJobId, err: string): Promise<void>;
 
+  /**
+   * Records that a proof request was cancelled. Unlike a result or error this is not terminal:
+   * re-enqueuing the same job id revives it, so the aborted state can survive a restart without
+   * permanently blocking the proof.
+   * @param id - The ID of the cancelled proof request
+   */
+  setProvingJobAborted(id: ProvingJobId): Promise<void>;
+
+  /**
+   * Clears any stored result for a proof request, returning it to the pending state while keeping
+   * the job itself. Used when reviving an aborted job so the revival is persisted and a restart
+   * cannot resurrect the stale aborted state.
+   * @param id - The ID of the proof request whose result should be cleared
+   */
+  deleteProvingJobResult(id: ProvingJobId): Promise<void>;
+
   /**
    * Closes the database
    */

package/src/proving_broker/proving_broker_instrumentation.ts

@@ -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;
@@ -39,6 +40,8 @@ export class ProvingBrokerInstrumentation {
 
     this.rejectedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_REJECTED_JOBS, provingJobAttrs);
 
+    this.abortedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_ABORTED_JOBS, provingJobAttrs);
+
     this.retriedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RETRIED_JOBS, provingJobAttrs);
 
     this.timedOutJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, provingJobAttrs);
@@ -72,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/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',