fix(prover-node): report awaiting-root and publishing-proof phases in EpochSession (A-1212)

yarn-project/prover-node/README.md

@@ -151,23 +151,35 @@ derived from the canonical content for that slot range.
 stateDiagram-v2
   [*] --> initialized
   initialized --> awaiting_checkpoints: start()
-  awaiting_checkpoints --> completed: publish succeeds
-  awaiting_checkpoints --> superseded: longer same-epoch candidate wins
-  awaiting_checkpoints --> failed: L1 submission errored
-  awaiting_checkpoints --> cancelled: cancel()
+  awaiting_checkpoints --> awaiting_root: sub-tree proofs ready, top-tree prove begins
+  awaiting_root --> publishing_proof: epoch proof ready, submit to L1
+  publishing_proof --> completed: publish succeeds
+  publishing_proof --> superseded: longer same-epoch candidate wins
+  publishing_proof --> failed: L1 submission errored
+  awaiting_checkpoints --> failed: top-tree prove errored
+  awaiting_root --> failed: top-tree prove errored
   initialized --> timed_out: deadline
   awaiting_checkpoints --> timed_out: deadline (EpochSession or candidate)
+  awaiting_root --> timed_out: deadline (EpochSession or candidate)
+  awaiting_checkpoints --> cancelled: cancel()
+  awaiting_root --> cancelled: cancel()
   completed --> [*]
   superseded --> [*]
   cancelled --> [*]
   timed_out --> [*]
   failed --> [*]
 ```
 
-The `awaiting-checkpoints` state covers the window between `start()` and the L1
-submission: a `TopTreeJob` is running over the `EpochSession`'s frozen checkpoint set,
-awaiting each checkpoint's sub-tree result (`CheckpointProver.whenBlockProofsReady`)
-and assembling the epoch proof.
+The non-terminal states track the window between `start()` and the L1 submission:
+
+- `awaiting-checkpoints` — a `TopTreeJob` is awaiting each checkpoint's sub-tree result
+  (`CheckpointProver.whenBlockProofsReady`) before the top-tree prove can begin.
+- `awaiting-root` — the sub-tree proofs are ready and the top-tree (root) prove is running,
+  assembling the epoch proof (set via the `TopTreeJob` `beforeProve` hook).
+- `publishing-proof` — the epoch proof is being submitted to L1 via `ProofPublishingService`.
+
+`cancel()` and the deadline can fire during any of these pre-submit phases; a terminal state
+set that way wins over the phase transitions (which are guarded by `isTerminal()`).
 
 The `EpochSession` does three sequential things: (1) run a `TopTreeJob` over the frozen
 checkpoint subset, (2) hand the resulting proof to `ProofPublishingService` as a

yarn-project/prover-node/src/job/epoch-session.test.ts

@@ -19,7 +19,13 @@ import { mock } from 'jest-mock-extended';
 import { ProverNodeJobMetrics } from '../metrics.js';
 import type { ProofPublishingService, PublishCandidate, PublishOutcome } from '../proof-publishing-service.js';
 import { CheckpointProver } from './checkpoint-prover.js';
-import { EpochSession, type EpochSessionDeps, type EpochSessionHooks, type SessionSpec } from './epoch-session.js';
+import {
+  type EpochProvingJobState,
+  EpochSession,
+  type EpochSessionDeps,
+  type EpochSessionHooks,
+  type SessionSpec,
+} from './epoch-session.js';
 import type { TopTreeProof } from './top-tree-job.js';
 
 describe('EpochSession', () => {
@@ -325,6 +331,34 @@ describe('EpochSession', () => {
     });
   });
 
+  // ---------------- state reporting ----------------
+
+  describe('state reporting', () => {
+    it('advances through awaiting-root (while proving) and publishing-proof (while submitting)', async () => {
+      let stateDuringProve: EpochProvingJobState | undefined;
+      let stateDuringSubmit: EpochProvingJobState | undefined;
+      const session = makeSession({
+        hooks: {
+          // beforeProve has already flipped the state by the time the prove runs.
+          topTreeProveOverride: () => {
+            stateDuringProve = session.getState();
+            return Promise.resolve(synthProof);
+          },
+        },
+      });
+      publishingService.submit.mockImplementation(() => {
+        stateDuringSubmit = session.getState();
+        return Promise.resolve('published');
+      });
+
+      const state = await session.start();
+
+      expect(stateDuringProve).toBe('awaiting-root');
+      expect(stateDuringSubmit).toBe('publishing-proof');
+      expect(state).toBe('completed');
+    });
+  });
+
   // ---------------- helpers ----------------
 
   /** Default session spec used by every test that doesn't override it. */

yarn-project/prover-node/src/job/epoch-session.ts

@@ -92,7 +92,7 @@ export type EpochSessionDeps = {
  *
  * Lifecycle (happy path):
  *
- *   initialized → awaiting-checkpoints → completed
+ *   initialized → awaiting-checkpoints → awaiting-root → publishing-proof → completed
  *
  * Terminal states map the publishing outcome: `published` → `completed`, `superseded` →
  * `superseded`, `failed` → `failed`, `expired` → `timed-out`, `withdrawn` → `cancelled`.
@@ -320,6 +320,12 @@ export class EpochSession implements Traceable {
       0,
     );
 
+    // Reflect the publish phase. Guard against a terminal state set concurrently by cancel() — the
+    // post-submit isTerminal() check below relies on cancel still winning.
+    if (!this.isTerminal()) {
+      this.state = 'publishing-proof';
+    }
+
     const outcome = await this.deps.publishingService.submit({
       id: this.uuid,
       epoch: this.spec.epochNumber,
@@ -410,15 +416,20 @@ export class EpochSession implements Traceable {
     }
   }
 
-  private toTopTreeHooks(): TopTreeJobHooks | undefined {
+  private toTopTreeHooks(): TopTreeJobHooks {
     const hooks = this.deps.hooks;
-    if (!hooks?.beforeTopTreeProve && !hooks?.afterTopTreeProve && !hooks?.topTreeProveOverride) {
-      return undefined;
-    }
     return {
-      beforeProve: hooks.beforeTopTreeProve,
-      afterProve: hooks.afterTopTreeProve,
-      proveOverride: hooks.topTreeProveOverride,
+      // `beforeProve` fires once the sub-tree (checkpoint block) proofs are ready and the root prove is
+      // about to start — the boundary between `awaiting-checkpoints` and proving the top tree. Don't
+      // clobber a terminal state set concurrently by cancel().
+      beforeProve: async () => {
+        if (!this.isTerminal()) {
+          this.state = 'awaiting-root';
+        }
+        await hooks?.beforeTopTreeProve?.();
+      },
+      afterProve: hooks?.afterTopTreeProve,
+      proveOverride: hooks?.topTreeProveOverride,
     };
   }
 }

yarn-project/stdlib/src/interfaces/prover-node.test.ts

@@ -40,7 +40,8 @@ class MockProverNode implements ProverNodeApi {
     return Promise.resolve([
       { uuid: 'uuid1', status: 'initialized', epochNumber: 10 },
       { uuid: 'uuid2', status: 'awaiting-checkpoints', epochNumber: 10 },
-      { uuid: 'uuid3', status: 'awaiting-predecessor', epochNumber: 10 },
+      { uuid: 'uuid3', status: 'awaiting-root', epochNumber: 10 },
+      { uuid: 'uuid3b', status: 'awaiting-predecessor', epochNumber: 10 },
       { uuid: 'uuid4', status: 'publishing-proof', epochNumber: 10 },
       { uuid: 'uuid5', status: 'completed', epochNumber: 10 },
       { uuid: 'uuid6', status: 'superseded', epochNumber: 10 },

yarn-project/stdlib/src/interfaces/prover-node.ts

@@ -8,6 +8,7 @@ import { type ComponentsVersions, getVersioningResponseHandler } from '../versio
 const EpochProvingJobState = [
   'initialized',
   'awaiting-checkpoints',
+  'awaiting-root',
   'awaiting-predecessor',
   'publishing-proof',
   'completed',