feat(batcher): get last max fee msg (#2127)

Co-authored-by: JuArce <52429267+JuArce@users.noreply.github.com>

Author: MarcosNicolau

crates/batcher/src/lib.rs

@@ -33,9 +33,9 @@ use aligned_sdk::common::constants::{
     RESPOND_TO_TASK_FEE_LIMIT_PERCENTAGE_MULTIPLIER,
 };
 use aligned_sdk::common::types::{
-    ClientMessage, GetNonceResponseMessage, NoncedVerificationData, ProofInvalidReason,
-    ProvingSystemId, SubmitProofMessage, SubmitProofResponseMessage, VerificationCommitmentBatch,
-    VerificationData, VerificationDataCommitment,
+    ClientMessage, GetLastMaxFeeResponseMessage, GetNonceResponseMessage, NoncedVerificationData,
+    ProofInvalidReason, ProvingSystemId, SubmitProofMessage, SubmitProofResponseMessage,
+    VerificationCommitmentBatch, VerificationData, VerificationDataCommitment,
 };
 
 use aws_sdk_s3::client::Client as S3Client;
@@ -916,6 +916,11 @@ impl Batcher {
                     .handle_submit_proof_msg(msg, ws_conn_sink)
                     .await
             }
+            ClientMessage::GetLastMaxFee(address) => {
+                self.clone()
+                    .handle_get_last_max_fee(address, ws_conn_sink)
+                    .await
+            }
         }
     }
 
@@ -1004,6 +1009,84 @@ impl Batcher {
         Ok(())
     }
 
+    async fn handle_get_last_max_fee(
+        self: Arc<Self>,
+        mut address: Address,
+        ws_conn_sink: WsMessageSink,
+    ) -> Result<(), Error> {
+        // If the address is not paying, we will return the last max fee of the aligned_payment_address
+        if !self.has_to_pay(&address) {
+            info!("Handling nonpaying message");
+            let Some(non_paying_config) = self.non_paying_config.as_ref() else {
+                warn!(
+                    "There isn't a non-paying configuration loaded. This message will be ignored"
+                );
+                send_message(
+                    ws_conn_sink.clone(),
+                    GetLastMaxFeeResponseMessage::InvalidRequest(
+                        "There isn't a non-paying configuration loaded.".to_string(),
+                    ),
+                )
+                .await;
+                return Ok(());
+            };
+            let replacement_addr = non_paying_config.replacement.address();
+            address = replacement_addr;
+        }
+
+        let user_states_guard = match timeout(MESSAGE_HANDLER_LOCK_TIMEOUT, self.user_states.read())
+            .await
+        {
+            Ok(guard) => guard,
+            Err(_) => {
+                warn!("User states read lock acquisition timed out in handle_get_last_max_fee_for_address_msg");
+                self.metrics.inc_message_handler_user_states_lock_timeouts();
+                send_message(ws_conn_sink, GetLastMaxFeeResponseMessage::ServerBusy).await;
+                return Ok(());
+            }
+        };
+
+        let Some(usr_ref) = user_states_guard.get(&address).cloned() else {
+            drop(user_states_guard);
+            send_message(
+                ws_conn_sink.clone(),
+                GetLastMaxFeeResponseMessage::LastMaxFee(U256::MAX),
+            )
+            .await;
+            return Ok(());
+        };
+
+        let Some(usr_lock) = self
+            .try_user_lock_with_timeout(address, usr_ref.lock())
+            .await
+        else {
+            drop(user_states_guard);
+            send_message(
+                ws_conn_sink.clone(),
+                GetLastMaxFeeResponseMessage::ServerBusy,
+            )
+            .await;
+            return Ok(());
+        };
+
+        let proofs_in_queue = usr_lock.proofs_in_batch;
+        let max_fee = if proofs_in_queue > 0 {
+            usr_lock.last_max_fee_limit
+        } else {
+            U256::MAX
+        };
+        drop(usr_lock);
+        drop(user_states_guard);
+
+        send_message(
+            ws_conn_sink.clone(),
+            GetLastMaxFeeResponseMessage::LastMaxFee(max_fee),
+        )
+        .await;
+
+        Ok(())
+    }
+
     /// Returns the Aligned-funded address that will be used to pay for proofs when users don't need to pay themselves.
     /// This function assumes that the non-paying configuration is set.
     fn aligned_payment_address(&self) -> Address {

crates/cli/src/main.rs

@@ -20,6 +20,7 @@ use aligned_sdk::communication::serialization::cbor_deserialize;
 use aligned_sdk::verification_layer;
 use aligned_sdk::verification_layer::estimate_fee;
 use aligned_sdk::verification_layer::get_chain_id;
+use aligned_sdk::verification_layer::get_last_max_fee;
 use aligned_sdk::verification_layer::get_nonce_from_batcher;
 use aligned_sdk::verification_layer::get_nonce_from_ethereum;
 use aligned_sdk::verification_layer::get_unlock_block_time;
@@ -46,6 +47,7 @@ use transaction::eip2718::TypedTransaction;
 use crate::AlignedCommands::DepositToBatcher;
 use crate::AlignedCommands::GetUserAmountOfQueuedProofs;
 use crate::AlignedCommands::GetUserBalance;
+use crate::AlignedCommands::GetUserLastMaxFee;
 use crate::AlignedCommands::GetUserNonce;
 use crate::AlignedCommands::GetUserNonceFromEthereum;
 use crate::AlignedCommands::GetVkCommitment;
@@ -94,6 +96,11 @@ pub enum AlignedCommands {
         name = "get-user-nonce-from-ethereum"
     )]
     GetUserNonceFromEthereum(GetUserNonceFromEthereumArgs),
