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        let connection_mode = 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        cmd_tx
489            .send(HandlerCommand::SetClient(client))
490            .map_err(|e| AxWsClientError::ChannelError(e.to_string()))?;
491
492        if !self.subscriptions.all_topics().is_empty() {
493            cmd_tx
494                .send(HandlerCommand::ReplaySubscriptions)
495                .map_err(|e| AxWsClientError::ChannelError(e.to_string()))?;
496        }
497
498        {
499            let mut command_sender = self.cmd_tx.write().await;
500            *command_sender = cmd_tx.clone();
501            self.connection_mode.store(connection_mode);
502        }
503
504        let signal = Arc::clone(&self.signal);
505        let subscriptions = self.subscriptions.clone();
506
507        if let Err(e) = self.task_handle.spawn(async move {
508            let mut handler =
509                AxMdWsFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
510
511            while let Some(msg) = handler.next().await {
512                if matches!(msg, AxDataWsMessage::Reconnected) {
513                    log::info!("WebSocket reconnected, subscriptions will be replayed");
514                }
515
516                if out_tx.send(msg).is_err() {
517                    log::debug!("Output channel closed");
518                    break;
519                }
520            }
521
522            log::debug!("Handler loop exited");
523        }) {
524            self.out_rx = None;
525            return Err(AxWsClientError::Transport(format!(
526                "Failed to start WebSocket handler task: {e}"
527            )));
528        }
529
530        if let Some(control) = &self.socket_control {
531            control.register(move || reconnect_handle.request_reconnect());
532        }
533
534        Ok(())
535    }
536
537    /// Subscribes to order book deltas for a symbol.
538    ///
539    /// Uses reference counting so the underlying AX subscription is only
540    /// removed when all data types have been unsubscribed.
541    ///
542    /// # Errors
543    ///
544    /// Returns an error if the subscription command cannot be sent.
545    pub async fn subscribe_book_deltas(
546        &self,
547        symbol: &str,
548        level: AxMarketDataLevel,
549    ) -> AxWsResult<()> {
550        let _guard = self.subscribe_lock.lock().await;
551
552        let current = self
553            .symbol_data_types
554            .load()
555            .get(symbol)
556            .cloned()
557            .unwrap_or_default();
558
559        if current.book_level == Some(level) {
560            log::debug!("Book deltas already subscribed for {symbol} at {level:?}, skipping");
561            return Ok(());
562        }
563
564        let old_spec = current.effective_subscription();
565        let mut next = current.clone();
566        next.book_level = Some(level);
567        let new_spec = next.effective_subscription();
568
569        self.update_data_subscription(symbol, old_spec, new_spec)
570            .await?;
571
572        self.symbol_data_types.rcu(|m| {
573            let entry = m.entry(symbol.to_string()).or_default();
574            entry.book_level = Some(level);
575        });
576
577        Ok(())
578    }
579
580    /// Subscribes to order book deltas for a symbol, opening `gate` once the subscription is
581    /// written.
582    ///
583    /// A stream already open with the same options is replaced so the book starts from a fresh
584    /// snapshot. Returns without subscribing when `cancel` has already fired, and stops waiting
585    /// once it fires. Cancellation skips the book write's unsent work; a write already admitted to the
586    /// network can still complete before later subscription changes.
587    pub(crate) async fn subscribe_book_deltas_gated(
588        &self,
589        symbol: Ustr,
590        level: AxMarketDataLevel,
591        cancel: CancellationToken,
592        gate: SnapshotGate,
593    ) -> AxWsResult<()> {
594        let (completion, receiver) = tokio::sync::oneshot::channel();
595
596        {
597            let _guard = self.subscribe_lock.lock().await;
598
599            if cancel.is_cancelled() {
600                return Ok(());
601            }
602
603            let current = self
604                .symbol_data_types
605                .load()
606                .get(symbol.as_str())
607                .cloned()
608                .unwrap_or_default();
609            let old_spec = current.effective_subscription();
610            let spec = current.book_subscription(level);
611            let replace = old_spec == Some(spec);
612
613            if let Some(old) = old_spec.filter(|_| !replace) {
614                log::debug!("Resubscribing {symbol}: {old:?} -> {spec:?}");
615                self.send_unsubscribe(symbol.as_str(), old).await?;
616            }
617
618            let topic = spec.topic(symbol.as_str());
619
620            if !replace {
621                self.subscriptions.mark_subscribe(&topic);
622            }
623
624            let command = HandlerCommand::WriteBook {
625                request_id: self.next_request_id(),
626                unsubscribe_request_id: replace.then(|| self.next_request_id()),
627                symbol,
628                spec,
629                cancel: cancel.clone(),
630                gate,
631                completion,
632            };
633
634            if let Err(e) = self.send_cmd(command).await {
635                if !replace {
636                    self.subscriptions.mark_unsubscribe(&topic);
637                }
638
639                return Err(e);
640            }
641
642            self.symbol_data_types.rcu(|m| {
643                m.entry(symbol.to_string()).or_default().book_level = Some(level);
644            });
645        }
646
647        tokio::select! {
648            biased;
649            () = cancel.cancelled() => Ok(()),
650            // The handler stopped after taking the command, so the subscription stays desired
651            result = receiver => result.map_err(|e| AxWsClientError::Transport(e.to_string()))?,
652        }
653    }
654
655    /// Replaces the symbol's market data subscription to request a fresh book snapshot, opening
656    /// `gate` once the subscription is written.
657    ///
658    /// Echoes the symbol's current subscription options and keeps the subscription desired for
659    /// reconnect replay.
660    pub(crate) async fn resubscribe_book_gated(
661        &self,
662        symbol: Ustr,
663        cancel: CancellationToken,
664        gate: SnapshotGate,
665    ) -> AxWsResult<()> {
666        let (completion, receiver) = tokio::sync::oneshot::channel();
667
668        {
669            // Serializes the subscription lookup with subscribe and unsubscribe changes
670            let _guard = self.subscribe_lock.lock().await;
671
672            if cancel.is_cancelled() {
673                return Err(AxWsClientError::ClientError(
674                    "Book recovery canceled".to_string(),
675                ));
676            }
677
678            let Some(spec) = self
679                .symbol_data_types
680                .load()
681                .get(symbol.as_str())
682                .and_then(|types| types.book_level.map(|level| types.book_subscription(level)))
683            else {
684                return Err(AxWsClientError::ClientError(format!(
685                    "Book deltas for {symbol} are no longer subscribed"
686                )));
687            };
688
689            self.send_cmd(HandlerCommand::WriteBook {
690                request_id: self.next_request_id(),
691                unsubscribe_request_id: Some(self.next_request_id()),
692                symbol,
693                spec,
694                cancel,
695                gate,
696                completion,
697            })
698            .await?;
699        }
700
701        receiver
702            .await
703            .map_err(|e| AxWsClientError::Transport(e.to_string()))?
704    }
705
706    /// Subscribes to quote data for a symbol.
707    ///
708    /// Uses reference counting so the underlying AX subscription is only
709    /// removed when all data types have been unsubscribed.
710    ///
711    /// # Errors
712    ///
713    /// Returns an error if the subscription command cannot be sent.
714    pub async fn subscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
715        let _guard = self.subscribe_lock.lock().await;
716
717        let current = self
718            .symbol_data_types
719            .load()
720            .get(symbol)
721            .cloned()
722            .unwrap_or_default();
723        let old_spec = current.effective_subscription();
724        let mut next = current.clone();
725        next.quotes = true;
726        let new_spec = next.effective_subscription();
727
728        self.update_data_subscription(symbol, old_spec, new_spec)
729            .await?;
730
731        self.symbol_data_types.rcu(|m| {
732            m.entry(symbol.to_string()).or_default().quotes = true;
733        });
734
735        Ok(())
736    }
737
738    /// Subscribes to trade data for a symbol.
739    ///
740    /// Uses reference counting so the underlying AX subscription is only
741    /// removed when all data types have been unsubscribed.
742    ///
743    /// # Errors
744    ///
745    /// Returns an error if the subscription command cannot be sent.
746    pub async fn subscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
747        let _guard = self.subscribe_lock.lock().await;
748
749        let current = self
750            .symbol_data_types
751            .load()
752            .get(symbol)
753            .cloned()
754            .unwrap_or_default();
755        let old_spec = current.effective_subscription();
756        let mut next = current.clone();
757        next.trades = true;
758        let new_spec = next.effective_subscription();
759
760        self.update_data_subscription(symbol, old_spec, new_spec)
761            .await?;
762
763        self.symbol_data_types.rcu(|m| {
764            m.entry(symbol.to_string()).or_default().trades = true;
765        });
766
767        Ok(())
768    }
769
770    /// Unsubscribes from order book deltas for a symbol.
771    ///
772    /// The underlying AX subscription is only removed when all data types
773    /// (quotes, trades, book) have been unsubscribed.
774    ///
775    /// # Errors
776    ///
777    /// Returns an error if the unsubscribe command cannot be sent.
778    pub async fn unsubscribe_book_deltas(&self, symbol: &str) -> AxWsResult<()> {
779        let _guard = self.subscribe_lock.lock().await;
780
781        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
782            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe book deltas");
783            return Ok(());
784        };
785        let old_spec = current.effective_subscription();
786        let mut next = current.clone();
787        next.book_level = None;
788        let new_spec = next.effective_subscription();
789
790        self.update_data_subscription(symbol, old_spec, new_spec)
791            .await?;
792
793        self.symbol_data_types.rcu(|m| {
794            if let Some(entry) = m.get_mut(symbol) {
795                entry.book_level = None;
796                if entry.is_empty() {
797                    m.remove(symbol);
798                }
799            }
800        });
801
802        Ok(())
803    }
804
805    /// Unsubscribes from quote data for a symbol.
806    ///
807    /// The underlying AX subscription is only removed when all data types
808    /// (quotes, trades, book) have been unsubscribed.
809    ///
810    /// # Errors
811    ///
812    /// Returns an error if the unsubscribe command cannot be sent.
813    pub async fn unsubscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
814        let _guard = self.subscribe_lock.lock().await;
815
816        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
817            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe quotes");
818            return Ok(());
819        };
820        let old_spec = current.effective_subscription();
821        let mut next = current.clone();
822        next.quotes = false;
823        let new_spec = next.effective_subscription();
824
825        self.update_data_subscription(symbol, old_spec, new_spec)
826            .await?;
827
828        self.symbol_data_types.rcu(|m| {
829            if let Some(entry) = m.get_mut(symbol) {
830                entry.quotes = false;
831                if entry.is_empty() {
832                    m.remove(symbol);
833                }
834            }
835        });
836
837        Ok(())
838    }
839
840    /// Unsubscribes from trade data for a symbol.
841    ///
842    /// The underlying AX subscription is only removed when all data types
843    /// (quotes, trades, book) have been unsubscribed.
844    ///
845    /// # Errors
846    ///
847    /// Returns an error if the unsubscribe command cannot be sent.
848    pub async fn unsubscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
849        let _guard = self.subscribe_lock.lock().await;
850
851        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
852            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe trades");
853            return Ok(());
854        };
855        let old_spec = current.effective_subscription();
856        let mut next = current.clone();
857        next.trades = false;
858        let new_spec = next.effective_subscription();
859
860        self.update_data_subscription(symbol, old_spec, new_spec)
861            .await?;
862
863        self.symbol_data_types.rcu(|m| {
864            if let Some(entry) = m.get_mut(symbol) {
865                entry.trades = false;
866                if entry.is_empty() {
867                    m.remove(symbol);
868                }
869            }
870        });
871
872        Ok(())
873    }
874
875    /// Subscribes to mark prices for a symbol.
876    ///
877    /// Ensures at least an L1 subscription so that ticker messages
878    /// (which carry the mark price field) are received.
879    ///
880    /// # Errors
881    ///
882    /// Returns an error if the subscription command cannot be sent.
883    pub async fn subscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
884        let _guard = self.subscribe_lock.lock().await;
885
886        let current = self
887            .symbol_data_types
888            .load()
889            .get(symbol)
890            .cloned()
891            .unwrap_or_default();
892        let old_spec = current.effective_subscription();
893        let mut next = current.clone();
894        next.mark_prices = true;
895        let new_spec = next.effective_subscription();
896
897        self.update_data_subscription(symbol, old_spec, new_spec)
898            .await?;
899
900        self.symbol_data_types.rcu(|m| {
901            m.entry(symbol.to_string()).or_default().mark_prices = true;
902        });
903
904        Ok(())
905    }
906
907    /// Unsubscribes from mark prices for a symbol.
908    ///
909    /// The underlying AX subscription is only removed when all data types
910    /// have been unsubscribed.
911    ///
912    /// # Errors
913    ///
914    /// Returns an error if the unsubscribe command cannot be sent.
915    pub async fn unsubscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
916        let _guard = self.subscribe_lock.lock().await;
917
918        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
919            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe mark prices");
920            return Ok(());
921        };
922        let old_spec = current.effective_subscription();
923        let mut next = current.clone();
924        next.mark_prices = false;
925        let new_spec = next.effective_subscription();
926
927        self.update_data_subscription(symbol, old_spec, new_spec)
928            .await?;
929
930        self.symbol_data_types.rcu(|m| {
931            if let Some(entry) = m.get_mut(symbol) {
932                entry.mark_prices = false;
933                if entry.is_empty() {
934                    m.remove(symbol);
935                }
936            }
937        });
938
939        Ok(())
940    }
941
942    /// Subscribes to instrument status for a symbol.
943    ///
944    /// Ensures at least an L1 subscription so that ticker messages
945    /// (which carry the instrument state field) are received.
946    ///
947    /// # Errors
948    ///
949    /// Returns an error if the subscription command cannot be sent.
950    pub async fn subscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
951        let _guard = self.subscribe_lock.lock().await;
952
953        let current = self
954            .symbol_data_types
955            .load()
956            .get(symbol)
957            .cloned()
958            .unwrap_or_default();
959        let old_spec = current.effective_subscription();
960        let mut next = current.clone();
961        next.instrument_status = true;
962        let new_spec = next.effective_subscription();
963
964        self.update_data_subscription(symbol, old_spec, new_spec)
965            .await?;
966
967        self.symbol_data_types.rcu(|m| {
968            m.entry(symbol.to_string()).or_default().instrument_status = true;
969        });
970
971        Ok(())
972    }
973
974    /// Unsubscribes from instrument status for a symbol.
975    ///
976    /// The underlying AX subscription is only removed when all data types
977    /// have been unsubscribed.
978    ///
979    /// # Errors
980    ///
981    /// Returns an error if the unsubscribe command cannot be sent.
982    pub async fn unsubscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
983        let _guard = self.subscribe_lock.lock().await;
984
985        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
986            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe instrument status");
987            return Ok(());
988        };
989        let old_spec = current.effective_subscription();
990        let mut next = current.clone();
991        next.instrument_status = false;
992        let new_spec = next.effective_subscription();
993
994        self.update_data_subscription(symbol, old_spec, new_spec)
995            .await?;
996
997        self.symbol_data_types.rcu(|m| {
998            if let Some(entry) = m.get_mut(symbol) {
999                entry.instrument_status = false;
1000                if entry.is_empty() {
1001                    m.remove(symbol);
1002                }
1003            }
1004        });
1005
1006        self.status_invalidations.lock().insert(Ustr::from(symbol));
1007
1008        Ok(())
1009    }
1010
1011    async fn update_data_subscription(
1012        &self,
1013        symbol: &str,
1014        old_spec: Option<AxMdSubscriptionSpec>,
1015        new_spec: Option<AxMdSubscriptionSpec>,
1016    ) -> AxWsResult<()> {
1017        if old_spec == new_spec {
1018            return Ok(());
1019        }
1020
1021        match (old_spec, new_spec) {
1022            (None, Some(spec)) => {
1023                log::debug!("Subscribing {symbol} at {spec:?}");
1024                self.send_subscribe(symbol, spec).await
1025            }
1026            (Some(old), None) => {
1027                log::debug!("Unsubscribing {symbol} (no remaining data types)");
1028                self.send_unsubscribe(symbol, old).await
1029            }
1030            (Some(old), Some(new)) => {
1031                log::debug!("Resubscribing {symbol}: {old:?} -> {new:?}");
1032                self.send_unsubscribe(symbol, old).await?;
1033                if let Err(e) = self.send_subscribe(symbol, new).await {
1034                    log::warn!("Resubscribe failed for {symbol} at {new:?}: {e}");
1035                    if let Err(restore_err) = self.send_subscribe(symbol, old).await {
1036                        log::error!(
1037                            "Failed to restore {symbol} at {old:?}: {restore_err}, \
1038                             reconnection required"
1039                        );
1040                        self.subscriptions.mark_subscribe(&old.topic(symbol));
1041                    }
1042                    return Err(e);
1043                }
1044                Ok(())
1045            }
1046            (None, None) => Ok(()),
1047        }
1048    }
1049
1050    async fn send_subscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
1051        let topic = spec.topic(symbol);
1052        let request_id = self.next_request_id();
1053
1054        self.subscriptions.mark_subscribe(&topic);
1055
1056        if let Err(e) = self
1057            .send_cmd(HandlerCommand::Subscribe {
1058                request_id,
1059                symbol: Ustr::from(symbol),
1060                spec,
1061            })
1062            .await
1063        {
1064            self.subscriptions.mark_unsubscribe(&topic);
1065            return Err(e);
1066        }
1067
1068        Ok(())
1069    }
1070
1071    async fn send_unsubscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
1072        let request_id = self.next_request_id();
1073        let topic = spec.topic(symbol);
1074        let was_pending = self
1075            .subscriptions
1076            .pending_subscribe_topics()
1077            .contains(&topic);
1078
1079        self.subscriptions.mark_unsubscribe(&topic);
1080
1081        if let Err(e) = self
1082            .send_cmd(HandlerCommand::Unsubscribe {
1083                request_id,
1084                symbol: Ustr::from(symbol),
1085                topic: topic.clone(),
1086            })
1087            .await
1088        {
1089            self.restore_unsubscribe_state(&topic, was_pending);
1090            return Err(e);
1091        }
1092
1093        Ok(())
1094    }
1095
1096    /// Subscribes to candle data for a symbol.
1097    ///
1098    /// Skips sending if already subscribed or subscription is pending.
1099    ///
1100    /// # Errors
1101    ///
1102    /// Returns an error if the subscription command cannot be sent.
1103    pub async fn subscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
1104        let _guard = self.subscribe_lock.lock().await;
1105        let topic = format!("candles:{symbol}:{width:?}");
1106
1107        // Skip if already subscribed or pending
1108        if self.is_subscribed_topic(&topic) {
1109            log::debug!("Already subscribed to {topic}, skipping");
1110            return Ok(());
1111        }
1112
1113        let request_id = self.next_request_id();
1114
1115        // Mark pending BEFORE sending to prevent race conditions with concurrent subscribes
1116        self.subscriptions.mark_subscribe(&topic);
1117
1118        if let Err(e) = self
1119            .send_cmd(HandlerCommand::SubscribeCandles {
1120                request_id,
1121                symbol: Ustr::from(symbol),
1122                width,
1123            })
1124            .await
1125        {
1126            // Rollback pending state on send failure
1127            self.subscriptions.mark_unsubscribe(&topic);
1128            return Err(e);
1129        }
1130
1131        Ok(())
1132    }
1133
1134    /// Unsubscribes from candle data for a symbol.
1135    ///
1136    /// # Errors
1137    ///
1138    /// Returns an error if the unsubscribe command cannot be sent.
1139    pub async fn unsubscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
1140        let _guard = self.subscribe_lock.lock().await;
1141        let request_id = self.next_request_id();
1142        let topic = format!("candles:{symbol}:{width:?}");
1143        let was_pending = self
1144            .subscriptions
1145            .pending_subscribe_topics()
1146            .contains(&topic);
1147
1148        if !self.is_subscribed_topic(&topic) {
1149            log::debug!("Not subscribed to {topic}, skipping unsubscribe");
1150            return Ok(());
1151        }
1152
1153        self.subscriptions.mark_unsubscribe(&topic);
1154
1155        if let Err(e) = self
1156            .send_cmd(HandlerCommand::UnsubscribeCandles {
1157                request_id,
1158                symbol: Ustr::from(symbol),
1159                width,
1160                topic: topic.clone(),
1161            })
1162            .await
1163        {
1164            self.restore_unsubscribe_state(&topic, was_pending);
1165            return Err(e);
1166        }
1167
1168        Ok(())
1169    }
1170
1171    fn restore_unsubscribe_state(&self, topic: &str, was_pending: bool) {
1172        self.subscriptions.confirm_unsubscribe(topic);
1173        self.subscriptions.mark_subscribe(topic);
1174        if !was_pending {
1175            self.subscriptions.confirm_subscribe(topic);
1176        }
1177    }
1178
1179    /// Returns a stream of WebSocket messages.
1180    ///
1181    /// # Panics
1182    ///
1183    /// Panics if called before `connect()` or if the stream has already been taken.
1184    pub fn stream(&mut self) -> impl futures_util::Stream<Item = AxDataWsMessage> + 'static {
1185        let rx = self
1186            .out_rx
1187            .take()
1188            .expect("Stream receiver already taken or client not connected - stream() can only be called once");
1189        let mut rx = Arc::try_unwrap(rx).expect(
1190            "Cannot take ownership of stream - client was cloned and other references exist",
1191        );
1192        async_stream::stream! {
1193            while let Some(msg) = rx.recv().await {
1194                yield msg;
1195            }
1196        }
1197    }
1198
1199    pub(crate) fn begin_shutdown(&self) {
1200        self.cancellation_token.load().cancel();
1201        self.signal.store(true, Ordering::Release);
1202    }
1203
1204    /// Disconnects the WebSocket connection gracefully.
1205    pub async fn disconnect(&self) {
1206        log::debug!("Disconnecting WebSocket");
1207        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1208    }
1209
1210    /// Closes the WebSocket connection and cleans up resources.
1211    ///
1212    /// # Errors
1213    ///
1214    /// Returns an error if the handler task fails or does not stop after abort.
1215    pub async fn close(&mut self) -> anyhow::Result<()> {
1216        let connect_lock = Arc::clone(&self.connect_lock);
1217        let _guard = connect_lock.lock().await;
1218        log::debug!("Closing WebSocket client");
1219
1220        // Send disconnect first to allow graceful cleanup before signal
1221        self.cancellation_token.load().cancel();
1222        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1223        tokio::time::sleep(Duration::from_millis(50)).await;
1224        self.signal.store(true, Ordering::Release);
1225
1226        let outcome = self
1227            .task_handle
1228            .finish(Duration::from_secs(2), Duration::from_secs(2))
1229            .await;
1230
1231        *self.reconnect_headers.lock() = None;
1232
1233        if let Some(control) = &self.socket_control {
1234            control.deregister();
1235        }
1236
1237        match outcome {
1238            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
1239            Some(TaskJoinOutcome::Failed(error)) => Err(anyhow::anyhow!(
1240                "Architect AX data WebSocket handler failed: {error}"
1241            )),
1242            Some(TaskJoinOutcome::Incomplete) => Err(anyhow::anyhow!(
1243                "Architect AX data WebSocket handler did not stop after abort"
1244            )),
1245        }
1246    }
1247
1248    async fn send_cmd(&self, cmd: HandlerCommand) -> AxWsResult<()> {
1249        let guard = self.cmd_tx.read().await;
1250        guard
1251            .send(cmd)
1252            .map_err(|e| AxWsClientError::ChannelError(e.to_string()))
1253    }
1254}
1255
1256impl Drop for AxMdWebSocketClient {
1257    fn drop(&mut self) {
1258        if Arc::strong_count(&self.task_handle) == 1 && !self.task_handle.is_empty() {
1259            self.cancellation_token.load().cancel();
1260            self.signal.store(true, Ordering::Release);
1261            self.task_handle.abort();
1262
1263            if let Some(control) = &self.socket_control {
1264                control.deregister();
1265            }
1266        }
1267    }
1268}
1269
1270#[cfg(test)]
1271mod tests {
1272    use futures_util::StreamExt;
1273    use nautilus_live::task::TaskGroup;
1274    use nautilus_model::identifiers::InstrumentId;
1275    use nautilus_testkit::connection::race_command_with_publication;
1276    use rstest::rstest;
1277
1278    use super::*;
1279    use crate::{
1280        book::{recovery::spawn_subscription_task, sync::BookSyncTracker},
1281        websocket::data::subscription::SubscriptionOrder,
1282    };
1283
1284    #[tokio::test]
1285    async fn command_during_sender_publication_reaches_initialized_handler() {
1286        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1287        let addr = listener.local_addr().unwrap();
1288        let (received_tx, mut received_rx) = tokio::sync::mpsc::unbounded_channel();
1289
1290        let server = tokio::spawn(async move {
1291            let (socket, _) = listener.accept().await.unwrap();
1292            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1293            while let Some(Ok(message)) = socket.next().await {
1294                if let tokio_tungstenite::tungstenite::Message::Text(text) = message {
1295                    let payload: serde_json::Value = serde_json::from_str(text.as_str()).unwrap();
1296                    let _ = received_tx.send(payload);
1297                }
1298            }
1299        });
1300
1301        let url = format!("ws://{addr}");
1302        let mut ws_client =
1303            AxMdWebSocketClient::new(url, "startup-token", 30, TransportBackend::default(), None);
1304        let observer = ws_client.clone();
1305
1306        let command = HandlerCommand::Subscribe {
1307            request_id: 71,
1308            symbol: Ustr::from("BTCUSD"),
1309            spec: AxMdSubscriptionSpec::new(AxMarketDataLevel::Trades, None, None),
1310        };
1311
1312        race_command_with_publication(
1313            Arc::clone(&ws_client.cmd_tx),
1314            ws_client.connect(),
1315            move || command,
1316            || observer.is_active(),
1317        )
1318        .await
1319        .expect("connect failed");
1320        let received = tokio::time::timeout(Duration::from_secs(2), received_rx.recv()).await;
1321        ws_client.close().await.expect("disconnect failed");
1322        server.abort();
1323
1324        let payload = received
1325            .expect("startup command must reach server")
1326            .unwrap();
1327        assert_eq!(
1328            payload,
1329            serde_json::json!({
1330                "rid": 71,
1331                "type": "subscribe",
1332                "symbol": "BTCUSD",
1333                "level": "TRADES"
1334            })
1335        );
1336    }
1337
1338    #[rstest]
1339    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
1340    #[case(
1341        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
1342        "wss://stream.example.com:9443"
1343    )]
1344    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
1345        let client = AxMdWebSocketClient::new(
1346            url.to_string(),
1347            "SYN-auth-token",
1348            30,
1349            TransportBackend::default(),
1350            None,
1351        );
1352        let debug = format!("{client:?}");
1353        assert!(debug.contains(&format!("url: {origin:?}")));
1354        assert!(!debug.contains("SYN-"));
1355        assert_eq!(client.url.expose(), url);
1356    }
1357
1358    #[rstest]
1359    fn test_auth_token_uses_secret_string_owner() {
1360        let client = AxMdWebSocketClient::new(
1361            "ws://localhost:9999/md/ws".to_string(),
1362            "initial-token".to_string(),
1363            30,
1364            TransportBackend::default(),
1365            None,
1366        );
1367
1368        client.set_auth_token("replacement-token".to_string());
1369
1370        let token = client.auth_token.lock();
1371        assert_eq!(
1372            token.as_ref().map(SecretString::expose_secret),
1373            Some("replacement-token")
1374        );
1375    }
1376
1377    #[tokio::test]
1378    async fn test_drop_aborts_handler_task() {
1379        let client = AxMdWebSocketClient::new(
1380            "ws://localhost:9999/md/ws".to_string(),
1381            "test_token".to_string(),
1382            30,
1383            TransportBackend::default(),
1384            None,
1385        );
1386        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1387        let handle = tokio::spawn(async move {
1388            started_tx.send(()).expect("started receiver");
1389            std::future::pending::<()>().await;
1390        });
1391        let abort_handle = handle.abort_handle();
1392        client.task_handle.insert(handle);
1393        started_rx.await.expect("handler task started");
1394
1395        drop(client);
1396
1397        tokio::time::timeout(Duration::from_secs(1), async {
1398            while !abort_handle.is_finished() {
1399                tokio::task::yield_now().await;
1400            }
1401        })
1402        .await
1403        .expect("handler task aborted");
1404    }
1405
1406    #[rstest]
1407    fn test_effective_subscription_empty_returns_none() {
1408        let sdt = SymbolDataTypes::default();
1409        assert_eq!(sdt.effective_subscription(), None);
1410        assert!(sdt.is_empty());
1411    }
1412
1413    #[rstest]
1414    fn test_effective_subscription_book_level_takes_precedence() {
1415        let sdt = SymbolDataTypes {
1416            book_level: Some(AxMarketDataLevel::Level2),
1417            quotes: true,
1418            ..Default::default()
1419        };
1420        assert_eq!(
1421            sdt.effective_subscription(),
1422            Some(AxMdSubscriptionSpec::new(
1423                AxMarketDataLevel::Level2,
1424                Some(false),
1425                Some(false),
1426            ))
1427        );
1428        assert!(!sdt.is_empty());
1429    }
1430
1431    #[rstest]
1432    #[case(
1433        true,
1434        false,
1435        false,
1436        false,
1437        AxMarketDataLevel::Level1,
1438        Some(false),
1439        Some(false)
1440    )]
1441    #[case(false, true, false, false, AxMarketDataLevel::Trades, None, None)]
1442    #[case(
1443        false,
1444        false,
1445        true,
1446        false,
1447        AxMarketDataLevel::Level1,
1448        Some(false),
1449        Some(true)
1450    )]
1451    #[case(
1452        false,
1453        false,
1454        false,
1455        true,
1456        AxMarketDataLevel::Level1,
1457        Some(false),
1458        Some(true)
1459    )]
1460    fn test_effective_subscription_for_single_data_type(
1461        #[case] quotes: bool,
1462        #[case] trades: bool,
1463        #[case] mark_prices: bool,
1464        #[case] instrument_status: bool,
1465        #[case] level: AxMarketDataLevel,
1466        #[case] include_trades: Option<bool>,
1467        #[case] include_ticker: Option<bool>,
1468    ) {
1469        let sdt = SymbolDataTypes {
1470            quotes,
1471            trades,
1472            mark_prices,
1473            instrument_status,
1474            book_level: None,
1475        };
1476        assert_eq!(
1477            sdt.effective_subscription(),
1478            Some(AxMdSubscriptionSpec::new(
1479                level,
1480                include_trades,
1481                include_ticker,
1482            ))
1483        );
1484        assert!(!sdt.is_empty());
1485    }
1486
1487    #[rstest]
1488    #[case(false)]
1489    #[case(true)]
1490    #[tokio::test]
1491    async fn test_unsubscribe_send_failure_restores_subscription(#[case] was_pending: bool) {
1492        let client = AxMdWebSocketClient::new(
1493            "ws://localhost:9999/md/ws".to_string(),
1494            "test_token".to_string(),
1495            30,
1496            TransportBackend::default(),
1497            None,
1498        );
1499        let symbol = "EURUSD-PERP";
1500        let spec = AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(false));
1501        let topic = spec.topic(symbol);
1502        client.subscriptions.mark_subscribe(&topic);
1503        if !was_pending {
1504            client.subscriptions.confirm_subscribe(&topic);
1505        }
1506
1507        let error = client.send_unsubscribe(symbol, spec).await.unwrap_err();
1508
1509        assert_eq!(error.to_string(), "Channel error: channel closed");
1510        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1511        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1512        assert_eq!(
1513            client.subscriptions.pending_subscribe_topics().len(),
1514            usize::from(was_pending)
1515        );
1516        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1517    }
1518
1519    #[rstest]
1520    #[case(false)]
1521    #[case(true)]
1522    #[tokio::test]
1523    async fn test_unsubscribe_candles_send_failure_restores_subscription(
1524        #[case] was_pending: bool,
1525    ) {
1526        let client = AxMdWebSocketClient::new(
1527            "ws://localhost:9999/md/ws".to_string(),
1528            "test_token".to_string(),
1529            30,
1530            TransportBackend::default(),
1531            None,
1532        );
1533        let symbol = "EURUSD-PERP";
1534        let width = AxCandleWidth::Minutes1;
1535        let topic = format!("candles:{symbol}:{width:?}");
1536        client.subscriptions.mark_subscribe(&topic);
1537        if !was_pending {
1538            client.subscriptions.confirm_subscribe(&topic);
1539        }
1540
1541        let error = client.unsubscribe_candles(symbol, width).await.unwrap_err();
1542
1543        assert_eq!(error.to_string(), "Channel error: channel closed");
1544        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1545        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1546        assert_eq!(
1547            client.subscriptions.pending_subscribe_topics().len(),
1548            usize::from(was_pending)
1549        );
1550        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1551    }
1552
1553    #[rstest]
1554    #[tokio::test]
1555    async fn test_unsubscribe_candles_skips_untracked_topic() {
1556        let client = AxMdWebSocketClient::new(
1557            "ws://localhost:9999/md/ws".to_string(),
1558            "test_token".to_string(),
1559            30,
1560            TransportBackend::default(),
1561            None,
1562        );
1563
1564        client
1565            .unsubscribe_candles("EURUSD-PERP", AxCandleWidth::Minutes1)
1566            .await
1567            .unwrap();
1568
1569        assert_eq!(client.subscription_count(), 0);
1570        assert!(client.subscriptions.all_topics().is_empty());
1571        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1572        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1573    }
1574
1575    const SYMBOL: &str = "EURUSD-PERP";
1576
1577    // Returns a client whose handler commands the test receives, with `current` already streaming
1578    async fn book_client(
1579        current: SymbolDataTypes,
1580    ) -> (
1581        AxMdWebSocketClient,
1582        tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
1583    ) {
1584        let client = AxMdWebSocketClient::new(
1585            "ws://localhost:9999/md/ws".to_string(),
1586            "test_token".to_string(),
1587            30,
1588            TransportBackend::default(),
1589            None,
1590        );
1591        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1592        *client.cmd_tx.write().await = cmd_tx;
1593
1594        if let Some(spec) = current.effective_subscription() {
1595            let topic = spec.topic(SYMBOL);
1596            client.subscriptions.mark_subscribe(&topic);
1597            client.subscriptions.confirm_subscribe(&topic);
1598            client.symbol_data_types.insert(SYMBOL.to_string(), current);
1599        }
1600
1601        (client, cmd_rx)
1602    }
1603
1604    #[rstest]
1605    #[case::no_stream(SymbolDataTypes::default(), None, false)]
1606    #[case::quote_stream(
1607        SymbolDataTypes { quotes: true, ..Default::default() },
1608        Some("EURUSD-PERP:Level1:false:false"),
1609        false
1610    )]
1611    #[case::book_stream(
1612        SymbolDataTypes { book_level: Some(AxMarketDataLevel::Level2), ..Default::default() },
1613        None,
1614        true
1615    )]
1616    #[tokio::test]
1617    async fn test_subscribe_book_deltas_gated_writes_book_stream(
1618        #[case] current: SymbolDataTypes,
1619        #[case] unsubscribed: Option<&str>,
1620        #[case] replaced: bool,
1621    ) {
1622        let quotes = current.quotes;
1623        let (client, mut cmd_rx) = book_client(current).await;
1624        let gate = SnapshotGate::default();
1625        gate.lock().close();
1626
1627        let subscribe = client.subscribe_book_deltas_gated(
1628            Ustr::from(SYMBOL),
1629            AxMarketDataLevel::Level2,
1630            CancellationToken::new(),
1631            gate,
1632        );
1633
1634        let handler = async {
1635            let mut topics = Vec::new();
1636
1637            loop {
1638                match cmd_rx.recv().await.expect("command queued") {
1639                    HandlerCommand::Unsubscribe { topic, .. } => topics.push(topic),
1640                    HandlerCommand::WriteBook {
1641                        spec,
1642                        unsubscribe_request_id,
1643                        completion,
1644                        ..
1645                    } => {
1646                        completion.send(Ok(())).unwrap();
1647                        return (topics, spec, unsubscribe_request_id.is_some());
1648                    }
1649                    other => panic!("unexpected command {other:?}"),
1650                }
1651            }
1652        };
1653
1654        let (result, (topics, spec, replaces)) = tokio::join!(subscribe, handler);
1655        let types = client
1656            .symbol_data_types
1657            .load()
1658            .get(SYMBOL)
1659            .cloned()
1660            .unwrap();
1661
1662        assert!(result.is_ok());
1663        assert_eq!(
1664            topics,
1665            unsubscribed
1666                .map(ToString::to_string)
1667                .into_iter()
1668                .collect::<Vec<_>>()
1669        );
1670        assert_eq!(
1671            spec,
1672            AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(false))
1673        );
1674        assert_eq!(replaces, replaced);
1675        assert_eq!(types.book_level, Some(AxMarketDataLevel::Level2));
1676        assert_eq!(types.quotes, quotes);
1677        assert_eq!(client.subscriptions.all_topics(), vec![spec.topic(SYMBOL)]);
1678    }
1679
1680    #[rstest]
1681    #[tokio::test]
1682    async fn test_book_subscription_write_failure_starts_recovery() {
1683        let (client, mut cmd_rx) = book_client(SymbolDataTypes::default()).await;
1684        let book_sync = BookSyncTracker::default();
1685        let tasks = TaskGroup::new();
1686        let instrument_id = InstrumentId::from("EURUSD-PERP.AX");
1687
1688        spawn_subscription_task(
1689            instrument_id,
1690            AxMarketDataLevel::Level2,
1691            SubscriptionOrder::default().next(),
1692            book_sync.clone(),
1693            client.clone(),
1694            Duration::from_secs(10),
1695            &tasks.spawner().unwrap(),
1696        );
1697
1698        let mut replaces = Vec::new();
1699
1700        for _ in 0..2 {
1701            let command = tokio::time::timeout(Duration::from_secs(5), cmd_rx.recv())
1702                .await
1703                .expect("book write queued")
1704                .expect("command channel open");
1705
1706            let HandlerCommand::WriteBook {
1707                unsubscribe_request_id,
1708                completion,
1709                ..
1710            } = command
1711            else {
1712                panic!("expected a book write, was {command:?}");
1713            };
1714
1715            replaces.push(unsubscribe_request_id.is_some());
1716
1717            // The initial write fails in transport, which recovery retries
1718            let _ = completion.send(Err(AxWsClientError::Transport("write failed".to_string())));
1719        }
1720
1721        let owned = book_sync.claim_recovery(instrument_id).is_none();
1722        tasks.begin_shutdown();
1723
1724        assert_eq!(replaces, [false, true]);
1725        assert!(owned, "the running recovery owns the book");
1726    }
1727
1728    #[rstest]
1729    #[tokio::test]
1730    async fn test_resubscribe_book_gated_echoes_current_subscription() {
1731        let current = SymbolDataTypes {
1732            mark_prices: true,
1733            book_level: Some(AxMarketDataLevel::Level2),
1734            ..Default::default()
1735        };
1736
1737        let (client, mut cmd_rx) = book_client(current).await;
1738
1739        let resubscribe = client.resubscribe_book_gated(
1740            Ustr::from(SYMBOL),
1741            CancellationToken::new(),
1742            SnapshotGate::default(),
1743        );
1744
1745        let handler = async {
1746            match cmd_rx.recv().await.expect("command queued") {
1747                HandlerCommand::WriteBook {
1748                    request_id,
1749                    unsubscribe_request_id,
1750                    symbol,
1751                    spec,
1752                    completion,
1753                    ..
1754                } => {
1755                    completion.send(Ok(())).unwrap();
1756                    (request_id, unsubscribe_request_id, symbol, spec)
1757                }
1758                other => panic!("unexpected command {other:?}"),
1759            }
1760        };
1761
1762        let (result, (request_id, unsubscribe_request_id, symbol, spec)) =
1763            tokio::join!(resubscribe, handler);
1764
1765        assert!(result.is_ok());
1766        assert_eq!(request_id, 1);
1767        assert_eq!(unsubscribe_request_id, Some(2));
1768        assert_eq!(symbol, Ustr::from(SYMBOL));
1769        assert_eq!(
1770            spec,
1771            AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(true))
1772        );
1773        assert!(cmd_rx.try_recv().is_err());
1774    }
1775
1776    #[rstest]
1777    #[case::not_subscribed(
1778        SymbolDataTypes { quotes: true, ..Default::default() },
1779        false,
1780        "Client error: Book deltas for EURUSD-PERP are no longer subscribed"
1781    )]
1782    #[case::cancelled(
1783        SymbolDataTypes { book_level: Some(AxMarketDataLevel::Level2), ..Default::default() },
1784        true,
1785        "Client error: Book recovery canceled"
1786    )]
1787    #[tokio::test]
1788    async fn test_resubscribe_book_gated_refuses_without_write(
1789        #[case] current: SymbolDataTypes,
1790        #[case] cancelled: bool,
1791        #[case] expected: &str,
1792    ) {
1793        let (client, mut cmd_rx) = book_client(current).await;
1794        let cancel = CancellationToken::new();
1795
1796        if cancelled {
1797            cancel.cancel();
1798        }
1799
1800        let error = client
1801            .resubscribe_book_gated(Ustr::from(SYMBOL), cancel, SnapshotGate::default())
1802            .await
1803            .unwrap_err();
1804
1805        assert_eq!(error.to_string(), expected);
1806        assert!(cmd_rx.try_recv().is_err());
1807    }
1808}
nautilus_architect_ax/websocket/data/client.rs

nautilus_architect_ax/websocket/data/
client.rs
