upgrade wallet version

bedrock/src/migration/controller.rs

@@ -1,6 +1,8 @@
 use crate::migration::error::MigrationError;
 use crate::migration::processor::{MigrationProcessor, ProcessorResult};
+use crate::migration::processors::enable_4337_module_processor::Enable4337ModuleProcessor;
 use crate::migration::processors::permit2_approval_processor::Permit2ApprovalProcessor;
+use crate::migration::processors::safe_upgrade_processor::SafeUpgradeProcessor;
 use crate::migration::state::{MigrationRecord, MigrationStatus};
 use crate::primitives::key_value_store::{DeviceKeyValueStore, KeyValueStoreError};
 use crate::smart_account::SafeSmartAccount;
@@ -69,6 +71,8 @@ impl MigrationController {
     /// Create a new [`MigrationController`] with default processors and optional additional ones.
     ///
     /// Default processors (loaded automatically):
+    /// - [`SafeUpgradeProcessor`]: Upgrades Safe wallets from v1.3.0 to v1.4.1
+    /// - [`Enable4337ModuleProcessor`]: Enables the Safe4337Module if not already enabled
     /// - [`Permit2ApprovalProcessor`]: Ensures max ERC20 approval to Permit2 on `WorldChain`
     ///
     /// Additional processors passed via `additional_processors` are appended after the defaults.
@@ -164,7 +168,11 @@ impl MigrationController {
     fn default_processors(
         safe_account: Arc<SafeSmartAccount>,
     ) -> Vec<Arc<dyn MigrationProcessor>> {
-        vec![Arc::new(Permit2ApprovalProcessor::new(safe_account))]
+        vec![
+            Arc::new(SafeUpgradeProcessor::new(Arc::clone(&safe_account))),
+            Arc::new(Enable4337ModuleProcessor::new(Arc::clone(&safe_account))),
+            Arc::new(Permit2ApprovalProcessor::new(safe_account)),
+        ]
     }
 
     /// Create a controller with processors injected in

bedrock/src/migration/mod.rs

@@ -76,6 +76,9 @@ mod state;
 /// Example processors showing how to implement migrations
 pub mod processors;
 
+/// Shared utilities for migration processors
+pub mod utils;
+
 // Public API exports
 pub use controller::{MigrationController, MigrationRunSummary};
 pub use error::MigrationError;

bedrock/src/migration/processors/enable_4337_module_processor.rs

@@ -0,0 +1,89 @@
+use async_trait::async_trait;
+use log::info;
+use std::sync::Arc;
+
+use crate::migration::error::MigrationError;
+use crate::migration::processor::{MigrationProcessor, ProcessorResult};
+use crate::migration::utils::poll_for_receipt;
+use crate::primitives::Network;
+use crate::smart_account::{Is4337Encodable, SafeSmartAccount};
+use crate::transactions::contracts::gnosis_safe::{
+    GnosisSafe, SafeEnableModule, SAFE_4337_MODULE,
+};
+use crate::transactions::rpc::{get_rpc_client, RpcProviderName};
+
+/// Migration processor that checks if the Safe4337Module is enabled on the wallet
+/// and enables it if not.
+///
+/// The 4337 module is required for the wallet to process ERC-4337 UserOperations.
+/// Some wallets may have been deployed without it or had it removed. This processor
+/// ensures the module is enabled by checking on-chain state and calling `enableModule`
+/// if needed.
+pub struct Enable4337ModuleProcessor {
+    safe_account: Arc<SafeSmartAccount>,
+}
+
+impl Enable4337ModuleProcessor {
+    /// Creates a new `Enable4337ModuleProcessor` with the given Safe smart account.
+    #[must_use]
+    pub fn new(safe_account: Arc<SafeSmartAccount>) -> Self {
+        Self { safe_account }
+    }
+}
+
+#[async_trait]
+impl MigrationProcessor for Enable4337ModuleProcessor {
+    fn migration_id(&self) -> String {
+        "wallet.safe.enable_4337_module.v1".to_string()
+    }
+
+    async fn is_applicable(&self) -> Result<bool, MigrationError> {
+        let rpc_client = get_rpc_client()
+            .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
+
+        let safe = GnosisSafe::new(self.safe_account.wallet_address);
+        let is_enabled = safe
+            .is_module_enabled(&rpc_client, SAFE_4337_MODULE)
+            .await
+            .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
+
+        if is_enabled {
+            info!("Safe4337Module is already enabled");
+        } else {
+            info!("Safe4337Module is NOT enabled, migration needed");
+        }
+
+        Ok(!is_enabled)
+    }
+
+    async fn execute(&self) -> Result<ProcessorResult, MigrationError> {
+        let enable_module =
+            SafeEnableModule::new(self.safe_account.wallet_address, SAFE_4337_MODULE);
+
+        let user_op_hash = match enable_module
+            .sign_and_execute(
+                &self.safe_account,
+                Network::WorldChain,
+                None,
+                None,
+                RpcProviderName::Any,
+            )
+            .await
+        {
+            Ok(hash) => {
+                info!("Submitted enableModule for 4337 module, userOpHash: {hash:?}");
+                hash
+            }
+            Err(e) => {
+                return Ok(ProcessorResult::Retryable {
+                    error_code: "RPC_ERROR".to_string(),
+                    error_message: format!(
+                        "Failed to submit enableModule transaction: {e}"
+                    ),
+                });
+            }
+        };
+
+        poll_for_receipt(user_op_hash, "enableModule").await
+    }
+}

bedrock/src/migration/processors/mod.rs

@@ -4,5 +4,11 @@
 /// that can be used as templates for actual migrations.
 mod example_processor;
 
+/// Processor that checks if the Safe4337Module is enabled and enables it if not.
+pub mod enable_4337_module_processor;
+
 /// Processor that ensures max ERC20 approval to Permit2 on `WorldChain` for supported tokens.
 pub mod permit2_approval_processor;
+
+/// Processor that upgrades Safe wallets from v1.3.0 to v1.4.1.
+pub mod safe_upgrade_processor;

bedrock/src/migration/processors/permit2_approval_processor.rs

@@ -6,6 +6,7 @@ use tokio::sync::Mutex;
 
 use crate::migration::error::MigrationError;
 use crate::migration::processor::{MigrationProcessor, ProcessorResult};
+use crate::migration::utils::poll_for_receipt;
 use crate::primitives::Network;
 use crate::smart_account::{Is4337Encodable, SafeSmartAccount};
 use crate::transactions::contracts::erc20::Erc20;
@@ -151,54 +152,6 @@ impl MigrationProcessor for Permit2ApprovalProcessor {
             }
         };
 
-        // Wait for the user operation to be mined before marking as success.
-        let rpc_client = get_rpc_client()
-            .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
-
-        let user_op_hash_hex = format!("{user_op_hash:#x}");
-        let delay_ms = 4000u64;
-
-        for attempt in 0..5 {
-            let response = rpc_client
-                .wa_get_user_operation_receipt(Network::WorldChain, &user_op_hash_hex)
-                .await
-                .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
-
-            match response.status.as_str() {
-                "mined_success" => {
-                    info!(
-                        "Permit2 approvals mined successfully for {names:?}, txHash: {:?}",
-                        response.transaction_hash
-                    );
-                    return Ok(ProcessorResult::Success);
-                }
-                "mined_revert" | "error" => {
-                    return Ok(ProcessorResult::Retryable {
-                        error_code: "MINED_REVERT".to_string(),
-                        error_message: format!(
-                            "Permit2 approval transaction failed for {names:?}, txHash: {:?}",
-                            response.transaction_hash
-                        ),
-                        });
-                }
-                _ => {
-                    // Still pending — keep polling unless this is the last attempt
-                    if attempt < 4 {
-                        tokio::time::sleep(tokio::time::Duration::from_millis(
-                            delay_ms,
-                        ))
-                        .await;
-                    }
-                }
-            }
-        }
-
-        // Still pending after all polling attempts — retry the whole migration later
-        Ok(ProcessorResult::Retryable {
-            error_code: "PENDING_TIMEOUT".to_string(),
-            error_message: format!(
-                "Permit2 approval for {names:?} still pending after polling, will retry"
-            ),
-        })
+        poll_for_receipt(user_op_hash, &format!("Permit2 approval for {names:?}")).await
     }
 }

