@aztec/prover-client 0.65.1 → 0.65.2

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

@@ -1,9 +1,9 @@
 import {
   type ProverAgentApi,
   type ProvingJob,
+  type ProvingJobInputs,
+  type ProvingJobResultsMap,
   type ProvingJobSource,
-  type ProvingRequest,
-  type ProvingRequestResultFor,
   ProvingRequestType,
   type ServerCircuitProver,
   makeProvingRequestResult,
@@ -12,6 +12,8 @@ import { createDebugLogger } from '@aztec/foundation/log';
 import { RunningPromise } from '@aztec/foundation/running-promise';
 import { elapsed } from '@aztec/foundation/timer';
 
+import { InlineProofStore } from '../proving_broker/proof_store.js';
+
 const PRINT_THRESHOLD_NS = 6e10; // 60 seconds
 
 /**
@@ -27,6 +29,7 @@ export class ProverAgent implements ProverAgentApi {
     }
   >();
   private runningPromise?: RunningPromise;
+  private proofInputsDatabase = new InlineProofStore();
 
   constructor(
     /** The prover implementation to defer jobs to */
@@ -101,12 +104,12 @@ export class ProverAgent implements ProverAgentApi {
             const promise = this.work(jobSource, job).finally(() => this.inFlightPromises.delete(job.id));
             this.inFlightPromises.set(job.id, {
               id: job.id,
-              type: job.request.type,
+              type: job.type,
               promise,
             });
           } catch (err) {
             this.log.warn(
-              `Error processing job! type=${ProvingRequestType[job.request.type]}: ${err}. ${(err as Error).stack}`,
+              `Error processing job! type=${ProvingRequestType[job.type]}: ${err}. ${(err as Error).stack}`,
             );
           }
         } catch (err) {
@@ -130,28 +133,24 @@ export class ProverAgent implements ProverAgentApi {
     this.log.info('Agent stopped');
   }
 
-  private async work<TRequest extends ProvingRequest>(
-    jobSource: ProvingJobSource,
-    job: ProvingJob<TRequest>,
-  ): Promise<void> {
+  private async work(jobSource: ProvingJobSource, job: ProvingJob): Promise<void> {
     try {
-      this.log.debug(`Picked up proving job id=${job.id} type=${ProvingRequestType[job.request.type]}`);
-      const type: TRequest['type'] = job.request.type;
-      const [time, result] = await elapsed(this.getProof(job.request));
+      this.log.debug(`Picked up proving job id=${job.id} type=${ProvingRequestType[job.type]}`);
+      const type = job.type;
+      const inputs = await this.proofInputsDatabase.getProofInput(job.inputsUri);
+      const [time, result] = await elapsed(this.getProof(inputs));
       if (this.#isRunning()) {
         this.log.verbose(`Processed proving job id=${job.id} type=${ProvingRequestType[type]} duration=${time}ms`);
         await jobSource.resolveProvingJob(job.id, makeProvingRequestResult(type, result));
       } else {
         this.log.verbose(
-          `Dropping proving job id=${job.id} type=${
-            ProvingRequestType[job.request.type]
-          } duration=${time}ms: agent stopped`,
+          `Dropping proving job id=${job.id} type=${ProvingRequestType[job.type]} duration=${time}ms: agent stopped`,
         );
       }
     } catch (err) {
-      const type = ProvingRequestType[job.request.type];
+      const type = ProvingRequestType[job.type];
       if (this.#isRunning()) {
-        if (job.request.type === ProvingRequestType.PUBLIC_VM && !process.env.AVM_PROVING_STRICT) {
+        if (job.type === ProvingRequestType.PUBLIC_VM && !process.env.AVM_PROVING_STRICT) {
           this.log.warn(`Expected error processing VM proving job id=${job.id} type=${type}: ${err}`);
         } else {
           this.log.error(`Error processing proving job id=${job.id} type=${type}: ${err}`, err);
@@ -164,10 +163,7 @@ export class ProverAgent implements ProverAgentApi {
     }
   }
 
-  private getProof<TRequest extends ProvingRequest>(
-    request: TRequest,
-  ): Promise<ProvingRequestResultFor<TRequest['type']>['result']>;
-  private getProof(request: ProvingRequest): Promise<ProvingRequestResultFor<typeof type>['result']> {
+  private getProof(request: ProvingJobInputs): Promise<ProvingJobResultsMap[ProvingRequestType]> {
     const { type, inputs } = request;
     switch (type) {
       case ProvingRequestType.PUBLIC_VM: {

package/src/proving_broker/caching_broker_facade.ts

@@ -0,0 +1,312 @@
+import {
+  type ProofAndVerificationKey,
+  type ProverCache,
+  type ProvingJobId,
+  type ProvingJobInputsMap,
+  type ProvingJobProducer,
+  type ProvingJobResultsMap,
+  ProvingRequestType,
+  type PublicInputsAndRecursiveProof,
+  type ServerCircuitProver,
+} from '@aztec/circuit-types';
+import {
+  type AVM_PROOF_LENGTH_IN_FIELDS,
+  type AvmCircuitInputs,
+  type BaseOrMergeRollupPublicInputs,
+  type BaseParityInputs,
+  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';
+import { sha256 } from '@aztec/foundation/crypto';
+import { createDebugLogger } from '@aztec/foundation/log';
+import { retryUntil } from '@aztec/foundation/retry';
+
+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 {
+  constructor(
+    private broker: ProvingJobProducer,
+    private cache: ProverCache = new InMemoryProverCache(),
+    private proofStore: ProofStore = new InlineProofStore(),
+    private waitTimeoutMs = MAX_WAIT_MS,
+    private pollIntervalMs = 1000,
+    private log = createDebugLogger('aztec:prover-client:caching-prover-broker'),
+  ) {}
+
+  private async enqueueAndWaitForJob<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+    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`);
+    }
+
+    if (!jobEnqueued) {
+      try {
+        const inputsUri = await this.proofStore.saveProofInput(id, type, inputs);
+        await this.broker.enqueueProvingJob({
+          id,
+          type,
+          inputsUri,
+        });
+        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;
+      }
+    }
+
+    // notify broker of cancelled job
+    const abortFn = async () => {
+      signal?.removeEventListener('abort', abortFn);
+      await this.broker.removeAndCancelProvingJob(id);
+    };
+
+    signal?.addEventListener('abort', abortFn);
+
+    try {
+      // loop here until the job settles
+      // NOTE: this could also terminate because the job was cancelled through event listener above
+      const result = await retryUntil(
+        async () => {
+          try {
+            return await this.broker.waitForJobToSettle(id);
+          } catch (err) {
+            // waitForJobToSettle can only fail for network errors
+            // keep retrying until we time out
+          }
+        },
+        `Proving job=${id} type=${ProvingRequestType[type]}`,
+        this.waitTimeoutMs / 1000,
+        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}`));
+        }
+      } else {
+        return Promise.reject(new Error(result.reason));
+      }
+    } finally {
+      signal?.removeEventListener('abort', abortFn);
+    }
+  }
+
+  getAvmProof(
+    inputs: AvmCircuitInputs,
+    signal?: AbortSignal,
+    _blockNumber?: number,
+  ): Promise<ProofAndVerificationKey<typeof AVM_PROOF_LENGTH_IN_FIELDS>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.PUBLIC_VM, inputs),
+      ProvingRequestType.PUBLIC_VM,
+      inputs,
+      signal,
+    );
+  }
+
+  getBaseParityProof(
+    inputs: BaseParityInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.BASE_PARITY, inputs),
+      ProvingRequestType.BASE_PARITY,
+      inputs,
+      signal,
+    );
+  }
+
+  getBlockMergeRollupProof(
+    input: BlockMergeRollupInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.BLOCK_MERGE_ROLLUP, input),
+      ProvingRequestType.BLOCK_MERGE_ROLLUP,
+      input,
+      signal,
+    );
+  }
+
+  getBlockRootRollupProof(
+    input: BlockRootRollupInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.BLOCK_ROOT_ROLLUP, input),
+      ProvingRequestType.BLOCK_ROOT_ROLLUP,
+      input,
+      signal,
+    );
+  }
+
+  getEmptyBlockRootRollupProof(
+    input: EmptyBlockRootRollupInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP, input),
+      ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP,
+      input,
+      signal,
+    );
+  }
+
+  getEmptyPrivateKernelProof(
+    inputs: PrivateKernelEmptyInputData,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<KernelCircuitPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.PRIVATE_KERNEL_EMPTY, inputs),
+      ProvingRequestType.PRIVATE_KERNEL_EMPTY,
+      inputs,
+      signal,
+    );
+  }
+
+  getMergeRollupProof(
+    input: MergeRollupInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.MERGE_ROLLUP, input),
+      ProvingRequestType.MERGE_ROLLUP,
+      input,
+      signal,
+    );
+  }
+  getPrivateBaseRollupProof(
+    baseRollupInput: PrivateBaseRollupInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.PRIVATE_BASE_ROLLUP, baseRollupInput),
+      ProvingRequestType.PRIVATE_BASE_ROLLUP,
+      baseRollupInput,
+      signal,
+    );
+  }
+
+  getPublicBaseRollupProof(
+    inputs: PublicBaseRollupInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.PUBLIC_BASE_ROLLUP, inputs),
+      ProvingRequestType.PUBLIC_BASE_ROLLUP,
+      inputs,
+      signal,
+    );
+  }
+
+  getRootParityProof(
+    inputs: RootParityInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof NESTED_RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.ROOT_PARITY, inputs),
+      ProvingRequestType.ROOT_PARITY,
+      inputs,
+      signal,
+    );
+  }
+
+  getRootRollupProof(
+    input: RootRollupInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<PublicInputsAndRecursiveProof<RootRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.ROOT_ROLLUP, input),
+      ProvingRequestType.ROOT_ROLLUP,
+      input,
+      signal,
+    );
+  }
+
+  getTubeProof(
+    tubeInput: TubeInputs,
+    signal?: AbortSignal,
+    _epochNumber?: number,
+  ): Promise<ProofAndVerificationKey<typeof TUBE_PROOF_LENGTH>> {
+    return this.enqueueAndWaitForJob(
+      this.generateId(ProvingRequestType.TUBE_PROOF, tubeInput),
+      ProvingRequestType.TUBE_PROOF,
+      tubeInput,
+      signal,
+    );
+  }
+
+  private generateId(type: ProvingRequestType, inputs: { toBuffer(): Buffer }) {
+    const inputsHash = sha256(inputs.toBuffer());
+    return `${ProvingRequestType[type]}:${inputsHash.toString('hex')}`;
+  }
+}

