nautilus_derive/websocket/
client.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! `tokio-tungstenite`-backed WebSocket client for the Derive JSON-RPC stream.
17//!
18//! [`DeriveWebSocketClient`] orchestrates the connection lifecycle and exposes
19//! a typed surface for `public/login` + the initial `ticker` channel. The
20//! actual I/O runs in `super::handler::FeedHandler`; the client communicates
21//! with it through an unbounded command channel and consumes
22//! [`DeriveWsMessage`] events.
23
24use std::{
25    fmt::Debug,
26    sync::{
27        Arc,
28        atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
29    },
30    time::Duration,
31};
32
33use alloy::signers::local::PrivateKeySigner;
34use arc_swap::ArcSwap;
35use dashmap::DashMap;
36use nautilus_common::live::get_runtime;
37use nautilus_core::UUID4;
38use nautilus_network::{
39    mode::ConnectionMode,
40    ratelimiter::{RateLimiter, clock::MonotonicClock, quota::Quota},
41    websocket::{
42        AuthTracker, TransportBackend, WebSocketClient, WebSocketConfig, channel_message_handler,
43    },
44};
45use serde::{Serialize, de::DeserializeOwned};
46use serde_json::Value;
47use ustr::Ustr;
48
49use super::{
50    error::{DeriveWsError, Result},
51    handler::{
52        DeriveWsMessage, FeedHandler, HandlerCommand, orderbook_subscribe_params,
53        ticker_subscribe_params, trades_subscribe_params,
54    },
55    messages::{
56        DeriveWsChannel, WsLoginParams, WsLoginResult, WsSubscribeParams, WsSubscribeResult,
57        WsUnsubscribeParams, WsUnsubscribeResult, methods, orderbook_channel, rate_limit_key_for,
58        ticker_channel, trades_channel,
59    },
60};
61use crate::{
62    common::{
63        consts::{
64            RECONNECT_BACKOFF_FACTOR, RECONNECT_BASE_BACKOFF, RECONNECT_JITTER_MS,
65            RECONNECT_MAX_BACKOFF, RECONNECT_TIMEOUT, WS_HEARTBEAT_SECS, WS_REQUEST_TIMEOUT,
66        },
67        enums::DeriveEnvironment,
68        rate_limit::{self, DERIVE_MATCHING_RATE_KEY},
69        urls,
70    },
71    http::{
72        models::{
73            DeriveEmptyResult, DeriveOpenOrdersResult, DeriveOrder, DeriveOrderResult,
74            DeriveReplaceResult,
75        },
76        query::{
77            DeriveCancelAllParams, DeriveCancelParams, DeriveCancelTriggerOrderParams,
78            DeriveGetTriggerOrdersParams, DeriveOrderParams, DeriveReplaceParams,
79            DeriveTriggerOrderParams,
80        },
81    },
82    signing::auth::build_ws_login,
83};
84
85/// Credentials for `public/login`. The session-key signer never escapes the
86/// client; only the wallet address is exposed via [`Debug`].
87#[derive(Clone)]
88pub struct DeriveWsCredentials {
89    /// Derive Chain smart-contract wallet address (`0x`-prefixed, 42 chars).
90    pub wallet_address: String,
91    /// secp256k1 session-key signer.
92    pub signer: PrivateKeySigner,
93}
94
95impl DeriveWsCredentials {
96    /// Constructs credentials by parsing `session_key_hex` into a signer.
97    ///
98    /// # Errors
99    ///
100    /// Returns [`DeriveWsError::Transport`] when the session-key hex cannot be parsed.
101    pub fn new(wallet_address: impl Into<String>, session_key_hex: &str) -> Result<Self> {
102        let signer: PrivateKeySigner = session_key_hex
103            .parse()
104            .map_err(|e| DeriveWsError::transport(format!("invalid session key: {e}")))?;
105        Ok(Self {
106            wallet_address: wallet_address.into(),
107            signer,
108        })
109    }
110}
111
112impl Debug for DeriveWsCredentials {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        f.debug_struct(stringify!(DeriveWsCredentials))
115            .field("wallet_address", &self.wallet_address)
116            .field("signer", &"***redacted***")
117            .finish()
118    }
119}
120
121// Rate limiter keyed by request kind (matching vs non-matching), shared with the
122// command handles so each frame is paced in the caller's task before it is
123// enqueued for the feed handler.
124type WsRateLimiter = RateLimiter<Ustr, MonotonicClock>;
125
126/// WebSocket client for the Derive JSON-RPC stream.
127///
128/// Construct with [`Self::new`] (public-only) or [`Self::with_credentials`]
129/// when private channels and signed actions are needed. Call [`Self::connect`]
130/// before any subscribe call; [`Self::disconnect`] tears the connection down.
131#[derive(Debug)]
132pub struct DeriveWebSocketClient {
133    url: String,
134    transport_backend: TransportBackend,
135    proxy_url: Option<String>,
136    connection_mode: Arc<ArcSwap<AtomicU8>>,
137    signal: Arc<AtomicBool>,
138    auth_tracker: AuthTracker,
139    credentials: Option<DeriveWsCredentials>,
140    next_id: Arc<AtomicU64>,
141    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
142    out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>>,
143    subscriptions: Arc<DashMap<String, ()>>,
144    task_handle: Option<tokio::task::JoinHandle<()>>,
145    request_timeout: Duration,
146    conn_id: Arc<ArcSwap<String>>,
147    rate_limiter: Arc<WsRateLimiter>,
148}
149
150/// Cloneable command handle for Derive public market data subscriptions.
151#[derive(Debug, Clone)]
152pub struct DeriveWebSocketSubscriptionHandle {
153    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
154    subscriptions: Arc<DashMap<String, ()>>,
155    request_timeout: Duration,
156    rate_limiter: Arc<WsRateLimiter>,
157}
158
159/// Cloneable handle for issuing signed `private/*` trading requests over the
160/// WebSocket transport.
161///
162/// Carries the same `cmd_tx` the owning [`DeriveWebSocketClient`] swaps on
163/// connect/reconnect, so a handle obtained at construction stays valid for the
164/// client's lifetime. The handle is transport-only: it sends the pre-signed
165/// body and surfaces the venue's JSON-RPC outcome. Session authorization is the
166/// client's responsibility (via `public/login`).
167#[derive(Debug, Clone)]
168pub struct DeriveWsExecutionHandle {
169    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
170    request_timeout: Duration,
171    conn_id: Arc<ArcSwap<String>>,
172    rate_limiter: Arc<WsRateLimiter>,
173}
174
175impl DeriveWebSocketClient {
176    /// Builds a public-only client. URL falls back to the environment default
177    /// when `url` is `None`.
178    #[must_use]
179    pub fn new(
180        url: Option<String>,
181        environment: DeriveEnvironment,
182        transport_backend: TransportBackend,
183        proxy_url: Option<String>,
184    ) -> Self {
185        let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
186        Self::build(url, transport_backend, proxy_url, None, None)
187    }
188
189    /// Builds a client that will issue `public/login` on connect and replay
190    /// it after each reconnect.
191    ///
192    /// `max_matching_requests_per_second` sets the matching-engine rate limit
193    /// for order writes; `None` applies the Trader-tier default. See
194    /// [`crate::common::rate_limit`].
195    #[must_use]
196    pub fn with_credentials(
197        url: Option<String>,
198        environment: DeriveEnvironment,
199        transport_backend: TransportBackend,
200        proxy_url: Option<String>,
201        credentials: DeriveWsCredentials,
202        max_matching_requests_per_second: Option<u32>,
203    ) -> Self {
204        let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
205        let matching_quota = rate_limit::matching_quota(max_matching_requests_per_second);
206        Self::build(
207            url,
208            transport_backend,
209            proxy_url,
210            Some(credentials),
211            Some(matching_quota),
212        )
213    }
214
215    fn build(
216        url: String,
217        transport_backend: TransportBackend,
218        proxy_url: Option<String>,
219        credentials: Option<DeriveWsCredentials>,
220        matching_quota: Option<Quota>,
221    ) -> Self {
222        let connection_mode = Arc::new(ArcSwap::new(Arc::new(AtomicU8::new(
223            ConnectionMode::Closed as u8,
224        ))));
225        // Placeholder channel; replaced by connect() before commands are issued.
226        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
227        // Matching-engine writes draw on the matching quota (authenticated
228        // clients only); logins, subscriptions, and reads fall through to the
229        // non-matching default. Handles pace each frame against this in the
230        // caller's task before enqueueing, so the feed handler never sleeps.
231        let mut keyed_quotas: Vec<(Ustr, Quota)> = Vec::new();
232        if let Some(quota) = matching_quota {
233            keyed_quotas.push((Ustr::from(DERIVE_MATCHING_RATE_KEY), quota));
234        }
235        let rate_limiter = Arc::new(RateLimiter::new_with_quota(
236            Some(rate_limit::non_matching_quota()),
237            keyed_quotas,
238        ));
239        Self {
240            url,
241            transport_backend,
242            proxy_url,
243            connection_mode,
244            signal: Arc::new(AtomicBool::new(false)),
245            auth_tracker: AuthTracker::new(),
246            credentials,
247            next_id: Arc::new(AtomicU64::new(1)),
248            cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
249            out_rx: None,
250            subscriptions: Arc::new(DashMap::new()),
251            task_handle: None,
252            request_timeout: WS_REQUEST_TIMEOUT,
253            conn_id: Arc::new(ArcSwap::from_pointee(UUID4::new().to_string())),
254            rate_limiter,
255        }
256    }
257
258    /// Returns the configured WebSocket URL.
259    #[must_use]
260    pub fn url(&self) -> &str {
261        &self.url
262    }
263
264    /// Returns `true` when credentials are configured and the venue has
265    /// confirmed the latest `public/login`. Cleared on reconnect.
266    #[must_use]
267    pub fn is_authenticated(&self) -> bool {
268        self.auth_tracker.is_authenticated()
269    }
270
271    /// Returns `true` while the underlying transport is in the active state.
272    #[must_use]
273    pub fn is_active(&self) -> bool {
274        self.connection_mode.load().load(Ordering::Relaxed) == ConnectionMode::Active as u8
275    }
276
277    /// Establishes the WebSocket connection and spawns the I/O handler task.
278    ///
279    /// When credentials are configured, issues `public/login` and awaits the
280    /// venue's acknowledgement before returning.
281    ///
282    /// # Errors
283    ///
284    /// Returns [`DeriveWsError::Transport`] for handshake failures and
285    /// propagates [`DeriveWsError::Auth`] / [`DeriveWsError::JsonRpc`] when
286    /// the login flow fails.
287    pub async fn connect(&mut self) -> Result<()> {
288        // Fast path requires authenticated session when creds are configured;
289        // otherwise fall through and rebuild so `Ok` always implies authenticated.
290        let auth_ok = self.credentials.is_none() || self.is_authenticated();
291        if self.is_active() && auth_ok && self.task_handle.is_some() {
292            log::warn!("Derive WebSocket already connected");
293            return Ok(());
294        }
295
296        // Tear down stale state so we don't orphan the old handler task on rebuild.
297        if self.task_handle.is_some() {
298            log::debug!("Tearing down stale Derive WebSocket state before connect");
299            self.teardown().await;
300        }
301
302        let (message_handler, raw_rx) = channel_message_handler();
303        let cfg = WebSocketConfig {
304            url: self.url.clone(),
305            headers: vec![],
306            heartbeat: Some(WS_HEARTBEAT_SECS),
307            heartbeat_msg: None,
308            reconnect_timeout_ms: Some(RECONNECT_TIMEOUT.as_millis() as u64),
309            reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
310            reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
311            reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
312            reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
313            reconnect_max_attempts: None,
314            idle_timeout_ms: None,
315            backend: self.transport_backend,
316            proxy_url: self.proxy_url.clone(),
317        };
318        // Rate limiting runs caller-side via `self.rate_limiter` before frames
319        // are enqueued, so the network client's own limiter is left unconfigured
320        // and never sleeps inside the single feed-handler task.
321        let client = WebSocketClient::connect(cfg, Some(message_handler), None, None, vec![], None)
322            .await
323            .map_err(|e| DeriveWsError::transport(e.to_string()))?;
324
325        // Register the tracker so the network controller clears
326        // `is_authenticated()` on dead-socket detection, not just on the
327        // later RECONNECTED sentinel.
328        client.set_auth_tracker(self.auth_tracker.clone(), false);
329
330        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
331        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<DeriveWsMessage>();
332
333        *self.cmd_tx.write().await = cmd_tx.clone();
334        self.out_rx = Some(out_rx);
335        self.conn_id.store(Arc::new(UUID4::new().to_string()));
336
337        self.connection_mode.store(client.connection_mode_atomic());
338        log::debug!("Derive WebSocket connected: {}", self.url);
339
340        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
341            return Err(DeriveWsError::transport(format!(
342                "failed to send SetClient command: {e}",
343            )));
344        }
345
346        let signal = Arc::clone(&self.signal);
347        let auth_tracker = self.auth_tracker.clone();
348        let next_id = Arc::clone(&self.next_id);
349        let credentials = self.credentials.clone();
350        let subscriptions = Arc::clone(&self.subscriptions);
351        let conn_id = Arc::clone(&self.conn_id);
352        let cmd_tx_for_loop = cmd_tx.clone();
353        let rate_limiter = Arc::clone(&self.rate_limiter);
354        let request_timeout = self.request_timeout;
355
356        let stream_handle = get_runtime().spawn(async move {
357            let mut handler =
358                FeedHandler::new(signal, cmd_rx, raw_rx, next_id, auth_tracker.clone());
359
360            loop {
361                match handler.next().await {
362                    Some(DeriveWsMessage::Reconnected) => {
363                        log::info!("Derive WebSocket re-establishing session after reconnect");
364                        conn_id.store(Arc::new(UUID4::new().to_string()));
365
366                        if out_tx.send(DeriveWsMessage::Reconnected).is_err() {
367                            log::debug!("Derive outer receiver dropped, exiting stream loop");
368                            break;
369                        }
370
371                        // Spawn so the loop keeps draining messages while
372                        // re-login + resubscribe are in flight.
373                        let cmd_tx_async = cmd_tx_for_loop.clone();
374                        let auth_tracker_async = auth_tracker.clone();
375                        let creds_async = credentials.clone();
376                        let subs_async = Arc::clone(&subscriptions);
377                        let rate_limiter_async = Arc::clone(&rate_limiter);
378
379                        get_runtime().spawn(async move {
380                            if let Some(creds) = creds_async
381                                && let Err(e) = login_via_handler(
382                                    &rate_limiter_async,
383                                    &cmd_tx_async,
384                                    &auth_tracker_async,
385                                    &creds,
386                                    request_timeout,
387                                )
388                                .await
389                            {
390                                log::error!("Derive WebSocket re-login failed: {e}");
391                            }
392                            // Snapshot channels before awaiting: a DashMap
393                            // shard guard held across `.await` can deadlock
394                            // on a single-worker runtime.
395                            let channels: Vec<String> =
396                                subs_async.iter().map(|e| e.key().clone()).collect();
397                            for channel in channels {
398                                if let Err(e) = subscribe_via_handler(
399                                    &rate_limiter_async,
400                                    &cmd_tx_async,
401                                    vec![channel.clone()],
402                                    request_timeout,
403                                )
404                                .await
405                                {
406                                    log::error!(
407                                        "Derive WebSocket resubscribe failed for {channel}: {e}",
408                                    );
409                                }
410                            }
411                        });
412                    }
413                    Some(msg) => {
414                        if out_tx.send(msg).is_err() {
415                            log::debug!("Derive outer receiver dropped, exiting stream loop");
416                            break;
417                        }
418                    }
419                    None => {
420                        log::debug!("Derive handler task ended");
421                        break;
422                    }
423                }
424            }
425        });
426        self.task_handle = Some(stream_handle);
427
428        if let Some(creds) = self.credentials.clone()
429            && let Err(e) = login_via_handler(
430                &self.rate_limiter,
431                &cmd_tx,
432                &self.auth_tracker,
433                &creds,
434                self.request_timeout,
435            )
436            .await
437        {
438            // Without teardown, a retry connect() would short-circuit on
439            // is_active() and return Ok without a valid session.
440            log::warn!("Derive WebSocket login failed; tearing down transport: {e}");
441            self.teardown().await;
442            return Err(e);
443        }
444
445        Ok(())
446    }
447
448    /// Signals the handler to disconnect, aborts the spawn task, and resets
449    /// the client's transport-related state. Shared by [`Self::disconnect`]
450    /// and the login-failure branch of [`Self::connect`].
451    async fn teardown(&mut self) {
452        self.signal.store(true, Ordering::Relaxed);
453
454        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
455            log::debug!(
456                "Failed to enqueue Disconnect command (handler may already be shut down): {e}",
457            );
458        }
459
460        if let Some(handle) = self.task_handle.take() {
461            let abort_handle = handle.abort_handle();
462            tokio::select! {
463                result = handle => match result {
464                    Ok(()) => log::debug!("Derive WebSocket task completed"),
465                    Err(e) if e.is_cancelled() => log::debug!("Derive WebSocket task cancelled"),
466                    Err(e) => log::error!("Derive WebSocket task error: {e:?}"),
467                },
468                () = tokio::time::sleep(Duration::from_secs(2)) => {
469                    log::warn!("Timeout waiting for Derive WebSocket task, aborting");
470                    abort_handle.abort();
471                }
472            }
473        }
474
475        // Subscriptions are also dropped: the venue session ended with the
476        // transport, so a fresh connect() must re-issue them.
477        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
478        *self.cmd_tx.write().await = placeholder_tx;
479        self.out_rx = None;
480        self.connection_mode
481            .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
482        self.auth_tracker.invalidate();
483        self.subscriptions.clear();
484        self.signal.store(false, Ordering::Relaxed);
485    }
486
487    /// Disconnects the WebSocket connection and awaits the handler task.
488    ///
489    /// # Errors
490    ///
491    /// Returns [`DeriveWsError::Transport`] when the disconnect command
492    /// cannot be enqueued; the handler still tears down on signal.
493    pub async fn disconnect(&mut self) -> Result<()> {
494        log::debug!("Disconnecting Derive WebSocket");
495        self.teardown().await;
496        Ok(())
497    }
498
499    /// Subscribes to `ticker_slim.{instrument_name}.{interval}`. `interval` is the
500    /// millisecond cadence string the venue exposes (e.g. `"100"`, `"1000"`).
501    ///
502    /// # Errors
503    ///
504    /// Propagates JSON-RPC errors raised by the venue and transport-level
505    /// failures.
506    pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
507        self.subscription_handle()
508            .subscribe_ticker(instrument_name, interval)
509            .await
510    }
511
512    /// Unsubscribes from `ticker_slim.{instrument_name}.{interval}`.
513    ///
514    /// # Errors
515    ///
516    /// Propagates JSON-RPC errors raised by the venue and transport-level
517    /// failures.
518    pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
519        self.subscription_handle()
520            .unsubscribe_ticker(instrument_name, interval)
521            .await
522    }
523
524    /// Subscribes to `orderbook.{instrument_name}.{group}.{depth}`.
525    ///
526    /// # Errors
527    ///
528    /// Propagates JSON-RPC errors raised by the venue and transport-level
529    /// failures.
530    pub async fn subscribe_orderbook(
531        &self,
532        instrument_name: &str,
533        group: &str,
534        depth: &str,
535    ) -> Result<()> {
536        self.subscription_handle()
537            .subscribe_orderbook(instrument_name, group, depth)
538            .await
539    }
540
541    /// Unsubscribes from `orderbook.{instrument_name}.{group}.{depth}`.
542    ///
543    /// # Errors
544    ///
545    /// Propagates JSON-RPC errors raised by the venue and transport-level
546    /// failures.
547    pub async fn unsubscribe_orderbook(
548        &self,
549        instrument_name: &str,
550        group: &str,
551        depth: &str,
552    ) -> Result<()> {
553        self.subscription_handle()
554            .unsubscribe_orderbook(instrument_name, group, depth)
555            .await
556    }
557
558    /// Subscribes to `trades.{instrument_type}.{currency}`.
559    ///
560    /// # Errors
561    ///
562    /// Propagates JSON-RPC errors raised by the venue and transport-level
563    /// failures.
564    pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
565        self.subscription_handle()
566            .subscribe_trades(instrument_type, currency)
567            .await
568    }
569
570    /// Unsubscribes from `trades.{instrument_type}.{currency}`.
571    ///
572    /// # Errors
573    ///
574    /// Propagates JSON-RPC errors raised by the venue and transport-level
575    /// failures.
576    pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
577        self.subscription_handle()
578            .unsubscribe_trades(instrument_type, currency)
579            .await
580    }
581
582    /// Subscribes to a list of channel topics in a single `subscribe` frame.
583    ///
584    /// Used by the execution client to bulk-subscribe to the private
585    /// `{subaccount_id}.orders`, `{subaccount_id}.trades`, and
586    /// `{subaccount_id}.balances` channels after login.
587    ///
588    /// # Errors
589    ///
590    /// Propagates JSON-RPC errors raised by the venue and transport-level
591    /// failures.
592    pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
593    where
594        C: Into<DeriveWsChannel>,
595    {
596        self.subscription_handle()
597            .subscribe_channels(channels)
598            .await
599    }
600
601    /// Unsubscribes from a list of channel topics in a single
602    /// `unsubscribe` frame.
603    ///
604    /// # Errors
605    ///
606    /// Propagates JSON-RPC errors raised by the venue and transport-level
607    /// failures.
608    pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
609    where
610        C: Into<DeriveWsChannel>,
611    {
612        self.subscription_handle()
613            .unsubscribe_channels(channels)
614            .await
615    }
616
617    /// Returns the next event emitted by the handler.
618    pub async fn next_event(&mut self) -> Option<DeriveWsMessage> {
619        if let Some(rx) = self.out_rx.as_mut() {
620            rx.recv().await
621        } else {
622            None
623        }
624    }
625
626    /// Returns the count of channels the client currently has confirmed
627    /// subscriptions for.
628    #[must_use]
629    pub fn subscription_count(&self) -> usize {
630        self.subscriptions.len()
631    }
632
633    /// Returns a cloneable handle for issuing subscription commands.
634    #[must_use]
635    pub fn subscription_handle(&self) -> DeriveWebSocketSubscriptionHandle {
636        DeriveWebSocketSubscriptionHandle {
637            cmd_tx: Arc::clone(&self.cmd_tx),
638            subscriptions: Arc::clone(&self.subscriptions),
639            request_timeout: self.request_timeout,
640            rate_limiter: Arc::clone(&self.rate_limiter),
641        }
642    }
643
644    /// Returns a cloneable handle for issuing signed `private/*` trading
645    /// requests.
646    ///
647    /// The handle shares the client's command channel, so it stays valid across
648    /// reconnects (the channel is swapped behind a shared lock). Obtain it once
649    /// and clone it into each order-submission task.
650    #[must_use]
651    pub fn execution_handle(&self) -> DeriveWsExecutionHandle {
652        DeriveWsExecutionHandle {
653            cmd_tx: Arc::clone(&self.cmd_tx),
654            request_timeout: self.request_timeout,
655            conn_id: Arc::clone(&self.conn_id),
656            rate_limiter: Arc::clone(&self.rate_limiter),
657        }
658    }
659
660    /// Takes the event receiver from the client.
661    ///
662    /// This lets the live data client own the receive loop while subscription
663    /// commands continue through [`Self::subscription_handle`].
664    pub fn take_event_receiver(
665        &mut self,
666    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>> {
667        self.out_rx.take()
668    }
669}
670
671impl DeriveWebSocketSubscriptionHandle {
672    /// Subscribes to `ticker_slim.{instrument_name}.{interval}`.
673    ///
674    /// # Errors
675    ///
676    /// Propagates JSON-RPC errors raised by the venue and transport-level
677    /// failures.
678    pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
679        let channel = ticker_channel(instrument_name, interval);
680        let params = ticker_subscribe_params(instrument_name, interval);
681        self.send_subscribe(channel, &params).await
682    }
683
684    /// Unsubscribes from `ticker_slim.{instrument_name}.{interval}`.
685    ///
686    /// # Errors
687    ///
688    /// Propagates JSON-RPC errors raised by the venue and transport-level
689    /// failures.
690    pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
691        let channel = ticker_channel(instrument_name, interval);
692        self.send_unsubscribe(channel).await
693    }
694
695    /// Subscribes to `orderbook.{instrument_name}.{group}.{depth}`.
696    ///
697    /// # Errors
698    ///
699    /// Propagates JSON-RPC errors raised by the venue and transport-level
700    /// failures.
701    pub async fn subscribe_orderbook(
702        &self,
703        instrument_name: &str,
704        group: &str,
705        depth: &str,
706    ) -> Result<()> {
707        let channel = orderbook_channel(instrument_name, group, depth);
708        let params = orderbook_subscribe_params(instrument_name, group, depth);
709        self.send_subscribe(channel, &params).await
710    }
711
712    /// Unsubscribes from `orderbook.{instrument_name}.{group}.{depth}`.
713    ///
714    /// # Errors
715    ///
716    /// Propagates JSON-RPC errors raised by the venue and transport-level
717    /// failures.
718    pub async fn unsubscribe_orderbook(
719        &self,
720        instrument_name: &str,
721        group: &str,
722        depth: &str,
723    ) -> Result<()> {
724        let channel = orderbook_channel(instrument_name, group, depth);
725        self.send_unsubscribe(channel).await
726    }
727
728    /// Subscribes to `trades.{instrument_type}.{currency}`.
729    ///
730    /// # Errors
731    ///
732    /// Propagates JSON-RPC errors raised by the venue and transport-level
733    /// failures.
734    pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
735        let channel = trades_channel(instrument_type, currency);
736        let params = trades_subscribe_params(instrument_type, currency);
737        self.send_subscribe(channel, &params).await
738    }
739
740    /// Unsubscribes from `trades.{instrument_type}.{currency}`.
741    ///
742    /// # Errors
743    ///
744    /// Propagates JSON-RPC errors raised by the venue and transport-level
745    /// failures.
746    pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
747        let channel = trades_channel(instrument_type, currency);
748        self.send_unsubscribe(channel).await
749    }
750
751    /// Subscribes to multiple channel topics in a single `subscribe` frame.
752    ///
753    /// # Errors
754    ///
755    /// Propagates JSON-RPC errors raised by the venue and transport-level
756    /// failures.
757    pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
758    where
759        C: Into<DeriveWsChannel>,
760    {
761        let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
762        if channels.is_empty() {
763            return Ok(());
764        }
765        let topics = channel_topics(&channels);
766        let params = WsSubscribeParams { channels };
767        let cmd_tx = self.cmd_tx.read().await.clone();
768        let _: WsSubscribeResult = send_request(
769            &self.rate_limiter,
770            &cmd_tx,
771            methods::PUBLIC_SUBSCRIBE,
772            &params,
773            self.request_timeout,
774        )
775        .await?;
776
777        for channel in topics {
778            self.subscriptions.insert(channel, ());
779        }
780        Ok(())
781    }
782
783    /// Unsubscribes from multiple channel topics in a single
784    /// `unsubscribe` frame.
785    ///
786    /// # Errors
787    ///
788    /// Propagates JSON-RPC errors raised by the venue and transport-level
789    /// failures.
790    pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
791    where
792        C: Into<DeriveWsChannel>,
793    {
794        let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
795        if channels.is_empty() {
796            return Ok(());
797        }
798        let topics = channel_topics(&channels);
799        let params = WsUnsubscribeParams { channels };
800        let cmd_tx = self.cmd_tx.read().await.clone();
801        let _: WsUnsubscribeResult = send_request(
802            &self.rate_limiter,
803            &cmd_tx,
804            methods::PUBLIC_UNSUBSCRIBE,
805            &params,
806            self.request_timeout,
807        )
808        .await?;
809
810        for channel in topics {
811            self.subscriptions.remove(&channel);
812        }
813        Ok(())
814    }
815
816    async fn send_subscribe(&self, channel: String, params: &WsSubscribeParams) -> Result<()> {
817        let cmd_tx = self.cmd_tx.read().await.clone();
818        let _: WsSubscribeResult = send_request(
819            &self.rate_limiter,
820            &cmd_tx,
821            methods::PUBLIC_SUBSCRIBE,
822            params,
823            self.request_timeout,
824        )
825        .await?;
826        self.subscriptions.insert(channel, ());
827        Ok(())
828    }
829
830    async fn send_unsubscribe(&self, channel: String) -> Result<()> {
831        let params = WsUnsubscribeParams {
832            channels: vec![DeriveWsChannel::from(channel.clone())],
833        };
834        let cmd_tx = self.cmd_tx.read().await.clone();
835        let _: WsUnsubscribeResult = send_request(
836            &self.rate_limiter,
837            &cmd_tx,
838            methods::PUBLIC_UNSUBSCRIBE,
839            &params,
840            self.request_timeout,
841        )
842        .await?;
843        self.subscriptions.remove(&channel);
844        Ok(())
845    }
846}
847
848impl DeriveWsExecutionHandle {
849    /// Returns the current WebSocket connection id used by trigger orders.
850    #[must_use]
851    pub fn conn_id(&self) -> String {
852        self.conn_id.load_full().as_ref().clone()
853    }
854
855    /// Submits a signed order via `private/order`.
856    ///
857    /// `params` must be the fully-built signed body from
858    /// [`crate::http::query::order_to_derive_payload`]. Returns the accepted
859    /// order echoed by the venue.
860    ///
861    /// # Errors
862    ///
863    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
864    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
865    /// outcome is ambiguous.
866    pub async fn submit_order(&self, params: &DeriveOrderParams) -> Result<DeriveOrder> {
867        let cmd_tx = self.cmd_tx.read().await.clone();
868        let result: DeriveOrderResult = send_request_typed(
869            &self.rate_limiter,
870            &cmd_tx,
871            methods::PRIVATE_ORDER,
872            params,
873            self.request_timeout,
874        )
875        .await?;
876        Ok(result.order)
877    }
878
879    /// Submits a signed trigger order via `private/trigger_order`.
880    ///
881    /// # Errors
882    ///
883    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
884    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
885    /// outcome is ambiguous.
886    pub async fn submit_trigger_order(
887        &self,
888        params: &DeriveTriggerOrderParams,
889    ) -> Result<DeriveOrder> {
890        let cmd_tx = self.cmd_tx.read().await.clone();
891        let result: DeriveOrderResult = send_request_typed(
892            &self.rate_limiter,
893            &cmd_tx,
894            methods::PRIVATE_TRIGGER_ORDER,
895            params,
896            self.request_timeout,
897        )
898        .await?;
899        Ok(result.order)
900    }
901
902    /// Modifies a working order by atomically cancelling it and submitting a
903    /// replacement (the venue's `private/replace`). Returns the new order
904    /// echoed by the venue.
905    ///
906    /// # Errors
907    ///
908    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
909    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
910    /// outcome is ambiguous.
911    pub async fn modify_order(&self, params: &DeriveReplaceParams) -> Result<DeriveOrder> {
912        let cmd_tx = self.cmd_tx.read().await.clone();
913        let result: DeriveReplaceResult = send_request_typed(
914            &self.rate_limiter,
915            &cmd_tx,
916            methods::PRIVATE_REPLACE,
917            params,
918            self.request_timeout,
919        )
920        .await?;
921        Ok(result.order)
922    }
923
924    /// Cancels a single order via `private/cancel`.
925    ///
926    /// # Errors
927    ///
928    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
929    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
930    /// outcome is ambiguous.
931    pub async fn cancel_order(&self, params: &DeriveCancelParams) -> Result<()> {
932        let cmd_tx = self.cmd_tx.read().await.clone();
933        let _: DeriveEmptyResult = send_request(
934            &self.rate_limiter,
935            &cmd_tx,
936            methods::PRIVATE_CANCEL,
937            params,
938            self.request_timeout,
939        )
940        .await?;
941        Ok(())
942    }
943
944    /// Cancels a single trigger order via `private/cancel_trigger_order`.
945    ///
946    /// # Errors
947    ///
948    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
949    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
950    /// outcome is ambiguous.
951    pub async fn cancel_trigger_order(
952        &self,
953        params: &DeriveCancelTriggerOrderParams,
954    ) -> Result<DeriveOrder> {
955        let cmd_tx = self.cmd_tx.read().await.clone();
956        send_request_typed(
957            &self.rate_limiter,
958            &cmd_tx,
959            methods::PRIVATE_CANCEL_TRIGGER_ORDER,
960            params,
961            self.request_timeout,
962        )
963        .await
964    }
965
966    /// Returns currently untriggered trigger orders via
967    /// `private/get_trigger_orders`.
968    ///
969    /// # Errors
970    ///
971    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
972    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
973    /// outcome is ambiguous.
974    pub async fn get_trigger_orders(
975        &self,
976        params: &DeriveGetTriggerOrdersParams,
977    ) -> Result<DeriveOpenOrdersResult> {
978        let cmd_tx = self.cmd_tx.read().await.clone();
979        send_request_typed(
980            &self.rate_limiter,
981            &cmd_tx,
982            methods::PRIVATE_GET_TRIGGER_ORDERS,
983            params,
984            self.request_timeout,
985        )
986        .await
987    }
988
989    /// Cancels every open order on the subaccount (the venue's
990    /// `private/cancel_all`), optionally scoped to an instrument.
991    ///
992    /// # Errors
993    ///
994    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
995    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
996    /// outcome is ambiguous.
997    pub async fn cancel_all_orders(&self, params: &DeriveCancelAllParams) -> Result<()> {
998        let cmd_tx = self.cmd_tx.read().await.clone();
999        let _: DeriveEmptyResult = send_request(
1000            &self.rate_limiter,
1001            &cmd_tx,
1002            methods::PRIVATE_CANCEL_ALL,
1003            params,
1004            self.request_timeout,
1005        )
1006        .await?;
1007        Ok(())
1008    }
1009}
1010
1011// Awaits the venue's raw `result`, bounded by `timeout`. A dropped responder
1012// (handler torn down on reconnect) surfaces as `RequestCancelled`, a timeout as
1013// `Timeout`; both leave a state-changing write's outcome ambiguous.
1014async fn send_raw<P>(
1015    rate_limiter: &WsRateLimiter,
1016    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1017    method: &'static str,
1018    params: &P,
1019    timeout: Duration,
1020) -> Result<Value>
1021where
1022    P: Serialize + ?Sized,
1023{
1024    let params = serde_json::to_value(params)?;
1025
1026    // Pace in the caller's task before enqueueing, so the shared feed handler
1027    // never sleeps mid-loop. Matching methods (order/cancel/replace) draw on the
1028    // matching quota; everything else on the non-matching default.
1029    //
1030    // Known limitation: matching writes arrive here already signed with a
1031    // `signature_expiry_sec`, so a long pace eats into that TTL. It only bites
1032    // under a pathological backlog (wait beyond the venue's ~300s TTL margin,
1033    // i.e. hundreds of orders queued at the 1/s Trader rate); moderate bursts
1034    // stay well inside it. Pacing above the signing layer would remove it and
1035    // is left as a follow-up.
1036    let rate_keys = [rate_limit_key_for(method)];
1037    rate_limiter.await_keys_ready(Some(&rate_keys)).await;
1038
1039    let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1040    cmd_tx
1041        .send(HandlerCommand::Request {
1042            method,
1043            params,
1044            response_tx,
1045        })
1046        .map_err(|e| DeriveWsError::transport(format!("failed to enqueue `{method}`: {e}")))?;
1047
1048    // On timeout the handler's `pending` entry leaks until the next reconnect's
1049    // `fail_pending` drains it; the later send to the dropped receiver is a
1050    // no-op logged at debug.
1051    match tokio::time::timeout(timeout, response_rx).await {
1052        Ok(Ok(outcome)) => outcome,
1053        Ok(Err(_)) => Err(DeriveWsError::RequestCancelled {
1054            method: method.to_owned(),
1055        }),
1056        Err(_) => Err(DeriveWsError::Timeout {
1057            method: method.to_owned(),
1058        }),
1059    }
1060}
1061
1062// Decodes the result, treating a null/absent `result` as `R::default()` (for
1063// login/subscribe/unsubscribe and the cancel family's `DeriveEmptyResult`).
1064async fn send_request<P, R>(
1065    rate_limiter: &WsRateLimiter,
1066    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1067    method: &'static str,
1068    params: &P,
1069    timeout: Duration,
1070) -> Result<R>
1071where
1072    P: Serialize + ?Sized,
1073    R: Default + DeserializeOwned,
1074{
1075    let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1076    let typed = if value.is_null() {
1077        R::default()
1078    } else {
1079        serde_json::from_value(value)?
1080    };
1081    Ok(typed)
1082}
1083
1084// Decodes the result with no `Default` fallback, for `private/order` and
1085// `private/replace` whose success result is always a populated object.
1086async fn send_request_typed<P, R>(
1087    rate_limiter: &WsRateLimiter,
1088    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1089    method: &'static str,
1090    params: &P,
1091    timeout: Duration,
1092) -> Result<R>
1093where
1094    P: Serialize + ?Sized,
1095    R: DeserializeOwned,
1096{
1097    let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1098    Ok(serde_json::from_value(value)?)
1099}
1100
1101fn channel_topics(channels: &[DeriveWsChannel]) -> Vec<String> {
1102    channels.iter().map(ToString::to_string).collect()
1103}
1104
1105async fn login_via_handler(
1106    rate_limiter: &WsRateLimiter,
1107    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1108    auth_tracker: &AuthTracker,
1109    creds: &DeriveWsCredentials,
1110    timeout: Duration,
1111) -> Result<()> {
1112    let login = build_ws_login(&creds.wallet_address, &creds.signer)?;
1113    let params = WsLoginParams {
1114        wallet: login.wallet,
1115        timestamp: login.timestamp,
1116        signature: login.signature,
1117    };
1118    let _receiver = auth_tracker.begin();
1119
1120    match send_request::<_, WsLoginResult>(
1121        rate_limiter,
1122        cmd_tx,
1123        methods::PUBLIC_LOGIN,
1124        &params,
1125        timeout,
1126    )
1127    .await
1128    {
1129        Ok(_) => {
1130            auth_tracker.succeed();
1131            log::debug!("Derive WebSocket authenticated");
1132            Ok(())
1133        }
1134        Err(e) => {
1135            auth_tracker.fail(e.to_string());
1136            Err(e)
1137        }
1138    }
1139}
1140
1141async fn subscribe_via_handler(
1142    rate_limiter: &WsRateLimiter,
1143    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1144    channels: Vec<String>,
1145    timeout: Duration,
1146) -> Result<()> {
1147    let params = WsSubscribeParams {
1148        channels: channels.into_iter().map(DeriveWsChannel::from).collect(),
1149    };
1150    let _: WsSubscribeResult = send_request(
1151        rate_limiter,
1152        cmd_tx,
1153        methods::PUBLIC_SUBSCRIBE,
1154        &params,
1155        timeout,
1156    )
1157    .await?;
1158    Ok(())
1159}
1160
1161#[cfg(test)]
1162mod tests {
1163    use rstest::rstest;
1164
1165    use super::*;
1166
1167    #[rstest]
1168    fn test_public_client_defaults_to_environment_url() {
1169        let client = DeriveWebSocketClient::new(
1170            None,
1171            DeriveEnvironment::Mainnet,
1172            TransportBackend::default(),
1173            None,
1174        );
1175        assert!(client.url().starts_with("wss://"));
1176        assert!(client.url().contains("api.lyra.finance"));
1177        assert!(!client.is_authenticated());
1178        assert!(!client.is_active());
1179        assert_eq!(client.subscription_count(), 0);
1180    }
1181
1182    #[rstest]
1183    fn test_testnet_client_routes_to_demo_url() {
1184        let client = DeriveWebSocketClient::new(
1185            None,
1186            DeriveEnvironment::Testnet,
1187            TransportBackend::default(),
1188            None,
1189        );
1190        assert!(client.url().contains("demo"));
1191    }
1192
1193    #[rstest]
1194    fn test_credentials_constructor_parses_session_key() {
1195        let creds = DeriveWsCredentials::new(
1196            "0x000000000000000000000000000000000000aaaa",
1197            "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
1198        )
1199        .unwrap();
1200        assert!(creds.wallet_address.starts_with("0x"));
1201        let client = DeriveWebSocketClient::with_credentials(
1202            None,
1203            DeriveEnvironment::Testnet,
1204            TransportBackend::default(),
1205            None,
1206            creds,
1207            None,
1208        );
1209        assert!(client.url().contains("demo"));
1210        assert!(!client.is_authenticated());
1211    }
1212
1213    #[rstest]
1214    fn test_credentials_debug_redacts_signer() {
1215        let creds = DeriveWsCredentials::new(
1216            "0xWALLET",
1217            "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
1218        )
1219        .unwrap();
1220        let debug = format!("{creds:?}");
1221        assert!(debug.contains("redacted"));
1222        assert!(debug.contains("0xWALLET"));
1223        assert!(!debug.contains("2ae8be44"));
1224    }
1225
1226    #[rstest]
1227    fn test_credentials_constructor_rejects_invalid_session_key() {
1228        let err = DeriveWsCredentials::new("0xWALLET", "not-a-hex-key").unwrap_err();
1229        assert!(err.to_string().contains("invalid session key"));
1230    }
1231
1232    #[rstest]
1233    #[tokio::test]
1234    async fn test_send_raw_times_out_when_no_response_arrives() {
1235        // Keep the receiver alive so the request enqueues, but never reply: the
1236        // bounded await must surface a Timeout rather than hang forever.
1237        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1238        let rate_limiter: WsRateLimiter = RateLimiter::new_with_quota(None, Vec::new());
1239        let err = send_raw(
1240            &rate_limiter,
1241            &cmd_tx,
1242            methods::PRIVATE_ORDER,
1243            &serde_json::json!({}),
1244            Duration::from_millis(50),
1245        )
1246        .await
1247        .expect_err("must time out");
1248
1249        match err {
1250            DeriveWsError::Timeout { method } => assert_eq!(method, methods::PRIVATE_ORDER),
1251            other => panic!("expected Timeout, was {other:?}"),
1252        }
1253    }
1254
1255    #[rstest]
1256    #[tokio::test]
1257    async fn test_send_request_typed_rejects_null_result() {
1258        // `private/order` and `private/replace` always return a populated
1259        // object on success; a null result is a protocol violation that must
1260        // surface as a serde error (classified ambiguous by the exec client).
1261        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
1262        tokio::spawn(async move {
1263            if let Some(HandlerCommand::Request { response_tx, .. }) = cmd_rx.recv().await {
1264                let _ = response_tx.send(Ok(Value::Null));
1265            }
1266        });
1267        let rate_limiter: WsRateLimiter = RateLimiter::new_with_quota(None, Vec::new());
1268        let result: Result<DeriveOrderResult> = send_request_typed(
1269            &rate_limiter,
1270            &cmd_tx,
1271            methods::PRIVATE_ORDER,
1272            &serde_json::json!({}),
1273            Duration::from_secs(1),
1274        )
1275        .await;
1276        assert!(matches!(result, Err(DeriveWsError::Serde(_))));
1277    }
1278}
nautilus_derive/websocket/client.rs

nautilus_derive/websocket/
client.rs
