nautilus_architect_ax/websocket/data/
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//! Market data WebSocket client for Ax.
17
18use std::{
19    fmt::{Debug, Display},
20    num::NonZeroU32,
21    sync::{
22        Arc,
23        atomic::{AtomicBool, AtomicI64, AtomicU8, Ordering},
24    },
25    time::Duration,
26};
27
28use ahash::AHashSet;
29use arc_swap::ArcSwap;
30use nautilus_core::{AtomicMap, string::secret::SecretString};
31use nautilus_live::{
32    SocketControl,
33    book::snapshot::SnapshotGate,
34    task::{SharedTaskSlot, TaskJoinOutcome},
35};
36use nautilus_network::{
37    ConnectionUrl,
38    http::create_standard_nautilus_headers,
39    mode::ConnectionMode,
40    websocket::{
41        InitialConnectRetryPolicy, PingHandler, ReconnectHeaders, SubscriptionState,
42        TransportBackend, WebSocketClient, WebSocketConfig, channel_message_handler,
43    },
44};
45use parking_lot::Mutex;
46use tokio_util::sync::CancellationToken;
47use ustr::Ustr;
48
49use super::{
50    AxMdSubscriptionSpec,
51    handler::{AxMdWsFeedHandler, HandlerCommand},
52};
53use crate::{
54    common::enums::{AxCandleWidth, AxMarketDataLevel},
55    websocket::messages::AxDataWsMessage,
56};
57
58/// Subscription topic delimiter for Ax.
59const AX_TOPIC_DELIMITER: char = ':';
60
61/// Result type for Ax WebSocket operations.
62pub type AxWsResult<T> = Result<T, AxWsClientError>;
63
64/// Error type for the Ax WebSocket client.
65#[derive(Debug, Clone)]
66pub enum AxWsClientError {
67    /// Transport/connection error.
68    Transport(String),
69    /// Channel send error.
70    ChannelError(String),
71    /// Client request error, such as a book that is no longer subscribed.
72    ClientError(String),
73    /// Operation timeout, such as a book snapshot that did not arrive.
74    OperationTimeout { timeout_ms: u64 },
75    /// Book frame that could not be converted to a snapshot.
76    InvalidSnapshot(String),
77}
78
79impl Display for AxWsClientError {
80    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
81        match self {
82            Self::Transport(msg) => write!(f, "Transport error: {msg}"),
83            Self::ChannelError(msg) => write!(f, "Channel error: {msg}"),
84            Self::ClientError(msg) => write!(f, "Client error: {msg}"),
85            Self::OperationTimeout { timeout_ms } => {
86                write!(f, "Operation timed out after {timeout_ms}ms")
87            }
88            Self::InvalidSnapshot(msg) => write!(f, "Invalid book snapshot: {msg}"),
89        }
90    }
91}
92
93impl std::error::Error for AxWsClientError {}
94
95#[derive(Debug, Default, Clone)]
96pub struct SymbolDataTypes {
97    pub quotes: bool,
98    pub trades: bool,
99    pub mark_prices: bool,
100    pub instrument_status: bool,
101    pub book_level: Option<AxMarketDataLevel>,
102}
103
104impl SymbolDataTypes {
105    fn effective_subscription(&self) -> Option<AxMdSubscriptionSpec> {
106        let book_level = self.book_level.or({
107            if self.quotes || self.ticker() {
108                Some(AxMarketDataLevel::Level1)
109            } else {
110                None
111            }
112        });
113
114        if let Some(level) = book_level {
115            return Some(self.book_subscription(level));
116        }
117
118        if self.trades {
119            return Some(AxMdSubscriptionSpec::new(
120                AxMarketDataLevel::Trades,
121                None,
122                None,
123            ));
124        }
125
126        None
127    }
128
129    fn book_subscription(&self, level: AxMarketDataLevel) -> AxMdSubscriptionSpec {
130        AxMdSubscriptionSpec::new(level, Some(self.trades), Some(self.ticker()))
131    }
132
133    fn ticker(&self) -> bool {
134        self.mark_prices || self.instrument_status
135    }
136
137    fn is_empty(&self) -> bool {
138        !self.quotes
139            && !self.trades
140            && !self.mark_prices
141            && !self.instrument_status
142            && self.book_level.is_none()
143    }
144}
145
146/// Market data WebSocket client for Ax.
147///
148/// Provides streaming market data including tickers, trades, order books, and candles.
149/// Requires Bearer token authentication obtained via the HTTP `/api/authenticate` endpoint.
150pub struct AxMdWebSocketClient {
151    url: ConnectionUrl,
152    heartbeat: Option<u64>,
153    auth_token: Arc<Mutex<Option<SecretString>>>,
154    reconnect_headers: Arc<Mutex<Option<ReconnectHeaders>>>,
155    connection_mode: Arc<ArcSwap<AtomicU8>>,
156    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
157    out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<AxDataWsMessage>>>,
158    signal: Arc<AtomicBool>,
159    cancellation_token: Arc<ArcSwap<CancellationToken>>,
160    task_handle: Arc<SharedTaskSlot<()>>,
161    connect_lock: Arc<tokio::sync::Mutex<()>>,
162    subscriptions: SubscriptionState,
163    request_id_counter: Arc<AtomicI64>,
164    subscribe_lock: Arc<tokio::sync::Mutex<()>>,
165    symbol_data_types: Arc<AtomicMap<String, SymbolDataTypes>>,
166    status_invalidations: Arc<Mutex<AHashSet<Ustr>>>,
167    transport_backend: TransportBackend,
168    proxy_url: Option<SecretString>,
169    socket_control: Option<SocketControl>,
170}
171
172impl Debug for AxMdWebSocketClient {
173    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
174        f.debug_struct(stringify!(AxMdWebSocketClient))
175            .field("url", &self.url)
176            .field("heartbeat", &self.heartbeat)
177            .field("confirmed_subscriptions", &self.subscriptions.len())
178            .finish()
179    }
180}
181
182impl Clone for AxMdWebSocketClient {
183    fn clone(&self) -> Self {
184        Self {
185            url: self.url.clone(),
186            heartbeat: self.heartbeat,
187            auth_token: Arc::clone(&self.auth_token),
188            reconnect_headers: Arc::clone(&self.reconnect_headers),
189            connection_mode: Arc::clone(&self.connection_mode),
190            cmd_tx: Arc::clone(&self.cmd_tx),
191            out_rx: None,
192            signal: Arc::clone(&self.signal),
193            cancellation_token: Arc::clone(&self.cancellation_token),
194            task_handle: Arc::clone(&self.task_handle),
195            connect_lock: Arc::clone(&self.connect_lock),
196            subscriptions: self.subscriptions.clone(),
197            subscribe_lock: Arc::clone(&self.subscribe_lock),
198            request_id_counter: Arc::clone(&self.request_id_counter),
199            symbol_data_types: Arc::clone(&self.symbol_data_types),
200            status_invalidations: Arc::clone(&self.status_invalidations),
201            transport_backend: self.transport_backend,
202            proxy_url: self.proxy_url.clone(),
203            socket_control: self.socket_control.clone(),
204        }
205    }
206}
207
208impl AxMdWebSocketClient {
209    fn initial_connect_retry_policy() -> InitialConnectRetryPolicy {
210        InitialConnectRetryPolicy {
211            max_attempts: NonZeroU32::new(5).expect("initial connect attempts must be non-zero"),
212            delay_initial: Duration::from_millis(500),
213            delay_max: Duration::from_secs(5),
214            backoff_factor: 2.0,
215            jitter_ms: 250,
216        }
217    }
218
219    /// Creates a new Ax market data WebSocket client.
220    ///
221    /// The `auth_token` is a Bearer token obtained from the HTTP `/api/authenticate` endpoint.
222    #[must_use]
223    pub fn new(
224        url: String,
225        auth_token: impl Into<SecretString>,
226        heartbeat: u64,
227        transport_backend: TransportBackend,
228        proxy_url: Option<String>,
229    ) -> Self {
230        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
231
232        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
233        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
234
235        Self {
236            url: url.into(),
237            heartbeat: Some(heartbeat),
238            auth_token: Arc::new(Mutex::new(Some(auth_token.into()))),
239            reconnect_headers: Arc::new(Mutex::new(None)),
240            connection_mode,
241            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
242            out_rx: None,
243            signal: Arc::new(AtomicBool::new(false)),
244            cancellation_token: Arc::new(ArcSwap::from_pointee(CancellationToken::new())),
245            task_handle: Arc::new(SharedTaskSlot::new()),
246            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
247            subscriptions: SubscriptionState::new(AX_TOPIC_DELIMITER),
248            request_id_counter: Arc::new(AtomicI64::new(1)),
249            subscribe_lock: Arc::new(tokio::sync::Mutex::new(())),
250            symbol_data_types: Arc::new(AtomicMap::new()),
251            status_invalidations: Arc::new(Mutex::new(AHashSet::new())),
252            transport_backend,
253            proxy_url: proxy_url.map(SecretString::from),
254            socket_control: None,
255        }
256    }
257
258    /// Creates a new Ax market data WebSocket client without authentication.
259    ///
260    /// Use [`set_auth_token`](Self::set_auth_token) to set the token before connecting.
261    #[must_use]
262    pub fn without_auth(
263        url: String,
264        heartbeat: u64,
265        transport_backend: TransportBackend,
266        proxy_url: Option<String>,
267    ) -> Self {
268        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
269
270        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
271        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
272
273        Self {
274            url: url.into(),
275            heartbeat: Some(heartbeat),
276            auth_token: Arc::new(Mutex::new(None)),
277            reconnect_headers: Arc::new(Mutex::new(None)),
278            connection_mode,
279            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
280            out_rx: None,
281            signal: Arc::new(AtomicBool::new(false)),
282            cancellation_token: Arc::new(ArcSwap::from_pointee(CancellationToken::new())),
283            task_handle: Arc::new(SharedTaskSlot::new()),
284            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
285            subscriptions: SubscriptionState::new(AX_TOPIC_DELIMITER),
286            request_id_counter: Arc::new(AtomicI64::new(1)),
287            subscribe_lock: Arc::new(tokio::sync::Mutex::new(())),
288            symbol_data_types: Arc::new(AtomicMap::new()),
289            status_invalidations: Arc::new(Mutex::new(AHashSet::new())),
290            transport_backend,
291            proxy_url: proxy_url.map(SecretString::from),
292            socket_control: None,
293        }
294    }
295
296    /// Configures socket state reporting and reconnect control.
297    #[must_use]
298    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
299        self.socket_control = Some(control);
300        self
301    }
302
303    /// Returns the WebSocket URL.
304    #[must_use]
305    pub fn url(&self) -> &str {
306        self.url.expose()
307    }
308
309    /// Sets the authentication token for subsequent connections.
310    ///
311    /// This should be called before `connect()` if authentication is required.
312    pub fn set_auth_token(&self, token: impl Into<SecretString>) {
313        *self.auth_token.lock() = Some(token.into());
314    }
315
316    /// Updates the token used by future automatic reconnect attempts.
317    ///
318    /// Updating the token does not interrupt the active WebSocket connection.
319    ///
320    /// # Errors
321    ///
322    /// Returns an error if the reconnect header cannot be updated.
323    pub fn update_auth_token(&self, token: SecretString) -> AxWsResult<()> {
324        let value = format!("Bearer {}", token.expose_secret());
325
326        if let Some(headers) = self.reconnect_headers.lock().as_ref() {
327            headers
328                .update("Authorization", &value)
329                .map_err(|e| AxWsClientError::Transport(e.to_string()))?;
330        }
331        self.set_auth_token(token);
332        Ok(())
333    }
334
335    /// Returns whether the client is currently connected and active.
336    #[must_use]
337    pub fn is_active(&self) -> bool {
338        let connection_mode_arc = self.connection_mode.load();
339        ConnectionMode::from_atomic(&connection_mode_arc).is_active()
340            && !self.signal.load(Ordering::Acquire)
341    }
342
343    /// Returns whether the client is closed.
344    #[must_use]
345    pub fn is_closed(&self) -> bool {
346        let connection_mode_arc = self.connection_mode.load();
347        ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
348            || self.signal.load(Ordering::Acquire)
349    }
350
351    /// Returns the number of confirmed subscriptions.
352    #[must_use]
353    pub fn subscription_count(&self) -> usize {
354        self.subscriptions.len()
355    }
356
357    /// Returns the symbol data types map (shared with handler).
358    #[must_use]
359    pub fn symbol_data_types(&self) -> Arc<AtomicMap<String, SymbolDataTypes>> {
360        Arc::clone(&self.symbol_data_types)
361    }
362
363    /// Returns the shared set of symbols whose instrument status cache has been invalidated.
364    pub fn status_invalidations(&self) -> Arc<Mutex<AHashSet<Ustr>>> {
365        Arc::clone(&self.status_invalidations)
366    }
367
368    fn next_request_id(&self) -> i64 {
369        self.request_id_counter.fetch_add(1, Ordering::Relaxed)
370    }
371
372    fn is_subscribed_topic(&self, topic: &str) -> bool {
373        let (channel, symbol) = topic
374            .split_once(AX_TOPIC_DELIMITER)
375            .map_or((topic, None), |(c, s)| (c, Some(s)));
376        let channel_ustr = Ustr::from(channel);
377        let symbol_ustr = symbol.map_or_else(|| Ustr::from(""), Ustr::from);
378        self.subscriptions
379            .is_subscribed(&channel_ustr, &symbol_ustr)
380    }
381
382    /// Establishes the WebSocket connection.
383    ///
384    /// # Errors
385    ///
386    /// Returns an error if the connection cannot be established or the initial handler command
387    /// cannot be sent.
388    pub async fn connect(&mut self) -> AxWsResult<()> {
389        let connect_lock = Arc::clone(&self.connect_lock);
390        let _guard = connect_lock.lock().await;
391
392        if !self.task_handle.is_empty() && !self.task_handle.is_finished() {
393            return Err(AxWsClientError::Transport(
394                "WebSocket handler is already running".to_string(),
395            ));
396        }
397
398        if let Some(outcome) = self
399            .task_handle
400            .finish(Duration::from_secs(2), Duration::from_secs(2))
401            .await
402        {
403            match outcome {
404                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
405                TaskJoinOutcome::Failed(error) => {
406                    return Err(AxWsClientError::Transport(format!(
407                        "Previous WebSocket handler failed: {error}"
408                    )));
409                }
410                TaskJoinOutcome::Incomplete => {
411                    return Err(AxWsClientError::Transport(
412                        "Previous WebSocket handler did not stop within shutdown bounds"
413                            .to_string(),
414                    ));
415                }
416            }
417        }
418
419        self.signal.store(false, Ordering::Release);
420        let cancellation_token = CancellationToken::new();
421        self.cancellation_token
422            .store(Arc::new(cancellation_token.clone()));
423
424        let (raw_handler, raw_rx) = channel_message_handler();
425
426        // No-op: ping responses are handled internally by the WebSocketClient
427        let ping_handler: PingHandler = Arc::new(move |_payload: Vec<u8>| {});
428
429        let mut headers = create_standard_nautilus_headers();
430
431        let auth_token = self.auth_token.lock().clone();
432
433        if let Some(token) = auth_token {
434            headers.push((
435                "Authorization".to_string(),
436                format!("Bearer {}", token.expose_secret()),
437            ));
438        }
439
440        let config = WebSocketConfig {
441            url: self.url.expose().to_owned(),
442            headers,
443            heartbeat_interval_secs: self.heartbeat,
444            heartbeat_payload: None, // Ax server sends heartbeats
445            connect_timeout_ms: Some(5_000),
446            reconnect_delay_initial_ms: Some(500),
447            reconnect_delay_max_ms: Some(5_000),
448            reconnect_backoff_factor: Some(1.5),
449            reconnect_jitter_ms: Some(250),
450            reconnect_max_attempts: None,
451            heartbeat_timeout_secs: None,
452            idle_timeout_ms: None,
453            writer_capacity: None,
454            backend: self.transport_backend,
455            proxy_url: self
456                .proxy_url
457                .as_ref()
458                .map(|url| url.expose_secret().to_owned()),
459            max_message_size_bytes: None,
460            max_frame_size_bytes: None,
461        };
462
463        let client = WebSocketClient::builder()
464            .config(config.clone())
465            .message_handler(raw_handler.clone())
466            .ping_handler(ping_handler.clone())
467            .initial_connect_retry_policy(Self::initial_connect_retry_policy())
468            .cancellation_token(cancellation_token)
469            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
470            .connect()
471            .await
472            .map_err(|e| {
473                AxWsClientError::Transport(format!(
474                    "WebSocket connection failed: {}: {}",
475                    self.url.origin(),
476                    e.redacted_message()
477                ))
478            })?;
479
480        self.connection_mode.store(client.connection_mode_atomic());
481        let reconnect_handle = client.reconnect_handle();
482        *self.reconnect_headers.lock() = Some(client.reconnect_headers());
483
484        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<AxDataWsMessage>();
485        self.out_rx = Some(Arc::new(out_rx));
486
487        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
488        *self.cmd_tx.write().await = cmd_tx.clone();
489
490        self.send_cmd(HandlerCommand::SetClient(client)).await?;
491
492        if !self.subscriptions.all_topics().is_empty() {
493            self.send_cmd(HandlerCommand::ReplaySubscriptions).await?;
494        }
495
496        let signal = Arc::clone(&self.signal);
497        let subscriptions = self.subscriptions.clone();
498
499        if let Err(e) = self.task_handle.spawn(async move {
500            let mut handler =
501                AxMdWsFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
502
503            while let Some(msg) = handler.next().await {
504                if matches!(msg, AxDataWsMessage::Reconnected) {
505                    log::info!("WebSocket reconnected, subscriptions will be replayed");
506                }
507
508                if out_tx.send(msg).is_err() {
509                    log::debug!("Output channel closed");
510                    break;
511                }
512            }
513
514            log::debug!("Handler loop exited");
515        }) {
516            self.out_rx = None;
517            return Err(AxWsClientError::Transport(format!(
518                "Failed to start WebSocket handler task: {e}"
519            )));
520        }
521
522        if let Some(control) = &self.socket_control {
523            control.register(move || reconnect_handle.request_reconnect());
524        }
525
526        Ok(())
527    }
528
529    /// Subscribes to order book deltas for a symbol.
530    ///
531    /// Uses reference counting so the underlying AX subscription is only
532    /// removed when all data types have been unsubscribed.
533    ///
534    /// # Errors
535    ///
536    /// Returns an error if the subscription command cannot be sent.
537    pub async fn subscribe_book_deltas(
538        &self,
539        symbol: &str,
540        level: AxMarketDataLevel,
541    ) -> AxWsResult<()> {
542        let _guard = self.subscribe_lock.lock().await;
543
544        let current = self
545            .symbol_data_types
546            .load()
547            .get(symbol)
548            .cloned()
549            .unwrap_or_default();
550
551        if current.book_level == Some(level) {
552            log::debug!("Book deltas already subscribed for {symbol} at {level:?}, skipping");
553            return Ok(());
554        }
555
556        let old_spec = current.effective_subscription();
557        let mut next = current.clone();
558        next.book_level = Some(level);
559        let new_spec = next.effective_subscription();
560
561        self.update_data_subscription(symbol, old_spec, new_spec)
562            .await?;
563
564        self.symbol_data_types.rcu(|m| {
565            let entry = m.entry(symbol.to_string()).or_default();
566            entry.book_level = Some(level);
567        });
568
569        Ok(())
570    }
571
572    /// Subscribes to order book deltas for a symbol, opening `gate` once the subscription is
573    /// written.
574    ///
575    /// A stream already open with the same options is replaced so the book starts from a fresh
576    /// snapshot. Returns without subscribing when `cancel` has already fired, and stops waiting
577    /// once it fires; a queued write still completes so the venue matches the subscription state.
578    pub(crate) async fn subscribe_book_deltas_gated(
579        &self,
580        symbol: Ustr,
581        level: AxMarketDataLevel,
582        cancel: CancellationToken,
583        gate: SnapshotGate,
584    ) -> AxWsResult<()> {
585        let (completion, receiver) = tokio::sync::oneshot::channel();
586
587        {
588            let _guard = self.subscribe_lock.lock().await;
589
590            if cancel.is_cancelled() {
591                return Ok(());
592            }
593
594            let current = self
595                .symbol_data_types
596                .load()
597                .get(symbol.as_str())
598                .cloned()
599                .unwrap_or_default();
600            let old_spec = current.effective_subscription();
601            let spec = current.book_subscription(level);
602            let replace = old_spec == Some(spec);
603
604            if let Some(old) = old_spec.filter(|_| !replace) {
605                log::debug!("Resubscribing {symbol}: {old:?} -> {spec:?}");
606                self.send_unsubscribe(symbol.as_str(), old).await?;
607            }
608
609            let topic = spec.topic(symbol.as_str());
610
611            if !replace {
612                self.subscriptions.mark_subscribe(&topic);
613            }
614
615            let command = HandlerCommand::WriteBook {
616                request_id: self.next_request_id(),
617                unsubscribe_request_id: replace.then(|| self.next_request_id()),
618                symbol,
619                spec,
620                cancel: cancel.clone(),
621                gate,
622                completion,
623            };
624
625            if let Err(e) = self.send_cmd(command).await {
626                if !replace {
627                    self.subscriptions.mark_unsubscribe(&topic);
628                }
629
630                return Err(e);
631            }
632
633            self.symbol_data_types.rcu(|m| {
634                m.entry(symbol.to_string()).or_default().book_level = Some(level);
635            });
636        }
637
638        tokio::select! {
639            biased;
640            () = cancel.cancelled() => Ok(()),
641            // The handler stopped after taking the command, so the subscription stays desired
642            result = receiver => result.map_err(|e| AxWsClientError::Transport(e.to_string()))?,
643        }
644    }
645
646    /// Replaces the symbol's market data subscription to request a fresh book snapshot, opening
647    /// `gate` once the subscription is written.
648    ///
649    /// Echoes the symbol's current subscription options and keeps the subscription desired for
650    /// reconnect replay.
651    pub(crate) async fn resubscribe_book_gated(
652        &self,
653        symbol: Ustr,
654        cancel: CancellationToken,
655        gate: SnapshotGate,
656    ) -> AxWsResult<()> {
657        let (completion, receiver) = tokio::sync::oneshot::channel();
658
659        {
660            // Serializes the subscription lookup with subscribe and unsubscribe changes
661            let _guard = self.subscribe_lock.lock().await;
662
663            if cancel.is_cancelled() {
664                return Err(AxWsClientError::ClientError(
665                    "Book recovery canceled".to_string(),
666                ));
667            }
668
669            let Some(spec) = self
670                .symbol_data_types
671                .load()
672                .get(symbol.as_str())
673                .and_then(|types| types.book_level.map(|level| types.book_subscription(level)))
674            else {
675                return Err(AxWsClientError::ClientError(format!(
676                    "Book deltas for {symbol} are no longer subscribed"
677                )));
678            };
679
680            self.send_cmd(HandlerCommand::WriteBook {
681                request_id: self.next_request_id(),
682                unsubscribe_request_id: Some(self.next_request_id()),
683                symbol,
684                spec,
685                cancel,
686                gate,
687                completion,
688            })
689            .await?;
690        }
691
692        receiver
693            .await
694            .map_err(|e| AxWsClientError::Transport(e.to_string()))?
695    }
696
697    /// Subscribes to quote data for a symbol.
698    ///
699    /// Uses reference counting so the underlying AX subscription is only
700    /// removed when all data types have been unsubscribed.
701    ///
702    /// # Errors
703    ///
704    /// Returns an error if the subscription command cannot be sent.
705    pub async fn subscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
706        let _guard = self.subscribe_lock.lock().await;
707
708        let current = self
709            .symbol_data_types
710            .load()
711            .get(symbol)
712            .cloned()
713            .unwrap_or_default();
714        let old_spec = current.effective_subscription();
715        let mut next = current.clone();
716        next.quotes = true;
717        let new_spec = next.effective_subscription();
718
719        self.update_data_subscription(symbol, old_spec, new_spec)
720            .await?;
721
722        self.symbol_data_types.rcu(|m| {
723            m.entry(symbol.to_string()).or_default().quotes = true;
724        });
725
726        Ok(())
727    }
728
729    /// Subscribes to trade data for a symbol.
730    ///
731    /// Uses reference counting so the underlying AX subscription is only
732    /// removed when all data types have been unsubscribed.
733    ///
734    /// # Errors
735    ///
736    /// Returns an error if the subscription command cannot be sent.
737    pub async fn subscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
738        let _guard = self.subscribe_lock.lock().await;
739
740        let current = self
741            .symbol_data_types
742            .load()
743            .get(symbol)
744            .cloned()
745            .unwrap_or_default();
746        let old_spec = current.effective_subscription();
747        let mut next = current.clone();
748        next.trades = true;
749        let new_spec = next.effective_subscription();
750
751        self.update_data_subscription(symbol, old_spec, new_spec)
752            .await?;
753
754        self.symbol_data_types.rcu(|m| {
755            m.entry(symbol.to_string()).or_default().trades = true;
756        });
757
758        Ok(())
759    }
760
761    /// Unsubscribes from order book deltas for a symbol.
762    ///
763    /// The underlying AX subscription is only removed when all data types
764    /// (quotes, trades, book) have been unsubscribed.
765    ///
766    /// # Errors
767    ///
768    /// Returns an error if the unsubscribe command cannot be sent.
769    pub async fn unsubscribe_book_deltas(&self, symbol: &str) -> AxWsResult<()> {
770        let _guard = self.subscribe_lock.lock().await;
771
772        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
773            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe book deltas");
774            return Ok(());
775        };
776        let old_spec = current.effective_subscription();
777        let mut next = current.clone();
778        next.book_level = None;
779        let new_spec = next.effective_subscription();
780
781        self.update_data_subscription(symbol, old_spec, new_spec)
782            .await?;
783
784        self.symbol_data_types.rcu(|m| {
785            if let Some(entry) = m.get_mut(symbol) {
786                entry.book_level = None;
787                if entry.is_empty() {
788                    m.remove(symbol);
789                }
790            }
791        });
792
793        Ok(())
794    }
795
796    /// Unsubscribes from quote data for a symbol.
797    ///
798    /// The underlying AX subscription is only removed when all data types
799    /// (quotes, trades, book) have been unsubscribed.
800    ///
801    /// # Errors
802    ///
803    /// Returns an error if the unsubscribe command cannot be sent.
804    pub async fn unsubscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
805        let _guard = self.subscribe_lock.lock().await;
806
807        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
808            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe quotes");
809            return Ok(());
810        };
811        let old_spec = current.effective_subscription();
812        let mut next = current.clone();
813        next.quotes = false;
814        let new_spec = next.effective_subscription();
815
816        self.update_data_subscription(symbol, old_spec, new_spec)
817            .await?;
818
819        self.symbol_data_types.rcu(|m| {
820            if let Some(entry) = m.get_mut(symbol) {
821                entry.quotes = false;
822                if entry.is_empty() {
823                    m.remove(symbol);
824                }
825            }
826        });
827
828        Ok(())
829    }
830
831    /// Unsubscribes from trade data for a symbol.
832    ///
833    /// The underlying AX subscription is only removed when all data types
834    /// (quotes, trades, book) have been unsubscribed.
835    ///
836    /// # Errors
837    ///
838    /// Returns an error if the unsubscribe command cannot be sent.
839    pub async fn unsubscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
840        let _guard = self.subscribe_lock.lock().await;
841
842        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
843            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe trades");
844            return Ok(());
845        };
846        let old_spec = current.effective_subscription();
847        let mut next = current.clone();
848        next.trades = false;
849        let new_spec = next.effective_subscription();
850
851        self.update_data_subscription(symbol, old_spec, new_spec)
852            .await?;
853
854        self.symbol_data_types.rcu(|m| {
855            if let Some(entry) = m.get_mut(symbol) {
856                entry.trades = false;
857                if entry.is_empty() {
858                    m.remove(symbol);
859                }
860            }
861        });
862
863        Ok(())
864    }
865
866    /// Subscribes to mark prices for a symbol.
867    ///
868    /// Ensures at least an L1 subscription so that ticker messages
869    /// (which carry the mark price field) are received.
870    ///
871    /// # Errors
872    ///
873    /// Returns an error if the subscription command cannot be sent.
874    pub async fn subscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
875        let _guard = self.subscribe_lock.lock().await;
876
877        let current = self
878            .symbol_data_types
879            .load()
880            .get(symbol)
881            .cloned()
882            .unwrap_or_default();
883        let old_spec = current.effective_subscription();
884        let mut next = current.clone();
885        next.mark_prices = true;
886        let new_spec = next.effective_subscription();
887
888        self.update_data_subscription(symbol, old_spec, new_spec)
889            .await?;
890
891        self.symbol_data_types.rcu(|m| {
892            m.entry(symbol.to_string()).or_default().mark_prices = true;
893        });
894
895        Ok(())
896    }
897
898    /// Unsubscribes from mark prices for a symbol.
899    ///
900    /// The underlying AX subscription is only removed when all data types
901    /// have been unsubscribed.
902    ///
903    /// # Errors
904    ///
905    /// Returns an error if the unsubscribe command cannot be sent.
906    pub async fn unsubscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
907        let _guard = self.subscribe_lock.lock().await;
908
909        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
910            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe mark prices");
911            return Ok(());
912        };
913        let old_spec = current.effective_subscription();
914        let mut next = current.clone();
915        next.mark_prices = false;
916        let new_spec = next.effective_subscription();
917
918        self.update_data_subscription(symbol, old_spec, new_spec)
919            .await?;
920
921        self.symbol_data_types.rcu(|m| {
922            if let Some(entry) = m.get_mut(symbol) {
923                entry.mark_prices = false;
924                if entry.is_empty() {
925                    m.remove(symbol);
926                }
927            }
928        });
929
930        Ok(())
931    }
932
933    /// Subscribes to instrument status for a symbol.
934    ///
935    /// Ensures at least an L1 subscription so that ticker messages
936    /// (which carry the instrument state field) are received.
937    ///
938    /// # Errors
939    ///
940    /// Returns an error if the subscription command cannot be sent.
941    pub async fn subscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
942        let _guard = self.subscribe_lock.lock().await;
943
944        let current = self
945            .symbol_data_types
946            .load()
947            .get(symbol)
948            .cloned()
949            .unwrap_or_default();
950        let old_spec = current.effective_subscription();
951        let mut next = current.clone();
952        next.instrument_status = true;
953        let new_spec = next.effective_subscription();
954
955        self.update_data_subscription(symbol, old_spec, new_spec)
956            .await?;
957
958        self.symbol_data_types.rcu(|m| {
959            m.entry(symbol.to_string()).or_default().instrument_status = true;
960        });
961
962        Ok(())
963    }
964
965    /// Unsubscribes from instrument status for a symbol.
966    ///
967    /// The underlying AX subscription is only removed when all data types
968    /// have been unsubscribed.
969    ///
970    /// # Errors
971    ///
972    /// Returns an error if the unsubscribe command cannot be sent.
973    pub async fn unsubscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
974        let _guard = self.subscribe_lock.lock().await;
975
976        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
977            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe instrument status");
978            return Ok(());
979        };
980        let old_spec = current.effective_subscription();
981        let mut next = current.clone();
982        next.instrument_status = false;
983        let new_spec = next.effective_subscription();
984
985        self.update_data_subscription(symbol, old_spec, new_spec)
986            .await?;
987
988        self.symbol_data_types.rcu(|m| {
989            if let Some(entry) = m.get_mut(symbol) {
990                entry.instrument_status = false;
991                if entry.is_empty() {
992                    m.remove(symbol);
993                }
994            }
995        });
996
997        self.status_invalidations.lock().insert(Ustr::from(symbol));
998
999        Ok(())
1000    }
1001
1002    async fn update_data_subscription(
1003        &self,
1004        symbol: &str,
1005        old_spec: Option<AxMdSubscriptionSpec>,
1006        new_spec: Option<AxMdSubscriptionSpec>,
1007    ) -> AxWsResult<()> {
1008        if old_spec == new_spec {
1009            return Ok(());
1010        }
1011
1012        match (old_spec, new_spec) {
1013            (None, Some(spec)) => {
1014                log::debug!("Subscribing {symbol} at {spec:?}");
1015                self.send_subscribe(symbol, spec).await
1016            }
1017            (Some(old), None) => {
1018                log::debug!("Unsubscribing {symbol} (no remaining data types)");
1019                self.send_unsubscribe(symbol, old).await
1020            }
1021            (Some(old), Some(new)) => {
1022                log::debug!("Resubscribing {symbol}: {old:?} -> {new:?}");
1023                self.send_unsubscribe(symbol, old).await?;
1024                if let Err(e) = self.send_subscribe(symbol, new).await {
1025                    log::warn!("Resubscribe failed for {symbol} at {new:?}: {e}");
1026                    if let Err(restore_err) = self.send_subscribe(symbol, old).await {
1027                        log::error!(
1028                            "Failed to restore {symbol} at {old:?}: {restore_err}, \
1029                             reconnection required"
1030                        );
1031                        self.subscriptions.mark_subscribe(&old.topic(symbol));
1032                    }
1033                    return Err(e);
1034                }
1035                Ok(())
1036            }
1037            (None, None) => Ok(()),
1038        }
1039    }
1040
1041    async fn send_subscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
1042        let topic = spec.topic(symbol);
1043        let request_id = self.next_request_id();
1044
1045        self.subscriptions.mark_subscribe(&topic);
1046
1047        if let Err(e) = self
1048            .send_cmd(HandlerCommand::Subscribe {
1049                request_id,
1050                symbol: Ustr::from(symbol),
1051                spec,
1052            })
1053            .await
1054        {
1055            self.subscriptions.mark_unsubscribe(&topic);
1056            return Err(e);
1057        }
1058
1059        Ok(())
1060    }
1061
1062    async fn send_unsubscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
1063        let request_id = self.next_request_id();
1064        let topic = spec.topic(symbol);
1065        let was_pending = self
1066            .subscriptions
1067            .pending_subscribe_topics()
1068            .contains(&topic);
1069
1070        self.subscriptions.mark_unsubscribe(&topic);
1071
1072        if let Err(e) = self
1073            .send_cmd(HandlerCommand::Unsubscribe {
1074                request_id,
1075                symbol: Ustr::from(symbol),
1076                topic: topic.clone(),
1077            })
1078            .await
1079        {
1080            self.restore_unsubscribe_state(&topic, was_pending);
1081            return Err(e);
1082        }
1083
1084        Ok(())
1085    }
1086
1087    /// Subscribes to candle data for a symbol.
1088    ///
1089    /// Skips sending if already subscribed or subscription is pending.
1090    ///
1091    /// # Errors
1092    ///
1093    /// Returns an error if the subscription command cannot be sent.
1094    pub async fn subscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
1095        let _guard = self.subscribe_lock.lock().await;
1096        let topic = format!("candles:{symbol}:{width:?}");
1097
1098        // Skip if already subscribed or pending
1099        if self.is_subscribed_topic(&topic) {
1100            log::debug!("Already subscribed to {topic}, skipping");
1101            return Ok(());
1102        }
1103
1104        let request_id = self.next_request_id();
1105
1106        // Mark pending BEFORE sending to prevent race conditions with concurrent subscribes
1107        self.subscriptions.mark_subscribe(&topic);
1108
1109        if let Err(e) = self
1110            .send_cmd(HandlerCommand::SubscribeCandles {
1111                request_id,
1112                symbol: Ustr::from(symbol),
1113                width,
1114            })
1115            .await
1116        {
1117            // Rollback pending state on send failure
1118            self.subscriptions.mark_unsubscribe(&topic);
1119            return Err(e);
1120        }
1121
1122        Ok(())
1123    }
1124
1125    /// Unsubscribes from candle data for a symbol.
1126    ///
1127    /// # Errors
1128    ///
1129    /// Returns an error if the unsubscribe command cannot be sent.
1130    pub async fn unsubscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
1131        let _guard = self.subscribe_lock.lock().await;
1132        let request_id = self.next_request_id();
1133        let topic = format!("candles:{symbol}:{width:?}");
1134        let was_pending = self
1135            .subscriptions
1136            .pending_subscribe_topics()
1137            .contains(&topic);
1138
1139        if !self.is_subscribed_topic(&topic) {
1140            log::debug!("Not subscribed to {topic}, skipping unsubscribe");
1141            return Ok(());
1142        }
1143
1144        self.subscriptions.mark_unsubscribe(&topic);
1145
1146        if let Err(e) = self
1147            .send_cmd(HandlerCommand::UnsubscribeCandles {
1148                request_id,
1149                symbol: Ustr::from(symbol),
1150                width,
1151                topic: topic.clone(),
1152            })
1153            .await
1154        {
1155            self.restore_unsubscribe_state(&topic, was_pending);
1156            return Err(e);
1157        }
1158
1159        Ok(())
1160    }
1161
1162    fn restore_unsubscribe_state(&self, topic: &str, was_pending: bool) {
1163        self.subscriptions.confirm_unsubscribe(topic);
1164        self.subscriptions.mark_subscribe(topic);
1165        if !was_pending {
1166            self.subscriptions.confirm_subscribe(topic);
1167        }
1168    }
1169
1170    /// Returns a stream of WebSocket messages.
1171    ///
1172    /// # Panics
1173    ///
1174    /// Panics if called before `connect()` or if the stream has already been taken.
1175    pub fn stream(&mut self) -> impl futures_util::Stream<Item = AxDataWsMessage> + 'static {
1176        let rx = self
1177            .out_rx
1178            .take()
1179            .expect("Stream receiver already taken or client not connected - stream() can only be called once");
1180        let mut rx = Arc::try_unwrap(rx).expect(
1181            "Cannot take ownership of stream - client was cloned and other references exist",
1182        );
1183        async_stream::stream! {
1184            while let Some(msg) = rx.recv().await {
1185                yield msg;
1186            }
1187        }
1188    }
1189
1190    pub(crate) fn begin_shutdown(&self) {
1191        self.cancellation_token.load().cancel();
1192        self.signal.store(true, Ordering::Release);
1193    }
1194
1195    /// Disconnects the WebSocket connection gracefully.
1196    pub async fn disconnect(&self) {
1197        log::debug!("Disconnecting WebSocket");
1198        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1199    }
1200
1201    /// Closes the WebSocket connection and cleans up resources.
1202    ///
1203    /// # Errors
1204    ///
1205    /// Returns an error if the handler task fails or does not stop after abort.
1206    pub async fn close(&mut self) -> anyhow::Result<()> {
1207        let connect_lock = Arc::clone(&self.connect_lock);
1208        let _guard = connect_lock.lock().await;
1209        log::debug!("Closing WebSocket client");
1210
1211        // Send disconnect first to allow graceful cleanup before signal
1212        self.cancellation_token.load().cancel();
1213        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1214        tokio::time::sleep(Duration::from_millis(50)).await;
1215        self.signal.store(true, Ordering::Release);
1216
1217        let outcome = self
1218            .task_handle
1219            .finish(Duration::from_secs(2), Duration::from_secs(2))
1220            .await;
1221
1222        *self.reconnect_headers.lock() = None;
1223
1224        if let Some(control) = &self.socket_control {
1225            control.deregister();
1226        }
1227
1228        match outcome {
1229            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
1230            Some(TaskJoinOutcome::Failed(error)) => Err(anyhow::anyhow!(
1231                "Architect AX data WebSocket handler failed: {error}"
1232            )),
1233            Some(TaskJoinOutcome::Incomplete) => Err(anyhow::anyhow!(
1234                "Architect AX data WebSocket handler did not stop after abort"
1235            )),
1236        }
1237    }
1238
1239    async fn send_cmd(&self, cmd: HandlerCommand) -> AxWsResult<()> {
1240        let guard = self.cmd_tx.read().await;
1241        guard
1242            .send(cmd)
1243            .map_err(|e| AxWsClientError::ChannelError(e.to_string()))
1244    }
1245}
1246
1247impl Drop for AxMdWebSocketClient {
1248    fn drop(&mut self) {
1249        if Arc::strong_count(&self.task_handle) == 1 && !self.task_handle.is_empty() {
1250            self.cancellation_token.load().cancel();
1251            self.signal.store(true, Ordering::Release);
1252            self.task_handle.abort();
1253
1254            if let Some(control) = &self.socket_control {
1255                control.deregister();
1256            }
1257        }
1258    }
1259}
1260
1261#[cfg(test)]
1262mod tests {
1263    use nautilus_live::task::TaskGroup;
1264    use nautilus_model::identifiers::InstrumentId;
1265    use rstest::rstest;
1266
1267    use super::*;
1268    use crate::{
1269        book::{recovery::spawn_subscription_task, sync::BookSyncTracker},
1270        websocket::data::subscription::SubscriptionOrder,
1271    };
1272
1273    #[rstest]
1274    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
1275    #[case(
1276        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
1277        "wss://stream.example.com:9443"
1278    )]
1279    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
1280        let client = AxMdWebSocketClient::new(
1281            url.to_string(),
1282            "SYN-auth-token",
1283            30,
1284            TransportBackend::default(),
1285            None,
1286        );
1287        let debug = format!("{client:?}");
1288        assert!(debug.contains(&format!("url: {origin:?}")));
1289        assert!(!debug.contains("SYN-"));
1290        assert_eq!(client.url.expose(), url);
1291    }
1292
1293    #[rstest]
1294    fn test_auth_token_uses_secret_string_owner() {
1295        let client = AxMdWebSocketClient::new(
1296            "ws://localhost:9999/md/ws".to_string(),
1297            "initial-token".to_string(),
1298            30,
1299            TransportBackend::default(),
1300            None,
1301        );
1302
1303        client.set_auth_token("replacement-token".to_string());
1304
1305        let token = client.auth_token.lock();
1306        assert_eq!(
1307            token.as_ref().map(SecretString::expose_secret),
1308            Some("replacement-token")
1309        );
1310    }
1311
1312    #[tokio::test]
1313    async fn test_drop_aborts_handler_task() {
1314        let client = AxMdWebSocketClient::new(
1315            "ws://localhost:9999/md/ws".to_string(),
1316            "test_token".to_string(),
1317            30,
1318            TransportBackend::default(),
1319            None,
1320        );
1321        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1322        let handle = tokio::spawn(async move {
1323            started_tx.send(()).expect("started receiver");
1324            std::future::pending::<()>().await;
1325        });
1326        let abort_handle = handle.abort_handle();
1327        client.task_handle.insert(handle);
1328        started_rx.await.expect("handler task started");
1329
1330        drop(client);
1331
1332        tokio::time::timeout(Duration::from_secs(1), async {
1333            while !abort_handle.is_finished() {
1334                tokio::task::yield_now().await;
1335            }
1336        })
1337        .await
1338        .expect("handler task aborted");
1339    }
1340
1341    #[rstest]
1342    fn test_effective_subscription_empty_returns_none() {
1343        let sdt = SymbolDataTypes::default();
1344        assert_eq!(sdt.effective_subscription(), None);
1345        assert!(sdt.is_empty());
1346    }
1347
1348    #[rstest]
1349    fn test_effective_subscription_book_level_takes_precedence() {
1350        let sdt = SymbolDataTypes {
1351            book_level: Some(AxMarketDataLevel::Level2),
1352            quotes: true,
1353            ..Default::default()
1354        };
1355        assert_eq!(
1356            sdt.effective_subscription(),
1357            Some(AxMdSubscriptionSpec::new(
1358                AxMarketDataLevel::Level2,
1359                Some(false),
1360                Some(false),
1361            ))
1362        );
1363        assert!(!sdt.is_empty());
1364    }
1365
1366    #[rstest]
1367    #[case(
1368        true,
1369        false,
1370        false,
1371        false,
1372        AxMarketDataLevel::Level1,
1373        Some(false),
1374        Some(false)
1375    )]
1376    #[case(false, true, false, false, AxMarketDataLevel::Trades, None, None)]
1377    #[case(
1378        false,
1379        false,
1380        true,
1381        false,
1382        AxMarketDataLevel::Level1,
1383        Some(false),
1384        Some(true)
1385    )]
1386    #[case(
1387        false,
1388        false,
1389        false,
1390        true,
1391        AxMarketDataLevel::Level1,
1392        Some(false),
1393        Some(true)
1394    )]
1395    fn test_effective_subscription_for_single_data_type(
1396        #[case] quotes: bool,
1397        #[case] trades: bool,
1398        #[case] mark_prices: bool,
1399        #[case] instrument_status: bool,
1400        #[case] level: AxMarketDataLevel,
1401        #[case] include_trades: Option<bool>,
1402        #[case] include_ticker: Option<bool>,
1403    ) {
1404        let sdt = SymbolDataTypes {
1405            quotes,
1406            trades,
1407            mark_prices,
1408            instrument_status,
1409            book_level: None,
1410        };
1411        assert_eq!(
1412            sdt.effective_subscription(),
1413            Some(AxMdSubscriptionSpec::new(
1414                level,
1415                include_trades,
1416                include_ticker,
1417            ))
1418        );
1419        assert!(!sdt.is_empty());
1420    }
1421
1422    #[rstest]
1423    #[case(false)]
1424    #[case(true)]
1425    #[tokio::test]
1426    async fn test_unsubscribe_send_failure_restores_subscription(#[case] was_pending: bool) {
1427        let client = AxMdWebSocketClient::new(
1428            "ws://localhost:9999/md/ws".to_string(),
1429            "test_token".to_string(),
1430            30,
1431            TransportBackend::default(),
1432            None,
1433        );
1434        let symbol = "EURUSD-PERP";
1435        let spec = AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(false));
1436        let topic = spec.topic(symbol);
1437        client.subscriptions.mark_subscribe(&topic);
1438        if !was_pending {
1439            client.subscriptions.confirm_subscribe(&topic);
1440        }
1441
1442        let error = client.send_unsubscribe(symbol, spec).await.unwrap_err();
1443
1444        assert_eq!(error.to_string(), "Channel error: channel closed");
1445        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1446        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1447        assert_eq!(
1448            client.subscriptions.pending_subscribe_topics().len(),
1449            usize::from(was_pending)
1450        );
1451        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1452    }
1453
1454    #[rstest]
1455    #[case(false)]
1456    #[case(true)]
1457    #[tokio::test]
1458    async fn test_unsubscribe_candles_send_failure_restores_subscription(
1459        #[case] was_pending: bool,
1460    ) {
1461        let client = AxMdWebSocketClient::new(
1462            "ws://localhost:9999/md/ws".to_string(),
1463            "test_token".to_string(),
1464            30,
1465            TransportBackend::default(),
1466            None,
1467        );
1468        let symbol = "EURUSD-PERP";
1469        let width = AxCandleWidth::Minutes1;
1470        let topic = format!("candles:{symbol}:{width:?}");
1471        client.subscriptions.mark_subscribe(&topic);
1472        if !was_pending {
1473            client.subscriptions.confirm_subscribe(&topic);
1474        }
1475
1476        let error = client.unsubscribe_candles(symbol, width).await.unwrap_err();
1477
1478        assert_eq!(error.to_string(), "Channel error: channel closed");
1479        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1480        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1481        assert_eq!(
1482            client.subscriptions.pending_subscribe_topics().len(),
1483            usize::from(was_pending)
1484        );
1485        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1486    }
1487
1488    #[rstest]
1489    #[tokio::test]
1490    async fn test_unsubscribe_candles_skips_untracked_topic() {
1491        let client = AxMdWebSocketClient::new(
1492            "ws://localhost:9999/md/ws".to_string(),
1493            "test_token".to_string(),
1494            30,
1495            TransportBackend::default(),
1496            None,
1497        );
1498
1499        client
1500            .unsubscribe_candles("EURUSD-PERP", AxCandleWidth::Minutes1)
1501            .await
1502            .unwrap();
1503
1504        assert_eq!(client.subscription_count(), 0);
1505        assert!(client.subscriptions.all_topics().is_empty());
1506        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1507        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1508    }
1509
1510    const SYMBOL: &str = "EURUSD-PERP";
1511
1512    // Returns a client whose handler commands the test receives, with `current` already streaming
1513    async fn book_client(
1514        current: SymbolDataTypes,
1515    ) -> (
1516        AxMdWebSocketClient,
1517        tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
1518    ) {
1519        let client = AxMdWebSocketClient::new(
1520            "ws://localhost:9999/md/ws".to_string(),
1521            "test_token".to_string(),
1522            30,
1523            TransportBackend::default(),
1524            None,
1525        );
1526        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1527        *client.cmd_tx.write().await = cmd_tx;
1528
1529        if let Some(spec) = current.effective_subscription() {
1530            let topic = spec.topic(SYMBOL);
1531            client.subscriptions.mark_subscribe(&topic);
1532            client.subscriptions.confirm_subscribe(&topic);
1533            client.symbol_data_types.insert(SYMBOL.to_string(), current);
1534        }
1535
1536        (client, cmd_rx)
1537    }
1538
1539    #[rstest]
1540    #[case::no_stream(SymbolDataTypes::default(), None, false)]
1541    #[case::quote_stream(
1542        SymbolDataTypes { quotes: true, ..Default::default() },
1543        Some("EURUSD-PERP:Level1:false:false"),
1544        false
1545    )]
1546    #[case::book_stream(
1547        SymbolDataTypes { book_level: Some(AxMarketDataLevel::Level2), ..Default::default() },
1548        None,
1549        true
1550    )]
1551    #[tokio::test]
1552    async fn test_subscribe_book_deltas_gated_writes_book_stream(
1553        #[case] current: SymbolDataTypes,
1554        #[case] unsubscribed: Option<&str>,
1555        #[case] replaced: bool,
1556    ) {
1557        let quotes = current.quotes;
1558        let (client, mut cmd_rx) = book_client(current).await;
1559        let gate = SnapshotGate::default();
1560        gate.lock().close();
1561
1562        let subscribe = client.subscribe_book_deltas_gated(
1563            Ustr::from(SYMBOL),
1564            AxMarketDataLevel::Level2,
1565            CancellationToken::new(),
1566            gate,
1567        );
1568
1569        let handler = async {
1570            let mut topics = Vec::new();
1571
1572            loop {
1573                match cmd_rx.recv().await.expect("command queued") {
1574                    HandlerCommand::Unsubscribe { topic, .. } => topics.push(topic),
1575                    HandlerCommand::WriteBook {
1576                        spec,
1577                        unsubscribe_request_id,
1578                        completion,
1579                        ..
1580                    } => {
1581                        completion.send(Ok(())).unwrap();
1582                        return (topics, spec, unsubscribe_request_id.is_some());
1583                    }
1584                    other => panic!("unexpected command {other:?}"),
1585                }
1586            }
1587        };
1588
1589        let (result, (topics, spec, replaces)) = tokio::join!(subscribe, handler);
1590        let types = client
1591            .symbol_data_types
1592            .load()
1593            .get(SYMBOL)
1594            .cloned()
1595            .unwrap();
1596
1597        assert!(result.is_ok());
1598        assert_eq!(
1599            topics,
1600            unsubscribed
1601                .map(ToString::to_string)
1602                .into_iter()
1603                .collect::<Vec<_>>()
1604        );
1605        assert_eq!(
1606            spec,
1607            AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(false))
1608        );
1609        assert_eq!(replaces, replaced);
1610        assert_eq!(types.book_level, Some(AxMarketDataLevel::Level2));
1611        assert_eq!(types.quotes, quotes);
1612        assert_eq!(client.subscriptions.all_topics(), vec![spec.topic(SYMBOL)]);
1613    }
1614
1615    #[rstest]
1616    #[tokio::test]
1617    async fn test_book_subscription_write_failure_starts_recovery() {
1618        let (client, mut cmd_rx) = book_client(SymbolDataTypes::default()).await;
1619        let book_sync = BookSyncTracker::default();
1620        let tasks = TaskGroup::new();
1621        let instrument_id = InstrumentId::from("EURUSD-PERP.AX");
1622
1623        spawn_subscription_task(
1624            instrument_id,
1625            AxMarketDataLevel::Level2,
1626            SubscriptionOrder::default().next(),
1627            book_sync.clone(),
1628            client.clone(),
1629            Duration::from_secs(10),
1630            &tasks.spawner().unwrap(),
1631        );
1632
1633        let mut replaces = Vec::new();
1634
1635        for _ in 0..2 {
1636            let command = tokio::time::timeout(Duration::from_secs(5), cmd_rx.recv())
1637                .await
1638                .expect("book write queued")
1639                .expect("command channel open");
1640
1641            let HandlerCommand::WriteBook {
1642                unsubscribe_request_id,
1643                completion,
1644                ..
1645            } = command
1646            else {
1647                panic!("expected a book write, was {command:?}");
1648            };
1649
1650            replaces.push(unsubscribe_request_id.is_some());
1651
1652            // The initial write fails in transport, which recovery retries
1653            let _ = completion.send(Err(AxWsClientError::Transport("write failed".to_string())));
1654        }
1655
1656        let owned = book_sync.claim_recovery(instrument_id).is_none();
1657        tasks.begin_shutdown();
1658
1659        assert_eq!(replaces, [false, true]);
1660        assert!(owned, "the running recovery owns the book");
1661    }
1662
1663    #[rstest]
1664    #[tokio::test]
1665    async fn test_resubscribe_book_gated_echoes_current_subscription() {
1666        let current = SymbolDataTypes {
1667            mark_prices: true,
1668            book_level: Some(AxMarketDataLevel::Level2),
1669            ..Default::default()
1670        };
1671
1672        let (client, mut cmd_rx) = book_client(current).await;
1673
1674        let resubscribe = client.resubscribe_book_gated(
1675            Ustr::from(SYMBOL),
1676            CancellationToken::new(),
1677            SnapshotGate::default(),
1678        );
1679
1680        let handler = async {
1681            match cmd_rx.recv().await.expect("command queued") {
1682                HandlerCommand::WriteBook {
1683                    request_id,
1684                    unsubscribe_request_id,
1685                    symbol,
1686                    spec,
1687                    completion,
1688                    ..
1689                } => {
1690                    completion.send(Ok(())).unwrap();
1691                    (request_id, unsubscribe_request_id, symbol, spec)
1692                }
1693                other => panic!("unexpected command {other:?}"),
1694            }
1695        };
1696
1697        let (result, (request_id, unsubscribe_request_id, symbol, spec)) =
1698            tokio::join!(resubscribe, handler);
1699
1700        assert!(result.is_ok());
1701        assert_eq!(request_id, 1);
1702        assert_eq!(unsubscribe_request_id, Some(2));
1703        assert_eq!(symbol, Ustr::from(SYMBOL));
1704        assert_eq!(
1705            spec,
1706            AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(true))
1707        );
1708        assert!(cmd_rx.try_recv().is_err());
1709    }
1710
1711    #[rstest]
1712    #[case::not_subscribed(
1713        SymbolDataTypes { quotes: true, ..Default::default() },
1714        false,
1715        "Client error: Book deltas for EURUSD-PERP are no longer subscribed"
1716    )]
1717    #[case::cancelled(
1718        SymbolDataTypes { book_level: Some(AxMarketDataLevel::Level2), ..Default::default() },
1719        true,
1720        "Client error: Book recovery canceled"
1721    )]
1722    #[tokio::test]
1723    async fn test_resubscribe_book_gated_refuses_without_write(
1724        #[case] current: SymbolDataTypes,
1725        #[case] cancelled: bool,
1726        #[case] expected: &str,
1727    ) {
1728        let (client, mut cmd_rx) = book_client(current).await;
1729        let cancel = CancellationToken::new();
1730
1731        if cancelled {
1732            cancel.cancel();
1733        }
1734
1735        let error = client
1736            .resubscribe_book_gated(Ustr::from(SYMBOL), cancel, SnapshotGate::default())
1737            .await
1738            .unwrap_err();
1739
1740        assert_eq!(error.to_string(), expected);
1741        assert!(cmd_rx.try_recv().is_err());
1742    }
1743}
nautilus_architect_ax/websocket/data/client.rs

nautilus_architect_ax/websocket/data/
client.rs
