@aztec/prover-client 0.67.1 → 0.68.1

package/src/prover-agent/prover-agent.ts

@@ -11,26 +11,29 @@ import {
 import { createLogger } from '@aztec/foundation/log';
 import { RunningPromise } from '@aztec/foundation/running-promise';
 import { elapsed } from '@aztec/foundation/timer';
+import { Attributes, type TelemetryClient, type Traceable, type Tracer, trackSpan } from '@aztec/telemetry-client';
+import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 
 import { InlineProofStore } from '../proving_broker/proof_store.js';
 
 const PRINT_THRESHOLD_NS = 6e10; // 60 seconds
 
+type InFlightPromise = {
+  id: string;
+  type: ProvingRequestType;
+  promise: Promise<any>;
+};
+
 /**
  * A helper class that encapsulates a circuit prover and connects it to a job source.
  */
-export class ProverAgent implements ProverAgentApi {
-  private inFlightPromises = new Map<
-    string,
-    {
-      id: string;
-      type: ProvingRequestType;
-      promise: Promise<any>;
-    }
-  >();
+export class ProverAgent implements ProverAgentApi, Traceable {
+  private inFlightPromises = new Map<string, InFlightPromise>();
   private runningPromise?: RunningPromise;
   private proofInputsDatabase = new InlineProofStore();
 
+  public readonly tracer: Tracer;
+
   constructor(
     /** The prover implementation to defer jobs to */
     private circuitProver: ServerCircuitProver,
@@ -38,8 +41,13 @@ export class ProverAgent implements ProverAgentApi {
     private maxConcurrency = 1,
     /** How long to wait between jobs */
     private pollIntervalMs = 100,
+    /** Telemetry client */
+    private telemetry: TelemetryClient = new NoopTelemetryClient(),
+    /** Logger */
     private log = createLogger('prover-client:prover-agent'),
-  ) {}
+  ) {
+    this.tracer = telemetry.getTracer('ProverAgent');
+  }
 
   setMaxConcurrency(maxConcurrency: number): Promise<void> {
     if (maxConcurrency < 1) {
@@ -74,49 +82,53 @@ export class ProverAgent implements ProverAgentApi {
 
     let lastPrint = process.hrtime.bigint();
 
-    this.runningPromise = new RunningPromise(async () => {
-      for (const jobId of this.inFlightPromises.keys()) {
-        await jobSource.heartbeat(jobId);
-      }
-
-      const now = process.hrtime.bigint();
-
-      if (now - lastPrint >= PRINT_THRESHOLD_NS) {
-        // only log if we're actually doing work
-        if (this.inFlightPromises.size > 0) {
-          const jobs = Array.from(this.inFlightPromises.values())
-            .map(job => `id=${job.id},type=${ProvingRequestType[job.type]}`)
-            .join(' ');
-          this.log.info(`Agent is running with ${this.inFlightPromises.size} in-flight jobs: ${jobs}`);
+    this.runningPromise = new RunningPromise(
+      async () => {
+        for (const jobId of this.inFlightPromises.keys()) {
+          await jobSource.heartbeat(jobId);
         }
-        lastPrint = now;
-      }
 
-      while (this.inFlightPromises.size < this.maxConcurrency) {
-        try {
-          const job = await jobSource.getProvingJob();
-          if (!job) {
-            // job source is fully drained, sleep for a bit and try again
-            return;
+        const now = process.hrtime.bigint();
+
+        if (now - lastPrint >= PRINT_THRESHOLD_NS) {
+          // only log if we're actually doing work
+          if (this.inFlightPromises.size > 0) {
+            const jobs = Array.from(this.inFlightPromises.values())
+              .map(job => `id=${job.id},type=${ProvingRequestType[job.type]}`)
+              .join(' ');
+            this.log.info(`Agent is running with ${this.inFlightPromises.size} in-flight jobs: ${jobs}`);
           }
+          lastPrint = now;
+        }
 
+        while (this.inFlightPromises.size < this.maxConcurrency) {
           try {
-            const promise = this.work(jobSource, job).finally(() => this.inFlightPromises.delete(job.id));
-            this.inFlightPromises.set(job.id, {
-              id: job.id,
-              type: job.type,
-              promise,
-            });
+            const job = await jobSource.getProvingJob();
+            if (!job) {
+              // job source is fully drained, sleep for a bit and try again
+              return;
+            }
+
+            try {
+              const promise = this.work(jobSource, job).finally(() => this.inFlightPromises.delete(job.id));
+              this.inFlightPromises.set(job.id, {
+                id: job.id,
+                type: job.type,
+                promise,
+              });
+            } catch (err) {
+              this.log.warn(
+                `Error processing job! type=${ProvingRequestType[job.type]}: ${err}. ${(err as Error).stack}`,
+              );
+            }
           } catch (err) {
-            this.log.warn(
-              `Error processing job! type=${ProvingRequestType[job.type]}: ${err}. ${(err as Error).stack}`,
-            );
+            this.log.error(`Error fetching job`, err);
           }
-        } catch (err) {
-          this.log.error(`Error fetching job`, err);
         }
-      }
-    }, this.pollIntervalMs);
+      },
+      this.log,
+      this.pollIntervalMs,
+    );
 
     this.runningPromise.start();
     this.log.info(`Agent started with concurrency=${this.maxConcurrency}`);
@@ -133,9 +145,16 @@ export class ProverAgent implements ProverAgentApi {
     this.log.info('Agent stopped');
   }
 
+  @trackSpan('ProverAgent.work', (_jobSoure, job) => ({
+    [Attributes.PROVING_JOB_ID]: job.id,
+    [Attributes.PROVING_JOB_TYPE]: ProvingRequestType[job.type],
+  }))
   private async work(jobSource: ProvingJobSource, job: ProvingJob): Promise<void> {
     try {
-      this.log.debug(`Picked up proving job id=${job.id} type=${ProvingRequestType[job.type]}`);
+      this.log.debug(`Picked up proving job ${job.id} ${ProvingRequestType[job.type]}`, {
+        jobId: job.id,
+        jobType: ProvingRequestType[job.type],
+      });
       const type = job.type;
       const inputs = await this.proofInputsDatabase.getProofInput(job.inputsUri);
       const [time, result] = await elapsed(this.getProof(inputs));

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

@@ -4,7 +4,6 @@ import {
   type EpochProver,
   type EpochProverManager,
   type ForkMerkleTreeOperations,
-  type ProverCache,
   type ProvingJobBroker,
   type ProvingJobConsumer,
   type ProvingJobProducer,
@@ -16,13 +15,10 @@ import { createLogger } from '@aztec/foundation/log';
 import { NativeACVMSimulator } from '@aztec/simulator';
 import { type TelemetryClient } from '@aztec/telemetry-client';
 
-import { join } from 'path';
-
 import { type ProverClientConfig } from '../config.js';
 import { ProvingOrchestrator } from '../orchestrator/orchestrator.js';
-import { CachingBrokerFacade } from '../proving_broker/caching_broker_facade.js';
+import { BrokerCircuitProverFacade } from '../proving_broker/broker_prover_facade.js';
 import { InlineProofStore } from '../proving_broker/proof_store.js';
-import { InMemoryProverCache } from '../proving_broker/prover_cache/memory.js';
 import { ProvingAgent } from '../proving_broker/proving_agent.js';
 
 /** Manages proving of epochs by orchestrating the proving of individual blocks relying on a pool of prover agents. */
@@ -30,8 +26,6 @@ export class ProverClient implements EpochProverManager {
   private running = false;
   private agents: ProvingAgent[] = [];
 
-  private cacheDir?: string;
-
   private constructor(
     private config: ProverClientConfig,
     private worldState: ForkMerkleTreeOperations,
@@ -42,13 +36,12 @@ export class ProverClient implements EpochProverManager {
   ) {
     // TODO(palla/prover-node): Cache the paddingTx here, and not in each proving orchestrator,
     // so it can be reused across multiple ones and not recomputed every time.
-    this.cacheDir = this.config.cacheDir ? join(this.config.cacheDir, `tx_prover_${this.config.proverId}`) : undefined;
   }
 
-  public createEpochProver(cache: ProverCache = new InMemoryProverCache()): EpochProver {
+  public createEpochProver(): EpochProver {
     return new ProvingOrchestrator(
       this.worldState,
-      new CachingBrokerFacade(this.orchestratorClient, cache),
+      new BrokerCircuitProverFacade(this.orchestratorClient),
       this.telemetry,
       this.config.proverId,
     );

package/src/proving_broker/{caching_broker_facade.ts → broker_prover_facade.ts}

@@ -1,6 +1,5 @@
 import {
   type ProofAndVerificationKey,
-  type ProverCache,
   type ProvingJobId,
   type ProvingJobInputsMap,
   type ProvingJobProducer,
@@ -12,47 +11,45 @@ import {
 import {
   type AVM_PROOF_LENGTH_IN_FIELDS,
   type AvmCircuitInputs,
-  type BaseOrMergeRollupPublicInputs,
   type BaseParityInputs,
+  type KernelCircuitPublicInputs,
+  type NESTED_RECURSIVE_PROOF_LENGTH,
+  type ParityPublicInputs,
+  type PrivateKernelEmptyInputData,
+  type RECURSIVE_PROOF_LENGTH,
+  type RootParityInputs,
+  type TUBE_PROOF_LENGTH,
+} from '@aztec/circuits.js';
+import {
+  type BaseOrMergeRollupPublicInputs,
   type BlockMergeRollupInputs,
   type BlockRootOrBlockMergePublicInputs,
   type BlockRootRollupInputs,
   type EmptyBlockRootRollupInputs,
-  type KernelCircuitPublicInputs,
   type MergeRollupInputs,
-  type NESTED_RECURSIVE_PROOF_LENGTH,
-  type ParityPublicInputs,
   type PrivateBaseRollupInputs,
-  type PrivateKernelEmptyInputData,
   type PublicBaseRollupInputs,
-  type RECURSIVE_PROOF_LENGTH,
-  type RootParityInputs,
   type RootRollupInputs,
   type RootRollupPublicInputs,
-  type TUBE_PROOF_LENGTH,
   type TubeInputs,
-} from '@aztec/circuits.js';
+} from '@aztec/circuits.js/rollup';
 import { sha256 } from '@aztec/foundation/crypto';
 import { createLogger } from '@aztec/foundation/log';
 import { retryUntil } from '@aztec/foundation/retry';
+import { truncate } from '@aztec/foundation/string';
 
 import { InlineProofStore, type ProofStore } from './proof_store.js';
-import { InMemoryProverCache } from './prover_cache/memory.js';
 
 // 20 minutes, roughly the length of an Aztec epoch. If a proof isn't ready in this amount of time then we've failed to prove the whole epoch
 const MAX_WAIT_MS = 1_200_000;
 
-/**
- * A facade around a job broker that generates stable job ids and caches results
- */
-export class CachingBrokerFacade implements ServerCircuitProver {
+export class BrokerCircuitProverFacade implements ServerCircuitProver {
   constructor(
     private broker: ProvingJobProducer,
-    private cache: ProverCache = new InMemoryProverCache(),
     private proofStore: ProofStore = new InlineProofStore(),
     private waitTimeoutMs = MAX_WAIT_MS,
     private pollIntervalMs = 1000,
-    private log = createLogger('prover-client:caching-prover-broker'),
+    private log = createLogger('prover-client:broker-circuit-prover-facade'),
   ) {}
 
   private async enqueueAndWaitForJob<T extends ProvingRequestType>(
@@ -62,49 +59,23 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber = 0,
     signal?: AbortSignal,
   ): Promise<ProvingJobResultsMap[T]> {
-    // first try the cache
-    let jobEnqueued = false;
-    try {
-      const cachedResult = await this.cache.getProvingJobStatus(id);
-      if (cachedResult.status !== 'not-found') {
-        this.log.debug(`Found cached result for job=${id}: status=${cachedResult.status}`);
-      }
-
-      if (cachedResult.status === 'fulfilled') {
-        const output = await this.proofStore.getProofOutput(cachedResult.value);
-        if (output.type === type) {
-          return output.result as ProvingJobResultsMap[T];
-        } else {
-          this.log.warn(`Cached result type mismatch for job=${id}. Expected=${type} but got=${output.type}`);
-        }
-      } else if (cachedResult.status === 'rejected') {
-        // prefer returning a rejected promises so that we don't trigger the catch block below
-        return Promise.reject(new Error(cachedResult.reason));
-      } else if (cachedResult.status === 'in-progress' || cachedResult.status === 'in-queue') {
-        jobEnqueued = true;
-      } else {
-        jobEnqueued = false;
-      }
-    } catch (err) {
-      this.log.warn(`Failed to get cached proving job id=${id}: ${err}. Re-running job`);
-    }
+    const inputsUri = await this.proofStore.saveProofInput(id, type, inputs);
+    await this.broker.enqueueProvingJob({
+      id,
+      type,
+      inputsUri,
+      epochNumber,
+    });
 
-    if (!jobEnqueued) {
-      try {
-        const inputsUri = await this.proofStore.saveProofInput(id, type, inputs);
-        await this.broker.enqueueProvingJob({
-          id,
-          type,
-          inputsUri,
-          epochNumber,
-        });
-        await this.cache.setProvingJobStatus(id, { status: 'in-queue' });
-      } catch (err) {
-        this.log.error(`Failed to enqueue proving job id=${id}: ${err}`);
-        await this.cache.setProvingJobStatus(id, { status: 'not-found' });
-        throw err;
-      }
-    }
+    this.log.verbose(
+      `Sent proving job to broker id=${id} type=${ProvingRequestType[type]} epochNumber=${epochNumber}`,
+      {
+        provingJobId: id,
+        provingJobType: ProvingRequestType[type],
+        epochNumber,
+        inputsUri: truncate(inputsUri),
+      },
+    );
 
     // notify broker of cancelled job
     const abortFn = async () => {
@@ -131,26 +102,18 @@ export class CachingBrokerFacade implements ServerCircuitProver {
         this.pollIntervalMs / 1000,
       );
 
-      try {
-        await this.cache.setProvingJobStatus(id, result);
-      } catch (err) {
-        this.log.warn(`Failed to cache proving job id=${id} resultStatus=${result.status}: ${err}`);
-      }
-
       if (result.status === 'fulfilled') {
         const output = await this.proofStore.getProofOutput(result.value);
         if (output.type === type) {
           return output.result as ProvingJobResultsMap[T];
         } else {
-          return Promise.reject(new Error(`Unexpected proof type: ${output.type}. Expected: ${type}`));
+          throw new Error(`Unexpected proof type: ${output.type}. Expected: ${type}`);
         }
       } else {
-        return Promise.reject(new Error(result.reason));
+        throw new Error(result.reason);
       }
     } finally {
       signal?.removeEventListener('abort', abortFn);
-      // we've saved the result in our cache. We can tell the broker to clear its state
-      await this.broker.cleanUpProvingJobState(id);
     }
   }
 
@@ -160,7 +123,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<ProofAndVerificationKey<typeof AVM_PROOF_LENGTH_IN_FIELDS>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PUBLIC_VM, inputs),
+      this.generateId(ProvingRequestType.PUBLIC_VM, inputs, epochNumber),
       ProvingRequestType.PUBLIC_VM,
       inputs,
       epochNumber,
@@ -174,7 +137,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.BASE_PARITY, inputs),
+      this.generateId(ProvingRequestType.BASE_PARITY, inputs, epochNumber),
       ProvingRequestType.BASE_PARITY,
       inputs,
       epochNumber,
@@ -188,7 +151,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.BLOCK_MERGE_ROLLUP, input),
+      this.generateId(ProvingRequestType.BLOCK_MERGE_ROLLUP, input, epochNumber),
       ProvingRequestType.BLOCK_MERGE_ROLLUP,
       input,
       epochNumber,
@@ -202,7 +165,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.BLOCK_ROOT_ROLLUP, input),
+      this.generateId(ProvingRequestType.BLOCK_ROOT_ROLLUP, input, epochNumber),
       ProvingRequestType.BLOCK_ROOT_ROLLUP,
       input,
       epochNumber,
@@ -216,7 +179,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP, input),
+      this.generateId(ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP, input, epochNumber),
       ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP,
       input,
       epochNumber,
@@ -230,7 +193,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<KernelCircuitPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PRIVATE_KERNEL_EMPTY, inputs),
+      this.generateId(ProvingRequestType.PRIVATE_KERNEL_EMPTY, inputs, epochNumber),
       ProvingRequestType.PRIVATE_KERNEL_EMPTY,
       inputs,
       epochNumber,
@@ -244,7 +207,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.MERGE_ROLLUP, input),
+      this.generateId(ProvingRequestType.MERGE_ROLLUP, input, epochNumber),
       ProvingRequestType.MERGE_ROLLUP,
       input,
       epochNumber,
@@ -257,7 +220,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PRIVATE_BASE_ROLLUP, baseRollupInput),
+      this.generateId(ProvingRequestType.PRIVATE_BASE_ROLLUP, baseRollupInput, epochNumber),
       ProvingRequestType.PRIVATE_BASE_ROLLUP,
       baseRollupInput,
       epochNumber,
@@ -271,7 +234,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PUBLIC_BASE_ROLLUP, inputs),
+      this.generateId(ProvingRequestType.PUBLIC_BASE_ROLLUP, inputs, epochNumber),
       ProvingRequestType.PUBLIC_BASE_ROLLUP,
       inputs,
       epochNumber,
@@ -285,7 +248,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof NESTED_RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.ROOT_PARITY, inputs),
+      this.generateId(ProvingRequestType.ROOT_PARITY, inputs, epochNumber),
       ProvingRequestType.ROOT_PARITY,
       inputs,
       epochNumber,
@@ -299,7 +262,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<RootRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.ROOT_ROLLUP, input),
+      this.generateId(ProvingRequestType.ROOT_ROLLUP, input, epochNumber),
       ProvingRequestType.ROOT_ROLLUP,
       input,
       epochNumber,
@@ -313,7 +276,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<ProofAndVerificationKey<typeof TUBE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.TUBE_PROOF, tubeInput),
+      this.generateId(ProvingRequestType.TUBE_PROOF, tubeInput, epochNumber),
       ProvingRequestType.TUBE_PROOF,
       tubeInput,
       epochNumber,
@@ -321,8 +284,8 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     );
   }
 
-  private generateId(type: ProvingRequestType, inputs: { toBuffer(): Buffer }) {
+  private generateId(type: ProvingRequestType, inputs: { toBuffer(): Buffer }, epochNumber = 0) {
     const inputsHash = sha256(inputs.toBuffer());
-    return `${ProvingRequestType[type]}:${inputsHash.toString('hex')}`;
+    return `${epochNumber}:${ProvingRequestType[type]}:${inputsHash.toString('hex')}`;
   }
 }

package/src/proving_broker/proving_agent.ts

@@ -10,8 +10,9 @@ import {
 } from '@aztec/circuit-types';
 import { createLogger } from '@aztec/foundation/log';
 import { RunningPromise } from '@aztec/foundation/running-promise';
+import { truncate } from '@aztec/foundation/string';
 import { Timer } from '@aztec/foundation/timer';
-import { type TelemetryClient } from '@aztec/telemetry-client';
+import { type TelemetryClient, type Traceable, type Tracer, trackSpan } from '@aztec/telemetry-client';
 
 import { type ProofStore } from './proof_store.js';
 import { ProvingAgentInstrumentation } from './proving_agent_instrumentation.js';
@@ -20,12 +21,14 @@ import { ProvingJobController, ProvingJobControllerStatus } from './proving_job_
 /**
  * A helper class that encapsulates a circuit prover and connects it to a job source.
  */
-export class ProvingAgent {
+export class ProvingAgent implements Traceable {
   private currentJobController?: ProvingJobController;
   private runningPromise: RunningPromise;
   private instrumentation: ProvingAgentInstrumentation;
   private idleTimer: Timer | undefined;
 
+  public readonly tracer: Tracer;
+
   constructor(
     /** The source of proving jobs */
     private broker: ProvingJobConsumer,
@@ -41,8 +44,9 @@ export class ProvingAgent {
     private pollIntervalMs = 1000,
     private log = createLogger('prover-client:proving-agent'),
   ) {
+    this.tracer = client.getTracer('ProvingAgent');
     this.instrumentation = new ProvingAgentInstrumentation(client);
-    this.runningPromise = new RunningPromise(this.safeWork, this.pollIntervalMs);
+    this.runningPromise = new RunningPromise(this.work.bind(this), this.log, this.pollIntervalMs);
   }
 
   public setCircuitProver(circuitProver: ServerCircuitProver): void {
@@ -63,76 +67,74 @@ export class ProvingAgent {
     await this.runningPromise.stop();
   }
 
-  private safeWork = async () => {
-    try {
-      // every tick we need to
-      // (1) either do a heartbeat, telling the broker that we're working
-      // (2) get a new job
-      // If during (1) the broker returns a new job that means we can cancel the current job and start the new one
-      let maybeJob: { job: ProvingJob; time: number } | undefined;
-      if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
-        maybeJob = await this.broker.reportProvingJobProgress(
-          this.currentJobController.getJobId(),
-          this.currentJobController.getStartedAt(),
-          { allowList: this.proofAllowList },
-        );
-      } else {
-        maybeJob = await this.broker.getProvingJob({ allowList: this.proofAllowList });
-      }
-
-      if (!maybeJob) {
-        return;
-      }
-
-      let abortedProofJobId: string | undefined;
-      let abortedProofName: string | undefined;
-      if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
-        abortedProofJobId = this.currentJobController.getJobId();
-        abortedProofName = this.currentJobController.getProofTypeName();
-        this.currentJobController?.abort();
-      }
-
-      const { job, time } = maybeJob;
-      let inputs: ProvingJobInputs;
-      try {
-        inputs = await this.proofStore.getProofInput(job.inputsUri);
-      } catch (err) {
-        await this.broker.reportProvingJobError(job.id, 'Failed to load proof inputs', true);
-        return;
-      }
-
-      this.currentJobController = new ProvingJobController(
-        job.id,
-        inputs,
-        time,
-        this.circuitProver,
-        this.handleJobResult,
+  @trackSpan('ProvingAgent.safeWork')
+  private async work() {
+    // every tick we need to
+    // (1) either do a heartbeat, telling the broker that we're working
+    // (2) get a new job
+    // If during (1) the broker returns a new job that means we can cancel the current job and start the new one
+    let maybeJob: { job: ProvingJob; time: number } | undefined;
+    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
+      maybeJob = await this.broker.reportProvingJobProgress(
+        this.currentJobController.getJobId(),
+        this.currentJobController.getStartedAt(),
+        { allowList: this.proofAllowList },
       );
+    } else {
+      maybeJob = await this.broker.getProvingJob({ allowList: this.proofAllowList });
+    }
+
+    if (!maybeJob) {
+      return;
+    }
+
+    let abortedProofJobId: string | undefined;
+    let abortedProofName: string | undefined;
+    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
+      abortedProofJobId = this.currentJobController.getJobId();
+      abortedProofName = this.currentJobController.getProofTypeName();
+      this.currentJobController?.abort();
+    }
 
-      if (abortedProofJobId) {
-        this.log.info(
-          `Aborting job id=${abortedProofJobId} type=${abortedProofName} to start new job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncateString(
-            job.inputsUri,
-          )}`,
-        );
-      } else {
-        this.log.info(
-          `Starting job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncateString(
-            job.inputsUri,
-          )}`,
-        );
-      }
-
-      if (this.idleTimer) {
-        this.instrumentation.recordIdleTime(this.idleTimer);
-      }
-      this.idleTimer = undefined;
-
-      this.currentJobController.start();
+    const { job, time } = maybeJob;
+    let inputs: ProvingJobInputs;
+    try {
+      inputs = await this.proofStore.getProofInput(job.inputsUri);
     } catch (err) {
-      this.log.error(`Error in ProvingAgent: ${String(err)}`);
+      await this.broker.reportProvingJobError(job.id, 'Failed to load proof inputs', true);
+      return;
     }
-  };
+
+    this.currentJobController = new ProvingJobController(
+      job.id,
+      inputs,
+      job.epochNumber,
+      time,
+      this.circuitProver,
+      this.handleJobResult,
+    );
+
+    if (abortedProofJobId) {
+      this.log.info(
+        `Aborting job id=${abortedProofJobId} type=${abortedProofName} to start new job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncate(
+          job.inputsUri,
+        )}`,
+      );
+    } else {
+      this.log.info(
+        `Starting job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncate(
+          job.inputsUri,
+        )}`,
+      );
+    }
+
+    if (this.idleTimer) {
+      this.instrumentation.recordIdleTime(this.idleTimer);
+    }
+    this.idleTimer = undefined;
+
+    this.currentJobController.start();
+  }
 
   handleJobResult = async <T extends ProvingRequestType>(
     jobId: ProvingJobId,
@@ -147,14 +149,8 @@ export class ProvingAgent {
       return this.broker.reportProvingJobError(jobId, err.message, retry);
     } else if (result) {
       const outputUri = await this.proofStore.saveProofOutput(jobId, type, result);
-      this.log.info(
-        `Job id=${jobId} type=${ProvingRequestType[type]} completed outputUri=${truncateString(outputUri)}`,
-      );
+      this.log.info(`Job id=${jobId} type=${ProvingRequestType[type]} completed outputUri=${truncate(outputUri)}`);
       return this.broker.reportProvingJobSuccess(jobId, outputUri);
     }
   };
 }
-
-function truncateString(str: string, length: number = 64): string {
-  return str.length > length ? str.slice(0, length) + '...' : str;
-}