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

nautilus_architect_ax/websocket/orders/
client.rs
