From 99062abb34f142d88db520371c651d86ae164c58 Mon Sep 17 00:00:00 2001
From: Dimitris Mitsios <di.mitsios@gmail.com>
Date: Tue, 14 Oct 2025 10:35:09 +0300
Subject: [PATCH] modified cli and ServiceMessages to accept proof generation
 message

---
 bin/ream/src/cli/mod.rs                     | 3 +++
 bin/ream/src/main.rs                        | 7 ++++---
 crates/common/chain/lean/src/messages.rs    | 5 +++++
 crates/common/chain/lean/src/service.rs     | 2 ++
 crates/common/validator/lean/src/service.rs | 9 +++++++++
 5 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/bin/ream/src/cli/mod.rs b/bin/ream/src/cli/mod.rs
index e7374358d..59de8e87b 100644
--- a/bin/ream/src/cli/mod.rs
+++ b/bin/ream/src/cli/mod.rs
@@ -48,6 +48,9 @@ pub struct Cli {
 
     #[arg(long, help = "Purges the database.")]
     pub purge_db: bool,
+
+    #[arg(long, short, help = "Proposer proves and gossips state transition, validators verify")]
+    pub proof_gen: bool,
 }
 
 #[derive(Debug, Subcommand)]
diff --git a/bin/ream/src/main.rs b/bin/ream/src/main.rs
index a0244038d..b5fffbebc 100644
--- a/bin/ream/src/main.rs
+++ b/bin/ream/src/main.rs
@@ -95,10 +95,11 @@ fn main() {
     }
 
     let ream_db = ReamDB::new(ream_dir.clone()).expect("unable to init Ream Database");
+    let proof_gen = cli.proof_gen;
 
     match cli.command {
         Commands::LeanNode(config) => {
-            executor_clone.spawn(async move { run_lean_node(*config, executor, ream_db).await });
+            executor_clone.spawn(async move { run_lean_node(*config, executor, ream_db, proof_gen).await });
         }
         Commands::BeaconNode(config) => {
             executor_clone.spawn(async move { run_beacon_node(*config, executor, ream_db).await });
@@ -139,7 +140,7 @@ fn main() {
 /// is used by all services.
 ///
 /// Besides the shared state, each service holds the channels to communicate with each other.
-pub async fn run_lean_node(config: LeanNodeConfig, executor: ReamExecutor, ream_db: ReamDB) {
+pub async fn run_lean_node(config: LeanNodeConfig, executor: ReamExecutor, ream_db: ReamDB, proof_gen: bool) {
     info!("starting up lean node...");
 
     // Initialize prometheus metrics
@@ -215,7 +216,7 @@ pub async fn run_lean_node(config: LeanNodeConfig, executor: ReamExecutor, ream_
     let keystores = load_validator_registry(&config.validator_registry_path, &config.node_id)
         .expect("Failed to load validator registry");
     let validator_service =
-        LeanValidatorService::new(lean_chain_reader.clone(), keystores, chain_sender).await;
+        LeanValidatorService::new(lean_chain_reader.clone(), keystores, chain_sender, proof_gen).await;
 
     let server_config = RpcServerConfig::new(
         config.http_address,
diff --git a/crates/common/chain/lean/src/messages.rs b/crates/common/chain/lean/src/messages.rs
index 187a86524..b08f5ebc6 100644
--- a/crates/common/chain/lean/src/messages.rs
+++ b/crates/common/chain/lean/src/messages.rs
@@ -37,4 +37,9 @@ pub enum LeanChainServiceMessage {
         is_trusted: bool,
         need_gossip: bool,
     },
+
+    ProduceProofBlock {
+        slot: u64,
+        sender: oneshot::Sender<Block>,
+    },
 }
diff --git a/crates/common/chain/lean/src/service.rs b/crates/common/chain/lean/src/service.rs
index 3f0b56b96..9e67be109 100644
--- a/crates/common/chain/lean/src/service.rs
+++ b/crates/common/chain/lean/src/service.rs
@@ -157,6 +157,8 @@ impl LeanChainService {
                                 warn!("Failed to send item to outbound gossip channel: {err:?}");
                             }
                         }
+
+                        LeanChainServiceMessage::ProduceProofBlock { slot, sender } => todo!(),
                     }
                 }
             }
diff --git a/crates/common/validator/lean/src/service.rs b/crates/common/validator/lean/src/service.rs
index a79222f65..710a54ad2 100644
--- a/crates/common/validator/lean/src/service.rs
+++ b/crates/common/validator/lean/src/service.rs
@@ -24,6 +24,7 @@ pub struct ValidatorService {
     lean_chain: LeanChainReader,
     keystores: Vec<LeanKeystore>,
     chain_sender: mpsc::UnboundedSender<LeanChainServiceMessage>,
+    proof: bool,
 }
 
 impl ValidatorService {
@@ -31,11 +32,13 @@ impl ValidatorService {
         lean_chain: LeanChainReader,
         keystores: Vec<LeanKeystore>,
         chain_sender: mpsc::UnboundedSender<LeanChainServiceMessage>,
+        proof: bool,
     ) -> Self {
         ValidatorService {
             lean_chain,
             keystores,
             chain_sender,
+            proof,
         }
     }
 
@@ -77,6 +80,12 @@ impl ValidatorService {
                                     keystore.validator_id,
                                 );
 
+                                if self.proof {
+                                    let (tx, rx) = oneshot::channel();
+                                    self.chain_sender
+                                        .send(LeanChainServiceMessage::ProduceProofBlock { slot, sender: tx })
+                                        .expect("Failed to send vote to LeanChainService");
+                                }
                                 // TODO: Sign the block with the keystore.
                                 let signed_block = SignedBlock {
                                     message: new_block,