Add compute node deregistration and stake reclamation

Makefile

@@ -188,3 +188,8 @@ submit-work:
 invalidate-work:
 	set -a; source ${ENV_FILE}; set +a; \
 	cargo run -p dev-utils --example invalidate_work -- --pool-id $${POOL_ID:-0} --penalty $${PENALTY} --work-key $${WORK_KEY} --key $${PRIVATE_KEY_VALIDATOR} --rpc-url $${RPC_URL}
+
+deregister-worker:
+	set -a; source ${ENV_FILE}; set +a; \
+	cargo run --bin worker -- deregister --compute-pool-id $${WORKER_COMPUTE_POOL_ID} --private-key-provider $${PRIVATE_KEY_PROVIDER} --private-key-node $${PRIVATE_KEY_NODE} --rpc-url $${RPC_URL}
+

shared/src/web3/contracts/implementations/prime_network_contract.rs

@@ -75,6 +75,26 @@ impl PrimeNetworkContract {
         Ok(add_node_tx)
     }
 
+    pub async fn remove_compute_node(
+        &self,
+        provider_address: Address,
+        node_address: Address,
+    ) -> Result<FixedBytes<32>, Box<dyn std::error::Error>> {
+        let remove_node_tx = self
+            .instance
+            .instance()
+            .function(
+                "removeComputeNode",
+                &[provider_address.into(), node_address.into()],
+            )?
+            .send()
+            .await?
+            .watch()
+            .await?;
+
+        Ok(remove_node_tx)
+    }
+
     pub async fn validate_node(
         &self,
         provider_address: Address,
@@ -225,6 +245,22 @@ impl PrimeNetworkContract {
 
         Ok(members_vec)
     }
+
+    pub async fn reclaim_stake(
+        &self,
+        amount: U256,
+    ) -> Result<FixedBytes<32>, Box<dyn std::error::Error>> {
+        let reclaim_tx = self
+            .instance
+            .instance()
+            .function("reclaimStake", &[amount.into()])?
+            .send()
+            .await?
+            .watch()
+            .await?;
+
+        Ok(reclaim_tx)
+    }
 }
 
 #[cfg(test)]

worker/src/cli/command.rs

@@ -140,6 +140,25 @@ pub enum Commands {
         #[arg(long)]
         private_key_node: Option<String>,
     },
