diff --git a/core/datasource/src/main/java/com/tangem/datasource/di/StakingStoreModule.kt b/core/datasource/src/main/java/com/tangem/datasource/di/StakingStoreModule.kt
index 59b45e9e38..655f4a41b3 100644
--- a/core/datasource/src/main/java/com/tangem/datasource/di/StakingStoreModule.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/di/StakingStoreModule.kt
@@ -4,9 +4,13 @@ import android.content.Context
 import androidx.datastore.core.DataStoreFactory
 import androidx.datastore.dataStoreFile
 import com.squareup.moshi.Moshi
+import com.tangem.datasource.api.stakekit.models.response.model.YieldBalanceWrapperDTO
 import com.tangem.datasource.local.datastore.RuntimeDataStore
+import com.tangem.datasource.local.datastore.RuntimeSharedStore
 import com.tangem.datasource.local.token.*
-import com.tangem.datasource.local.token.utils.YieldBalancesSerializer
+import com.tangem.datasource.utils.MoshiDataStoreSerializer
+import com.tangem.datasource.utils.mapWithStringKeyTypes
+import com.tangem.datasource.utils.setTypes
 import com.tangem.utils.coroutines.CoroutineDispatcherProvider
 import dagger.Module
 import dagger.Provides
@@ -35,11 +39,16 @@ internal object StakingStoreModule {
         dispatchers: CoroutineDispatcherProvider,
     ): StakingBalanceStore {
         return DefaultStakingBalanceStore(
-            dataStore = DataStoreFactory.create(
-                serializer = YieldBalancesSerializer(moshi),
+            persistenceStore = DataStoreFactory.create(
+                serializer = MoshiDataStoreSerializer(
+                    moshi = moshi,
+                    types = mapWithStringKeyTypes(valueTypes = setTypes<YieldBalanceWrapperDTO>()),
+                    defaultValue = emptyMap(),
+                ),
                 produceFile = { context.dataStoreFile(fileName = "yield_balances") },
                 scope = CoroutineScope(context = dispatchers.io + SupervisorJob()),
             ),
+            runtimeStore = RuntimeSharedStore(),
         )
     }
 
diff --git a/core/datasource/src/main/java/com/tangem/datasource/local/token/DefaultStakingBalanceStore.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/DefaultStakingBalanceStore.kt
index 28f5e8298f..f5798a708a 100644
--- a/core/datasource/src/main/java/com/tangem/datasource/local/token/DefaultStakingBalanceStore.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/DefaultStakingBalanceStore.kt
@@ -2,52 +2,82 @@ package com.tangem.datasource.local.token
 
 import androidx.datastore.core.DataStore
 import com.tangem.datasource.api.stakekit.models.response.model.YieldBalanceWrapperDTO
-import com.tangem.datasource.local.token.entity.YieldBalanceWrappersDTO
+import com.tangem.datasource.local.datastore.RuntimeSharedStore
+import com.tangem.datasource.local.token.converter.YieldBalanceConverter
+import com.tangem.domain.staking.model.stakekit.YieldBalance
 import com.tangem.domain.wallets.models.UserWalletId
 import com.tangem.utils.extensions.addOrReplace
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.firstOrNull
-import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.launch
 
+internal typealias YieldBalanceWrappersDTO = Map<String, Set<YieldBalanceWrapperDTO>>
+internal typealias YieldBalanceListByWalletId = Map<UserWalletId, Set<YieldBalance>>
+
+/**
+ * Default implementation of [StakingBalanceStore]
+ *
+ * @property persistenceStore persistence store
+ * @property runtimeStore     runtime store
+ */
 internal class DefaultStakingBalanceStore(
-    private val dataStore: DataStore<YieldBalanceWrappersDTO>,
+    private val persistenceStore: DataStore<YieldBalanceWrappersDTO>,
+    private val runtimeStore: RuntimeSharedStore<YieldBalanceListByWalletId>,
 ) : StakingBalanceStore {
 
-    override fun get(userWalletId: UserWalletId): Flow<Set<YieldBalanceWrapperDTO>> {
-        return dataStore.data.map { it[userWalletId.stringValue].orEmpty() }
-    }
-
-    override suspend fun getSyncOrNull(userWalletId: UserWalletId): Set<YieldBalanceWrapperDTO>? {
-        return dataStore.data.firstOrNull()
-            ?.get(userWalletId.stringValue)
-    }
-
-    override suspend fun store(userWalletId: UserWalletId, items: Set<YieldBalanceWrapperDTO>) {
-        dataStore.updateData { current ->
-            current.toMutableMap().apply {
-                this[userWalletId.stringValue] = items
+    override fun get(userWalletId: UserWalletId): Flow<Set<YieldBalance>> = channelFlow {
+        val cachedBalances = persistenceStore.data
+            .map {
+                val wrappers = it[userWalletId.stringValue].orEmpty()
+                YieldBalanceConverter(isCached = true).convertSet(input = wrappers)
             }
+            .firstOrNull()
+            .orEmpty()
+
+        if (cachedBalances.isNotEmpty()) {
+            send(cachedBalances)
+        }
+
+        runtimeStore.get()
+            .map { it[userWalletId].orEmpty() }
+            .onEach {
+                val mergedBalances = mergeYieldBalances(cachedBalances = cachedBalances, runtimeBalances = it)
+
+                send(mergedBalances)
+            }
+            .launchIn(scope = this)
+    }
+
+    override fun get(userWalletId: UserWalletId, address: String, integrationId: String): Flow<YieldBalance?> {
+        return get(userWalletId).map { balances ->
+            balances.getBalance(address = address, integrationId = integrationId)
         }
     }
 
-    override fun get(
-        userWalletId: UserWalletId,
-        address: String,
-        integrationId: String,
-    ): Flow<YieldBalanceWrapperDTO?> {
-        return get(userWalletId)
-            .map { balances ->
-                balances.firstOrNull { it.integrationId == integrationId && it.addresses.address == address }
-            }
+    override suspend fun getSyncOrNull(userWalletId: UserWalletId): Set<YieldBalance>? {
+        return runtimeStore.getSyncOrNull()?.getValue(userWalletId)
     }
 
     override suspend fun getSyncOrNull(
         userWalletId: UserWalletId,
         address: String,
         integrationId: String,
-    ): YieldBalanceWrapperDTO? {
-        return getSyncOrNull(userWalletId)
-            ?.firstOrNull { it.integrationId == integrationId && it.addresses.address == address }
+    ): YieldBalance? {
+        val balances = getSyncOrNull(userWalletId) ?: return null
+
+        return balances.getBalance(address, integrationId)
+    }
+
+    override suspend fun store(userWalletId: UserWalletId, items: Set<YieldBalanceWrapperDTO>) {
+        coroutineScope {
+            launch {
+                storeInRuntimeStore(
+                    userWalletId = userWalletId,
+                    items = YieldBalanceConverter(isCached = false).convertSet(input = items),
+                )
+            }
+            launch { storeInPersistenceStore(userWalletId = userWalletId, items = items) }
+        }
     }
 
     override suspend fun store(
@@ -56,10 +86,85 @@ internal class DefaultStakingBalanceStore(
         address: String,
         item: YieldBalanceWrapperDTO,
     ) {
-        val balances = getSyncOrNull(userWalletId)
-            ?.addOrReplace(item) { it.integrationId == integrationId && it.addresses.address == address }
-            ?: setOf(item)
+        coroutineScope {
+            launch {
+                storeInRuntimeStore(userWalletId, integrationId, address, item)
+                storeInPersistenceStore(userWalletId, integrationId, address, item)
+            }
+        }
+    }
 
-        store(userWalletId, balances)
+    private suspend fun storeInRuntimeStore(
+        userWalletId: UserWalletId,
+        integrationId: String,
+        address: String,
+        item: YieldBalanceWrapperDTO,
+    ) {
+        val newBalance = YieldBalanceConverter(isCached = false).convert(value = item)
+
+        val balances = getSyncOrNull(userWalletId)
+            ?.addOrReplace(newBalance) { it.integrationId == integrationId && it.address == address }
+            ?: setOf(newBalance)
+
+        storeInRuntimeStore(userWalletId = userWalletId, items = balances)
+    }
+
+    private suspend fun storeInRuntimeStore(userWalletId: UserWalletId, items: Set<YieldBalance>) {
+        runtimeStore.update(default = emptyMap()) {
+            it.toMutableMap().apply {
+                this[userWalletId] = items
+            }
+        }
+    }
+
+    private suspend fun storeInPersistenceStore(userWalletId: UserWalletId, items: Set<YieldBalanceWrapperDTO>) {
+        persistenceStore.updateData { current ->
+            current.toMutableMap().apply {
+                this[userWalletId.stringValue] = items
+            }
+        }
+    }
+
+    private suspend fun storeInPersistenceStore(
+        userWalletId: UserWalletId,
+        integrationId: String,
+        address: String,
+        item: YieldBalanceWrapperDTO,
+    ) {
+        persistenceStore.updateData { current ->
+            current.toMutableMap().apply {
+                this[userWalletId.stringValue] = current[userWalletId.stringValue]
+                    ?.addOrReplace(item) { it.integrationId == integrationId && it.addresses.address == address }
+                    ?: setOf(item)
+            }
+        }
+    }
+
+    private fun mergeYieldBalances(
+        cachedBalances: Set<YieldBalance>,
+        runtimeBalances: Set<YieldBalance>,
+    ): Set<YieldBalance> {
+        return runtimeBalances
+            .map { runtime ->
+                if (runtime is YieldBalance.Error) {
+                    cachedBalances.getBalance(address = runtime.address, integrationId = runtime.integrationId)
+                        ?: runtime
+                } else {
+                    runtime
+                }
+            }
+            .toSet()
+    }
+
+    private fun Set<YieldBalance>.getBalance(address: String?, integrationId: String?): YieldBalance? {
+        return firstOrNull { yieldBalance ->
+            val data = yieldBalance as? YieldBalance.Data
+            val balance = data?.balance
+
+            val isCorrectAddress = address != null && address == data?.address
+            val isCorrectIntegration = integrationId != null && balance?.integrationId == integrationId
+
+            isCorrectIntegration && isCorrectAddress
+        }
     }
 }
\ No newline at end of file
diff --git a/core/datasource/src/main/java/com/tangem/datasource/local/token/StakingBalanceStore.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/StakingBalanceStore.kt
index 8e10d5f395..0638d1ac9f 100644
--- a/core/datasource/src/main/java/com/tangem/datasource/local/token/StakingBalanceStore.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/StakingBalanceStore.kt
@@ -1,24 +1,29 @@
 package com.tangem.datasource.local.token
 
 import com.tangem.datasource.api.stakekit.models.response.model.YieldBalanceWrapperDTO
+import com.tangem.domain.staking.model.stakekit.YieldBalance
+import com.tangem.domain.staking.model.stakekit.YieldBalanceList
 import com.tangem.domain.wallets.models.UserWalletId
 import kotlinx.coroutines.flow.Flow
 
+/** Staking balance store */
 interface StakingBalanceStore {
 
-    fun get(userWalletId: UserWalletId): Flow<Set<YieldBalanceWrapperDTO>>
+    /** Get flow of [YieldBalanceList] by [userWalletId] */
+    fun get(userWalletId: UserWalletId): Flow<Set<YieldBalance>>
 
-    suspend fun getSyncOrNull(userWalletId: UserWalletId): Set<YieldBalanceWrapperDTO>?
+    /** Get flow of [YieldBalance] by [userWalletId], [address] and [integrationId] */
+    fun get(userWalletId: UserWalletId, address: String, integrationId: String): Flow<YieldBalance?>
 
+    /** Get [YieldBalanceList] synchronously or null by [userWalletId] */
+    suspend fun getSyncOrNull(userWalletId: UserWalletId): Set<YieldBalance>?
+
+    /** Get [YieldBalance] synchronously or null by [userWalletId], [address] and [integrationId] */
+    suspend fun getSyncOrNull(userWalletId: UserWalletId, address: String, integrationId: String): YieldBalance?
+
+    /** Store [items] by [userWalletId] */
     suspend fun store(userWalletId: UserWalletId, items: Set<YieldBalanceWrapperDTO>)
 
-    fun get(userWalletId: UserWalletId, address: String, integrationId: String): Flow<YieldBalanceWrapperDTO?>
-
-    suspend fun getSyncOrNull(
-        userWalletId: UserWalletId,
-        address: String,
-        integrationId: String,
-    ): YieldBalanceWrapperDTO?
-
+    /** Store [item] by [userWalletId], [integrationId] and [address] */
     suspend fun store(userWalletId: UserWalletId, integrationId: String, address: String, item: YieldBalanceWrapperDTO)
 }
\ No newline at end of file
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/BalanceTypeConverter.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/BalanceTypeConverter.kt
similarity index 86%
rename from data/staking/src/main/java/com/tangem/data/staking/converters/BalanceTypeConverter.kt
rename to core/datasource/src/main/java/com/tangem/datasource/local/token/converter/BalanceTypeConverter.kt
index a7260e3f31..a1e9edcb49 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/BalanceTypeConverter.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/BalanceTypeConverter.kt
@@ -1,10 +1,10 @@
-package com.tangem.data.staking.converters
+package com.tangem.datasource.local.token.converter
 
 import com.tangem.datasource.api.stakekit.models.response.model.BalanceDTO.BalanceTypeDTO
 import com.tangem.domain.staking.model.stakekit.BalanceType
 import com.tangem.utils.converter.Converter
 
-internal class BalanceTypeConverter : Converter<BalanceTypeDTO, BalanceType> {
+internal object BalanceTypeConverter : Converter<BalanceTypeDTO, BalanceType> {
 
     override fun convert(value: BalanceTypeDTO): BalanceType {
         return when (value) {
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/action/PendingActionConverter.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/PendingActionConverter.kt
similarity index 83%
rename from data/staking/src/main/java/com/tangem/data/staking/converters/action/PendingActionConverter.kt
rename to core/datasource/src/main/java/com/tangem/datasource/local/token/converter/PendingActionConverter.kt
index c9ddee7f77..716d5fc960 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/action/PendingActionConverter.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/PendingActionConverter.kt
@@ -1,16 +1,14 @@
-package com.tangem.data.staking.converters.action
+package com.tangem.datasource.local.token.converter
 
 import com.tangem.datasource.api.stakekit.models.response.model.BalanceDTO
 import com.tangem.domain.staking.model.stakekit.PendingAction
 import com.tangem.utils.converter.Converter
 
-internal class PendingActionConverter : Converter<BalanceDTO.PendingAction, PendingAction> {
-
-    private val stakingActionTypeConverter by lazy(LazyThreadSafetyMode.NONE) { StakingActionTypeConverter() }
+internal object PendingActionConverter : Converter<BalanceDTO.PendingAction, PendingAction> {
 
     override fun convert(value: BalanceDTO.PendingAction): PendingAction {
         return PendingAction(
-            type = stakingActionTypeConverter.convert(value.type),
+            type = StakingActionTypeConverter.convert(value.type),
             passthrough = value.passthrough,
             args = with(value.args) {
                 PendingAction.PendingActionArgs(
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/action/StakingActionTypeConverter.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/StakingActionTypeConverter.kt
similarity index 91%
rename from data/staking/src/main/java/com/tangem/data/staking/converters/action/StakingActionTypeConverter.kt
rename to core/datasource/src/main/java/com/tangem/datasource/local/token/converter/StakingActionTypeConverter.kt
index 95829958e7..d95777148f 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/action/StakingActionTypeConverter.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/StakingActionTypeConverter.kt
@@ -1,11 +1,11 @@
-package com.tangem.data.staking.converters.action
+package com.tangem.datasource.local.token.converter
 
 import com.tangem.datasource.api.stakekit.models.response.model.action.StakingActionTypeDTO
 import com.tangem.domain.staking.model.stakekit.action.StakingActionType
 import com.tangem.utils.converter.Converter
 
 @Suppress("CyclomaticComplexMethod")
-class StakingActionTypeConverter : Converter<StakingActionTypeDTO, StakingActionType> {
+object StakingActionTypeConverter : Converter<StakingActionTypeDTO, StakingActionType> {
 
     override fun convert(value: StakingActionTypeDTO): StakingActionType {
         return when (value) {
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/StakingNetworkTypeConverter.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/StakingNetworkTypeConverter.kt
similarity index 98%
rename from data/staking/src/main/java/com/tangem/data/staking/converters/StakingNetworkTypeConverter.kt
rename to core/datasource/src/main/java/com/tangem/datasource/local/token/converter/StakingNetworkTypeConverter.kt
index 82ebb3dbbb..e69821a778 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/StakingNetworkTypeConverter.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/StakingNetworkTypeConverter.kt
@@ -1,11 +1,11 @@
-package com.tangem.data.staking.converters
+package com.tangem.datasource.local.token.converter
 
 import com.tangem.datasource.api.stakekit.models.response.model.NetworkTypeDTO
 import com.tangem.domain.staking.model.stakekit.NetworkType
 import com.tangem.utils.converter.TwoWayConverter
 
 @Suppress("CyclomaticComplexMethod", "LongMethod")
-class StakingNetworkTypeConverter : TwoWayConverter<NetworkTypeDTO, NetworkType> {
+object StakingNetworkTypeConverter : TwoWayConverter<NetworkTypeDTO, NetworkType> {
 
     override fun convert(value: NetworkTypeDTO): NetworkType {
         return when (value) {
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/TokenConverter.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/TokenConverter.kt
similarity index 74%
rename from data/staking/src/main/java/com/tangem/data/staking/converters/TokenConverter.kt
rename to core/datasource/src/main/java/com/tangem/datasource/local/token/converter/TokenConverter.kt
index 2f67326a0c..4a363aa324 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/TokenConverter.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/TokenConverter.kt
@@ -1,17 +1,15 @@
-package com.tangem.data.staking.converters
+package com.tangem.datasource.local.token.converter
 
 import com.tangem.datasource.api.stakekit.models.response.model.TokenDTO
 import com.tangem.domain.staking.model.stakekit.Token
 import com.tangem.utils.converter.TwoWayConverter
 
-class TokenConverter(
-    private val stakingNetworkTypeConverter: StakingNetworkTypeConverter,
-) : TwoWayConverter<TokenDTO, Token> {
+object TokenConverter : TwoWayConverter<TokenDTO, Token> {
 
     override fun convert(value: TokenDTO): Token {
         return Token(
             name = value.name,
-            network = stakingNetworkTypeConverter.convert(value.network),
+            network = StakingNetworkTypeConverter.convert(value.network),
             symbol = value.symbol,
             decimals = value.decimals,
             address = value.address,
@@ -24,7 +22,7 @@ class TokenConverter(
     override fun convertBack(value: Token): TokenDTO {
         return TokenDTO(
             name = value.name,
-            network = stakingNetworkTypeConverter.convertBack(value.network),
+            network = StakingNetworkTypeConverter.convertBack(value.network),
             symbol = value.symbol,
             decimals = value.decimals,
             address = value.address,
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/YieldBalanceConverter.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/YieldBalanceConverter.kt
similarity index 64%
rename from data/staking/src/main/java/com/tangem/data/staking/converters/YieldBalanceConverter.kt
rename to core/datasource/src/main/java/com/tangem/datasource/local/token/converter/YieldBalanceConverter.kt
index c0807165b7..288a1f1323 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/YieldBalanceConverter.kt
+++ b/core/datasource/src/main/java/com/tangem/datasource/local/token/converter/YieldBalanceConverter.kt
@@ -1,36 +1,36 @@
-package com.tangem.data.staking.converters
+package com.tangem.datasource.local.token.converter
 
-import com.tangem.data.staking.converters.action.PendingActionConverter
 import com.tangem.datasource.api.stakekit.models.response.model.YieldBalanceWrapperDTO
 import com.tangem.domain.staking.model.stakekit.BalanceItem
 import com.tangem.domain.staking.model.stakekit.YieldBalance
 import com.tangem.domain.staking.model.stakekit.YieldBalanceItem
 import com.tangem.utils.converter.Converter
 
-internal class YieldBalanceConverter : Converter<YieldBalanceWrapperDTO, YieldBalance> {
+internal class YieldBalanceConverter(private val isCached: Boolean) : Converter<YieldBalanceWrapperDTO, YieldBalance> {
 
-    private val pendingActionConverter by lazy(LazyThreadSafetyMode.NONE) { PendingActionConverter() }
-    private val networkTypeConverter by lazy(LazyThreadSafetyMode.NONE) { StakingNetworkTypeConverter() }
-    private val tokenConverter by lazy(LazyThreadSafetyMode.NONE) { TokenConverter(networkTypeConverter) }
-    private val balanceTypeConverter by lazy(LazyThreadSafetyMode.NONE) { BalanceTypeConverter() }
     override fun convert(value: YieldBalanceWrapperDTO): YieldBalance {
         return if (value.balances.isEmpty()) {
-            YieldBalance.Empty
+            YieldBalance.Empty(
+                integrationId = value.integrationId,
+                address = value.addresses.address,
+                isCached = isCached,
+            )
         } else {
             YieldBalance.Data(
+                integrationId = value.integrationId,
                 address = value.addresses.address,
                 balance = YieldBalanceItem(
                     items = value.balances.map { item ->
                         BalanceItem(
                             groupId = item.groupId,
-                            token = tokenConverter.convert(item.tokenDTO),
-                            type = balanceTypeConverter.convert(item.type),
+                            token = TokenConverter.convert(item.tokenDTO),
+                            type = BalanceTypeConverter.convert(item.type),
                             amount = item.amount,
                             rawCurrencyId = item.tokenDTO.coinGeckoId,
                             // tron-specific. operates validatorAddresses instead of validatorAddress
                             validatorAddress = item.validatorAddress ?: item.validatorAddresses?.get(0),
                             date = item.date?.toDateTime(),
-                            pendingActions = pendingActionConverter
+                            pendingActions = PendingActionConverter
                                 .convertList(item.pendingActions)
                                 .sortedBy { it.passthrough },
                             isPending = false,
@@ -39,6 +39,7 @@ internal class YieldBalanceConverter : Converter<YieldBalanceWrapperDTO, YieldBa
                         .sortedWith(compareBy({ it.type }, { it.amount })),
                     integrationId = value.integrationId,
                 ),
+                isCached = isCached,
             )
         }
     }
diff --git a/core/datasource/src/main/java/com/tangem/datasource/local/token/entity/YieldBalanceWrappersDTO.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/entity/YieldBalanceWrappersDTO.kt
deleted file mode 100644
index 2490b5a01b..0000000000
--- a/core/datasource/src/main/java/com/tangem/datasource/local/token/entity/YieldBalanceWrappersDTO.kt
+++ /dev/null
@@ -1,5 +0,0 @@
-package com.tangem.datasource.local.token.entity
-
-import com.tangem.datasource.api.stakekit.models.response.model.YieldBalanceWrapperDTO
-
-internal typealias YieldBalanceWrappersDTO = Map<String, Set<YieldBalanceWrapperDTO>>
\ No newline at end of file
diff --git a/core/datasource/src/main/java/com/tangem/datasource/local/token/utils/YieldBalancesSerializer.kt b/core/datasource/src/main/java/com/tangem/datasource/local/token/utils/YieldBalancesSerializer.kt
deleted file mode 100644
index 421fd14414..0000000000
--- a/core/datasource/src/main/java/com/tangem/datasource/local/token/utils/YieldBalancesSerializer.kt
+++ /dev/null
@@ -1,36 +0,0 @@
-package com.tangem.datasource.local.token.utils
-
-import androidx.datastore.core.Serializer
-import com.squareup.moshi.Moshi
-import com.squareup.moshi.Types
-import com.tangem.datasource.api.stakekit.models.response.model.YieldBalanceWrapperDTO
-import com.tangem.datasource.local.token.entity.YieldBalanceWrappersDTO
-import java.io.InputStream
-import java.io.OutputStream
-
-internal class YieldBalancesSerializer(moshi: Moshi) : Serializer<YieldBalanceWrappersDTO> {
-
-    private val adapter by lazy {
-        val type = Types.newParameterizedType(
-            Map::class.java,
-            String::class.java,
-            Types.newParameterizedType(Set::class.java, YieldBalanceWrapperDTO::class.java),
-        )
-
-        moshi.adapter<YieldBalanceWrappersDTO>(type)
-    }
-
-    override val defaultValue: YieldBalanceWrappersDTO = emptyMap()
-
-    override suspend fun readFrom(input: InputStream): YieldBalanceWrappersDTO {
-        return input.bufferedReader().use { reader ->
-            adapter.fromJson(reader.readText()) ?: defaultValue
-        }
-    }
-
-    override suspend fun writeTo(t: YieldBalanceWrappersDTO, output: OutputStream) {
-        output.bufferedWriter().use { writer ->
-            writer.write(adapter.toJson(t))
-        }
-    }
-}
\ No newline at end of file
diff --git a/data/staking/src/main/java/com/tangem/data/staking/DefaultStakingRepository.kt b/data/staking/src/main/java/com/tangem/data/staking/DefaultStakingRepository.kt
index 4694b9f4d6..4e4c6e7833 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/DefaultStakingRepository.kt
+++ b/data/staking/src/main/java/com/tangem/data/staking/DefaultStakingRepository.kt
@@ -15,10 +15,10 @@ import com.tangem.common.extensions.hexToBytes
 import com.tangem.common.extensions.toCompressedPublicKey
 import com.tangem.data.common.api.safeApiCall
 import com.tangem.data.common.cache.CacheRegistry
-import com.tangem.data.staking.converters.*
+import com.tangem.data.staking.converters.YieldBalanceListConverter
+import com.tangem.data.staking.converters.YieldConverter
 import com.tangem.data.staking.converters.action.ActionStatusConverter
 import com.tangem.data.staking.converters.action.EnterActionResponseConverter
-import com.tangem.data.staking.converters.action.StakingActionTypeConverter
 import com.tangem.data.staking.converters.transaction.GasEstimateConverter
 import com.tangem.data.staking.converters.transaction.StakingTransactionConverter
 import com.tangem.data.staking.converters.transaction.StakingTransactionStatusConverter
@@ -32,6 +32,8 @@ import com.tangem.datasource.api.stakekit.models.response.model.action.StakingAc
 import com.tangem.datasource.api.stakekit.models.response.model.transaction.tron.TronStakeKitTransaction
 import com.tangem.datasource.local.token.StakingBalanceStore
 import com.tangem.datasource.local.token.StakingYieldsStore
+import com.tangem.datasource.local.token.converter.StakingNetworkTypeConverter
+import com.tangem.datasource.local.token.converter.TokenConverter
 import com.tangem.domain.staking.model.StakingApproval
 import com.tangem.domain.staking.model.StakingAvailability
 import com.tangem.domain.staking.model.StakingEntryInfo
@@ -77,36 +79,19 @@ internal class DefaultStakingRepository(
     moshi: Moshi,
 ) : StakingRepository {
 
-    private val stakingNetworkTypeConverter = StakingNetworkTypeConverter()
-    private val networkTypeConverter = StakingNetworkTypeConverter()
     private val transactionStatusConverter = StakingTransactionStatusConverter()
     private val transactionTypeConverter = StakingTransactionTypeConverter()
     private val actionStatusConverter = ActionStatusConverter()
-    private val stakingActionTypeConverter = StakingActionTypeConverter()
-    private val tokenConverter = TokenConverter(
-        stakingNetworkTypeConverter = stakingNetworkTypeConverter,
-    )
-    private val yieldConverter = YieldConverter(
-        tokenConverter = tokenConverter,
-    )
-    private val gasEstimateConverter = GasEstimateConverter(
-        tokenConverter = tokenConverter,
-    )
+
     private val transactionConverter = StakingTransactionConverter(
-        networkTypeConverter = networkTypeConverter,
         transactionStatusConverter = transactionStatusConverter,
         transactionTypeConverter = transactionTypeConverter,
-        gasEstimateConverter = gasEstimateConverter,
     )
     private val enterActionResponseConverter = EnterActionResponseConverter(
         actionStatusConverter = actionStatusConverter,
-        stakingActionTypeConverter = stakingActionTypeConverter,
         transactionConverter = transactionConverter,
     )
 
-    private val yieldBalanceConverter = YieldBalanceConverter()
-    private val yieldBalanceListConverter = YieldBalanceListConverter(yieldBalanceConverter)
-
     private val tronStakeKitTransactionAdapter by lazy { moshi.adapter(TronStakeKitTransaction::class.java) }
     private val networkTypeAdapter by lazy { moshi.adapter(NetworkTypeDTO::class.java) }
     private val stakingActionStatusAdapter by lazy { moshi.adapter(StakingActionStatusDTO::class.java) }
@@ -157,7 +142,7 @@ internal class DefaultStakingRepository(
         return withContext(dispatchers.io) {
             val address = walletManagersFacade.getDefaultAddress(userWalletId, cryptoCurrency.network).orEmpty()
 
-            val networkTypeDto = networkTypeConverter.convertBack(networkType)
+            val networkTypeDto = StakingNetworkTypeConverter.convertBack(networkType)
             val networkTypeString = networkTypeDto.extractJsonName()
 
             val actionStatusDTO = actionStatusConverter.convertBack(stakingActionStatus)
@@ -301,7 +286,7 @@ internal class DefaultStakingRepository(
                 )
             }
 
-            gasEstimateConverter.convert(gasEstimateDTO.getOrThrow())
+            GasEstimateConverter.convert(gasEstimateDTO.getOrThrow())
         }
     }
 
@@ -382,7 +367,7 @@ internal class DefaultStakingRepository(
                 .distinctUntilChanged()
                 .collectLatest {
                     if (it != null) {
-                        send(yieldBalanceConverter.convert(it))
+                        send(it)
                     } else {
                         error("No yield balance available for currency ${cryptoCurrency.id.value}")
                     }
@@ -408,10 +393,8 @@ internal class DefaultStakingRepository(
         val integrationId = integrationIdMap[getIntegrationKey(cryptoCurrency.id)]
             ?: error("Could not get integrationId")
 
-        val result = stakingBalanceStore.getSyncOrNull(userWalletId, address, integrationId)
-            ?: return@withContext YieldBalance.Error
-
-        yieldBalanceConverter.convert(result)
+        stakingBalanceStore.getSyncOrNull(userWalletId, address, integrationId)
+            ?: YieldBalance.Error(integrationId, address)
     }
 
     override suspend fun fetchMultiYieldBalance(
@@ -434,7 +417,7 @@ internal class DefaultStakingRepository(
                     return@invokeOnExpire
                 }
 
-                val yields = yieldConverter.convertListIgnoreErrors(
+                val yields = YieldConverter.convertListIgnoreErrors(
                     input = yieldDTOs,
                     onError = { Timber.e("Error converting one of the items in enabled yields: $it") },
                 )
@@ -485,7 +468,7 @@ internal class DefaultStakingRepository(
         cryptoCurrencies: List<CryptoCurrency>,
     ): Flow<YieldBalanceList> {
         return stakingBalanceStore.get(userWalletId)
-            .map(yieldBalanceListConverter::convert)
+            .map(YieldBalanceListConverter::convert)
             .flowOn(dispatchers.io)
     }
 
@@ -495,7 +478,7 @@ internal class DefaultStakingRepository(
     ): Flow<YieldBalanceList> = channelFlow {
         stakingBalanceStore.get(userWalletId)
             .onEach {
-                val balances = yieldBalanceListConverter.convert(it)
+                val balances = YieldBalanceListConverter.convert(it)
                 send(balances)
             }
             .launchIn(scope = this + dispatchers.io)
@@ -510,15 +493,19 @@ internal class DefaultStakingRepository(
         cryptoCurrencies: List<CryptoCurrency>,
     ): YieldBalanceList = withContext(dispatchers.io) {
         fetchMultiYieldBalance(userWalletId, cryptoCurrencies)
-        val result = stakingBalanceStore.getSyncOrNull(userWalletId) ?: return@withContext YieldBalanceList.Error
-        yieldBalanceListConverter.convert(result)
+
+        stakingBalanceStore.getSyncOrNull(userWalletId)?.let(YieldBalanceListConverter::convert)
+            ?: YieldBalanceList.Error
     }
 
     override suspend fun isAnyTokenStaked(userWalletId: UserWalletId): Boolean {
         return withContext(dispatchers.io) {
             stakingBalanceStore.getSyncOrNull(userWalletId)
                 ?.let {
-                    it.isNotEmpty() && it.any { yieldBalance -> yieldBalance.balances.isNotEmpty() }
+                    it.isNotEmpty() &&
+                        it.any { yieldBalance ->
+                            (yieldBalance as? YieldBalance.Data)?.balance?.items?.isNotEmpty() == true
+                        }
                 }
                 ?: false
         }
@@ -537,7 +524,7 @@ internal class DefaultStakingRepository(
             ),
             args = ActionRequestBodyArgs(
                 amount = params.amount.toPlainString(),
-                inputToken = tokenConverter.convertBack(params.token),
+                inputToken = TokenConverter.convertBack(params.token),
                 validatorAddress = params.validatorAddress,
                 validatorAddresses = listOf(params.validatorAddress), // check on other networks
                 tronResource = getTronResource(network),
@@ -609,7 +596,7 @@ internal class DefaultStakingRepository(
     }
 
     private suspend fun getEnabledYieldsSync(): List<Yield> {
-        return yieldConverter.convertListIgnoreErrors(
+        return YieldConverter.convertListIgnoreErrors(
             input = stakingYieldsStore.getSync(),
             onError = { Timber.e("Error converting one of the items in enabled yields: $it") },
         )
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/YieldBalanceListConverter.kt b/data/staking/src/main/java/com/tangem/data/staking/converters/YieldBalanceListConverter.kt
index 0f3d08a7c1..63bc4d1655 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/YieldBalanceListConverter.kt
+++ b/data/staking/src/main/java/com/tangem/data/staking/converters/YieldBalanceListConverter.kt
@@ -1,20 +1,16 @@
 package com.tangem.data.staking.converters
 
-import com.tangem.datasource.api.stakekit.models.response.model.YieldBalanceWrapperDTO
+import com.tangem.domain.staking.model.stakekit.YieldBalance
 import com.tangem.domain.staking.model.stakekit.YieldBalanceList
 import com.tangem.utils.converter.Converter
 
-internal class YieldBalanceListConverter(
-    private val yieldBalanceConverter: YieldBalanceConverter,
-) : Converter<Set<YieldBalanceWrapperDTO>, YieldBalanceList> {
+internal object YieldBalanceListConverter : Converter<Set<YieldBalance>, YieldBalanceList> {
 
-    override fun convert(value: Set<YieldBalanceWrapperDTO>): YieldBalanceList {
+    override fun convert(value: Set<YieldBalance>): YieldBalanceList {
         return if (value.isEmpty()) {
             YieldBalanceList.Empty
         } else {
-            YieldBalanceList.Data(
-                balances = value.map(yieldBalanceConverter::convert),
-            )
+            YieldBalanceList.Data(balances = value.toList())
         }
     }
 }
\ No newline at end of file
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/YieldConverter.kt b/data/staking/src/main/java/com/tangem/data/staking/converters/YieldConverter.kt
index cc61ab8eae..9d8c056e83 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/YieldConverter.kt
+++ b/data/staking/src/main/java/com/tangem/data/staking/converters/YieldConverter.kt
@@ -4,6 +4,7 @@ import com.tangem.datasource.api.stakekit.models.response.model.AddressArgumentD
 import com.tangem.datasource.api.stakekit.models.response.model.YieldDTO
 import com.tangem.datasource.api.stakekit.models.response.model.YieldDTO.MetadataDTO.RewardScheduleDTO
 import com.tangem.datasource.api.stakekit.models.response.model.YieldDTO.ValidatorDTO.ValidatorStatusDTO
+import com.tangem.datasource.local.token.converter.TokenConverter
 import com.tangem.domain.staking.model.stakekit.AddressArgument
 import com.tangem.domain.staking.model.stakekit.Yield
 import com.tangem.domain.staking.model.stakekit.Yield.Metadata.RewardSchedule
@@ -11,15 +12,20 @@ import com.tangem.domain.staking.model.stakekit.Yield.Validator.ValidatorStatus
 import com.tangem.utils.converter.Converter
 import kotlinx.collections.immutable.toImmutableList
 
-class YieldConverter(
-    private val tokenConverter: TokenConverter,
-) : Converter<YieldDTO, Yield> {
+internal object YieldConverter : Converter<YieldDTO, Yield> {
+
+    private val PARTNERS = listOf(
+        "cosmosvaloper1wrx0x9m9ykdhw9sg04v7uljme53wuj03aa5d4f",
+        "H2tJNyMHnRF6ahCQLQ1sSycM4FGchymuzyYzUqKEuydk",
+    )
+
+    private val PARTNERS_NAMES = listOf("Meria")
 
     override fun convert(value: YieldDTO): Yield {
         return Yield(
             id = value.id.asMandatory("id"),
-            token = tokenConverter.convert(value.token.asMandatory("token")),
-            tokens = value.tokens.asMandatory("tokens").map { tokenConverter.convert(it) },
+            token = TokenConverter.convert(value.token.asMandatory("token")),
+            tokens = value.tokens.asMandatory("tokens").map(TokenConverter::convert),
             args = convertArgs(value.args.asMandatory("args")),
             status = convertStatus(value.status.asMandatory("status")),
             apy = value.apy.asMandatory("apy"),
@@ -85,9 +91,9 @@ class YieldConverter(
             logoUri = metadataDTO.logoUri.asMandatory("logoUri"),
             description = metadataDTO.description.asMandatory("description"),
             documentation = metadataDTO.documentation,
-            gasFeeToken = tokenConverter.convert(metadataDTO.gasFeeTokenDTO.asMandatory("gasFeeTokenDTO")),
-            token = tokenConverter.convert(metadataDTO.tokenDTO.asMandatory("tokenDTO")),
-            tokens = metadataDTO.tokensDTO.asMandatory("tokensDTO").map { tokenConverter.convert(it) },
+            gasFeeToken = TokenConverter.convert(metadataDTO.gasFeeTokenDTO.asMandatory("gasFeeTokenDTO")),
+            token = TokenConverter.convert(metadataDTO.tokenDTO.asMandatory("tokenDTO")),
+            tokens = metadataDTO.tokensDTO.asMandatory("tokensDTO").map(TokenConverter::convert),
             type = metadataDTO.type.asMandatory("type"),
             rewardSchedule = convertRewardSchedule(metadataDTO.rewardSchedule.asMandatory("rewardSchedule")),
             cooldownPeriod = metadataDTO.cooldownPeriod?.let { convertPeriod(it) },
@@ -183,14 +189,4 @@ class YieldConverter(
     private fun isStrategicPartner(validatorAddress: String?, validatorName: String): Boolean {
         return PARTNERS.any { it == validatorAddress } || PARTNERS_NAMES.any { it.equals(validatorName, true) }
     }
-
-    private companion object {
-        val PARTNERS = listOf(
-            "cosmosvaloper1wrx0x9m9ykdhw9sg04v7uljme53wuj03aa5d4f",
-            "H2tJNyMHnRF6ahCQLQ1sSycM4FGchymuzyYzUqKEuydk",
-        )
-        val PARTNERS_NAMES = listOf(
-            "Meria",
-        )
-    }
 }
\ No newline at end of file
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/action/EnterActionResponseConverter.kt b/data/staking/src/main/java/com/tangem/data/staking/converters/action/EnterActionResponseConverter.kt
index b46fd75f85..162b984c15 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/action/EnterActionResponseConverter.kt
+++ b/data/staking/src/main/java/com/tangem/data/staking/converters/action/EnterActionResponseConverter.kt
@@ -2,12 +2,12 @@ package com.tangem.data.staking.converters.action
 
 import com.tangem.data.staking.converters.transaction.StakingTransactionConverter
 import com.tangem.datasource.api.stakekit.models.response.ActionDTO
+import com.tangem.datasource.local.token.converter.StakingActionTypeConverter
 import com.tangem.domain.staking.model.stakekit.action.StakingAction
 import com.tangem.utils.converter.Converter
 
 class EnterActionResponseConverter(
     private val actionStatusConverter: ActionStatusConverter,
-    private val stakingActionTypeConverter: StakingActionTypeConverter,
     private val transactionConverter: StakingTransactionConverter,
 ) : Converter<ActionDTO, StakingAction> {
 
@@ -16,7 +16,7 @@ class EnterActionResponseConverter(
             id = value.id,
             integrationId = value.integrationId,
             status = actionStatusConverter.convert(value.status),
-            type = stakingActionTypeConverter.convert(value.type),
+            type = StakingActionTypeConverter.convert(value.type),
             currentStepIndex = value.currentStepIndex,
             amount = value.amount,
             validatorAddress = value.validatorAddress,
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/GasEstimateConverter.kt b/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/GasEstimateConverter.kt
index 6220823faf..a9fb76fc0a 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/GasEstimateConverter.kt
+++ b/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/GasEstimateConverter.kt
@@ -1,18 +1,16 @@
 package com.tangem.data.staking.converters.transaction
 
-import com.tangem.data.staking.converters.TokenConverter
 import com.tangem.datasource.api.stakekit.models.response.model.transaction.StakingGasEstimateDTO
+import com.tangem.datasource.local.token.converter.TokenConverter
 import com.tangem.domain.staking.model.stakekit.transaction.StakingGasEstimate
 import com.tangem.utils.converter.Converter
 
-class GasEstimateConverter(
-    private val tokenConverter: TokenConverter,
-) : Converter<StakingGasEstimateDTO, StakingGasEstimate> {
+internal object GasEstimateConverter : Converter<StakingGasEstimateDTO, StakingGasEstimate> {
 
     override fun convert(value: StakingGasEstimateDTO): StakingGasEstimate {
         return StakingGasEstimate(
             amount = value.amount,
-            token = tokenConverter.convert(value.token),
+            token = TokenConverter.convert(value.token),
             gasLimit = value.gasLimit,
         )
     }
diff --git a/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/StakingTransactionConverter.kt b/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/StakingTransactionConverter.kt
index ef86c8ef8b..bbb6b0c175 100644
--- a/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/StakingTransactionConverter.kt
+++ b/data/staking/src/main/java/com/tangem/data/staking/converters/transaction/StakingTransactionConverter.kt
@@ -1,21 +1,19 @@
 package com.tangem.data.staking.converters.transaction
 
-import com.tangem.data.staking.converters.StakingNetworkTypeConverter
 import com.tangem.datasource.api.stakekit.models.response.model.transaction.StakingTransactionDTO
+import com.tangem.datasource.local.token.converter.StakingNetworkTypeConverter
 import com.tangem.domain.staking.model.stakekit.transaction.StakingTransaction
 import com.tangem.utils.converter.Converter
 
 class StakingTransactionConverter(
-    private val networkTypeConverter: StakingNetworkTypeConverter,
     private val transactionStatusConverter: StakingTransactionStatusConverter,
     private val transactionTypeConverter: StakingTransactionTypeConverter,
-    private val gasEstimateConverter: GasEstimateConverter,
 ) : Converter<StakingTransactionDTO, StakingTransaction> {
 
     override fun convert(value: StakingTransactionDTO): StakingTransaction {
         return StakingTransaction(
             id = value.id,
-            network = networkTypeConverter.convert(value.network),
+            network = StakingNetworkTypeConverter.convert(value.network),
             status = transactionStatusConverter.convert(value.status),
             type = transactionTypeConverter.convert(value.type),
             hash = value.hash,
@@ -23,7 +21,7 @@ class StakingTransactionConverter(
             unsignedTransaction = value.unsignedTransaction,
             stepIndex = value.stepIndex,
             error = value.error,
-            gasEstimate = value.gasEstimate?.let { gasEstimateConverter.convert(it) },
+            gasEstimate = value.gasEstimate?.let(GasEstimateConverter::convert),
             stakeId = value.stakeId,
             explorerUrl = value.explorerUrl,
             ledgerHwAppId = value.ledgerHwAppId,
diff --git a/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalance.kt b/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalance.kt
index ef9946cdb7..c36f0519fa 100644
--- a/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalance.kt
+++ b/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalance.kt
@@ -6,9 +6,14 @@ import java.math.BigDecimal
 
 sealed class YieldBalance {
 
+    abstract val integrationId: String?
+    abstract val address: String?
+
     data class Data(
+        override val integrationId: String?,
+        override val address: String,
         val balance: YieldBalanceItem,
-        val address: String,
+        val isCached: Boolean,
     ) : YieldBalance() {
         fun getTotalWithRewardsStakingBalance(): BigDecimal {
             return balance.items.sumOf { it.amount }
@@ -34,9 +39,13 @@ sealed class YieldBalance {
         }
     }
 
-    data object Empty : YieldBalance()
+    data class Empty(
+        override val integrationId: String?,
+        override val address: String,
+        val isCached: Boolean,
+    ) : YieldBalance()
 
-    data object Error : YieldBalance()
+    data class Error(override val integrationId: String?, override val address: String?) : YieldBalance()
 }
 
 data class YieldBalanceItem(
diff --git a/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalanceList.kt b/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalanceList.kt
index c1f4ec5e38..23aebbbdb4 100644
--- a/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalanceList.kt
+++ b/domain/staking/models/src/main/kotlin/com/tangem/domain/staking/model/stakekit/YieldBalanceList.kt
@@ -2,9 +2,7 @@ package com.tangem.domain.staking.model.stakekit
 
 sealed class YieldBalanceList {
 
-    data class Data(
-        val balances: List<YieldBalance>,
-    ) : YieldBalanceList() {
+    data class Data(val balances: List<YieldBalance>) : YieldBalanceList() {
 
         fun getBalance(address: String?, integrationId: String?): YieldBalance {
             return balances.firstOrNull { yieldBalance ->
@@ -15,7 +13,7 @@ sealed class YieldBalanceList {
                 val isCorrectIntegration = integrationId != null && balance?.integrationId == integrationId
 
                 isCorrectIntegration && isCorrectAddress
-            } ?: YieldBalance.Error
+            } ?: YieldBalance.Error(integrationId = integrationId, address = address)
         }
     }
 
diff --git a/domain/tokens/src/test/kotlin/com/tangem/domain/tokens/repository/MockStakingRepository.kt b/domain/tokens/src/test/kotlin/com/tangem/domain/tokens/repository/MockStakingRepository.kt
index d88dda9050..dbb5002018 100644
--- a/domain/tokens/src/test/kotlin/com/tangem/domain/tokens/repository/MockStakingRepository.kt
+++ b/domain/tokens/src/test/kotlin/com/tangem/domain/tokens/repository/MockStakingRepository.kt
@@ -138,13 +138,13 @@ class MockStakingRepository : StakingRepository {
         userWalletId: UserWalletId,
         cryptoCurrency: CryptoCurrency,
     ): Flow<YieldBalance> = channelFlow {
-        send(YieldBalance.Error)
+        send(YieldBalance.Error(integrationId = null, address = null))
     }
 
     override suspend fun getSingleYieldBalanceSync(
         userWalletId: UserWalletId,
         cryptoCurrency: CryptoCurrency,
-    ): YieldBalance = YieldBalance.Error
+    ): YieldBalance = YieldBalance.Error(integrationId = null, address = null)
 
     override suspend fun fetchMultiYieldBalance(
         userWalletId: UserWalletId,
@@ -158,7 +158,11 @@ class MockStakingRepository : StakingRepository {
         userWalletId: UserWalletId,
         cryptoCurrencies: List<CryptoCurrency>,
     ): Flow<YieldBalanceList> {
-        return flowOf(YieldBalanceList.Data(listOf(YieldBalance.Error)))
+        return flowOf(
+            YieldBalanceList.Data(
+                balances = listOf(YieldBalance.Error(integrationId = null, address = null)),
+            ),
+        )
     }
 
     override fun getMultiYieldBalanceUpdatesLegacy(
@@ -170,7 +174,7 @@ class MockStakingRepository : StakingRepository {
         userWalletId: UserWalletId,
         cryptoCurrencies: List<CryptoCurrency>,
     ): YieldBalanceList = YieldBalanceList.Data(
-        balances = listOf(YieldBalance.Error),
+        balances = listOf(YieldBalance.Error(integrationId = null, address = null)),
     )
 
     override suspend fun createAction(