nautilus_architect_ax/websocket/orders/
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//! Orders WebSocket client for Ax.
17
18use std::{
19    fmt::Debug,
20    num::NonZeroU32,
21    sync::{
22        Arc,
23        atomic::{AtomicBool, AtomicI64, AtomicU8, Ordering},
24    },
25    time::Duration,
26};
27
28use arc_swap::ArcSwap;
29use dashmap::{DashMap, mapref::entry::Entry};
30use nautilus_common::cache::InstrumentLookupError;
31use nautilus_core::{
32    AtomicMap,
33    nanos::UnixNanos,
34    string::secret::SecretString,
35    time::{AtomicTime, get_atomic_clock_realtime},
36};
37use nautilus_live::{
38    SocketControl,
39    task::{SharedTaskSlot, TaskJoinOutcome},
40};
41use nautilus_model::{
42    enums::{OrderSide, TimeInForce},
43    identifiers::{AccountId, ClientOrderId, InstrumentId, StrategyId, TraderId, VenueOrderId},
44    instruments::{Instrument, InstrumentAny},
45    types::{Price, Quantity},
46};
47use nautilus_network::{
48    ConnectionUrl,
49    http::create_standard_nautilus_headers,
50    mode::ConnectionMode,
51    websocket::{
52        AuthTracker, InitialConnectRetryPolicy, PingHandler, ReconnectHeaders, TransportBackend,
53        WebSocketClient, WebSocketConfig, channel_message_handler,
54    },
55};
56use parking_lot::Mutex;
57use tokio_util::sync::CancellationToken;
58use ustr::Ustr;
59
60use super::handler::{AxOrdersWsFeedHandler, HandlerCommand, WsOrderInfo};
61use crate::{
62    common::{
63        consts::AX_NAUTILUS_TAG,
64        enums::{AxOrderRequestType, AxOrderSide, AxTimeInForce},
65        parse::{client_order_id_to_cid, quantity_to_contracts},
66    },
67    websocket::messages::{AxOrdersWsMessage, AxWsPlaceOrder, OrderMetadata},
68};
69
70/// Result type for Ax orders WebSocket operations.
71pub type AxOrdersWsResult<T> = Result<T, AxOrdersWsClientError>;
72
73/// Shared caches for order state tracking between the client and consumers.
74#[derive(Debug, Clone)]
75pub struct OrdersCaches {
76    /// Maps client order IDs to order metadata.
77    pub orders_metadata: Arc<DashMap<ClientOrderId, OrderMetadata>>,
78    /// Maps venue order IDs to client order IDs.
79    pub venue_to_client_id: Arc<DashMap<VenueOrderId, ClientOrderId>>,
80    /// Maps AX cid values to client order IDs.
81    pub cid_to_client_order_id: Arc<DashMap<u64, ClientOrderId>>,
82}
83
84impl Default for OrdersCaches {
85    fn default() -> Self {
86        Self {
87            orders_metadata: Arc::new(DashMap::new()),
88            venue_to_client_id: Arc::new(DashMap::new()),
89            cid_to_client_order_id: Arc::new(DashMap::new()),
90        }
91    }
92}
93
94/// Error type for the Ax orders WebSocket client.
95#[derive(Debug, Clone)]
96pub enum AxOrdersWsClientError {
97    /// Transport/connection error.
98    Transport(String),
99    /// Channel send error.
100    ChannelError(String),
101    /// Authentication error.
102    AuthenticationError(String),
103    /// Client-side validation error.
104    ClientError(String),
105}
106
107impl core::fmt::Display for AxOrdersWsClientError {
108    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
109        match self {
110            Self::Transport(msg) => write!(f, "Transport error: {msg}"),
111            Self::ChannelError(msg) => write!(f, "Channel error: {msg}"),
112            Self::AuthenticationError(msg) => write!(f, "Authentication error: {msg}"),
113            Self::ClientError(msg) => write!(f, "Client error: {msg}"),
114        }
115    }
116}
117
118impl std::error::Error for AxOrdersWsClientError {}
119
120impl From<&'static str> for AxOrdersWsClientError {
121    fn from(msg: &'static str) -> Self {
122        Self::ClientError(msg.to_string())
123    }
124}
125
126/// Orders WebSocket client for Ax.
127///
128/// Provides authenticated order management including placing, canceling,
129/// and monitoring order status via WebSocket.
130pub struct AxOrdersWebSocketClient {
131    clock: &'static AtomicTime,
132    url: ConnectionUrl,
133    heartbeat: Option<u64>,
134    reconnect_headers: Arc<Mutex<Option<ReconnectHeaders>>>,
135    connection_mode: Arc<ArcSwap<AtomicU8>>,
136    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
137    out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<AxOrdersWsMessage>>>,
138    signal: Arc<AtomicBool>,
139    cancellation_token: Arc<ArcSwap<CancellationToken>>,
140    task_handle: Arc<SharedTaskSlot<()>>,
141    connect_lock: Arc<tokio::sync::Mutex<()>>,
142    auth_tracker: AuthTracker,
143    instruments_cache: Arc<AtomicMap<Ustr, InstrumentAny>>,
144    caches: OrdersCaches,
145    request_id_counter: Arc<AtomicI64>,
146    account_id: AccountId,
147    trader_id: TraderId,
148    transport_backend: TransportBackend,
149    proxy_url: Option<SecretString>,
150    socket_control: Option<SocketControl>,
151}
152
153impl Debug for AxOrdersWebSocketClient {
154    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
155        f.debug_struct(stringify!(AxOrdersWebSocketClient))
156            .field("url", &self.url)
157            .field("heartbeat", &self.heartbeat)
158            .field("account_id", &self.account_id)
159            .finish()
160    }
161}
162
163impl Clone for AxOrdersWebSocketClient {
164    fn clone(&self) -> Self {
165        Self {
166            clock: self.clock,
167            url: self.url.clone(),
168            heartbeat: self.heartbeat,
169            reconnect_headers: Arc::clone(&self.reconnect_headers),
170            connection_mode: Arc::clone(&self.connection_mode),
171            cmd_tx: Arc::clone(&self.cmd_tx),
172            out_rx: None, // Each clone gets its own receiver
173            signal: Arc::clone(&self.signal),
174            cancellation_token: Arc::clone(&self.cancellation_token),
175            task_handle: Arc::clone(&self.task_handle),
176            connect_lock: Arc::clone(&self.connect_lock),
177            auth_tracker: self.auth_tracker.clone(),
178            instruments_cache: Arc::clone(&self.instruments_cache),
179            caches: self.caches.clone(),
180            request_id_counter: Arc::clone(&self.request_id_counter),
181            account_id: self.account_id,
182            trader_id: self.trader_id,
183            transport_backend: self.transport_backend,
184            proxy_url: self.proxy_url.clone(),
185            socket_control: self.socket_control.clone(),
186        }
187    }
188}
189
190impl AxOrdersWebSocketClient {
191    fn initial_connect_retry_policy() -> InitialConnectRetryPolicy {
192        InitialConnectRetryPolicy {
193            max_attempts: NonZeroU32::new(5).expect("initial connect attempts must be non-zero"),
194            delay_initial: Duration::from_millis(500),
195            delay_max: Duration::from_secs(5),
196            backoff_factor: 2.0,
197            jitter_ms: 250,
198        }
199    }
200
201    /// Creates a new Ax orders WebSocket client.
202    #[must_use]
203    pub fn new(
204        url: String,
205        account_id: AccountId,
206        trader_id: TraderId,
207        heartbeat: u64,
208        transport_backend: TransportBackend,
209        proxy_url: Option<String>,
210    ) -> Self {
211        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
212
213        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
214        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
215
216        Self {
217            clock: get_atomic_clock_realtime(),
218            url: url.into(),
219            heartbeat: Some(heartbeat),
220            reconnect_headers: Arc::new(Mutex::new(None)),
221            connection_mode,
222            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
223            out_rx: None,
224            signal: Arc::new(AtomicBool::new(false)),
225            cancellation_token: Arc::new(ArcSwap::from_pointee(CancellationToken::new())),
226            task_handle: Arc::new(SharedTaskSlot::new()),
227            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
228            auth_tracker: AuthTracker::default(),
229            instruments_cache: Arc::new(AtomicMap::new()),
230            caches: OrdersCaches::default(),
231            request_id_counter: Arc::new(AtomicI64::new(1)),
232            account_id,
233            trader_id,
234            transport_backend,
235            proxy_url: proxy_url.map(SecretString::from),
236            socket_control: None,
237        }
238    }
239
240    /// Configures socket state reporting and reconnect control.
241    #[must_use]
242    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
243        self.socket_control = Some(control);
244        self
245    }
246
247    fn generate_ts_init(&self) -> UnixNanos {
248        self.clock.get_time_ns()
249    }
250
251    /// Returns the WebSocket URL.
252    #[must_use]
253    pub fn url(&self) -> &str {
254        self.url.expose()
255    }
256
257    /// Returns the account ID.
258    #[must_use]
259    pub fn account_id(&self) -> AccountId {
260        self.account_id
261    }
262
263    /// Returns whether the client is currently connected and active.
264    #[must_use]
265    pub fn is_active(&self) -> bool {
266        let connection_mode_arc = self.connection_mode.load();
267        ConnectionMode::from_atomic(&connection_mode_arc).is_active()
268            && !self.signal.load(Ordering::Acquire)
269    }
270
271    /// Returns whether the client is closed.
272    #[must_use]
273    pub fn is_closed(&self) -> bool {
274        let connection_mode_arc = self.connection_mode.load();
275        ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
276            || self.signal.load(Ordering::Acquire)
277    }
278
279    /// Generates a unique request ID.
280    fn next_request_id(&self) -> i64 {
281        self.request_id_counter.fetch_add(1, Ordering::Relaxed)
282    }
283
284    /// Caches an instrument for use during message parsing.
285    pub fn cache_instrument(&self, instrument: InstrumentAny) {
286        let symbol = instrument.symbol().inner();
287        self.instruments_cache.insert(symbol, instrument);
288    }
289
290    /// Caches multiple instruments for use during message parsing.
291    pub fn cache_instruments(&self, instruments: &[InstrumentAny]) {
292        self.instruments_cache.rcu(|m| {
293            for inst in instruments {
294                m.insert(inst.symbol().inner(), inst.clone());
295            }
296        });
297    }
298
299    /// Updates the token used by future automatic reconnect attempts.
300    ///
301    /// Updating the token does not interrupt the active WebSocket connection.
302    ///
303    /// # Errors
304    ///
305    /// Returns an error if the reconnect header cannot be updated.
306    pub fn update_auth_token(&self, token: &str) -> AxOrdersWsResult<()> {
307        let value = format!("Bearer {token}");
308
309        if let Some(headers) = self.reconnect_headers.lock().as_ref() {
310            headers
311                .update("Authorization", &value)
312                .map_err(|e| AxOrdersWsClientError::Transport(e.to_string()))?;
313        }
314        Ok(())
315    }
316
317    /// Returns a cached instrument by symbol.
318    #[must_use]
319    pub fn get_cached_instrument(&self, symbol: &Ustr) -> Option<InstrumentAny> {
320        self.instruments_cache.get_cloned(symbol)
321    }
322
323    /// Returns the shared order caches.
324    #[must_use]
325    pub fn caches(&self) -> &OrdersCaches {
326        &self.caches
327    }
328
329    /// Returns the instruments cache.
330    #[must_use]
331    pub fn instruments_cache(&self) -> Arc<AtomicMap<Ustr, InstrumentAny>> {
332        Arc::clone(&self.instruments_cache)
333    }
334
335    /// Returns the orders metadata cache.
336    #[must_use]
337    pub fn orders_metadata(&self) -> &Arc<DashMap<ClientOrderId, OrderMetadata>> {
338        &self.caches.orders_metadata
339    }
340
341    /// Returns the cid to client order ID mapping for order correlation.
342    #[must_use]
343    pub fn cid_to_client_order_id(&self) -> &Arc<DashMap<u64, ClientOrderId>> {
344        &self.caches.cid_to_client_order_id
345    }
346
347    /// Resolves a cid to a ClientOrderId if the mapping exists.
348    #[must_use]
349    pub fn resolve_cid(&self, cid: u64) -> Option<ClientOrderId> {
350        self.caches.cid_to_client_order_id.get(&cid).map(|v| *v)
351    }
352
353    /// Registers an external order with the WebSocket handler for event tracking.
354    ///
355    /// This allows the handler to create proper events (e.g., OrderCanceled, OrderFilled)
356    /// for orders that were reconciled externally and not submitted through this client.
357    ///
358    /// Returns `false` if the instrument is not cached (registration skipped).
359    pub fn register_external_order(
360        &self,
361        client_order_id: ClientOrderId,
362        venue_order_id: VenueOrderId,
363        instrument_id: InstrumentId,
364        strategy_id: StrategyId,
365    ) -> bool {
366        if self.caches.orders_metadata.contains_key(&client_order_id) {
367            return true;
368        }
369
370        // Required for correct precision on fills
371        let symbol = instrument_id.symbol.inner();
372        let Some(instrument) = self.get_cached_instrument(&symbol) else {
373            log::warn!(
374                "Cannot register external order {client_order_id}: \
375                 instrument {instrument_id} not in cache"
376            );
377            return false;
378        };
379
380        let metadata = OrderMetadata {
381            trader_id: self.trader_id,
382            strategy_id,
383            instrument_id,
384            client_order_id,
385            venue_order_id: Some(venue_order_id),
386            ts_init: self.generate_ts_init(),
387            size_precision: instrument.size_precision(),
388            price_precision: instrument.price_precision(),
389            quote_currency: instrument.quote_currency(),
390        };
391
392        self.caches
393            .orders_metadata
394            .insert(client_order_id, metadata);
395        self.caches
396            .venue_to_client_id
397            .insert(venue_order_id, client_order_id);
398
399        log::debug!(
400            "Registered external order {client_order_id} ({venue_order_id}) for {instrument_id} [{strategy_id}]"
401        );
402
403        true
404    }
405
406    /// Establishes the WebSocket connection with authentication.
407    ///
408    /// # Arguments
409    ///
410    /// * `bearer_token` - The bearer token for authentication.
411    ///
412    /// # Errors
413    ///
414    /// Returns an error if the connection cannot be established.
415    pub async fn connect(&mut self, bearer_token: &str) -> AxOrdersWsResult<()> {
416        let connect_lock = Arc::clone(&self.connect_lock);
417        let _guard = connect_lock.lock().await;
418
419        if !self.task_handle.is_empty() && !self.task_handle.is_finished() {
420            return Err(AxOrdersWsClientError::ClientError(
421                "WebSocket handler is already running".to_string(),
422            ));
423        }
424
425        if let Some(outcome) = self
426            .task_handle
427            .finish(Duration::from_secs(2), Duration::from_secs(2))
428            .await
429        {
430            match outcome {
431                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
432                TaskJoinOutcome::Failed(error) => {
433                    return Err(AxOrdersWsClientError::ClientError(format!(
434                        "Previous WebSocket handler failed: {error}"
435                    )));
436                }
437                TaskJoinOutcome::Incomplete => {
438                    return Err(AxOrdersWsClientError::ClientError(
439                        "Previous WebSocket handler did not stop within shutdown bounds"
440                            .to_string(),
441                    ));
442                }
443            }
444        }
445
446        self.signal.store(false, Ordering::Release);
447        let cancellation_token = CancellationToken::new();
448        self.cancellation_token
449            .store(Arc::new(cancellation_token.clone()));
450
451        let (raw_handler, raw_rx) = channel_message_handler();
452
453        // No-op ping handler: handler owns the WebSocketClient and responds to pings directly
454        let ping_handler: PingHandler = Arc::new(move |_payload: Vec<u8>| {
455            // Handler responds to pings internally via select! loop
456        });
457
458        let mut headers = create_standard_nautilus_headers();
459        headers.push((
460            "Authorization".to_string(),
461            format!("Bearer {bearer_token}"),
462        ));
463
464        let config = WebSocketConfig {
465            url: self.url.expose().to_owned(),
466            headers,
467            heartbeat_interval_secs: self.heartbeat,
468            heartbeat_payload: None, // Ax server sends heartbeats
469            connect_timeout_ms: Some(5_000),
470            reconnect_delay_initial_ms: Some(500),
471            reconnect_delay_max_ms: Some(5_000),
472            reconnect_backoff_factor: Some(1.5),
473            reconnect_jitter_ms: Some(250),
474            reconnect_max_attempts: None,
475            heartbeat_timeout_secs: None,
476            idle_timeout_ms: None,
477            writer_capacity: None,
478            backend: self.transport_backend,
479            proxy_url: self
480                .proxy_url
481                .as_ref()
482                .map(|url| url.expose_secret().to_owned()),
483            max_message_size_bytes: None,
484            max_frame_size_bytes: None,
485        };
486
487        let client = WebSocketClient::builder()
488            .config(config.clone())
489            .message_handler(raw_handler.clone())
490            .ping_handler(ping_handler.clone())
491            .initial_connect_retry_policy(Self::initial_connect_retry_policy())
492            .cancellation_token(cancellation_token)
493            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
494            .connect()
495            .await
496            .map_err(|e| {
497                AxOrdersWsClientError::Transport(format!(
498                    "WebSocket connection failed: {}: {}",
499                    self.url.origin(),
500                    e.redacted_message()
501                ))
502            })?;
503
504        self.connection_mode.store(client.connection_mode_atomic());
505        let reconnect_handle = client.reconnect_handle();
506        *self.reconnect_headers.lock() = Some(client.reconnect_headers());
507
508        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<AxOrdersWsMessage>();
509        self.out_rx = Some(Arc::new(out_rx));
510
511        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
512        *self.cmd_tx.write().await = cmd_tx.clone();
513
514        self.send_cmd(HandlerCommand::SetClient(client)).await?;
515
516        self.send_cmd(HandlerCommand::SessionAuthenticated).await?;
517
518        let signal = Arc::clone(&self.signal);
519        let auth_tracker = self.auth_tracker.clone();
520        let orders_metadata = Arc::clone(&self.caches.orders_metadata);
521        let venue_to_client_order_id = Arc::clone(&self.caches.venue_to_client_id);
522        let cid_to_client_order_id = Arc::clone(&self.caches.cid_to_client_order_id);
523
524        if let Err(e) = self.task_handle.spawn(async move {
525            let mut handler = AxOrdersWsFeedHandler::new(
526                signal.clone(),
527                cmd_rx,
528                raw_rx,
529                auth_tracker.clone(),
530                orders_metadata,
531                venue_to_client_order_id,
532                cid_to_client_order_id,
533            );
534
535            while let Some(msg) = handler.next().await {
536                if matches!(msg, AxOrdersWsMessage::Reconnected) {
537                    log::info!("WebSocket reconnected, authentication will be restored");
538                }
539
540                if out_tx.send(msg).is_err() {
541                    log::debug!("Output channel closed");
542                    break;
543                }
544            }
545
546            log::debug!("Handler loop exited");
547        }) {
548            self.out_rx = None;
549            return Err(AxOrdersWsClientError::Transport(format!(
550                "Failed to start WebSocket handler task: {e}"
551            )));
552        }
553
554        if let Some(control) = &self.socket_control {
555            control.register(move || reconnect_handle.request_reconnect());
556        }
557
558        Ok(())
559    }
560
561    /// Submits the AX priced order shape using Nautilus domain types.
562    ///
563    /// This method handles conversion from Nautilus domain types to AX-specific
564    /// types and stores order metadata for event correlation.
565    ///
566    /// # Errors
567    ///
568    /// Returns an error if:
569    /// - The time-in-force is not supported.
570    /// - The instrument is not found in the cache.
571    /// - The order command cannot be sent.
572    #[expect(clippy::too_many_arguments)]
573    pub async fn submit_order(
574        &self,
575        trader_id: TraderId,
576        strategy_id: StrategyId,
577        instrument_id: InstrumentId,
578        client_order_id: ClientOrderId,
579        order_side: OrderSide,
580        quantity: Quantity,
581        time_in_force: TimeInForce,
582        price: Price,
583        post_only: bool,
584    ) -> AxOrdersWsResult<i64> {
585        // Get instrument from cache for precision
586        let symbol = instrument_id.symbol.inner();
587        let instrument = self.get_cached_instrument(&symbol).ok_or_else(|| {
588            AxOrdersWsClientError::ClientError(
589                InstrumentLookupError::not_found(instrument_id).to_string(),
590            )
591        })?;
592
593        let ax_side = AxOrderSide::from(order_side);
594
595        let qty_contracts = quantity_to_contracts(quantity)
596            .map_err(|e| AxOrdersWsClientError::ClientError(e.to_string()))?;
597
598        let request_id = self.next_request_id();
599        let ax_tif = AxTimeInForce::try_from(time_in_force)?;
600        let cid = client_order_id_to_cid(&client_order_id);
601
602        reserve_cid_mapping(&self.caches, cid, client_order_id)?;
603
604        // Store order metadata for event correlation (after validation to avoid stale entries)
605        let metadata = OrderMetadata {
606            trader_id,
607            strategy_id,
608            instrument_id,
609            client_order_id,
610            venue_order_id: None,
611            ts_init: self.generate_ts_init(),
612            size_precision: instrument.size_precision(),
613            price_precision: instrument.price_precision(),
614            quote_currency: instrument.quote_currency(),
615        };
616        self.caches
617            .orders_metadata
618            .insert(client_order_id, metadata);
619
620        let order = AxWsPlaceOrder {
621            rid: request_id,
622            t: AxOrderRequestType::PlaceOrder,
623            s: symbol,
624            d: ax_side,
625            q: qty_contracts,
626            p: price.as_decimal(),
627            tif: ax_tif,
628            po: post_only,
629            tag: Some(AX_NAUTILUS_TAG.to_string()),
630            cid: Some(cid),
631        };
632
633        let order_info = WsOrderInfo {
634            client_order_id,
635            symbol,
636            cid,
637        };
638
639        let result = self
640            .send_cmd(HandlerCommand::PlaceOrder {
641                request_id,
642                order,
643                order_info,
644            })
645            .await;
646
647        if result.is_err() {
648            self.caches.orders_metadata.remove(&client_order_id);
649            self.caches.cid_to_client_order_id.remove(&cid);
650        }
651
652        result?;
653        Ok(request_id)
654    }
655
656    /// Cancels an order via WebSocket.
657    ///
658    /// Requires a known `venue_order_id`.
659    ///
660    /// # Errors
661    ///
662    /// Returns an error if the cancel command cannot be sent.
663    pub async fn cancel_order(
664        &self,
665        client_order_id: ClientOrderId,
666        venue_order_id: Option<VenueOrderId>,
667    ) -> AxOrdersWsResult<i64> {
668        let order_id = venue_order_id.map(|v| v.to_string()).ok_or_else(|| {
669            AxOrdersWsClientError::ClientError(format!(
670                "Cannot cancel order {client_order_id}: missing venue_order_id"
671            ))
672        })?;
673
674        let request_id = self.next_request_id();
675
676        self.send_cmd(HandlerCommand::CancelOrder {
677            request_id,
678            order_id,
679        })
680        .await?;
681
682        Ok(request_id)
683    }
684
685    /// Requests open orders via WebSocket.
686    ///
687    /// # Errors
688    ///
689    /// Returns an error if the request command cannot be sent.
690    pub async fn get_open_orders(&self) -> AxOrdersWsResult<i64> {
691        let request_id = self.next_request_id();
692
693        self.send_cmd(HandlerCommand::GetOpenOrders { request_id })
694            .await?;
695
696        Ok(request_id)
697    }
698
699    /// Returns a stream of WebSocket messages.
700    ///
701    /// # Panics
702    ///
703    /// Panics if called before `connect()` or if the stream has already been taken.
704    pub fn stream(&mut self) -> impl futures_util::Stream<Item = AxOrdersWsMessage> + 'static {
705        let rx = self
706            .out_rx
707            .take()
708            .expect("Stream receiver already taken or client not connected - stream() can only be called once");
709        let mut rx = Arc::try_unwrap(rx).expect(
710            "Cannot take ownership of stream - client was cloned and other references exist",
711        );
712        async_stream::stream! {
713            while let Some(msg) = rx.recv().await {
714                yield msg;
715            }
716        }
717    }
718
719    pub(crate) fn begin_shutdown(&self) {
720        self.cancellation_token.load().cancel();
721        self.signal.store(true, Ordering::Release);
722    }
723
724    /// Disconnects the WebSocket connection gracefully.
725    pub async fn disconnect(&self) {
726        log::debug!("Disconnecting WebSocket");
727        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
728    }
729
730    /// Closes the WebSocket connection and cleans up resources.
731    ///
732    /// # Errors
733    ///
734    /// Returns an error if the handler task fails or does not stop after abort.
735    pub async fn close(&mut self) -> anyhow::Result<()> {
736        let connect_lock = Arc::clone(&self.connect_lock);
737        let _guard = connect_lock.lock().await;
738        log::debug!("Closing WebSocket client");
739
740        // Send disconnect first to allow graceful cleanup before signal
741        self.cancellation_token.load().cancel();
742        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
743        tokio::time::sleep(Duration::from_millis(50)).await;
744        self.signal.store(true, Ordering::Release);
745
746        let outcome = self
747            .task_handle
748            .finish(Duration::from_secs(2), Duration::from_secs(2))
749            .await;
750        *self.reconnect_headers.lock() = None;
751
752        if let Some(control) = &self.socket_control {
753            control.deregister();
754        }
755
756        match outcome {
757            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
758            Some(TaskJoinOutcome::Failed(error)) => Err(anyhow::anyhow!(
759                "Architect AX orders WebSocket handler failed: {error}"
760            )),
761            Some(TaskJoinOutcome::Incomplete) => Err(anyhow::anyhow!(
762                "Architect AX orders WebSocket handler did not stop after abort"
763            )),
764        }
765    }
766
767    async fn send_cmd(&self, cmd: HandlerCommand) -> AxOrdersWsResult<()> {
768        let guard = self.cmd_tx.read().await;
769        guard
770            .send(cmd)
771            .map_err(|e| AxOrdersWsClientError::ChannelError(e.to_string()))
772    }
773}
774
775impl Drop for AxOrdersWebSocketClient {
776    fn drop(&mut self) {
777        if Arc::strong_count(&self.task_handle) == 1 && !self.task_handle.is_empty() {
778            self.cancellation_token.load().cancel();
779            self.signal.store(true, Ordering::Release);
780            self.task_handle.abort();
781
782            if let Some(control) = &self.socket_control {
783                control.deregister();
784            }
785        }
786    }
787}
788
789fn reserve_cid_mapping(
790    caches: &OrdersCaches,
791    cid: u64,
792    client_order_id: ClientOrderId,
793) -> AxOrdersWsResult<()> {
794    match caches.cid_to_client_order_id.entry(cid) {
795        Entry::Vacant(entry) => {
796            entry.insert(client_order_id);
797            Ok(())
798        }
799        Entry::Occupied(entry) => Err(AxOrdersWsClientError::ClientError(format!(
800            "AX cid {cid} is already mapped to {}",
801            entry.get(),
802        ))),
803    }
804}
805
806#[cfg(test)]
807mod tests {
808    use std::sync::Arc;
809
810    use rstest::rstest;
811
812    use super::*;
813
814    #[rstest]
815    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
816    #[case(
817        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
818        "wss://stream.example.com:9443"
819    )]
820    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
821        let client = AxOrdersWebSocketClient::new(
822            url.to_string(),
823            AccountId::from("AX-001"),
824            TraderId::from("TRADER-001"),
825            30,
826            TransportBackend::default(),
827            None,
828        );
829        let debug = format!("{client:?}");
830        assert!(debug.contains(&format!("url: {origin:?}")));
831        assert!(!debug.contains("SYN-"));
832        assert_eq!(client.url.expose(), url);
833    }
834
835    #[tokio::test]
836    async fn test_drop_aborts_handler_task() {
837        let client = AxOrdersWebSocketClient::new(
838            "wss://example.com/orders/ws".to_string(),
839            AccountId::from("AX-001"),
840            TraderId::from("TRADER-001"),
841            30,
842            TransportBackend::default(),
843            None,
844        );
845        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
846        let handle = tokio::spawn(async move {
847            started_tx.send(()).expect("started receiver");
848            std::future::pending::<()>().await;
849        });
850        let abort_handle = handle.abort_handle();
851        client.task_handle.insert(handle);
852        started_rx.await.expect("handler task started");
853
854        drop(client);
855
856        tokio::time::timeout(Duration::from_secs(1), async {
857            while !abort_handle.is_finished() {
858                tokio::task::yield_now().await;
859            }
860        })
861        .await
862        .expect("handler task aborted");
863    }
864
865    #[rstest]
866    fn test_reserve_cid_mapping_rejects_collision() {
867        let caches = OrdersCaches::default();
868        let cid = 123;
869        let first_client_order_id = ClientOrderId::from("CID-123-A");
870        let second_client_order_id = ClientOrderId::from("CID-123-B");
871
872        reserve_cid_mapping(&caches, cid, first_client_order_id).unwrap();
873        let result = reserve_cid_mapping(&caches, cid, second_client_order_id);
874
875        assert!(matches!(
876            result,
877            Err(AxOrdersWsClientError::ClientError(msg))
878                if msg == "AX cid 123 is already mapped to CID-123-A"
879        ));
880        assert_eq!(
881            caches
882                .cid_to_client_order_id
883                .get(&cid)
884                .map(|client_order_id| *client_order_id),
885            Some(first_client_order_id),
886        );
887    }
888
889    #[tokio::test]
890    async fn test_cancel_order_rejects_without_venue_order_id() {
891        let client = AxOrdersWebSocketClient::new(
892            "wss://example.com/orders/ws".to_string(),
893            AccountId::from("AX-001"),
894            TraderId::from("TRADER-001"),
895            30,
896            TransportBackend::default(),
897            None,
898        );
899        let client_order_id = ClientOrderId::from("CID-123");
900
901        let result = client.cancel_order(client_order_id, None).await;
902
903        assert!(matches!(
904            result,
905            Err(AxOrdersWsClientError::ClientError(msg))
906            if msg.contains("missing venue_order_id")
907        ));
908    }
909
910    #[tokio::test]
911    async fn test_cancel_order_sends_known_venue_order_id() {
912        let mut client = AxOrdersWebSocketClient::new(
913            "wss://example.com/orders/ws".to_string(),
914            AccountId::from("AX-001"),
915            TraderId::from("TRADER-001"),
916            30,
917            TransportBackend::default(),
918            None,
919        );
920
921        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
922        client.cmd_tx = Arc::new(tokio::sync::RwLock::new(cmd_tx));
923
924        let client_order_id = ClientOrderId::from("CID-456");
925        let venue_order_id = VenueOrderId::from("V-ORDER-789");
926
927        let request_id = client
928            .cancel_order(client_order_id, Some(venue_order_id))
929            .await
930            .unwrap();
931
932        assert_eq!(request_id, 1);
933        let cmd = cmd_rx.recv().await.unwrap();
934        match cmd {
935            HandlerCommand::CancelOrder {
936                request_id,
937                order_id,
938            } => {
939                assert_eq!(request_id, 1);
940                assert_eq!(order_id, "V-ORDER-789");
941            }
942            other => panic!("unexpected command: {other:?}"),
943        }
944    }
945}
nautilus_architect_ax/websocket/orders/client.rs

nautilus_architect_ax/websocket/orders/
client.rs