+
+    /// Deregister worker from compute pool
+    Deregister {
+        /// Private key for the provider
+        #[arg(long)]
+        private_key_provider: Option<String>,
+
+        /// Private key for the node
+        #[arg(long)]
+        private_key_node: Option<String>,
+
+        /// RPC URL
+        #[arg(long, default_value = "http://localhost:8545")]
+        rpc_url: String,
+
+        /// Compute pool ID
+        #[arg(long)]
+        compute_pool_id: u64,
+    },
 }
 
 pub async fn execute_command(
@@ -769,6 +788,124 @@ pub async fn execute_command(
 
             println!("\nSignature: {}", hex::encode(combined_signature));
 
+            Ok(())
+        }
+        Commands::Deregister {
+            private_key_provider,
+            private_key_node,
+            rpc_url,
+            compute_pool_id,
+        } => {
+            let private_key_provider = if let Some(key) = private_key_provider {
+                key.clone()
+            } else {
+                std::env::var("PRIVATE_KEY_PROVIDER").expect("PRIVATE_KEY_PROVIDER must be set")
+            };
+
+            let private_key_node = if let Some(key) = private_key_node {
+                key.clone()
+            } else {
+                std::env::var("PRIVATE_KEY_NODE").expect("PRIVATE_KEY_NODE must be set")
+            };
+
+            let provider_wallet_instance = Arc::new(
+                match Wallet::new(&private_key_provider, Url::parse(rpc_url).unwrap()) {
+                    Ok(wallet) => wallet,
+                    Err(err) => {
+                        Console::user_error(&format!("Failed to create wallet: {}", err));
+                        std::process::exit(1);
+                    }
+                },
+            );
+
+            let node_wallet_instance = Arc::new(
+                match Wallet::new(&private_key_node, Url::parse(rpc_url).unwrap()) {
+                    Ok(wallet) => wallet,
+                    Err(err) => {
+                        Console::user_error(&format!("❌ Failed to create wallet: {}", err));
+                        std::process::exit(1);
+                    }
+                },
+            );
+            let state = Arc::new(SystemState::new(None, true, None));
+            /*
+             Initialize dependencies - services, contracts, operations
+            */
+
+            let contracts = Arc::new(
+                ContractBuilder::new(&provider_wallet_instance)
+                    .with_compute_registry()
+                    .with_ai_token()
+                    .with_prime_network()
+                    .with_compute_pool()
+                    .with_stake_manager()
+                    .build()
+                    .unwrap(),
+            );
+
+            let compute_node_ops = ComputeNodeOperations::new(
+                &provider_wallet_instance,
+                &node_wallet_instance,
+                contracts.clone(),
+                state.clone(),
+            );
+
+            let provider_ops =
+                ProviderOperations::new(provider_wallet_instance.clone(), contracts.clone(), false);
+
+            let compute_node_exists = match compute_node_ops.check_compute_node_exists().await {
+                Ok(exists) => exists,
+                Err(e) => {
+                    Console::user_error(&format!(
+                        "❌ Failed to check if compute node exists: {}",
+                        e
+                    ));
+                    std::process::exit(1);
+                }
+            };
+
+            let pool_id = U256::from(*compute_pool_id as u32);
+
+            if compute_node_exists {
+                match contracts
+                    .compute_pool
+                    .leave_compute_pool(
+                        pool_id,
+                        provider_wallet_instance.wallet.default_signer().address(),
+                        node_wallet_instance.wallet.default_signer().address(),
+                    )
+                    .await
+                {
+                    Ok(result) => {
+                        Console::success(&format!("Leave compute pool tx: {:?}", result));
+                    }
+                    Err(e) => {
+                        Console::user_error(&format!("❌ Failed to leave compute pool: {}", e));
+                        std::process::exit(1);
+                    }
+                }
+                match compute_node_ops.remove_compute_node().await {
+                    Ok(_removed_node) => {
+                        Console::success("Compute node removed");
+                        match provider_ops.reclaim_stake(U256::from(0)).await {
+                            Ok(_) => {
+                                Console::success("Successfully reclaimed stake");
+                            }
+                            Err(e) => {
+                                Console::user_error(&format!("❌ Failed to reclaim stake: {}", e));
+                                std::process::exit(1);
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        Console::user_error(&format!("❌ Failed to remove compute node: {}", e));
+                        std::process::exit(1);
+                    }
+                }
+            } else {
+                Console::success("Compute node is not registered");
+            }
+
             Ok(())
         }
     }

worker/src/operations/compute_node.rs

@@ -138,4 +138,23 @@ impl<'c> ComputeNodeOperations<'c> {
         Console::success(&format!("Add node tx: {:?}", add_node_tx));
         Ok(true)
     }
+
+    pub async fn remove_compute_node(&self) -> Result<bool, Box<dyn std::error::Error>> {
+        Console::title("🔄 Removing compute node");
+
+        if !self.check_compute_node_exists().await? {
+            return Ok(false);
+        }
+
+        Console::progress("Removing compute node");
+        let provider_address = self.provider_wallet.wallet.default_signer().address();
+        let node_address = self.node_wallet.wallet.default_signer().address();
+        let remove_node_tx = self
+            .contracts
+            .prime_network
+            .remove_compute_node(provider_address, node_address)
+            .await?;
+        Console::success(&format!("Remove node tx: {:?}", remove_node_tx));
+        Ok(true)
+    }
 }

worker/src/operations/provider.rs

@@ -395,6 +395,23 @@ impl ProviderOperations {
         Console::success("Provider stake increased successfully");
         Ok(())
     }
+
+    pub async fn reclaim_stake(&self, amount: U256) -> Result<(), ProviderError> {
+        Console::progress("Reclaiming stake");
+        let reclaim_tx = match self.contracts.prime_network.reclaim_stake(amount).await {
+            Ok(tx) => tx,
+            Err(e) => {
+                println!("Failed to reclaim stake: {:?}", e);
+                return Err(ProviderError::Other);
+            }
+        };
+        Console::info(
+            "Stake reclaim transaction completed: ",
+            &format!("{:?}", reclaim_tx),
+        );
+        Console::success("Provider stake reclaimed successfully");
+        Ok(())
+    }
 }
 
 #[derive(Debug)]