Manage nonce internally

.github/workflows/test.yml

@@ -14,7 +14,7 @@ jobs:
   rust:
     runs-on: ubuntu-latest
     container:
-      image: ghcr.io/dojoengine/dojo-dev:v1.0.0-alpha.4
+      image: ghcr.io/dojoengine/dojo-dev:v1.0.0-alpha.11
     steps:
       - uses: actions/checkout@v4
       - run: git config --system --add safe.directory '*'

packages/account_sdk/src/controller.rs

@@ -47,6 +47,7 @@ pub struct Controller {
     contract: Option<Box<abigen::controller::Controller<Self>>>,
     pub factory: ControllerFactory,
     pub(crate) storage: Storage,
+    nonce: Felt,
 }
 
 impl Controller {
@@ -85,6 +86,7 @@ impl Controller {
             contract: None,
             factory,
             storage: Storage::default(),
+            nonce: Felt::ZERO,
         };
 
         let contract = Box::new(abigen::controller::Controller::new(
@@ -184,43 +186,68 @@ impl Controller {
             return self.execute_from_outside(calls).await;
         }
 
-        let result = self.execute_v1(calls.clone()).max_fee(max_fee).send().await;
-
-        match result {
-            Ok(tx_result) => {
-                // Update is_registered to true after successful execution with a session
-                if let Some((key, metadata)) =
-                    self.session_metadata(&Policy::from_calls(&calls), None)
-                {
-                    if !metadata.is_registered {
-                        let mut updated_metadata = metadata;
-                        updated_metadata.is_registered = true;
-                        self.storage.set_session(&key, updated_metadata)?;
+        let mut retry_count = 0;
+        let max_retries = 1;
+
+        loop {
+            let nonce = self.get_nonce().await?;
+            let result = self
+                .execute_v1(calls.clone())
+                .nonce(nonce)
+                .max_fee(max_fee)
+                .send()
+                .await;
+
+            match result {
+                Ok(tx_result) => {
+                    // Update nonce
+                    self.nonce = self.nonce + 1;
+
+                    // Update is_registered to true after successful execution with a session
+                    if let Some((key, metadata)) =
+                        self.session_metadata(&Policy::from_calls(&calls), None)
+                    {
+                        if !metadata.is_registered {
+                            let mut updated_metadata = metadata;
+                            updated_metadata.is_registered = true;
+                            self.storage.set_session(&key, updated_metadata)?;
+                        }
                     }
+                    return Ok(tx_result);
                 }
-                Ok(tx_result)
-            }
-            Err(e) => {
-                if let AccountError::Provider(ProviderError::StarknetError(
-                    StarknetError::TransactionExecutionError(data),
-                )) = &e
-                {
-                    if data
-                        .execution_error
-                        .contains(&format!("{:x} is not deployed.", self.address))
-                    {
-                        let balance = self.eth_balance().await?;
-                        let mut fee_estimate = self.deploy().estimate_fee().await?;
-                        fee_estimate.overall_fee += WEBAUTHN_GAS * fee_estimate.gas_price;
-                        Err(ControllerError::NotDeployed {
-                            fee_estimate: Box::new(fee_estimate),
-                            balance,
-                        })
-                    } else {
-                        Err(ControllerError::AccountError(e))
+                Err(e) => {
+                    match &e {
+                        AccountError::Provider(ProviderError::StarknetError(
+                            StarknetError::TransactionExecutionError(data),
+                        )) if data
+                            .execution_error
+                            .contains(&format!("{:x} is not deployed.", self.address)) =>
+                        {
+                            let balance = self.eth_balance().await?;
+                            let mut fee_estimate = self.deploy().estimate_fee().await?;
+                            fee_estimate.overall_fee += WEBAUTHN_GAS * fee_estimate.gas_price;
+                            return Err(ControllerError::NotDeployed {
+                                fee_estimate: Box::new(fee_estimate),
+                                balance,
+                            });
+                        }
+                        AccountError::Provider(ProviderError::StarknetError(
+                            StarknetError::InvalidTransactionNonce,
+                        )) => {
+                            if retry_count < max_retries {
+                                // Refetch nonce from the provider
+                                let new_nonce = self
+                                    .provider
+                                    .get_nonce(self.block_id(), self.address())
+                                    .await?;
+                                self.nonce = new_nonce;
+                                retry_count += 1;
+                                continue;
+                            }
+                        }
+                        _ => {}
                     }
-                } else {
-                    Err(ControllerError::AccountError(e))
+                    return Err(ControllerError::AccountError(e));
                 }
             }
         }
@@ -265,6 +292,26 @@ impl Controller {
         let signature = self.owner.sign(&hash).await?;
         Ok(Vec::<SignerSignature>::cairo_serialize(&vec![signature]))
     }
+
+    async fn get_nonce(&self) -> Result<Felt, ProviderError> {
+        let current_nonce = self.nonce;
+
+        if current_nonce == Felt::ZERO {
+            match self
+                .provider
+                .get_nonce(self.block_id(), self.address())
+                .await
+            {
+                Ok(nonce) => Ok(nonce),
+                Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => {
+                    Ok(Felt::ZERO)
+                }
+                Err(e) => Err(e),
+            }
+        } else {
+            Ok(current_nonce)
+        }
+    }
 }
 
 impl_account!(Controller, |account: &Controller, context| {
@@ -289,9 +336,7 @@ impl ConnectedAccount for Controller {
     }
 
     async fn get_nonce(&self) -> Result<Felt, ProviderError> {
-        self.provider
-            .get_nonce(self.block_id(), self.address())
-            .await
+        self.get_nonce().await
     }
 }
 

packages/account_sdk/src/controller_test.rs

@@ -1,11 +1,16 @@
 use std::time::Duration;
 
 use crate::{
+    abigen::erc_20::Erc20,
     artifacts::{Version, CONTROLLERS},
     controller::Controller,
     signers::Signer,
-    tests::{runners::katana::KatanaRunner, transaction_waiter::TransactionWaiter},
+    tests::{
+        account::FEE_TOKEN_ADDRESS, runners::katana::KatanaRunner,
+        transaction_waiter::TransactionWaiter,
+    },
 };
+use cainome::cairo_serde::{ContractAddress, U256};
 use starknet::{accounts::Account, macros::felt, providers::Provider, signers::SigningKey};
 use starknet_crypto::Felt;
 
@@ -71,3 +76,83 @@ async fn test_deploy_controller() {
         "Deployed address doesn't match expected address"
     );
 }
+
+#[tokio::test]
+async fn test_controller_nonce_mismatch_recovery() {
+    let username = "testuser".to_string();
+    let signer = Signer::new_starknet_random();
+
+    let runner = KatanaRunner::load();
+    let mut controller1 = runner
+        .deploy_controller(username.clone(), signer.clone(), Version::LATEST)
+        .await;
+
+    let chain_id = runner.client().chain_id().await.unwrap();
+
+    // Create the second controller with the same credentials and address
+    let mut controller2 = Controller::new(
+        "app_id".to_string(),
+        username.clone(),
+        controller1.class_hash,
+        runner.rpc_url.clone(),
+        signer.clone(),
+        controller1.address,
+        chain_id,
+    );
+
+    // Send a transaction using controller1
+    let recipient = ContractAddress(felt!("0x18301129"));
+    let amount = U256 { low: 0, high: 0 };
+    let erc20 = Erc20::new(*FEE_TOKEN_ADDRESS, &controller1);
+
+    let tx1 = erc20.transfer_getcall(&recipient, &amount);
+    let max_fee = controller1
+        .estimate_invoke_fee(vec![tx1.clone()])
+        .await
+        .unwrap();
+    let res = Controller::execute(&mut controller1, vec![tx1.clone()], max_fee.overall_fee)
+        .await
+        .unwrap();
+
+    TransactionWaiter::new(res.transaction_hash, runner.client())
+        .wait()
+        .await
+        .unwrap();
+
+    // Now send a transaction using controller2, which should have a stale nonce
+    let erc20 = Erc20::new(*FEE_TOKEN_ADDRESS, &controller2);
+    let tx2 = erc20.transfer_getcall(&recipient, &amount);
+
+    let tx2_result =
+        Controller::execute(&mut controller2, vec![tx2.clone()], max_fee.overall_fee).await;
+
+    // Verify that it succeeds after recovering from nonce mismatch
+    assert!(
+        tx2_result.is_ok(),
+        "Controller did not recover from nonce mismatch: {:?}",
+        tx2_result.err()
+    );
+
+    TransactionWaiter::new(tx2_result.unwrap().transaction_hash, runner.client())
+        .wait()
+        .await
+        .unwrap();
+
+    let res = Controller::execute(&mut controller1, vec![tx1], max_fee.overall_fee)
+        .await
+        .unwrap();
+
+    TransactionWaiter::new(res.transaction_hash, runner.client())
+        .wait()
+        .await
+        .unwrap();
+
+    let tx2_result = Controller::execute(&mut controller2, vec![tx2], max_fee.overall_fee).await;
+
+    // Verify that it succeeds after recovering from nonce mismatch
+    assert!(
+        tx2_result.is_ok(),
+        "Controller did not recover from nonce mismatch: {:?}",
+        tx2_result.err()
+    );
+}

packages/account_sdk/src/session.rs

@@ -90,12 +90,10 @@ impl Controller {
             }),
         )?;
 
-        let txn = self
+        let call = self
             .contract()
-            .register_session(&session.raw(), &self.owner_guid())
-            .max_fee(max_fee)
-            .send()
-            .await?;
+            .register_session_getcall(&session.raw(), &self.owner_guid());
+        let txn = self.execute(vec![call], max_fee).await?;
 
         self.storage.set_session(
             &Selectors::session(&self.address, &self.app_id, &self.chain_id),