From 22e99c6f1bb843d4467c869995ca532b0aec6e24 Mon Sep 17 00:00:00 2001
From: Santiago Palladino <santiago@aztecprotocol.com>
Date: Wed, 28 Aug 2024 12:11:44 -0300
Subject: [PATCH] feat: Add status check to prover agent

Adds a status check and an endpoint for querying current jobs to the
prover agent:

```
$ curl -XPOST -d'{"method": "prover_getCurrentJobs"}' http://localhost:8085/
{"result":[{"id":"ab3f1629","type":"ROOT_PARITY"}]}
```
---
 docker-compose.provernet.yml                  |  5 ++++
 .../aztec/src/cli/cmds/start_prover_agent.ts  |  8 +++++--
 .../src/prover-agent/prover-agent.ts          |  4 ++++
 .../prover-client/src/prover-agent/rpc.ts     | 24 +++++++++++++++++++
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/docker-compose.provernet.yml b/docker-compose.provernet.yml
index ab9154d2c9d..7fdf67c3b6a 100644
--- a/docker-compose.provernet.yml
+++ b/docker-compose.provernet.yml
@@ -145,6 +145,11 @@ services:
         condition: service_healthy
     command: [ "start", "--prover" ]
     restart: on-failure:5
+    healthcheck:
+      test: [ "CMD", "curl", "-fSs", "http://127.0.0.1:$AZTEC_PORT/status" ]
+      interval: 3s
+      timeout: 30s
+      start_period: 20s
 
   # Simple watcher that logs the latest block numbers every few seconds using the CLI and the bot's PXE
   aztec-block-watcher:
diff --git a/yarn-project/aztec/src/cli/cmds/start_prover_agent.ts b/yarn-project/aztec/src/cli/cmds/start_prover_agent.ts
index 0a6534fae54..8fa31fd6428 100644
--- a/yarn-project/aztec/src/cli/cmds/start_prover_agent.ts
+++ b/yarn-project/aztec/src/cli/cmds/start_prover_agent.ts
@@ -1,7 +1,11 @@
 import { BBNativeRollupProver, TestCircuitProver } from '@aztec/bb-prover';
 import { type ServerCircuitProver } from '@aztec/circuit-types';
 import { type ProverClientConfig, proverClientConfigMappings } from '@aztec/prover-client';
-import { ProverAgent, createProvingJobSourceClient } from '@aztec/prover-client/prover-agent';
+import {
+  ProverAgent,
+  createProverAgentRpcServer,
+  createProvingJobSourceClient,
+} from '@aztec/prover-client/prover-agent';
 import {
   type TelemetryClientConfig,
   createAndStartTelemetryClient,
@@ -43,5 +47,5 @@ export const startProverAgent: ServiceStarter = async (options, signalHandlers,
 
   signalHandlers.push(() => agent.stop());
 
-  return Promise.resolve([]);
+  return [{ prover: createProverAgentRpcServer(agent) }];
 };
diff --git a/yarn-project/prover-client/src/prover-agent/prover-agent.ts b/yarn-project/prover-client/src/prover-agent/prover-agent.ts
index 2dd1c5b4727..fec9869e78e 100644
--- a/yarn-project/prover-client/src/prover-agent/prover-agent.ts
+++ b/yarn-project/prover-client/src/prover-agent/prover-agent.ts
@@ -53,6 +53,10 @@ export class ProverAgent {
     return this.runningPromise?.isRunning() ?? false;
   }
 
+  getCurrentJobs(): { id: string; type: string }[] {
+    return Array.from(this.inFlightPromises.values()).map(({ id, type }) => ({ id, type: ProvingRequestType[type] }));
+  }
+
   start(jobSource: ProvingJobSource): void {
     if (this.runningPromise) {
       throw new Error('Agent is already running');
diff --git a/yarn-project/prover-client/src/prover-agent/rpc.ts b/yarn-project/prover-client/src/prover-agent/rpc.ts
index 8f18f7b2889..4474882f187 100644
--- a/yarn-project/prover-client/src/prover-agent/rpc.ts
+++ b/yarn-project/prover-client/src/prover-agent/rpc.ts
@@ -1,12 +1,14 @@
 import { type ProvingJobSource } from '@aztec/circuit-types';
 import {
   AvmCircuitInputs,
+  AztecAddress,
   BaseOrMergeRollupPublicInputs,
   BaseParityInputs,
   BaseRollupInputs,
   BlockMergeRollupInputs,
   BlockRootOrBlockMergePublicInputs,
   BlockRootRollupInputs,
+  EthAddress,
   Fr,
   Header,
   KernelCircuitPublicInputs,
@@ -28,6 +30,7 @@ import {
 import { createJsonRpcClient, makeFetch } from '@aztec/foundation/json-rpc/client';
 import { JsonRpcServer } from '@aztec/foundation/json-rpc/server';
 
+import { type ProverAgent } from './prover-agent.js';
 import { ProvingError } from './proving-error.js';
 
 export function createProvingJobSourceServer(queue: ProvingJobSource): JsonRpcServer {
@@ -104,3 +107,24 @@ export function createProvingJobSourceClient(
     fetch,
   ) as ProvingJobSource;
 }
+
+/**
+ * Wrap a ProverAgent instance with a JSON RPC HTTP server.
+ * @param node - The ProverNode
+ * @returns An JSON-RPC HTTP server
+ */
+export function createProverAgentRpcServer(agent: ProverAgent) {
+  const rpc = new JsonRpcServer(
+    agent,
+    {
+      AztecAddress,
+      EthAddress,
+      Fr,
+      Header,
+    },
+    {},
+    // disable methods
+    ['start', 'stop', 'setCircuitProver', 'work', 'getProof'],
+  );
+  return rpc;
+}