1use std::{
26 fmt::Debug,
27 num::NonZeroU32,
28 sync::{
29 Arc, LazyLock,
30 atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
31 },
32 time::Duration,
33};
34
35use ahash::{AHashMap, AHashSet};
36use arc_swap::ArcSwap;
37use dashmap::DashMap;
38use futures_util::Stream;
39use nautilus_common::live::dst::time;
40use nautilus_core::{
41 AtomicMap, AtomicTime, UnixNanos,
42 env::{get_env_var, get_or_env_var},
43 string::secret::{REDACTED, SecretString},
44 time::get_atomic_clock_realtime,
45};
46use nautilus_live::{
47 SocketControl,
48 task::{TaskGroup, TaskShutdownError},
49};
50use nautilus_model::{
51 data::BarType,
52 enums::{OrderSide, OrderType, PositionSide, TimeInForce, TriggerType},
53 identifiers::{AccountId, ClientOrderId, InstrumentId, StrategyId, TraderId, VenueOrderId},
54 instruments::{Instrument, InstrumentAny},
55 types::{Price, Quantity},
56};
57use nautilus_network::{
58 http::create_standard_nautilus_headers,
59 mode::ConnectionMode,
60 ratelimiter::quota::Quota,
61 websocket::{
62 AUTHENTICATION_TIMEOUT_SECS, AuthTracker, SubscriptionState, TEXT_PING, TransportBackend,
63 WebSocketClient, WebSocketConfig, channel_message_handler,
64 },
65};
66use parking_lot::Mutex;
67use serde_json::Value;
68use tokio_tungstenite::tungstenite::Error;
69use tokio_util::sync::CancellationToken;
70use ustr::Ustr;
71
72use super::{
73 enums::OKXWsChannel,
74 error::OKXWsError,
75 handler::{HandlerCommand, OKXWsFeedHandler, SnapshotGate},
76 messages::{
77 OKXAuthentication, OKXAuthenticationArg, OKXSubscriptionArg, OKXWsMessage, OKXWsRequest,
78 WsAmendOrderParamsBuilder, WsAttachAlgoOrdParams, WsCancelOrderParamsBuilder,
79 WsMassCancelParams, WsPostAlgoOrderParamsBuilder, WsPostOrderParamsBuilder,
80 },
81 subscription::topic_from_subscription_arg,
82};
83use crate::common::{
84 consts::{
85 OKX_NAUTILUS_BROKER_ID, OKX_SUPPORTED_ORDER_TYPES, OKX_SUPPORTED_TIME_IN_FORCE,
86 OKX_WS_PUBLIC_URL, OKX_WS_TOPIC_DELIMITER, okx_reduce_only_wire_value, select_book_channel,
87 spot_trade_quote_ccy_wire_value,
88 },
89 credential::Credential,
90 enums::{
91 OKXBookChannel, OKXGreeksType, OKXInstrumentType, OKXOrderType, OKXPositionSide,
92 OKXTargetCurrency, OKXTradeMode, OKXTriggerType, OKXVipLevel,
93 conditional_order_to_algo_type, is_conditional_order,
94 },
95 parse::{
96 bar_spec_as_okx_channel, okx_instrument_type, okx_instrument_type_from_symbol,
97 parse_base_quote_from_symbol,
98 },
99};
100
101pub static OKX_WS_CONNECTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
105 Quota::per_second(NonZeroU32::new(3).expect("non-zero")).expect("valid constant")
106});
107
108pub static OKX_WS_SUBSCRIPTION_QUOTA: LazyLock<Quota> =
113 LazyLock::new(|| Quota::per_hour(NonZeroU32::new(480).expect("non-zero")));
114
115pub static OKX_WS_ORDER_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
117 Quota::per_second(NonZeroU32::new(30).expect("non-zero")).expect("valid constant")
118});
119
120pub static OKX_WS_BATCH_ORDER_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
122 Quota::per_second(NonZeroU32::new(7).expect("non-zero")).expect("valid constant")
123});
124
125pub static OKX_WS_MASS_CANCEL_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
127 Quota::per_second(NonZeroU32::new(2).expect("non-zero")).expect("valid constant")
128});
129
130pub static OKX_WS_ALGO_ORDER_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
132 Quota::per_second(NonZeroU32::new(10).expect("non-zero")).expect("valid constant")
133});
134
135pub static OKX_WS_ALGO_CANCEL_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
137 Quota::per_second(NonZeroU32::new(1).expect("non-zero")).expect("valid constant")
138});
139
140const OKX_WS_SUBSCRIPTION_ARGS_MAX_PER_MESSAGE: usize = 256;
142
143const RECONNECT_AUTH_RETRY_INITIAL: Duration = Duration::from_secs(1);
144const RECONNECT_AUTH_RETRY_MAX: Duration = Duration::from_secs(30);
145
146pub static OKX_RATE_LIMIT_KEY_SUBSCRIPTION: LazyLock<[Ustr; 1]> =
151 LazyLock::new(|| [Ustr::from("subscription")]);
152
153pub static OKX_RATE_LIMIT_KEY_ORDER: LazyLock<[Ustr; 1]> = LazyLock::new(|| [Ustr::from("order")]);
157
158pub static OKX_RATE_LIMIT_KEY_BATCH_ORDER: LazyLock<[Ustr; 1]> =
162 LazyLock::new(|| [Ustr::from("batch-order")]);
163
164pub static OKX_RATE_LIMIT_KEY_CANCEL: LazyLock<[Ustr; 1]> =
168 LazyLock::new(|| [Ustr::from("cancel")]);
169
170pub static OKX_RATE_LIMIT_KEY_BATCH_CANCEL: LazyLock<[Ustr; 1]> =
174 LazyLock::new(|| [Ustr::from("batch-cancel")]);
175
176pub static OKX_RATE_LIMIT_KEY_MASS_CANCEL: LazyLock<[Ustr; 1]> =
180 LazyLock::new(|| [Ustr::from("mass-cancel")]);
181
182pub static OKX_RATE_LIMIT_KEY_AMEND: LazyLock<[Ustr; 1]> = LazyLock::new(|| [Ustr::from("amend")]);
186
187pub static OKX_RATE_LIMIT_KEY_BATCH_AMEND: LazyLock<[Ustr; 1]> =
191 LazyLock::new(|| [Ustr::from("batch-amend")]);
192
193pub static OKX_RATE_LIMIT_KEY_ALGO_ORDER: LazyLock<[Ustr; 1]> =
197 LazyLock::new(|| [Ustr::from("algo-order")]);
198
199pub static OKX_RATE_LIMIT_KEY_ALGO_CANCEL: LazyLock<[Ustr; 1]> =
203 LazyLock::new(|| [Ustr::from("algo-cancel")]);
204
205#[derive(Debug, Clone)]
209#[allow(dead_code)]
210#[allow(
211 clippy::struct_field_names,
212 reason = "fields follow the codebase-wide trader_id/strategy_id/instrument_id naming family"
213)]
214pub(crate) struct PendingOrderInfo {
215 pub trader_id: TraderId,
216 pub strategy_id: StrategyId,
217 pub instrument_id: InstrumentId,
218}
219
220#[derive(Clone)]
222pub struct OKXWebSocketClient {
223 clock: &'static AtomicTime,
224 url: String,
225 vip_level: Arc<AtomicU8>,
226 credential: Option<Credential>,
227 heartbeat: Option<u64>,
228 auth_timeout_secs: u64,
229 auth_tracker: AuthTracker,
230 signal: Arc<AtomicBool>,
231 connection_mode: Arc<ArcSwap<AtomicU8>>,
232 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
233 out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<OKXWsMessage>>>,
234 handler_tasks: Arc<TaskGroup>,
235 connect_lock: Arc<tokio::sync::Mutex<()>>,
236 handler_abort: Arc<Mutex<CancellationToken>>,
237 subscriptions_inst_type: Arc<DashMap<OKXWsChannel, AHashSet<OKXInstrumentType>>>,
238 subscriptions_inst_family: Arc<DashMap<OKXWsChannel, AHashSet<Ustr>>>,
239 subscriptions_inst_id: Arc<DashMap<OKXWsChannel, AHashSet<Ustr>>>,
240 subscriptions_bare: Arc<DashMap<OKXWsChannel, bool>>,
241 subscriptions_state: SubscriptionState,
242 request_id_counter: Arc<AtomicU64>,
243 instruments_cache: Arc<AtomicMap<Ustr, InstrumentAny>>,
244 inst_id_code_cache: Arc<AtomicMap<Ustr, u64>>,
245 trade_quote_ccy_lists: Arc<AtomicMap<Ustr, Vec<Ustr>>>,
246 spot_trade_quote_ccy: Arc<Mutex<Option<Ustr>>>,
247 pub(crate) pending_orders: Arc<DashMap<String, PendingOrderInfo>>,
248 pub(crate) pending_cancels: Arc<DashMap<String, PendingOrderInfo>>,
249 pub(crate) pending_amends: Arc<DashMap<String, PendingOrderInfo>>,
250 option_greeks_subs: Arc<AtomicMap<InstrumentId, AHashSet<OKXGreeksType>>>,
251 index_pair_subscribers: Arc<DashMap<Ustr, usize>>,
258 index_pair_transition: Arc<tokio::sync::Mutex<()>>,
263 transport_backend: TransportBackend,
265 proxy_url: Option<SecretString>,
267 cancellation_token: CancellationToken,
268 socket_control: Option<Arc<SocketControl>>,
269}
270
271struct ConnectRollback {
272 handler_tasks: Arc<TaskGroup>,
273 signal: Arc<AtomicBool>,
274 handler_abort: CancellationToken,
275 socket_control: Option<Arc<SocketControl>>,
276 armed: bool,
277}
278
279impl ConnectRollback {
280 fn disarm(&mut self) {
281 self.armed = false;
282 }
283}
284
285impl Drop for ConnectRollback {
286 fn drop(&mut self) {
287 if !self.armed {
288 return;
289 }
290
291 self.handler_tasks.begin_shutdown();
292 self.signal.store(true, Ordering::Release);
293 self.handler_abort.cancel();
294
295 if let Some(control) = &self.socket_control {
296 control.deregister();
297 }
298 }
299}
300
301impl Default for OKXWebSocketClient {
302 fn default() -> Self {
303 Self::new(
304 None,
305 None,
306 None,
307 None,
308 None,
309 None,
310 None,
311 TransportBackend::default(),
312 None,
313 )
314 .unwrap()
315 }
316}
317
318impl Debug for OKXWebSocketClient {
319 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
320 f.debug_struct(stringify!(OKXWebSocketClient))
321 .field("url", &self.url)
322 .field("credential", &self.credential.as_ref().map(|_| REDACTED))
323 .field("heartbeat", &self.heartbeat)
324 .finish_non_exhaustive()
325 }
326}
327
328impl OKXWebSocketClient {
329 #[allow(clippy::too_many_arguments)]
335 pub fn new(
336 url: Option<String>,
337 api_key: Option<String>,
338 api_secret: Option<String>,
339 api_passphrase: Option<String>,
340 _account_id: Option<AccountId>,
341 heartbeat: Option<u64>,
342 auth_timeout_secs: Option<u64>,
343 transport_backend: TransportBackend,
344 proxy_url: Option<String>,
345 ) -> anyhow::Result<Self> {
346 let url = url.unwrap_or(OKX_WS_PUBLIC_URL.to_string());
347 let credential = match (api_key, api_secret, api_passphrase) {
348 (Some(key), Some(secret), Some(passphrase)) => {
349 Some(Credential::new(key, secret, passphrase))
350 }
351 (None, None, None) => None,
352 _ => anyhow::bail!(
353 "`api_key`, `api_secret`, `api_passphrase` credentials must be provided together"
354 ),
355 };
356
357 let signal = Arc::new(AtomicBool::new(false));
358 let subscriptions_inst_type = Arc::new(DashMap::new());
359 let subscriptions_inst_family = Arc::new(DashMap::new());
360 let subscriptions_inst_id = Arc::new(DashMap::new());
361 let subscriptions_bare = Arc::new(DashMap::new());
362 let subscriptions_state = SubscriptionState::new(OKX_WS_TOPIC_DELIMITER);
363
364 Ok(Self {
365 clock: get_atomic_clock_realtime(),
366 url,
367 vip_level: Arc::new(AtomicU8::new(0)),
368 credential,
369 heartbeat,
370 auth_timeout_secs: auth_timeout_secs.unwrap_or(AUTHENTICATION_TIMEOUT_SECS),
371 auth_tracker: AuthTracker::new(),
372 signal,
373 connection_mode: Arc::new(ArcSwap::from_pointee(AtomicU8::new(
374 ConnectionMode::Closed.as_u8(),
375 ))),
376 cmd_tx: {
377 let (tx, _) = tokio::sync::mpsc::unbounded_channel();
379 Arc::new(tokio::sync::RwLock::new(tx))
380 },
381 out_rx: None,
382 handler_tasks: Arc::new(TaskGroup::new()),
383 connect_lock: Arc::new(tokio::sync::Mutex::new(())),
384 handler_abort: Arc::new(Mutex::new(CancellationToken::new())),
385 subscriptions_inst_type,
386 subscriptions_inst_family,
387 subscriptions_inst_id,
388 subscriptions_bare,
389 subscriptions_state,
390 request_id_counter: Arc::new(AtomicU64::new(1)),
391 instruments_cache: Arc::new(AtomicMap::new()),
392 inst_id_code_cache: Arc::new(AtomicMap::new()),
393 trade_quote_ccy_lists: Arc::new(AtomicMap::new()),
394 spot_trade_quote_ccy: Arc::new(Mutex::new(None)),
395 pending_orders: Arc::new(DashMap::new()),
396 pending_cancels: Arc::new(DashMap::new()),
397 pending_amends: Arc::new(DashMap::new()),
398 option_greeks_subs: Arc::new(AtomicMap::new()),
399 index_pair_subscribers: Arc::new(DashMap::new()),
400 index_pair_transition: Arc::new(tokio::sync::Mutex::new(())),
401 transport_backend,
402 proxy_url: proxy_url.map(SecretString::from),
403 cancellation_token: CancellationToken::new(),
404 socket_control: None,
405 })
406 }
407
408 #[must_use]
410 pub fn with_socket_control(mut self, control: SocketControl) -> Self {
411 self.socket_control = Some(Arc::new(control));
412 self
413 }
414
415 #[allow(clippy::too_many_arguments)]
422 pub fn with_credentials(
423 url: Option<String>,
424 api_key: Option<String>,
425 api_secret: Option<String>,
426 api_passphrase: Option<String>,
427 account_id: Option<AccountId>,
428 heartbeat: Option<u64>,
429 auth_timeout_secs: Option<u64>,
430 transport_backend: TransportBackend,
431 proxy_url: Option<String>,
432 ) -> anyhow::Result<Self> {
433 let url = url.unwrap_or(OKX_WS_PUBLIC_URL.to_string());
434 let api_key = get_or_env_var(api_key, "OKX_API_KEY")?;
435 let api_secret = get_or_env_var(api_secret, "OKX_API_SECRET")?;
436 let api_passphrase = get_or_env_var(api_passphrase, "OKX_API_PASSPHRASE")?;
437
438 Self::new(
439 Some(url),
440 Some(api_key),
441 Some(api_secret),
442 Some(api_passphrase),
443 account_id,
444 heartbeat,
445 auth_timeout_secs,
446 transport_backend,
447 proxy_url,
448 )
449 }
450
451 pub fn from_env() -> anyhow::Result<Self> {
458 let url = get_env_var("OKX_WS_URL")?;
459 let api_key = get_env_var("OKX_API_KEY")?;
460 let api_secret = get_env_var("OKX_API_SECRET")?;
461 let api_passphrase = get_env_var("OKX_API_PASSPHRASE")?;
462
463 Self::new(
464 Some(url),
465 Some(api_key),
466 Some(api_secret),
467 Some(api_passphrase),
468 None,
469 None,
470 None,
471 TransportBackend::default(),
472 None,
473 )
474 }
475
476 pub fn cancel_all_requests(&self) {
478 self.cancellation_token.cancel();
479 }
480
481 pub fn cancellation_token(&self) -> &CancellationToken {
483 &self.cancellation_token
484 }
485
486 pub fn url(&self) -> &str {
488 self.url.as_str()
489 }
490
491 pub fn api_key(&self) -> Option<&str> {
493 self.credential.as_ref().map(Credential::api_key)
494 }
495
496 #[must_use]
498 pub fn api_key_masked(&self) -> Option<String> {
499 self.credential.as_ref().map(Credential::api_key_masked)
500 }
501
502 pub fn is_active(&self) -> bool {
504 let connection_mode_arc = self.connection_mode.load();
505 ConnectionMode::from_atomic(&connection_mode_arc).is_active()
506 && !self.signal.load(Ordering::Acquire)
507 }
508
509 pub fn is_closed(&self) -> bool {
511 let connection_mode_arc = self.connection_mode.load();
512 ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
513 || self.signal.load(Ordering::Acquire)
514 }
515
516 pub(crate) fn has_task(&self) -> bool {
518 !self.handler_tasks.is_empty()
519 }
520
521 pub fn cache_instruments(&self, instruments: &[InstrumentAny]) {
525 self.instruments_cache.rcu(|m| {
526 for inst in instruments {
527 m.insert(inst.symbol().inner(), inst.clone());
528 }
529 });
530 }
531
532 pub fn cache_instrument(&self, instrument: InstrumentAny) {
536 self.instruments_cache
537 .insert(instrument.symbol().inner(), instrument);
538 }
539
540 pub fn instruments_snapshot(&self) -> AHashMap<Ustr, InstrumentAny> {
542 (**self.instruments_cache.load()).clone()
543 }
544
545 pub fn instruments_cache_arc(&self) -> Arc<AtomicMap<Ustr, InstrumentAny>> {
547 Arc::clone(&self.instruments_cache)
548 }
549
550 pub fn cache_inst_id_code(&self, inst_id: Ustr, inst_id_code: u64) {
554 self.inst_id_code_cache.insert(inst_id, inst_id_code);
555 }
556
557 pub fn cache_inst_id_codes(&self, mappings: impl IntoIterator<Item = (Ustr, u64)>) {
561 let entries: Vec<_> = mappings.into_iter().collect();
562 self.inst_id_code_cache.rcu(|m| {
563 for (inst_id, inst_id_code) in &entries {
564 m.insert(*inst_id, *inst_id_code);
565 }
566 });
567 }
568
569 #[must_use]
573 pub fn get_inst_id_code(&self, inst_id: &Ustr) -> Option<u64> {
574 self.inst_id_code_cache.load().get(inst_id).copied()
575 }
576
577 pub fn set_spot_trade_quote_ccy(&self, ccy: Option<String>) {
579 *self.spot_trade_quote_ccy.lock() = ccy.map(|value| Ustr::from(value.as_str()));
580 }
581
582 pub fn cache_trade_quote_ccy_lists(
584 &self,
585 mappings: impl IntoIterator<Item = (Ustr, Vec<Ustr>)>,
586 ) {
587 let entries: Vec<_> = mappings.into_iter().collect();
588 self.trade_quote_ccy_lists.rcu(|m| {
589 for (inst_id, list) in &entries {
590 m.insert(*inst_id, list.clone());
591 }
592 });
593 }
594
595 fn resolve_spot_trade_quote_ccy(
596 &self,
597 instrument_type: OKXInstrumentType,
598 inst_id: Ustr,
599 ) -> Result<Option<Ustr>, OKXWsError> {
600 let configured = *self.spot_trade_quote_ccy.lock();
601 let available = self
602 .trade_quote_ccy_lists
603 .load()
604 .get(&inst_id)
605 .cloned()
606 .unwrap_or_default();
607 spot_trade_quote_ccy_wire_value(
608 instrument_type,
609 configured.as_ref().map(Ustr::as_str),
610 &available,
611 )
612 .map_err(OKXWsError::ClientError)
613 }
614
615 fn inst_id_symbol_and_code_from_snapshot(
616 inst_id_codes: &AHashMap<Ustr, u64>,
617 inst_id: &InstrumentId,
618 action: &str,
619 ) -> Result<(Ustr, u64), OKXWsError> {
620 let inst_id_symbol = inst_id.symbol.inner();
621 let inst_id_code = inst_id_codes.get(&inst_id_symbol).copied().ok_or_else(|| {
622 OKXWsError::ClientError(format!(
623 "No instIdCode cached for {inst_id}, cannot {action} order"
624 ))
625 })?;
626 Ok((inst_id_symbol, inst_id_code))
627 }
628
629 pub fn set_vip_level(&self, vip_level: OKXVipLevel) {
633 self.vip_level.store(vip_level as u8, Ordering::Relaxed);
634 }
635
636 pub fn vip_level(&self) -> OKXVipLevel {
638 let level = self.vip_level.load(Ordering::Relaxed);
639 OKXVipLevel::from(level)
640 }
641
642 pub async fn connect(&mut self) -> anyhow::Result<()> {
648 let connect_lock = Arc::clone(&self.connect_lock);
649 let _connect_guard = connect_lock.lock().await;
650
651 if !self.handler_tasks.is_empty() && !self.handler_tasks.all_finished() {
652 anyhow::bail!("Cannot connect while previous WebSocket handler task is still running");
653 }
654
655 if !self.handler_tasks.is_open() || !self.handler_tasks.is_empty() {
656 self.handler_tasks.begin_shutdown();
657 self.handler_tasks
658 .finish_shutdown(Duration::from_secs(2), Duration::from_secs(2))
659 .await
660 .map_err(|e| anyhow::anyhow!("Previous WebSocket handler failed: {e}"))?;
661 self.handler_tasks.start_generation().map_err(|e| {
662 anyhow::anyhow!("Failed to start WebSocket handler task generation: {e}")
663 })?;
664 }
665 let handler_spawner = self.handler_tasks.spawner().map_err(|e| {
666 anyhow::anyhow!("Failed to acquire WebSocket handler task spawner: {e}")
667 })?;
668 let handler_abort = CancellationToken::new();
669 *self.handler_abort.lock() = handler_abort.clone();
670 let mut rollback = ConnectRollback {
671 handler_tasks: Arc::clone(&self.handler_tasks),
672 signal: Arc::clone(&self.signal),
673 handler_abort: handler_abort.clone(),
674 socket_control: self.socket_control.clone(),
675 armed: true,
676 };
677
678 self.signal.store(false, Ordering::Release);
680
681 let (message_handler, raw_rx) = channel_message_handler();
682
683 let headers = create_standard_nautilus_headers();
689
690 let config = WebSocketConfig {
691 url: self.url.clone(),
692 headers,
693 heartbeat_interval_secs: self.heartbeat,
694 heartbeat_payload: Some(TEXT_PING.to_string()),
695 connect_timeout_ms: None,
696 reconnect_delay_initial_ms: None,
697 reconnect_delay_max_ms: None,
698 reconnect_backoff_factor: None,
699 reconnect_jitter_ms: None,
700 reconnect_max_attempts: None,
701 heartbeat_timeout_secs: None,
702 idle_timeout_ms: None,
703 writer_capacity: None,
704 backend: self.transport_backend,
705 proxy_url: self
706 .proxy_url
707 .as_ref()
708 .map(|value| value.expose_secret().to_owned()),
709 max_message_size_bytes: None,
710 max_frame_size_bytes: None,
711 };
712
713 let keyed_quotas = vec![
714 (
715 OKX_RATE_LIMIT_KEY_SUBSCRIPTION[0].to_string(),
716 *OKX_WS_SUBSCRIPTION_QUOTA,
717 ),
718 (OKX_RATE_LIMIT_KEY_ORDER[0].to_string(), *OKX_WS_ORDER_QUOTA),
719 (
720 OKX_RATE_LIMIT_KEY_BATCH_ORDER[0].to_string(),
721 *OKX_WS_BATCH_ORDER_QUOTA,
722 ),
723 (
724 OKX_RATE_LIMIT_KEY_CANCEL[0].to_string(),
725 *OKX_WS_ORDER_QUOTA,
726 ),
727 (
728 OKX_RATE_LIMIT_KEY_BATCH_CANCEL[0].to_string(),
729 *OKX_WS_BATCH_ORDER_QUOTA,
730 ),
731 (
732 OKX_RATE_LIMIT_KEY_MASS_CANCEL[0].to_string(),
733 *OKX_WS_MASS_CANCEL_QUOTA,
734 ),
735 (OKX_RATE_LIMIT_KEY_AMEND[0].to_string(), *OKX_WS_ORDER_QUOTA),
736 (
737 OKX_RATE_LIMIT_KEY_BATCH_AMEND[0].to_string(),
738 *OKX_WS_BATCH_ORDER_QUOTA,
739 ),
740 (
741 OKX_RATE_LIMIT_KEY_ALGO_ORDER[0].to_string(),
742 *OKX_WS_ALGO_ORDER_QUOTA,
743 ),
744 (
745 OKX_RATE_LIMIT_KEY_ALGO_CANCEL[0].to_string(),
746 *OKX_WS_ALGO_CANCEL_QUOTA,
747 ),
748 ];
749
750 let client = WebSocketClient::builder()
751 .config(config)
752 .message_handler(message_handler)
753 .keyed_quotas(keyed_quotas)
754 .default_quota(*OKX_WS_CONNECTION_QUOTA)
755 .maybe_state_sink(self.socket_control.as_ref().map(|control| control.sink()))
756 .connect()
757 .await?;
758
759 self.connection_mode.store(client.connection_mode_atomic());
761 let reconnect_handle = client.reconnect_handle();
762
763 let (msg_tx, rx) = tokio::sync::mpsc::unbounded_channel::<OKXWsMessage>();
764
765 self.out_rx = Some(Arc::new(rx));
766
767 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
768 *self.cmd_tx.write().await = cmd_tx.clone();
769
770 let signal = self.signal.clone();
771 let auth_tracker = self.auth_tracker.clone();
772 let subscriptions_state = self.subscriptions_state.clone();
773 let clock = self.clock;
774
775 let handler_task = {
776 let auth_tracker = auth_tracker.clone();
777 let signal = signal.clone();
778 let credential = self.credential.clone();
779 let cmd_tx_for_reconnect = cmd_tx.clone();
780 let subscriptions_bare = self.subscriptions_bare.clone();
781 let subscriptions_inst_type = self.subscriptions_inst_type.clone();
782 let subscriptions_inst_family = self.subscriptions_inst_family.clone();
783 let subscriptions_inst_id = self.subscriptions_inst_id.clone();
784 let cmd_tx_for_auth = Arc::clone(&self.cmd_tx);
785 let auth_tracker_for_auth = self.auth_tracker.clone();
786 let clock_for_auth = self.clock;
787 let auth_timeout_secs = self.auth_timeout_secs;
788 let reconnect_auth_generation = Arc::new(AtomicU64::new(0));
789 let handler_spawner_for_auth = handler_spawner.clone();
790 let mut has_reconnected = false;
791
792 async move {
793 let mut handler = OKXWsFeedHandler::new(
794 signal.clone(),
795 cmd_rx,
796 raw_rx,
797 msg_tx,
798 auth_tracker.clone(),
799 subscriptions_state.clone(),
800 clock,
801 );
802
803 let resubscribe_all = || {
804 let args = subscription_args(
805 &subscriptions_inst_type,
806 &subscriptions_inst_family,
807 &subscriptions_inst_id,
808 &subscriptions_bare,
809 );
810
811 for chunk in subscription_arg_batches(&args) {
812 if let Err(e) = cmd_tx_for_reconnect.send(HandlerCommand::Subscribe {
813 args: chunk.to_vec(),
814 }) {
815 log::error!("Failed to send resubscribe command: error={e}");
816 }
817 }
818 };
819
820 loop {
821 let message = tokio::select! {
822 biased;
823 () = handler_abort.cancelled() => {
824 log::debug!("Handler task aborted");
825 break;
826 }
827 message = handler.next() => message,
828 };
829
830 match message {
831 Some(OKXWsMessage::Reconnected) => {
832 if signal.load(Ordering::Acquire) {
833 continue;
834 }
835
836 has_reconnected = true;
837
838 subscriptions_state.reset_after_reconnect();
839
840 if let Some(credential) = credential.clone() {
841 log::debug!("Re-authenticating after reconnection");
842 let generation =
843 reconnect_auth_generation.fetch_add(1, Ordering::AcqRel) + 1;
844 let signal = signal.clone();
845 let abort = handler_abort.clone();
846 let generation_flag = Arc::clone(&reconnect_auth_generation);
847 let cmd_tx = Arc::clone(&cmd_tx_for_auth);
848 let auth_tracker = auth_tracker_for_auth.clone();
849
850 if let Err(e) = handler_spawner_for_auth.spawn(async move {
851 retry_reconnect_authentication(
852 credential,
853 auth_tracker,
854 cmd_tx,
855 clock_for_auth,
856 auth_timeout_secs,
857 signal,
858 abort,
859 generation_flag,
860 generation,
861 )
862 .await;
863 }) {
864 log::error!(
865 "Failed to spawn re-authentication retry task: error={e}"
866 );
867 }
868 }
869
870 if credential.is_none() {
873 log::debug!(
874 "No authentication required, resubscribing immediately"
875 );
876 resubscribe_all();
877 }
878
879 if handler.send(OKXWsMessage::Reconnected).is_err() {
881 log_receiver_dropped(&signal, "Reconnected");
882 break;
883 }
884 }
885 Some(OKXWsMessage::Authenticated) => {
886 if has_reconnected {
887 resubscribe_all();
888 }
889 }
890 Some(msg) => {
891 if handler.send(msg).is_err() {
892 log_receiver_dropped(&signal, "message");
893 break;
894 }
895 }
896 None => {
897 if handler.is_stopped() {
898 log::debug!("Stop signal received, ending message processing");
899 break;
900 }
901 log::debug!("WebSocket stream closed");
902 break;
903 }
904 }
905 }
906
907 log::debug!("Handler task exiting");
908 }
909 };
910
911 if let Err(e) = handler_spawner.spawn(handler_task) {
912 self.out_rx = None;
913 anyhow::bail!("Failed to register WebSocket handler task: {e}");
914 }
915
916 let set_client_result = {
917 let cmd_tx = self.cmd_tx.read().await;
918 cmd_tx.send(HandlerCommand::SetClient(client))
919 };
920
921 if let Err(e) = set_client_result {
922 self.handler_tasks.begin_shutdown();
923 self.signal.store(true, Ordering::Release);
924 let handler_abort = self.handler_abort.lock().clone();
925 handler_abort.cancel();
926 let shutdown_result = self.close_stream_task(Duration::from_secs(2)).await;
927 self.out_rx = None;
928 anyhow::bail!(match shutdown_result {
929 Ok(()) => format!("Failed to send WebSocket client to handler: {e}"),
930 Err(shutdown_error) => format!(
931 "Failed to send WebSocket client to handler: {e}; handler shutdown failed: \
932 {shutdown_error}"
933 ),
934 });
935 }
936
937 if let Some(control) = &self.socket_control {
938 control.register(move || reconnect_handle.request_reconnect());
939 }
940 log::debug!("Sent WebSocket client to handler");
941
942 if self.credential.is_some()
943 && let Err(e) = self.authenticate().await
944 {
945 self.handler_tasks.begin_shutdown();
946 self.request_close().await;
947 let shutdown_result = self.close_stream_task(Duration::from_secs(2)).await;
948
949 if let Some(control) = &self.socket_control {
950 control.deregister();
951 }
952 self.out_rx = None;
953
954 match shutdown_result {
955 Ok(()) => anyhow::bail!("Authentication failed: {e}"),
956 Err(shutdown_error) => anyhow::bail!(
957 "Authentication failed: {e}; handler shutdown failed: {shutdown_error}"
958 ),
959 }
960 }
961
962 rollback.disarm();
963 Ok(())
964 }
965
966 async fn authenticate(&self) -> Result<(), Error> {
968 let credential = self.credential.as_ref().ok_or_else(|| {
969 Error::Io(std::io::Error::other(
970 "API credentials not available to authenticate",
971 ))
972 })?;
973
974 match authenticate_session(
975 credential,
976 &self.auth_tracker,
977 &self.cmd_tx,
978 self.clock,
979 self.auth_timeout_secs,
980 )
981 .await
982 {
983 Ok(()) => Ok(()),
984 Err(e) if auth_attempt_superseded(&e) => {
985 if self
988 .auth_tracker
989 .wait_for_authenticated(Duration::from_secs(self.auth_timeout_secs))
990 .await
991 {
992 Ok(())
993 } else {
994 Err(e)
995 }
996 }
997 Err(e) => Err(e),
998 }
999 }
1000
1001 pub fn stream(&mut self) -> impl Stream<Item = OKXWsMessage> + 'static {
1009 let rx = self
1010 .out_rx
1011 .take()
1012 .expect("Data stream receiver already taken or not connected");
1013 let mut rx = Arc::try_unwrap(rx).expect("Cannot take ownership - other references exist");
1014 async_stream::stream! {
1015 while let Some(data) = rx.recv().await {
1016 yield data;
1017 }
1018 }
1019 }
1020
1021 pub async fn wait_until_active(&self, timeout_secs: f64) -> Result<(), OKXWsError> {
1027 let timeout = time::Duration::from_secs_f64(timeout_secs);
1028
1029 time::timeout(timeout, async {
1030 while !self.is_active() {
1031 time::sleep(time::Duration::from_millis(10)).await;
1032 }
1033 })
1034 .await
1035 .map_err(|_| {
1036 OKXWsError::ClientError(format!(
1037 "WebSocket connection timeout after {timeout_secs} seconds"
1038 ))
1039 })?;
1040
1041 Ok(())
1042 }
1043
1044 pub(crate) fn begin_shutdown(&self) {
1045 self.handler_tasks.begin_shutdown();
1046 self.signal.store(true, Ordering::Release);
1047
1048 let handler_abort = self.handler_abort.lock().clone();
1049 handler_abort.cancel();
1050 }
1051
1052 pub(crate) async fn request_close(&self) {
1054 self.signal.store(true, Ordering::Release);
1055
1056 if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
1057 log::debug!("Handler channel closed before disconnect command was sent: {e}");
1058 } else {
1059 log::debug!("Sent disconnect command to handler");
1060 }
1061 }
1062
1063 pub async fn close(&mut self) -> Result<(), Error> {
1070 let connect_lock = Arc::clone(&self.connect_lock);
1071 let _connect_guard = connect_lock.lock().await;
1072
1073 self.close_locked().await
1074 }
1075
1076 async fn close_locked(&self) -> Result<(), Error> {
1077 log::debug!("Starting close process");
1078
1079 self.handler_tasks.begin_shutdown();
1080 self.request_close().await;
1081
1082 let task_result = self.close_stream_task(Duration::from_secs(2)).await;
1083
1084 self.index_pair_subscribers.clear();
1088
1089 if let Some(control) = &self.socket_control {
1090 control.deregister();
1091 }
1092
1093 log::debug!("Close process completed");
1094
1095 task_result
1096 }
1097
1098 async fn close_stream_task(&self, timeout: Duration) -> Result<(), Error> {
1099 match self.handler_tasks.finish_shutdown(timeout, timeout).await {
1100 Ok(()) => Ok(()),
1101 Err(error @ TaskShutdownError::Timeout { .. }) => Err(Error::Io(std::io::Error::new(
1102 std::io::ErrorKind::TimedOut,
1103 format!("Timed out joining WebSocket handler task after abort: {error}"),
1104 ))),
1105 Err(e) => Err(Error::Io(std::io::Error::other(format!(
1106 "WebSocket handler shutdown failed: {e}"
1107 )))),
1108 }
1109 }
1110
1111 pub fn get_subscriptions(&self, instrument_id: InstrumentId) -> Vec<OKXWsChannel> {
1113 let symbol = instrument_id.symbol.inner();
1114 let mut channels = Vec::new();
1115
1116 for entry in self.subscriptions_inst_id.iter() {
1117 let (channel, instruments) = entry.pair();
1118 if instruments.contains(&symbol) {
1119 channels.push(channel.clone());
1120 }
1121 }
1122
1123 channels
1124 }
1125
1126 fn generate_unique_request_id(&self) -> String {
1127 self.request_id_counter
1128 .fetch_add(1, Ordering::SeqCst)
1129 .to_string()
1130 }
1131
1132 async fn subscribe(&self, args: Vec<OKXSubscriptionArg>) -> Result<(), OKXWsError> {
1133 self.cmd_tx
1135 .read()
1136 .await
1137 .send(HandlerCommand::Subscribe { args: args.clone() })
1138 .map_err(|e| {
1139 OKXWsError::ClientError(format!("Failed to send subscribe command: {e}"))
1140 })?;
1141
1142 self.record_subscriptions(&args);
1143 Ok(())
1144 }
1145
1146 fn record_subscriptions(&self, args: &[OKXSubscriptionArg]) {
1147 for arg in args {
1148 let topic = topic_from_subscription_arg(arg);
1149 self.subscriptions_state.mark_subscribe(&topic);
1150
1151 if arg.inst_type.is_none() && arg.inst_family.is_none() && arg.inst_id.is_none() {
1153 self.subscriptions_bare.insert(arg.channel.clone(), true);
1154 } else {
1155 if let Some(inst_type) = &arg.inst_type {
1156 self.subscriptions_inst_type
1157 .entry(arg.channel.clone())
1158 .or_default()
1159 .insert(*inst_type);
1160 }
1161
1162 if let Some(inst_family) = &arg.inst_family {
1163 self.subscriptions_inst_family
1164 .entry(arg.channel.clone())
1165 .or_default()
1166 .insert(*inst_family);
1167 }
1168
1169 if let Some(inst_id) = &arg.inst_id {
1170 self.subscriptions_inst_id
1171 .entry(arg.channel.clone())
1172 .or_default()
1173 .insert(*inst_id);
1174 }
1175 }
1176 }
1177 }
1178
1179 #[expect(clippy::collapsible_if)]
1180 async fn unsubscribe(&self, args: Vec<OKXSubscriptionArg>) -> Result<(), OKXWsError> {
1181 self.cmd_tx
1183 .read()
1184 .await
1185 .send(HandlerCommand::Unsubscribe { args: args.clone() })
1186 .map_err(|e| {
1187 OKXWsError::ClientError(format!("Failed to send unsubscribe command: {e}"))
1188 })?;
1189
1190 for arg in &args {
1191 let topic = topic_from_subscription_arg(arg);
1192 self.subscriptions_state.mark_unsubscribe(&topic);
1193
1194 if arg.inst_type.is_none() && arg.inst_family.is_none() && arg.inst_id.is_none() {
1195 self.subscriptions_bare.remove(&arg.channel);
1196 } else {
1197 if let Some(inst_type) = &arg.inst_type {
1198 if let Some(mut entry) = self.subscriptions_inst_type.get_mut(&arg.channel) {
1199 entry.remove(inst_type);
1200 if entry.is_empty() {
1201 drop(entry);
1202 self.subscriptions_inst_type.remove(&arg.channel);
1203 }
1204 }
1205 }
1206
1207 if let Some(inst_family) = &arg.inst_family {
1208 if let Some(mut entry) = self.subscriptions_inst_family.get_mut(&arg.channel) {
1209 entry.remove(inst_family);
1210 if entry.is_empty() {
1211 drop(entry);
1212 self.subscriptions_inst_family.remove(&arg.channel);
1213 }
1214 }
1215 }
1216
1217 if let Some(inst_id) = &arg.inst_id {
1218 if let Some(mut entry) = self.subscriptions_inst_id.get_mut(&arg.channel) {
1219 entry.remove(inst_id);
1220 if entry.is_empty() {
1221 drop(entry);
1222 self.subscriptions_inst_id.remove(&arg.channel);
1223 }
1224 }
1225 }
1226 }
1227 }
1228
1229 Ok(())
1230 }
1231
1232 async fn subscribe_inst_id(
1233 &self,
1234 channel: OKXWsChannel,
1235 inst_id: Ustr,
1236 ) -> Result<(), OKXWsError> {
1237 self.subscribe(vec![OKXSubscriptionArg {
1238 channel,
1239 inst_type: None,
1240 inst_family: None,
1241 inst_id: Some(inst_id),
1242 }])
1243 .await
1244 }
1245
1246 async fn unsubscribe_inst_id(
1247 &self,
1248 channel: OKXWsChannel,
1249 inst_id: Ustr,
1250 ) -> Result<(), OKXWsError> {
1251 self.unsubscribe(vec![OKXSubscriptionArg {
1252 channel,
1253 inst_type: None,
1254 inst_family: None,
1255 inst_id: Some(inst_id),
1256 }])
1257 .await
1258 }
1259
1260 pub async fn unsubscribe_all(&self) -> Result<(), OKXWsError> {
1269 let all_args = subscription_args(
1270 &self.subscriptions_inst_type,
1271 &self.subscriptions_inst_family,
1272 &self.subscriptions_inst_id,
1273 &self.subscriptions_bare,
1274 );
1275
1276 if all_args.is_empty() {
1277 log::debug!("No active subscriptions to unsubscribe from");
1278 return Ok(());
1279 }
1280
1281 log::debug!("Batched unsubscribe from {} channels", all_args.len());
1282
1283 for chunk in subscription_arg_batches(&all_args) {
1284 self.unsubscribe(chunk.to_vec()).await?;
1285 }
1286
1287 self.index_pair_subscribers.clear();
1291
1292 Ok(())
1293 }
1294
1295 pub async fn subscribe_instruments(
1307 &self,
1308 instrument_type: OKXInstrumentType,
1309 ) -> Result<(), OKXWsError> {
1310 let arg = OKXSubscriptionArg {
1311 channel: OKXWsChannel::Instruments,
1312 inst_type: Some(instrument_type),
1313 inst_family: None,
1314 inst_id: None,
1315 };
1316 self.subscribe(vec![arg]).await
1317 }
1318
1319 pub async fn subscribe_instrument(
1333 &self,
1334 instrument_id: InstrumentId,
1335 ) -> Result<(), OKXWsError> {
1336 let inst_type = okx_instrument_type_from_symbol(instrument_id.symbol.as_str());
1337 log::debug!("Subscribing to instrument type {inst_type:?} for {instrument_id}");
1338 self.subscribe_instruments(inst_type).await
1339 }
1340
1341 pub async fn subscribe_book(&self, instrument_id: InstrumentId) -> anyhow::Result<()> {
1350 self.subscribe_book_with_depth(instrument_id, 0).await
1351 }
1352
1353 pub async fn subscribe_book_rpi(&self, instrument_id: InstrumentId) -> Result<(), OKXWsError> {
1359 self.subscribe_inst_id(OKXWsChannel::BooksRpi, instrument_id.symbol.inner())
1360 .await
1361 }
1362
1363 pub(crate) async fn subscribe_book_channel(
1364 &self,
1365 instrument_id: InstrumentId,
1366 channel: OKXBookChannel,
1367 cancel: CancellationToken,
1368 gate: SnapshotGate,
1369 ) -> Result<(), OKXWsError> {
1370 let (completion, receiver) = tokio::sync::oneshot::channel();
1371 let sender = self.cmd_tx.read().await;
1372
1373 if cancel.is_cancelled() {
1374 return Ok(());
1375 }
1376
1377 let subscription = OKXSubscriptionArg {
1378 channel: ws_channel_for_book(channel),
1379 inst_type: None,
1380 inst_family: None,
1381 inst_id: Some(instrument_id.symbol.inner()),
1382 };
1383
1384 self.record_subscriptions(std::slice::from_ref(&subscription));
1385 sender
1386 .send(HandlerCommand::SubscribeBook {
1387 subscription,
1388 cancel: cancel.clone(),
1389 gate,
1390 completion,
1391 })
1392 .map_err(|e| OKXWsError::HandlerUnavailable(e.to_string()))?;
1393
1394 drop(sender);
1395 tokio::select! {
1396 biased;
1397 () = cancel.cancelled() => Ok(()),
1398 result = receiver => result.map_err(|e| OKXWsError::HandlerUnavailable(e.to_string()))?,
1399 }
1400 }
1401
1402 pub(crate) async fn resubscribe_book_channel(
1404 &self,
1405 instrument_id: InstrumentId,
1406 channel: OKXBookChannel,
1407 cancel: CancellationToken,
1408 gate: SnapshotGate,
1409 ) -> Result<(), OKXWsError> {
1410 let (completion, receiver) = tokio::sync::oneshot::channel();
1411 let sender = self.cmd_tx.read().await;
1412
1413 if cancel.is_cancelled() {
1414 return Err(OKXWsError::ClientError("Book recovery canceled".into()));
1415 }
1416
1417 sender
1418 .send(HandlerCommand::Resubscribe {
1419 subscription: OKXSubscriptionArg {
1420 channel: ws_channel_for_book(channel),
1421 inst_type: None,
1422 inst_family: None,
1423 inst_id: Some(instrument_id.symbol.inner()),
1424 },
1425 cancel,
1426 gate,
1427 completion,
1428 })
1429 .map_err(|e| OKXWsError::HandlerUnavailable(e.to_string()))?;
1430
1431 drop(sender);
1432 receiver
1433 .await
1434 .map_err(|e| OKXWsError::HandlerUnavailable(e.to_string()))?
1435 }
1436
1437 pub async fn subscribe_book_depth5(
1449 &self,
1450 instrument_id: InstrumentId,
1451 ) -> Result<(), OKXWsError> {
1452 self.subscribe_inst_id(OKXWsChannel::Books5, instrument_id.symbol.inner())
1453 .await
1454 }
1455
1456 pub async fn subscribe_book50_l2_tbt(
1468 &self,
1469 instrument_id: InstrumentId,
1470 ) -> Result<(), OKXWsError> {
1471 self.subscribe_inst_id(OKXWsChannel::Books50Tbt, instrument_id.symbol.inner())
1472 .await
1473 }
1474
1475 pub async fn subscribe_book_l2_tbt(
1487 &self,
1488 instrument_id: InstrumentId,
1489 ) -> Result<(), OKXWsError> {
1490 self.subscribe_inst_id(OKXWsChannel::BooksTbt, instrument_id.symbol.inner())
1491 .await
1492 }
1493
1494 pub async fn subscribe_book_with_depth(
1508 &self,
1509 instrument_id: InstrumentId,
1510 depth: u16,
1511 ) -> anyhow::Result<()> {
1512 let vip = self.vip_level();
1513
1514 if !matches!(depth, 0 | 50 | 400) {
1515 anyhow::bail!("Invalid depth {depth}, must be 0, 50, or 400");
1516 }
1517
1518 if depth == 50 && vip < OKXVipLevel::Vip4 {
1519 anyhow::bail!("VIP level {vip} insufficient for 50 depth subscription (requires VIP4)");
1520 }
1521
1522 let channel = select_book_channel(depth as usize, vip);
1523 self.subscribe_inst_id(ws_channel_for_book(channel), instrument_id.symbol.inner())
1524 .await?;
1525 Ok(())
1526 }
1527
1528 pub async fn subscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), OKXWsError> {
1541 self.subscribe_inst_id(OKXWsChannel::BboTbt, instrument_id.symbol.inner())
1542 .await
1543 }
1544
1545 pub async fn subscribe_trades(
1559 &self,
1560 instrument_id: InstrumentId,
1561 aggregated: bool,
1562 ) -> Result<(), OKXWsError> {
1563 let channel = if aggregated {
1564 OKXWsChannel::TradesAll
1565 } else {
1566 OKXWsChannel::Trades
1567 };
1568 self.subscribe_inst_id(channel, instrument_id.symbol.inner())
1569 .await
1570 }
1571
1572 pub async fn subscribe_ticker(&self, instrument_id: InstrumentId) -> Result<(), OKXWsError> {
1584 self.subscribe_inst_id(OKXWsChannel::Tickers, instrument_id.symbol.inner())
1585 .await
1586 }
1587
1588 pub async fn subscribe_mark_prices(
1600 &self,
1601 instrument_id: InstrumentId,
1602 ) -> Result<(), OKXWsError> {
1603 self.subscribe_inst_id(OKXWsChannel::MarkPrice, instrument_id.symbol.inner())
1604 .await
1605 }
1606
1607 pub async fn subscribe_index_prices(
1619 &self,
1620 instrument_id: InstrumentId,
1621 ) -> Result<(), OKXWsError> {
1622 let symbol = instrument_id.symbol.inner();
1624 let (base, quote) = parse_base_quote_from_symbol(symbol.as_str())
1625 .map_err(|e| OKXWsError::ClientError(e.to_string()))?;
1626 let base_pair = Ustr::from(&format!("{base}-{quote}"));
1627
1628 let _guard = self.index_pair_transition.lock().await;
1634
1635 let is_first = {
1640 let mut count = self.index_pair_subscribers.entry(base_pair).or_insert(0);
1641 *count += 1;
1642 *count == 1
1643 };
1644
1645 if !is_first {
1646 return Ok(());
1647 }
1648
1649 let arg = OKXSubscriptionArg {
1650 channel: OKXWsChannel::IndexTickers,
1651 inst_type: None,
1652 inst_family: None,
1653 inst_id: Some(base_pair),
1654 };
1655
1656 match self.subscribe(vec![arg]).await {
1657 Ok(()) => Ok(()),
1658 Err(e) => {
1659 self.index_pair_subscribers.remove(&base_pair);
1668 Err(e)
1669 }
1670 }
1671 }
1672
1673 pub async fn subscribe_option_summary(&self, inst_family: Ustr) -> Result<(), OKXWsError> {
1686 let arg = OKXSubscriptionArg {
1687 channel: OKXWsChannel::OptionSummary,
1688 inst_type: None,
1689 inst_family: Some(inst_family),
1690 inst_id: None,
1691 };
1692 self.subscribe(vec![arg]).await
1693 }
1694
1695 pub async fn subscribe_event_contract_markets(&self) -> Result<(), OKXWsError> {
1705 let arg = OKXSubscriptionArg {
1706 channel: OKXWsChannel::EventContractMarkets,
1707 inst_type: Some(OKXInstrumentType::Events),
1708 inst_family: None,
1709 inst_id: None,
1710 };
1711 self.subscribe(vec![arg]).await
1712 }
1713
1714 pub fn option_greeks_subs(&self) -> &Arc<AtomicMap<InstrumentId, AHashSet<OKXGreeksType>>> {
1718 &self.option_greeks_subs
1719 }
1720
1721 pub fn add_option_greeks_sub(&self, instrument_id: InstrumentId) {
1724 let both: AHashSet<OKXGreeksType> =
1725 [OKXGreeksType::Bs, OKXGreeksType::Pa].into_iter().collect();
1726 self.option_greeks_subs.insert(instrument_id, both);
1727 }
1728
1729 pub fn add_option_greeks_sub_with_conventions(
1732 &self,
1733 instrument_id: InstrumentId,
1734 conventions: AHashSet<OKXGreeksType>,
1735 ) {
1736 let set = if conventions.is_empty() {
1737 [OKXGreeksType::Bs, OKXGreeksType::Pa].into_iter().collect()
1738 } else {
1739 conventions
1740 };
1741 self.option_greeks_subs.insert(instrument_id, set);
1742 }
1743
1744 pub fn remove_option_greeks_sub(&self, instrument_id: &InstrumentId) {
1746 self.option_greeks_subs.remove(instrument_id);
1747 }
1748
1749 pub async fn subscribe_funding_rates(
1761 &self,
1762 instrument_id: InstrumentId,
1763 ) -> Result<(), OKXWsError> {
1764 self.subscribe_inst_id(OKXWsChannel::FundingRate, instrument_id.symbol.inner())
1765 .await
1766 }
1767
1768 pub async fn subscribe_bars(&self, bar_type: BarType) -> Result<(), OKXWsError> {
1780 let channel = bar_spec_as_okx_channel(bar_type.spec())
1782 .map_err(|e| OKXWsError::ClientError(e.to_string()))?;
1783 self.subscribe_inst_id(channel, bar_type.instrument_id().symbol.inner())
1784 .await
1785 }
1786
1787 pub async fn unsubscribe_instruments(
1793 &self,
1794 instrument_type: OKXInstrumentType,
1795 ) -> Result<(), OKXWsError> {
1796 let arg = OKXSubscriptionArg {
1797 channel: OKXWsChannel::Instruments,
1798 inst_type: Some(instrument_type),
1799 inst_family: None,
1800 inst_id: None,
1801 };
1802 self.unsubscribe(vec![arg]).await
1803 }
1804
1805 #[allow(
1815 clippy::unused_async,
1816 clippy::unused_async_trait_impl,
1817 reason = "async signature is kept for parity with the other unsubscribe methods and callers"
1818 )]
1819 pub async fn unsubscribe_instrument(
1820 &self,
1821 instrument_id: InstrumentId,
1822 ) -> Result<(), OKXWsError> {
1823 log::debug!("Instrument unsubscribe is a no-op (shared per-type channel): {instrument_id}");
1824 Ok(())
1825 }
1826
1827 pub async fn unsubscribe_book(&self, instrument_id: InstrumentId) -> Result<(), OKXWsError> {
1833 self.unsubscribe_inst_id(OKXWsChannel::Books, instrument_id.symbol.inner())
1834 .await
1835 }
1836
1837 pub async fn unsubscribe_book_rpi(
1843 &self,
1844 instrument_id: InstrumentId,
1845 ) -> Result<(), OKXWsError> {
1846 self.unsubscribe_inst_id(OKXWsChannel::BooksRpi, instrument_id.symbol.inner())
1847 .await
1848 }
1849
1850 pub async fn unsubscribe_book_depth5(
1856 &self,
1857 instrument_id: InstrumentId,
1858 ) -> Result<(), OKXWsError> {
1859 self.unsubscribe_inst_id(OKXWsChannel::Books5, instrument_id.symbol.inner())
1860 .await
1861 }
1862
1863 pub async fn unsubscribe_book50_l2_tbt(
1869 &self,
1870 instrument_id: InstrumentId,
1871 ) -> Result<(), OKXWsError> {
1872 self.unsubscribe_inst_id(OKXWsChannel::Books50Tbt, instrument_id.symbol.inner())
1873 .await
1874 }
1875
1876 pub async fn unsubscribe_book_l2_tbt(
1882 &self,
1883 instrument_id: InstrumentId,
1884 ) -> Result<(), OKXWsError> {
1885 self.unsubscribe_inst_id(OKXWsChannel::BooksTbt, instrument_id.symbol.inner())
1886 .await
1887 }
1888
1889 pub async fn unsubscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), OKXWsError> {
1895 self.unsubscribe_inst_id(OKXWsChannel::BboTbt, instrument_id.symbol.inner())
1896 .await
1897 }
1898
1899 pub async fn unsubscribe_ticker(&self, instrument_id: InstrumentId) -> Result<(), OKXWsError> {
1905 self.unsubscribe_inst_id(OKXWsChannel::Tickers, instrument_id.symbol.inner())
1906 .await
1907 }
1908
1909 pub async fn unsubscribe_mark_prices(
1915 &self,
1916 instrument_id: InstrumentId,
1917 ) -> Result<(), OKXWsError> {
1918 self.unsubscribe_inst_id(OKXWsChannel::MarkPrice, instrument_id.symbol.inner())
1919 .await
1920 }
1921
1922 pub async fn unsubscribe_index_prices(
1935 &self,
1936 instrument_id: InstrumentId,
1937 ) -> Result<(), OKXWsError> {
1938 let symbol = instrument_id.symbol.inner();
1939 let (base, quote) = parse_base_quote_from_symbol(symbol.as_str())
1940 .map_err(|e| OKXWsError::ClientError(e.to_string()))?;
1941 let base_pair = Ustr::from(&format!("{base}-{quote}"));
1942
1943 let _guard = self.index_pair_transition.lock().await;
1946
1947 let is_last = {
1948 let Some(mut count) = self.index_pair_subscribers.get_mut(&base_pair) else {
1949 return Ok(());
1951 };
1952 *count = count.saturating_sub(1);
1953 *count == 0
1954 };
1955
1956 if !is_last {
1957 return Ok(());
1958 }
1959
1960 self.index_pair_subscribers
1961 .remove_if(&base_pair, |_, count| *count == 0);
1962
1963 let arg = OKXSubscriptionArg {
1964 channel: OKXWsChannel::IndexTickers,
1965 inst_type: None,
1966 inst_family: None,
1967 inst_id: Some(base_pair),
1968 };
1969 self.unsubscribe(vec![arg]).await
1970 }
1971
1972 pub async fn unsubscribe_option_summary(&self, inst_family: Ustr) -> Result<(), OKXWsError> {
1978 let arg = OKXSubscriptionArg {
1979 channel: OKXWsChannel::OptionSummary,
1980 inst_type: None,
1981 inst_family: Some(inst_family),
1982 inst_id: None,
1983 };
1984 self.unsubscribe(vec![arg]).await
1985 }
1986
1987 pub async fn unsubscribe_event_contract_markets(&self) -> Result<(), OKXWsError> {
1993 let arg = OKXSubscriptionArg {
1994 channel: OKXWsChannel::EventContractMarkets,
1995 inst_type: Some(OKXInstrumentType::Events),
1996 inst_family: None,
1997 inst_id: None,
1998 };
1999 self.unsubscribe(vec![arg]).await
2000 }
2001
2002 pub async fn unsubscribe_funding_rates(
2008 &self,
2009 instrument_id: InstrumentId,
2010 ) -> Result<(), OKXWsError> {
2011 self.unsubscribe_inst_id(OKXWsChannel::FundingRate, instrument_id.symbol.inner())
2012 .await
2013 }
2014
2015 pub async fn unsubscribe_trades(
2021 &self,
2022 instrument_id: InstrumentId,
2023 aggregated: bool,
2024 ) -> Result<(), OKXWsError> {
2025 let channel = if aggregated {
2026 OKXWsChannel::TradesAll
2027 } else {
2028 OKXWsChannel::Trades
2029 };
2030 self.unsubscribe_inst_id(channel, instrument_id.symbol.inner())
2031 .await
2032 }
2033
2034 pub async fn unsubscribe_bars(&self, bar_type: BarType) -> Result<(), OKXWsError> {
2040 let channel = bar_spec_as_okx_channel(bar_type.spec())
2041 .map_err(|e| OKXWsError::ClientError(e.to_string()))?;
2042 self.unsubscribe_inst_id(channel, bar_type.instrument_id().symbol.inner())
2043 .await
2044 }
2045
2046 pub async fn subscribe_orders(
2052 &self,
2053 instrument_type: OKXInstrumentType,
2054 ) -> Result<(), OKXWsError> {
2055 let arg = OKXSubscriptionArg {
2056 channel: OKXWsChannel::Orders,
2057 inst_type: Some(instrument_type),
2058 inst_family: None,
2059 inst_id: None,
2060 };
2061 self.subscribe(vec![arg]).await
2062 }
2063
2064 pub async fn unsubscribe_orders(
2070 &self,
2071 instrument_type: OKXInstrumentType,
2072 ) -> Result<(), OKXWsError> {
2073 let arg = OKXSubscriptionArg {
2074 channel: OKXWsChannel::Orders,
2075 inst_type: Some(instrument_type),
2076 inst_family: None,
2077 inst_id: None,
2078 };
2079 self.unsubscribe(vec![arg]).await
2080 }
2081
2082 pub async fn subscribe_spread_orders(&self) -> Result<(), OKXWsError> {
2088 let arg = OKXSubscriptionArg {
2089 channel: OKXWsChannel::SprdOrders,
2090 inst_type: None,
2091 inst_family: None,
2092 inst_id: None,
2093 };
2094 self.subscribe(vec![arg]).await
2095 }
2096
2097 pub async fn unsubscribe_spread_orders(&self) -> Result<(), OKXWsError> {
2103 let arg = OKXSubscriptionArg {
2104 channel: OKXWsChannel::SprdOrders,
2105 inst_type: None,
2106 inst_family: None,
2107 inst_id: None,
2108 };
2109 self.unsubscribe(vec![arg]).await
2110 }
2111
2112 pub async fn subscribe_spread_quotes(
2118 &self,
2119 instrument_id: InstrumentId,
2120 ) -> Result<(), OKXWsError> {
2121 self.subscribe_inst_id(OKXWsChannel::SprdBboTbt, instrument_id.symbol.inner())
2122 .await
2123 }
2124
2125 pub async fn subscribe_spread_book(
2131 &self,
2132 instrument_id: InstrumentId,
2133 ) -> Result<(), OKXWsError> {
2134 self.subscribe_inst_id(OKXWsChannel::SprdBooks5, instrument_id.symbol.inner())
2135 .await
2136 }
2137
2138 pub async fn subscribe_spread_trades(
2144 &self,
2145 instrument_id: InstrumentId,
2146 ) -> Result<(), OKXWsError> {
2147 self.subscribe_inst_id(OKXWsChannel::SprdPublicTrades, instrument_id.symbol.inner())
2148 .await
2149 }
2150
2151 pub async fn unsubscribe_spread_quotes(
2157 &self,
2158 instrument_id: InstrumentId,
2159 ) -> Result<(), OKXWsError> {
2160 self.unsubscribe_inst_id(OKXWsChannel::SprdBboTbt, instrument_id.symbol.inner())
2161 .await
2162 }
2163
2164 pub async fn unsubscribe_spread_book(
2170 &self,
2171 instrument_id: InstrumentId,
2172 ) -> Result<(), OKXWsError> {
2173 self.unsubscribe_inst_id(OKXWsChannel::SprdBooks5, instrument_id.symbol.inner())
2174 .await
2175 }
2176
2177 pub async fn unsubscribe_spread_trades(
2183 &self,
2184 instrument_id: InstrumentId,
2185 ) -> Result<(), OKXWsError> {
2186 self.unsubscribe_inst_id(OKXWsChannel::SprdPublicTrades, instrument_id.symbol.inner())
2187 .await
2188 }
2189
2190 pub async fn subscribe_orders_algo(
2196 &self,
2197 instrument_type: OKXInstrumentType,
2198 ) -> Result<(), OKXWsError> {
2199 let arg = OKXSubscriptionArg {
2200 channel: OKXWsChannel::OrdersAlgo,
2201 inst_type: Some(instrument_type),
2202 inst_family: None,
2203 inst_id: None,
2204 };
2205 self.subscribe(vec![arg]).await
2206 }
2207
2208 pub async fn unsubscribe_orders_algo(
2214 &self,
2215 instrument_type: OKXInstrumentType,
2216 ) -> Result<(), OKXWsError> {
2217 let arg = OKXSubscriptionArg {
2218 channel: OKXWsChannel::OrdersAlgo,
2219 inst_type: Some(instrument_type),
2220 inst_family: None,
2221 inst_id: None,
2222 };
2223 self.unsubscribe(vec![arg]).await
2224 }
2225
2226 pub async fn subscribe_algo_advance(
2232 &self,
2233 instrument_type: OKXInstrumentType,
2234 ) -> Result<(), OKXWsError> {
2235 let arg = OKXSubscriptionArg {
2236 channel: OKXWsChannel::AlgoAdvance,
2237 inst_type: Some(instrument_type),
2238 inst_family: None,
2239 inst_id: None,
2240 };
2241 self.subscribe(vec![arg]).await
2242 }
2243
2244 pub async fn unsubscribe_algo_advance(
2250 &self,
2251 instrument_type: OKXInstrumentType,
2252 ) -> Result<(), OKXWsError> {
2253 let arg = OKXSubscriptionArg {
2254 channel: OKXWsChannel::AlgoAdvance,
2255 inst_type: Some(instrument_type),
2256 inst_family: None,
2257 inst_id: None,
2258 };
2259 self.unsubscribe(vec![arg]).await
2260 }
2261
2262 pub async fn subscribe_account(&self) -> Result<(), OKXWsError> {
2268 let arg = OKXSubscriptionArg {
2269 channel: OKXWsChannel::Account,
2270 inst_type: None,
2271 inst_family: None,
2272 inst_id: None,
2273 };
2274 self.subscribe(vec![arg]).await
2275 }
2276
2277 pub async fn unsubscribe_account(&self) -> Result<(), OKXWsError> {
2283 let arg = OKXSubscriptionArg {
2284 channel: OKXWsChannel::Account,
2285 inst_type: None,
2286 inst_family: None,
2287 inst_id: None,
2288 };
2289 self.unsubscribe(vec![arg]).await
2290 }
2291
2292 pub async fn subscribe_positions(
2302 &self,
2303 inst_type: OKXInstrumentType,
2304 ) -> Result<(), OKXWsError> {
2305 let arg = OKXSubscriptionArg {
2306 channel: OKXWsChannel::Positions,
2307 inst_type: Some(inst_type),
2308 inst_family: None,
2309 inst_id: None,
2310 };
2311 self.subscribe(vec![arg]).await
2312 }
2313
2314 pub async fn unsubscribe_positions(
2320 &self,
2321 inst_type: OKXInstrumentType,
2322 ) -> Result<(), OKXWsError> {
2323 let arg = OKXSubscriptionArg {
2324 channel: OKXWsChannel::Positions,
2325 inst_type: Some(inst_type),
2326 inst_family: None,
2327 inst_id: None,
2328 };
2329 self.unsubscribe(vec![arg]).await
2330 }
2331
2332 pub async fn subscribe_liquidation_warning(
2342 &self,
2343 instrument_type: OKXInstrumentType,
2344 ) -> Result<(), OKXWsError> {
2345 let arg = OKXSubscriptionArg {
2346 channel: OKXWsChannel::LiquidationWarning,
2347 inst_type: Some(instrument_type),
2348 inst_family: None,
2349 inst_id: None,
2350 };
2351 self.subscribe(vec![arg]).await
2352 }
2353
2354 pub async fn unsubscribe_liquidation_warning(
2360 &self,
2361 instrument_type: OKXInstrumentType,
2362 ) -> Result<(), OKXWsError> {
2363 let arg = OKXSubscriptionArg {
2364 channel: OKXWsChannel::LiquidationWarning,
2365 inst_type: Some(instrument_type),
2366 inst_family: None,
2367 inst_id: None,
2368 };
2369 self.unsubscribe(vec![arg]).await
2370 }
2371
2372 async fn ws_batch_place_orders(
2378 &self,
2379 args: Vec<Value>,
2380 client_order_ids: Vec<ClientOrderId>,
2381 ) -> Result<(), OKXWsError> {
2382 let request_id = self.generate_unique_request_id();
2383 let request = OKXWsRequest::<Value> {
2384 id: Some(request_id.clone()),
2385 op: super::enums::OKXWsOperation::BatchOrders,
2386 exp_time: None,
2387 args,
2388 };
2389
2390 let payload = serde_json::to_string(&request)
2391 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize batch orders: {e}")))?;
2392
2393 let cmd = HandlerCommand::Send {
2394 payload,
2395 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_BATCH_ORDER.to_vec()),
2396 request_id: Some(request_id),
2397 client_order_ids,
2398 op: Some(super::enums::OKXWsOperation::BatchOrders),
2399 };
2400
2401 self.send_cmd(cmd).await
2402 }
2403
2404 async fn ws_batch_cancel_orders(
2410 &self,
2411 args: Vec<Value>,
2412 client_order_ids: Vec<ClientOrderId>,
2413 ) -> Result<(), OKXWsError> {
2414 let request_id = self.generate_unique_request_id();
2415 let request = OKXWsRequest::<Value> {
2416 id: Some(request_id.clone()),
2417 op: super::enums::OKXWsOperation::BatchCancelOrders,
2418 exp_time: None,
2419 args,
2420 };
2421
2422 let payload = serde_json::to_string(&request)
2423 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize batch cancel: {e}")))?;
2424
2425 let cmd = HandlerCommand::Send {
2426 payload,
2427 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_BATCH_CANCEL.to_vec()),
2428 request_id: Some(request_id),
2429 client_order_ids,
2430 op: Some(super::enums::OKXWsOperation::BatchCancelOrders),
2431 };
2432
2433 self.send_cmd(cmd).await
2434 }
2435
2436 async fn ws_batch_amend_orders(
2442 &self,
2443 args: Vec<Value>,
2444 client_order_ids: Vec<ClientOrderId>,
2445 ) -> Result<(), OKXWsError> {
2446 let request_id = self.generate_unique_request_id();
2447 let request = OKXWsRequest::<Value> {
2448 id: Some(request_id.clone()),
2449 op: super::enums::OKXWsOperation::BatchAmendOrders,
2450 exp_time: None,
2451 args,
2452 };
2453
2454 let payload = serde_json::to_string(&request)
2455 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize batch amend: {e}")))?;
2456
2457 let cmd = HandlerCommand::Send {
2458 payload,
2459 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_BATCH_AMEND.to_vec()),
2460 request_id: Some(request_id),
2461 client_order_ids,
2462 op: Some(super::enums::OKXWsOperation::BatchAmendOrders),
2463 };
2464
2465 self.send_cmd(cmd).await
2466 }
2467
2468 #[expect(clippy::too_many_arguments)]
2480 pub async fn submit_order(
2481 &self,
2482 trader_id: TraderId,
2483 strategy_id: StrategyId,
2484 instrument_id: InstrumentId,
2485 td_mode: OKXTradeMode,
2486 client_order_id: ClientOrderId,
2487 order_side: OrderSide,
2488 order_type: OrderType,
2489 quantity: Quantity,
2490 time_in_force: Option<TimeInForce>,
2491 price: Option<Price>,
2492 trigger_price: Option<Price>,
2493 post_only: Option<bool>,
2494 reduce_only: Option<bool>,
2495 quote_quantity: Option<bool>,
2496 position_side: Option<PositionSide>,
2497 attach_algo_ords: Option<Vec<WsAttachAlgoOrdParams>>,
2498 px_usd: Option<String>,
2499 px_vol: Option<String>,
2500 outcome: Option<String>,
2501 slippage_pct: Option<String>,
2502 rpi: Option<bool>,
2503 rpi_taker_access: Option<bool>,
2504 rpi_px_round: Option<bool>,
2505 ) -> Result<(), OKXWsError> {
2506 let rpi = rpi.unwrap_or(false);
2507
2508 if !OKX_SUPPORTED_ORDER_TYPES.contains(&order_type) {
2509 return Err(OKXWsError::ClientError(format!(
2510 "Unsupported order type: {order_type:?}",
2511 )));
2512 }
2513
2514 if let Some(tif) = time_in_force
2515 && !OKX_SUPPORTED_TIME_IN_FORCE.contains(&tif)
2516 {
2517 return Err(OKXWsError::ClientError(format!(
2518 "Unsupported time in force: {tif:?}",
2519 )));
2520 }
2521
2522 let mut builder = WsPostOrderParamsBuilder::default();
2523
2524 let inst_id_code = self
2525 .get_inst_id_code(&instrument_id.symbol.inner())
2526 .ok_or_else(|| {
2527 OKXWsError::ClientError(format!(
2528 "No instIdCode cached for {instrument_id}, cannot submit order"
2529 ))
2530 })?;
2531 builder.inst_id_code(inst_id_code);
2532
2533 builder.td_mode(td_mode);
2534 builder.cl_ord_id(client_order_id.as_str());
2535
2536 let (instrument_type, quote_currency) = {
2537 let instruments = self.instruments_cache.load();
2538 let symbol = instrument_id.symbol.inner();
2539 let instrument = instruments.get(&symbol).ok_or_else(|| {
2540 OKXWsError::ClientError(format!("Unknown instrument {instrument_id}"))
2541 })?;
2542 let instrument_type = okx_instrument_type(instrument)
2543 .map_err(|e| OKXWsError::ClientError(e.to_string()))?;
2544 (instrument_type, instrument.quote_currency())
2545 };
2546
2547 if instrument_type == OKXInstrumentType::Option
2549 && matches!(order_type, OrderType::Market | OrderType::MarketToLimit)
2550 {
2551 return Err(OKXWsError::ClientError(
2552 "Market orders are not supported for OKX options, use Limit orders instead"
2553 .to_string(),
2554 ));
2555 }
2556
2557 match instrument_type {
2558 OKXInstrumentType::Spot | OKXInstrumentType::Margin => {
2559 builder.ccy(quote_currency.to_string());
2561 }
2562 OKXInstrumentType::Swap | OKXInstrumentType::Futures => {
2563 builder.ccy(quote_currency.to_string());
2565
2566 if position_side.is_none() {
2569 builder.pos_side(OKXPositionSide::Net);
2570 }
2571 }
2572 OKXInstrumentType::Option => {
2573 builder.ccy(quote_currency.to_string());
2574
2575 if position_side.is_none() {
2576 builder.pos_side(OKXPositionSide::Net);
2577 }
2578 }
2579 OKXInstrumentType::Events => {}
2580 _ => {
2581 builder.ccy(quote_currency.to_string());
2582
2583 if position_side.is_none() {
2584 builder.pos_side(OKXPositionSide::Net);
2585 }
2586 }
2587 }
2588
2589 if let Some(reduce_only) = okx_reduce_only_wire_value(
2590 instrument_type,
2591 td_mode,
2592 order_side,
2593 position_side,
2594 reduce_only,
2595 )
2596 .map_err(OKXWsError::ClientError)?
2597 {
2598 builder.reduce_only(reduce_only);
2599 }
2600
2601 if let Some(attach_algo_ords) = attach_algo_ords {
2602 builder.attach_algo_ords(attach_algo_ords);
2603 }
2604
2605 if instrument_type == OKXInstrumentType::Spot
2612 && order_type == OrderType::Market
2613 && td_mode == OKXTradeMode::Cash
2614 {
2615 match quote_quantity {
2616 Some(true) => {
2617 builder.tgt_ccy(OKXTargetCurrency::QuoteCcy);
2618 }
2619 Some(false) if order_side == OrderSide::Buy => {
2621 builder.tgt_ccy(OKXTargetCurrency::BaseCcy);
2622 }
2623 Some(false) | None => {}
2625 }
2626 }
2627
2628 builder.side(order_side);
2629
2630 if let Some(pos_side) = position_side {
2631 builder.pos_side(pos_side);
2632 }
2633
2634 if rpi && order_type != OrderType::Limit {
2638 return Err(OKXWsError::ClientError(
2639 "OKX RPI orders require a limit order".to_string(),
2640 ));
2641 }
2642
2643 let (okx_ord_type, price) = if rpi {
2644 (OKXOrderType::Rpi, price)
2645 } else if post_only.unwrap_or(false) {
2646 (OKXOrderType::PostOnly, price)
2647 } else if let Some(tif) = time_in_force {
2648 match (order_type, tif) {
2649 (OrderType::Market, TimeInForce::Fok) => {
2650 return Err(OKXWsError::ClientError(
2651 "Market orders with FOK time-in-force are not supported by OKX. Use Limit order with FOK instead.".to_string()
2652 ));
2653 }
2654 (OrderType::Market, TimeInForce::Ioc) => {
2655 if matches!(
2657 instrument_type,
2658 OKXInstrumentType::Spot | OKXInstrumentType::Option
2659 ) {
2660 (OKXOrderType::Market, price)
2661 } else {
2662 (OKXOrderType::OptimalLimitIoc, price)
2663 }
2664 }
2665 (OrderType::Limit, TimeInForce::Fok) => {
2666 if instrument_type == OKXInstrumentType::Option {
2668 (OKXOrderType::OpFok, price)
2669 } else {
2670 (OKXOrderType::Fok, price)
2671 }
2672 }
2673 (OrderType::Limit, TimeInForce::Ioc) => (OKXOrderType::Ioc, price),
2674 _ => (OKXOrderType::from(order_type), price),
2675 }
2676 } else {
2677 (OKXOrderType::from(order_type), price)
2678 };
2679
2680 log::debug!(
2681 "Order type mapping: order_type={order_type:?}, time_in_force={time_in_force:?}, post_only={post_only:?} -> okx_ord_type={okx_ord_type:?}"
2682 );
2683
2684 if instrument_type == OKXInstrumentType::Events && outcome.is_none() {
2685 return Err(OKXWsError::ClientError(
2686 "OKX event contract orders require `outcome`".to_string(),
2687 ));
2688 }
2689
2690 if let Some(outcome) = outcome {
2691 builder.outcome(outcome);
2692 }
2693
2694 if let Some(slippage) = slippage_pct {
2695 builder.slippage_pct(slippage);
2696 }
2697
2698 if let Some(rpi_taker_access) = rpi_taker_access {
2699 builder.rpi_taker_access(rpi_taker_access);
2700 }
2701
2702 if let Some(rpi_px_round) = rpi_px_round {
2703 builder.rpi_px_round(rpi_px_round);
2704 }
2705
2706 if let Some(trade_quote_ccy) =
2707 self.resolve_spot_trade_quote_ccy(instrument_type, instrument_id.symbol.inner())?
2708 {
2709 builder.trade_quote_ccy(trade_quote_ccy);
2710 }
2711
2712 builder.ord_type(okx_ord_type);
2713 builder.sz(quantity.to_string());
2714
2715 if let Some(usd) = px_usd {
2717 builder.px_usd(usd);
2718 } else if let Some(vol) = px_vol {
2719 builder.px_vol(vol);
2720 } else if let Some(tp) = trigger_price {
2721 builder.px(tp.to_string());
2722 } else if let Some(p) = price {
2723 builder.px(p.to_string());
2724 }
2725
2726 builder.tag(OKX_NAUTILUS_BROKER_ID);
2727
2728 let params = builder
2729 .build()
2730 .map_err(|e| OKXWsError::ClientError(format!("Build order params error: {e}")))?;
2731
2732 let request_id = self.generate_unique_request_id();
2733 let request = OKXWsRequest {
2734 id: Some(request_id.clone()),
2735 op: super::enums::OKXWsOperation::Order,
2736 exp_time: None,
2737 args: vec![params],
2738 };
2739
2740 let payload = serde_json::to_string(&request)
2741 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize order: {e}")))?;
2742
2743 let cl_ord_key = client_order_id.to_string();
2744 self.pending_orders.insert(
2745 cl_ord_key.clone(),
2746 PendingOrderInfo {
2747 trader_id,
2748 strategy_id,
2749 instrument_id,
2750 },
2751 );
2752
2753 let cmd = HandlerCommand::Send {
2754 payload,
2755 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_ORDER.to_vec()),
2756 request_id: Some(request_id),
2757 client_order_ids: vec![client_order_id],
2758 op: Some(super::enums::OKXWsOperation::Order),
2759 };
2760
2761 let result = self.send_cmd(cmd).await;
2762
2763 if result.is_err() {
2764 self.pending_orders.remove(&cl_ord_key);
2765 }
2766
2767 result
2768 }
2769
2770 #[expect(clippy::too_many_arguments)]
2786 pub async fn modify_order(
2787 &self,
2788 trader_id: TraderId,
2789 strategy_id: StrategyId,
2790 instrument_id: InstrumentId,
2791 client_order_id: Option<ClientOrderId>,
2792 price: Option<Price>,
2793 quantity: Option<Quantity>,
2794 venue_order_id: Option<VenueOrderId>,
2795 new_px_usd: Option<String>,
2796 new_px_vol: Option<String>,
2797 rpi_taker_access: Option<bool>,
2798 rpi_px_round: Option<bool>,
2799 ) -> Result<(), OKXWsError> {
2800 let mut builder = WsAmendOrderParamsBuilder::default();
2801
2802 let inst_id_code = self
2803 .get_inst_id_code(&instrument_id.symbol.inner())
2804 .ok_or_else(|| {
2805 OKXWsError::ClientError(format!(
2806 "No instIdCode cached for {instrument_id}, cannot amend order"
2807 ))
2808 })?;
2809 builder.inst_id_code(inst_id_code);
2810
2811 if let Some(venue_order_id) = venue_order_id {
2812 builder.ord_id(venue_order_id.as_str());
2813 }
2814
2815 let cl_ord_key = client_order_id.map(|id| id.to_string());
2816
2817 if let Some(client_order_id) = client_order_id {
2818 builder.cl_ord_id(client_order_id.as_str());
2819 self.pending_amends.insert(
2820 client_order_id.to_string(),
2821 PendingOrderInfo {
2822 trader_id,
2823 strategy_id,
2824 instrument_id,
2825 },
2826 );
2827 }
2828
2829 if let Some(usd) = new_px_usd {
2831 builder.new_px_usd(usd);
2832 } else if let Some(vol) = new_px_vol {
2833 builder.new_px_vol(vol);
2834 } else if let Some(price) = price {
2835 builder.new_px(price.to_string());
2836 }
2837
2838 if let Some(quantity) = quantity {
2839 builder.new_sz(quantity.to_string());
2840 }
2841
2842 if let Some(rpi_taker_access) = rpi_taker_access {
2843 builder.rpi_taker_access(rpi_taker_access);
2844 }
2845
2846 if let Some(rpi_px_round) = rpi_px_round {
2847 builder.rpi_px_round(rpi_px_round);
2848 }
2849
2850 let params = builder
2851 .build()
2852 .map_err(|e| OKXWsError::ClientError(format!("Build amend params error: {e}")))?;
2853
2854 let request_id = self.generate_unique_request_id();
2855 let request = OKXWsRequest {
2856 id: Some(request_id.clone()),
2857 op: super::enums::OKXWsOperation::AmendOrder,
2858 exp_time: None,
2859 args: vec![params],
2860 };
2861
2862 let payload = serde_json::to_string(&request)
2863 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize amend: {e}")))?;
2864
2865 let cmd = HandlerCommand::Send {
2866 payload,
2867 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_AMEND.to_vec()),
2868 request_id: Some(request_id),
2869 client_order_ids: client_order_id.into_iter().collect(),
2870 op: Some(super::enums::OKXWsOperation::AmendOrder),
2871 };
2872
2873 let result = self.send_cmd(cmd).await;
2874
2875 if let (Err(_), Some(key)) = (&result, &cl_ord_key) {
2876 self.pending_amends.remove(key);
2877 }
2878
2879 result
2880 }
2881
2882 pub async fn cancel_order(
2893 &self,
2894 trader_id: TraderId,
2895 strategy_id: StrategyId,
2896 instrument_id: InstrumentId,
2897 client_order_id: Option<ClientOrderId>,
2898 venue_order_id: Option<VenueOrderId>,
2899 ) -> Result<(), OKXWsError> {
2900 let mut builder = WsCancelOrderParamsBuilder::default();
2901
2902 let inst_id_code = self
2903 .get_inst_id_code(&instrument_id.symbol.inner())
2904 .ok_or_else(|| {
2905 OKXWsError::ClientError(format!(
2906 "No instIdCode cached for {instrument_id}, cannot cancel order"
2907 ))
2908 })?;
2909 builder.inst_id_code(inst_id_code);
2910
2911 if let Some(venue_order_id) = venue_order_id {
2912 builder.ord_id(venue_order_id.as_str());
2913 }
2914
2915 let cl_ord_key = client_order_id.map(|id| id.to_string());
2916
2917 if let Some(client_order_id) = client_order_id {
2918 builder.cl_ord_id(client_order_id.as_str());
2919 self.pending_cancels.insert(
2920 client_order_id.to_string(),
2921 PendingOrderInfo {
2922 trader_id,
2923 strategy_id,
2924 instrument_id,
2925 },
2926 );
2927 }
2928
2929 let params = builder
2930 .build()
2931 .map_err(|e| OKXWsError::ClientError(format!("Build cancel params error: {e}")))?;
2932
2933 let request_id = self.generate_unique_request_id();
2934 let request = OKXWsRequest {
2935 id: Some(request_id.clone()),
2936 op: super::enums::OKXWsOperation::CancelOrder,
2937 exp_time: None,
2938 args: vec![params],
2939 };
2940
2941 let payload = serde_json::to_string(&request)
2942 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize cancel: {e}")))?;
2943
2944 let cmd = HandlerCommand::Send {
2945 payload,
2946 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_CANCEL.to_vec()),
2947 request_id: Some(request_id),
2948 client_order_ids: client_order_id.into_iter().collect(),
2949 op: Some(super::enums::OKXWsOperation::CancelOrder),
2950 };
2951
2952 let result = self.send_cmd(cmd).await;
2953
2954 if let (Err(_), Some(key)) = (&result, &cl_ord_key) {
2955 self.pending_cancels.remove(key);
2956 }
2957
2958 result
2959 }
2960
2961 pub async fn mass_cancel_orders(&self, instrument_id: InstrumentId) -> Result<(), OKXWsError> {
2971 let (inst_type, inst_family) = {
2972 let instrument = self
2973 .instruments_cache
2974 .get_cloned(&instrument_id.symbol.inner())
2975 .ok_or_else(|| {
2976 OKXWsError::ClientError(format!("Unknown instrument {instrument_id}"))
2977 })?;
2978
2979 let inst_type = okx_instrument_type(&instrument)
2980 .map_err(|e| OKXWsError::ClientError(e.to_string()))?;
2981
2982 let symbol = instrument.symbol().inner();
2983 let inst_family = match &instrument {
2984 InstrumentAny::CurrencyPair(_) => symbol.as_str().to_string(),
2985 InstrumentAny::CryptoPerpetual(_) => symbol
2986 .as_str()
2987 .strip_suffix("-SWAP")
2988 .unwrap_or(symbol.as_str())
2989 .to_string(),
2990 InstrumentAny::CryptoFuture(_) => {
2991 let s = symbol.as_str();
2992 if let Some(idx) = s.rfind('-') {
2993 s[..idx].to_string()
2994 } else {
2995 s.to_string()
2996 }
2997 }
2998 _ => {
2999 return Err(OKXWsError::ClientError(
3000 "Unsupported instrument type for mass cancel".to_string(),
3001 ));
3002 }
3003 };
3004
3005 (inst_type, inst_family)
3006 };
3007
3008 let params = WsMassCancelParams {
3009 inst_type,
3010 inst_family: Ustr::from(&inst_family),
3011 };
3012
3013 let request_id = self.generate_unique_request_id();
3014 let request = OKXWsRequest {
3015 id: Some(request_id.clone()),
3016 op: super::enums::OKXWsOperation::MassCancel,
3017 exp_time: None,
3018 args: vec![
3019 serde_json::to_value(params).map_err(|e| OKXWsError::JsonError(e.to_string()))?,
3020 ],
3021 };
3022
3023 let payload = serde_json::to_string(&request)
3024 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize mass cancel: {e}")))?;
3025
3026 let cmd = HandlerCommand::Send {
3027 payload,
3028 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_MASS_CANCEL.to_vec()),
3029 request_id: Some(request_id),
3030 client_order_ids: Vec::new(),
3031 op: Some(super::enums::OKXWsOperation::MassCancel),
3032 };
3033
3034 self.send_cmd(cmd).await
3035 }
3036
3037 #[expect(clippy::type_complexity)]
3044 pub async fn batch_submit_orders(
3045 &self,
3046 orders: Vec<(
3047 OKXInstrumentType,
3048 InstrumentId,
3049 OKXTradeMode,
3050 ClientOrderId,
3051 OrderSide,
3052 Option<PositionSide>,
3053 OrderType,
3054 Quantity,
3055 Option<Price>,
3056 Option<Price>,
3057 Option<bool>,
3058 Option<bool>,
3059 Option<String>,
3060 Option<bool>,
3061 Option<bool>,
3062 Option<bool>,
3063 )>,
3064 ) -> Result<(), OKXWsError> {
3065 let client_order_ids: Vec<ClientOrderId> = orders.iter().map(|o| o.3).collect();
3066 let args: Vec<Value> = {
3067 let mut args = Vec::with_capacity(orders.len());
3068 let inst_id_codes = self.inst_id_code_cache.load();
3069 let instruments = self.instruments_cache.load();
3070
3071 for (
3072 inst_type,
3073 inst_id,
3074 td_mode,
3075 cl_ord_id,
3076 ord_side,
3077 pos_side,
3078 ord_type,
3079 qty,
3080 pr,
3081 tp,
3082 post_only,
3083 reduce_only,
3084 outcome,
3085 rpi,
3086 rpi_taker_access,
3087 rpi_px_round,
3088 ) in orders
3089 {
3090 let rpi = rpi.unwrap_or(false);
3091 let mut builder = WsPostOrderParamsBuilder::default();
3092
3093 let (inst_id_symbol, inst_id_code) = Self::inst_id_symbol_and_code_from_snapshot(
3094 &inst_id_codes,
3095 &inst_id,
3096 "submit",
3097 )?;
3098 builder.inst_id_code(inst_id_code);
3099
3100 builder.td_mode(td_mode);
3101 builder.cl_ord_id(cl_ord_id.as_str());
3102 builder.side(ord_side);
3103
3104 if inst_type != OKXInstrumentType::Events
3105 && let Some(instrument) = instruments.get(&inst_id_symbol)
3106 {
3107 builder.ccy(instrument.quote_currency().to_string());
3108 }
3109
3110 if let Some(ps) = pos_side {
3111 builder.pos_side(OKXPositionSide::from(ps));
3112 } else if matches!(
3113 inst_type,
3114 OKXInstrumentType::Swap
3115 | OKXInstrumentType::Futures
3116 | OKXInstrumentType::Option
3117 ) {
3118 builder.pos_side(OKXPositionSide::Net);
3119 }
3120
3121 if rpi && ord_type != OrderType::Limit {
3122 return Err(OKXWsError::ClientError(
3123 "OKX RPI batch orders require limit orders".to_string(),
3124 ));
3125 }
3126
3127 let okx_ord_type = if rpi {
3128 OKXOrderType::Rpi
3129 } else if post_only.unwrap_or(false) {
3130 OKXOrderType::PostOnly
3131 } else {
3132 match ord_type {
3133 OrderType::Market => OKXOrderType::Market,
3134 OrderType::Limit => OKXOrderType::Limit,
3135 OrderType::MarketToLimit => OKXOrderType::Ioc,
3136 _ => {
3137 return Err(OKXWsError::ClientError(format!(
3138 "Unsupported order type for batch submit: {ord_type:?}"
3139 )));
3140 }
3141 }
3142 };
3143
3144 builder.ord_type(okx_ord_type);
3145 builder.sz(qty.to_string());
3146
3147 if let Some(p) = pr {
3148 builder.px(p.to_string());
3149 } else if let Some(p) = tp {
3150 builder.px(p.to_string());
3151 }
3152
3153 if let Some(reduce_only) =
3154 okx_reduce_only_wire_value(inst_type, td_mode, ord_side, pos_side, reduce_only)
3155 .map_err(OKXWsError::ClientError)?
3156 {
3157 builder.reduce_only(reduce_only);
3158 }
3159
3160 if inst_type == OKXInstrumentType::Events && outcome.is_none() {
3161 return Err(OKXWsError::ClientError(
3162 "OKX event contract orders require `outcome`".to_string(),
3163 ));
3164 }
3165
3166 if let Some(outcome) = outcome {
3167 builder.outcome(outcome);
3168 }
3169
3170 if let Some(rpi_taker_access) = rpi_taker_access {
3171 builder.rpi_taker_access(rpi_taker_access);
3172 }
3173
3174 if let Some(rpi_px_round) = rpi_px_round {
3175 builder.rpi_px_round(rpi_px_round);
3176 }
3177
3178 if let Some(trade_quote_ccy) =
3179 self.resolve_spot_trade_quote_ccy(inst_type, inst_id_symbol)?
3180 {
3181 builder.trade_quote_ccy(trade_quote_ccy);
3182 }
3183
3184 builder.tag(OKX_NAUTILUS_BROKER_ID);
3185
3186 let params = builder.build().map_err(|e| {
3187 OKXWsError::ClientError(format!("Build order params error: {e}"))
3188 })?;
3189 let val = serde_json::to_value(params)
3190 .map_err(|e| OKXWsError::JsonError(e.to_string()))?;
3191 args.push(val);
3192 }
3193 args
3194 };
3195
3196 self.ws_batch_place_orders(args, client_order_ids).await
3197 }
3198
3199 #[expect(clippy::type_complexity)]
3206 pub async fn batch_modify_orders(
3207 &self,
3208 orders: Vec<(
3209 OKXInstrumentType,
3210 InstrumentId,
3211 ClientOrderId,
3212 Option<String>,
3213 Option<Price>,
3214 Option<Quantity>,
3215 Option<bool>,
3216 Option<bool>,
3217 )>,
3218 ) -> Result<(), OKXWsError> {
3219 let client_order_ids: Vec<ClientOrderId> = orders.iter().map(|o| o.2).collect();
3220 let args: Vec<Value> = {
3221 let mut args = Vec::with_capacity(orders.len());
3222 let inst_id_codes = self.inst_id_code_cache.load();
3223
3224 for (
3225 _inst_type,
3226 inst_id,
3227 cl_ord_id,
3228 request_id,
3229 pr,
3230 sz,
3231 rpi_taker_access,
3232 rpi_px_round,
3233 ) in orders
3234 {
3235 let mut builder = WsAmendOrderParamsBuilder::default();
3236
3237 let (_, inst_id_code) =
3238 Self::inst_id_symbol_and_code_from_snapshot(&inst_id_codes, &inst_id, "amend")?;
3239 builder.inst_id_code(inst_id_code);
3240
3241 builder.cl_ord_id(cl_ord_id.as_str());
3242
3243 if let Some(request_id) = request_id {
3244 builder.req_id(request_id);
3245 }
3246
3247 if let Some(p) = pr {
3248 builder.new_px(p.to_string());
3249 }
3250
3251 if let Some(q) = sz {
3252 builder.new_sz(q.to_string());
3253 }
3254
3255 if let Some(rpi_taker_access) = rpi_taker_access {
3256 builder.rpi_taker_access(rpi_taker_access);
3257 }
3258
3259 if let Some(rpi_px_round) = rpi_px_round {
3260 builder.rpi_px_round(rpi_px_round);
3261 }
3262
3263 let params = builder.build().map_err(|e| {
3264 OKXWsError::ClientError(format!("Build amend batch params error: {e}"))
3265 })?;
3266 let val = serde_json::to_value(params)
3267 .map_err(|e| OKXWsError::JsonError(e.to_string()))?;
3268 args.push(val);
3269 }
3270 args
3271 };
3272
3273 self.ws_batch_amend_orders(args, client_order_ids).await
3274 }
3275
3276 pub async fn batch_cancel_orders(
3289 &self,
3290 orders: Vec<(InstrumentId, Option<ClientOrderId>, Option<VenueOrderId>)>,
3291 ) -> Result<(), OKXWsError> {
3292 let client_order_ids: Vec<ClientOrderId> = orders
3293 .iter()
3294 .filter_map(|(_, cl_ord_id, _)| *cl_ord_id)
3295 .collect();
3296 let args: Vec<Value> = {
3297 let mut args = Vec::with_capacity(orders.len());
3298 let inst_id_codes = self.inst_id_code_cache.load();
3299
3300 for (inst_id, cl_ord_id, ord_id) in orders {
3301 let mut builder = WsCancelOrderParamsBuilder::default();
3302
3303 let (_, inst_id_code) = Self::inst_id_symbol_and_code_from_snapshot(
3304 &inst_id_codes,
3305 &inst_id,
3306 "cancel",
3307 )?;
3308 builder.inst_id_code(inst_id_code);
3309
3310 if let Some(c) = cl_ord_id {
3311 builder.cl_ord_id(c.as_str());
3312 }
3313
3314 if let Some(o) = ord_id {
3315 builder.ord_id(o.as_str());
3316 }
3317
3318 let params = builder.build().map_err(|e| {
3319 OKXWsError::ClientError(format!("Build cancel batch params error: {e}"))
3320 })?;
3321 let val = serde_json::to_value(params)
3322 .map_err(|e| OKXWsError::JsonError(e.to_string()))?;
3323 args.push(val);
3324 }
3325 args
3326 };
3327
3328 self.ws_batch_cancel_orders(args, client_order_ids).await
3329 }
3330
3331 #[expect(clippy::too_many_arguments)]
3342 pub async fn submit_algo_order(
3343 &self,
3344 _trader_id: TraderId,
3345 _strategy_id: StrategyId,
3346 instrument_id: InstrumentId,
3347 td_mode: OKXTradeMode,
3348 client_order_id: ClientOrderId,
3349 order_side: OrderSide,
3350 order_type: OrderType,
3351 quantity: Quantity,
3352 trigger_price: Option<Price>,
3353 trigger_type: Option<TriggerType>,
3354 limit_price: Option<Price>,
3355 reduce_only: Option<bool>,
3356 callback_ratio: Option<String>,
3357 callback_spread: Option<String>,
3358 activation_price: Option<Price>,
3359 ) -> Result<(), OKXWsError> {
3360 if !is_conditional_order(order_type) {
3361 return Err(OKXWsError::ClientError(format!(
3362 "Order type {order_type:?} is not a conditional order"
3363 )));
3364 }
3365
3366 let mut builder = WsPostAlgoOrderParamsBuilder::default();
3367
3368 if !matches!(order_side, OrderSide::Buy | OrderSide::Sell) {
3369 return Err(OKXWsError::ClientError(
3370 "Invalid order side for OKX".to_string(),
3371 ));
3372 }
3373
3374 let inst_id_code = self
3375 .get_inst_id_code(&instrument_id.symbol.inner())
3376 .ok_or_else(|| {
3377 OKXWsError::ClientError(format!(
3378 "No instIdCode cached for {instrument_id}, cannot submit algo order"
3379 ))
3380 })?;
3381 builder.inst_id_code(inst_id_code);
3382
3383 builder.td_mode(td_mode);
3384 builder.cl_ord_id(client_order_id.as_str());
3385 builder.side(order_side);
3386 builder.ord_type(
3387 conditional_order_to_algo_type(order_type)
3388 .map_err(|e| OKXWsError::ClientError(e.to_string()))?,
3389 );
3390 builder.sz(quantity.to_string());
3391
3392 if let Some(tp) = trigger_price {
3393 builder.trigger_px(tp.to_string());
3394 }
3395
3396 let okx_trigger_type = trigger_type.map_or(OKXTriggerType::Last, Into::into);
3398 builder.trigger_px_type(okx_trigger_type);
3399
3400 if matches!(order_type, OrderType::StopLimit | OrderType::LimitIfTouched)
3402 && let Some(price) = limit_price
3403 {
3404 builder.order_px(price.to_string());
3405 }
3406
3407 if let Some(reduce) = reduce_only {
3408 builder.reduce_only(reduce);
3409 }
3410
3411 if let Some(ratio) = callback_ratio {
3412 builder.callback_ratio(ratio);
3413 }
3414
3415 if let Some(spread) = callback_spread {
3416 builder.callback_spread(spread);
3417 }
3418
3419 if let Some(active) = activation_price {
3420 builder.active_px(active.to_string());
3421 }
3422
3423 builder.tag(OKX_NAUTILUS_BROKER_ID);
3424
3425 let params = builder
3426 .build()
3427 .map_err(|e| OKXWsError::ClientError(format!("Build algo order params error: {e}")))?;
3428
3429 let request_id = self.generate_unique_request_id();
3430 let request = OKXWsRequest {
3431 id: Some(request_id.clone()),
3432 op: super::enums::OKXWsOperation::OrderAlgo,
3433 exp_time: None,
3434 args: vec![params],
3435 };
3436
3437 let payload = serde_json::to_string(&request)
3438 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize algo order: {e}")))?;
3439
3440 let cmd = HandlerCommand::Send {
3441 payload,
3442 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_ALGO_ORDER.to_vec()),
3443 request_id: Some(request_id),
3444 client_order_ids: vec![client_order_id],
3445 op: Some(super::enums::OKXWsOperation::OrderAlgo),
3446 };
3447
3448 self.send_cmd(cmd).await
3449 }
3450
3451 pub async fn cancel_algo_order(
3462 &self,
3463 _trader_id: TraderId,
3464 _strategy_id: StrategyId,
3465 instrument_id: InstrumentId,
3466 client_order_id: Option<ClientOrderId>,
3467 algo_order_id: Option<String>,
3468 ) -> Result<(), OKXWsError> {
3469 let mut builder = super::messages::WsCancelAlgoOrderParamsBuilder::default();
3470
3471 let inst_id_code = self
3472 .get_inst_id_code(&instrument_id.symbol.inner())
3473 .ok_or_else(|| {
3474 OKXWsError::ClientError(format!(
3475 "No instIdCode cached for {instrument_id}, cannot cancel algo order"
3476 ))
3477 })?;
3478 builder.inst_id_code(inst_id_code);
3479
3480 if let Some(algo_id) = algo_order_id {
3481 builder.algo_id(algo_id);
3482 }
3483
3484 if let Some(cl_ord_id) = client_order_id {
3485 builder.algo_cl_ord_id(cl_ord_id.to_string());
3486 }
3487
3488 let params = builder
3489 .build()
3490 .map_err(|e| OKXWsError::ClientError(format!("Build cancel algo params error: {e}")))?;
3491
3492 let request_id = self.generate_unique_request_id();
3493 let request = OKXWsRequest {
3494 id: Some(request_id.clone()),
3495 op: super::enums::OKXWsOperation::CancelAlgos,
3496 exp_time: None,
3497 args: vec![params],
3498 };
3499
3500 let payload = serde_json::to_string(&request)
3501 .map_err(|e| OKXWsError::JsonError(format!("Failed to serialize cancel algo: {e}")))?;
3502
3503 let cmd = HandlerCommand::Send {
3504 payload,
3505 rate_limit_keys: Some(OKX_RATE_LIMIT_KEY_ALGO_CANCEL.to_vec()),
3506 request_id: Some(request_id),
3507 client_order_ids: client_order_id.into_iter().collect(),
3508 op: Some(super::enums::OKXWsOperation::CancelAlgos),
3509 };
3510
3511 self.send_cmd(cmd).await
3512 }
3513
3514 async fn send_cmd(&self, cmd: HandlerCommand) -> Result<(), OKXWsError> {
3516 self.cmd_tx
3517 .read()
3518 .await
3519 .send(cmd)
3520 .map_err(|e| OKXWsError::HandlerUnavailable(e.to_string()))
3521 }
3522}
3523
3524fn authentication_timestamp(now: UnixNanos) -> String {
3525 now.as_seconds().to_string()
3526}
3527
3528fn subscription_args(
3529 subscriptions_inst_type: &DashMap<OKXWsChannel, AHashSet<OKXInstrumentType>>,
3530 subscriptions_inst_family: &DashMap<OKXWsChannel, AHashSet<Ustr>>,
3531 subscriptions_inst_id: &DashMap<OKXWsChannel, AHashSet<Ustr>>,
3532 subscriptions_bare: &DashMap<OKXWsChannel, bool>,
3533) -> Vec<OKXSubscriptionArg> {
3534 let mut args = Vec::new();
3535
3536 for entry in subscriptions_inst_type {
3537 let (channel, inst_types) = entry.pair();
3538 for inst_type in inst_types {
3539 args.push(OKXSubscriptionArg {
3540 channel: channel.clone(),
3541 inst_type: Some(*inst_type),
3542 inst_family: None,
3543 inst_id: None,
3544 });
3545 }
3546 }
3547
3548 for entry in subscriptions_inst_family {
3549 let (channel, inst_families) = entry.pair();
3550 for inst_family in inst_families {
3551 args.push(OKXSubscriptionArg {
3552 channel: channel.clone(),
3553 inst_type: None,
3554 inst_family: Some(*inst_family),
3555 inst_id: None,
3556 });
3557 }
3558 }
3559
3560 for entry in subscriptions_inst_id {
3561 let (channel, inst_ids) = entry.pair();
3562 for inst_id in inst_ids {
3563 args.push(OKXSubscriptionArg {
3564 channel: channel.clone(),
3565 inst_type: None,
3566 inst_family: None,
3567 inst_id: Some(*inst_id),
3568 });
3569 }
3570 }
3571
3572 for entry in subscriptions_bare {
3573 args.push(OKXSubscriptionArg {
3574 channel: entry.key().clone(),
3575 inst_type: None,
3576 inst_family: None,
3577 inst_id: None,
3578 });
3579 }
3580
3581 args.sort_unstable_by_key(topic_from_subscription_arg);
3582 args
3583}
3584
3585fn subscription_arg_batches(
3586 args: &[OKXSubscriptionArg],
3587) -> impl Iterator<Item = &[OKXSubscriptionArg]> {
3588 args.chunks(OKX_WS_SUBSCRIPTION_ARGS_MAX_PER_MESSAGE)
3589}
3590
3591fn auth_attempt_superseded(error: &Error) -> bool {
3592 error.to_string().contains("superseded")
3593}
3594
3595fn reconnect_auth_retry_delay(attempt: u32) -> Duration {
3596 let shift = attempt.saturating_sub(1).min(31);
3597 let factor = 1u32.checked_shl(shift).unwrap_or(u32::MAX);
3598
3599 RECONNECT_AUTH_RETRY_INITIAL
3600 .saturating_mul(factor)
3601 .min(RECONNECT_AUTH_RETRY_MAX)
3602}
3603
3604async fn authenticate_session(
3605 credential: &Credential,
3606 auth_tracker: &AuthTracker,
3607 cmd_tx: &tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>,
3608 clock: &'static AtomicTime,
3609 auth_timeout_secs: u64,
3610) -> Result<(), Error> {
3611 let rx = auth_tracker.begin();
3612 let timestamp = authentication_timestamp(clock.get_time_ns());
3613 let signature = credential.sign(×tamp, "GET", "/users/self/verify", "");
3614
3615 let auth_message = OKXAuthentication {
3616 op: "login",
3617 args: vec![OKXAuthenticationArg {
3618 api_key: SecretString::from(credential.api_key()),
3619 passphrase: SecretString::from(credential.api_passphrase()),
3620 timestamp,
3621 sign: SecretString::from(signature),
3622 }],
3623 };
3624
3625 let payload = serde_json::to_string(&auth_message)
3626 .map(SecretString::from)
3627 .map_err(|e| {
3628 Error::Io(std::io::Error::other(format!(
3629 "Failed to serialize auth message: {e}"
3630 )))
3631 })?;
3632
3633 cmd_tx
3634 .read()
3635 .await
3636 .send(HandlerCommand::Authenticate { payload })
3637 .map_err(|e| {
3638 Error::Io(std::io::Error::other(format!(
3639 "Failed to send authenticate command: {e}"
3640 )))
3641 })?;
3642
3643 match auth_tracker
3644 .wait_for_result::<OKXWsError>(Duration::from_secs(auth_timeout_secs), rx)
3645 .await
3646 {
3647 Ok(()) => {
3648 log::debug!("WebSocket authenticated");
3649 Ok(())
3650 }
3651 Err(e) => {
3652 let auth_error = Error::Io(std::io::Error::other(e.to_string()));
3653 if !auth_attempt_superseded(&auth_error) {
3654 log::warn!("WebSocket authentication failed: error={e}");
3655 }
3656
3657 Err(auth_error)
3658 }
3659 }
3660}
3661
3662#[allow(
3663 clippy::too_many_arguments,
3664 reason = "retry loop needs the session pieces without cloning the websocket client"
3665)]
3666async fn retry_reconnect_authentication(
3667 credential: Credential,
3668 auth_tracker: AuthTracker,
3669 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
3670 clock: &'static AtomicTime,
3671 auth_timeout_secs: u64,
3672 signal: Arc<AtomicBool>,
3673 abort: CancellationToken,
3674 generation_flag: Arc<AtomicU64>,
3675 generation: u64,
3676) {
3677 let mut attempt = 0u32;
3678
3679 loop {
3680 if signal.load(Ordering::Acquire) || abort.is_cancelled() {
3681 break;
3682 }
3683
3684 if generation_flag.load(Ordering::Acquire) != generation {
3685 break;
3686 }
3687
3688 if cmd_tx.read().await.is_closed() {
3689 break;
3690 }
3691
3692 if auth_tracker.is_authenticated() {
3693 break;
3694 }
3695
3696 attempt = attempt.saturating_add(1);
3697
3698 match authenticate_session(
3699 &credential,
3700 &auth_tracker,
3701 &cmd_tx,
3702 clock,
3703 auth_timeout_secs,
3704 )
3705 .await
3706 {
3707 Ok(()) => {
3708 log::debug!("Re-authenticated after reconnection");
3709 break;
3710 }
3711 Err(e) => {
3712 if generation_flag.load(Ordering::Acquire) != generation {
3713 break;
3714 }
3715
3716 if auth_tracker.is_authenticated() {
3717 break;
3718 }
3719
3720 let delay = reconnect_auth_retry_delay(attempt);
3721
3722 log::warn!(
3723 "Re-authentication after reconnection failed, retrying in {delay:?}: {e}"
3724 );
3725
3726 tokio::select! {
3727 biased;
3728 () = abort.cancelled() => break,
3729 () = time::sleep(delay) => {}
3730 }
3731 }
3732 }
3733 }
3734}
3735
3736pub(crate) fn ws_channel_for_book(channel: OKXBookChannel) -> OKXWsChannel {
3737 match channel {
3738 OKXBookChannel::Book => OKXWsChannel::Books,
3739 OKXBookChannel::BookL2Tbt => OKXWsChannel::BooksTbt,
3740 OKXBookChannel::Books50L2Tbt => OKXWsChannel::Books50Tbt,
3741 OKXBookChannel::BooksRpi => OKXWsChannel::BooksRpi,
3742 OKXBookChannel::SprdBooks5 => OKXWsChannel::SprdBooks5,
3743 }
3744}
3745
3746fn log_receiver_dropped(signal: &AtomicBool, item: &str) {
3747 if signal.load(Ordering::Acquire) {
3748 log::debug!("Receiver dropped after stop signal while forwarding {item}");
3749 } else {
3750 log::error!("Failed to send {item} through channel: receiver dropped");
3751 }
3752}
3753
3754#[cfg(test)]
3755mod tests {
3756 use nautilus_core::time::get_atomic_clock_realtime;
3757 use nautilus_live::{SocketReconnectRegistry, SocketReconnectRequestOutcome};
3758 use nautilus_model::{identifiers::ClientId, instruments::stubs::crypto_perpetual_ethusdt};
3759 use nautilus_network::RECONNECTED;
3760 use rstest::rstest;
3761 use tokio_tungstenite::tungstenite::Message;
3762
3763 use super::*;
3764 use crate::{
3765 common::{
3766 consts::{OKX_POST_ONLY_CANCEL_SOURCE, OKX_VENUE},
3767 enums::{
3768 OKXExecType, OKXOrderCategory, OKXOrderStatus, OKXPriceType, OKXQuickMarginType,
3769 OKXSelfTradePreventionMode, OKXSide,
3770 },
3771 },
3772 websocket::{
3773 handler::is_post_only_auto_cancel,
3774 messages::{OKXOrderMsg, OKXWebSocketError, OKXWsFrame},
3775 },
3776 };
3777
3778 struct DropSignal(Option<tokio::sync::oneshot::Sender<()>>);
3779
3780 impl Drop for DropSignal {
3781 fn drop(&mut self) {
3782 if let Some(sender) = self.0.take() {
3783 let _ = sender.send(());
3784 }
3785 }
3786 }
3787
3788 struct BlockingDrop(Arc<(parking_lot::Mutex<bool>, parking_lot::Condvar)>);
3789
3790 impl Drop for BlockingDrop {
3791 fn drop(&mut self) {
3792 let (lock, condvar) = &*self.0;
3793 let mut released = lock.lock();
3794 condvar.wait_while(&mut released, |released| !*released);
3795 }
3796 }
3797
3798 #[rstest]
3799 #[case(OKXBookChannel::Book, OKXWsChannel::Books)]
3800 #[case(OKXBookChannel::BookL2Tbt, OKXWsChannel::BooksTbt)]
3801 #[case(OKXBookChannel::Books50L2Tbt, OKXWsChannel::Books50Tbt)]
3802 #[case(OKXBookChannel::BooksRpi, OKXWsChannel::BooksRpi)]
3803 #[case(OKXBookChannel::SprdBooks5, OKXWsChannel::SprdBooks5)]
3804 fn test_ws_channel_for_book(#[case] channel: OKXBookChannel, #[case] expected: OKXWsChannel) {
3805 assert_eq!(ws_channel_for_book(channel), expected);
3806 }
3807
3808 #[rstest]
3809 fn test_timestamp_format_for_websocket_auth() {
3810 let now = UnixNanos::new(1_700_000_000_999_999_999);
3811
3812 assert_eq!(authentication_timestamp(now), "1700000000");
3813 }
3814
3815 #[rstest]
3816 fn test_subscription_args_are_sorted_by_topic() {
3817 let client = OKXWebSocketClient::default();
3818 client
3819 .subscriptions_inst_type
3820 .entry(OKXWsChannel::Instruments)
3821 .or_default()
3822 .extend([OKXInstrumentType::Swap, OKXInstrumentType::Spot]);
3823 client
3824 .subscriptions_inst_family
3825 .entry(OKXWsChannel::OpenInterest)
3826 .or_default()
3827 .insert(Ustr::from("BTC-USD"));
3828 client
3829 .subscriptions_inst_id
3830 .entry(OKXWsChannel::Tickers)
3831 .or_default()
3832 .extend([Ustr::from("ETH-USDT"), Ustr::from("BTC-USDT")]);
3833 client.subscriptions_bare.insert(OKXWsChannel::Status, true);
3834
3835 let topics = subscription_args(
3836 &client.subscriptions_inst_type,
3837 &client.subscriptions_inst_family,
3838 &client.subscriptions_inst_id,
3839 &client.subscriptions_bare,
3840 )
3841 .iter()
3842 .map(topic_from_subscription_arg)
3843 .collect::<Vec<_>>();
3844
3845 assert_eq!(
3846 topics,
3847 [
3848 "Instruments:Spot",
3849 "Instruments:Swap",
3850 "OpenInterest:BTC-USD",
3851 "Status",
3852 "Tickers:BTC-USDT",
3853 "Tickers:ETH-USDT",
3854 ]
3855 );
3856 }
3857
3858 #[rstest]
3859 fn test_new_without_credentials() {
3860 let client = OKXWebSocketClient::default();
3861 assert!(client.credential.is_none());
3862 assert_eq!(client.api_key(), None);
3863 }
3864
3865 #[rstest]
3866 fn test_instruments_cache_arc_observes_post_clone_writes() {
3867 let client = OKXWebSocketClient::default();
3868 let cache = client.instruments_cache_arc();
3869 assert!(cache.load().is_empty());
3870
3871 let instrument = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
3872 let symbol = instrument.symbol().inner();
3873 client.cache_instruments(std::slice::from_ref(&instrument));
3874
3875 let loaded = cache.load();
3876 assert_eq!(loaded.len(), 1);
3877 let stored = loaded.get(&symbol).expect("instrument not refreshed");
3878 assert_eq!(stored.id(), instrument.id());
3879 }
3880
3881 #[rstest]
3882 fn test_add_option_greeks_sub_defaults_to_both_conventions() {
3883 let client = OKXWebSocketClient::default();
3884 let instrument_id = InstrumentId::from("BTC-USD-250328-92000-C.OKX");
3885
3886 client.add_option_greeks_sub(instrument_id);
3887
3888 let subs = client.option_greeks_subs().load();
3889 let stored = subs.get(&instrument_id).expect("instrument not registered");
3890 assert_eq!(stored.len(), 2);
3891 assert!(stored.contains(&OKXGreeksType::Bs));
3892 assert!(stored.contains(&OKXGreeksType::Pa));
3893 }
3894
3895 #[rstest]
3896 #[case::bs_only(vec![OKXGreeksType::Bs])]
3897 #[case::pa_only(vec![OKXGreeksType::Pa])]
3898 #[case::both(vec![OKXGreeksType::Bs, OKXGreeksType::Pa])]
3899 fn test_add_option_greeks_sub_with_conventions_stores_requested_set(
3900 #[case] conventions: Vec<OKXGreeksType>,
3901 ) {
3902 let client = OKXWebSocketClient::default();
3903 let instrument_id = InstrumentId::from("BTC-USD-250328-92000-C.OKX");
3904 let set: AHashSet<OKXGreeksType> = conventions.iter().copied().collect();
3905
3906 client.add_option_greeks_sub_with_conventions(instrument_id, set.clone());
3907
3908 let subs = client.option_greeks_subs().load();
3909 let stored = subs.get(&instrument_id).expect("instrument not registered");
3910 assert_eq!(stored, &set);
3911 }
3912
3913 #[rstest]
3914 fn test_add_option_greeks_sub_with_empty_conventions_falls_back_to_both() {
3915 let client = OKXWebSocketClient::default();
3916 let instrument_id = InstrumentId::from("BTC-USD-250328-92000-C.OKX");
3917
3918 client.add_option_greeks_sub_with_conventions(instrument_id, AHashSet::new());
3919
3920 let subs = client.option_greeks_subs().load();
3921 let stored = subs.get(&instrument_id).expect("instrument not registered");
3922 assert_eq!(stored.len(), 2);
3923 }
3924
3925 #[rstest]
3926 fn test_remove_option_greeks_sub_clears_entry() {
3927 let client = OKXWebSocketClient::default();
3928 let instrument_id = InstrumentId::from("BTC-USD-250328-92000-C.OKX");
3929
3930 client.add_option_greeks_sub(instrument_id);
3931 client.remove_option_greeks_sub(&instrument_id);
3932
3933 let subs = client.option_greeks_subs().load();
3934 assert!(!subs.contains_key(&instrument_id));
3935 }
3936
3937 #[rstest]
3938 fn test_new_with_credentials() {
3939 let client = OKXWebSocketClient::new(
3940 None,
3941 Some("test_key".to_string()),
3942 Some("test_secret".to_string()),
3943 Some("test_passphrase".to_string()),
3944 None,
3945 None,
3946 None,
3947 TransportBackend::default(),
3948 None,
3949 )
3950 .unwrap();
3951 assert!(client.credential.is_some());
3952 assert_eq!(client.api_key(), Some("test_key"));
3953 }
3954
3955 #[rstest]
3956 fn test_new_partial_credentials_fails() {
3957 let result = OKXWebSocketClient::new(
3958 None,
3959 Some("test_key".to_string()),
3960 None,
3961 Some("test_passphrase".to_string()),
3962 None,
3963 None,
3964 None,
3965 TransportBackend::default(),
3966 None,
3967 );
3968 result.unwrap_err();
3969 }
3970
3971 #[rstest]
3972 fn test_request_id_generation() {
3973 let client = OKXWebSocketClient::default();
3974
3975 let initial_counter = client.request_id_counter.load(Ordering::SeqCst);
3976
3977 let id1 = client.request_id_counter.fetch_add(1, Ordering::SeqCst);
3978 let id2 = client.request_id_counter.fetch_add(1, Ordering::SeqCst);
3979
3980 assert_eq!(id1, initial_counter);
3981 assert_eq!(id2, initial_counter + 1);
3982 assert_eq!(
3983 client.request_id_counter.load(Ordering::SeqCst),
3984 initial_counter + 2
3985 );
3986 }
3987
3988 #[rstest]
3989 fn test_client_state_management() {
3990 let client = OKXWebSocketClient::default();
3991
3992 assert!(client.is_closed());
3993 assert!(!client.is_active());
3994
3995 let client_with_heartbeat = OKXWebSocketClient::new(
3996 None,
3997 None,
3998 None,
3999 None,
4000 None,
4001 Some(30),
4002 None,
4003 TransportBackend::default(),
4004 None,
4005 )
4006 .unwrap();
4007
4008 assert!(client_with_heartbeat.heartbeat.is_some());
4009 assert_eq!(client_with_heartbeat.heartbeat.unwrap(), 30);
4010 }
4011
4012 #[rstest]
4013 #[tokio::test]
4014 async fn begin_shutdown_stops_handler_before_bounded_close() {
4015 let client_id = ClientId::from("OKX-TEST");
4016 let endpoint = Ustr::from("okx-test-stream");
4017 let registry = SocketReconnectRegistry::default();
4018 let control =
4019 SocketControl::with_registry(client_id, Some(*OKX_VENUE), endpoint, ®istry);
4020 let _sink = control.sink();
4021 control.register(|| SocketReconnectRequestOutcome::Accepted);
4022 let mut client = OKXWebSocketClient::default().with_socket_control(control);
4023 client
4024 .connection_mode
4025 .load()
4026 .store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
4027 let (drop_tx, drop_rx) = tokio::sync::oneshot::channel();
4028 let signal = DropSignal(Some(drop_tx));
4029 let handler_abort = CancellationToken::new();
4030 *client.handler_abort.lock() = handler_abort.clone();
4031 client
4032 .handler_tasks
4033 .spawn(async move {
4034 let _signal = signal;
4035 handler_abort.cancelled().await;
4036 })
4037 .expect("handler task should register");
4038
4039 assert!(registry.handle(client_id, endpoint).is_some());
4040 client.begin_shutdown();
4041
4042 tokio::time::timeout(Duration::from_secs(1), drop_rx)
4043 .await
4044 .expect("begin shutdown must drop the handler task")
4045 .expect("drop signal");
4046 assert!(client.is_closed());
4047 assert!(!client.handler_tasks.is_open());
4048 assert!(registry.handle(client_id, endpoint).is_some());
4049
4050 client.close().await.expect("bounded close");
4051 assert!(!client.has_task());
4052 assert!(registry.handle(client_id, endpoint).is_none());
4053 }
4054
4055 #[rstest]
4056 #[tokio::test]
4057 async fn connect_rollback_closes_handler_admission_and_deregisters_socket() {
4058 let client_id = ClientId::from("OKX-CONNECT-ROLLBACK");
4059 let endpoint = Ustr::from("okx-connect-rollback");
4060 let registry = SocketReconnectRegistry::default();
4061 let control =
4062 SocketControl::with_registry(client_id, Some(*OKX_VENUE), endpoint, ®istry);
4063 control.register(|| SocketReconnectRequestOutcome::Accepted);
4064 let handler_tasks = Arc::new(TaskGroup::new());
4065 let signal = Arc::new(AtomicBool::new(false));
4066 let handler_abort = CancellationToken::new();
4067
4068 let rollback = ConnectRollback {
4069 handler_tasks: Arc::clone(&handler_tasks),
4070 signal: Arc::clone(&signal),
4071 handler_abort: handler_abort.clone(),
4072 socket_control: Some(Arc::new(control)),
4073 armed: true,
4074 };
4075
4076 drop(rollback);
4077
4078 assert!(!handler_tasks.is_open());
4079 assert!(signal.load(Ordering::Acquire));
4080 assert!(handler_abort.is_cancelled());
4081 assert!(registry.handle(client_id, endpoint).is_none());
4082 handler_tasks
4083 .finish_shutdown(Duration::ZERO, Duration::from_secs(1))
4084 .await
4085 .expect("empty handler scope should drain");
4086 }
4087
4088 #[rstest]
4089 #[tokio::test]
4090 async fn request_close_signals_before_handler_shutdown() {
4091 let mut client = OKXWebSocketClient::default();
4092 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
4093 client.cmd_tx = Arc::new(tokio::sync::RwLock::new(cmd_tx));
4094 client.signal.store(false, Ordering::Release);
4095
4096 client.request_close().await;
4097
4098 assert!(client.signal.load(Ordering::Acquire));
4099 assert!(matches!(cmd_rx.try_recv(), Ok(HandlerCommand::Disconnect)));
4100 }
4101
4102 #[rstest]
4103 #[tokio::test]
4104 async fn close_joins_handler_shared_with_clone() {
4105 let mut client = OKXWebSocketClient::default();
4106 let (drop_tx, drop_rx) = tokio::sync::oneshot::channel();
4107 let signal = DropSignal(Some(drop_tx));
4108 client
4109 .handler_tasks
4110 .spawn(async move {
4111 let _signal = signal;
4112 std::future::pending::<()>().await;
4113 })
4114 .expect("handler task should register");
4115 let retained = client.clone();
4116
4117 client.close().await.expect("close with retained clone");
4118
4119 tokio::time::timeout(Duration::from_secs(1), drop_rx)
4120 .await
4121 .expect("close must drop the handler task")
4122 .expect("drop signal");
4123 assert!(!client.has_task());
4124 assert!(!retained.has_task());
4125 }
4126
4127 #[rstest]
4128 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4129 async fn timeout_retains_unfinished_handler_task() {
4130 let mut client = OKXWebSocketClient::default();
4131 let release = Arc::new((parking_lot::Mutex::new(false), parking_lot::Condvar::new()));
4132 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
4133 let blocking_drop = BlockingDrop(Arc::clone(&release));
4134 client
4135 .handler_tasks
4136 .spawn(async move {
4137 let _blocking_drop = blocking_drop;
4138 started_tx.send(()).expect("started receiver");
4139 std::future::pending::<()>().await;
4140 })
4141 .expect("handler task should register");
4142 started_rx.await.expect("blocking task started");
4143 client.begin_shutdown();
4144
4145 let result = client.close_stream_task(Duration::from_millis(10)).await;
4146 let retained = client.has_task();
4147 let reconnect_result = client.connect().await;
4148
4149 let (lock, condvar) = &*release;
4150 *lock.lock() = true;
4151 condvar.notify_all();
4152
4153 client
4154 .close_stream_task(Duration::from_secs(1))
4155 .await
4156 .expect("blocking handler task terminated");
4157
4158 assert!(result.is_err());
4159 assert!(retained);
4160 assert_eq!(
4161 reconnect_result
4162 .expect_err("reconnect with unfinished handler")
4163 .to_string(),
4164 "Cannot connect while previous WebSocket handler task is still running"
4165 );
4166 assert!(!client.has_task());
4167 }
4168
4169 #[rstest]
4170 fn test_websocket_error_handling() {
4171 let clock = get_atomic_clock_realtime();
4172 let ts = clock.get_time_ns().as_u64();
4173
4174 let error = OKXWebSocketError {
4175 code: "60012".to_string(),
4176 message: "Invalid request".to_string(),
4177 conn_id: None,
4178 timestamp: ts,
4179 };
4180
4181 assert_eq!(error.code, "60012");
4182 assert_eq!(error.message, "Invalid request");
4183 assert_eq!(error.timestamp, ts);
4184
4185 let nautilus_msg = OKXWsMessage::Error(error);
4186 match nautilus_msg {
4187 OKXWsMessage::Error(e) => {
4188 assert_eq!(e.code, "60012");
4189 assert_eq!(e.message, "Invalid request");
4190 }
4191 _ => panic!("Expected Error variant"),
4192 }
4193 }
4194
4195 #[rstest]
4196 fn test_request_id_generation_sequence() {
4197 let client = OKXWebSocketClient::default();
4198
4199 let initial_counter = client
4200 .request_id_counter
4201 .load(std::sync::atomic::Ordering::SeqCst);
4202 let mut ids = Vec::new();
4203
4204 for _ in 0..10 {
4205 let id = client
4206 .request_id_counter
4207 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
4208 ids.push(id);
4209 }
4210
4211 for (i, &id) in ids.iter().enumerate() {
4212 assert_eq!(id, initial_counter + i as u64);
4213 }
4214
4215 assert_eq!(
4216 client
4217 .request_id_counter
4218 .load(std::sync::atomic::Ordering::SeqCst),
4219 initial_counter + 10
4220 );
4221 }
4222
4223 #[rstest]
4224 fn test_client_state_transitions() {
4225 let client = OKXWebSocketClient::default();
4226
4227 assert!(client.is_closed());
4228 assert!(!client.is_active());
4229
4230 let client_with_heartbeat = OKXWebSocketClient::new(
4231 None,
4232 None,
4233 None,
4234 None,
4235 None,
4236 Some(30), None,
4238 TransportBackend::default(),
4239 None,
4240 )
4241 .unwrap();
4242
4243 assert!(client_with_heartbeat.heartbeat.is_some());
4244 assert_eq!(client_with_heartbeat.heartbeat.unwrap(), 30);
4245 }
4246
4247 #[rstest]
4248 fn test_websocket_error_scenarios() {
4249 let clock = get_atomic_clock_realtime();
4250 let ts = clock.get_time_ns().as_u64();
4251
4252 let error_scenarios = vec![
4253 ("60012", "Invalid request", None),
4254 ("60009", "Invalid API key", Some("conn-123".to_string())),
4255 ("60014", "Too many requests", None),
4256 ("50001", "Order not found", None),
4257 ];
4258
4259 for (code, message, conn_id) in error_scenarios {
4260 let error = OKXWebSocketError {
4261 code: code.to_string(),
4262 message: message.to_string(),
4263 conn_id: conn_id.clone(),
4264 timestamp: ts,
4265 };
4266
4267 assert_eq!(error.code, code);
4268 assert_eq!(error.message, message);
4269 assert_eq!(error.conn_id, conn_id);
4270 assert_eq!(error.timestamp, ts);
4271
4272 let nautilus_msg = OKXWsMessage::Error(error);
4273 match nautilus_msg {
4274 OKXWsMessage::Error(e) => {
4275 assert_eq!(e.code, code);
4276 assert_eq!(e.message, message);
4277 assert_eq!(e.conn_id, conn_id);
4278 }
4279 _ => panic!("Expected Error variant"),
4280 }
4281 }
4282 }
4283
4284 #[rstest]
4285 fn test_feed_handler_reconnection_detection() {
4286 let msg = Message::Text(RECONNECTED.to_string().into());
4287 let result = OKXWsFeedHandler::parse_raw_message(msg);
4288 assert!(matches!(result, Some(OKXWsFrame::Reconnected)));
4289 }
4290
4291 #[rstest]
4292 fn test_feed_handler_normal_message_processing() {
4293 let ping_msg = Message::Text(TEXT_PING.to_string().into());
4294 let result = OKXWsFeedHandler::parse_raw_message(ping_msg);
4295 assert!(matches!(result, Some(OKXWsFrame::Ping)));
4296
4297 let sub_msg = r#"{
4298 "event": "subscribe",
4299 "arg": {
4300 "channel": "tickers",
4301 "instType": "SPOT"
4302 },
4303 "connId": "a4d3ae55"
4304 }"#;
4305
4306 let sub_result =
4307 OKXWsFeedHandler::parse_raw_message(Message::Text(sub_msg.to_string().into()));
4308 assert!(matches!(sub_result, Some(OKXWsFrame::Subscription { .. })));
4309 }
4310
4311 #[rstest]
4312 fn test_feed_handler_close_message() {
4313 let result = OKXWsFeedHandler::parse_raw_message(Message::Close(None));
4314 assert!(result.is_none());
4315 }
4316
4317 #[rstest]
4318 fn test_reconnection_message_constant() {
4319 assert_eq!(RECONNECTED, "__RECONNECTED__");
4320 }
4321
4322 #[rstest]
4323 fn reconnect_auth_retry_delay_grows_and_caps() {
4324 assert_eq!(reconnect_auth_retry_delay(1), Duration::from_secs(1));
4325 assert_eq!(reconnect_auth_retry_delay(2), Duration::from_secs(2));
4326 assert_eq!(reconnect_auth_retry_delay(3), Duration::from_secs(4));
4327 assert_eq!(reconnect_auth_retry_delay(6), Duration::from_secs(30));
4328 assert_eq!(
4329 reconnect_auth_retry_delay(u32::MAX),
4330 Duration::from_secs(30)
4331 );
4332 }
4333
4334 #[rstest]
4335 fn resubscribe_batches_subscription_args_into_venue_message_size() {
4336 let args: Vec<_> = (0..300)
4337 .map(|i| OKXSubscriptionArg {
4338 channel: OKXWsChannel::Tickers,
4339 inst_type: None,
4340 inst_family: None,
4341 inst_id: Some(Ustr::from(&format!("INST-{i}"))),
4342 })
4343 .collect();
4344
4345 let batches: Vec<_> = subscription_arg_batches(&args).collect();
4346
4347 assert_eq!(OKX_WS_SUBSCRIPTION_ARGS_MAX_PER_MESSAGE, 256);
4348 assert_eq!(batches.len(), 2);
4349 assert_eq!(batches[0].len(), 256);
4350 assert_eq!(batches[1].len(), 44);
4351 }
4352
4353 #[rstest]
4354 #[tokio::test]
4355 async fn authenticate_session_sends_login_and_completes_on_success() {
4356 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
4357 let cmd_tx = tokio::sync::RwLock::new(tx);
4358 let auth_tracker = AuthTracker::new();
4359 let credential = Credential::new("key".into(), "secret".into(), "pass".into());
4360 let clock = get_atomic_clock_realtime();
4361 let tracker = auth_tracker.clone();
4362
4363 let task = tokio::spawn(async move {
4364 authenticate_session(&credential, &auth_tracker, &cmd_tx, clock, 5).await
4365 });
4366
4367 match rx.recv().await.expect("authenticate command") {
4368 HandlerCommand::Authenticate { .. } => {}
4369 other => panic!("Expected HandlerCommand::Authenticate, was {other:?}"),
4370 }
4371
4372 tracker.succeed();
4373 task.await
4374 .expect("authenticate task")
4375 .expect("authentication should succeed");
4376 }
4377
4378 #[rstest]
4379 fn auth_attempt_superseded_matches_tracker_message() {
4380 let error = Error::Io(std::io::Error::other("Authentication attempt superseded"));
4381 assert!(auth_attempt_superseded(&error));
4382 let other = Error::Io(std::io::Error::other("Authentication timed out"));
4383 assert!(!auth_attempt_superseded(&other));
4384 }
4385
4386 #[rstest]
4387 #[tokio::test]
4388 async fn retry_reconnect_authentication_stops_when_already_authenticated() {
4389 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
4390 let cmd_tx = Arc::new(tokio::sync::RwLock::new(tx));
4391 let auth_tracker = AuthTracker::new();
4392 auth_tracker.succeed();
4393
4394 tokio::time::timeout(
4395 Duration::from_millis(200),
4396 retry_reconnect_authentication(
4397 Credential::new("key".into(), "secret".into(), "pass".into()),
4398 auth_tracker,
4399 cmd_tx,
4400 get_atomic_clock_realtime(),
4401 5,
4402 Arc::new(AtomicBool::new(false)),
4403 CancellationToken::new(),
4404 Arc::new(AtomicU64::new(1)),
4405 1,
4406 ),
4407 )
4408 .await
4409 .expect("authenticated retry must return without sending login");
4410 assert!(
4411 rx.try_recv().is_err(),
4412 "already-authenticated retry must not send a login command"
4413 );
4414 }
4415
4416 #[rstest]
4417 #[tokio::test]
4418 async fn retry_reconnect_authentication_stops_when_aborted() {
4419 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
4420 let cmd_tx = Arc::new(tokio::sync::RwLock::new(tx));
4421 let abort = CancellationToken::new();
4422 abort.cancel();
4423
4424 tokio::time::timeout(
4425 Duration::from_millis(200),
4426 retry_reconnect_authentication(
4427 Credential::new("key".into(), "secret".into(), "pass".into()),
4428 AuthTracker::new(),
4429 cmd_tx,
4430 get_atomic_clock_realtime(),
4431 5,
4432 Arc::new(AtomicBool::new(false)),
4433 abort,
4434 Arc::new(AtomicU64::new(1)),
4435 1,
4436 ),
4437 )
4438 .await
4439 .expect("aborted retry must return without waiting for authentication");
4440 }
4441
4442 #[rstest]
4443 fn test_multiple_reconnection_signals() {
4444 for _ in 0..3 {
4445 let msg = Message::Text(RECONNECTED.to_string().into());
4446 let result = OKXWsFeedHandler::parse_raw_message(msg);
4447 assert!(matches!(result, Some(OKXWsFrame::Reconnected)));
4448 }
4449 }
4450
4451 #[tokio::test]
4452 async fn test_wait_until_active_timeout() {
4453 let client = OKXWebSocketClient::new(
4454 None,
4455 Some("test_key".to_string()),
4456 Some("test_secret".to_string()),
4457 Some("test_passphrase".to_string()),
4458 Some(AccountId::from("test-account")),
4459 None,
4460 None,
4461 TransportBackend::default(),
4462 None,
4463 )
4464 .unwrap();
4465
4466 let result = client.wait_until_active(0.1).await;
4467
4468 assert!(result.is_err());
4469 assert!(!client.is_active());
4470 }
4471
4472 fn sample_canceled_order_msg() -> OKXOrderMsg {
4473 OKXOrderMsg {
4474 acc_fill_sz: Some("0".to_string()),
4475 avg_px: "0".to_string(),
4476 c_time: 0,
4477 cancel_source: None,
4478 cancel_source_reason: None,
4479 category: OKXOrderCategory::Normal,
4480 ccy: Ustr::from("USDT"),
4481 cl_ord_id: "order-1".to_string(),
4482 algo_cl_ord_id: None,
4483 attach_algo_cl_ord_id: None,
4484 attach_algo_ords: Vec::new(),
4485 outcome: None,
4486 fee: None,
4487 fee_ccy: Ustr::from("USDT"),
4488 fill_px: "0".to_string(),
4489 fill_sz: "0".to_string(),
4490 fill_time: 0,
4491 inst_id: Ustr::from("ETH-USDT-SWAP"),
4492 inst_type: OKXInstrumentType::Swap,
4493 lever: "1".to_string(),
4494 ord_id: Ustr::from("123456"),
4495 ord_type: OKXOrderType::Limit,
4496 pnl: "0".to_string(),
4497 pos_side: OKXPositionSide::Net,
4498 px: "0".to_string(),
4499 reduce_only: "false".to_string(),
4500 side: OKXSide::Buy,
4501 state: OKXOrderStatus::Canceled,
4502 exec_type: OKXExecType::None,
4503 sz: "1".to_string(),
4504 td_mode: OKXTradeMode::Cross,
4505 tgt_ccy: None,
4506 trade_id: String::new(),
4507 algo_id: None,
4508 fill_fee: None,
4509 fill_fee_ccy: None,
4510 fill_mark_px: None,
4511 fill_mark_vol: None,
4512 fill_px_vol: None,
4513 fill_px_usd: None,
4514 fill_fwd_px: None,
4515 fill_notional_usd: None,
4516 fill_pnl: None,
4517 is_tp_limit: None,
4518 linked_algo_ord: None,
4519 notional_usd: None,
4520 px_type: OKXPriceType::None,
4521 px_usd: None,
4522 px_vol: None,
4523 quick_mgn_type: OKXQuickMarginType::None,
4524 rebate: None,
4525 rebate_ccy: None,
4526 sl_ord_px: None,
4527 sl_trigger_px: None,
4528 sl_trigger_px_type: None,
4529 source: None,
4530 stp_id: None,
4531 stp_mode: OKXSelfTradePreventionMode::None,
4532 tag: None,
4533 tp_ord_px: None,
4534 tp_trigger_px: None,
4535 tp_trigger_px_type: None,
4536 amend_result: None,
4537 req_id: None,
4538 code: None,
4539 msg: None,
4540 u_time: 0,
4541 }
4542 }
4543
4544 #[rstest]
4545 fn test_is_post_only_auto_cancel_detects_cancel_source() {
4546 let mut msg = sample_canceled_order_msg();
4547 msg.cancel_source = Some(OKX_POST_ONLY_CANCEL_SOURCE.to_string());
4548
4549 assert!(is_post_only_auto_cancel(&msg));
4550 }
4551
4552 #[rstest]
4553 fn test_is_post_only_auto_cancel_detects_reason() {
4554 let mut msg = sample_canceled_order_msg();
4555 msg.cancel_source_reason = Some("POST_ONLY would take liquidity".to_string());
4556
4557 assert!(is_post_only_auto_cancel(&msg));
4558 }
4559
4560 #[rstest]
4561 fn test_is_post_only_auto_cancel_false_without_markers() {
4562 let msg = sample_canceled_order_msg();
4563
4564 assert!(!is_post_only_auto_cancel(&msg));
4565 }
4566
4567 #[rstest]
4568 fn test_is_post_only_auto_cancel_false_for_order_type_only() {
4569 let mut msg = sample_canceled_order_msg();
4570 msg.ord_type = OKXOrderType::PostOnly;
4571
4572 assert!(!is_post_only_auto_cancel(&msg));
4573 }
4574
4575 #[tokio::test]
4576 async fn test_batch_cancel_orders_with_multiple_orders() {
4577 use nautilus_model::identifiers::{ClientOrderId, InstrumentId, VenueOrderId};
4578
4579 let client = OKXWebSocketClient::new(
4580 Some("wss://test.okx.com".to_string()),
4581 None,
4582 None,
4583 None,
4584 None,
4585 None,
4586 None,
4587 TransportBackend::default(),
4588 None,
4589 )
4590 .expect("Failed to create client");
4591
4592 let instrument_id = InstrumentId::from("BTC-USDT.OKX");
4593 let client_order_id1 = ClientOrderId::new("order1");
4594 let client_order_id2 = ClientOrderId::new("order2");
4595 let venue_order_id1 = VenueOrderId::new("venue1");
4596 let venue_order_id2 = VenueOrderId::new("venue2");
4597
4598 let orders = vec![
4599 (instrument_id, Some(client_order_id1), Some(venue_order_id1)),
4600 (instrument_id, Some(client_order_id2), Some(venue_order_id2)),
4601 ];
4602
4603 let result = client.batch_cancel_orders(orders).await;
4604 assert!(result.is_err());
4605 }
4606
4607 #[tokio::test]
4608 async fn test_batch_cancel_orders_with_only_client_order_id() {
4609 use nautilus_model::identifiers::{ClientOrderId, InstrumentId};
4610
4611 let client = OKXWebSocketClient::new(
4612 Some("wss://test.okx.com".to_string()),
4613 None,
4614 None,
4615 None,
4616 None,
4617 None,
4618 None,
4619 TransportBackend::default(),
4620 None,
4621 )
4622 .expect("Failed to create client");
4623
4624 let instrument_id = InstrumentId::from("BTC-USDT.OKX");
4625 let client_order_id = ClientOrderId::new("order1");
4626
4627 let orders = vec![(instrument_id, Some(client_order_id), None)];
4628
4629 let result = client.batch_cancel_orders(orders).await;
4630
4631 assert!(result.is_err());
4632 }
4633
4634 #[tokio::test]
4635 async fn test_batch_cancel_orders_with_only_venue_order_id() {
4636 use nautilus_model::identifiers::{InstrumentId, VenueOrderId};
4637
4638 let client = OKXWebSocketClient::new(
4639 Some("wss://test.okx.com".to_string()),
4640 None,
4641 None,
4642 None,
4643 None,
4644 None,
4645 None,
4646 TransportBackend::default(),
4647 None,
4648 )
4649 .expect("Failed to create client");
4650
4651 let instrument_id = InstrumentId::from("BTC-USDT.OKX");
4652 let venue_order_id = VenueOrderId::new("venue1");
4653
4654 let orders = vec![(instrument_id, None, Some(venue_order_id))];
4655
4656 let result = client.batch_cancel_orders(orders).await;
4657
4658 assert!(result.is_err());
4659 }
4660
4661 #[tokio::test]
4662 async fn test_batch_cancel_orders_with_both_ids() {
4663 use nautilus_model::identifiers::{ClientOrderId, InstrumentId, VenueOrderId};
4664
4665 let client = OKXWebSocketClient::new(
4666 Some("wss://test.okx.com".to_string()),
4667 None,
4668 None,
4669 None,
4670 None,
4671 None,
4672 None,
4673 TransportBackend::default(),
4674 None,
4675 )
4676 .expect("Failed to create client");
4677
4678 let instrument_id = InstrumentId::from("BTC-USDT-SWAP.OKX");
4679 let client_order_id = ClientOrderId::new("order1");
4680 let venue_order_id = VenueOrderId::new("venue1");
4681
4682 let orders = vec![(instrument_id, Some(client_order_id), Some(venue_order_id))];
4683
4684 let result = client.batch_cancel_orders(orders).await;
4685
4686 assert!(result.is_err());
4687 }
4688
4689 #[tokio::test]
4690 async fn test_cancel_order_fails_without_inst_id_code() {
4691 use nautilus_model::identifiers::{ClientOrderId, InstrumentId, StrategyId, TraderId};
4692
4693 let client = OKXWebSocketClient::default();
4694 let instrument_id = InstrumentId::from("BTC-USDT-SWAP.OKX");
4695
4696 let result = client
4697 .cancel_order(
4698 TraderId::from("TESTER-001"),
4699 StrategyId::from("S-001"),
4700 instrument_id,
4701 Some(ClientOrderId::new("O-001")),
4702 None,
4703 )
4704 .await;
4705
4706 assert!(result.is_err());
4707 let err = result.unwrap_err().to_string();
4708 assert!(
4709 err.contains("No instIdCode cached for BTC-USDT-SWAP.OKX"),
4710 "Expected instIdCode error, found: {err}"
4711 );
4712 }
4713
4714 #[tokio::test]
4715 async fn test_submit_order_fails_without_inst_id_code() {
4716 use nautilus_model::{
4717 enums::{OrderSide, OrderType},
4718 identifiers::{ClientOrderId, InstrumentId, StrategyId, TraderId},
4719 types::Quantity,
4720 };
4721
4722 use crate::common::enums::OKXTradeMode;
4723
4724 let client = OKXWebSocketClient::default();
4725 let instrument_id = InstrumentId::from("ETH-USDT-SWAP.OKX");
4726
4727 let result = client
4728 .submit_order(
4729 TraderId::from("TESTER-001"),
4730 StrategyId::from("S-001"),
4731 instrument_id,
4732 OKXTradeMode::Cross,
4733 ClientOrderId::new("O-001"),
4734 OrderSide::Buy,
4735 OrderType::Limit,
4736 Quantity::from("0.01"),
4737 None,
4738 None,
4739 None,
4740 None,
4741 None,
4742 None,
4743 None,
4744 None,
4745 None,
4746 None,
4747 None,
4748 None,
4749 None,
4750 None,
4751 None,
4752 )
4753 .await;
4754
4755 assert!(result.is_err());
4756 let err = result.unwrap_err().to_string();
4757 assert!(
4758 err.contains("No instIdCode cached for ETH-USDT-SWAP.OKX"),
4759 "Expected instIdCode error, found: {err}"
4760 );
4761 }
4762
4763 #[tokio::test]
4764 async fn test_cancel_order_passes_inst_id_code_lookup_when_cached() {
4765 use nautilus_model::identifiers::{ClientOrderId, InstrumentId, StrategyId, TraderId};
4766 use ustr::Ustr;
4767
4768 let client = OKXWebSocketClient::default();
4769 let instrument_id = InstrumentId::from("BTC-USDT-SWAP.OKX");
4770
4771 client.cache_inst_id_code(Ustr::from("BTC-USDT-SWAP"), 10459);
4773
4774 let result = client
4775 .cancel_order(
4776 TraderId::from("TESTER-001"),
4777 StrategyId::from("S-001"),
4778 instrument_id,
4779 Some(ClientOrderId::new("O-001")),
4780 None,
4781 )
4782 .await;
4783
4784 assert!(result.is_err());
4786 let err = result.unwrap_err().to_string();
4787 assert!(
4788 !err.contains("No instIdCode cached"),
4789 "Should pass instIdCode lookup, found: {err}"
4790 );
4791 }
4792
4793 #[rstest]
4794 fn test_inst_id_code_cache_keeps_usd_and_usdc_instruments_distinct() {
4795 use ustr::Ustr;
4796
4797 let client = OKXWebSocketClient::default();
4798 client.cache_inst_id_code(Ustr::from("BTC-USD"), 10401);
4799 client.cache_inst_id_code(Ustr::from("BTC-USDC"), 20459);
4800
4801 assert_eq!(client.get_inst_id_code(&Ustr::from("BTC-USD")), Some(10401));
4802 assert_eq!(
4803 client.get_inst_id_code(&Ustr::from("BTC-USDC")),
4804 Some(20459)
4805 );
4806 assert_ne!(
4807 client.get_inst_id_code(&Ustr::from("BTC-USD")),
4808 client.get_inst_id_code(&Ustr::from("BTC-USDC"))
4809 );
4810 }
4811
4812 #[rstest]
4813 fn test_race_unsubscribe_failure_recovery() {
4814 let client = OKXWebSocketClient::new(
4820 Some("wss://test.okx.com".to_string()),
4821 None,
4822 None,
4823 None,
4824 None,
4825 None,
4826 None,
4827 TransportBackend::default(),
4828 None,
4829 )
4830 .expect("Failed to create client");
4831
4832 let topic = "trades:BTC-USDT-SWAP";
4833
4834 client.subscriptions_state.mark_subscribe(topic);
4836 client.subscriptions_state.confirm_subscribe(topic);
4837 assert_eq!(client.subscriptions_state.len(), 1);
4838
4839 client.subscriptions_state.mark_unsubscribe(topic);
4841 assert_eq!(client.subscriptions_state.len(), 0);
4842 assert_eq!(
4843 client.subscriptions_state.pending_unsubscribe_topics(),
4844 vec![topic]
4845 );
4846
4847 client.subscriptions_state.confirm_unsubscribe(topic); client.subscriptions_state.mark_subscribe(topic); client.subscriptions_state.confirm_subscribe(topic); assert_eq!(client.subscriptions_state.len(), 1);
4855 assert!(
4856 client
4857 .subscriptions_state
4858 .pending_unsubscribe_topics()
4859 .is_empty()
4860 );
4861 assert!(
4862 client
4863 .subscriptions_state
4864 .pending_subscribe_topics()
4865 .is_empty()
4866 );
4867
4868 let all = client.subscriptions_state.all_topics();
4870 assert_eq!(all.len(), 1);
4871 assert!(all.contains(&topic.to_string()));
4872 }
4873
4874 #[rstest]
4875 fn test_race_resubscribe_before_unsubscribe_ack() {
4876 let client = OKXWebSocketClient::new(
4880 Some("wss://test.okx.com".to_string()),
4881 None,
4882 None,
4883 None,
4884 None,
4885 None,
4886 None,
4887 TransportBackend::default(),
4888 None,
4889 )
4890 .expect("Failed to create client");
4891
4892 let topic = "books:BTC-USDT";
4893
4894 client.subscriptions_state.mark_subscribe(topic);
4896 client.subscriptions_state.confirm_subscribe(topic);
4897 assert_eq!(client.subscriptions_state.len(), 1);
4898
4899 client.subscriptions_state.mark_unsubscribe(topic);
4901 assert_eq!(client.subscriptions_state.len(), 0);
4902 assert_eq!(
4903 client.subscriptions_state.pending_unsubscribe_topics(),
4904 vec![topic]
4905 );
4906
4907 client.subscriptions_state.mark_subscribe(topic);
4909 assert_eq!(
4910 client.subscriptions_state.pending_subscribe_topics(),
4911 vec![topic]
4912 );
4913
4914 client.subscriptions_state.confirm_unsubscribe(topic);
4916 assert!(
4917 client
4918 .subscriptions_state
4919 .pending_unsubscribe_topics()
4920 .is_empty()
4921 );
4922 assert_eq!(
4923 client.subscriptions_state.pending_subscribe_topics(),
4924 vec![topic]
4925 );
4926
4927 client.subscriptions_state.confirm_subscribe(topic);
4929 assert_eq!(client.subscriptions_state.len(), 1);
4930 assert!(
4931 client
4932 .subscriptions_state
4933 .pending_subscribe_topics()
4934 .is_empty()
4935 );
4936
4937 let all = client.subscriptions_state.all_topics();
4939 assert_eq!(all.len(), 1);
4940 assert!(all.contains(&topic.to_string()));
4941 }
4942
4943 #[rstest]
4944 fn test_race_late_subscribe_confirmation_after_unsubscribe() {
4945 let client = OKXWebSocketClient::new(
4948 Some("wss://test.okx.com".to_string()),
4949 None,
4950 None,
4951 None,
4952 None,
4953 None,
4954 None,
4955 TransportBackend::default(),
4956 None,
4957 )
4958 .expect("Failed to create client");
4959
4960 let topic = "tickers:ETH-USDT";
4961
4962 client.subscriptions_state.mark_subscribe(topic);
4964 assert_eq!(
4965 client.subscriptions_state.pending_subscribe_topics(),
4966 vec![topic]
4967 );
4968
4969 client.subscriptions_state.mark_unsubscribe(topic);
4971 assert!(
4972 client
4973 .subscriptions_state
4974 .pending_subscribe_topics()
4975 .is_empty()
4976 ); assert_eq!(
4978 client.subscriptions_state.pending_unsubscribe_topics(),
4979 vec![topic]
4980 );
4981
4982 client.subscriptions_state.confirm_subscribe(topic);
4984 assert_eq!(client.subscriptions_state.len(), 0); assert_eq!(
4986 client.subscriptions_state.pending_unsubscribe_topics(),
4987 vec![topic]
4988 );
4989
4990 client.subscriptions_state.confirm_unsubscribe(topic);
4992
4993 assert!(client.subscriptions_state.is_empty());
4995 assert!(client.subscriptions_state.all_topics().is_empty());
4996 }
4997
4998 #[rstest]
4999 fn test_race_reconnection_with_pending_states() {
5000 let client = OKXWebSocketClient::new(
5002 Some("wss://test.okx.com".to_string()),
5003 Some("test_key".to_string()),
5004 Some("test_secret".to_string()),
5005 Some("test_passphrase".to_string()),
5006 Some(AccountId::new("OKX-TEST")),
5007 None,
5008 None,
5009 TransportBackend::default(),
5010 None,
5011 )
5012 .expect("Failed to create client");
5013
5014 let trade_btc = "trades:BTC-USDT-SWAP";
5017 client.subscriptions_state.mark_subscribe(trade_btc);
5018 client.subscriptions_state.confirm_subscribe(trade_btc);
5019
5020 let trade_eth = "trades:ETH-USDT-SWAP";
5022 client.subscriptions_state.mark_subscribe(trade_eth);
5023
5024 let book_btc = "books:BTC-USDT";
5026 client.subscriptions_state.mark_subscribe(book_btc);
5027 client.subscriptions_state.confirm_subscribe(book_btc);
5028 client.subscriptions_state.mark_unsubscribe(book_btc);
5029
5030 let topics_to_restore = client.subscriptions_state.all_topics();
5032
5033 assert_eq!(topics_to_restore.len(), 2);
5035 assert!(topics_to_restore.contains(&trade_btc.to_string()));
5036 assert!(topics_to_restore.contains(&trade_eth.to_string()));
5037 assert!(!topics_to_restore.contains(&book_btc.to_string())); }
5039
5040 #[rstest]
5041 fn test_race_duplicate_subscribe_messages_idempotent() {
5042 let client = OKXWebSocketClient::new(
5045 Some("wss://test.okx.com".to_string()),
5046 None,
5047 None,
5048 None,
5049 None,
5050 None,
5051 None,
5052 TransportBackend::default(),
5053 None,
5054 )
5055 .expect("Failed to create client");
5056
5057 let topic = "trades:BTC-USDT-SWAP";
5058
5059 client.subscriptions_state.mark_subscribe(topic);
5061 client.subscriptions_state.confirm_subscribe(topic);
5062 assert_eq!(client.subscriptions_state.len(), 1);
5063
5064 client.subscriptions_state.mark_subscribe(topic);
5066 assert!(
5067 client
5068 .subscriptions_state
5069 .pending_subscribe_topics()
5070 .is_empty()
5071 ); assert_eq!(client.subscriptions_state.len(), 1); client.subscriptions_state.confirm_subscribe(topic);
5076 assert_eq!(client.subscriptions_state.len(), 1);
5077
5078 let all = client.subscriptions_state.all_topics();
5080 assert_eq!(all.len(), 1);
5081 assert_eq!(all[0], topic);
5082 }
5083}