bedrock/src/migration/processors/safe_upgrade_processor.rs

@@ -0,0 +1,86 @@
+use async_trait::async_trait;
+use log::info;
+use std::sync::Arc;
+
+use crate::migration::error::MigrationError;
+use crate::migration::processor::{MigrationProcessor, ProcessorResult};
+use crate::migration::utils::poll_for_receipt;
+use crate::primitives::Network;
+use crate::smart_account::{Is4337Encodable, SafeSmartAccount};
+use crate::transactions::contracts::gnosis_safe::{
+    GnosisSafe, SafeWalletVersionUpgrade, SAFE_VERSION_130,
+};
+use crate::transactions::rpc::{get_rpc_client, RpcProviderName};
+
+/// Migration processor that checks the Gnosis Safe wallet version and upgrades
+/// from v1.3.0 to v1.4.1 if needed.
+///
+/// Uses a `delegatecall` to the `WC_MIGRATION_WALLET_UPGRADE` contract which
+/// handles updating the Safe proxy's singleton address.
+pub struct SafeUpgradeProcessor {
+    safe_account: Arc<SafeSmartAccount>,
+}
+
+impl SafeUpgradeProcessor {
+    /// Creates a new `SafeUpgradeProcessor` with the given Safe smart account.
+    #[must_use]
+    pub fn new(safe_account: Arc<SafeSmartAccount>) -> Self {
+        Self { safe_account }
+    }
+}
+
+#[async_trait]
+impl MigrationProcessor for SafeUpgradeProcessor {
+    fn migration_id(&self) -> String {
+        "wallet.safe.upgrade.v1".to_string()
+    }
+
+    async fn is_applicable(&self) -> Result<bool, MigrationError> {
+        let rpc_client = get_rpc_client()
+            .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
+
+        let safe = GnosisSafe::new(self.safe_account.wallet_address);
+        let version = safe
+            .fetch_version(&rpc_client)
+            .await
+            .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
+
+        if version == SAFE_VERSION_130 {
+            info!("Safe is on v1.3.0, upgrade to v1.4.1 needed");
+            Ok(true)
+        } else {
+            info!("Safe is on v{version}, no upgrade needed");
+            Ok(false)
+        }
+    }
+
+    async fn execute(&self) -> Result<ProcessorResult, MigrationError> {
+        let upgrade = SafeWalletVersionUpgrade;
+
+        let user_op_hash = match upgrade
+            .sign_and_execute(
+                &self.safe_account,
+                Network::WorldChain,
+                None,
+                None,
+                RpcProviderName::Any,
+            )
+            .await
+        {
+            Ok(hash) => {
+                info!("Submitted Safe upgrade transaction, userOpHash: {hash:?}");
+                hash
+            }
+            Err(e) => {
+                return Ok(ProcessorResult::Retryable {
+                    error_code: "RPC_ERROR".to_string(),
+                    error_message: format!(
+                        "Failed to submit Safe upgrade transaction: {e}"
+                    ),
+                });
+            }
+        };
+
+        poll_for_receipt(user_op_hash, "Safe upgrade").await
+    }
+}

