From 22c571a398246ee9359c7e3e347cf2818205e8db Mon Sep 17 00:00:00 2001
From: Jkorf <jankorf91@gmail.com>
Date: Tue, 22 Oct 2024 09:14:09 +0200
Subject: [PATCH] Fixed websocket subscription request revitalization

---
 .../Clients/SpotApi/KrakenSocketClientSpotApi.cs      |  7 ++++++-
 .../Subscriptions/Spot/KrakenBalanceSubscription.cs   | 11 +++++------
 .../Subscriptions/Spot/KrakenOrderSubscription.cs     | 11 +++++------
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/Kraken.Net/Clients/SpotApi/KrakenSocketClientSpotApi.cs b/Kraken.Net/Clients/SpotApi/KrakenSocketClientSpotApi.cs
index 8216183..e848275 100644
--- a/Kraken.Net/Clients/SpotApi/KrakenSocketClientSpotApi.cs
+++ b/Kraken.Net/Clients/SpotApi/KrakenSocketClientSpotApi.cs
@@ -609,7 +609,12 @@ public async Task<CallResult<KrakenCancelAfterResult>> CancelAllOrdersAfterAsync
         /// <inheritdoc />
         protected override async Task<CallResult> RevitalizeRequestAsync(Subscription subscription)
         {
-            var krakenSubscription = (KrakenSubscription)subscription;
+            var krakenSubscription = subscription as KrakenSubscription;
+            if (krakenSubscription == null)
+            {
+                // Heartbeat subscription
+                return new CallResult(null);
+            }
 
             if (!krakenSubscription.TokenRequired)
                 return new CallResult(null);
diff --git a/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenBalanceSubscription.cs b/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenBalanceSubscription.cs
index 2234858..59d5d18 100644
--- a/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenBalanceSubscription.cs
+++ b/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenBalanceSubscription.cs
@@ -8,7 +8,7 @@
 
 namespace Kraken.Net.Objects.Sockets.Subscriptions.Spot
 {
-    internal class KrakenBalanceSubscription : Subscription<KrakenSocketResponseV2<KrakenSocketSubResponse>, KrakenSocketResponseV2<KrakenSocketSubResponse>>
+    internal class KrakenBalanceSubscription : KrakenSubscription
     {
         private static readonly MessagePath _typePath = MessagePath.Get().Property("type");
 
@@ -16,14 +16,13 @@ internal class KrakenBalanceSubscription : Subscription<KrakenSocketResponseV2<K
         private readonly Action<DataEvent<IEnumerable<KrakenBalanceUpdate>>> _updateHandler;
 
         private bool? _snapshot;
-        private string _token;
 
         public override HashSet<string> ListenerIdentifiers { get; set; }
 
-        public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, Action<DataEvent<IEnumerable<KrakenBalanceSnapshot>>>? snapshotHandler, Action<DataEvent<IEnumerable<KrakenBalanceUpdate>>> updateHandler) : base(logger, false)
+        public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, Action<DataEvent<IEnumerable<KrakenBalanceSnapshot>>>? snapshotHandler, Action<DataEvent<IEnumerable<KrakenBalanceUpdate>>> updateHandler) : base(logger, true)
         {
             _snapshot = snapshot;
-            _token = token;
+            Token = token;
 
             _snapshotHandler = snapshotHandler;
             _updateHandler = updateHandler;
@@ -50,7 +49,7 @@ public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, A
                     {
                         Channel = "balances",
                         Snapshot = _snapshot,
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }
@@ -66,7 +65,7 @@ public KrakenBalanceSubscription(ILogger logger, bool? snapshot, string token, A
                     {
                         Channel = "balances",
                         Snapshot = _snapshot,
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }
diff --git a/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenOrderSubscription.cs b/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenOrderSubscription.cs
index da348a6..338d00b 100644
--- a/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenOrderSubscription.cs
+++ b/Kraken.Net/Objects/Sockets/Subscriptions/Spot/KrakenOrderSubscription.cs
@@ -7,7 +7,7 @@
 
 namespace Kraken.Net.Objects.Sockets.Subscriptions.Spot
 {
-    internal class KrakenOrderSubscription : Subscription<KrakenSocketResponseV2<KrakenSocketSubResponse>, KrakenSocketResponseV2<KrakenSocketSubResponse>>
+    internal class KrakenOrderSubscription : KrakenSubscription
     {
         private static readonly MessagePath _typePath = MessagePath.Get().Property("type");
 
@@ -15,15 +15,14 @@ internal class KrakenOrderSubscription : Subscription<KrakenSocketResponseV2<Kra
 
         private bool? _snapshotOrder;
         private bool? _snapshotTrades;
-        private string _token;
 
         public override HashSet<string> ListenerIdentifiers { get; set; }
 
-        public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapshotTrades, string token, Action<DataEvent<IEnumerable<KrakenOrderUpdate>>> updateHandler) : base(logger, false)
+        public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapshotTrades, string token, Action<DataEvent<IEnumerable<KrakenOrderUpdate>>> updateHandler) : base(logger, true)
         {
             _snapshotOrder = snapshotOrder;
             _snapshotTrades = snapshotTrades;
-            _token = token;
+            Token = token;
 
             _updateHandler = updateHandler;
 
@@ -47,7 +46,7 @@ public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapsh
                         Channel = "executions",
                         SnapshotOrders = _snapshotOrder,
                         SnapshotTrades = _snapshotTrades,
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }
@@ -62,7 +61,7 @@ public KrakenOrderSubscription(ILogger logger, bool? snapshotOrder, bool? snapsh
                     Parameters = new KrakenSocketSubRequest
                     {
                         Channel = "executions",
-                        Token = _token
+                        Token = Token
                     }
                 }, false);
         }