From 09d42f1df748a02a8929a1002b2f698ed9b03007 Mon Sep 17 00:00:00 2001
From: Vlad Stan <stan.v.vlad@gmail.com>
Date: Tue, 31 Oct 2023 21:23:18 +0200
Subject: [PATCH] fix: subscription close before disconnect

---
 nostr/client/client.py | 2 +-
 nostr/relay_manager.py | 9 ++++++---
 router.py              | 9 +++++----
 views_api.py           | 4 ++--
 4 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/nostr/client/client.py b/nostr/client/client.py
index 5fa3a00..4624ff3 100644
--- a/nostr/client/client.py
+++ b/nostr/client/client.py
@@ -25,7 +25,7 @@ def reconnect(self, relays):
 
     def close(self):
         try:
-            self.relay_manager.close_subscriptions()
+            self.relay_manager.close_all_subscriptions()
             self.relay_manager.close_connections()
 
             self.running = False
diff --git a/nostr/relay_manager.py b/nostr/relay_manager.py
index b32a1ac..ff7ca9c 100644
--- a/nostr/relay_manager.py
+++ b/nostr/relay_manager.py
@@ -81,11 +81,14 @@ def close_subscription(self, id: str):
         except Exception as e:
             logger.debug(e)
 
-    def close_subscriptions(self):
-        all_subscriptions = list(self._cached_subscriptions.keys())
-        for id in all_subscriptions:
+    def close_subscriptions(self, subscriptions: List[str]):
+        for id in subscriptions:
             self.close_subscription(id)
 
+    def close_all_subscriptions(self):
+        all_subscriptions = list(self._cached_subscriptions.keys())
+        self.close_subscriptions(all_subscriptions)
+
     def check_and_restart_relays(self):
         stopped_relays = [r for r in self.relays.values() if r.shutdown]
         for relay in stopped_relays:
diff --git a/router.py b/router.py
index cef82e2..e6ccdef 100644
--- a/router.py
+++ b/router.py
@@ -31,7 +31,7 @@ def start(self):
         self.tasks.append(asyncio.create_task(self._client_to_nostr()))
         self.tasks.append(asyncio.create_task(self._nostr_to_client()))
 
-    def stop(self):
+    async def stop(self):
         nostr_client.relay_manager.close_subscriptions(self.subscriptions)
         self.connected = False
 
@@ -42,7 +42,7 @@ def stop(self):
                 pass
 
         try:
-            self.websocket.close()
+            await self.websocket.close()
         except Exception as _:
             pass
 
@@ -53,8 +53,9 @@ async def _client_to_nostr(self):
         while self.connected:
             try:
                 json_str = await self.websocket.receive_text()
-            except WebSocketDisconnect:
-                self.stop()
+            except WebSocketDisconnect as e:
+                logger.debug(e)
+                await self.stop()
                 break
 
             try:
diff --git a/views_api.py b/views_api.py
index b0ebee6..14642d0 100644
--- a/views_api.py
+++ b/views_api.py
@@ -117,7 +117,7 @@ async def api_test_endpoint(data: TestMessage) -> TestMessageResponse:
 async def api_stop():
     for router in all_routers:
         try:
-            router.stop()
+            await router.stop()
             all_routers.remove(router)
         except Exception as e:
             logger.error(e)
@@ -146,6 +146,6 @@ async def ws_relay(websocket: WebSocket) -> None:
     while router.connected:
         await asyncio.sleep(10)
 
-    router.stop()
+    await router.stop()
     all_routers.remove(router)