bedrock/src/migration/utils/mod.rs

@@ -0,0 +1,79 @@
+//! Shared utilities for migration processors.
+
+use alloy::primitives::FixedBytes;
+use log::info;
+
+use crate::migration::error::MigrationError;
+use crate::migration::processor::ProcessorResult;
+use crate::primitives::Network;
+use crate::transactions::rpc::get_rpc_client;
+
+/// Number of polling attempts before giving up.
+const POLL_ATTEMPTS: u32 = 5;
+
+/// Delay between polling attempts in milliseconds.
+const POLL_DELAY_MS: u64 = 4000;
+
+/// Polls `wa_getUserOperationReceipt` on WorldChain until the operation is mined
+/// or the maximum number of attempts is exhausted.
+///
+/// # Arguments
+/// * `user_op_hash` - The hash returned by `sign_and_execute`.
+/// * `label` - A human-readable label for log messages (e.g. `"enableModule"`).
+///
+/// # Returns
+/// - `ProcessorResult::Success` if the operation was mined successfully.
+/// - `ProcessorResult::Retryable` if the operation reverted, errored, or is still pending.
+///
+/// # Errors
+/// Returns a `MigrationError` if the RPC client cannot be obtained or an RPC call fails.
+pub async fn poll_for_receipt(
+    user_op_hash: FixedBytes<32>,
+    label: &str,
+) -> Result<ProcessorResult, MigrationError> {
+    let rpc_client = get_rpc_client()
+        .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
+
+    let user_op_hash_hex = format!("{user_op_hash:#x}");
+
+    for attempt in 0..POLL_ATTEMPTS {
+        let response = rpc_client
+            .wa_get_user_operation_receipt(Network::WorldChain, &user_op_hash_hex)
+            .await
+            .map_err(|e| MigrationError::InvalidOperation(e.to_string()))?;
+
+        match response.status.as_str() {
+            "mined_success" => {
+                info!(
+                    "{label} mined successfully, txHash: {:?}",
+                    response.transaction_hash
+                );
+                return Ok(ProcessorResult::Success);
+            }
+            "mined_revert" | "error" => {
+                return Ok(ProcessorResult::Retryable {
+                    error_code: "MINED_REVERT".to_string(),
+                    error_message: format!(
+                        "{label} transaction reverted, txHash: {:?}",
+                        response.transaction_hash
+                    ),
+                });
+            }
+            _ => {
+                if attempt < POLL_ATTEMPTS - 1 {
+                    tokio::time::sleep(tokio::time::Duration::from_millis(
+                        POLL_DELAY_MS,
+                    ))
+                    .await;
+                }
+            }
+        }
+    }
+
+    Ok(ProcessorResult::Retryable {
+        error_code: "PENDING_TIMEOUT".to_string(),
+        error_message: format!(
+            "{label} still pending after polling, will retry"
+        ),
+    })
+}

bedrock/src/smart_account/nonce.rs

@@ -48,6 +48,10 @@ pub enum TransactionTypeId {
     WLDVaultMigration = 139,
     /// USD Vault migration to ERC-4626 vault
     USDVaultMigration = 140,
+    /// Safe `enableModule for 4337` call
+    SafeEnable4337Module = 141,
+    /// Safe implementation upgrade (v1.3.0 → v1.4.1)
+    SafeWalletVersionUpgrade = 142,
 }
 
 impl TransactionTypeId {