fix: refresh account state before follow-up buys

Author: Pigbibi

application/rebalance_service.py

@@ -15,6 +15,45 @@ def record_skip_log(skip_logs, *, translator, with_prefix, kind, detail):
     print(with_prefix(message), flush=True)
 
 
+def resolve_rebalance_plan(
+    *,
+    indicators,
+    account_state,
+    trend_ma_window,
+    translator,
+    cash_reserve_ratio,
+    min_trade_ratio,
+    min_trade_floor,
+    rebalance_threshold_ratio,
+    small_account_deploy_ratio,
+    mid_account_deploy_ratio,
+    large_account_deploy_ratio,
+    trade_layer_decay_coeff,
+    income_layer_start_usd,
+    income_layer_max_ratio,
+    income_layer_qqqi_weight,
+    income_layer_spyi_weight,
+):
+    return build_rebalance_plan(
+        indicators,
+        account_state,
+        trend_ma_window=trend_ma_window,
+        translator=translator,
+        cash_reserve_ratio=cash_reserve_ratio,
+        min_trade_ratio=min_trade_ratio,
+        min_trade_floor=min_trade_floor,
+        rebalance_threshold_ratio=rebalance_threshold_ratio,
+        small_account_deploy_ratio=small_account_deploy_ratio,
+        mid_account_deploy_ratio=mid_account_deploy_ratio,
+        large_account_deploy_ratio=large_account_deploy_ratio,
+        trade_layer_decay_coeff=trade_layer_decay_coeff,
+        income_layer_start_usd=income_layer_start_usd,
+        income_layer_max_ratio=income_layer_max_ratio,
+        income_layer_qqqi_weight=income_layer_qqqi_weight,
+        income_layer_spyi_weight=income_layer_spyi_weight,
+    )
+
+
 def run_strategy(
     *,
     project_id,
@@ -69,9 +108,9 @@ def run_strategy(
 
     strategy_assets = ["SOXL", "SOXX", "BOXX", "QQQI", "SPYI"]
     account_state = fetch_strategy_account_state(quote_context, trade_context, strategy_assets)
-    plan = build_rebalance_plan(
-        indicators,
-        account_state,
+    plan = resolve_rebalance_plan(
+        indicators=indicators,
+        account_state=account_state,
         trend_ma_window=trend_ma_window,
         translator=translator,
         cash_reserve_ratio=cash_reserve_ratio,
@@ -91,6 +130,7 @@ def run_strategy(
     logs = []
     skip_logs = []
     action_done = False
+    sell_submitted = False
     threshold_value = plan["threshold_value"]
     limit_order_symbols = set(plan["limit_order_symbols"])
 
@@ -135,6 +175,7 @@ def run_strategy(
 
                 if submitted:
                     action_done = True
+                    sell_submitted = True
             elif plan["sellable_quantities"][symbol] <= 0 and plan["quantities"][symbol] > 0:
                 record_skip_log(
                     skip_logs,
@@ -148,6 +189,29 @@ def run_strategy(
                     ),
                 )
 
+    if sell_submitted:
+        account_state = fetch_strategy_account_state(quote_context, trade_context, strategy_assets)
+        plan = resolve_rebalance_plan(
+            indicators=indicators,
+            account_state=account_state,
+            trend_ma_window=trend_ma_window,
+            translator=translator,
+            cash_reserve_ratio=cash_reserve_ratio,
+            min_trade_ratio=min_trade_ratio,
+            min_trade_floor=min_trade_floor,
+            rebalance_threshold_ratio=rebalance_threshold_ratio,
+            small_account_deploy_ratio=small_account_deploy_ratio,
+            mid_account_deploy_ratio=mid_account_deploy_ratio,
+            large_account_deploy_ratio=large_account_deploy_ratio,
+            trade_layer_decay_coeff=trade_layer_decay_coeff,
+            income_layer_start_usd=income_layer_start_usd,
+            income_layer_max_ratio=income_layer_max_ratio,
+            income_layer_qqqi_weight=income_layer_qqqi_weight,
+            income_layer_spyi_weight=income_layer_spyi_weight,
+        )
+        threshold_value = plan["threshold_value"]
+        limit_order_symbols = set(plan["limit_order_symbols"])
+
     investable_cash = plan["investable_cash"]
     for symbol in plan["strategy_assets"]:
         diff = plan["targets"][symbol] - plan["market_values"][symbol]

tests/test_rebalance_service.py

@@ -19,8 +19,17 @@
 
 
 class RebalanceServiceNotificationTests(unittest.TestCase):
-    def _run_strategy(self, plan, *, prices):
+    def _run_strategy(
+        self,
+        plan,
+        *,
+        prices,
+        refreshed_plan=None,
+        account_states=None,
+        estimate_max_purchase_quantity_value=0,
+    ):
         sent_messages = []
+        observed_account_states = []
 
         def fake_send_tg_message(message):
             sent_messages.append(message)
@@ -43,7 +52,19 @@ def fake_submit_order_with_alert(
             logs.append(f"{log_message} [order_id=test-order]")
             return True
 
-        with patch.object(rebalance_service, "build_rebalance_plan", return_value=plan):
+        plan_side_effect = [plan, refreshed_plan or plan]
+
+        account_state_values = list(account_states or [{}, {}])
+
+        def fake_fetch_strategy_account_state(quote_context, trade_context, strategy_assets):
+            del quote_context, trade_context, strategy_assets
+            if not account_state_values:
+                raise AssertionError("unexpected extra fetch_strategy_account_state call")
+            value = account_state_values.pop(0)
+            observed_account_states.append(value)
+            return value
+
+        with patch.object(rebalance_service, "build_rebalance_plan", side_effect=plan_side_effect):
             rebalance_service.run_strategy(
                 project_id="project-1",
                 secret_name="secret-1",
@@ -72,13 +93,13 @@ def fake_submit_order_with_alert(
                 refresh_token_if_needed=lambda *args, **kwargs: "live-token",
                 build_contexts=lambda app_key, app_secret, token: ("quote-context", "trade-context"),
                 calculate_rotation_indicators=lambda quote_context, trend_window: {"soxl": {"price": 1, "ma_trend": 2}},
-                fetch_strategy_account_state=lambda quote_context, trade_context, strategy_assets: {},
+                fetch_strategy_account_state=fake_fetch_strategy_account_state,
                 fetch_last_price=lambda quote_context, symbol: prices[symbol],
-                estimate_max_purchase_quantity=lambda *args, **kwargs: 0,
+                estimate_max_purchase_quantity=lambda *args, **kwargs: estimate_max_purchase_quantity_value,
                 submit_order_with_alert=fake_submit_order_with_alert,
             )
 
-        return sent_messages
+        return sent_messages, observed_account_states
 
     def test_sell_then_buy_skip_is_sent_in_single_summary_message(self):
         plan = {
@@ -100,7 +121,7 @@ def test_sell_then_buy_skip_is_sent_in_single_summary_message(self):
             "total_strategy_equity": 60000.0,
             "portfolio_rows": (("SOXL", "SOXX"),),
         }
-        sent_messages = self._run_strategy(
+        sent_messages, _ = self._run_strategy(
             plan,
             prices={"SOXL.US": 45.94, "SOXX.US": 322.74},
         )
@@ -131,7 +152,7 @@ def test_buy_skip_without_orders_is_sent_in_single_heartbeat_message(self):
             "total_strategy_equity": 60000.0,
             "portfolio_rows": (("SOXX",),),
         }
-        sent_messages = self._run_strategy(
+        sent_messages, _ = self._run_strategy(
             plan,
             prices={"SOXX.US": 322.74},
         )
@@ -142,6 +163,49 @@ def test_buy_skip_without_orders_is_sent_in_single_heartbeat_message(self):
         self.assertIn("跳过项", sent_messages[0])
         self.assertIn("SOXX.US", sent_messages[0])
 
+    def test_refreshes_account_state_after_sell_and_can_place_followup_buy(self):
+        initial_plan = {
+            "strategy_assets": ["SOXL", "SOXX"],
+            "targets": {"SOXL": 0.0, "SOXX": 34718.05},
+            "market_values": {"SOXL": 31928.30, "SOXX": 0.0},
+            "sellable_quantities": {"SOXL": 695, "SOXX": 0},
+            "quantities": {"SOXL": 695, "SOXX": 0},
+            "current_min_trade": 100.0,
+            "limit_order_symbols": ("SOXL", "SOXX"),
+            "threshold_value": 100.0,
+            "investable_cash": 101.95,
+            "market_status": "🛡️ DE-LEVER (SOXX)",
+            "deploy_ratio_text": "57.9%",
+            "income_ratio_text": "0.0%",
+            "income_locked_ratio_text": "38.3%",
+            "signal_message": "SOXL 跌破 150 日均线，切换至 SOXX，交易层风险仓位 57.9%",
+            "available_cash": 101.95,
+            "total_strategy_equity": 60000.0,
+            "portfolio_rows": (("SOXL", "SOXX"),),
+        }
+        refreshed_plan = {
+            **initial_plan,
+            "market_values": {"SOXL": 0.0, "SOXX": 0.0},
+            "sellable_quantities": {"SOXL": 0, "SOXX": 0},
+            "quantities": {"SOXL": 0, "SOXX": 0},
+            "investable_cash": 40000.0,
+            "available_cash": 40000.0,
+        }
+        sent_messages, observed_account_states = self._run_strategy(
+            initial_plan,
+            refreshed_plan=refreshed_plan,
+            account_states=[{"phase": "before_sell"}, {"phase": "after_sell"}],
+            prices={"SOXL.US": 45.94, "SOXX.US": 322.74},
+            estimate_max_purchase_quantity_value=200,
+        )
+
+        self.assertEqual(observed_account_states, [{"phase": "before_sell"}, {"phase": "after_sell"}])
+        self.assertEqual(len(sent_messages), 1)
+        self.assertIn("🔔 【调仓指令】", sent_messages[0])
+        self.assertIn("限价卖出", sent_messages[0])
+        self.assertIn("限价买入", sent_messages[0])
+        self.assertNotIn("买入跳过", sent_messages[0])
+
 
 if __name__ == "__main__":
     unittest.main()