iwa 0.0.32__py3-none-any.whl → 0.0.58__py3-none-any.whl

iwa/core/services/transfer/native.py

@@ -85,24 +85,14 @@ class NativeTransferMixin:
     ) -> Optional[str]:
         """Send native currency via EOA using unified TransactionService."""
         # Build transaction dict
-        try:
-            gas_price = chain_interface.web3.eth.gas_price
-            gas_estimate = chain_interface.web3.eth.estimate_gas({
+        tx = chain_interface.calculate_transaction_params(
+            built_method=None,
+            tx_params={
                 "from": from_account.address,
                 "to": to_address,
                 "value": amount_wei,
-            })
-        except Exception as e:
-            logger.error(f"Failed to estimate gas for native transfer: {e}")
-            return None
-
-        tx = {
-            "from": from_account.address,
-            "to": to_address,
-            "value": amount_wei,
-            "gas": gas_estimate,
-            "gasPrice": gas_price,
-        }
+            }
+        )
 
         # Use unified TransactionService
         success, receipt = self.transaction_service.sign_and_send(
@@ -189,17 +179,13 @@ class NativeTransferMixin:
         logger.info(f"Wrapping {amount_eth:.4f} xDAI → WXDAI...")
 
         try:
-            tx = contract.functions.deposit().build_transaction(
-                {
-                    "from": account.address,
-                    "value": amount_wei,
-                    "gas": 100000,
-                    "gasPrice": chain_interface.web3._web3.eth.gas_price,
-                    "nonce": chain_interface.web3._web3.eth.get_transaction_count(account.address),
-                }
+            tx_params = chain_interface.calculate_transaction_params(
+                built_method=contract.functions.deposit(),
+                tx_params={"from": account.address, "value": amount_wei},
             )
+            transaction = contract.functions.deposit().build_transaction(tx_params)
 
-            signed = self.key_storage.sign_transaction(tx, account.address)
+            signed = self.key_storage.sign_transaction(transaction, account.address)
             tx_hash = chain_interface.web3._web3.eth.send_raw_transaction(signed.raw_transaction)
             receipt = chain_interface.web3._web3.eth.wait_for_transaction_receipt(
                 tx_hash, timeout=60
@@ -270,14 +256,11 @@ class NativeTransferMixin:
         logger.info(f"Unwrapping {amount_eth:.4f} WXDAI → xDAI...")
 
         try:
-            tx = contract.functions.withdraw(amount_wei).build_transaction(
-                {
-                    "from": account.address,
-                    "gas": 100000,
-                    "gasPrice": chain_interface.web3._web3.eth.gas_price,
-                    "nonce": chain_interface.web3._web3.eth.get_transaction_count(account.address),
-                }
+            tx_params = chain_interface.calculate_transaction_params(
+                built_method=contract.functions.withdraw(amount_wei),
+                tx_params={"from": account.address},
             )
+            tx = contract.functions.withdraw(amount_wei).build_transaction(tx_params)
 
             signed = self.key_storage.sign_transaction(tx, account.address)
             tx_hash = chain_interface.web3._web3.eth.send_raw_transaction(signed.raw_transaction)

iwa/core/services/transfer/swap.py

@@ -104,6 +104,7 @@ class SwapMixin:
             sell_token_name=sell_token_name,
             buy_token_name=buy_token_name,
             order_type=order_type,
+            wait_for_execution=True,
         )
 
         if result:

iwa/core/tests/test_gnosis_fee.py

@@ -0,0 +1,87 @@
+"""Test Gnosis fee calculation fix."""
+
+import unittest
+from unittest.mock import MagicMock
+
+from iwa.core.chain.interface import ChainInterface
+
+
+class TestGnosisFeeFix(unittest.TestCase):
+    """Test fee calculation for Gnosis chain."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.chain_interface = ChainInterface("gnosis")
+        # Mock web3 to avoid real connection
+        self.chain_interface.web3 = MagicMock()
+        self.chain_interface.web3.eth = MagicMock()
+
+    def test_fee_too_low_fix(self):
+        """Test that maxPriorityFeePerGas is forced to at least 1 wei on Gnosis."""
+        # 1. Setup EIP-1559 environment (block has baseFeePerGas)
+        mock_block = {"baseFeePerGas": 5000}
+        self.chain_interface.web3.eth.get_block.return_value = mock_block
+
+        # 2. Simulate RPC returning 0 priority fee (cause of the error)
+        self.chain_interface.web3.eth.max_priority_fee = 0
+
+        # 3. Setup dummy function for gas estimation
+        mock_func = MagicMock()
+        mock_func.estimate_gas.return_value = 100_000
+
+        # 4. Call calculation
+        tx_params = {"from": "0x123", "value": 0}
+        params = self.chain_interface.calculate_transaction_params(mock_func, tx_params)
+
+        # 5. Verify the fix
+        # Should have EIP-1559 fields
+        self.assertIn("maxFeePerGas", params)
+        self.assertIn("maxPriorityFeePerGas", params)
+        self.assertNotIn("gasPrice", params)
+
+        # CRITICAL ASSERTION: maxPriorityFeePerGas must be >= 1
+        # If the fix works, it should be 1. If it fails (old behavior), it would be 0.
+        self.assertEqual(params["maxPriorityFeePerGas"], 1, "Priority fee should be forced to 1 wei")
+
+        # Verify max fee calculation: (base * 1.5) + priority
+        expected_max_fee = int(5000 * 1.5) + 1
+        self.assertEqual(params["maxFeePerGas"], expected_max_fee)
+
+    def test_legacy_fallback(self):
+        """Test fallback to legacy gasPrice if baseFeePerGas is missing."""
+        # Setup legacy block (no baseFeePerGas)
+        self.chain_interface.web3.eth.get_block.return_value = {}
+        self.chain_interface.web3.eth.gas_price = 2000000000
+
+        mock_func = MagicMock()
+        mock_func.estimate_gas.return_value = 100_000
+
+        tx_params = {"from": "0x123", "value": 0}
+        params = self.chain_interface.calculate_transaction_params(mock_func, tx_params)
+
+        self.assertIn("gasPrice", params)
+        self.assertNotIn("maxFeePerGas", params)
+        self.assertEqual(params["gasPrice"], 2000000000)
+
+    def test_other_chain_behavior(self):
+        """Test that other chains (e.g. Ethereum) don't necessarily upgrade 0 to 1 (unless generic rule applies)."""
+        # Our fix in interface.py applies the fallback logic:
+        # if max_priority_fee < 1: max_priority_fee = 1
+        # This is now generic in the cleaned up code (lines 449-450: if max_priority_fee < 1: max_priority_fee = 1)
+        # So it should apply to ALL chains that support EIP-1559.
+
+        # We'll use Ethereum to verify generic behavior
+        eth_interface = ChainInterface("ethereum")
+        eth_interface.web3 = MagicMock()
+        eth_interface.web3.eth = MagicMock()
+
+        mock_block = {"baseFeePerGas": 100_000}
+        eth_interface.web3.eth.get_block.return_value = mock_block
+        eth_interface.web3.eth.max_priority_fee = 0
+
+        mock_func = MagicMock()
+        mock_func.estimate_gas.return_value = 21000
+
+        params = eth_interface.calculate_transaction_params(mock_func, {"from": "0x123"})
+
+        self.assertEqual(params["maxPriorityFeePerGas"], 1, "Generic fallback should apply to all chains")

iwa/core/tests/test_ipfs.py

@@ -0,0 +1,85 @@
+"""Tests for IPFS module."""
+
+from unittest.mock import ANY, MagicMock, patch
+
+import pytest
+
+import iwa.core.ipfs as ipfs_module
+from iwa.core.ipfs import push_to_ipfs_sync
+
+
+@pytest.fixture
+def mock_config():
+    """Mock config."""
+    with patch("iwa.core.ipfs.Config") as mock_c:
+        mock_c.return_value.core.ipfs_api_url = "http://fake-ipfs:5001"
+        yield mock_c
+
+
+@pytest.fixture
+def mock_cid_decode():
+    """Mock CID.decode."""
+    with patch("iwa.core.ipfs.CID") as mock_cid:
+        # returns an object that has version, codec, hashfun.name, raw_digest attributes
+        mock_decoded = MagicMock()
+        mock_decoded.version = 1
+        mock_decoded.codec = "raw"
+        mock_decoded.hashfun.name = "sha2-256"
+        mock_decoded.raw_digest = b"digest"
+
+        mock_cid.decode.return_value = mock_decoded
+        yield mock_cid
+
+
+def test_push_to_ipfs_sync_uses_session(mock_config, mock_cid_decode):
+    """Test push_to_ipfs_sync uses persistent session."""
+    # Reset global session for test
+    ipfs_module._SYNC_SESSION = None
+
+    # Mock requests.Session
+    with patch("requests.Session") as mock_session_cls:
+        mock_session = MagicMock()
+        mock_session_cls.return_value = mock_session
+
+        # Mock response
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"Hash": "QmTestHash"}
+        mock_session.post.return_value = mock_response
+
+        # Call function
+        cid_str, cid_hex = push_to_ipfs_sync(b"test data")
+
+        # Verify Session was created
+        mock_session_cls.assert_called_once()
+
+        # Verify post called on session
+        mock_session.post.assert_called_once()
+
+        # Verify session is stored globally
+        assert ipfs_module._SYNC_SESSION == mock_session
+
+        # Verify second call reuses session
+        push_to_ipfs_sync(b"test data 2")
+        mock_session_cls.assert_called_once() # Should still be 1 call
+        assert mock_session.post.call_count == 2
+
+
+def test_push_to_ipfs_sync_retry_config(mock_cid_decode):
+    """Test session matches retry config."""
+    # Reset
+    ipfs_module._SYNC_SESSION = None
+
+    with patch("requests.Session") as mock_session_cls:
+        mock_session = MagicMock()
+        mock_session_cls.return_value = mock_session
+        mock_session.post.return_value.status_code = 200
+        mock_session.post.return_value.json.return_value = {"Hash": "QmHash"}
+
+        push_to_ipfs_sync(b"data")
+
+        # Verify adapter mounting
+        # Since we mock the class return value, we check calls on the return value
+        assert mock_session.mount.call_count == 2
+        mock_session.mount.assert_any_call("https://", ANY)
+        mock_session.mount.assert_any_call("http://", ANY)

iwa/core/tests/test_pricing.py

@@ -0,0 +1,65 @@
+"""Tests for Pricing module."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import requests
+
+from iwa.core.pricing import PriceService
+
+
+@pytest.fixture
+def mock_secrets():
+    """Mock secrets."""
+    with patch("iwa.core.pricing.secrets") as mock_s:
+        mock_s.coingecko_api_key.get_secret_value.return_value = "fake_key"
+        yield mock_s
+
+
+@pytest.fixture
+def price_service(mock_secrets):
+    """PriceService fixture."""
+    return PriceService()
+
+
+def test_init_session(price_service):
+    """Test session initialization."""
+    assert isinstance(price_service.session, requests.Session)
+
+    # Verify adapters are mounted
+    assert "https://" in price_service.session.adapters
+    assert "http://" in price_service.session.adapters
+
+    # Verify retry configuration in adapter
+    adapter = price_service.session.adapters["https://"]
+    assert adapter.max_retries.total == 3
+    assert adapter.max_retries.status_forcelist == [429, 500, 502, 503, 504]
+
+
+def test_close(price_service):
+    """Test close method."""
+    price_service.session = MagicMock()
+    price_service.close()
+    price_service.session.close.assert_called_once()
+
+
+def test_get_token_price_uses_session(price_service):
+    """Test get_token_price uses session."""
+    price_service.session = MagicMock()
+
+    # Mock response
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"autonolas": {"eur": 5.0}}
+    price_service.session.get.return_value = mock_response
+
+    price = price_service.get_token_price("autonolas", "eur")
+
+    assert price == 5.0
+    price_service.session.get.assert_called()
+
+    # Verify call args
+    args, kwargs = price_service.session.get.call_args
+    assert "api.coingecko.com" in args[0]
+    assert kwargs["params"]["ids"] == "autonolas"
+    assert kwargs["params"]["vs_currencies"] == "eur"

iwa/core/tests/test_regression_fixes.py

@@ -0,0 +1,100 @@
+"""Regression tests for recent fixes."""
+
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from iwa.core.keys import KeyStorage
+from iwa.core.models import EthereumAddress
+from iwa.core.services.transaction import TransactionService
+
+
+class TestRegressionFixes(unittest.TestCase):
+    """Test regression fixes for fees and serialization."""
+
+    def test_transaction_service_fee_autofill(self):
+        """Test that TransactionService auto-fills fees if missing."""
+        # 1. Setup mocks
+        mock_key_storage = MagicMock()
+        mock_account_service = MagicMock()
+
+        # Mock account resolution
+        mock_account = MagicMock()
+        mock_account.address = "0x1234567890123456789012345678901234567890"
+        mock_account.tag = "signer"
+        mock_account_service.resolve_account.return_value = mock_account
+
+        service = TransactionService(mock_key_storage, mock_account_service)
+
+        # 2. Mock ChainInterface to return specific fees
+        mock_chain_interface = MagicMock()
+        mock_chain_interface.chain.chain_id = 100
+        mock_chain_interface.web3.eth.get_transaction_count.return_value = 1
+
+        # This is what we are testing: get_suggested_fees() provides the safety net
+        mock_chain_interface.get_suggested_fees.return_value = {
+            "maxFeePerGas": 1500,
+            "maxPriorityFeePerGas": 10
+        }
+
+        with patch("iwa.core.services.transaction.ChainInterfaces") as mock_interfaces:
+            mock_interfaces.return_value.get.return_value = mock_chain_interface
+
+            # 3. Prepare transaction WITHOUT fees
+            tx = {
+                "to": "0x09312C66A14a024B4e903D986Ca7E2C0dDD06227",
+                "value": 1000,
+                "gas": 21000
+            }
+
+            # 4. Run internal preparation
+            service._prepare_transaction(tx, "signer", mock_chain_interface)
+
+            # 5. Verify fees were auto-filled
+            self.assertEqual(tx["maxFeePerGas"], 1500)
+            self.assertEqual(tx["maxPriorityFeePerGas"], 10)
+            self.assertIn("nonce", tx)
+            self.assertEqual(tx["chainId"], 100)
+
+    def test_key_storage_mode_json_serialization(self):
+        """Test that KeyStorage uses mode='json' to serialize EthereumAddress correctly."""
+        import tempfile
+
+        from iwa.core.models import StoredAccount
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir) / "wallet.json"
+
+            # 1. Initialize KeyStorage
+            storage = KeyStorage(path=tmp_path, password="test")
+
+            # 2. Add an account with a real EthereumAddress object
+            addr_str = "0x1234567890123456789012345678901234567890"
+            addr_obj = EthereumAddress(addr_str)
+            acc = StoredAccount(address=addr_obj, tag="test-tag")
+            storage.accounts = {addr_obj: acc}
+
+            # 3. Mock json.dump to see what's being written
+            with patch("json.dump") as mock_dump:
+                # We need to mock open() as well to prevent real file creation if not needed,
+                # but since we are in a temp dir, it's fine.
+                storage.save()
+
+                # 4. Capture the data passed to json.dump
+                self.assertTrue(mock_dump.called)
+                dumped_data = mock_dump.call_args[0][0]
+
+                # 5. Verify the address is a plain string in the dump
+                accounts = dumped_data["accounts"]
+
+                for key in accounts.keys():
+                    # Key serialization in Pydantic v2 model_dump(mode='json')
+                    self.assertIsInstance(key, str, f"Key {key} should be a string")
+                    self.assertEqual(key.lower(), addr_str.lower())
+
+                    # Also check the address field inside the value
+                    self.assertIsInstance(accounts[key]["address"], str)
+                    self.assertEqual(accounts[key]["address"].lower(), addr_str.lower())
+
+if __name__ == "__main__":
+    unittest.main()

iwa/core/wallet.py

@@ -1,13 +1,13 @@
 """Wallet module."""
 
 from concurrent.futures import ThreadPoolExecutor
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 from web3.types import Wei
 
 from iwa.core.chain import SupportedChain
 from iwa.core.db import init_db
-from iwa.core.keys import KeyStorage
+from iwa.core.keys import EncryptedAccount, KeyStorage
 from iwa.core.models import EthereumAddress, StoredSafeAccount
 from iwa.core.services import (
     AccountService,
@@ -51,7 +51,7 @@ class Wallet:
         init_db()
 
     @property
-    def master_account(self) -> Optional[StoredSafeAccount]:
+    def master_account(self) -> Optional[Union[EncryptedAccount, StoredSafeAccount]]:
         """Get master account"""
         return self.account_service.master_account
 

iwa/plugins/gnosis/cow/quotes.py

@@ -33,7 +33,7 @@ async def get_max_sell_amount_wei(
     safe_address: ChecksumAddress | None = None,
     app_data: str | None = None,
     env: "Envs" = "prod",
-    slippage_tolerance: float = 0.005,
+    slippage_tolerance: float = 0.015,
 ) -> int:
     """Calculate the estimated sell amount needed to buy a fixed amount of tokens."""
     if app_data is None:
@@ -95,7 +95,7 @@ async def get_max_buy_amount_wei(
     safe_address: ChecksumAddress | None = None,
     app_data: str | None = None,
     env: "Envs" = "prod",
-    slippage_tolerance: float = 0.005,
+    slippage_tolerance: float = 0.015,
 ) -> int:
     """Calculate the maximum buy amount for a given sell amount."""
     if app_data is None:

iwa/plugins/gnosis/cow/swap.py

@@ -96,13 +96,9 @@ class CowSwap:
 
         logger.info(f"Checking order status for UID: {order.uid}")
 
-        max_retries = 6
-        sleep_between_retries = 25
-        retries = 0
-
-        while retries < max_retries:
-            retries += 1
+        sleep_between_retries = 15
 
+        while True:
             try:
                 # Use a thread executor for blocking requests.get
                 loop = asyncio.get_event_loop()
@@ -116,47 +112,37 @@ class CowSwap:
 
             if response.status_code != HTTP_OK:
                 logger.debug(
-                    f"Order status check {retries}/{max_retries}: HTTP {response.status_code}. Retry in {sleep_between_retries}s"
+                    f"Order status check: HTTP {response.status_code}. Retry in {sleep_between_retries}s"
                 )
                 await asyncio.sleep(sleep_between_retries)
                 continue
 
             order_data = response.json()
-
             status = order_data.get("status", "unknown")
             valid_to = int(order_data.get("validTo", 0))
             current_time = int(time.time())
 
-            if status == "expired" or (valid_to > 0 and current_time > valid_to):
-                logger.error(
-                    f"Order expired without execution (Status: {status}, ValidTo: {valid_to}, Now: {current_time})."
-                )
-                return None
-
+            # Log current status
             executed_sell = int(order_data.get("executedSellAmount", "0"))
             executed_buy = int(order_data.get("executedBuyAmount", "0"))
 
-            if executed_sell > 0 or executed_buy > 0:
+            if executed_sell > 0 or executed_buy > 0 or status == "fulfilled":
                 logger.info("Order executed successfully.")
-                sell_price = order_data.get("quote", {}).get("sellTokenPrice", None)
-                buy_price = order_data.get("quote", {}).get("buyTokenPrice", None)
-
-                if sell_price is not None:
-                    logger.debug(f"Sell price: ${float(sell_price):.2f}")
+                return order_data
 
-                if buy_price is not None:
-                    logger.debug(f"Buy price: ${float(buy_price):.2f}")
+            if status in ["expired", "cancelled"]:
+                logger.error(f"Order {status} without execution.")
+                return None
 
-                return order_data
+            if valid_to > 0 and current_time > valid_to + 60:
+                logger.error(
+                    f"Order timeout: current time {current_time} exceeded valid_to {valid_to} by >60s."
+                )
+                return None
 
-            logger.info(
-                f"Order pending... ({retries}/{max_retries}). Retry in {sleep_between_retries}s"
-            )
+            logger.info(f"Order status: {status}. Waiting {sleep_between_retries}s...")
             await asyncio.sleep(sleep_between_retries)
 
-        logger.warning("Max retries reached. Order status unknown.")
-        return None
-
     async def swap(
         self,
         amount_wei: Wei,
@@ -255,7 +241,7 @@ class CowSwap:
         safe_address: ChecksumAddress | None = None,
         app_data: str | None = None,
         env: "Envs" = "prod",
-        slippage_tolerance: float = 0.005,
+        slippage_tolerance: float = 0.015,
     ) -> int:
         """Calculate the estimated sell amount needed to buy a fixed amount of tokens."""
         return await get_max_sell_amount_wei(
@@ -278,7 +264,7 @@ class CowSwap:
         safe_address: ChecksumAddress | None = None,
         app_data: str | None = None,
         env: "Envs" = "prod",
-        slippage_tolerance: float = 0.005,
+        slippage_tolerance: float = 0.015,
     ) -> int:
         """Calculate the maximum buy amount for a given sell amount."""
         return await get_max_buy_amount_wei(
@@ -304,7 +290,7 @@ class CowSwap:
         app_data: str | None = None,
         valid_to: int | None = None,
         env: "Envs" = "prod",
-        slippage_tolerance: float = 0.005,
+        slippage_tolerance: float = 0.015,
         partially_fillable: bool = False,
     ) -> "CompletedOrder":
         """Execute a 'Buy' order (Exact Output) on CoW Protocol."""

iwa/plugins/gnosis/tests/test_cow.py

@@ -89,8 +89,8 @@ def test_init(cowswap, mock_chain):
 async def test_get_max_sell_amount_wei(cowswap, mock_cowpy_modules):
     """Test get_max_sell_amount_wei."""
     amount = await cowswap.get_max_sell_amount_wei(100, "0xSell", "0xBuy")
-    # mocked sellAmount root is "100", slippage is 0.005 default -> 100 * 1.005 = 100
-    assert amount == 100
+    # mocked sellAmount root is "100", slippage is 0.015 default -> 100 * 1.015 = 101.5 -> int 101
+    assert amount == 101
     mock_cowpy_modules["get_order_quote"].assert_called_once()
 
 
@@ -98,8 +98,8 @@ async def test_get_max_sell_amount_wei(cowswap, mock_cowpy_modules):
 async def test_get_max_buy_amount_wei(cowswap, mock_cowpy_modules):
     """Test get_max_buy_amount_wei."""
     amount = await cowswap.get_max_buy_amount_wei(100, "0xSell", "0xBuy")
-    # mocked buyAmount root is "90", slippage 0.005 -> 90 * 0.995 = 89.55 -> int 89
-    assert amount == 89
+    # mocked buyAmount root is "90", slippage 0.015 -> 90 * 0.985 = 88.65 -> int 88
+    assert amount == 88
     mock_cowpy_modules["get_order_quote"].assert_called_once()
 
 
@@ -213,15 +213,24 @@ async def test_check_cowswap_order_expired(cowswap):
 
 @pytest.mark.asyncio
 async def test_check_cowswap_order_timeout(cowswap):
-    """Test check_cowswap_order timeout."""
+    """Test check_cowswap_order timeout after exceeding valid_to."""
     mock_order = MagicMock()
     mock_order.url = "http://api/order"
 
     with patch("requests.get") as mock_get:
         mock_get.return_value.status_code = 200
-        mock_get.return_value.json.return_value = {"status": "open", "executedSellAmount": "0"}
+        # Order is always open
+        mock_get.return_value.json.return_value = {
+            "status": "open",
+            "executedSellAmount": "0",
+            "validTo": 1000,
+        }
 
-        # Speed up retry sleep (asyncio.sleep)
-        with patch("asyncio.sleep", new_callable=AsyncMock):
-            result = await cowswap.check_cowswap_order(mock_order)
-            assert result is None
+        # Mock time to start at 900 and then jump to 1100 to trigger timeout
+        with patch("time.time") as mock_time:
+            mock_time.side_effect = [900, 1100]
+            # Speed up retry sleep (asyncio.sleep)
+            with patch("asyncio.sleep", new_callable=AsyncMock):
+                result = await cowswap.check_cowswap_order(mock_order)
+                assert result is None
+                assert mock_time.call_count >= 2

iwa/plugins/olas/importer.py

@@ -767,7 +767,7 @@ class OlasServiceImporter:
             safe_result = self._import_safe(
                 address=service.safe_address,
                 signers=self._get_agent_signers(service),
-                tag_suffix="safe", # e.g. trader_zeta_safe
+                tag_suffix="multisig", # e.g. trader_zeta_safe
                 service_name=service.service_name
             )
             if safe_result[0]:
@@ -785,7 +785,7 @@ class OlasServiceImporter:
             safe_result = self._import_safe(
                 address=service.service_owner_multisig_address,
                 signers=owner_signers,
-                tag_suffix="owner_safe", # e.g. trader_zeta_owner_safe
+                tag_suffix="owner_multisig", # e.g. trader_zeta_owner_safe
                 service_name=service.service_name
             )
             if safe_result[0]:
@@ -878,8 +878,7 @@ class OlasServiceImporter:
                 self.key_storage._password,
                 tag,
             )
-            self.key_storage.accounts[encrypted.address] = encrypted
-            self.key_storage.save()
+            self.key_storage.register_account(encrypted)
             logger.info(f"Imported key {key.address} as '{tag}'")
             return True, "ok"
         except Exception as e:
@@ -926,7 +925,7 @@ class OlasServiceImporter:
         self,
         address: str,
         signers: List[str] = None,
-        tag_suffix: str = "safe",
+        tag_suffix: str = "multisig",
         service_name: Optional[str] = None
     ) -> Tuple[bool, str]:
         """Import a generic Safe."""
@@ -960,8 +959,7 @@ class OlasServiceImporter:
             signers=signers or [],
         )
 
-        self.key_storage.accounts[address] = safe_account
-        self.key_storage.save()
+        self.key_storage.register_account(safe_account)
         logger.info(f"Imported Safe {address} as '{tag}'")
         return True, "ok"