package/src/proving_broker/factory.ts

@@ -0,0 +1,21 @@
+import { type ProverBrokerConfig } from '@aztec/circuit-types';
+import { AztecLmdbStore } from '@aztec/kv-store/lmdb';
+
+import { ProvingBroker } from './proving_broker.js';
+import { InMemoryBrokerDatabase } from './proving_broker_database/memory.js';
+import { KVBrokerDatabase } from './proving_broker_database/persisted.js';
+
+export async function createAndStartProvingBroker(config: ProverBrokerConfig): Promise<ProvingBroker> {
+  const database = config.proverBrokerDataDirectory
+    ? new KVBrokerDatabase(AztecLmdbStore.open(config.proverBrokerDataDirectory))
+    : new InMemoryBrokerDatabase();
+
+  const broker = new ProvingBroker(database, {
+    jobTimeoutMs: config.proverBrokerJobTimeoutMs,
+    maxRetries: config.proverBrokerJobMaxRetries,
+    timeoutIntervalMs: config.proverBrokerPollIntervalMs,
+  });
+
+  await broker.start();
+  return broker;
+}

package/src/proving_broker/index.ts

@@ -0,0 +1,8 @@
+export * from './proving_agent.js';
+export * from './proving_broker.js';
+export * from './rpc.js';
+export * from './proving_broker_database.js';
+export * from './proving_broker_database/memory.js';
+export * from './proving_broker_database/persisted.js';
+export * from './proof_store.js';
+export * from './factory.js';

