From ae353aa0d5e954a6be19ecbdaada1f2811589e58 Mon Sep 17 00:00:00 2001
From: cryptooda <crypto.oda@gmail.com>
Date: Fri, 9 Jan 2026 03:34:35 +0700
Subject: [PATCH] Implement balance update callback in TradingBotBase for
 immediate sync

- Removed the private _currentBalance field and replaced it with direct access to Config.BotTradingBalance.
- Added OnBalanceUpdatedCallback to TradingBotBase for immediate synchronization and database saving when the balance is updated.
- Updated LiveTradingBotGrain to set the callback for balance updates, ensuring accurate state management.
- Modified PostgreSqlBotRepository to save the updated bot trading balance during entity updates.
---
 .../Bots/Grains/LiveTradingBotGrain.cs        | 12 ++++++++
 .../Bots/TradingBotBase.cs                    | 30 +++++++++++++++----
 .../PostgreSql/PostgreSqlBotRepository.cs     |  1 +
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/src/Managing.Application/Bots/Grains/LiveTradingBotGrain.cs b/src/Managing.Application/Bots/Grains/LiveTradingBotGrain.cs
index 5d4821bc..f38615cf 100644
--- a/src/Managing.Application/Bots/Grains/LiveTradingBotGrain.cs
+++ b/src/Managing.Application/Bots/Grains/LiveTradingBotGrain.cs
@@ -593,6 +593,18 @@ public class LiveTradingBotGrain : Grain, ILiveTradingBotGrain, IRemindable
         // Load state into the trading bot instance
         LoadStateIntoTradingBot(tradingBot);
 
+        // Set up callback for immediate balance sync and save when balance is updated
+        tradingBot.OnBalanceUpdatedCallback = async () =>
+        {
+            SyncStateFromBase();
+            await _state.WriteStateAsync();
+            
+            // Save to database immediately so GetStrategiesPaginated shows correct balance
+            var botRegistry = GrainFactory.GetGrain<ILiveBotRegistryGrain>(0);
+            var status = await botRegistry.GetBotStatus(this.GetPrimaryKey());
+            await SaveBotAsync(status);
+        };
+
         return tradingBot;
     }
 
diff --git a/src/Managing.Application/Bots/TradingBotBase.cs b/src/Managing.Application/Bots/TradingBotBase.cs
index 5c976051..136fb614 100644
--- a/src/Managing.Application/Bots/TradingBotBase.cs
+++ b/src/Managing.Application/Bots/TradingBotBase.cs
@@ -40,7 +40,6 @@ public abstract class TradingBotBase : ITradingBot
     public Dictionary<string, LightSignal> Signals { get; set; }
     public Dictionary<Guid, Position> Positions { get; set; }
     public Dictionary<DateTime, decimal> WalletBalances { get; set; }
-    private decimal _currentBalance;
     public DateTime PreloadSince { get; set; }
     public int PreloadedCandlesCount { get; set; }
     public long ExecutionCount { get; set; } = 0;
@@ -51,6 +50,12 @@ public abstract class TradingBotBase : ITradingBot
     // OPTIMIZATION 2: Cache open position state to avoid expensive Positions.Any() calls
     private bool _hasOpenPosition = false;
 
+    /// <summary>
+    /// Callback to notify the grain when balance is updated (for immediate sync and save to database).
+    /// Set by the grain after creating the bot instance.
+    /// </summary>
+    public Func<Task>? OnBalanceUpdatedCallback { get; set; }
+
     public TradingBotBase(
         ILogger<TradingBotBase> logger,
         IServiceScopeFactory scopeFactory,
@@ -65,7 +70,6 @@ public abstract class TradingBotBase : ITradingBot
         Signals = new Dictionary<string, LightSignal>();
         Positions = new Dictionary<Guid, Position>();
         WalletBalances = new Dictionary<DateTime, decimal>();
-        _currentBalance = config.BotTradingBalance;
         PreloadSince = CandleHelpers.GetBotPreloadSinceFromTimeframe(config.Timeframe);
     }
 
@@ -433,13 +437,13 @@ public abstract class TradingBotBase : ITradingBot
 
         if (WalletBalances.Count == 0)
         {
-            WalletBalances[date] = _currentBalance;
+            WalletBalances[date] = Config.BotTradingBalance;
             return;
         }
 
         if (!WalletBalances.ContainsKey(date))
         {
-            WalletBalances[date] = _currentBalance;
+            WalletBalances[date] = Config.BotTradingBalance;
         }
     }
 
@@ -1270,13 +1274,15 @@ public abstract class TradingBotBase : ITradingBot
 
                 if (position.ProfitAndLoss != null)
                 {
-                    // Update the current balance when position closes
-                    _currentBalance += position.ProfitAndLoss.Net;
+                    // Update the balance when position closes
                     Config.BotTradingBalance += position.ProfitAndLoss.Net;
 
                     await LogDebugAsync(
                         string.Format("💰 Balance Updated\nNew bot trading balance: `${0:F2}`",
                             Config.BotTradingBalance));
+
+                    // For live trading, immediately sync and save to database
+                    await OnBalanceUpdatedAsync();
                 }
             }
             else
@@ -2060,6 +2066,18 @@ public abstract class TradingBotBase : ITradingBot
         await CancelAllOrders();
     }
 
+    /// <summary>
+    /// Called when the bot trading balance is updated (e.g., after a position closes).
+    /// Calls the OnBalanceUpdatedCallback if set (by the grain for live trading).
+    /// </summary>
+    protected virtual async Task OnBalanceUpdatedAsync()
+    {
+        if (OnBalanceUpdatedCallback != null)
+        {
+            await OnBalanceUpdatedCallback();
+        }
+    }
+
     // Interface implementation
     public async Task LogInformation(string message)
     {
diff --git a/src/Managing.Infrastructure.Database/PostgreSql/PostgreSqlBotRepository.cs b/src/Managing.Infrastructure.Database/PostgreSql/PostgreSqlBotRepository.cs
index 5d854dc5..b40fcc15 100644
--- a/src/Managing.Infrastructure.Database/PostgreSql/PostgreSqlBotRepository.cs
+++ b/src/Managing.Infrastructure.Database/PostgreSql/PostgreSqlBotRepository.cs
@@ -83,6 +83,7 @@ public class PostgreSqlBotRepository : IBotRepository
         existingEntity.AccumulatedRunTimeSeconds = bot.AccumulatedRunTimeSeconds;
         existingEntity.MasterBotUserId =
             bot.MasterBotUserId ?? existingEntity.MasterBotUserId;
+        existingEntity.BotTradingBalance = bot.BotTradingBalance;
 
         await _context.SaveChangesAsync().ConfigureAwait(false);
     }