nautilus_bitmex/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//! Provides the WebSocket client integration for the [BitMEX](https://bitmex.com) WebSocket API.
17//!
18//! This module defines and implements a [`BitmexWebSocketClient`] for
19//! connecting to BitMEX WebSocket streams. It handles authentication (when credentials
20//! are provided), manages subscriptions to market data and account update channels,
21//! and emits venue-specific message types for consumers to parse.
22
23use std::{
24    sync::{
25        Arc,
26        atomic::{AtomicBool, AtomicU8, Ordering},
27    },
28    time::Duration,
29};
30
31use arc_swap::ArcSwap;
32use dashmap::DashMap;
33use futures_util::Stream;
34use nautilus_common::live::get_runtime;
35use nautilus_core::{
36    consts::NAUTILUS_USER_AGENT,
37    env::{get_env_var, get_or_env_var_opt},
38};
39use nautilus_live::SocketControl;
40use nautilus_model::{
41    data::bar::BarType,
42    identifiers::{AccountId, InstrumentId},
43    instruments::{Instrument, InstrumentAny},
44};
45use nautilus_network::{
46    http::USER_AGENT,
47    mode::ConnectionMode,
48    websocket::{
49        AUTHENTICATION_TIMEOUT_SECS, AuthTracker, SubscriptionState, TransportBackend,
50        WebSocketClient, WebSocketConfig, channel_message_handler,
51    },
52};
53use tokio_tungstenite::tungstenite::Message;
54use ustr::Ustr;
55
56use super::{
57    enums::{BitmexWsAuthAction, BitmexWsAuthChannel, BitmexWsOperation, BitmexWsTopic},
58    error::BitmexWsError,
59    handler::{BitmexWsFeedHandler, HandlerCommand},
60    messages::{BitmexAuthentication, BitmexSubscription, BitmexWsMessage},
61    parse::{is_index_symbol, topic_from_bar_spec},
62};
63use crate::common::{
64    consts::{BITMEX_WS_TOPIC_DELIMITER, BITMEX_WS_URL},
65    credential::{Credential, credential_env_vars},
66    enums::BitmexEnvironment,
67};
68
69/// Provides a WebSocket client for connecting to the [BitMEX](https://bitmex.com) real-time API.
70///
71/// Key runtime patterns:
72/// - Authentication handshakes are managed by the internal auth tracker, ensuring resubscriptions
73///   occur only after BitMEX acknowledges `authKey` messages.
74/// - The subscription state maintains pending and confirmed topics so reconnection replay is
75///   deterministic and per-topic errors are surfaced.
76#[derive(Clone, Debug)]
77pub struct BitmexWebSocketClient {
78    url: String,
79    credential: Option<Credential>,
80    heartbeat: Option<u64>,
81    auth_timeout_secs: u64,
82    account_id: AccountId,
83    auth_tracker: AuthTracker,
84    signal: Arc<AtomicBool>,
85    connection_mode: Arc<ArcSwap<AtomicU8>>,
86    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
87    out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<BitmexWsMessage>>>,
88    task_handle: Option<Arc<tokio::task::JoinHandle<()>>>,
89    subscriptions: SubscriptionState,
90    tracked_subscriptions: Arc<DashMap<String, ()>>,
91    instruments: Arc<DashMap<Ustr, InstrumentAny>>,
92    transport_backend: TransportBackend,
93    proxy_url: Option<String>,
94    socket_control: Option<SocketControl>,
95}
96
97impl BitmexWebSocketClient {
98    /// Creates a new [`BitmexWebSocketClient`] instance.
99    ///
100    /// # Errors
101    ///
102    /// Returns an error if only one of `api_key` or `api_secret` is provided (both or neither required).
103    #[expect(clippy::too_many_arguments)]
104    pub fn new(
105        url: Option<String>,
106        api_key: Option<String>,
107        api_secret: Option<String>,
108        account_id: Option<AccountId>,
109        heartbeat: u64,
110        auth_timeout_secs: Option<u64>,
111        transport_backend: TransportBackend,
112        proxy_url: Option<String>,
113    ) -> anyhow::Result<Self> {
114        let credential = match (api_key, api_secret) {
115            (Some(key), Some(secret)) => Some(Credential::new(key, secret)),
116            (None, None) => None,
117            _ => anyhow::bail!("Both `api_key` and `api_secret` must be provided together"),
118        };
119
120        let account_id = account_id.unwrap_or(AccountId::from("BITMEX-master"));
121
122        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
123        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
124
125        // Placeholder channel until connect() creates the real one
126        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
127
128        Ok(Self {
129            url: url.unwrap_or(BITMEX_WS_URL.to_string()),
130            credential,
131            heartbeat: Some(heartbeat),
132            auth_timeout_secs: auth_timeout_secs.unwrap_or(AUTHENTICATION_TIMEOUT_SECS),
133            account_id,
134            auth_tracker: AuthTracker::new(),
135            signal: Arc::new(AtomicBool::new(false)),
136            connection_mode,
137            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
138            out_rx: None,
139            task_handle: None,
140            subscriptions: SubscriptionState::new(BITMEX_WS_TOPIC_DELIMITER),
141            tracked_subscriptions: Arc::new(DashMap::new()),
142            instruments: Arc::new(DashMap::new()),
143            transport_backend,
144            proxy_url,
145            socket_control: None,
146        })
147    }
148
149    /// Configures socket state reporting and reconnect control.
150    #[must_use]
151    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
152        self.socket_control = Some(control);
153        self
154    }
155
156    /// Creates a new [`BitmexWebSocketClient`] with environment variable credential resolution.
157    ///
158    /// If `api_key` or `api_secret` are not provided, they will be loaded from
159    /// environment variables based on the `environment`:
160    /// - Testnet: `BITMEX_TESTNET_API_KEY`, `BITMEX_TESTNET_API_SECRET`
161    /// - Mainnet: `BITMEX_API_KEY`, `BITMEX_API_SECRET`
162    ///
163    /// # Errors
164    ///
165    /// Returns an error if only one of `api_key` or `api_secret` is provided.
166    #[expect(clippy::too_many_arguments)]
167    pub fn new_with_env(
168        url: Option<String>,
169        api_key: Option<String>,
170        api_secret: Option<String>,
171        account_id: Option<AccountId>,
172        heartbeat: u64,
173        auth_timeout_secs: Option<u64>,
174        environment: BitmexEnvironment,
175        transport_backend: TransportBackend,
176        proxy_url: Option<String>,
177    ) -> anyhow::Result<Self> {
178        let (api_key_env, api_secret_env) = credential_env_vars(environment);
179
180        let key = get_or_env_var_opt(api_key, api_key_env);
181        let secret = get_or_env_var_opt(api_secret, api_secret_env);
182
183        Self::new(
184            url,
185            key,
186            secret,
187            account_id,
188            heartbeat,
189            auth_timeout_secs,
190            transport_backend,
191            proxy_url,
192        )
193    }
194
195    /// Creates a new authenticated [`BitmexWebSocketClient`] using environment variables.
196    ///
197    /// # Errors
198    ///
199    /// Returns an error if environment variables are not set or credentials are invalid.
200    pub fn from_env() -> anyhow::Result<Self> {
201        let url = get_env_var("BITMEX_WS_URL")?;
202        let (key_var, secret_var) = credential_env_vars(BitmexEnvironment::Mainnet);
203        let api_key = get_env_var(key_var)?;
204        let api_secret = get_env_var(secret_var)?;
205
206        Self::new(
207            Some(url),
208            Some(api_key),
209            Some(api_secret),
210            None,
211            5,
212            None,
213            TransportBackend::default(),
214            None,
215        )
216    }
217
218    /// Returns the websocket url being used by the client.
219    #[must_use]
220    pub const fn url(&self) -> &str {
221        self.url.as_str()
222    }
223
224    /// Returns the public API key being used by the client.
225    #[must_use]
226    pub fn api_key(&self) -> Option<&str> {
227        self.credential.as_ref().map(|c| c.api_key())
228    }
229
230    /// Returns a masked version of the API key for logging purposes.
231    #[must_use]
232    pub fn api_key_masked(&self) -> Option<String> {
233        self.credential.as_ref().map(|c| c.api_key_masked())
234    }
235
236    /// Returns a value indicating whether the client is active.
237    #[must_use]
238    pub fn is_active(&self) -> bool {
239        let connection_mode_arc = self.connection_mode.load();
240        ConnectionMode::from_atomic(&connection_mode_arc).is_active()
241            && !self.signal.load(Ordering::Relaxed)
242    }
243
244    /// Returns a value indicating whether the client is closed.
245    #[must_use]
246    pub fn is_closed(&self) -> bool {
247        let connection_mode_arc = self.connection_mode.load();
248        ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
249            || self.signal.load(Ordering::Relaxed)
250    }
251
252    /// Returns the account ID.
253    #[must_use]
254    pub fn account_id(&self) -> AccountId {
255        self.account_id
256    }
257
258    /// Sets the account ID.
259    pub fn set_account_id(&mut self, account_id: AccountId) {
260        self.account_id = account_id;
261    }
262
263    /// Bulk-replaces the instrument cache with the given instruments.
264    pub fn cache_instruments(&self, instruments: &[InstrumentAny]) {
265        self.instruments.clear();
266        for inst in instruments {
267            self.instruments
268                .insert(inst.raw_symbol().inner(), inst.clone());
269        }
270    }
271
272    /// Upserts a single instrument into the cache.
273    pub fn cache_instrument(&self, instrument: InstrumentAny) {
274        self.instruments
275            .insert(instrument.raw_symbol().inner(), instrument);
276    }
277
278    /// Retrieves an instrument from the cache by symbol.
279    #[must_use]
280    pub fn get_instrument(&self, symbol: &Ustr) -> Option<InstrumentAny> {
281        self.instruments
282            .get(symbol)
283            .map(|entry| entry.value().clone())
284    }
285
286    /// Connect to the BitMEX WebSocket server.
287    ///
288    /// # Errors
289    ///
290    /// Returns an error if the WebSocket connection fails or authentication fails (if credentials provided).
291    pub async fn connect(&mut self) -> Result<(), BitmexWsError> {
292        let (client, raw_rx) = self.connect_inner().await?;
293
294        // Reset shutdown signal so is_active() works after close+reconnect
295        self.signal.store(false, Ordering::Relaxed);
296
297        // Replace connection state so all clones see the underlying WebSocketClient's state
298        self.connection_mode.store(client.connection_mode_atomic());
299        let reconnect_handle = client.reconnect_handle();
300
301        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<BitmexWsMessage>();
302        self.out_rx = Some(Arc::new(out_rx));
303
304        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
305        *self.cmd_tx.write().await = cmd_tx.clone();
306
307        // Send WebSocketClient to handler
308        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
309            return Err(BitmexWsError::ClientError(format!(
310                "Failed to send WebSocketClient to handler: {e}"
311            )));
312        }
313
314        if let Some(control) = &self.socket_control {
315            control.register(move || reconnect_handle.request_reconnect());
316        }
317
318        let signal = self.signal.clone();
319        let credential = self.credential.clone();
320        let auth_tracker = self.auth_tracker.clone();
321        let subscriptions = self.subscriptions.clone();
322        let cmd_tx_for_reconnect = cmd_tx.clone();
323
324        let stream_handle = get_runtime().spawn(async move {
325            let mut handler = BitmexWsFeedHandler::new(
326                signal.clone(),
327                cmd_rx,
328                raw_rx,
329                out_tx,
330                auth_tracker.clone(),
331                subscriptions.clone(),
332            );
333
334            // Resubscribe all tracked subscriptions after reconnection.
335            let resubscribe_all = || {
336                // Use SubscriptionState as source of truth for what to restore
337                let topics = subscriptions.all_topics();
338
339                if topics.is_empty() {
340                    return;
341                }
342
343                log::debug!(
344                    "Resubscribing to confirmed subscriptions: count={}",
345                    topics.len()
346                );
347
348                for topic in &topics {
349                    subscriptions.mark_subscribe(topic.as_str());
350                }
351
352                // Serialize subscription messages
353                let mut payloads = Vec::with_capacity(topics.len());
354                for topic in &topics {
355                    let message = BitmexSubscription {
356                        op: BitmexWsOperation::Subscribe,
357                        args: vec![Ustr::from(topic.as_ref())],
358                    };
359
360                    if let Ok(payload) = serde_json::to_string(&message) {
361                        payloads.push(payload);
362                    }
363                }
364
365                if let Err(e) =
366                    cmd_tx_for_reconnect.send(HandlerCommand::Subscribe { topics: payloads })
367                {
368                    log::error!("Failed to send resubscribe command: {e}");
369                }
370            };
371
372            let mut waiting_for_reconnect_auth = false;
373
374            // Run message processing with reconnection handling
375            loop {
376                match handler.next().await {
377                    Some(BitmexWsMessage::Reconnected) => {
378                        if signal.load(Ordering::Relaxed) {
379                            continue;
380                        }
381
382                        log::info!("WebSocket reconnected");
383
384                        subscriptions.reset_after_reconnect();
385
386                        if let Some(cred) = &credential {
387                            log::debug!("Re-authenticating after reconnection");
388                            waiting_for_reconnect_auth = true;
389
390                            let expires = (jiff::Timestamp::now()
391                                + jiff::SignedDuration::from_secs(30))
392                            .as_second();
393                            let signature = cred.sign("GET", "/realtime", expires, "");
394
395                            let auth_message = BitmexAuthentication {
396                                op: BitmexWsAuthAction::AuthKeyExpires,
397                                args: (cred.api_key().to_string(), expires, signature),
398                            };
399
400                            if let Ok(payload) = serde_json::to_string(&auth_message) {
401                                if let Err(e) = cmd_tx_for_reconnect
402                                    .send(HandlerCommand::Authenticate { payload })
403                                {
404                                    log::error!("Failed to send reconnection auth command: {e}");
405                                }
406                            } else {
407                                log::error!("Failed to serialize reconnection auth message");
408                            }
409                        }
410
411                        // Unauthenticated sessions resubscribe immediately after reconnection,
412                        // authenticated sessions wait for Authenticated message
413                        if credential.is_none() {
414                            log::debug!("No authentication required, resubscribing immediately");
415                            resubscribe_all();
416                        }
417
418                        if handler.send(BitmexWsMessage::Reconnected).is_err() {
419                            if handler.is_stopped() {
420                                log::debug!("Failed to forward reconnect event (receiver dropped)");
421                            } else {
422                                log::error!("Failed to forward reconnect event (receiver dropped)");
423                            }
424                            break;
425                        }
426                    }
427                    Some(BitmexWsMessage::Authenticated) => {
428                        if waiting_for_reconnect_auth {
429                            log::debug!("Authenticated after reconnection, resubscribing");
430                            resubscribe_all();
431                            waiting_for_reconnect_auth = false;
432                        }
433                    }
434                    Some(msg) => {
435                        if handler.send(msg).is_err() {
436                            if handler.is_stopped() {
437                                log::debug!("Failed to send message (receiver dropped)");
438                            } else {
439                                log::error!("Failed to send message (receiver dropped)");
440                            }
441                            break;
442                        }
443                    }
444                    None => {
445                        // Stream ended - check if it's a stop signal
446                        if handler.is_stopped() {
447                            log::debug!("Stop signal received, ending message processing");
448                            break;
449                        }
450                        // Otherwise it's an unexpected stream end
451                        log::warn!("WebSocket stream ended unexpectedly");
452                        break;
453                    }
454                }
455            }
456
457            log::debug!("Handler task exiting");
458        });
459
460        self.task_handle = Some(Arc::new(stream_handle));
461
462        if self.credential.is_some()
463            && let Err(e) = self.authenticate().await
464        {
465            if let Some(handle) = self.task_handle.take() {
466                handle.abort();
467            }
468            self.signal.store(true, Ordering::Relaxed);
469            return Err(e);
470        }
471
472        // Subscribe to instrument topic
473        let instrument_topic = BitmexWsTopic::Instrument.as_ref().to_string();
474        self.subscriptions.mark_subscribe(&instrument_topic);
475        self.tracked_subscriptions.insert(instrument_topic, ());
476
477        let subscribe_msg = BitmexSubscription {
478            op: BitmexWsOperation::Subscribe,
479            args: vec![Ustr::from(BitmexWsTopic::Instrument.as_ref())],
480        };
481
482        match serde_json::to_string(&subscribe_msg) {
483            Ok(subscribe_json) => {
484                if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Subscribe {
485                    topics: vec![subscribe_json],
486                }) {
487                    log::error!("Failed to send subscribe command for instruments: {e}");
488                } else {
489                    log::debug!("Subscribed to all instruments");
490                }
491            }
492            Err(e) => {
493                log::error!("Failed to serialize subscribe message: {e}");
494            }
495        }
496
497        Ok(())
498    }
499
500    /// Connect to the WebSocket and return a message receiver.
501    ///
502    /// # Errors
503    ///
504    /// Returns an error if the WebSocket connection fails or if authentication fails (when credentials are provided).
505    async fn connect_inner(
506        &self,
507    ) -> Result<
508        (
509            WebSocketClient,
510            tokio::sync::mpsc::UnboundedReceiver<Message>,
511        ),
512        BitmexWsError,
513    > {
514        let (message_handler, rx) = channel_message_handler();
515
516        // Inbound Ping frames are answered by the transport, so no ping handler is needed;
517        // the reader routes them away from the message channel and the handler never sees them.
518
519        let config = WebSocketConfig {
520            url: self.url.clone(),
521            headers: vec![(USER_AGENT.to_string(), NAUTILUS_USER_AGENT.to_string())],
522            heartbeat_interval_secs: self.heartbeat,
523            heartbeat_payload: None,
524            connect_timeout_ms: Some(5_000),
525            reconnect_delay_initial_ms: None, // Use default
526            reconnect_delay_max_ms: None,     // Use default
527            reconnect_backoff_factor: None,   // Use default
528            reconnect_jitter_ms: None,        // Use default
529            reconnect_max_attempts: None,
530            heartbeat_timeout_secs: None,
531            idle_timeout_ms: None,
532            backend: self.transport_backend,
533            proxy_url: self.proxy_url.clone(),
534        };
535
536        let keyed_quotas = vec![];
537        let client = WebSocketClient::builder()
538            .config(config)
539            .message_handler(message_handler)
540            .keyed_quotas(keyed_quotas)
541            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
542            .connect()
543            .await
544            .map_err(|e| BitmexWsError::ClientError(e.to_string()))?;
545
546        Ok((client, rx))
547    }
548
549    /// Authenticate the WebSocket connection using the provided credentials.
550    ///
551    /// # Errors
552    ///
553    /// Returns an error if the WebSocket is not connected, if authentication fails,
554    /// or if credentials are not available.
555    async fn authenticate(&self) -> Result<(), BitmexWsError> {
556        let credential = match &self.credential {
557            Some(credential) => credential,
558            None => {
559                return Err(BitmexWsError::AuthenticationError(
560                    "API credentials not available to authenticate".to_string(),
561                ));
562            }
563        };
564
565        let receiver = self.auth_tracker.begin();
566
567        let expires = (jiff::Timestamp::now() + jiff::SignedDuration::from_secs(30)).as_second();
568        let signature = credential.sign("GET", "/realtime", expires, "");
569
570        let auth_message = BitmexAuthentication {
571            op: BitmexWsAuthAction::AuthKeyExpires,
572            args: (credential.api_key().to_string(), expires, signature),
573        };
574
575        let auth_json = serde_json::to_string(&auth_message).map_err(|e| {
576            let msg = format!("Failed to serialize auth message: {e}");
577            self.auth_tracker.fail(msg.clone());
578            BitmexWsError::AuthenticationError(msg)
579        })?;
580
581        // Send Authenticate command to handler
582        self.cmd_tx
583            .read()
584            .await
585            .send(HandlerCommand::Authenticate { payload: auth_json })
586            .map_err(|e| {
587                let msg = format!("Failed to send authenticate command: {e}");
588                self.auth_tracker.fail(msg.clone());
589                BitmexWsError::AuthenticationError(msg)
590            })?;
591
592        self.auth_tracker
593            .wait_for_result::<BitmexWsError>(Duration::from_secs(self.auth_timeout_secs), receiver)
594            .await
595    }
596
597    /// Wait until the WebSocket connection is active.
598    ///
599    /// # Errors
600    ///
601    /// Returns an error if the connection times out.
602    pub async fn wait_until_active(&self, timeout_secs: f64) -> Result<(), BitmexWsError> {
603        let timeout = Duration::from_secs_f64(timeout_secs);
604
605        tokio::time::timeout(timeout, async {
606            while !self.is_active() {
607                tokio::time::sleep(Duration::from_millis(10)).await;
608            }
609        })
610        .await
611        .map_err(|_| {
612            BitmexWsError::ClientError(format!(
613                "WebSocket connection timeout after {timeout_secs} seconds"
614            ))
615        })?;
616
617        Ok(())
618    }
619
620    /// Provides the internal stream as a channel-based stream.
621    ///
622    /// # Panics
623    ///
624    /// This function panics:
625    /// - If the websocket is not connected.
626    /// - If `stream` has already been called somewhere else (stream receiver is then taken).
627    pub fn stream(&mut self) -> impl Stream<Item = BitmexWsMessage> + use<> {
628        let rx = self
629            .out_rx
630            .take()
631            .expect("Stream receiver already taken or not connected");
632        let mut rx = Arc::try_unwrap(rx).expect("Cannot take ownership - other references exist");
633        async_stream::stream! {
634            while let Some(msg) = rx.recv().await {
635                yield msg;
636            }
637        }
638    }
639
640    /// Closes the client.
641    ///
642    /// # Errors
643    ///
644    /// Returns an error if the WebSocket is not connected or if closing fails.
645    pub async fn close(&mut self) -> Result<(), BitmexWsError> {
646        log::debug!("Starting close process");
647
648        self.signal.store(true, Ordering::Relaxed);
649
650        // Send Disconnect command to handler
651        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
652            log::debug!(
653                "Failed to send disconnect command (handler may already be shut down): {e}"
654            );
655        }
656
657        // Clean up task handle with timeout
658        if let Some(task_handle) = self.task_handle.take() {
659            match Arc::try_unwrap(task_handle) {
660                Ok(handle) => {
661                    log::debug!("Waiting for task handle to complete");
662                    match tokio::time::timeout(Duration::from_secs(2), handle).await {
663                        Ok(Ok(())) => log::debug!("Task handle completed successfully"),
664                        Ok(Err(e)) => log::error!("Task handle encountered an error: {e:?}"),
665                        Err(_) => {
666                            log::warn!(
667                                "Timeout waiting for task handle, task may still be running"
668                            );
669                            // The task will be dropped and should clean up automatically
670                        }
671                    }
672                }
673                Err(arc_handle) => {
674                    log::debug!(
675                        "Cannot take ownership of task handle - other references exist, aborting task"
676                    );
677                    arc_handle.abort();
678                }
679            }
680        } else {
681            log::debug!("No task handle to await");
682        }
683
684        log::debug!("Closed");
685
686        if let Some(control) = &self.socket_control {
687            control.deregister();
688        }
689
690        Ok(())
691    }
692
693    /// Subscribe to the specified topics.
694    ///
695    /// # Errors
696    ///
697    /// Returns an error if the WebSocket is not connected or if sending the subscription message fails.
698    pub async fn subscribe(&self, topics: Vec<String>) -> Result<(), BitmexWsError> {
699        log::debug!("Subscribing to topics: {topics:?}");
700
701        for topic in &topics {
702            self.subscriptions.mark_subscribe(topic.as_str());
703            self.tracked_subscriptions.insert(topic.clone(), ());
704        }
705
706        // Serialize subscription messages
707        let mut payloads = Vec::with_capacity(topics.len());
708        for topic in &topics {
709            let message = BitmexSubscription {
710                op: BitmexWsOperation::Subscribe,
711                args: vec![Ustr::from(topic.as_ref())],
712            };
713            let payload = serde_json::to_string(&message).map_err(|e| {
714                BitmexWsError::SubscriptionError(format!("Failed to serialize subscription: {e}"))
715            })?;
716            payloads.push(payload);
717        }
718
719        // Send Subscribe command to handler
720        let cmd = HandlerCommand::Subscribe { topics: payloads };
721
722        self.send_cmd(cmd).await.map_err(|e| {
723            BitmexWsError::SubscriptionError(format!("Failed to send subscribe command: {e}"))
724        })
725    }
726
727    /// Unsubscribe from the specified topics.
728    ///
729    /// # Errors
730    ///
731    /// Returns an error if the WebSocket is not connected or if sending the unsubscription message fails.
732    async fn unsubscribe(&self, topics: Vec<String>) -> Result<(), BitmexWsError> {
733        log::debug!("Attempting to unsubscribe from topics: {topics:?}");
734
735        if self.signal.load(Ordering::Relaxed) {
736            log::debug!("Shutdown signal detected, skipping unsubscribe");
737            return Ok(());
738        }
739
740        for topic in &topics {
741            self.subscriptions.mark_unsubscribe(topic.as_str());
742            self.tracked_subscriptions.remove(topic);
743        }
744
745        // Serialize unsubscription messages
746        let mut payloads = Vec::with_capacity(topics.len());
747        for topic in &topics {
748            let message = BitmexSubscription {
749                op: BitmexWsOperation::Unsubscribe,
750                args: vec![Ustr::from(topic.as_ref())],
751            };
752
753            if let Ok(payload) = serde_json::to_string(&message) {
754                payloads.push(payload);
755            }
756        }
757
758        // Send Unsubscribe command to handler
759        let cmd = HandlerCommand::Unsubscribe { topics: payloads };
760
761        if let Err(e) = self.send_cmd(cmd).await {
762            log::debug!("Failed to send unsubscribe command: {e}");
763        }
764
765        Ok(())
766    }
767
768    /// Get the current number of active subscriptions.
769    #[must_use]
770    pub fn subscription_count(&self) -> usize {
771        self.subscriptions.len()
772    }
773
774    pub fn get_subscriptions(&self, instrument_id: InstrumentId) -> Vec<String> {
775        let symbol = instrument_id.symbol.inner();
776        let confirmed = self.subscriptions.confirmed();
777        let mut channels = Vec::with_capacity(confirmed.len());
778
779        for (channel, symbols) in confirmed.iter() {
780            if symbols.contains(&symbol) {
781                // Return the full topic string (e.g., "orderBookL2:XBTUSD")
782                channels.push(format!("{channel}:{symbol}"));
783            } else {
784                let has_channel_marker = symbols.iter().any(|s| s.is_empty());
785                if has_channel_marker
786                    && (*channel == BitmexWsAuthChannel::Execution.as_ref()
787                        || *channel == BitmexWsAuthChannel::Order.as_ref())
788                {
789                    // These are account-level subscriptions without symbols
790                    channels.push(channel.to_string());
791                }
792            }
793        }
794
795        channels
796    }
797
798    /// Subscribe to instrument updates for all instruments on the venue.
799    ///
800    /// # Errors
801    ///
802    /// Returns an error if the WebSocket is not connected or if the subscription fails.
803    pub async fn subscribe_instruments(&self) -> Result<(), BitmexWsError> {
804        // Already subscribed automatically on connection
805        log::debug!("Already subscribed to all instruments on connection, skipping");
806        Ok(())
807    }
808
809    /// Subscribe to instrument updates (mark/index prices) for the specified instrument.
810    ///
811    /// # Errors
812    ///
813    /// Returns an error if the WebSocket is not connected or if the subscription fails.
814    pub async fn subscribe_instrument(
815        &self,
816        instrument_id: InstrumentId,
817    ) -> Result<(), BitmexWsError> {
818        // Already subscribed to all instruments on connection
819        log::debug!(
820            "Already subscribed to all instruments on connection (includes {instrument_id}), skipping"
821        );
822        Ok(())
823    }
824
825    /// Subscribe to order book updates for the specified instrument.
826    ///
827    /// # Errors
828    ///
829    /// Returns an error if the WebSocket is not connected or if the subscription fails.
830    pub async fn subscribe_book(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
831        let topic = BitmexWsTopic::OrderBookL2;
832        let symbol = instrument_id.symbol.inner();
833        self.subscribe(vec![format!("{topic}:{symbol}")]).await
834    }
835
836    /// Subscribe to order book L2 (25 levels) updates for the specified instrument.
837    ///
838    /// # Errors
839    ///
840    /// Returns an error if the WebSocket is not connected or if the subscription fails.
841    pub async fn subscribe_book_25(
842        &self,
843        instrument_id: InstrumentId,
844    ) -> Result<(), BitmexWsError> {
845        let topic = BitmexWsTopic::OrderBookL2_25;
846        let symbol = instrument_id.symbol.inner();
847        self.subscribe(vec![format!("{topic}:{symbol}")]).await
848    }
849
850    /// Subscribe to order book depth 10 updates for the specified instrument.
851    ///
852    /// # Errors
853    ///
854    /// Returns an error if the WebSocket is not connected or if the subscription fails.
855    pub async fn subscribe_book_depth10(
856        &self,
857        instrument_id: InstrumentId,
858    ) -> Result<(), BitmexWsError> {
859        let topic = BitmexWsTopic::OrderBook10;
860        let symbol = instrument_id.symbol.inner();
861        self.subscribe(vec![format!("{topic}:{symbol}")]).await
862    }
863
864    /// Subscribe to quote updates for the specified instrument.
865    ///
866    /// Note: Index symbols (starting with '.') do not have quotes and will be silently ignored.
867    ///
868    /// # Errors
869    ///
870    /// Returns an error if the WebSocket is not connected or if the subscription fails.
871    pub async fn subscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
872        let symbol = instrument_id.symbol.inner();
873
874        // Index symbols don't have quotes (bid/ask), only a single price
875        if is_index_symbol(&instrument_id.symbol.inner()) {
876            log::warn!("Ignoring quote subscription for index symbol: {symbol}");
877            return Ok(());
878        }
879
880        let topic = BitmexWsTopic::Quote;
881        self.subscribe(vec![format!("{topic}:{symbol}")]).await
882    }
883
884    /// Subscribe to trade updates for the specified instrument.
885    ///
886    /// Note: Index symbols (starting with '.') do not have trades and will be silently ignored.
887    ///
888    /// # Errors
889    ///
890    /// Returns an error if the WebSocket is not connected or if the subscription fails.
891    pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
892        let symbol = instrument_id.symbol.inner();
893
894        // Index symbols don't have trades
895        if is_index_symbol(&symbol) {
896            log::warn!("Ignoring trade subscription for index symbol: {symbol}");
897            return Ok(());
898        }
899
900        let topic = BitmexWsTopic::Trade;
901        self.subscribe(vec![format!("{topic}:{symbol}")]).await
902    }
903
904    /// Subscribe to mark price updates for the specified instrument.
905    ///
906    /// # Errors
907    ///
908    /// Returns an error if the WebSocket is not connected or if the subscription fails.
909    pub async fn subscribe_mark_prices(
910        &self,
911        instrument_id: InstrumentId,
912    ) -> Result<(), BitmexWsError> {
913        self.subscribe_instrument(instrument_id).await
914    }
915
916    /// Subscribe to index price updates for the specified instrument.
917    ///
918    /// # Errors
919    ///
920    /// Returns an error if the WebSocket is not connected or if the subscription fails.
921    pub async fn subscribe_index_prices(
922        &self,
923        instrument_id: InstrumentId,
924    ) -> Result<(), BitmexWsError> {
925        self.subscribe_instrument(instrument_id).await
926    }
927
928    /// Subscribe to funding rate updates for the specified instrument.
929    ///
930    /// # Errors
931    ///
932    /// Returns an error if the WebSocket is not connected or if the subscription fails.
933    pub async fn subscribe_funding_rates(
934        &self,
935        instrument_id: InstrumentId,
936    ) -> Result<(), BitmexWsError> {
937        let topic = BitmexWsTopic::Funding;
938        let symbol = instrument_id.symbol.inner();
939        self.subscribe(vec![format!("{topic}:{symbol}")]).await
940    }
941
942    /// Subscribe to bar updates for the specified bar type.
943    ///
944    /// # Errors
945    ///
946    /// Returns an error if the WebSocket is not connected or if the subscription fails.
947    pub async fn subscribe_bars(&self, bar_type: BarType) -> Result<(), BitmexWsError> {
948        let topic = topic_from_bar_spec(bar_type.spec());
949        let symbol = bar_type.instrument_id().symbol.inner();
950        self.subscribe(vec![format!("{topic}:{symbol}")]).await
951    }
952
953    /// Unsubscribe from instrument updates for all instruments on the venue.
954    ///
955    /// # Errors
956    ///
957    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
958    pub async fn unsubscribe_instruments(&self) -> Result<(), BitmexWsError> {
959        // No-op: instruments are required for proper operation
960        log::debug!(
961            "Instruments subscription maintained for proper operation, skipping unsubscribe"
962        );
963        Ok(())
964    }
965
966    /// Unsubscribe from instrument updates (mark/index prices) for the specified instrument.
967    ///
968    /// # Errors
969    ///
970    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
971    pub async fn unsubscribe_instrument(
972        &self,
973        instrument_id: InstrumentId,
974    ) -> Result<(), BitmexWsError> {
975        // No-op: instruments are required for proper operation
976        log::debug!(
977            "Instruments subscription maintained for proper operation (includes {instrument_id}), skipping unsubscribe"
978        );
979        Ok(())
980    }
981
982    /// Unsubscribe from order book updates for the specified instrument.
983    ///
984    /// # Errors
985    ///
986    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
987    pub async fn unsubscribe_book(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
988        let topic = BitmexWsTopic::OrderBookL2;
989        let symbol = instrument_id.symbol.inner();
990        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
991    }
992
993    /// Unsubscribe from order book L2 (25 levels) updates for the specified instrument.
994    ///
995    /// # Errors
996    ///
997    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
998    pub async fn unsubscribe_book_25(
999        &self,
1000        instrument_id: InstrumentId,
1001    ) -> Result<(), BitmexWsError> {
1002        let topic = BitmexWsTopic::OrderBookL2_25;
1003        let symbol = instrument_id.symbol.inner();
1004        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1005    }
1006
1007    /// Unsubscribe from order book depth 10 updates for the specified instrument.
1008    ///
1009    /// # Errors
1010    ///
1011    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1012    pub async fn unsubscribe_book_depth10(
1013        &self,
1014        instrument_id: InstrumentId,
1015    ) -> Result<(), BitmexWsError> {
1016        let topic = BitmexWsTopic::OrderBook10;
1017        let symbol = instrument_id.symbol.inner();
1018        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1019    }
1020
1021    /// Unsubscribe from quote updates for the specified instrument.
1022    ///
1023    /// # Errors
1024    ///
1025    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1026    pub async fn unsubscribe_quotes(
1027        &self,
1028        instrument_id: InstrumentId,
1029    ) -> Result<(), BitmexWsError> {
1030        let symbol = instrument_id.symbol.inner();
1031
1032        // Index symbols don't have quotes
1033        if is_index_symbol(&symbol) {
1034            return Ok(());
1035        }
1036
1037        let topic = BitmexWsTopic::Quote;
1038        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1039    }
1040
1041    /// Unsubscribe from trade updates for the specified instrument.
1042    ///
1043    /// # Errors
1044    ///
1045    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1046    pub async fn unsubscribe_trades(
1047        &self,
1048        instrument_id: InstrumentId,
1049    ) -> Result<(), BitmexWsError> {
1050        let symbol = instrument_id.symbol.inner();
1051
1052        // Index symbols don't have trades
1053        if is_index_symbol(&symbol) {
1054            return Ok(());
1055        }
1056
1057        let topic = BitmexWsTopic::Trade;
1058        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1059    }
1060
1061    /// Unsubscribe from mark price updates for the specified instrument.
1062    ///
1063    /// # Errors
1064    ///
1065    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1066    pub async fn unsubscribe_mark_prices(
1067        &self,
1068        instrument_id: InstrumentId,
1069    ) -> Result<(), BitmexWsError> {
1070        // No-op: instrument channel shared with index prices
1071        log::debug!(
1072            "Mark prices for {instrument_id} uses shared instrument channel, skipping unsubscribe"
1073        );
1074        Ok(())
1075    }
1076
1077    /// Unsubscribe from index price updates for the specified instrument.
1078    ///
1079    /// # Errors
1080    ///
1081    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1082    pub async fn unsubscribe_index_prices(
1083        &self,
1084        instrument_id: InstrumentId,
1085    ) -> Result<(), BitmexWsError> {
1086        // No-op: instrument channel shared with mark prices
1087        log::debug!(
1088            "Index prices for {instrument_id} uses shared instrument channel, skipping unsubscribe"
1089        );
1090        Ok(())
1091    }
1092
1093    /// Unsubscribe from funding rate updates for the specified instrument.
1094    ///
1095    /// # Errors
1096    ///
1097    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1098    pub async fn unsubscribe_funding_rates(
1099        &self,
1100        instrument_id: InstrumentId,
1101    ) -> Result<(), BitmexWsError> {
1102        // No-op: unsubscribing during shutdown causes race conditions
1103        log::debug!(
1104            "Funding rates for {instrument_id}, skipping unsubscribe to avoid shutdown race"
1105        );
1106        Ok(())
1107    }
1108
1109    /// Unsubscribe from bar updates for the specified bar type.
1110    ///
1111    /// # Errors
1112    ///
1113    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1114    pub async fn unsubscribe_bars(&self, bar_type: BarType) -> Result<(), BitmexWsError> {
1115        let topic = topic_from_bar_spec(bar_type.spec());
1116        let symbol = bar_type.instrument_id().symbol.inner();
1117        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1118    }
1119
1120    /// Subscribe to order updates for the authenticated account.
1121    ///
1122    /// # Errors
1123    ///
1124    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1125    pub async fn subscribe_orders(&self) -> Result<(), BitmexWsError> {
1126        if self.credential.is_none() {
1127            return Err(BitmexWsError::MissingCredentials);
1128        }
1129        self.subscribe(vec![BitmexWsAuthChannel::Order.to_string()])
1130            .await
1131    }
1132
1133    /// Subscribe to execution updates for the authenticated account.
1134    ///
1135    /// # Errors
1136    ///
1137    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1138    pub async fn subscribe_executions(&self) -> Result<(), BitmexWsError> {
1139        if self.credential.is_none() {
1140            return Err(BitmexWsError::MissingCredentials);
1141        }
1142        self.subscribe(vec![BitmexWsAuthChannel::Execution.to_string()])
1143            .await
1144    }
1145
1146    /// Subscribe to position updates for the authenticated account.
1147    ///
1148    /// # Errors
1149    ///
1150    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1151    pub async fn subscribe_positions(&self) -> Result<(), BitmexWsError> {
1152        if self.credential.is_none() {
1153            return Err(BitmexWsError::MissingCredentials);
1154        }
1155        self.subscribe(vec![BitmexWsAuthChannel::Position.to_string()])
1156            .await
1157    }
1158
1159    /// Subscribe to margin updates for the authenticated account.
1160    ///
1161    /// # Errors
1162    ///
1163    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1164    pub async fn subscribe_margin(&self) -> Result<(), BitmexWsError> {
1165        if self.credential.is_none() {
1166            return Err(BitmexWsError::MissingCredentials);
1167        }
1168        self.subscribe(vec![BitmexWsAuthChannel::Margin.to_string()])
1169            .await
1170    }
1171
1172    /// Subscribe to wallet updates for the authenticated account.
1173    ///
1174    /// # Errors
1175    ///
1176    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1177    pub async fn subscribe_wallet(&self) -> Result<(), BitmexWsError> {
1178        if self.credential.is_none() {
1179            return Err(BitmexWsError::MissingCredentials);
1180        }
1181        self.subscribe(vec![BitmexWsAuthChannel::Wallet.to_string()])
1182            .await
1183    }
1184
1185    /// Unsubscribe from order updates for the authenticated account.
1186    ///
1187    /// # Errors
1188    ///
1189    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1190    pub async fn unsubscribe_orders(&self) -> Result<(), BitmexWsError> {
1191        self.unsubscribe(vec![BitmexWsAuthChannel::Order.to_string()])
1192            .await
1193    }
1194
1195    /// Unsubscribe from execution updates for the authenticated account.
1196    ///
1197    /// # Errors
1198    ///
1199    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1200    pub async fn unsubscribe_executions(&self) -> Result<(), BitmexWsError> {
1201        self.unsubscribe(vec![BitmexWsAuthChannel::Execution.to_string()])
1202            .await
1203    }
1204
1205    /// Unsubscribe from position updates for the authenticated account.
1206    ///
1207    /// # Errors
1208    ///
1209    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1210    pub async fn unsubscribe_positions(&self) -> Result<(), BitmexWsError> {
1211        self.unsubscribe(vec![BitmexWsAuthChannel::Position.to_string()])
1212            .await
1213    }
1214
1215    /// Unsubscribe from margin updates for the authenticated account.
1216    ///
1217    /// # Errors
1218    ///
1219    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1220    pub async fn unsubscribe_margin(&self) -> Result<(), BitmexWsError> {
1221        self.unsubscribe(vec![BitmexWsAuthChannel::Margin.to_string()])
1222            .await
1223    }
1224
1225    /// Unsubscribe from wallet updates for the authenticated account.
1226    ///
1227    /// # Errors
1228    ///
1229    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1230    pub async fn unsubscribe_wallet(&self) -> Result<(), BitmexWsError> {
1231        self.unsubscribe(vec![BitmexWsAuthChannel::Wallet.to_string()])
1232            .await
1233    }
1234
1235    /// Sends a command to the handler.
1236    async fn send_cmd(&self, cmd: HandlerCommand) -> Result<(), BitmexWsError> {
1237        self.cmd_tx
1238            .read()
1239            .await
1240            .send(cmd)
1241            .map_err(|e| BitmexWsError::ClientError(format!("Handler not available: {e}")))
1242    }
1243}
1244
1245#[cfg(test)]
1246mod tests {
1247    use rstest::rstest;
1248
1249    use super::*;
1250
1251    #[rstest]
1252    fn test_reconnect_topics_restoration_logic() {
1253        // Create real client with credentials
1254        let client = BitmexWebSocketClient::new(
1255            Some("ws://test.com".to_string()),
1256            Some("test_key".to_string()),
1257            Some("test_secret".to_string()),
1258            Some(AccountId::new("BITMEX-TEST")),
1259            5,
1260            None,
1261            TransportBackend::default(),
1262            None,
1263        )
1264        .unwrap();
1265
1266        // Populate subscriptions like they would be during normal operation
1267        for topic in [
1268            format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref()),
1269            format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref()),
1270            format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref()),
1271            BitmexWsAuthChannel::Order.as_ref().to_string(),
1272            BitmexWsAuthChannel::Position.as_ref().to_string(),
1273        ] {
1274            client.subscriptions.mark_subscribe(&topic);
1275            client.subscriptions.confirm_subscribe(&topic);
1276        }
1277
1278        // Test the actual reconnection topic building logic
1279        let topics_to_restore = client.subscriptions.all_topics();
1280
1281        // Verify it builds the correct restoration topics
1282        assert!(topics_to_restore.contains(&format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref())));
1283        assert!(topics_to_restore.contains(&format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref())));
1284        assert!(
1285            topics_to_restore.contains(&format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref()))
1286        );
1287        assert!(topics_to_restore.contains(&BitmexWsAuthChannel::Order.as_ref().to_string()));
1288        assert!(topics_to_restore.contains(&BitmexWsAuthChannel::Position.as_ref().to_string()));
1289        assert_eq!(topics_to_restore.len(), 5);
1290    }
1291
1292    #[rstest]
1293    fn test_reconnect_auth_message_building() {
1294        // Test with credentials
1295        let client_with_creds = BitmexWebSocketClient::new(
1296            Some("ws://test.com".to_string()),
1297            Some("test_key".to_string()),
1298            Some("test_secret".to_string()),
1299            Some(AccountId::new("BITMEX-TEST")),
1300            5,
1301            None,
1302            TransportBackend::default(),
1303            None,
1304        )
1305        .unwrap();
1306
1307        // Test the actual auth message building logic from lines 220-228
1308        if let Some(cred) = &client_with_creds.credential {
1309            let expires =
1310                (jiff::Timestamp::now() + jiff::SignedDuration::from_secs(30)).as_second();
1311            let signature = cred.sign("GET", "/realtime", expires, "");
1312
1313            let auth_message = BitmexAuthentication {
1314                op: BitmexWsAuthAction::AuthKeyExpires,
1315                args: (cred.api_key().to_string(), expires, signature),
1316            };
1317
1318            // Verify auth message structure
1319            assert_eq!(auth_message.op, BitmexWsAuthAction::AuthKeyExpires);
1320            assert_eq!(auth_message.args.0, "test_key");
1321            assert!(auth_message.args.1 > 0); // expires should be positive
1322            assert!(!auth_message.args.2.is_empty()); // signature should exist
1323        } else {
1324            panic!("Client should have credentials");
1325        }
1326
1327        // Test without credentials
1328        let client_no_creds = BitmexWebSocketClient::new(
1329            Some("ws://test.com".to_string()),
1330            None,
1331            None,
1332            Some(AccountId::new("BITMEX-TEST")),
1333            5,
1334            None,
1335            TransportBackend::default(),
1336            None,
1337        )
1338        .unwrap();
1339
1340        assert!(client_no_creds.credential.is_none());
1341    }
1342
1343    #[rstest]
1344    fn test_subscription_state_after_unsubscribe() {
1345        let client = BitmexWebSocketClient::new(
1346            Some("ws://test.com".to_string()),
1347            Some("test_key".to_string()),
1348            Some("test_secret".to_string()),
1349            Some(AccountId::new("BITMEX-TEST")),
1350            5,
1351            None,
1352            TransportBackend::default(),
1353            None,
1354        )
1355        .unwrap();
1356
1357        // Set up initial subscriptions
1358        for topic in [
1359            format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref()),
1360            format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref()),
1361            format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref()),
1362        ] {
1363            client.subscriptions.mark_subscribe(&topic);
1364            client.subscriptions.confirm_subscribe(&topic);
1365        }
1366
1367        // Simulate unsubscribe logic
1368        let topic = format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref());
1369        client.subscriptions.mark_unsubscribe(&topic);
1370        client.subscriptions.confirm_unsubscribe(&topic);
1371
1372        // Build restoration topics after unsubscribe
1373        let topics_to_restore = client.subscriptions.all_topics();
1374
1375        // Should have XBTUSD trade but not ETHUSD trade
1376        let trade_xbt = format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref());
1377        let trade_eth = format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref());
1378        let book_xbt = format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref());
1379
1380        assert!(topics_to_restore.contains(&trade_xbt));
1381        assert!(!topics_to_restore.contains(&trade_eth));
1382        assert!(topics_to_restore.contains(&book_xbt));
1383        assert_eq!(topics_to_restore.len(), 2);
1384    }
1385
1386    #[rstest]
1387    fn test_race_unsubscribe_failure_recovery() {
1388        // Simulates the race condition where venue rejects an unsubscribe request.
1389        // The adapter must perform the 3-step recovery:
1390        // 1. confirm_unsubscribe() - clear pending_unsubscribe
1391        // 2. mark_subscribe() - mark as subscribing again
1392        // 3. confirm_subscribe() - restore to confirmed state
1393        let client = BitmexWebSocketClient::new(
1394            Some("ws://test.com".to_string()),
1395            None,
1396            None,
1397            Some(AccountId::new("BITMEX-TEST")),
1398            5,
1399            None,
1400            TransportBackend::default(),
1401            None,
1402        )
1403        .unwrap();
1404
1405        let topic = format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref());
1406
1407        // Initial subscribe flow
1408        client.subscriptions.mark_subscribe(&topic);
1409        client.subscriptions.confirm_subscribe(&topic);
1410        assert_eq!(client.subscriptions.len(), 1);
1411
1412        // User unsubscribes
1413        client.subscriptions.mark_unsubscribe(&topic);
1414        assert_eq!(client.subscriptions.len(), 0);
1415        assert_eq!(
1416            client.subscriptions.pending_unsubscribe_topics(),
1417            vec![topic.clone()]
1418        );
1419
1420        // Venue REJECTS the unsubscribe (error message)
1421        // Adapter must perform 3-step recovery (from lines 1884-1891)
1422        client.subscriptions.confirm_unsubscribe(&topic); // Step 1: clear pending_unsubscribe
1423        client.subscriptions.mark_subscribe(&topic); // Step 2: mark as subscribing
1424        client.subscriptions.confirm_subscribe(&topic); // Step 3: confirm subscription
1425
1426        // Verify recovery: topic should be back in confirmed state
1427        assert_eq!(client.subscriptions.len(), 1);
1428        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1429        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1430
1431        // Verify topic is in all_topics() for reconnect
1432        let all = client.subscriptions.all_topics();
1433        assert_eq!(all.len(), 1);
1434        assert!(all.contains(&topic));
1435    }
1436
1437    #[rstest]
1438    fn test_race_resubscribe_before_unsubscribe_ack() {
1439        // Simulates: User unsubscribes, then immediately resubscribes before
1440        // the unsubscribe ACK arrives from the venue.
1441        // This is the race condition fixed in the subscription tracker.
1442        let client = BitmexWebSocketClient::new(
1443            Some("ws://test.com".to_string()),
1444            None,
1445            None,
1446            Some(AccountId::new("BITMEX-TEST")),
1447            5,
1448            None,
1449            TransportBackend::default(),
1450            None,
1451        )
1452        .unwrap();
1453
1454        let topic = format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref());
1455
1456        // Initial subscribe
1457        client.subscriptions.mark_subscribe(&topic);
1458        client.subscriptions.confirm_subscribe(&topic);
1459        assert_eq!(client.subscriptions.len(), 1);
1460
1461        // User unsubscribes
1462        client.subscriptions.mark_unsubscribe(&topic);
1463        assert_eq!(client.subscriptions.len(), 0);
1464        assert_eq!(
1465            client.subscriptions.pending_unsubscribe_topics(),
1466            vec![topic.clone()]
1467        );
1468
1469        // User immediately changes mind and resubscribes (before unsubscribe ACK)
1470        client.subscriptions.mark_subscribe(&topic);
1471        assert_eq!(
1472            client.subscriptions.pending_subscribe_topics(),
1473            vec![topic.clone()]
1474        );
1475
1476        // NOW the unsubscribe ACK arrives - should NOT clear pending_subscribe
1477        client.subscriptions.confirm_unsubscribe(&topic);
1478        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1479        assert_eq!(
1480            client.subscriptions.pending_subscribe_topics(),
1481            vec![topic.clone()]
1482        );
1483
1484        // Subscribe ACK arrives
1485        client.subscriptions.confirm_subscribe(&topic);
1486        assert_eq!(client.subscriptions.len(), 1);
1487        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1488
1489        // Verify final state is correct
1490        let all = client.subscriptions.all_topics();
1491        assert_eq!(all.len(), 1);
1492        assert!(all.contains(&topic));
1493    }
1494
1495    #[rstest]
1496    fn test_race_channel_level_reconnection_with_pending_states() {
1497        // Simulates reconnection with mixed pending states including channel-level subscriptions.
1498        let client = BitmexWebSocketClient::new(
1499            Some("ws://test.com".to_string()),
1500            Some("test_key".to_string()),
1501            Some("test_secret".to_string()),
1502            Some(AccountId::new("BITMEX-TEST")),
1503            5,
1504            None,
1505            TransportBackend::default(),
1506            None,
1507        )
1508        .unwrap();
1509
1510        // Set up mixed state before reconnection
1511        // Confirmed: trade:XBTUSD
1512        let trade_xbt = format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref());
1513        client.subscriptions.mark_subscribe(&trade_xbt);
1514        client.subscriptions.confirm_subscribe(&trade_xbt);
1515
1516        // Confirmed: order (channel-level, no symbol)
1517        let order_channel = BitmexWsAuthChannel::Order.as_ref();
1518        client.subscriptions.mark_subscribe(order_channel);
1519        client.subscriptions.confirm_subscribe(order_channel);
1520
1521        // Pending subscribe: trade:ETHUSD
1522        let trade_eth = format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref());
1523        client.subscriptions.mark_subscribe(&trade_eth);
1524
1525        // Pending unsubscribe: orderBookL2:XBTUSD (user cancelled)
1526        let book_xbt = format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref());
1527        client.subscriptions.mark_subscribe(&book_xbt);
1528        client.subscriptions.confirm_subscribe(&book_xbt);
1529        client.subscriptions.mark_unsubscribe(&book_xbt);
1530
1531        // Get topics for reconnection
1532        let topics_to_restore = client.subscriptions.all_topics();
1533
1534        // Should include: confirmed + pending_subscribe (NOT pending_unsubscribe)
1535        assert_eq!(topics_to_restore.len(), 3);
1536        assert!(topics_to_restore.contains(&trade_xbt));
1537        assert!(topics_to_restore.contains(&order_channel.to_string()));
1538        assert!(topics_to_restore.contains(&trade_eth));
1539        assert!(!topics_to_restore.contains(&book_xbt)); // Excluded
1540
1541        // Verify channel-level marker is handled correctly
1542        // order channel should not have ':' delimiter
1543        for topic in &topics_to_restore {
1544            if topic == order_channel {
1545                assert!(
1546                    !topic.contains(':'),
1547                    "Channel-level topic should not have delimiter"
1548                );
1549            }
1550        }
1551    }
1552}
nautilus_bitmex/websocket/client.rs

nautilus_bitmex/websocket/
client.rs
