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

nautilus_architect_ax/websocket/orders/
client.rs
