@aztec/prover-client 0.0.1-commit.c31f2472 → 0.0.1-commit.c52d6e7

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,
@@ -32,11 +33,28 @@ type InProgressMetadata = {
 
 type EnqueuedProvingJob = Pick<ProvingJob, 'id' | 'epochNumber'>;
 
+/** In-memory scheduling metadata for a job. The large `inputsUri` is deliberately excluded — it lives
+ * in the database and is read on demand when dispatching to an agent. */
+type ProvingJobMetadata = Pick<ProvingJob, 'id' | 'type' | 'epochNumber'>;
+
+/** Settled state kept in memory: the status plus the small payloads (rejected reason, aborted). The
+ * large fulfilled proof `value` is NOT kept here — it lives in the database (read on demand) and,
+ * transiently, in `pendingResults` while its database write is in flight. */
+type SettledRecord = { status: 'fulfilled' } | { status: 'rejected'; reason: string } | { status: 'aborted' };
+
+function toJobMetadata(job: ProvingJob): ProvingJobMetadata {
+  return { id: job.id, type: job.type, epochNumber: job.epochNumber };
+}
+
+function toSettledRecord(result: ProvingJobSettledResult): SettledRecord {
+  return result.status === 'fulfilled' ? { status: 'fulfilled' } : result;
+}
+
 /**
  * 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),
@@ -68,11 +86,17 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     [ProvingRequestType.PARITY_ROOT]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
   };
 
-  // holds a copy of the database in memory in order to quickly fulfill requests
-  // this is fine because this broker is the only one that can modify the database
-  private jobsCache = new Map<ProvingJobId, ProvingJob>();
-  // as above, but for results
-  private resultsCache = new Map<ProvingJobId, ProvingJobSettledResult>();
+  // Scheduling metadata for every known job (id, type, epochNumber). The large `inputsUri` is NOT held
+  // here — it lives in the database and is read on demand when a job is dispatched to an agent.
+  private jobsCache = new Map<ProvingJobId, ProvingJobMetadata>();
+  // Settled status for every settled job. The large fulfilled proof `value` is NOT held here — it lives
+  // in the database (read on demand in getProvingJobStatus) and, transiently, in `pendingResults` while
+  // its database write is in flight. Small payloads (rejected reason, aborted) are kept inline.
+  private resultsCache = new Map<ProvingJobId, SettledRecord>();
+  // Transient hold of a fulfilled result's full value, from the moment it settles until its database
+  // write commits, so status reads are read-your-writes without retaining every proof. Evicted once the
+  // value is durable in the database; kept only if that write failed (preserving the result in memory).
+  private pendingResults = new Map<ProvingJobId, ProvingJobSettledResult>();
 
   // tracks when each job was enqueued
   private enqueuedAt = new Map<ProvingJobId, Timer>();
@@ -86,8 +110,9 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   // keep track of which proving job has been retried
   private retries = new Map<ProvingJobId, number>();
 
-  // a map of promises that will be resolved when a job is settled
-  private promises = new Map<ProvingJobId, PromiseWithResolvers<ProvingJobSettledResult>>();
+  // a map of promises that are resolved when a job settles. Carries no payload (settled status/results
+  // are read from resultsCache/pendingResults/DB) so it never pins a proof value in memory.
+  private promises = new Map<ProvingJobId, PromiseWithResolvers<void>>();
 
   private cleanupPromise: RunningPromise;
   private msTimeSource = () => Date.now();
@@ -114,6 +139,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
   private started = false;
 
+  private debugReplayEnabled: boolean;
+
   public constructor(
     private database: ProvingBrokerDatabase,
     {
@@ -121,6 +148,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       proverBrokerPollIntervalMs,
       proverBrokerJobMaxRetries,
       proverBrokerMaxEpochsToKeepResultsFor,
+      proverBrokerDebugReplayEnabled,
     }: Required<
       Pick<
         ProverBrokerConfig,
@@ -128,6 +156,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
         | 'proverBrokerPollIntervalMs'
         | 'proverBrokerJobMaxRetries'
         | 'proverBrokerMaxEpochsToKeepResultsFor'
+        | 'proverBrokerDebugReplayEnabled'
       >
     > = defaultProverBrokerConfig,
     client: TelemetryClient = getTelemetryClient(),
@@ -139,6 +168,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) => {
@@ -169,12 +199,12 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
         status: result ? result.status : 'pending',
       });
 
-      this.jobsCache.set(item.id, item);
+      this.jobsCache.set(item.id, toJobMetadata(item));
       this.promises.set(item.id, promiseWithResolvers());
 
       if (result) {
-        this.promises.get(item.id)!.resolve(result);
-        this.resultsCache.set(item.id, result);
+        this.promises.get(item.id)!.resolve();
+        this.resultsCache.set(item.id, toSettledRecord(result));
       } else {
         this.enqueueJobInternal(item);
       }
@@ -205,7 +235,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   }
 
   public getProvingJobStatus(id: ProvingJobId): Promise<ProvingJobStatus> {
-    return Promise.resolve(this.#getProvingJobStatus(id));
+    return this.#getProvingJobStatus(id);
   }
 
   public getCompletedJobs(ids: ProvingJobId[]): Promise<ProvingJobId[]> {
@@ -213,7 +243,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   }
 
   public getProvingJob(filter?: ProvingJobFilter): Promise<GetProvingJobResponse | undefined> {
-    return Promise.resolve(this.#getProvingJob(filter));
+    return this.#getProvingJob(filter);
   }
 
   public reportProvingJobSuccess(
@@ -238,20 +268,79 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     startedAt: number,
     filter?: ProvingJobFilter,
   ): Promise<{ job: ProvingJob; time: number } | undefined> {
-    return Promise.resolve(this.#reportProvingJobProgress(id, startedAt, filter));
+    return 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, toJobMetadata(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);
+    // The status returned to the caller reflects the job's state at the start of this call: `not-found`
+    // for a brand-new job (it did not exist yet), or the cached status for one already known. Crucially
+    // this must NOT `await` before the synchronous `jobsCache`/`resultsCache` gate below, or the
+    // enqueue/revive lock would no longer hold. So the not-found default is set synchronously here and
+    // the cached branch computes the real status (which may read the DB) only when it returns.
+    let jobStatus: ProvingJobStatus = { status: 'not-found' };
     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;
+      // Identity check is metadata-only: `inputsUri` now lives in the database, not memory, and job ids
+      // are content-addressed (same id ⇒ same inputs), so a mismatch shows up in id/type/epochNumber.
+      assert.deepStrictEqual(toJobMetadata(job), existing, 'Duplicate proving job ID');
+
+      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, so
+        // `enqueueJobInternal` below mints a fresh one for the retry.) This holds unchanged with the
+        // slimmer caches: `jobsCache` (now metadata) and `resultsCache` (now settled status) are still
+        // synchronous in-memory maps, so the tear-down + re-set span still contains no await.
+        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, toJobMetadata(job));
+        await this.database.deleteProvingJobResult(job.id);
+        // The job is re-set in the cache and about to be re-enqueued below: its start status is in-queue.
+        jobStatus = { status: 'in-queue' };
+      } 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 the job's current status. This reads the fulfilled proof value from the DB when needed;
+        // the await is fine here because this branch does not enqueue, so it is outside the lock span.
+        return await this.#getProvingJobStatus(job.id);
+      }
     }
 
     if (this.isJobStale(job)) {
@@ -264,7 +353,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     this.logger.info(`New proving job id=${job.id} epochNumber=${job.epochNumber}`, { provingJobId: job.id });
     try {
       // do this first so it acts as a "lock". If this job is enqueued again while we're saving it the if at the top will catch it.
-      this.jobsCache.set(job.id, job);
+      this.jobsCache.set(job.id, toJobMetadata(job));
       await this.database.addProvingJob(job);
       this.enqueueJobInternal(job);
       this.instrumentation.incTotalJobs(job.type);
@@ -277,43 +366,80 @@ 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.
+    this.resultsCache.set(id, { status: 'aborted' });
+    this.promises.get(id)?.resolve();
+    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();
+      }
       this.promises.delete(id);
       this.resultsCache.delete(id);
+      this.pendingResults.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 {
-    const result = this.resultsCache.get(id);
-    if (result) {
-      return result;
-    } else {
-      // no result yet, check if we know the item
-      const item = this.jobsCache.get(id);
-
-      if (!item) {
-        return { status: 'not-found' };
+  async #getProvingJobStatus(id: ProvingJobId): Promise<ProvingJobStatus> {
+    const settled = this.resultsCache.get(id);
+    if (settled) {
+      if (settled.status !== 'fulfilled') {
+        // rejected/aborted carry their (small) payload inline.
+        return settled;
       }
+      // A fulfilled job's proof value is not held in memory: serve it from the transient hold if its
+      // database write is still in flight, otherwise read it back from the database.
+      const full = this.pendingResults.get(id) ?? (await this.database.getProvingJobResult(id));
+      if (full?.status === 'fulfilled') {
+        return full;
+      }
+      // Settled as fulfilled but the value is nowhere to be found — should not happen (it is held in
+      // pendingResults until its DB write commits). Report not-found rather than a torn result.
+      this.logger.error(`Fulfilled proving job id=${id} has no retrievable result value`, { provingJobId: id });
+      return { status: 'not-found' };
+    }
 
-      return { status: this.inProgress.has(id) ? 'in-progress' : 'in-queue' };
+    // no result yet, check if we know the item
+    const item = this.jobsCache.get(id);
+    if (!item) {
+      return { status: 'not-found' };
     }
+    return { status: this.inProgress.has(id) ? 'in-progress' : 'in-queue' };
   }
 
   #getCompletedJobs(ids: ProvingJobId[]): Promise<ProvingJobId[]> {
@@ -323,7 +449,9 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return Promise.resolve(notifications.concat(completedJobs));
   }
 
-  #getProvingJob(filter: ProvingJobFilter = { allowList: [] }): { job: ProvingJob; time: number } | undefined {
+  async #getProvingJob(
+    filter: ProvingJobFilter = { allowList: [] },
+  ): Promise<{ job: ProvingJob; time: number } | undefined> {
     const allowedProofs: ProvingRequestType[] =
       Array.isArray(filter.allowList) && filter.allowList.length > 0
         ? [...filter.allowList]
@@ -338,21 +466,40 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       // this can happen if the broker crashes and restarts
       // it's possible agents will report progress or results for jobs that are in the queue (after the restart)
       while ((enqueuedJob = queue.getImmediate())) {
-        const job = this.jobsCache.get(enqueuedJob.id);
-        if (job && !this.inProgress.has(enqueuedJob.id) && !this.resultsCache.has(enqueuedJob.id)) {
+        const meta = this.jobsCache.get(enqueuedJob.id);
+        if (meta && !this.inProgress.has(enqueuedJob.id) && !this.resultsCache.has(enqueuedJob.id)) {
           const time = this.msTimeSource();
-          this.inProgress.set(job.id, {
-            id: job.id,
+          // Claim the job synchronously (before the await below) so a concurrent dispatch can't re-pick
+          // it. The same id can sit in a queue more than once (an abort leaves a stale entry that a
+          // later revive re-enqueues alongside), so this in-progress claim — not the queue pop — is what
+          // dedups those copies across concurrently-polling agents. It must be set before the await.
+          this.inProgress.set(meta.id, {
+            id: meta.id,
             startedAt: time,
             lastUpdatedAt: time,
           });
-          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);
+
+          // The large inputs are not kept in memory; read them from the database. They were durably
+          // persisted (addProvingJob's write commits) before the job became dispatchable, so this hits.
+          const inputsUri = await this.database.getProvingJobInputs(meta.id);
+          if (!inputsUri) {
+            // The job was cleaned up (or its inputs lost) after we claimed it — release the claim and
+            // keep draining the queue for another candidate. Nothing else was mutated yet, so there is
+            // no wait metric or enqueued-at timer to unwind.
+            this.inProgress.delete(meta.id);
+            this.logger.warn(`No inputs found for proving job id=${meta.id}; skipping dispatch`, {
+              provingJobId: meta.id,
+            });
+            continue;
           }
 
+          // The dispatch is now committed: record the queue wait and clear the enqueued-at timer.
+          const enqueuedAt = this.enqueuedAt.get(meta.id);
+          if (enqueuedAt) {
+            this.instrumentation.recordJobWait(meta.type, enqueuedAt);
+            this.enqueuedAt.delete(meta.id);
+          }
+          const job: ProvingJob = { id: meta.id, type: meta.type, epochNumber: meta.epochNumber, inputsUri };
           return { job, time };
         }
       }
@@ -419,11 +566,10 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       },
     );
 
-    // save the result to the cache and notify clients of the job status
-    // this should work even if our database breaks because the result is cached in memory
-    const result: ProvingJobSettledResult = { status: 'rejected', reason: String(err) };
-    this.resultsCache.set(id, result);
-    this.promises.get(id)!.resolve(result);
+    // save the (small) rejection reason to the cache and notify clients of the job status. The reason
+    // is kept in memory, so status reads work even if the database write below fails.
+    this.resultsCache.set(id, { status: 'rejected', reason: String(err) });
+    this.promises.get(id)!.resolve();
     this.completedJobNotifications.push(id);
 
     this.instrumentation.incRejectedJobs(item.type);
@@ -445,20 +591,20 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return this.#getProvingJob(filter);
   }
 
-  #reportProvingJobProgress(
+  async #reportProvingJobProgress(
     id: ProvingJobId,
     startedAt: number,
     filter?: ProvingJobFilter,
-  ): { job: ProvingJob; time: number } | undefined {
+  ): Promise<{ job: ProvingJob; time: number } | undefined> {
     const job = this.jobsCache.get(id);
     if (!job) {
       this.logger.warn(`Proving job id=${id} does not exist`, { provingJobId: id });
-      return this.#getProvingJob(filter);
+      return await this.#getProvingJob(filter);
     }
 
     if (this.resultsCache.has(id)) {
       this.logger.warn(`Proving job id=${id} has already been completed`, { provingJobId: id });
-      return this.#getProvingJob(filter);
+      return await this.#getProvingJob(filter);
     }
 
     const metadata = this.inProgress.get(id);
@@ -498,7 +644,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       { provingJobId: id },
     );
 
-    return this.#getProvingJob(filter);
+    return await this.#getProvingJob(filter);
   }
 
   async #reportProvingJobSuccess(
@@ -532,12 +678,12 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       { provingJobId: id },
     );
 
-    // save result to our local cache and notify clients
-    // if save to database fails, that's ok because we have the result in memory
-    // if the broker crashes and needs the result again, we're covered because we can just recompute it
-    const result: ProvingJobSettledResult = { status: 'fulfilled', value };
-    this.resultsCache.set(id, result);
-    this.promises.get(id)!.resolve(result);
+    // Mark settled synchronously (the gate that makes concurrent settles no-op, unchanged from before),
+    // holding only the status in resultsCache. The large proof value is kept transiently in
+    // pendingResults so status reads are read-your-writes until the database write below commits.
+    this.resultsCache.set(id, { status: 'fulfilled' });
+    this.pendingResults.set(id, { status: 'fulfilled', value });
+    this.promises.get(id)!.resolve();
     this.completedJobNotifications.push(id);
 
     this.instrumentation.incResolvedJobs(item.type);
@@ -548,7 +694,11 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
     try {
       await this.database.setProvingJobResult(id, value);
+      // The value is durable in the database now; drop the transient in-memory copy.
+      this.pendingResults.delete(id);
     } catch (saveErr) {
+      // Keep the value in pendingResults so status/result reads still succeed despite the failed write —
+      // this preserves the previous "works even if the database breaks" behaviour for this edge.
       this.logger.error(`Failed to save proving job result id=${id}`, saveErr, {
         provingJobId: id,
       });
@@ -560,21 +710,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,15 +748,30 @@ 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,
+          });
+          this.resultsCache.set(id, { status: 'rejected', reason: 'Timed out' });
+          this.promises.get(id)?.resolve();
+          this.completedJobNotifications.push(id);
+          this.instrumentation.incRejectedJobs(item.type);
+        }
       }
     }
   }
 
-  private enqueueJobInternal(job: ProvingJob): void {
+  private enqueueJobInternal(job: ProvingJobMetadata): void {
     if (!this.promises.has(job.id)) {
       this.promises.set(job.id, promiseWithResolvers());
     }
@@ -618,7 +783,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     this.epochHeight = Math.max(this.epochHeight, job.epochNumber);
   }
 
-  private isJobStale(job: ProvingJob) {
+  private isJobStale(job: ProvingJobMetadata) {
     return job.epochNumber < this.oldestEpochToKeep();
   }
 

package/src/proving_broker/proving_broker_database/memory.ts

@@ -17,8 +17,12 @@ export class InMemoryBrokerDatabase implements ProvingBrokerDatabase {
     return this.jobs.get(id);
   }
 
-  getProvingJobResult(id: ProvingJobId): ProvingJobSettledResult | undefined {
-    return this.results.get(id);
+  getProvingJobInputs(id: ProvingJobId): Promise<ProofUri | undefined> {
+    return Promise.resolve(this.jobs.get(id)?.inputsUri);
+  }
+
+  getProvingJobResult(id: ProvingJobId): Promise<ProvingJobSettledResult | undefined> {
+    return Promise.resolve(this.results.get(id));
   }
 
   addProvingJob(job: ProvingJob): Promise<void> {
@@ -36,6 +40,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);