update error handling for connecting to websockets

clients/cli/src/prover.rs

@@ -72,27 +72,59 @@ async fn main() {
         .expect("error generating public parameters");
 
     // If the prover_id file is found, use the contents, otherwise generate a new random id
-    // and store it.
-    let mut prover_id = format!(
+    // and store it. e.g., "happy-cloud-42"
+    let default_prover_id: String = format!(
         "{}-{}-{}",
         random_word::gen(Lang::En),
         random_word::gen(Lang::En),
         rand::thread_rng().next_u32() % 100,
     );
-    match home::home_dir() {
+
+    // setting the prover-id we will use (either from the file or generated)
+    let prover_id: String = match home::home_dir() {
         Some(path) if !path.as_os_str().is_empty() => {
             let nexus_dir = Path::new(&path).join(".nexus");
-            prover_id = match fs::read(nexus_dir.join("prover-id")) {
-                Ok(buf) => String::from_utf8(buf).unwrap(),
-                Err(_) => {
-                    let _ = fs::create_dir(nexus_dir.clone());
-                    fs::write(nexus_dir.join("prover-id"), prover_id.clone()).unwrap();
-                    prover_id
+
+            // Try to read the prover-id file
+            match fs::read(nexus_dir.join("prover-id")) {
+                // 1. If file exists and can be read:
+                Ok(buf) => match String::from_utf8(buf) {
+                    Ok(id) => id.trim().to_string(), // Trim whitespace
+                    Err(e) => {
+                        eprintln!("Failed to read prover-id file. Using default: {}", e);
+                        default_prover_id // Fall back to generated ID, if file has invalid UTF-8
+                    },
+                },
+                // 2. If file doesn't exist or can't be read:
+                Err(e) => {
+                    eprintln!("Could not read prover-id file: {}", e);
+
+                    // if the error is because the file doesn't exist
+                    // Try to save the generated prover-id to the file
+                    if e.kind() == std::io::ErrorKind::NotFound {
+
+                        // Try to create the .nexus directory
+                        match fs::create_dir(nexus_dir.clone()) {
+                            Ok(_) => {
+                                // Only try to write file if directory was created successfully
+                                if let Err(e) = fs::write(nexus_dir.join("prover-id"), &default_prover_id) {
+                                    eprintln!("Warning: Could not save prover-id: {}", e);
+                                }
+                            },
+                            Err(e) => {
+                                eprintln!("Failed to create .nexus directory: {}", e);
+                            },
+                        }
+                    }
+
+                    // Use the previously generated prover-id
+                    default_prover_id
                 }
             }
         }
         _ => {
-            println!("Unable to get home dir.");
+            println!("Unable to determine home directory. Using temporary prover-id.");
+            default_prover_id
         }
     };
 
@@ -103,16 +135,60 @@ async fn main() {
         json!({"prover_id": prover_id}),
     );
 
-    let (mut client, _) = tokio_tungstenite::connect_async(&ws_addr_string)
-        .await
-        .unwrap();
+    // This function connects to the Orchestrator via websockets
+    // and returns the connected client
+    async fn connect_to_orchestrator(ws_addr: &str) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>, Box<dyn std::error::Error>> {
         
+        // Connect to the Orchestrator via websockets
+        let (client, _) = tokio_tungstenite::connect_async(ws_addr)
+        .await
+        // If the connection fails, print an error and return the error
+        .map_err(|e| {
+            eprintln!("Failed to connect to orchestrator at {}: {}", ws_addr, e);
+            e
+        })?;
+
+        // Return the connected client
+        Ok(client)
+    }
+
+    /// This function wraps connect_to_orchestrator and retries
+    /// with exponential backoff if the connection fails
+    async fn connect_to_orchestrator_with_retry(ws_addr: &str) -> WebSocketStream<MaybeTlsStream<TcpStream>> {
+        let mut attempt = 1;
+
+        loop {
+            match connect_to_orchestrator(ws_addr).await {
+                Ok(client) => {
                     track(
                         "connected".into(),
                         "Connected.".into(),
                         &ws_addr_string,
                         json!({"prover_id": prover_id}),
                     );
+                    return client;
+                },
+                Err(e) => {
+
+                    eprintln!(
+                        "Could not connect to orchestrator (attempt {}). Retrying in {} seconds...", 
+                        attempt,
+                        2u64.pow(attempt.min(6)), // Cap exponential backoff at 64 seconds
+                    );
+                    
+                    // Exponential backoff
+                    tokio::time::sleep(
+                        tokio::time::Duration::from_secs(2u64.pow(attempt.min(6)))
+                    ).await;
+
+                    attempt += 1;
+                }
+            }
+        }
+    }
+
+    // Connect to the Orchestrator with exponential backoff
+    let mut client = connect_to_orchestrator_with_retry(&ws_addr_string).await;    
 
     let registration = ProverRequest {
         contents: Some(prover_request::Contents::Registration(