+    #[clap(
+        about = "Gets user current last max fee from the batcher. This is the max limit you can send in your next proof.",
+        name = "get-user-last-max-fee"
+    )]
+    GetUserLastMaxFee(GetUserLastMaxFeeArgs),
     #[clap(
         about = "Gets the number of proofs a user has queued in the Batcher.",
         name = "get-user-amount-of-queued-proofs"
@@ -316,6 +323,15 @@ pub struct GetUserNonceArgs {
     address: String,
 }
 
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+pub struct GetUserLastMaxFeeArgs {
+    #[clap(flatten)]
+    network: NetworkArg,
+    #[arg(name = "The user's Ethereum address", required = true)]
+    address: String,
+}
+
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 pub struct GetUserNonceFromEthereumArgs {
@@ -1041,6 +1057,22 @@ async fn main() -> Result<(), AlignedError> {
                 }
             }
         }
+        GetUserLastMaxFee(args) => {
+            let address = H160::from_str(&args.address).unwrap();
+            match get_last_max_fee(args.network.into(), address).await {
+                Ok(last_max_fee) => {
+                    let last_max_fee_ether = ethers::utils::format_ether(last_max_fee);
+                    info!(
+                        "Last max fee for address {} is {}eth",
+                        address, last_max_fee_ether
+                    );
+                }
+                Err(e) => {
+                    error!("Error while getting last max fee: {:?}", e);
+                    return Ok(());
+                }
+            }
+        }
         GetUserAmountOfQueuedProofs(args) => {
             let address = H160::from_str(&args.address).unwrap();
             let network: Network = args.network.into();

crates/sdk/src/common/errors.rs

@@ -266,6 +266,16 @@ pub enum GetNonceError {
     GenericError(String),
 }
 
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub enum GetLastMaxFeeError {
+    ConnectionFailed(String),
+    SerializationError(String),
+    UnexpectedResponse(String),
+    InvalidRequest(String),
+    ProtocolMismatch { current: u16, expected: u16 },
+    GenericError(String),
+}
+
 #[derive(Debug)]
 pub enum ChainIdError {
     EthereumProviderError(String),

crates/sdk/src/common/types.rs

@@ -295,6 +295,7 @@ pub enum ClientMessage {
     // Needs to be wrapped in box as the message is 3x bigger than the others
     // see https://rust-lang.github.io/rust-clippy/master/index.html#large_enum_variant
     SubmitProof(Box<SubmitProofMessage>),
+    GetLastMaxFee(Address),
 }
 
 impl Display for ClientMessage {
@@ -303,6 +304,7 @@ impl Display for ClientMessage {
         match self {
             ClientMessage::GetNonceForAddress(_) => write!(f, "GetNonceForAddress"),
             ClientMessage::SubmitProof(_) => write!(f, "SubmitProof"),
+            ClientMessage::GetLastMaxFee(_) => write!(f, "GetLastMaxFee"),
         }
     }
 }
@@ -465,6 +467,13 @@ pub enum GetNonceResponseMessage {
     ServerBusy,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum GetLastMaxFeeResponseMessage {
+    LastMaxFee(U256),
+    InvalidRequest(String),
+    ServerBusy,
+}
+
 #[derive(Debug, Clone)]
 pub enum Network {
     Devnet,

crates/sdk/src/verification_layer/mod.rs

@@ -5,10 +5,11 @@ use crate::{
             DEFAULT_MAX_FEE_BATCH_SIZE, GAS_PRICE_PERCENTAGE_MULTIPLIER,
             INSTANT_MAX_FEE_BATCH_SIZE, PERCENTAGE_DIVIDER,
         },
-        errors::{self, GetNonceError, PaymentError},
+        errors::{self, GetLastMaxFeeError, GetNonceError, PaymentError},
         types::{
-            AlignedVerificationData, ClientMessage, FeeEstimationType, GetNonceResponseMessage,
-            Network, ProvingSystemId, VerificationData,
+            AlignedVerificationData, ClientMessage, FeeEstimationType,
+            GetLastMaxFeeResponseMessage, GetNonceResponseMessage, Network, ProvingSystemId,
+            VerificationData,
         },
     },
     communication::{
@@ -650,6 +651,87 @@ pub async fn get_nonce_from_ethereum(
     }
 }
 
+/// Retrieves the `max_fee` of the proof with the highest nonce in the batcher queue for a given address.
+///
+/// This value represents the maximum fee limit that can be used when submitting the next proof.  
+/// To increase the fee limit for a new proof, you must first bump the fee of the previous proofs,
+/// and continue doing so recursively until you reach the proof with the highest nonce (this one).
+///
+/// Read more here: https://docs.alignedlayer.com/architecture/1_proof_verification_layer/1_batcher#max-fee-priority-queue
+///
+/// # Arguments
+/// * `network` - The network from which to retrieve the last `max_fee`.
+/// * `address` - The user address whose last `max_fee` will be retrieved.
+///
+/// # Returns
+/// * `Ok(U256)` - The `max_fee` of the proof with the highest nonce for the given user.  
+/// * `Ok(U256::MAX)` - If the user has no proofs in the queue.  
+pub async fn get_last_max_fee(
+    network: Network,
+    address: Address,
+) -> Result<U256, GetLastMaxFeeError> {
+    let (ws_stream, _) = connect_async(network.get_batcher_url())
+        .await
+        .map_err(|_| {
+            GetLastMaxFeeError::ConnectionFailed("Ws connection to batcher failed".to_string())
+        })?;
+
+    debug!("WebSocket handshake has been successfully completed");
+    let (mut ws_write, mut ws_read) = ws_stream.split();
+    check_protocol_version(&mut ws_read)
+        .map_err(|e| match e {
+            errors::SubmitError::ProtocolVersionMismatch { current, expected } => {
+                GetLastMaxFeeError::ProtocolMismatch { current, expected }
+            }
+            _ => GetLastMaxFeeError::UnexpectedResponse(
+                "Unexpected response, expected protocol version".to_string(),
+            ),
+        })
+        .await?;
+
+    let msg = ClientMessage::GetLastMaxFee(address);
+
+    let msg_bin = cbor_serialize(&msg).map_err(|_| {
+        GetLastMaxFeeError::SerializationError("Failed to serialize msg".to_string())
+    })?;
+    ws_write
+        .send(Message::Binary(msg_bin.clone()))
+        .await
+        .map_err(|_| {
+            GetLastMaxFeeError::ConnectionFailed(
+                "Ws connection failed to send message to batcher".to_string(),
+            )
+        })?;
+
+    let mut response_stream: ResponseStream =
+        ws_read.try_filter(|msg| futures_util::future::ready(msg.is_binary()));
+
+    let msg = match response_stream.next().await {
+        Some(Ok(msg)) => msg,
+        _ => {
+            return Err(GetLastMaxFeeError::ConnectionFailed(
+                "Connection was closed without close message before receiving all messages"
+                    .to_string(),
+            ));
+        }
+    };
+
+    let _ = ws_write.close().await;
+
+    match cbor_deserialize(msg.into_data().as_slice()) {
+        Ok(GetLastMaxFeeResponseMessage::LastMaxFee(last_max_fee)) => Ok(last_max_fee),
+        Ok(GetLastMaxFeeResponseMessage::InvalidRequest(e)) => {
+            Err(GetLastMaxFeeError::InvalidRequest(e))
+        }
+        Ok(GetLastMaxFeeResponseMessage::ServerBusy) => Err(GetLastMaxFeeError::GenericError(
+            "Server is busy processing requests, please retry".to_string(),
+        )),
+        Err(_) => Err(GetLastMaxFeeError::SerializationError(
+            "Failed to deserialize batcher message".to_string(),
+        )),
+    }
+}
+
 /// Returns the chain ID of the Ethereum network.
 ///
 /// # Arguments

docs/3_guides/1.2_SDK_api_reference.md

@@ -274,6 +274,43 @@ pub async fn get_nonce_from_batcher(
 
 - `EthRpcError` if the batcher has an error in the Ethereum call when retrieving the nonce if not already cached.
 
+### `get_last_max_fee`
+
+Retrieves the `max_fee` of the proof with the highest nonce in the batcher queue for a given address.
+
+This value represents the maximum fee limit that can be used when submitting the next proof. To increase the fee limit for a new proof, you must first bump the fee of the previous proofs queued in the batcher.
+
+Read more here: https://docs.alignedlayer.com/architecture/1_proof_verification_layer/1_batcher#max-fee-priority-queue
+
+```rust
+pub async fn get_last_max_fee(
+    network: Network,
+    address: Address,
+) -> Result<U256, GetLastMaxFeeError>
+```
+
+#### Arguments
+
+- `network` - The network from which to retrieve the last `max_fee`.
+- `address` - The user address whose last `max_fee` will be retrieved.
+
+#### Returns
+
+- `Result<U256, GetLastMaxFeeError>` - The `max_fee` of the proof with the highest nonce for the given user, or `U256::MAX` if the user has no proofs in the queue.
+
+#### Errors
+
+- `ConnectionFailed` if there is an error connecting to the batcher.
+- `ProtocolMismatch` if the protocol version doesn't match.
+- `SerializationError` if there is an error serializing/deserializing the message.
+- `InvalidRequest` if the request is invalid.
+- `UnexpectedResponse` if the batcher responds with an unexpected message.
+- `GenericError` if the error doesn't match any of the previous ones.
+
+#### Notes
+
+- Returns `U256::MAX` (2^256 - 1) when no proofs are present in the queue for the user.
+
 ### `get_chain_id`
 
 Returns the chain ID for a given rpc url.