From 802590c2b94ded9a78b3a9d9d1289a35418aefd6 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 17 Apr 2026 10:57:42 +0800 Subject: [PATCH] refactor: extract ResolveClaw402WalletKey to store layer and expand OKX margin mode tests - Move duplicated claw402 wallet resolution logic into store.AIModelStore.ResolveClaw402WalletKey - api/strategy.go and manager/trader_manager.go now delegate to the shared method - Add detailed doc comment on OKX SetMarginMode explaining the local-state-only approach and why the legacy /api/v5/account/set-isolated-mode endpoint is not called - Add 3 new test cases: cross mode leverage, OpenShort tdMode, SetTakeProfit tdMode --- api/strategy.go | 33 +--------------- manager/trader_manager.go | 19 ++++----- store/ai_model.go | 37 +++++++++++++++++ trader/okx/trader_account.go | 14 ++++++- trader/okx/trader_margin_mode_test.go | 57 +++++++++++++++++++++++++++ 5 files changed, 115 insertions(+), 45 deletions(-) diff --git a/api/strategy.go b/api/strategy.go index f2813efb..64105fa5 100644 --- a/api/strategy.go +++ b/api/strategy.go @@ -708,36 +708,5 @@ func (s *Server) runRealAITest(userID, modelID, systemPrompt, userPrompt string) } func (s *Server) resolveStrategyDataWalletKey(userID, selectedModelID string) (string, error) { - if selectedModelID != "" { - model, err := s.store.AIModel().Get(userID, selectedModelID) - if err != nil { - return "", fmt.Errorf("failed to load selected AI model") - } - - if model.Provider == "claw402" { - walletKey := string(model.APIKey) - if walletKey == "" { - return "", fmt.Errorf("selected claw402 model is missing wallet private key") - } - return walletKey, nil - } - } - - models, err := s.store.AIModel().List(userID) - if err != nil { - return "", fmt.Errorf("failed to load AI models") - } - - for _, model := range models { - if model == nil || model.Provider != "claw402" { - continue - } - - walletKey := string(model.APIKey) - if walletKey != "" { - return walletKey, nil - } - } - - return "", nil + return s.store.AIModel().ResolveClaw402WalletKey(userID, selectedModelID) } diff --git a/manager/trader_manager.go b/manager/trader_manager.go index 8c65941d..36b745b8 100644 --- a/manager/trader_manager.go +++ b/manager/trader_manager.go @@ -744,6 +744,7 @@ func (tm *TraderManager) addTraderFromStore(traderCfg *store.Trader, aiModelCfg } func resolveTraderDataWalletKey(st *store.Store, userID string, selectedModel *store.AIModel) string { + // Fast path: selected model is itself a claw402 model. if selectedModel != nil && selectedModel.Provider == "claw402" { if walletKey := string(selectedModel.APIKey); walletKey != "" { return walletKey @@ -754,20 +755,14 @@ func resolveTraderDataWalletKey(st *store.Store, userID string, selectedModel *s return "" } - models, err := st.AIModel().List(userID) + // Fallback: find any configured claw402 model for this user so that paid + // NofxAI data sources work even when a non-claw402 model (e.g. deepseek) is + // selected as the AI brain. + preferredID := "" + walletKey, err := st.AIModel().ResolveClaw402WalletKey(userID, preferredID) if err != nil { logger.Warnf("⚠️ Failed to load claw402 wallet for trader data routing: %v", err) return "" } - - for _, model := range models { - if model == nil || model.Provider != "claw402" { - continue - } - if walletKey := string(model.APIKey); walletKey != "" { - return walletKey - } - } - - return "" + return walletKey } diff --git a/store/ai_model.go b/store/ai_model.go index e1af9d0f..ef577cc1 100644 --- a/store/ai_model.go +++ b/store/ai_model.go @@ -253,6 +253,43 @@ func (s *AIModelStore) Update(userID, id string, enabled bool, apiKey, customAPI } // Create creates an AI model +// ResolveClaw402WalletKey returns the claw402 wallet private key for a user. +// If preferredModelID is non-empty and points to a claw402 model, its key is returned first. +// Otherwise the first enabled claw402 model in the user's model list is used. +// Returns ("", nil) when no claw402 model is configured — callers should treat this as +// "no paid data routing" rather than an error. +func (s *AIModelStore) ResolveClaw402WalletKey(userID, preferredModelID string) (string, error) { + if preferredModelID != "" { + model, err := s.Get(userID, preferredModelID) + if err != nil { + return "", fmt.Errorf("failed to load selected AI model") + } + if model.Provider == "claw402" { + walletKey := string(model.APIKey) + if walletKey == "" { + return "", fmt.Errorf("selected claw402 model is missing wallet private key") + } + return walletKey, nil + } + } + + models, err := s.List(userID) + if err != nil { + return "", fmt.Errorf("failed to load AI models") + } + + for _, model := range models { + if model == nil || model.Provider != "claw402" { + continue + } + if walletKey := string(model.APIKey); walletKey != "" { + return walletKey, nil + } + } + + return "", nil +} + func (s *AIModelStore) Create(userID, id, name, provider string, enabled bool, apiKey, customAPIURL string) error { model := &AIModel{ ID: id, diff --git a/trader/okx/trader_account.go b/trader/okx/trader_account.go index 416a109d..9220b3d8 100644 --- a/trader/okx/trader_account.go +++ b/trader/okx/trader_account.go @@ -80,7 +80,19 @@ func (t *OKXTrader) GetBalance() (map[string]interface{}, error) { return result, nil } -// SetMarginMode sets margin mode +// SetMarginMode configures the margin mode (cross/isolated) that will be applied +// to all subsequent leverage and order requests for this trader instance. +// +// OKX V5 unified accounts do not expose a per-symbol mode-switch endpoint that +// works reliably — the legacy /api/v5/account/set-isolated-mode endpoint returns +// error 51000 ("Parameter isoMode error") when called on a unified account. +// Instead, OKX applies the mode per-request via the mgnMode field on +// /api/v5/account/set-leverage and via the tdMode field on order placement. +// +// This implementation therefore stores the configured mode locally and injects it +// into each subsequent API request, rather than making an API call here. +// NOTE: unlike Binance/Bybit implementations of this interface, no network call +// is made — the method only updates local state. func (t *OKXTrader) SetMarginMode(symbol string, isCrossMargin bool) error { t.isCrossMargin = isCrossMargin mgnMode := t.marginMode() diff --git a/trader/okx/trader_margin_mode_test.go b/trader/okx/trader_margin_mode_test.go index b633a364..e4bbb694 100644 --- a/trader/okx/trader_margin_mode_test.go +++ b/trader/okx/trader_margin_mode_test.go @@ -188,3 +188,60 @@ func TestOKXPlaceLimitOrderUsesConfiguredMarginMode(t *testing.T) { t.Fatalf("expected isolated tdMode, got %#v", orderRequests[0].Body["tdMode"]) } } + +func TestOKXCrossMarginModeUsedInLeverage(t *testing.T) { + rt := &recordingTransport{} + trader := newTestOKXTrader(rt, true) // cross margin + + if err := trader.SetLeverage("BTCUSDT", 10); err != nil { + t.Fatalf("SetLeverage failed: %v", err) + } + + leverageRequests := rt.requestsForPath(okxLeveragePath) + if len(leverageRequests) != 2 { + t.Fatalf("expected 2 leverage requests, got %d", len(leverageRequests)) + } + + for _, req := range leverageRequests { + if req.Body["mgnMode"] != "cross" { + t.Fatalf("expected cross leverage mode, got %#v", req.Body["mgnMode"]) + } + } +} + +func TestOKXOpenShortUsesConfiguredMarginMode(t *testing.T) { + rt := &recordingTransport{} + trader := newTestOKXTrader(rt, false) // isolated + + if _, err := trader.OpenShort("BTCUSDT", 0.1, 5); err != nil { + t.Fatalf("OpenShort failed: %v", err) + } + + orderRequests := rt.requestsForPath(okxOrderPath) + if len(orderRequests) == 0 { + t.Fatal("expected at least one order request") + } + + lastOrder := orderRequests[len(orderRequests)-1] + if lastOrder.Body["tdMode"] != "isolated" { + t.Fatalf("expected isolated tdMode for OpenShort, got %#v", lastOrder.Body["tdMode"]) + } +} + +func TestOKXSetTakeProfitUsesConfiguredMarginMode(t *testing.T) { + rt := &recordingTransport{} + trader := newTestOKXTrader(rt, false) // isolated + + if err := trader.SetTakeProfit("BTCUSDT", "LONG", 0.1, 100000); err != nil { + t.Fatalf("SetTakeProfit failed: %v", err) + } + + algoRequests := rt.requestsForPath(okxAlgoOrderPath) + if len(algoRequests) != 1 { + t.Fatalf("expected 1 algo order request, got %d", len(algoRequests)) + } + + if algoRequests[0].Body["tdMode"] != "isolated" { + t.Fatalf("expected isolated tdMode for SetTakeProfit, got %#v", algoRequests[0].Body["tdMode"]) + } +}