package/src/proving_broker/proof_store.ts

@@ -0,0 +1,106 @@
+import {
+  type ProofUri,
+  type ProvingJobId,
+  ProvingJobInputs,
+  type ProvingJobInputsMap,
+  ProvingJobResult,
+  type ProvingJobResultsMap,
+  type ProvingRequestType,
+} from '@aztec/circuit-types';
+import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
+import { type ZodFor } from '@aztec/foundation/schemas';
+
+/**
+ * A database for storing proof inputs and outputs.
+ */
+export interface ProofStore {
+  /**
+   * Save a proof input to the database.
+   * @param jobId - The ID of the job the proof input is associated with.
+   * @param type - The type of the proving request.
+   * @param inputs - The proof input to save.
+   * @returns The URI of the saved proof input.
+   */
+  saveProofInput<T extends ProvingRequestType>(
+    jobId: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+  ): Promise<ProofUri>;
+
+  /**
+   * Save a proof output to the database.
+   * @param jobId - The ID of the job the proof input is associated with.
+   * @param type - The type of the proving request.
+   * @param result - The proof output to save.
+   * @returns The URI of the saved proof output.
+   */
+  saveProofOutput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    result: ProvingJobResultsMap[T],
+  ): Promise<ProofUri>;
+
+  /**
+   * Retrieve a proof input from the database.
+   * @param uri - The URI of the proof input to retrieve.
+   * @returns The proof input.
+   */
+  getProofInput(uri: ProofUri): Promise<ProvingJobInputs>;
+
+  /**
+   * Retrieve a proof output from the database.
+   * @param uri - The URI of the proof output to retrieve.
+   * @returns The proof output.
+   */
+  getProofOutput(uri: ProofUri): Promise<ProvingJobResult>;
+}
+
+// use an ASCII encoded data uri https://datatracker.ietf.org/doc/html/rfc2397#section-2
+// we do this to avoid double encoding to base64 (since the inputs already serialize to a base64 string)
+const PREFIX = 'data:application/json;charset=utf-8';
+const SEPARATOR = ',';
+
+/**
+ * An implementation of a proof input/output database that stores data inline in the URI.
+ */
+export class InlineProofStore implements ProofStore {
+  saveProofInput<T extends ProvingRequestType>(
+    _id: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+  ): Promise<ProofUri> {
+    const jobInputs = { type, inputs } as ProvingJobInputs;
+    return Promise.resolve(this.encode(jobInputs));
+  }
+
+  saveProofOutput<T extends ProvingRequestType>(
+    _id: ProvingJobId,
+    type: T,
+    result: ProvingJobResultsMap[T],
+  ): Promise<ProofUri> {
+    const jobResult = { type, result } as ProvingJobResult;
+    return Promise.resolve(this.encode(jobResult));
+  }
+
+  getProofInput(uri: ProofUri): Promise<ProvingJobInputs> {
+    return Promise.resolve(this.decode(uri, ProvingJobInputs));
+  }
+
+  getProofOutput(uri: ProofUri): Promise<ProvingJobResult> {
+    return Promise.resolve(this.decode(uri, ProvingJobResult));
+  }
+
+  private encode(obj: object): ProofUri {
+    const encoded = encodeURIComponent(jsonStringify(obj));
+    return (PREFIX + SEPARATOR + encoded) as ProofUri;
+  }
+
+  private decode<T>(uri: ProofUri, schema: ZodFor<T>): T {
+    const [prefix, data] = uri.split(SEPARATOR);
+    if (prefix !== PREFIX) {
+      throw new Error('Invalid proof input URI: ' + prefix);
+    }
+
+    return jsonParseWithSchema(decodeURIComponent(data), schema);
+  }
+}

package/src/proving_broker/prover_cache/memory.ts

@@ -0,0 +1,20 @@
+import type { ProverCache, ProvingJobStatus } from '@aztec/circuit-types';
+
+export class InMemoryProverCache implements ProverCache {
+  private proofs: Record<string, ProvingJobStatus> = {};
+
+  constructor() {}
+
+  setProvingJobStatus(jobId: string, status: ProvingJobStatus): Promise<void> {
+    this.proofs[jobId] = status;
+    return Promise.resolve();
+  }
+
+  getProvingJobStatus(jobId: string): Promise<ProvingJobStatus> {
+    return Promise.resolve(this.proofs[jobId] ?? { status: 'not-found' });
+  }
+
+  close(): Promise<void> {
+    return Promise.resolve();
+  }
+}