From b6cfd2f85fc83b40b02c056f6e325bbd2cfa7808 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Thu, 12 Jun 2025 12:48:51 +0000
Subject: [PATCH] proto.rune.tw.ctl: support zeroing balances manually

---
 mmgen/proto/rune/tw/ctl.py | 17 +++++++++++++++++
 mmgen/tw/store.py          | 12 +-----------
 mmgen/tw/view.py           |  2 ++
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/mmgen/proto/rune/tw/ctl.py b/mmgen/proto/rune/tw/ctl.py
index 5003ea54..a319fdf3 100755
--- a/mmgen/proto/rune/tw/ctl.py
+++ b/mmgen/proto/rune/tw/ctl.py
@@ -12,7 +12,24 @@
 proto.rune.tw.ctl: THORChain tracking wallet control class
 """
 
+from ....util import ymsg
 from ....tw.store import TwCtlWithStore
 
 class THORChainTwCtl(TwCtlWithStore):
+
 	use_cached_balances = True
+
+	async def rpc_get_balance(self, addr, block='latest'):
+		assert self.rpc.is_remote, 'tw.store.rpc_get_balance(): RPC is not remote!'
+		try:
+			return self.rpc.get_balance(addr, block=block)
+		except Exception as e:
+			ymsg(f'{type(e).__name__}: {e}')
+			ymsg(f'Unable to get balance for address ‘{addr}’')
+			if self.get_cached_balance(addr, self.cur_balances, self.data_root):
+				from ....ui import keypress_confirm
+				if keypress_confirm(self.cfg, 'Zero the balance for this address?'):
+					return self.proto.coin_amt('0')
+			else:
+				import asyncio
+				await asyncio.sleep(3)
diff --git a/mmgen/tw/store.py b/mmgen/tw/store.py
index e9b61c4f..483367ab 100755
--- a/mmgen/tw/store.py
+++ b/mmgen/tw/store.py
@@ -17,7 +17,7 @@ from pathlib import Path
 
 from ..base_obj import AsyncInit
 from ..obj import TwComment
-from ..util import msg, ymsg, die, cached_property
+from ..util import msg, die, cached_property
 from ..addr import is_coin_addr, is_mmgen_id, CoinAddr
 
 from .shared import TwMMGenID, TwLabel
@@ -216,16 +216,6 @@ class TwCtlWithStore(TwCtl, metaclass=AsyncInit):
 				if self.aggressive_sync:
 					self.write()
 
-	async def rpc_get_balance(self, addr, block='latest'):
-		assert self.rpc.is_remote, 'tw.store.rpc_get_balance(): RPC is not remote!'
-		try:
-			return self.rpc.get_balance(addr, block=block)
-		except Exception as e:
-			ymsg(f'{type(e).__name__}: {e}')
-			ymsg(f'Unable to get balance for address ‘{addr}’')
-			import asyncio
-			await asyncio.sleep(3)
-
 	def get_cached_balance(self, addr, session_cache, data_root):
 		if addr in session_cache:
 			return self.proto.coin_amt(session_cache[addr])
diff --git a/mmgen/tw/view.py b/mmgen/tw/view.py
index 0c59c5ae..a207f606 100755
--- a/mmgen/tw/view.py
+++ b/mmgen/tw/view.py
@@ -743,6 +743,8 @@ class TwView(MMGenObject, metaclass=AsyncInit):
 				await parent.get_data()
 				parent.oneshot_msg = yellow(
 					f'{parent.proto.dcoin} balance for {parent.item_desc} #{idx} refreshed')
+				if res == 0:
+					return False # zeroing balance may mess up display
 
 		async def i_addr_delete(self, parent, idx):
 			if not parent.keypress_confirm(