Skip to main content

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;
36#[cfg(test)]
37use nautilus_common::live::get_runtime;
38use nautilus_core::{
39    UUID4,
40    string::secret::{REDACTED, SecretString},
41};
42use nautilus_live::{
43    SocketControl,
44    task::{SharedTaskSlot, TaskJoinOutcome, TaskSlot, finish_task},
45};
46use nautilus_network::{
47    http::create_standard_nautilus_headers,
48    mode::ConnectionMode,
49    ratelimiter::clock::MonotonicClock,
50    websocket::{
51        AuthTracker, TransportBackend, WebSocketClient, WebSocketConfig, channel_message_handler,
52    },
53};
54use serde::{Serialize, de::DeserializeOwned};
55use serde_json::Value;
56use ustr::Ustr;
57use zeroize::Zeroizing;
58
59use super::{
60    error::{DeriveWsError, Result},
61    handler::{
62        DeriveWsMessage, FeedHandler, HandlerCommand, orderbook_subscribe_params,
63        ticker_subscribe_params, trades_subscribe_params,
64    },
65    messages::{
66        DeriveWsChannel, WsLoginParams, WsLoginResult, WsSubscribeParams, WsSubscribeResult,
67        WsUnsubscribeParams, WsUnsubscribeResult, methods, orderbook_channel, ticker_channel,
68        trades_channel,
69    },
70};
71use crate::{
72    common::{
73        consts::{
74            RECONNECT_BACKOFF_FACTOR, RECONNECT_BASE_BACKOFF, RECONNECT_JITTER_MS,
75            RECONNECT_MAX_BACKOFF, RECONNECT_TIMEOUT, WS_HEARTBEAT_SECS, WS_HEARTBEAT_TIMEOUT,
76            WS_REQUEST_TIMEOUT,
77        },
78        enums::DeriveEnvironment,
79        rate_limit::{
80            DeriveRateLimiter, FixedWindowLimiter, FixedWindowLimits, RateClass,
81            rate_class_for_method,
82        },
83        urls,
84    },
85    http::{
86        models::{
87            DeriveCancelByInstrumentResult, DeriveCancelByLabelResult, DeriveEmptyResult,
88            DeriveOpenOrdersResult, DeriveOrder, DeriveOrderResult, DeriveReplaceOutcome,
89            DeriveReplaceResult,
90        },
91        query::{
92            DeriveCancelAllParams, DeriveCancelByInstrumentParams, DeriveCancelByLabelParams,
93            DeriveCancelParams, DeriveCancelTriggerOrderParams, DeriveGetTriggerOrdersParams,
94            DeriveOrderParams, DeriveReplaceParams, DeriveTriggerOrderParams,
95        },
96    },
97    signing::auth::build_ws_login,
98};
99
100/// Credentials for `public/login`. The session-key signer never escapes the
101/// client; only the wallet address is exposed via [`Debug`].
102#[derive(Clone)]
103pub struct DeriveWsCredentials {
104    /// Derive Chain smart-contract wallet address (`0x`-prefixed, 42 chars).
105    pub wallet_address: String,
106    /// secp256k1 session-key signer.
107    pub signer: PrivateKeySigner,
108}
109
110impl DeriveWsCredentials {
111    /// Constructs credentials by parsing `session_key_hex` into a signer.
112    ///
113    /// # Errors
114    ///
115    /// Returns [`DeriveWsError::Transport`] when the session-key hex cannot be parsed.
116    pub fn new(wallet_address: impl Into<String>, session_key_hex: &str) -> Result<Self> {
117        let signer: PrivateKeySigner = session_key_hex
118            .parse()
119            .map_err(|e| DeriveWsError::transport(format!("invalid session key: {e}")))?;
120        Ok(Self {
121            wallet_address: wallet_address.into(),
122            signer,
123        })
124    }
125}
126
127impl Debug for DeriveWsCredentials {
128    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129        f.debug_struct(stringify!(DeriveWsCredentials))
130            .field("wallet_address", &self.wallet_address)
131            .field("signer", &REDACTED)
132            .finish()
133    }
134}
135
136// Fixed-window rate limiter shared with the command handles so each frame is
137// paced in the caller's task before it is enqueued for the feed handler.
138type WsRateLimiter = DeriveRateLimiter;
139
140const MAX_SESSION_RECOVERY_ATTEMPTS: u32 = 3;
141const SUBSCRIPTION_ACCEPTED_STATUSES: &[&str] = &["ok"];
142const SUBSCRIPTION_REPLAY_ACCEPTED_STATUSES: &[&str] = &["ok", "already subscribed"];
143pub(super) const UNAUTHENTICATED_CONNECTION_EPOCH: u64 = u64::MAX;
144
145/// WebSocket client for the Derive JSON-RPC stream.
146///
147/// Construct with [`Self::new`] (public-only) or [`Self::with_credentials`]
148/// when private channels and signed actions are needed. Call [`Self::connect`]
149/// before any subscribe call; [`Self::disconnect`] tears the connection down.
150#[derive(Debug)]
151pub struct DeriveWebSocketClient {
152    url: String,
153    transport_backend: TransportBackend,
154    proxy_url: Option<SecretString>,
155    connection_mode: Arc<ArcSwap<AtomicU8>>,
156    connection_epoch: Arc<ArcSwap<AtomicU64>>,
157    signal: Arc<AtomicBool>,
158    auth_tracker: AuthTracker,
159    authenticated_epoch: Arc<AtomicU64>,
160    credentials: Option<DeriveWsCredentials>,
161    next_id: Arc<AtomicU64>,
162    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
163    out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>>,
164    subscriptions: Arc<DashMap<String, ()>>,
165    subscription_lock: Arc<tokio::sync::Mutex<()>>,
166    task_handle: TaskSlot<()>,
167    send_task: Arc<SharedTaskSlot<()>>,
168    shutdown_errors: Vec<String>,
169    request_timeout: Duration,
170    conn_id: Arc<ArcSwap<String>>,
171    rate_limiter: Arc<WsRateLimiter>,
172    socket_control: Option<SocketControl>,
173}
174
175#[derive(Clone, Debug)]
176pub(crate) struct DeriveWebSocketShutdownHandle {
177    signal: Arc<AtomicBool>,
178}
179
180impl DeriveWebSocketShutdownHandle {
181    pub(crate) fn begin_shutdown(&self) {
182        self.signal.store(true, Ordering::Release);
183    }
184}
185
186struct DeriveWebSocketSetupGuard {
187    shutdown: DeriveWebSocketShutdownHandle,
188    armed: bool,
189}
190
191impl DeriveWebSocketSetupGuard {
192    fn new(shutdown: DeriveWebSocketShutdownHandle) -> Self {
193        Self {
194            shutdown,
195            armed: true,
196        }
197    }
198
199    fn disarm(mut self) {
200        self.armed = false;
201    }
202}
203
204impl Drop for DeriveWebSocketSetupGuard {
205    fn drop(&mut self) {
206        if self.armed {
207            self.shutdown.begin_shutdown();
208        }
209    }
210}
211
212/// Cloneable command handle for Derive public market data subscriptions.
213#[derive(Debug, Clone)]
214pub struct DeriveWebSocketSubscriptionHandle {
215    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
216    subscriptions: Arc<DashMap<String, ()>>,
217    subscription_lock: Arc<tokio::sync::Mutex<()>>,
218    request_timeout: Duration,
219    rate_limiter: Arc<WsRateLimiter>,
220}
221
222/// Cloneable handle for issuing signed `private/*` trading requests over the
223/// WebSocket transport.
224///
225/// Carries the same `cmd_tx` the owning [`DeriveWebSocketClient`] swaps on
226/// connect/reconnect, so a handle obtained at construction stays valid for the
227/// client's lifetime. The handle is transport-only: it sends the pre-signed
228/// body and surfaces the venue's JSON-RPC outcome. Session authorization is the
229/// client's responsibility (via `public/login`).
230#[derive(Debug, Clone)]
231pub struct DeriveWsExecutionHandle {
232    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
233    auth_tracker: AuthTracker,
234    request_timeout: Duration,
235    conn_id: Arc<ArcSwap<String>>,
236    rate_limiter: Arc<WsRateLimiter>,
237}
238
239#[derive(Debug)]
240pub(crate) struct MatchingRateLimitReservation {
241    method: &'static str,
242    instrument_name: Ustr,
243    window: u32,
244}
245
246impl DeriveWebSocketClient {
247    /// Builds a public-only client. URL falls back to the environment default
248    /// when `url` is `None`.
249    #[must_use]
250    pub fn new(
251        url: Option<String>,
252        environment: DeriveEnvironment,
253        transport_backend: TransportBackend,
254        proxy_url: Option<String>,
255    ) -> Self {
256        let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
257        Self::build(
258            url,
259            transport_backend,
260            proxy_url,
261            None,
262            FixedWindowLimits::websocket(None, None),
263        )
264    }
265
266    /// Builds a client that will issue `public/login` on connect and replay
267    /// it after each reconnect.
268    ///
269    /// `max_matching_requests_per_second` sets the account-wide matching
270    /// allowance for order writes and `max_per_instrument_matching_requests_per_second`
271    /// the independent per-instrument allowance; `None` applies the Trader-tier
272    /// default of each. See [`crate::common::rate_limit`].
273    #[must_use]
274    pub fn with_credentials(
275        url: Option<String>,
276        environment: DeriveEnvironment,
277        transport_backend: TransportBackend,
278        proxy_url: Option<String>,
279        credentials: DeriveWsCredentials,
280        max_matching_requests_per_second: Option<u32>,
281        max_per_instrument_matching_requests_per_second: Option<u32>,
282    ) -> Self {
283        let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
284        let limits = FixedWindowLimits::websocket(
285            max_matching_requests_per_second,
286            max_per_instrument_matching_requests_per_second,
287        );
288        Self::build(url, transport_backend, proxy_url, Some(credentials), limits)
289    }
290
291    fn build(
292        url: String,
293        transport_backend: TransportBackend,
294        proxy_url: Option<String>,
295        credentials: Option<DeriveWsCredentials>,
296        limits: FixedWindowLimits,
297    ) -> Self {
298        let connection_mode = Arc::new(ArcSwap::new(Arc::new(AtomicU8::new(
299            ConnectionMode::Closed as u8,
300        ))));
301        let connection_epoch = Arc::new(ArcSwap::new(Arc::new(AtomicU64::new(0))));
302
303        // Placeholder channel; replaced by connect() before commands are issued.
304        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
305
306        // Matching writes draw on the account-wide and per-instrument
307        // allowances; custom cancellation methods have their own windows and
308        // login, subscription, and reads use the non-matching allowance.
309        // Handles pace each frame in the caller's task before enqueueing, so
310        // the feed handler never sleeps.
311        let rate_limiter = Arc::new(FixedWindowLimiter::new(limits, MonotonicClock {}));
312        Self {
313            url,
314            transport_backend,
315            proxy_url: proxy_url.map(SecretString::from),
316            connection_mode,
317            connection_epoch,
318            signal: Arc::new(AtomicBool::new(false)),
319            auth_tracker: AuthTracker::new(),
320            authenticated_epoch: Arc::new(AtomicU64::new(UNAUTHENTICATED_CONNECTION_EPOCH)),
321            credentials,
322            next_id: Arc::new(AtomicU64::new(1)),
323            cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
324            out_rx: None,
325            subscriptions: Arc::new(DashMap::new()),
326            subscription_lock: Arc::new(tokio::sync::Mutex::new(())),
327            task_handle: TaskSlot::new(),
328            send_task: Arc::new(SharedTaskSlot::new()),
329            shutdown_errors: Vec::new(),
330            request_timeout: WS_REQUEST_TIMEOUT,
331            conn_id: Arc::new(ArcSwap::from_pointee(UUID4::new().to_string())),
332            rate_limiter,
333            socket_control: None,
334        }
335    }
336
337    /// Configures socket state reporting and reconnect control.
338    #[must_use]
339    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
340        self.socket_control = Some(control);
341        self
342    }
343
344    /// Returns the configured WebSocket URL.
345    #[must_use]
346    pub fn url(&self) -> &str {
347        &self.url
348    }
349
350    /// Sets the per-operation WebSocket timeout (login, subscribe, reads, writes).
351    ///
352    /// Must be called before `connect()`. Defaults to `WS_REQUEST_TIMEOUT`.
353    pub fn set_request_timeout(&mut self, timeout: Duration) {
354        self.request_timeout = timeout;
355    }
356
357    /// Returns `true` when credentials are configured and the venue has
358    /// confirmed the latest `public/login`. Cleared on reconnect.
359    #[must_use]
360    pub fn is_authenticated(&self) -> bool {
361        self.auth_tracker.is_authenticated()
362            && self.is_active()
363            && self.authenticated_epoch.load(Ordering::Acquire)
364                == self.connection_epoch.load().load(Ordering::Acquire)
365    }
366
367    /// Returns `true` while the underlying transport is in the active state.
368    #[must_use]
369    pub fn is_active(&self) -> bool {
370        self.connection_mode.load().load(Ordering::Relaxed) == ConnectionMode::Active as u8
371    }
372
373    /// Establishes the WebSocket connection and spawns the I/O handler task.
374    ///
375    /// When credentials are configured, issues `public/login` and awaits the
376    /// venue's acknowledgement before returning.
377    ///
378    /// # Errors
379    ///
380    /// Returns [`DeriveWsError::Transport`] for handshake failures and
381    /// propagates [`DeriveWsError::Auth`] / [`DeriveWsError::JsonRpc`] when
382    /// the login flow fails.
383    pub async fn connect(&mut self) -> Result<()> {
384        // Fast path requires authenticated session when creds are configured;
385        // otherwise fall through and rebuild so `Ok` always implies authenticated.
386        let auth_ok = self.credentials.is_none() || self.is_authenticated();
387        if self.is_active() && auth_ok && self.task_handle.is_some() {
388            log::warn!("Derive WebSocket already connected");
389            return Ok(());
390        }
391
392        // Tear down stale state so we don't orphan the old handler task on rebuild.
393        if self.task_handle.is_some() || !self.send_task.is_empty() {
394            log::debug!("Tearing down stale Derive WebSocket state before connect");
395            self.teardown().await?;
396        }
397
398        self.signal.store(false, Ordering::Release);
399        let setup_guard = DeriveWebSocketSetupGuard::new(self.shutdown_handle());
400
401        self.authenticated_epoch
402            .store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
403
404        let (message_handler, raw_rx) = channel_message_handler();
405        let headers = create_standard_nautilus_headers();
406
407        let cfg = WebSocketConfig {
408            url: self.url.clone(),
409            headers,
410            heartbeat_interval_secs: Some(WS_HEARTBEAT_SECS),
411            heartbeat_payload: None,
412            connect_timeout_ms: Some(RECONNECT_TIMEOUT.as_millis() as u64),
413            reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
414            reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
415            reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
416            reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
417            reconnect_max_attempts: None,
418            heartbeat_timeout_secs: Some(WS_HEARTBEAT_TIMEOUT.as_secs()),
419            idle_timeout_ms: None,
420            backend: self.transport_backend,
421            proxy_url: self
422                .proxy_url
423                .as_ref()
424                .map(|value| value.expose_secret().to_owned()),
425        };
426        // Rate limiting runs caller-side via `self.rate_limiter` before frames
427        // are enqueued, so the network client's own limiter is left unconfigured
428        // and never sleeps inside the single feed-handler task.
429        let client = WebSocketClient::builder()
430            .config(cfg)
431            .message_handler(message_handler)
432            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
433            .connect()
434            .await
435            .map_err(|e| DeriveWsError::transport(e.to_string()))?;
436
437        // Register the tracker so the network controller clears
438        // `is_authenticated()` on dead-socket detection, not just on the
439        // later RECONNECTED sentinel.
440        client.set_auth_tracker(self.auth_tracker.clone(), false);
441
442        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
443        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<DeriveWsMessage>();
444
445        *self.cmd_tx.write().await = cmd_tx.clone();
446        self.out_rx = Some(out_rx);
447        self.conn_id.store(Arc::new(UUID4::new().to_string()));
448
449        let connection_mode = client.connection_mode_atomic();
450        let connection_epoch = client.connection_epoch_atomic();
451        let reconnect_handle = client.reconnect_handle();
452        self.connection_mode.store(Arc::clone(&connection_mode));
453        self.connection_epoch.store(Arc::clone(&connection_epoch));
454        log::debug!("Derive WebSocket connected: {}", self.url);
455
456        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
457            return Err(DeriveWsError::transport(format!(
458                "failed to send SetClient command: {e}",
459            )));
460        }
461
462        let signal = Arc::clone(&self.signal);
463        let auth_tracker = self.auth_tracker.clone();
464        let authenticated_epoch = Arc::clone(&self.authenticated_epoch);
465        let next_id = Arc::clone(&self.next_id);
466        let credentials = self.credentials.clone();
467        let subscriptions = Arc::clone(&self.subscriptions);
468        let subscription_lock = Arc::clone(&self.subscription_lock);
469        let conn_id = Arc::clone(&self.conn_id);
470        let cmd_tx_for_loop = cmd_tx.clone();
471        let rate_limiter = Arc::clone(&self.rate_limiter);
472        let request_timeout = self.request_timeout;
473        let recovery_connection_mode = Arc::clone(&connection_mode);
474        let recovery_connection_epoch = Arc::clone(&connection_epoch);
475        let send_task = Arc::clone(&self.send_task);
476
477        if let Err(e) = self.task_handle.spawn(async move {
478            let (recovery_tx, mut recovery_rx) = tokio::sync::mpsc::unbounded_channel::<u64>();
479            let recovery_out_tx = out_tx.clone();
480            let recovery_cmd_tx = cmd_tx_for_loop.clone();
481            let recovery_auth_tracker = auth_tracker.clone();
482            let recovery_authenticated_epoch = Arc::clone(&authenticated_epoch);
483            let recovery_subscriptions = Arc::clone(&subscriptions);
484            let recovery_subscription_lock = Arc::clone(&subscription_lock);
485            let recovery_rate_limiter = Arc::clone(&rate_limiter);
486            let recovery_credentials = credentials.clone();
487            let recovery_mode = Arc::clone(&recovery_connection_mode);
488            let recovery_epoch = Arc::clone(&recovery_connection_epoch);
489            let mut recovery_tasks = tokio::task::JoinSet::new();
490
491            recovery_tasks.spawn(async move {
492                while let Some(mut requested_epoch) = recovery_rx.recv().await {
493                    while let Ok(epoch) = recovery_rx.try_recv() {
494                        requested_epoch = requested_epoch.max(epoch);
495                    }
496
497                    loop {
498                        match recover_session(
499                            &recovery_rate_limiter,
500                            &recovery_cmd_tx,
501                            &recovery_auth_tracker,
502                            &recovery_authenticated_epoch,
503                            &recovery_mode,
504                            &recovery_epoch,
505                            recovery_credentials.as_ref(),
506                            &recovery_subscriptions,
507                            &recovery_subscription_lock,
508                            request_timeout,
509                        )
510                        .await
511                        {
512                            Ok(recovered_epoch) => {
513                                while let Ok(epoch) = recovery_rx.try_recv() {
514                                    requested_epoch = requested_epoch.max(epoch);
515                                }
516
517                                if requested_epoch > recovered_epoch
518                                    || !session_connection_is_active(
519                                        &recovery_mode,
520                                        &recovery_epoch,
521                                        recovered_epoch,
522                                    )
523                                {
524                                    continue;
525                                }
526
527                                if recovery_out_tx.send(DeriveWsMessage::Reconnected).is_err() {
528                                    log::debug!("Derive outer receiver dropped during recovery");
529                                }
530                                break;
531                            }
532                            Err(e) => {
533                                let mut retry = false;
534
535                                while let Ok(epoch) = recovery_rx.try_recv() {
536                                    requested_epoch = requested_epoch.max(epoch);
537                                    retry = true;
538                                }
539
540                                if retry {
541                                    continue;
542                                }
543                                log::error!("Derive WebSocket session recovery failed: {e}");
544                                let _ = recovery_out_tx
545                                    .send(DeriveWsMessage::SessionRecoveryFailed(e.to_string()));
546                                let _ = recovery_cmd_tx.send(HandlerCommand::Disconnect);
547                                return;
548                            }
549                        }
550                    }
551                }
552            });
553
554            let mut handler = FeedHandler::new_with_send_task(
555                signal,
556                cmd_rx,
557                raw_rx,
558                next_id,
559                auth_tracker.clone(),
560                Arc::clone(&authenticated_epoch),
561                send_task,
562            );
563
564            loop {
565                match handler.next().await {
566                    Some(DeriveWsMessage::Reconnected) => {
567                        log::info!("Derive WebSocket re-establishing session after reconnect");
568                        conn_id.store(Arc::new(UUID4::new().to_string()));
569                        let epoch = recovery_connection_epoch.load(Ordering::Acquire);
570                        if recovery_tx.send(epoch).is_err() {
571                            log::error!("Derive WebSocket recovery task stopped unexpectedly");
572                            let _ = cmd_tx_for_loop.send(HandlerCommand::Disconnect);
573                        }
574                    }
575                    Some(msg) => {
576                        if out_tx.send(msg).is_err() {
577                            log::debug!("Derive outer receiver dropped, exiting stream loop");
578                            break;
579                        }
580                    }
581                    None => {
582                        log::debug!("Derive handler task ended");
583                        break;
584                    }
585                }
586            }
587        }) {
588            let shutdown_result = self.teardown().await;
589            return Err(DeriveWsError::transport(match shutdown_result {
590                Ok(()) => format!("failed to start WebSocket handler task: {e}"),
591                Err(shutdown_error) => format!(
592                    "failed to start WebSocket handler task: {e}; startup rollback failed: \
593                     {shutdown_error}"
594                ),
595            }));
596        }
597
598        if let Some(control) = &self.socket_control {
599            control.register(move || reconnect_handle.request_reconnect());
600        }
601
602        if let Some(creds) = self.credentials.clone()
603            && let Err(e) = login_via_handler(
604                &self.rate_limiter,
605                &cmd_tx,
606                &self.auth_tracker,
607                &self.authenticated_epoch,
608                &connection_mode,
609                &connection_epoch,
610                &creds,
611                self.request_timeout,
612            )
613            .await
614        {
615            // Without teardown, a retry connect() would short-circuit on
616            // is_active() and return Ok without a valid session.
617            log::warn!("Derive WebSocket login failed; tearing down transport: {e}");
618            self.teardown().await?;
619            return Err(e);
620        }
621
622        setup_guard.disarm();
623        Ok(())
624    }
625
626    pub(crate) fn begin_shutdown(&self) {
627        self.signal.store(true, Ordering::Release);
628    }
629
630    pub(crate) fn shutdown_handle(&self) -> DeriveWebSocketShutdownHandle {
631        DeriveWebSocketShutdownHandle {
632            signal: Arc::clone(&self.signal),
633        }
634    }
635
636    /// Signals the handler to disconnect, aborts the spawn task, and resets
637    /// the client's transport-related state. Shared by [`Self::disconnect`]
638    /// and the login-failure branch of [`Self::connect`].
639    async fn teardown(&mut self) -> Result<()> {
640        self.begin_shutdown();
641
642        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
643            log::debug!(
644                "Failed to enqueue Disconnect command (handler may already be shut down): {e}",
645            );
646        }
647
648        match finish_task(
649            &mut self.task_handle,
650            Duration::from_secs(2),
651            Duration::from_secs(2),
652        )
653        .await
654        {
655            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => {}
656            Some(TaskJoinOutcome::Failed(error)) => self
657                .shutdown_errors
658                .push(format!("WebSocket handler task failed: {error}")),
659            Some(TaskJoinOutcome::Incomplete) => self
660                .shutdown_errors
661                .push("WebSocket handler task did not stop after abort".to_string()),
662        }
663
664        if self.task_handle.is_some() {
665            if let Some(control) = &self.socket_control {
666                control.deregister();
667            }
668            return self.take_shutdown_result();
669        }
670
671        match self
672            .send_task
673            .finish(Duration::from_secs(1), Duration::from_secs(2))
674            .await
675        {
676            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => {}
677            Some(TaskJoinOutcome::Failed(error)) => self
678                .shutdown_errors
679                .push(format!("WebSocket send worker failed: {error}")),
680            Some(TaskJoinOutcome::Incomplete) => self
681                .shutdown_errors
682                .push("WebSocket send worker did not stop after abort".to_string()),
683        }
684
685        if !self.send_task.is_empty() {
686            if let Some(control) = &self.socket_control {
687                control.deregister();
688            }
689            return self.take_shutdown_result();
690        }
691
692        // Subscriptions are also dropped: the venue session ended with the
693        // transport, so a fresh connect() must re-issue them.
694        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
695        *self.cmd_tx.write().await = placeholder_tx;
696        self.out_rx = None;
697        self.connection_mode
698            .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
699        self.connection_epoch.store(Arc::new(AtomicU64::new(0)));
700        self.auth_tracker.invalidate();
701        self.authenticated_epoch
702            .store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
703        self.subscriptions.clear();
704        self.signal.store(false, Ordering::Relaxed);
705
706        if let Some(control) = &self.socket_control {
707            control.deregister();
708        }
709
710        self.take_shutdown_result()
711    }
712
713    fn take_shutdown_result(&mut self) -> Result<()> {
714        if self.shutdown_errors.is_empty() {
715            Ok(())
716        } else {
717            Err(DeriveWsError::transport(
718                std::mem::take(&mut self.shutdown_errors).join("; "),
719            ))
720        }
721    }
722
723    /// Disconnects the WebSocket connection and awaits the handler task.
724    ///
725    /// # Errors
726    ///
727    /// Returns [`DeriveWsError::Transport`] when the disconnect command
728    /// cannot be enqueued; the handler still tears down on signal.
729    pub async fn disconnect(&mut self) -> Result<()> {
730        log::debug!("Disconnecting Derive WebSocket");
731        self.begin_shutdown();
732        self.teardown().await
733    }
734
735    /// Subscribes to `ticker_slim.{instrument_name}.{interval}`. `interval` is the
736    /// millisecond cadence string the venue exposes (e.g. `"100"`, `"1000"`).
737    ///
738    /// # Errors
739    ///
740    /// Propagates JSON-RPC errors raised by the venue and transport-level
741    /// failures.
742    pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
743        self.subscription_handle()
744            .subscribe_ticker(instrument_name, interval)
745            .await
746    }
747
748    /// Unsubscribes from `ticker_slim.{instrument_name}.{interval}`.
749    ///
750    /// # Errors
751    ///
752    /// Propagates JSON-RPC errors raised by the venue and transport-level
753    /// failures.
754    pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
755        self.subscription_handle()
756            .unsubscribe_ticker(instrument_name, interval)
757            .await
758    }
759
760    /// Subscribes to `orderbook.{instrument_name}.{group}.{depth}`.
761    ///
762    /// # Errors
763    ///
764    /// Propagates JSON-RPC errors raised by the venue and transport-level
765    /// failures.
766    pub async fn subscribe_orderbook(
767        &self,
768        instrument_name: &str,
769        group: &str,
770        depth: &str,
771    ) -> Result<()> {
772        self.subscription_handle()
773            .subscribe_orderbook(instrument_name, group, depth)
774            .await
775    }
776
777    /// Unsubscribes from `orderbook.{instrument_name}.{group}.{depth}`.
778    ///
779    /// # Errors
780    ///
781    /// Propagates JSON-RPC errors raised by the venue and transport-level
782    /// failures.
783    pub async fn unsubscribe_orderbook(
784        &self,
785        instrument_name: &str,
786        group: &str,
787        depth: &str,
788    ) -> Result<()> {
789        self.subscription_handle()
790            .unsubscribe_orderbook(instrument_name, group, depth)
791            .await
792    }
793
794    /// Subscribes to `trades.{instrument_type}.{currency}`.
795    ///
796    /// # Errors
797    ///
798    /// Propagates JSON-RPC errors raised by the venue and transport-level
799    /// failures.
800    pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
801        self.subscription_handle()
802            .subscribe_trades(instrument_type, currency)
803            .await
804    }
805
806    /// Unsubscribes from `trades.{instrument_type}.{currency}`.
807    ///
808    /// # Errors
809    ///
810    /// Propagates JSON-RPC errors raised by the venue and transport-level
811    /// failures.
812    pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
813        self.subscription_handle()
814            .unsubscribe_trades(instrument_type, currency)
815            .await
816    }
817
818    /// Subscribes to a list of channel topics in a single `subscribe` frame.
819    ///
820    /// Used by the execution client to bulk-subscribe to the private
821    /// `{subaccount_id}.orders`, `{subaccount_id}.trades`, and
822    /// `{subaccount_id}.balances` channels after login.
823    ///
824    /// # Errors
825    ///
826    /// Propagates JSON-RPC errors raised by the venue and transport-level
827    /// failures.
828    pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
829    where
830        C: Into<DeriveWsChannel>,
831    {
832        self.subscription_handle()
833            .subscribe_channels(channels)
834            .await
835    }
836
837    /// Unsubscribes from a list of channel topics in a single
838    /// `unsubscribe` frame.
839    ///
840    /// # Errors
841    ///
842    /// Propagates JSON-RPC errors raised by the venue and transport-level
843    /// failures.
844    pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
845    where
846        C: Into<DeriveWsChannel>,
847    {
848        self.subscription_handle()
849            .unsubscribe_channels(channels)
850            .await
851    }
852
853    /// Returns the next event emitted by the handler.
854    pub async fn next_event(&mut self) -> Option<DeriveWsMessage> {
855        if let Some(rx) = self.out_rx.as_mut() {
856            rx.recv().await
857        } else {
858            None
859        }
860    }
861
862    /// Returns the count of channels the client currently has confirmed
863    /// subscriptions for.
864    #[must_use]
865    pub fn subscription_count(&self) -> usize {
866        self.subscriptions.len()
867    }
868
869    /// Returns a cloneable handle for issuing subscription commands.
870    #[must_use]
871    pub fn subscription_handle(&self) -> DeriveWebSocketSubscriptionHandle {
872        DeriveWebSocketSubscriptionHandle {
873            cmd_tx: Arc::clone(&self.cmd_tx),
874            subscriptions: Arc::clone(&self.subscriptions),
875            subscription_lock: Arc::clone(&self.subscription_lock),
876            request_timeout: self.request_timeout,
877            rate_limiter: Arc::clone(&self.rate_limiter),
878        }
879    }
880
881    /// Returns a cloneable handle for issuing signed `private/*` trading
882    /// requests.
883    ///
884    /// The handle shares the client's command channel, so it stays valid across
885    /// reconnects (the channel is swapped behind a shared lock). Obtain it once
886    /// and clone it into each order-submission task.
887    #[must_use]
888    pub fn execution_handle(&self) -> DeriveWsExecutionHandle {
889        DeriveWsExecutionHandle {
890            cmd_tx: Arc::clone(&self.cmd_tx),
891            auth_tracker: self.auth_tracker.clone(),
892            request_timeout: self.request_timeout,
893            conn_id: Arc::clone(&self.conn_id),
894            rate_limiter: Arc::clone(&self.rate_limiter),
895        }
896    }
897
898    /// Takes the event receiver from the client.
899    ///
900    /// This lets the live data client own the receive loop while subscription
901    /// commands continue through [`Self::subscription_handle`].
902    pub fn take_event_receiver(
903        &mut self,
904    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>> {
905        self.out_rx.take()
906    }
907}
908
909impl Drop for DeriveWebSocketClient {
910    fn drop(&mut self) {
911        self.signal.store(true, Ordering::Relaxed);
912
913        if let Some(handle) = self.task_handle.as_ref() {
914            handle.abort();
915        }
916
917        if let Some(control) = &self.socket_control {
918            control.deregister();
919        }
920    }
921}
922
923impl DeriveWebSocketSubscriptionHandle {
924    pub(crate) fn has_subscription(&self, channel: &str) -> bool {
925        self.subscriptions.contains_key(channel)
926    }
927
928    pub(crate) fn forget_subscription(&self, channel: &str) {
929        self.subscriptions.remove(channel);
930    }
931
932    pub(crate) fn remember_subscription(&self, channel: &str) {
933        self.subscriptions.insert(channel.to_string(), ());
934    }
935
936    /// Subscribes to `ticker_slim.{instrument_name}.{interval}`.
937    ///
938    /// # Errors
939    ///
940    /// Propagates JSON-RPC errors raised by the venue and transport-level
941    /// failures.
942    pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
943        let channel = ticker_channel(instrument_name, interval);
944        let params = ticker_subscribe_params(instrument_name, interval);
945        self.send_subscribe(channel, &params).await
946    }
947
948    /// Unsubscribes from `ticker_slim.{instrument_name}.{interval}`.
949    ///
950    /// # Errors
951    ///
952    /// Propagates JSON-RPC errors raised by the venue and transport-level
953    /// failures.
954    pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
955        let channel = ticker_channel(instrument_name, interval);
956        self.send_unsubscribe(channel).await
957    }
958
959    /// Subscribes to `orderbook.{instrument_name}.{group}.{depth}`.
960    ///
961    /// # Errors
962    ///
963    /// Propagates JSON-RPC errors raised by the venue and transport-level
964    /// failures.
965    pub async fn subscribe_orderbook(
966        &self,
967        instrument_name: &str,
968        group: &str,
969        depth: &str,
970    ) -> Result<()> {
971        let channel = orderbook_channel(instrument_name, group, depth);
972        let params = orderbook_subscribe_params(instrument_name, group, depth);
973        self.send_subscribe(channel, &params).await
974    }
975
976    /// Unsubscribes from `orderbook.{instrument_name}.{group}.{depth}`.
977    ///
978    /// # Errors
979    ///
980    /// Propagates JSON-RPC errors raised by the venue and transport-level
981    /// failures.
982    pub async fn unsubscribe_orderbook(
983        &self,
984        instrument_name: &str,
985        group: &str,
986        depth: &str,
987    ) -> Result<()> {
988        let channel = orderbook_channel(instrument_name, group, depth);
989        self.send_unsubscribe(channel).await
990    }
991
992    /// Subscribes to `trades.{instrument_type}.{currency}`.
993    ///
994    /// # Errors
995    ///
996    /// Propagates JSON-RPC errors raised by the venue and transport-level
997    /// failures.
998    pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
999        let channel = trades_channel(instrument_type, currency);
1000        let params = trades_subscribe_params(instrument_type, currency);
1001        self.send_subscribe(channel, &params).await
1002    }
1003
1004    /// Unsubscribes from `trades.{instrument_type}.{currency}`.
1005    ///
1006    /// # Errors
1007    ///
1008    /// Propagates JSON-RPC errors raised by the venue and transport-level
1009    /// failures.
1010    pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
1011        let channel = trades_channel(instrument_type, currency);
1012        self.send_unsubscribe(channel).await
1013    }
1014
1015    /// Subscribes to multiple channel topics in a single `subscribe` frame.
1016    ///
1017    /// # Errors
1018    ///
1019    /// Propagates JSON-RPC errors raised by the venue and transport-level
1020    /// failures.
1021    pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
1022    where
1023        C: Into<DeriveWsChannel>,
1024    {
1025        let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
1026        if channels.is_empty() {
1027            return Ok(());
1028        }
1029        let _guard = self.subscription_lock.lock().await;
1030        let params = WsSubscribeParams { channels };
1031        let cmd_tx = self.cmd_tx.read().await.clone();
1032        let result: WsSubscribeResult = send_request(
1033            &self.rate_limiter,
1034            &cmd_tx,
1035            methods::PUBLIC_SUBSCRIBE,
1036            &params,
1037            self.request_timeout,
1038        )
1039        .await?;
1040
1041        let (confirmed, failure) =
1042            subscription_outcome(&params.channels, &result, SUBSCRIPTION_ACCEPTED_STATUSES);
1043
1044        for channel in confirmed {
1045            self.subscriptions.insert(channel, ());
1046        }
1047        failure.map_or(Ok(()), Err)
1048    }
1049
1050    /// Unsubscribes from multiple channel topics in a single
1051    /// `unsubscribe` frame.
1052    ///
1053    /// # Errors
1054    ///
1055    /// Propagates JSON-RPC errors raised by the venue and transport-level
1056    /// failures.
1057    pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
1058    where
1059        C: Into<DeriveWsChannel>,
1060    {
1061        let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
1062        if channels.is_empty() {
1063            return Ok(());
1064        }
1065        let _guard = self.subscription_lock.lock().await;
1066        let topics = channel_topics(&channels);
1067        let params = WsUnsubscribeParams { channels };
1068        let cmd_tx = self.cmd_tx.read().await.clone();
1069        let _: WsUnsubscribeResult = send_request(
1070            &self.rate_limiter,
1071            &cmd_tx,
1072            methods::PUBLIC_UNSUBSCRIBE,
1073            &params,
1074            self.request_timeout,
1075        )
1076        .await?;
1077
1078        for channel in topics {
1079            self.subscriptions.remove(&channel);
1080        }
1081
1082        Ok(())
1083    }
1084
1085    async fn send_subscribe(&self, channel: String, params: &WsSubscribeParams) -> Result<()> {
1086        let _guard = self.subscription_lock.lock().await;
1087        let cmd_tx = self.cmd_tx.read().await.clone();
1088        let result: WsSubscribeResult = send_request(
1089            &self.rate_limiter,
1090            &cmd_tx,
1091            methods::PUBLIC_SUBSCRIBE,
1092            params,
1093            self.request_timeout,
1094        )
1095        .await?;
1096
1097        let (confirmed, failure) =
1098            subscription_outcome(&params.channels, &result, SUBSCRIPTION_ACCEPTED_STATUSES);
1099
1100        if confirmed.iter().any(|topic| topic == &channel) {
1101            self.subscriptions.insert(channel, ());
1102        }
1103        failure.map_or(Ok(()), Err)
1104    }
1105
1106    async fn send_unsubscribe(&self, channel: String) -> Result<()> {
1107        let _guard = self.subscription_lock.lock().await;
1108        let params = WsUnsubscribeParams {
1109            channels: vec![DeriveWsChannel::from(channel.clone())],
1110        };
1111        let cmd_tx = self.cmd_tx.read().await.clone();
1112        let _: WsUnsubscribeResult = send_request(
1113            &self.rate_limiter,
1114            &cmd_tx,
1115            methods::PUBLIC_UNSUBSCRIBE,
1116            &params,
1117            self.request_timeout,
1118        )
1119        .await?;
1120
1121        self.subscriptions.remove(&channel);
1122
1123        Ok(())
1124    }
1125}
1126
1127impl DeriveWsExecutionHandle {
1128    /// Returns the current WebSocket connection id used by trigger orders.
1129    #[must_use]
1130    pub fn conn_id(&self) -> String {
1131        self.conn_id.load_full().as_ref().clone()
1132    }
1133
1134    /// Submits a signed order via `private/order`.
1135    ///
1136    /// `params` must be the fully-built signed body from
1137    /// [`crate::http::query::order_to_derive_payload`]. Returns the accepted
1138    /// order echoed by the venue.
1139    ///
1140    /// # Errors
1141    ///
1142    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1143    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1144    /// outcome is ambiguous.
1145    pub async fn submit_order(&self, params: &DeriveOrderParams) -> Result<DeriveOrder> {
1146        let reservation = self
1147            .reserve_matching_request(methods::PRIVATE_ORDER, &params.instrument_name)
1148            .await?;
1149        self.submit_order_after_rate_limit(params, reservation)
1150            .await
1151    }
1152
1153    pub(crate) async fn submit_order_after_rate_limit(
1154        &self,
1155        params: &DeriveOrderParams,
1156        reservation: MatchingRateLimitReservation,
1157    ) -> Result<DeriveOrder> {
1158        self.ensure_authenticated(methods::PRIVATE_ORDER)?;
1159        debug_assert_eq!(reservation.method, methods::PRIVATE_ORDER);
1160        self.refresh_matching_reservation(&reservation).await;
1161        let cmd_tx = self.cmd_tx.read().await.clone();
1162        let result: DeriveOrderResult = send_request_typed_after_rate_limit(
1163            &self.rate_limiter,
1164            &cmd_tx,
1165            methods::PRIVATE_ORDER,
1166            params,
1167            self.request_timeout,
1168        )
1169        .await?;
1170        Ok(result.order)
1171    }
1172
1173    /// Submits a signed trigger order via `private/trigger_order`.
1174    ///
1175    /// # Errors
1176    ///
1177    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1178    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1179    /// outcome is ambiguous.
1180    pub async fn submit_trigger_order(
1181        &self,
1182        params: &DeriveTriggerOrderParams,
1183    ) -> Result<DeriveOrder> {
1184        let reservation = self
1185            .reserve_matching_request(
1186                methods::PRIVATE_TRIGGER_ORDER,
1187                &params.order.instrument_name,
1188            )
1189            .await?;
1190        self.submit_trigger_order_after_rate_limit(params, reservation)
1191            .await
1192    }
1193
1194    pub(crate) async fn submit_trigger_order_after_rate_limit(
1195        &self,
1196        params: &DeriveTriggerOrderParams,
1197        reservation: MatchingRateLimitReservation,
1198    ) -> Result<DeriveOrder> {
1199        self.ensure_authenticated(methods::PRIVATE_TRIGGER_ORDER)?;
1200        debug_assert_eq!(reservation.method, methods::PRIVATE_TRIGGER_ORDER);
1201        self.refresh_matching_reservation(&reservation).await;
1202        let cmd_tx = self.cmd_tx.read().await.clone();
1203        let result: DeriveOrderResult = send_request_typed_after_rate_limit(
1204            &self.rate_limiter,
1205            &cmd_tx,
1206            methods::PRIVATE_TRIGGER_ORDER,
1207            params,
1208            self.request_timeout,
1209        )
1210        .await?;
1211        Ok(result.order)
1212    }
1213
1214    /// Modifies a working order by cancelling it and submitting a replacement
1215    /// through the venue's `private/replace`.
1216    ///
1217    /// # Errors
1218    ///
1219    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1220    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1221    /// outcome is ambiguous.
1222    pub async fn modify_order(&self, params: &DeriveReplaceParams) -> Result<DeriveReplaceOutcome> {
1223        let reservation = self
1224            .reserve_matching_request(methods::PRIVATE_REPLACE, &params.order.instrument_name)
1225            .await?;
1226        self.modify_order_after_rate_limit(params, reservation)
1227            .await
1228    }
1229
1230    pub(crate) async fn modify_order_after_rate_limit(
1231        &self,
1232        params: &DeriveReplaceParams,
1233        reservation: MatchingRateLimitReservation,
1234    ) -> Result<DeriveReplaceOutcome> {
1235        self.ensure_authenticated(methods::PRIVATE_REPLACE)?;
1236        debug_assert_eq!(reservation.method, methods::PRIVATE_REPLACE);
1237        self.refresh_matching_reservation(&reservation).await;
1238        let cmd_tx = self.cmd_tx.read().await.clone();
1239        let result: DeriveReplaceResult = send_request_typed_after_rate_limit(
1240            &self.rate_limiter,
1241            &cmd_tx,
1242            methods::PRIVATE_REPLACE,
1243            params,
1244            self.request_timeout,
1245        )
1246        .await?;
1247        result
1248            .into_outcome(&params.order_id_to_cancel, &params.order.label)
1249            .map_err(|message| {
1250                DeriveWsError::Serde(<serde_json::Error as serde::de::Error>::custom(message))
1251            })
1252    }
1253
1254    /// Cancels a single order via `private/cancel`.
1255    ///
1256    /// # Errors
1257    ///
1258    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1259    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1260    /// outcome is ambiguous.
1261    pub async fn cancel_order(&self, params: &DeriveCancelParams) -> Result<()> {
1262        self.require_authenticated(methods::PRIVATE_CANCEL).await?;
1263        let cmd_tx = self.cmd_tx.read().await.clone();
1264        let _: DeriveEmptyResult = send_request_for_instrument(
1265            &self.rate_limiter,
1266            &cmd_tx,
1267            methods::PRIVATE_CANCEL,
1268            params,
1269            self.request_timeout,
1270            params.instrument_name,
1271        )
1272        .await?;
1273        Ok(())
1274    }
1275
1276    /// Cancels every open order for one instrument via `private/cancel_by_instrument`.
1277    ///
1278    /// # Errors
1279    ///
1280    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1281    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1282    /// outcome is ambiguous.
1283    pub async fn cancel_by_instrument(
1284        &self,
1285        params: &DeriveCancelByInstrumentParams,
1286    ) -> Result<DeriveCancelByInstrumentResult> {
1287        self.require_authenticated(methods::PRIVATE_CANCEL_BY_INSTRUMENT)
1288            .await?;
1289        let cmd_tx = self.cmd_tx.read().await.clone();
1290        send_request_typed_for_instrument(
1291            &self.rate_limiter,
1292            &cmd_tx,
1293            methods::PRIVATE_CANCEL_BY_INSTRUMENT,
1294            params,
1295            self.request_timeout,
1296            params.instrument_name,
1297        )
1298        .await
1299    }
1300
1301    /// Cancels a single trigger order via `private/cancel_trigger_order`.
1302    ///
1303    /// # Errors
1304    ///
1305    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1306    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1307    /// outcome is ambiguous.
1308    pub async fn cancel_trigger_order(
1309        &self,
1310        params: &DeriveCancelTriggerOrderParams,
1311    ) -> Result<DeriveOrder> {
1312        self.require_authenticated(methods::PRIVATE_CANCEL_TRIGGER_ORDER)
1313            .await?;
1314        let cmd_tx = self.cmd_tx.read().await.clone();
1315        send_request_typed(
1316            &self.rate_limiter,
1317            &cmd_tx,
1318            methods::PRIVATE_CANCEL_TRIGGER_ORDER,
1319            params,
1320            self.request_timeout,
1321        )
1322        .await
1323    }
1324
1325    /// Cancels every open order with the given label via
1326    /// `private/cancel_by_label`.
1327    ///
1328    /// # Errors
1329    ///
1330    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1331    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1332    /// outcome is ambiguous.
1333    pub async fn cancel_by_label(
1334        &self,
1335        params: &DeriveCancelByLabelParams,
1336    ) -> Result<DeriveCancelByLabelResult> {
1337        self.require_authenticated(methods::PRIVATE_CANCEL_BY_LABEL)
1338            .await?;
1339        let cmd_tx = self.cmd_tx.read().await.clone();
1340        send_request_typed(
1341            &self.rate_limiter,
1342            &cmd_tx,
1343            methods::PRIVATE_CANCEL_BY_LABEL,
1344            params,
1345            self.request_timeout,
1346        )
1347        .await
1348    }
1349
1350    /// Returns currently untriggered trigger orders via
1351    /// `private/get_trigger_orders`.
1352    ///
1353    /// # Errors
1354    ///
1355    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1356    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1357    /// outcome is ambiguous.
1358    pub async fn get_trigger_orders(
1359        &self,
1360        params: &DeriveGetTriggerOrdersParams,
1361    ) -> Result<DeriveOpenOrdersResult> {
1362        self.require_authenticated(methods::PRIVATE_GET_TRIGGER_ORDERS)
1363            .await?;
1364        let cmd_tx = self.cmd_tx.read().await.clone();
1365        send_request_typed(
1366            &self.rate_limiter,
1367            &cmd_tx,
1368            methods::PRIVATE_GET_TRIGGER_ORDERS,
1369            params,
1370            self.request_timeout,
1371        )
1372        .await
1373    }
1374
1375    /// Cancels every open order on the subaccount (the venue's
1376    /// `private/cancel_all`), optionally scoped to an instrument.
1377    ///
1378    /// # Errors
1379    ///
1380    /// Returns [`DeriveWsError::JsonRpc`] for venue rejections and
1381    /// [`DeriveWsError::Transport`] / [`DeriveWsError::Timeout`] when the
1382    /// outcome is ambiguous.
1383    pub async fn cancel_all_orders(&self, params: &DeriveCancelAllParams) -> Result<()> {
1384        self.require_authenticated(methods::PRIVATE_CANCEL_ALL)
1385            .await?;
1386        let cmd_tx = self.cmd_tx.read().await.clone();
1387        let _: DeriveEmptyResult = send_request(
1388            &self.rate_limiter,
1389            &cmd_tx,
1390            methods::PRIVATE_CANCEL_ALL,
1391            params,
1392            self.request_timeout,
1393        )
1394        .await?;
1395        Ok(())
1396    }
1397
1398    pub(crate) async fn reserve_matching_request(
1399        &self,
1400        operation: &'static str,
1401        instrument_name: &Ustr,
1402    ) -> Result<MatchingRateLimitReservation> {
1403        self.require_authenticated(operation).await?;
1404        debug_assert_eq!(rate_class_for_method(operation), RateClass::Matching);
1405        let window = self
1406            .rate_limiter
1407            .await_class_ready(RateClass::Matching, Some(instrument_name))
1408            .await;
1409        self.ensure_authenticated(operation)?;
1410        Ok(MatchingRateLimitReservation {
1411            method: operation,
1412            instrument_name: *instrument_name,
1413            window,
1414        })
1415    }
1416
1417    // Signing between the reservation and the reserved send can cross a
1418    // window boundary; a rolled window re-acquires so the departure draws on
1419    // its own window's cells.
1420    async fn refresh_matching_reservation(&self, reservation: &MatchingRateLimitReservation) {
1421        self.rate_limiter
1422            .ensure_window_current(
1423                RateClass::Matching,
1424                Some(&reservation.instrument_name),
1425                reservation.window,
1426            )
1427            .await;
1428    }
1429
1430    fn ensure_authenticated(&self, operation: &'static str) -> Result<()> {
1431        if self.auth_tracker.is_authenticated() {
1432            return Ok(());
1433        }
1434
1435        Err(DeriveWsError::Authentication {
1436            operation: operation.to_string(),
1437            reason: "WebSocket session is not authenticated".to_string(),
1438        })
1439    }
1440
1441    async fn require_authenticated(&self, operation: &'static str) -> Result<()> {
1442        if self
1443            .auth_tracker
1444            .wait_for_authenticated(self.request_timeout)
1445            .await
1446        {
1447            return Ok(());
1448        }
1449
1450        Err(DeriveWsError::Authentication {
1451            operation: operation.to_string(),
1452            reason: "WebSocket session is not authenticated".to_string(),
1453        })
1454    }
1455}
1456
1457// Awaits the venue's raw `result`, bounded by `timeout`. A dropped responder
1458// (handler torn down on reconnect) surfaces as `RequestCancelled`, a timeout as
1459// `Timeout`; both leave a state-changing write's outcome ambiguous.
1460#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1461enum RequestRateLimit {
1462    /// Pace the request now, against the class buckets plus the carried
1463    /// instrument's per-instrument bucket when present.
1464    Await(Option<Ustr>),
1465    /// A matching reservation already consumed the cells; do not pace or
1466    /// consume again.
1467    Reserved,
1468}
1469
1470async fn send_raw<P>(
1471    rate_limiter: &WsRateLimiter,
1472    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1473    method: &'static str,
1474    params: &P,
1475    timeout: Duration,
1476) -> Result<Value>
1477where
1478    P: Serialize + ?Sized,
1479{
1480    send_raw_with_rate_limit(
1481        rate_limiter,
1482        cmd_tx,
1483        method,
1484        params,
1485        timeout,
1486        RequestRateLimit::Await(None),
1487        None,
1488    )
1489    .await
1490}
1491
1492// Awaits the venue's raw `result` for a matching write that carries an
1493// instrument, pacing it against the account-wide and per-instrument buckets.
1494async fn send_raw_for_instrument<P>(
1495    rate_limiter: &WsRateLimiter,
1496    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1497    method: &'static str,
1498    params: &P,
1499    timeout: Duration,
1500    instrument_name: Ustr,
1501) -> Result<Value>
1502where
1503    P: Serialize + ?Sized,
1504{
1505    send_raw_with_rate_limit(
1506        rate_limiter,
1507        cmd_tx,
1508        method,
1509        params,
1510        timeout,
1511        RequestRateLimit::Await(Some(instrument_name)),
1512        None,
1513    )
1514    .await
1515}
1516
1517async fn send_raw_after_rate_limit<P>(
1518    rate_limiter: &WsRateLimiter,
1519    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1520    method: &'static str,
1521    params: &P,
1522    timeout: Duration,
1523) -> Result<Value>
1524where
1525    P: Serialize + ?Sized,
1526{
1527    send_raw_with_rate_limit(
1528        rate_limiter,
1529        cmd_tx,
1530        method,
1531        params,
1532        timeout,
1533        RequestRateLimit::Reserved,
1534        None,
1535    )
1536    .await
1537}
1538
1539async fn send_raw_with_rate_limit<P>(
1540    rate_limiter: &WsRateLimiter,
1541    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1542    method: &'static str,
1543    params: &P,
1544    timeout: Duration,
1545    rate_limit: RequestRateLimit,
1546    connection_epoch: Option<u64>,
1547) -> Result<Value>
1548where
1549    P: Serialize + ?Sized,
1550{
1551    let params = SecretString::from(serde_json::to_string(params)?);
1552
1553    if let RequestRateLimit::Await(instrument_name) = rate_limit {
1554        rate_limiter
1555            .await_class_ready(rate_class_for_method(method), instrument_name.as_ref())
1556            .await;
1557    }
1558
1559    let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1560    cmd_tx
1561        .send(HandlerCommand::Request {
1562            method,
1563            params,
1564            connection_epoch,
1565            response_tx,
1566        })
1567        .map_err(|e| DeriveWsError::transport(format!("failed to enqueue `{method}`: {e}")))?;
1568
1569    // On timeout the handler's `pending` entry leaks until the next reconnect's
1570    // `fail_pending` drains it; the later send to the dropped receiver is a
1571    // no-op logged at debug.
1572    match tokio::time::timeout(timeout, response_rx).await {
1573        Ok(Ok(outcome)) => outcome,
1574        Ok(Err(_)) => Err(DeriveWsError::RequestCancelled {
1575            method: method.to_owned(),
1576        }),
1577        Err(_) => Err(DeriveWsError::Timeout {
1578            method: method.to_owned(),
1579        }),
1580    }
1581}
1582
1583// Decodes the result, treating a null/absent `result` as `R::default()` (for
1584// login/subscribe/unsubscribe and the cancel family's `DeriveEmptyResult`).
1585async fn send_request<P, R>(
1586    rate_limiter: &WsRateLimiter,
1587    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1588    method: &'static str,
1589    params: &P,
1590    timeout: Duration,
1591) -> Result<R>
1592where
1593    P: Serialize + ?Sized,
1594    R: Default + DeserializeOwned,
1595{
1596    let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1597    decode_default_result(value)
1598}
1599
1600// Same as `send_request` for a matching write that carries an instrument, so
1601// the venue's per-instrument allowance is paced alongside the global one.
1602async fn send_request_for_instrument<P, R>(
1603    rate_limiter: &WsRateLimiter,
1604    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1605    method: &'static str,
1606    params: &P,
1607    timeout: Duration,
1608    instrument_name: Ustr,
1609) -> Result<R>
1610where
1611    P: Serialize + ?Sized,
1612    R: Default + DeserializeOwned,
1613{
1614    let value = send_raw_for_instrument(
1615        rate_limiter,
1616        cmd_tx,
1617        method,
1618        params,
1619        timeout,
1620        instrument_name,
1621    )
1622    .await?;
1623    decode_default_result(value)
1624}
1625
1626// Keep strict result decoding while reserving both matching buckets
1627async fn send_request_typed_for_instrument<P, R>(
1628    rate_limiter: &WsRateLimiter,
1629    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1630    method: &'static str,
1631    params: &P,
1632    timeout: Duration,
1633    instrument_name: Ustr,
1634) -> Result<R>
1635where
1636    P: Serialize + ?Sized,
1637    R: DeserializeOwned,
1638{
1639    let value = send_raw_for_instrument(
1640        rate_limiter,
1641        cmd_tx,
1642        method,
1643        params,
1644        timeout,
1645        instrument_name,
1646    )
1647    .await?;
1648    Ok(serde_json::from_value(value)?)
1649}
1650
1651fn decode_default_result<R>(value: Value) -> Result<R>
1652where
1653    R: Default + DeserializeOwned,
1654{
1655    if value.is_null() {
1656        Ok(R::default())
1657    } else {
1658        Ok(serde_json::from_value(value)?)
1659    }
1660}
1661
1662async fn send_request_on_connection<P, R>(
1663    rate_limiter: &WsRateLimiter,
1664    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1665    method: &'static str,
1666    params: &P,
1667    timeout: Duration,
1668    connection_epoch: u64,
1669) -> Result<R>
1670where
1671    P: Serialize + ?Sized,
1672    R: Default + DeserializeOwned,
1673{
1674    let value = send_raw_with_rate_limit(
1675        rate_limiter,
1676        cmd_tx,
1677        method,
1678        params,
1679        timeout,
1680        RequestRateLimit::Await(None),
1681        Some(connection_epoch),
1682    )
1683    .await?;
1684
1685    decode_default_result(value)
1686}
1687
1688// Decodes the result with no `Default` fallback, for `private/order` and
1689// `private/replace` whose success result is always a populated object.
1690async fn send_request_typed<P, R>(
1691    rate_limiter: &WsRateLimiter,
1692    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1693    method: &'static str,
1694    params: &P,
1695    timeout: Duration,
1696) -> Result<R>
1697where
1698    P: Serialize + ?Sized,
1699    R: DeserializeOwned,
1700{
1701    let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1702    Ok(serde_json::from_value(value)?)
1703}
1704
1705async fn send_request_typed_after_rate_limit<P, R>(
1706    rate_limiter: &WsRateLimiter,
1707    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1708    method: &'static str,
1709    params: &P,
1710    timeout: Duration,
1711) -> Result<R>
1712where
1713    P: Serialize + ?Sized,
1714    R: DeserializeOwned,
1715{
1716    let value = send_raw_after_rate_limit(rate_limiter, cmd_tx, method, params, timeout).await?;
1717    Ok(serde_json::from_value(value)?)
1718}
1719
1720fn channel_topics(channels: &[DeriveWsChannel]) -> Vec<String> {
1721    channels.iter().map(ToString::to_string).collect()
1722}
1723
1724#[expect(
1725    clippy::too_many_arguments,
1726    reason = "authentication state is passed explicitly for epoch fencing"
1727)]
1728async fn login_via_handler(
1729    rate_limiter: &WsRateLimiter,
1730    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1731    auth_tracker: &AuthTracker,
1732    authenticated_epoch: &AtomicU64,
1733    connection_mode: &AtomicU8,
1734    connection_epoch: &AtomicU64,
1735    creds: &DeriveWsCredentials,
1736    timeout: Duration,
1737) -> Result<()> {
1738    let _receiver = auth_tracker.begin();
1739    let expected_epoch = connection_epoch.load(Ordering::Acquire);
1740
1741    match send_login_request(rate_limiter, cmd_tx, creds, timeout, expected_epoch).await {
1742        Ok(())
1743            if complete_session_authentication(
1744                auth_tracker,
1745                authenticated_epoch,
1746                connection_mode,
1747                connection_epoch,
1748                expected_epoch,
1749            ) =>
1750        {
1751            log::debug!("Derive WebSocket authenticated");
1752
1753            Ok(())
1754        }
1755        Ok(()) => {
1756            let e = DeriveWsError::transport(
1757                "connection changed while completing WebSocket authentication",
1758            );
1759            authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1760            auth_tracker.fail(e.to_string());
1761            Err(e)
1762        }
1763        Err(e) => {
1764            authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1765            auth_tracker.fail(e.to_string());
1766            Err(e)
1767        }
1768    }
1769}
1770
1771async fn send_login_request(
1772    rate_limiter: &WsRateLimiter,
1773    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1774    creds: &DeriveWsCredentials,
1775    timeout: Duration,
1776    connection_epoch: u64,
1777) -> Result<()> {
1778    let login = build_ws_login(&creds.wallet_address, &creds.signer)?;
1779    let params = Zeroizing::new(WsLoginParams {
1780        wallet: login.wallet,
1781        timestamp: login.timestamp,
1782        signature: login.signature,
1783    });
1784    let result = send_request_on_connection::<_, WsLoginResult>(
1785        rate_limiter,
1786        cmd_tx,
1787        methods::PUBLIC_LOGIN,
1788        &*params,
1789        timeout,
1790        connection_epoch,
1791    )
1792    .await?;
1793
1794    if matches!(result, WsLoginResult::Success { success: false }) {
1795        return Err(DeriveWsError::Authentication {
1796            operation: methods::PUBLIC_LOGIN.to_string(),
1797            reason: "venue returned an unsuccessful login result".to_string(),
1798        });
1799    }
1800
1801    Ok(())
1802}
1803
1804#[expect(
1805    clippy::too_many_arguments,
1806    reason = "recovery state is passed explicitly for epoch fencing"
1807)]
1808async fn recover_session(
1809    rate_limiter: &WsRateLimiter,
1810    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1811    auth_tracker: &AuthTracker,
1812    authenticated_epoch: &AtomicU64,
1813    connection_mode: &AtomicU8,
1814    connection_epoch: &AtomicU64,
1815    creds: Option<&DeriveWsCredentials>,
1816    subscriptions: &DashMap<String, ()>,
1817    subscription_lock: &tokio::sync::Mutex<()>,
1818    timeout: Duration,
1819) -> Result<u64> {
1820    let _guard = subscription_lock.lock().await;
1821    let _receiver = creds.map(|_| auth_tracker.begin());
1822
1823    for attempt in 1..=MAX_SESSION_RECOVERY_ATTEMPTS {
1824        let expected_epoch = wait_for_session_connection(connection_mode, connection_epoch).await?;
1825
1826        let result = async {
1827            if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1828                return Err(DeriveWsError::transport(
1829                    "connection changed before WebSocket session recovery",
1830                ));
1831            }
1832
1833            if let Some(creds) = creds {
1834                send_login_request(rate_limiter, cmd_tx, creds, timeout, expected_epoch).await?;
1835            }
1836            let channels: Vec<String> = subscriptions
1837                .iter()
1838                .map(|entry| entry.key().clone())
1839                .collect();
1840            subscribe_via_handler(rate_limiter, cmd_tx, channels, timeout, expected_epoch).await?;
1841
1842            if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1843                return Err(DeriveWsError::transport(
1844                    "connection changed during WebSocket session recovery",
1845                ));
1846            }
1847
1848            Ok(())
1849        }
1850        .await;
1851
1852        match result {
1853            Ok(()) => {
1854                if creds.is_some()
1855                    && !complete_session_authentication(
1856                        auth_tracker,
1857                        authenticated_epoch,
1858                        connection_mode,
1859                        connection_epoch,
1860                        expected_epoch,
1861                    )
1862                {
1863                    continue;
1864                }
1865
1866                if creds.is_some() {
1867                    log::info!("Derive WebSocket session re-authenticated");
1868                }
1869
1870                return Ok(expected_epoch);
1871            }
1872            Err(e) if attempt < MAX_SESSION_RECOVERY_ATTEMPTS => {
1873                let multiplier = 1_u32 << (attempt - 1);
1874                let delay = RECONNECT_BASE_BACKOFF
1875                    .saturating_mul(multiplier)
1876                    .min(RECONNECT_MAX_BACKOFF);
1877                log::warn!(
1878                    "Derive WebSocket session recovery attempt {attempt}/{MAX_SESSION_RECOVERY_ATTEMPTS} failed: {e}; retrying in {delay:?}",
1879                );
1880                tokio::time::sleep(delay).await;
1881            }
1882            Err(e) => {
1883                authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1884
1885                if creds.is_some() {
1886                    auth_tracker.fail(e.to_string());
1887                }
1888                return Err(e);
1889            }
1890        }
1891    }
1892
1893    let e = DeriveWsError::transport("WebSocket session changed while recovery completed");
1894    authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1895
1896    if creds.is_some() {
1897        auth_tracker.fail(e.to_string());
1898    }
1899    Err(e)
1900}
1901
1902async fn wait_for_session_connection(
1903    connection_mode: &AtomicU8,
1904    connection_epoch: &AtomicU64,
1905) -> Result<u64> {
1906    loop {
1907        match ConnectionMode::from_atomic(connection_mode) {
1908            ConnectionMode::Active => return Ok(connection_epoch.load(Ordering::Acquire)),
1909            ConnectionMode::Reconnect => tokio::time::sleep(RECONNECT_BASE_BACKOFF).await,
1910            ConnectionMode::Disconnect | ConnectionMode::Closed => {
1911                return Err(DeriveWsError::transport(
1912                    "WebSocket closed during session recovery",
1913                ));
1914            }
1915        }
1916    }
1917}
1918
1919fn complete_session_authentication(
1920    auth_tracker: &AuthTracker,
1921    authenticated_epoch: &AtomicU64,
1922    connection_mode: &AtomicU8,
1923    connection_epoch: &AtomicU64,
1924    expected_epoch: u64,
1925) -> bool {
1926    if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1927        return false;
1928    }
1929
1930    authenticated_epoch.store(expected_epoch, Ordering::Release);
1931    auth_tracker.succeed();
1932
1933    if session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1934        true
1935    } else {
1936        authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1937        auth_tracker.invalidate();
1938        false
1939    }
1940}
1941
1942fn session_connection_is_active(
1943    connection_mode: &AtomicU8,
1944    connection_epoch: &AtomicU64,
1945    expected_epoch: u64,
1946) -> bool {
1947    ConnectionMode::from_atomic(connection_mode).is_active()
1948        && connection_epoch.load(Ordering::Acquire) == expected_epoch
1949}
1950
1951async fn subscribe_via_handler(
1952    rate_limiter: &WsRateLimiter,
1953    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1954    channels: Vec<String>,
1955    timeout: Duration,
1956    connection_epoch: u64,
1957) -> Result<()> {
1958    if channels.is_empty() {
1959        return Ok(());
1960    }
1961
1962    let params = WsSubscribeParams {
1963        channels: channels.into_iter().map(DeriveWsChannel::from).collect(),
1964    };
1965    let result: WsSubscribeResult = send_request_on_connection(
1966        rate_limiter,
1967        cmd_tx,
1968        methods::PUBLIC_SUBSCRIBE,
1969        &params,
1970        timeout,
1971        connection_epoch,
1972    )
1973    .await?;
1974
1975    let (_, failure) = subscription_outcome(
1976        &params.channels,
1977        &result,
1978        SUBSCRIPTION_REPLAY_ACCEPTED_STATUSES,
1979    );
1980    failure.map_or(Ok(()), Err)
1981}
1982
1983fn subscription_outcome(
1984    requested: &[DeriveWsChannel],
1985    result: &WsSubscribeResult,
1986    accepted_statuses: &[&str],
1987) -> (Vec<String>, Option<DeriveWsError>) {
1988    let mut confirmed = Vec::with_capacity(requested.len());
1989    let mut failures = Vec::new();
1990
1991    for channel in requested {
1992        let topic = channel.to_string();
1993        match result.status.get(channel) {
1994            Some(status) if accepted_statuses.contains(&status.as_str()) => confirmed.push(topic),
1995            Some(status) => failures.push(format!("{topic}: {status}")),
1996            None if result.channels.contains(channel) => confirmed.push(topic),
1997            None => failures.push(format!("{topic}: missing channel status")),
1998        }
1999    }
2000
2001    let failure = (!failures.is_empty()).then(|| DeriveWsError::Subscription {
2002        details: failures.join(", "),
2003    });
2004    (confirmed, failure)
2005}
2006
2007#[cfg(test)]
2008mod tests {
2009    use rstest::rstest;
2010
2011    use super::*;
2012    use crate::common::rate_limit::RateBucket;
2013
2014    #[rstest]
2015    fn test_public_client_defaults_to_environment_url() {
2016        let client = DeriveWebSocketClient::new(
2017            None,
2018            DeriveEnvironment::Mainnet,
2019            TransportBackend::default(),
2020            None,
2021        );
2022        assert!(client.url().starts_with("wss://"));
2023        assert!(client.url().contains("api.lyra.finance"));
2024        assert!(!client.is_authenticated());
2025        assert!(!client.is_active());
2026        assert_eq!(client.subscription_count(), 0);
2027    }
2028
2029    #[tokio::test]
2030    async fn test_execution_auth_barrier_waits_for_authentication() {
2031        let client = DeriveWebSocketClient::with_credentials(
2032            None,
2033            DeriveEnvironment::Mainnet,
2034            TransportBackend::default(),
2035            None,
2036            DeriveWsCredentials::new(
2037                "0x000000000000000000000000000000000000aaaa",
2038                "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2039            )
2040            .unwrap(),
2041            None,
2042            None,
2043        );
2044        let execution = client.execution_handle();
2045        let auth_tracker = execution.auth_tracker.clone();
2046        let _receiver = auth_tracker.begin();
2047        let tracker_for_task = auth_tracker.clone();
2048
2049        get_runtime().spawn(async move {
2050            tokio::time::sleep(Duration::from_millis(10)).await;
2051            tracker_for_task.succeed();
2052        });
2053
2054        execution
2055            .require_authenticated(methods::PRIVATE_ORDER)
2056            .await
2057            .expect("barrier should wait for successful authentication");
2058    }
2059
2060    #[tokio::test]
2061    async fn test_execution_auth_barrier_fails_on_terminal_auth_failure() {
2062        let client = DeriveWebSocketClient::with_credentials(
2063            None,
2064            DeriveEnvironment::Mainnet,
2065            TransportBackend::default(),
2066            None,
2067            DeriveWsCredentials::new(
2068                "0x000000000000000000000000000000000000aaaa",
2069                "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2070            )
2071            .unwrap(),
2072            None,
2073            None,
2074        );
2075        let execution = client.execution_handle();
2076        let _receiver = execution.auth_tracker.begin();
2077        execution.auth_tracker.fail("bad signature");
2078
2079        let error = execution
2080            .require_authenticated(methods::PRIVATE_ORDER)
2081            .await
2082            .expect_err("terminal auth failure must reject private operations");
2083
2084        assert!(matches!(error, DeriveWsError::Authentication { .. }));
2085    }
2086
2087    #[tokio::test]
2088    async fn test_session_recovery_waits_for_reconnecting_transport() {
2089        let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect as u8));
2090        let connection_epoch = Arc::new(AtomicU64::new(1));
2091        let mode_for_task = Arc::clone(&connection_mode);
2092        let epoch_for_task = Arc::clone(&connection_epoch);
2093
2094        get_runtime().spawn(async move {
2095            tokio::time::sleep(Duration::from_millis(10)).await;
2096            epoch_for_task.store(2, Ordering::Release);
2097            mode_for_task.store(ConnectionMode::Active as u8, Ordering::Release);
2098        });
2099
2100        let epoch = tokio::time::timeout(
2101            Duration::from_secs(1),
2102            wait_for_session_connection(&connection_mode, &connection_epoch),
2103        )
2104        .await
2105        .expect("session recovery should resume after reconnect")
2106        .expect("active replacement connection should be accepted");
2107
2108        assert_eq!(epoch, 2);
2109    }
2110
2111    #[rstest]
2112    fn test_testnet_client_routes_to_demo_url() {
2113        let client = DeriveWebSocketClient::new(
2114            None,
2115            DeriveEnvironment::Testnet,
2116            TransportBackend::default(),
2117            None,
2118        );
2119        assert!(client.url().contains("demo"));
2120    }
2121
2122    #[rstest]
2123    fn test_credentials_constructor_parses_session_key() {
2124        let creds = DeriveWsCredentials::new(
2125            "0x000000000000000000000000000000000000aaaa",
2126            "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2127        )
2128        .unwrap();
2129        assert!(creds.wallet_address.starts_with("0x"));
2130        let client = DeriveWebSocketClient::with_credentials(
2131            None,
2132            DeriveEnvironment::Testnet,
2133            TransportBackend::default(),
2134            None,
2135            creds,
2136            None,
2137            None,
2138        );
2139        assert!(client.url().contains("demo"));
2140        assert!(!client.is_authenticated());
2141    }
2142
2143    #[rstest]
2144    fn test_credentials_debug_redacts_signer() {
2145        let creds = DeriveWsCredentials::new(
2146            "0xWALLET",
2147            "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2148        )
2149        .unwrap();
2150        let debug = format!("{creds:?}");
2151        assert!(debug.contains("redacted"));
2152        assert!(debug.contains("0xWALLET"));
2153        assert!(!debug.contains("2ae8be44"));
2154    }
2155
2156    #[rstest]
2157    fn test_credentials_constructor_rejects_invalid_session_key() {
2158        let err = DeriveWsCredentials::new("0xWALLET", "not-a-hex-key").unwrap_err();
2159        assert!(err.to_string().contains("invalid session key"));
2160    }
2161
2162    #[rstest]
2163    #[tokio::test]
2164    async fn test_send_raw_times_out_when_no_response_arrives() {
2165        // Keep the receiver alive so the request enqueues, but never reply: the
2166        // bounded await must surface a Timeout rather than hang forever.
2167        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
2168        let rate_limiter: WsRateLimiter =
2169            FixedWindowLimiter::new(FixedWindowLimits::websocket(None, None), MonotonicClock {});
2170        let err = send_raw(
2171            &rate_limiter,
2172            &cmd_tx,
2173            methods::PRIVATE_ORDER,
2174            &serde_json::json!({}),
2175            Duration::from_millis(50),
2176        )
2177        .await
2178        .expect_err("must time out");
2179
2180        match err {
2181            DeriveWsError::Timeout { method } => assert_eq!(method, methods::PRIVATE_ORDER),
2182            other => panic!("expected Timeout, was {other:?}"),
2183        }
2184    }
2185
2186    #[rstest]
2187    #[tokio::test]
2188    async fn test_send_request_typed_rejects_null_result() {
2189        // `private/order` and `private/replace` always return a populated
2190        // object on success; a null result is a protocol violation that must
2191        // surface as a serde error (classified ambiguous by the exec client).
2192        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
2193        tokio::spawn(async move {
2194            if let Some(HandlerCommand::Request { response_tx, .. }) = cmd_rx.recv().await {
2195                let _ = response_tx.send(Ok(Value::Null));
2196            }
2197        });
2198        let rate_limiter: WsRateLimiter =
2199            FixedWindowLimiter::new(FixedWindowLimits::websocket(None, None), MonotonicClock {});
2200        let result: Result<DeriveOrderResult> = send_request_typed(
2201            &rate_limiter,
2202            &cmd_tx,
2203            methods::PRIVATE_ORDER,
2204            &serde_json::json!({}),
2205            Duration::from_secs(1),
2206        )
2207        .await;
2208        assert!(matches!(result, Err(DeriveWsError::Serde(_))));
2209    }
2210
2211    #[rstest]
2212    #[tokio::test]
2213    async fn test_reserved_send_does_not_wait_for_or_consume_second_quota_cell() {
2214        let rate_limiter: WsRateLimiter =
2215            FixedWindowLimiter::new(FixedWindowLimits::websocket(None, None), MonotonicClock {});
2216
2217        for _ in 0..5 {
2218            rate_limiter
2219                .check_bucket(RateBucket::Matching)
2220                .expect("reservation consumes the matching window");
2221        }
2222        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
2223        tokio::spawn(async move {
2224            if let Some(HandlerCommand::Request { response_tx, .. }) = cmd_rx.recv().await {
2225                let _ = response_tx.send(Ok(serde_json::json!({"accepted": true})));
2226            }
2227        });
2228
2229        let response = tokio::time::timeout(
2230            Duration::from_millis(100),
2231            send_raw_after_rate_limit(
2232                &rate_limiter,
2233                &cmd_tx,
2234                methods::PRIVATE_ORDER,
2235                &serde_json::json!({}),
2236                Duration::from_secs(1),
2237            ),
2238        )
2239        .await
2240        .expect("reserved send must not wait for quota")
2241        .expect("reserved send succeeds");
2242
2243        assert_eq!(response, serde_json::json!({"accepted": true}));
2244        assert!(
2245            rate_limiter.check_bucket(RateBucket::Matching).is_err(),
2246            "reserved send must not consume a second cell",
2247        );
2248    }
2249}