update lastest code

clients/cli/src/prover.rs

@@ -32,7 +32,7 @@ use nexus_core::{
         init_circuit_trace, key::CanonicalSerialize, pp::gen_vm_pp, prove_seq_step, types::*,
     },
 };
-use std::env;
+// use std::env;
 use zstd::stream::Encoder;
 use rand::{ RngCore };
 
@@ -125,10 +125,25 @@ async fn main() {
         )),
     };
 
-    client
+    let mut retries = 0;
-        .send(Message::Binary(registration.encode_to_vec()))
+    let max_retries = 5;
-        .await
+
-        .unwrap();
+    loop {
+        if let Err(e) = client.send(Message::Binary(registration.encode_to_vec())).await {
+            eprintln!("Failed to send message: {:?}, attempt {}/{}", e, retries + 1, max_retries);
+
+            retries += 1;
+            if retries >= max_retries {
+                eprintln!("Max retries reached, exiting...");
+                break;
+            }
+
+            // Optionally, add a delay before retrying
+            tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
+        } else {
+            break;
+        }
+    }
 
     track(
         "register".into(),
@@ -261,10 +276,26 @@ async fn main() {
                 }),
             );
             progress_time = SystemTime::now();
-            client
+
-                .send(Message::Binary(progress.encode_to_vec()))
+            let mut retries = 0;
-                .await
+            let max_retries = 5;
-                .unwrap();
+            loop {
+                if let Err(e) = client.send(Message::Binary(progress.encode_to_vec())).await {
+                    eprintln!("Failed to send message: {:?}, attempt {}/{}", e, retries + 1, max_retries);
+        
+                    retries += 1;
+                    if retries >= max_retries {
+                        eprintln!("Max retries reached, exiting...");
+                        break;
+                    }
+        
+                    // Optionally, add a delay before retrying
+                    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
+                } else {
+                    break;
+                }
+            }
+
             if step == end - 1 {
                 let mut buf = Vec::new();
                 let mut writer = Box::new(&mut buf);
@@ -323,4 +354,4 @@ async fn main() {
         &ws_addr_string,
         json!({ "prover_id": prover_id }),
     );
-}
+}