add back tx tracking and clean blockexplorer api

lnbits/app.py

@@ -477,7 +477,6 @@ async def check_and_register_extensions(app: FastAPI) -> None:
 
 
 def register_async_tasks() -> None:
-
     task_manager.init()
 
     # listen to all incoming payments and dispatch payment notifications

lnbits/core/views/blockexplorer_api.py

@@ -8,19 +8,19 @@ from pydantic.types import UUID4
 
 from lnbits.decorators import check_access_token, check_user_exists
 from lnbits.settings import settings
+from lnbits.task_manager import OnchainAddressEvent, OnchainTxEvent, task_manager
 from lnbits.utils.electrum import (
     AddressResponse,
+    Balance,
     BlockHeader,
     BlockInfo,
     ElectrumClient,
     ElectrumError,
     FeeResponse,
-    HistoryEntry,
     Transaction,
     parse_block_header,
     parse_raw_tx,
     scripthash_from_address,
-    scripthash_from_scriptpubkey,
 )
 
 blockexplorer_router = APIRouter(
@@ -51,46 +51,6 @@ def _client() -> ElectrumClient:
     return ElectrumClient(settings.lnbits_blockexplorer_electrum_url)
 
 
-def _watch_scripthash(tx: Transaction) -> str | None:
-    """Return scripthash of first spendable output to watch for tx confirmation."""
-    for out in tx.vout:
-        if out.scriptPubKey.type != "nulldata":
-            return scripthash_from_scriptpubkey(bytes.fromhex(out.scriptPubKey.hex))
-    return None
-
-
-async def _tx_status(
-    c: ElectrumClient, txid: str, scripthash: str | None
-) -> dict[str, Any]:
-    if scripthash:
-        try:
-            history: list[HistoryEntry] = await c.get_history(scripthash)
-            for entry in history:
-                if entry.tx_hash == txid:
-                    return {
-                        "txid": txid,
-                        "confirmed": entry.height > 0,
-                        "height": entry.height if entry.height > 0 else None,
-                        "fee": entry.fee,
-                    }
-        except ElectrumError:
-            # History too large; check mempool to determine confirmation status.
-            try:
-                mempool = await c.get_mempool(scripthash)
-                for mem in mempool:
-                    if mem.tx_hash == txid:
-                        return {
-                            "txid": txid,
-                            "confirmed": False,
-                            "height": None,
-                            "fee": mem.fee,
-                        }
-                return {"txid": txid, "confirmed": True, "height": None, "fee": None}
-            except ElectrumError:
-                pass
-    return {"txid": txid, "confirmed": False, "height": None, "fee": None}
-
-
 # ---- REST ----
 
 
@@ -213,63 +173,61 @@ async def ws_blocks(websocket: WebSocket) -> None:
         logger.debug(f"ws_blocks error: {e}")
 
 
+def _address_event_to_response(event: OnchainAddressEvent) -> AddressResponse:
+    return AddressResponse(
+        balance=Balance(confirmed=event.confirmed, unconfirmed=event.unconfirmed),
+        history=event.history,
+        history_error=event.history_error,
+    )
+
+
 @blockexplorer_router.websocket("/ws/address/{address}")
 async def ws_address(websocket: WebSocket, address: str) -> None:
     if not settings.lnbits_blockexplorer_enabled:
         await websocket.close(code=1008)
         return
     try:
-        scripthash = scripthash_from_address(address)
+        scripthash_from_address(address)  # validate
     except ValueError as e:
         await websocket.close(code=1008, reason=str(e))
         return
     await websocket.accept()
+
+    queue: asyncio.Queue[OnchainAddressEvent] = asyncio.Queue()
+    task_manager.register_ws_address_queue(address, queue)
     try:
-        async with _client() as c:
-            await c.subscribe_scripthash(scripthash)
-            balance_res, history_res = await asyncio.gather(
-                c.get_balance(scripthash),
-                c.get_history(scripthash),
-                return_exceptions=True,
-            )
-            if isinstance(balance_res, BaseException):
-                await websocket.send_json({"error": str(balance_res)})
-                return
-            history = [] if isinstance(history_res, BaseException) else history_res
-            history_error = (
-                str(history_res) if isinstance(history_res, BaseException) else None
-            )
-            resp = AddressResponse(
-                balance=balance_res, history=history, history_error=history_error
-            )
-            await websocket.send_json(resp.dict())
+        await _ws_address_loop(websocket, address, queue)
+    finally:
+        task_manager.unregister_ws_address_queue(address, queue)
 
-            async def on_address_change(params: list[Any]) -> None:
-                try:
-                    bal_r, hist_r = await asyncio.gather(
-                        c.get_balance(scripthash),
-                        c.get_history(scripthash),
-                        return_exceptions=True,
-                    )
-                    if isinstance(bal_r, BaseException):
-                        return
-                    hist = [] if isinstance(hist_r, BaseException) else hist_r
-                    h_err = str(hist_r) if isinstance(hist_r, BaseException) else None
-                    await websocket.send_json(
-                        AddressResponse(
-                            balance=bal_r, history=hist, history_error=h_err
-                        ).dict()
-                    )
-                except Exception as e:
-                    logger.debug(f"ws_address send error: {e}")
 
-            c.on("blockchain.scripthash.subscribe", on_address_change)
-            while True:
-                msg = await websocket.receive()
-                if msg["type"] == "websocket.disconnect":
+async def _ws_address_loop(
+    websocket: WebSocket,
+    address: str,
+    queue: asyncio.Queue[OnchainAddressEvent],
+) -> None:
+    try:
+        while True:
+            recv_task = asyncio.create_task(websocket.receive())
+            event_task = asyncio.create_task(queue.get())
+            done, pending = await asyncio.wait(
+                [recv_task, event_task], return_when=asyncio.FIRST_COMPLETED
+            )
+            for t in pending:
+                t.cancel()
+            if recv_task in done:
+                if recv_task.result().get("type") == "websocket.disconnect":
                     break
-    except Exception as e:
-        logger.debug(f"ws_address error: {e}")
+            if event_task in done:
+                try:
+                    await websocket.send_json(
+                        _address_event_to_response(event_task.result()).dict()
+                    )
+                except Exception as exc:
+                    logger.debug(f"ws_address send error: {exc}")
+                    break
+    except Exception as exc:
+        logger.debug(f"ws_address error: {exc}")
 
 
 @blockexplorer_router.websocket("/ws/tx/{txid}")
@@ -278,29 +236,39 @@ async def ws_tx(websocket: WebSocket, txid: str) -> None:
         await websocket.close(code=1008)
         return
     await websocket.accept()
+
+    queue: asyncio.Queue[OnchainTxEvent] = asyncio.Queue()
+    task_manager.register_ws_tx_queue(txid, queue)
     try:
-        async with _client() as c:
-            try:
-                raw = await c.get_transaction(txid)
-            except ElectrumError as e:
-                await websocket.send_json({"error": str(e)})
-                return
-            tx = parse_raw_tx(raw)
-            watch = _watch_scripthash(tx)
-            await websocket.send_json(await _tx_status(c, txid, watch))
-            if watch:
-                await c.subscribe_scripthash(watch)
+        await _ws_tx_loop(websocket, queue)
+    finally:
+        task_manager.unregister_ws_tx_queue(txid, queue)
 
-                async def on_tx_change(params: list[Any]) -> None:
-                    try:
-                        await websocket.send_json(await _tx_status(c, txid, watch))
-                    except Exception as e:
-                        logger.debug(f"ws_tx send error: {e}")
 
-                c.on("blockchain.scripthash.subscribe", on_tx_change)
-            while True:
-                msg = await websocket.receive()
-                if msg["type"] == "websocket.disconnect":
+async def _ws_tx_loop(
+    websocket: WebSocket,
+    queue: asyncio.Queue[OnchainTxEvent],
+) -> None:
+    try:
+        while True:
+            recv_task = asyncio.create_task(websocket.receive())
+            event_task = asyncio.create_task(queue.get())
+            done, pending = await asyncio.wait(
+                [recv_task, event_task], return_when=asyncio.FIRST_COMPLETED
+            )
+            for t in pending:
+                t.cancel()
+            if recv_task in done:
+                if recv_task.result().get("type") == "websocket.disconnect":
                     break
-    except Exception as e:
-        logger.debug(f"ws_tx error: {e}")
+            if event_task in done:
+                event: OnchainTxEvent = event_task.result()
+                try:
+                    await websocket.send_json(event.dict())
+                except Exception as exc:
+                    logger.debug(f"ws_tx send error: {exc}")
+                    break
+                if event.confirmed:
+                    break
+    except Exception as exc:
+        logger.debug(f"ws_tx error: {exc}")

lnbits/task_manager.py

@@ -3,14 +3,18 @@ import traceback
 import uuid
 from collections.abc import Callable, Coroutine
 from datetime import datetime, timezone
-from typing import Any
 
 from loguru import logger
 from pydantic import BaseModel
 
 from lnbits.core.models import Payment
 from lnbits.settings import settings
-from lnbits.utils.electrum import ElectrumClient, ElectrumError, scripthash_from_address
+from lnbits.utils.electrum import (
+    AddressTracker,
+    OnchainAddressEvent,
+    OnchainTxEvent,
+    TransactionTracker,
+)
 
 
 class PublicTask(BaseModel):
@@ -20,13 +24,6 @@ class PublicTask(BaseModel):
     created_at: datetime
 
 
-class OnchainAddressEvent(BaseModel):
-    address: str
-    confirmed: int  # satoshis
-    unconfirmed: int  # satoshis
-    txids: list[str]
-
-
 class Task:
     """Model used on the backend to keep track of background tasks."""
 
@@ -58,7 +55,9 @@ class TaskManager:
     tasks: list[Task] = []
     invoice_queue: asyncio.Queue[Payment] = asyncio.Queue()
     internal_invoice_queue: asyncio.Queue[Payment] = asyncio.Queue()
-    _tracked_addresses: dict[str, str] = {}  # address -> task_name
+    _tracked_addresses: dict[str, int] = {}  # address -> ref count
+    _ws_address_queues: dict[str, list[asyncio.Queue[OnchainAddressEvent]]] = {}
+    _ws_tx_queues: dict[str, list[asyncio.Queue[OnchainTxEvent]]] = {}
 
     def init(self) -> None:
         self.create_permanent_task(
@@ -171,20 +170,76 @@ class TaskManager:
         )
 
     def track_address(self, address: str) -> None:
-        """Start tracking a Bitcoin address via Electrum."""
-        if address in self._tracked_addresses:
-            return
-        task_name = f"onchain_address_{address}"
-        self._tracked_addresses[address] = task_name
-        self.create_task(self._address_tracker(address), name=task_name)
+        """Start tracking a Bitcoin address via Electrum (ref-counted)."""
+        count = self._tracked_addresses.get(address, 0)
+        self._tracked_addresses[address] = count + 1
+        if count == 0:
+            self.create_task(
+                self._address_tracker(address), name=f"onchain_address_{address}"
+            )
 
     def untrack_address(self, address: str) -> None:
-        """Stop tracking a Bitcoin address."""
-        task_name = self._tracked_addresses.pop(address, None)
-        if task_name:
-            task = self.get_task(task_name)
+        """Decrement ref count; cancel the tracker when the last caller unregisters."""
+        count = self._tracked_addresses.get(address, 0)
+        if count <= 1:
+            self._tracked_addresses.pop(address, None)
+            task = self.get_task(f"onchain_address_{address}")
             if task:
                 self.cancel_task(task)
+        else:
+            self._tracked_addresses[address] = count - 1
+
+    def register_ws_address_queue(
+        self, address: str, queue: asyncio.Queue[OnchainAddressEvent]
+    ) -> None:
+        """Register a per-connection queue for a watched address."""
+        self._ws_address_queues.setdefault(address, []).append(queue)
+        self.track_address(address)
+
+    def unregister_ws_address_queue(
+        self, address: str, queue: asyncio.Queue[OnchainAddressEvent]
+    ) -> None:
+        """Deregister a per-connection queue and decrement the address ref count."""
+        queues = self._ws_address_queues.get(address, [])
+        if queue in queues:
+            queues.remove(queue)
+        if not queues:
+            self._ws_address_queues.pop(address, None)
+        self.untrack_address(address)
+
+    def register_ws_tx_queue(
+        self, txid: str, queue: asyncio.Queue[OnchainTxEvent]
+    ) -> None:
+        """Register a per-connection queue for a watched transaction."""
+        self._ws_tx_queues.setdefault(txid, []).append(queue)
+        if len(self._ws_tx_queues[txid]) == 1:
+            self.create_task(
+                self._transaction_tracker_dispatch(txid), name=f"ws_tx_{txid}"
+            )
+
+    def unregister_ws_tx_queue(
+        self, txid: str, queue: asyncio.Queue[OnchainTxEvent]
+    ) -> None:
+        """Deregister a per-connection queue; cancel tracker when last one leaves."""
+        queues = self._ws_tx_queues.get(txid, [])
+        if queue in queues:
+            queues.remove(queue)
+        if not queues:
+            self._ws_tx_queues.pop(txid, None)
+            task = self.get_task(f"ws_tx_{txid}")
+            if task:
+                self.cancel_task(task)
+
+    def track_transaction(
+        self,
+        txid: str,
+        callback: Callable[[OnchainTxEvent], Coroutine],
+    ) -> Task:
+        """Track a transaction until confirmed, calling callback on each change."""
+        return self.create_task(
+            self._transaction_tracker(txid, callback),
+            name=f"onchain_tx_{txid}",
+        )
 
     async def _heart_beat(self) -> None:
         """A heartbeat that removes done tasks logs the number of tasks."""
@@ -256,12 +311,18 @@ class TaskManager:
             logger.debug(f"Enqueing payment to task {task.name}")
             task.invoice_queue.put_nowait(payment)
 
-    def _dispatch_onchain_event(self, event: OnchainAddressEvent) -> None:
-        """Dispatches an onchain address event to all registered onchain listeners."""
+    async def _dispatch_onchain_event(self, event: OnchainAddressEvent) -> None:
+        """Dispatches an onchain address event to listeners and WS queues."""
         for task in self.tasks:
-            if not task.onchain_queue:
-                continue
-            task.onchain_queue.put_nowait(event)
+            if task.onchain_queue:
+                task.onchain_queue.put_nowait(event)
+        for q in list(self._ws_address_queues.get(event.address, [])):
+            q.put_nowait(event)
+
+    async def _dispatch_onchain_tx_event(self, event: OnchainTxEvent) -> None:
+        """Dispatches a tx event to all WS queues watching that txid."""
+        for q in list(self._ws_tx_queues.get(event.txid, [])):
+            q.put_nowait(event)
 
     async def _invoice_listener_consumer(self) -> None:
         payment = await self.invoice_queue.get()
@@ -274,60 +335,26 @@ class TaskManager:
         self._invoice_dispatcher(payment)
 
     async def _address_tracker(self, address: str) -> None:
-        """Track an address via Electrum subscription; dispatch OnchainAddressEvents."""
-        electrum_url = settings.lnbits_blockexplorer_electrum_url
-        scripthash = scripthash_from_address(address)
+        await AddressTracker(settings.lnbits_blockexplorer_electrum_url).track(
+            address,
+            self._dispatch_onchain_event,
+            lambda: address in self._tracked_addresses and settings.lnbits_running,
+        )
 
-        while address in self._tracked_addresses and settings.lnbits_running:
-            try:
-                async with ElectrumClient(electrum_url) as client:
-                    await self._fetch_and_dispatch_address(client, address, scripthash)
-
-                    async def on_status_change(params: list[Any]) -> None:
-                        if params and params[0] == scripthash:
-                            await self._fetch_and_dispatch_address(
-                                client, address, scripthash
-                            )
-
-                    await client.subscribe_scripthash(scripthash, on_status_change)
-
-                    while (
-                        address in self._tracked_addresses and settings.lnbits_running
-                    ):
-                        await asyncio.sleep(30)
-            except asyncio.CancelledError:
-                raise
-            except Exception as exc:
-                if not settings.lnbits_running:
-                    return
-                logger.warning(f"Address tracker {address}: {exc!s}, retrying in 5s")
-                await asyncio.sleep(5)
-
-    async def _fetch_and_dispatch_address(
-        self, client: Any, address: str, scripthash: str
+    async def _transaction_tracker(
+        self, txid: str, callback: Callable[[OnchainTxEvent], Coroutine]
     ) -> None:
-        """Fetch balance + history for a scripthash and dispatch an onchain event."""
-        balance = await client.get_balance(scripthash)
-        txids: list[str] = []
-        try:
-            history = await client.get_history(scripthash)
-            txids = [e.tx_hash for e in history]
-        except ElectrumError:
-            pass
-        try:
-            mempool = await client.get_mempool(scripthash)
-            for e in mempool:
-                if e.tx_hash not in txids:
-                    txids.append(e.tx_hash)
-        except ElectrumError:
-            pass
-        self._dispatch_onchain_event(
-            OnchainAddressEvent(
-                address=address,
-                confirmed=balance.confirmed,
-                unconfirmed=balance.unconfirmed,
-                txids=txids,
-            )
+        await TransactionTracker(settings.lnbits_blockexplorer_electrum_url).track(
+            txid,
+            callback,
+            lambda: settings.lnbits_running,
+        )
+
+    async def _transaction_tracker_dispatch(self, txid: str) -> None:
+        await TransactionTracker(settings.lnbits_blockexplorer_electrum_url).track(
+            txid,
+            self._dispatch_onchain_tx_event,
+            lambda: txid in self._ws_tx_queues and settings.lnbits_running,
         )
 
 

lnbits/utils/electrum.py

@@ -12,7 +12,7 @@ import itertools
 import json
 import ssl
 import struct
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from typing import Any
 from urllib.parse import urlparse
 
@@ -712,3 +712,218 @@ class ElectrumClient:
         """Returns mempool fee histogram as FeeHistogramEntry(fee_rate, vsize) list."""
         data = await self._call("mempool.get_fee_histogram")
         return [FeeHistogramEntry(fee_rate=r[0], vsize=r[1]) for r in data]
+
+
+# ---------------------------------------------------------------------------
+# Address tracking
+# ---------------------------------------------------------------------------
+
+
+class OnchainAddressEvent(BaseModel):
+    address: str
+    confirmed: int  # satoshis
+    unconfirmed: int  # satoshis
+    history: list[HistoryEntry] = []
+    history_error: str | None = None
+
+    @property
+    def txids(self) -> list[str]:
+        return [e.tx_hash for e in self.history]
+
+
+class AddressTracker:
+    """
+    Subscribes to a Bitcoin address via Electrum and calls a callback on every
+    balance/history change.  Reconnects automatically on failure.
+
+    Args:
+        url:       Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``).
+        callback:  Async function receiving an :class:`OnchainAddressEvent`.
+        is_active: Callable returning ``False`` when tracking should stop.
+    """
+
+    def __init__(self, url: str) -> None:
+        self.url = url
+
+    async def track(
+        self,
+        address: str,
+        callback: Callable[[OnchainAddressEvent], Coroutine[Any, Any, None]],
+        is_active: Callable[[], bool],
+    ) -> None:
+        scripthash = scripthash_from_address(address)
+        while is_active():
+            try:
+                async with ElectrumClient(self.url) as client:
+                    await self._fetch_and_dispatch(
+                        client, address, scripthash, callback
+                    )
+
+                    async def on_status_change(params: list[Any]) -> None:
+                        if params and params[0] == scripthash:
+                            await self._fetch_and_dispatch(
+                                client, address, scripthash, callback
+                            )
+
+                    await client.subscribe_scripthash(scripthash, on_status_change)
+                    while is_active():
+                        await asyncio.sleep(30)
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                if not is_active():
+                    return
+                logger.warning(f"AddressTracker {address}: {exc!s}, retrying in 5s")
+                await asyncio.sleep(5)
+
+    @staticmethod
+    async def _fetch_and_dispatch(
+        client: ElectrumClient,
+        address: str,
+        scripthash: str,
+        callback: Callable[[OnchainAddressEvent], Coroutine[Any, Any, None]],
+    ) -> None:
+        balance = await client.get_balance(scripthash)
+        history: list[HistoryEntry] = []
+        history_error: str | None = None
+        try:
+            history = await client.get_history(scripthash)
+        except ElectrumError as exc:
+            history_error = str(exc)
+        try:
+            seen = {e.tx_hash for e in history}
+            for m in await client.get_mempool(scripthash):
+                if m.tx_hash not in seen:
+                    history.append(HistoryEntry(tx_hash=m.tx_hash, height=0, fee=m.fee))
+        except ElectrumError:
+            pass
+        await callback(
+            OnchainAddressEvent(
+                address=address,
+                confirmed=balance.confirmed,
+                unconfirmed=balance.unconfirmed,
+                history=history,
+                history_error=history_error,
+            )
+        )
+
+
+# ---------------------------------------------------------------------------
+# Transaction tracking
+# ---------------------------------------------------------------------------
+
+
+class OnchainTxEvent(BaseModel):
+    txid: str
+    confirmed: bool
+    height: int | None = None
+    fee: int | None = None
+
+
+def tx_watch_scripthash(tx: Transaction) -> str | None:
+    """Return the scripthash of the first spendable output, used to subscribe
+    for confirmation notifications."""
+    for out in tx.vout:
+        if out.scriptPubKey.type != "nulldata":
+            return scripthash_from_scriptpubkey(bytes.fromhex(out.scriptPubKey.hex))
+    return None
+
+
+class TransactionTracker:
+    """
+    Subscribes to a Bitcoin transaction via Electrum and calls a callback on
+    each status change (unconfirmed → confirmed).  Stops automatically once
+    the transaction is confirmed or ``is_active()`` returns ``False``.
+
+    Args:
+        url: Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``).
+    """
+
+    def __init__(self, url: str) -> None:
+        self.url = url
+
+    async def track(
+        self,
+        txid: str,
+        callback: Callable[[OnchainTxEvent], Coroutine[Any, Any, None]],
+        is_active: Callable[[], bool],
+    ) -> None:
+        while is_active():
+            try:
+                confirmed = await self._track_once(txid, callback, is_active)
+                if confirmed:
+                    return
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                if not is_active():
+                    return
+                logger.warning(
+                    f"TransactionTracker {txid[:8]}: {exc!s}, retrying in 5s"
+                )
+                await asyncio.sleep(5)
+
+    async def _track_once(
+        self,
+        txid: str,
+        callback: Callable[[OnchainTxEvent], Coroutine[Any, Any, None]],
+        is_active: Callable[[], bool],
+    ) -> bool:
+        """One connection attempt; returns True if the tx is confirmed."""
+        async with ElectrumClient(self.url) as client:
+            try:
+                raw = await client.get_transaction(txid)
+            except ElectrumError as exc:
+                logger.warning(f"TransactionTracker {txid[:8]}: {exc!s}")
+                await asyncio.sleep(10)
+                return False
+
+            scripthash = tx_watch_scripthash(parse_raw_tx(raw))
+            event = await self._fetch_status(client, txid, scripthash)
+            await callback(event)
+            if event.confirmed:
+                return True
+
+            confirmed_event = asyncio.Event()
+
+            async def on_change(
+                params: list[Any],
+                _sh: str | None = scripthash,
+                _done: asyncio.Event = confirmed_event,
+            ) -> None:
+                if params and params[0] == _sh:
+                    ev = await self._fetch_status(client, txid, _sh)
+                    await callback(ev)
+                    if ev.confirmed:
+                        _done.set()
+
+            if scripthash:
+                await client.subscribe_scripthash(scripthash, on_change)
+
+            while is_active() and not confirmed_event.is_set():
+                await asyncio.sleep(5)
+            return confirmed_event.is_set()
+
+    @staticmethod
+    async def _fetch_status(
+        client: ElectrumClient, txid: str, scripthash: str | None
+    ) -> OnchainTxEvent:
+        if scripthash:
+            try:
+                for entry in await client.get_history(scripthash):
+                    if entry.tx_hash == txid:
+                        return OnchainTxEvent(
+                            txid=txid,
+                            confirmed=entry.height > 0,
+                            height=entry.height if entry.height > 0 else None,
+                            fee=entry.fee,
+                        )
+            except ElectrumError:
+                try:
+                    for m in await client.get_mempool(scripthash):
+                        if m.tx_hash == txid:
+                            return OnchainTxEvent(txid=txid, confirmed=False, fee=m.fee)
+                    return OnchainTxEvent(txid=txid, confirmed=True)
+                except ElectrumError:
+                    pass
+        return OnchainTxEvent(txid=txid, confirmed=False)