[Fix] Select fee based on consensus height

leo/cli/commands/deploy.rs

@@ -193,7 +193,7 @@ fn handle_deploy<A: Aleo<Network = N, BaseField = N::Field>, N: Network>(
             }
             None => {
                 // Make sure the user has enough public balance to pay for the deployment.
-                check_balance(&private_key, endpoint, &network.to_string(), context.clone(), total_cost)?;
+                check_balance(&private_key, endpoint, &network.to_string(), &context, total_cost)?;
                 let fee_authorization = vm.authorize_fee_public(
                     &private_key,
                     total_cost,

leo/cli/commands/execute.rs

@@ -42,6 +42,7 @@ use snarkvm::{
         VM,
         Value,
         execution_cost_v1,
+        execution_cost_v2,
         query::Query as SnarkVMQuery,
         store::{
             ConsensusStore,
@@ -213,8 +214,33 @@ fn handle_execute<A: Aleo>(
         // Check the transaction cost.
         let (mut total_cost, (storage_cost, finalize_cost)) = if let ExecuteTransaction(_, execution, _) = &transaction
         {
-            // TODO: Update to V2 after block migration.
-            execution_cost_v1(&vm.process().read(), execution)?
+            // Attempt to get the height of the latest block to determine which version of the execution cost to use.
+            if let Ok(height) = get_latest_block_height(endpoint, &network.to_string(), &context) {
+                if height < A::Network::CONSENSUS_V2_HEIGHT {
+                    execution_cost_v1(&vm.process().read(), execution)?
+                } else {
+                    execution_cost_v2(&vm.process().read(), execution)?
+                }
+            }
+            // Otherwise, default to the one provided in `fee_options`.
+            else {
+                // Get the consensus version from the command.
+                let version = match command.fee_options.consensus_version {
+                    Some(1) => 1,
+                    None | Some(2) => 2,
+                    Some(version) => {
+                        panic!("Invalid consensus version: {version}. Please specify a valid version.")
+                    }
+                };
+                // Print a warning message.
+                println!("Failed to get the latest block height. Defaulting to V{version}.",);
+                // Use the provided version.
+                match version {
+                    1 => execution_cost_v1(&vm.process().read(), execution)?,
+                    2 => execution_cost_v2(&vm.process().read(), execution)?,
+                    _ => unreachable!(),
+                }
+            }
         } else {
             panic!("All transactions should be of type Execute.")
         };
@@ -231,7 +257,7 @@ fn handle_execute<A: Aleo>(
 
         // Check if the public balance is sufficient.
         if fee_record.is_none() {
-            check_balance::<A::Network>(&private_key, endpoint, &network.to_string(), context, total_cost)?;
+            check_balance::<A::Network>(&private_key, endpoint, &network.to_string(), &context, total_cost)?;
         }
 
         // Broadcast the execution transaction.

leo/cli/commands/mod.rs

@@ -220,6 +220,8 @@ pub struct FeeOptions {
         long
     )]
     record: Option<String>,
+    #[clap(long, help = "Consensus version to use for the transaction.")]
+    pub(crate) consensus_version: Option<u8>,
 }
 
 /// Parses the record string. If the string is a ciphertext, then attempt to decrypt it. Lifted from snarkOS.
@@ -241,7 +243,7 @@ fn check_balance<N: Network>(
     private_key: &PrivateKey<N>,
     endpoint: &str,
     network: &str,
-    context: Context,
+    context: &Context,
     total_cost: u64,
 ) -> Result<()> {
     // Derive the account address.
@@ -251,7 +253,7 @@ fn check_balance<N: Network>(
         endpoint: Some(endpoint.to_string()),
         network: Some(network.to_string()),
         command: QueryCommands::Program {
-            command: crate::cli::commands::query::LeoProgram {
+            command: query::LeoProgram {
                 name: "credits".to_string(),
                 mappings: false,
                 mapping_value: Some(vec!["account".to_string(), address.to_string()]),
@@ -276,6 +278,30 @@ fn check_balance<N: Network>(
     }
 }
 
+// A helper function to query for the latest block height.
+fn get_latest_block_height(endpoint: &str, network: &str, context: &Context) -> Result<u32> {
+    // Query the latest block height.
+    let height = LeoQuery {
+        endpoint: Some(endpoint.to_string()),
+        network: Some(network.to_string()),
+        command: QueryCommands::Block {
+            command: query::LeoBlock {
+                id: None,
+                latest: false,
+                latest_hash: false,
+                latest_height: true,
+                range: None,
+                transactions: false,
+                to_height: false,
+            },
+        },
+    }
+    .execute(Context::new(context.path.clone(), context.home.clone(), true)?)?;
+    // Parse the height.
+    let height = height.parse::<u32>().map_err(CliError::string_parse_error)?;
+    Ok(height)
+}
+
 /// Determine if the transaction should be broadcast or displayed to user.
 fn handle_broadcast<N: Network>(endpoint: &String, transaction: Transaction<N>, operation: &String) -> Result<()> {
     println!("Broadcasting transaction to {}...\n", endpoint.clone());

leo/cli/commands/query/mod.rs

@@ -17,26 +17,26 @@
 use super::*;
 use snarkvm::prelude::{CanaryV0, MainnetV0, TestnetV0};
 
-mod block;
-use block::LeoBlock;
+pub mod block;
+pub use block::LeoBlock;
 
 pub mod program;
 pub use program::LeoProgram;
 
-mod state_root;
-use state_root::StateRoot;
+pub mod state_root;
+pub use state_root::StateRoot;
 
-mod committee;
-use committee::LeoCommittee;
+pub mod committee;
+pub use committee::LeoCommittee;
 
-mod mempool;
-use mempool::LeoMempool;
+pub mod mempool;
+pub use mempool::LeoMempool;
 
-mod peers;
-use peers::LeoPeers;
+pub mod peers;
+pub use peers::LeoPeers;
 
-mod transaction;
-use transaction::LeoTransaction;
+pub mod transaction;
+pub use transaction::LeoTransaction;
 
 mod utils;
 use utils::*;