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

nautilus_bitmex/websocket/
client.rs
