nautilus_okx/websocket/
handler.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! WebSocket message handler for OKX.
17//!
18//! The handler is a thin I/O boundary between the network layer and the client. It owns the
19//! `WebSocketClient`, deserializes raw venue messages into `OKXWsMessage` events, and handles
20//! subscription management, authentication, and retry logic.
21//!
22//! All domain parsing (venue types to Nautilus types) occurs outside the handler:
23//! - Data parsing in `PyOKXWebSocketClient` (uses an instruments cache)
24//! - Execution parsing in `execution.rs` (uses the system Cache)
25
26use std::{
27    collections::VecDeque,
28    sync::{
29        Arc,
30        atomic::{AtomicBool, Ordering},
31    },
32};
33
34use nautilus_model::identifiers::ClientOrderId;
35use nautilus_network::{
36    RECONNECTED,
37    retry::{RetryError, RetryManager, create_websocket_retry_manager},
38    websocket::{AuthTracker, SubscriptionState, TEXT_PING, TEXT_PONG, WebSocketClient},
39};
40use serde_json::{Map, Value};
41use tokio_tungstenite::tungstenite::Message;
42use ustr::Ustr;
43
44use super::{
45    enums::{OKXSubscriptionEvent, OKXWsChannel, OKXWsOperation},
46    error::OKXWsError,
47    messages::{
48        OKXOrderMsg, OKXSubscription, OKXSubscriptionArg, OKXWebSocketArg, OKXWebSocketError,
49        OKXWsFrame, OKXWsMessage,
50    },
51    subscription::topic_from_websocket_arg,
52};
53use crate::{
54    common::{
55        consts::{OKX_FIELD_SMSG, OKX_SUCCESS_CODE, should_retry_error_code},
56        enums::{OKXOrderStatus, OKXOrderType},
57        parse::prefer_rpi_response_fields,
58    },
59    websocket::client::OKX_RATE_LIMIT_KEY_SUBSCRIPTION,
60};
61
62/// Commands sent from the outer client to the inner message handler.
63#[derive(Debug)]
64pub enum HandlerCommand {
65    /// Set the WebSocketClient for the handler to use.
66    SetClient(WebSocketClient),
67    /// Disconnect the WebSocket connection.
68    Disconnect,
69    /// Send authentication payload to the WebSocket.
70    Authenticate { payload: String },
71    /// Subscribe to the given channels.
72    Subscribe { args: Vec<OKXSubscriptionArg> },
73    /// Unsubscribe from the given channels.
74    Unsubscribe { args: Vec<OKXSubscriptionArg> },
75    /// Send a pre-serialized payload (used for order operations).
76    Send {
77        payload: String,
78        rate_limit_keys: Option<Vec<Ustr>>,
79        request_id: Option<String>,
80        client_order_ids: Vec<ClientOrderId>,
81        op: Option<OKXWsOperation>,
82    },
83}
84
85pub(super) struct OKXWsFeedHandler {
86    signal: Arc<AtomicBool>,
87    inner: Option<WebSocketClient>,
88    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
89    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
90    out_tx: tokio::sync::mpsc::UnboundedSender<OKXWsMessage>,
91    auth_tracker: AuthTracker,
92    subscriptions_state: SubscriptionState,
93    retry_manager: RetryManager<OKXWsError>,
94    pending_messages: VecDeque<OKXWsMessage>,
95}
96
97impl OKXWsFeedHandler {
98    /// Creates a new [`OKXWsFeedHandler`] instance.
99    pub(super) fn new(
100        signal: Arc<AtomicBool>,
101        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
102        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
103        out_tx: tokio::sync::mpsc::UnboundedSender<OKXWsMessage>,
104        auth_tracker: AuthTracker,
105        subscriptions_state: SubscriptionState,
106    ) -> Self {
107        Self {
108            signal,
109            inner: None,
110            cmd_rx,
111            raw_rx,
112            out_tx,
113            auth_tracker,
114            subscriptions_state,
115            retry_manager: create_websocket_retry_manager(),
116            pending_messages: VecDeque::new(),
117        }
118    }
119
120    pub(super) fn is_stopped(&self) -> bool {
121        self.signal.load(Ordering::Acquire)
122    }
123
124    pub(super) fn send(&self, msg: OKXWsMessage) -> Result<(), ()> {
125        self.out_tx.send(msg).map_err(|_| ())
126    }
127
128    async fn send_with_retry(
129        &self,
130        payload: String,
131        rate_limit_keys: Option<&[Ustr]>,
132    ) -> Result<(), OKXWsError> {
133        if let Some(client) = &self.inner {
134            let keys_owned: Option<Vec<Ustr>> = rate_limit_keys.map(|k| k.to_vec());
135            self.retry_manager
136                .execute_with_retry(
137                    "websocket_send",
138                    || {
139                        let payload = payload.clone();
140                        let keys = keys_owned.clone();
141                        async move {
142                            client
143                                .send_text(payload, keys.as_deref())
144                                .await
145                                .map_err(|e| OKXWsError::SendFailed(e.to_string()))
146                        }
147                    },
148                    should_retry_okx_error,
149                    create_okx_retry_error,
150                )
151                .await
152        } else {
153            Err(OKXWsError::NoActiveClient)
154        }
155    }
156
157    pub(super) async fn send_pong(&self) -> anyhow::Result<()> {
158        match self.send_with_retry(TEXT_PONG.to_string(), None).await {
159            Ok(()) => {
160                log::trace!("Sent pong response to OKX text ping");
161                Ok(())
162            }
163            Err(e) => {
164                log::warn!("Failed to send pong after retries: error={e}");
165                Err(anyhow::anyhow!("Failed to send pong: {e}"))
166            }
167        }
168    }
169
170    pub(super) async fn next(&mut self) -> Option<OKXWsMessage> {
171        if let Some(message) = self.pending_messages.pop_front() {
172            return Some(message);
173        }
174
175        loop {
176            tokio::select! {
177                Some(cmd) = self.cmd_rx.recv() => {
178                    match cmd {
179                        HandlerCommand::SetClient(client) => {
180                            log::debug!("Handler received WebSocket client");
181                            self.inner = Some(client);
182                        }
183                        HandlerCommand::Disconnect => {
184                            log::debug!("Handler disconnecting WebSocket client");
185                            self.inner = None;
186                            return None;
187                        }
188                        HandlerCommand::Authenticate { payload } => {
189                            if let Err(e) = self.send_with_retry(
190                                payload,
191                                Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()),
192                            ).await {
193                                log::error!(
194                                    "Failed to send authentication message after retries: error={e}"
195                                );
196                            }
197                        }
198                        HandlerCommand::Subscribe { args } => {
199                            if let Err(e) = self.handle_subscribe(args).await {
200                                log::error!("Failed to handle subscribe command: error={e}");
201                            }
202                        }
203                        HandlerCommand::Unsubscribe { args } => {
204                            if let Err(e) = self.handle_unsubscribe(args).await {
205                                log::error!("Failed to handle unsubscribe command: error={e}");
206                            }
207                        }
208                        HandlerCommand::Send {
209                            payload,
210                            rate_limit_keys,
211                            request_id,
212                            client_order_ids,
213                            op,
214                        } => {
215                            if let Err(e) = self.send_with_retry(
216                                payload,
217                                rate_limit_keys.as_deref(),
218                            ).await {
219                                log::error!("Failed to send message after retries: error={e}");
220
221                                if let Some(request_id) = request_id {
222                                    self.pending_messages.push_back(OKXWsMessage::SendFailed {
223                                        request_id,
224                                        client_order_ids,
225                                        op,
226                                        error: e,
227                                    });
228                                }
229                            }
230                        }
231                    }
232                }
233
234                () = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
235                    if self.signal.load(Ordering::Acquire) {
236                        log::debug!("Stop signal received during idle period");
237                        return None;
238                    }
239                }
240
241                msg = self.raw_rx.recv() => {
242                    let event = match msg {
243                        Some(msg) => match Self::parse_raw_message(msg) {
244                            Some(event) => event,
245                            None => continue,
246                        },
247                        None => {
248                            log::debug!("WebSocket stream closed");
249                            return None;
250                        }
251                    };
252
253                    match event {
254                        OKXWsFrame::Ping => {
255                            if let Err(e) = self.send_pong().await {
256                                log::warn!("Failed to send pong response: error={e}");
257                            }
258                        }
259                        OKXWsFrame::Login {
260                            code, msg, conn_id, ..
261                        } => {
262                            if code == OKX_SUCCESS_CODE {
263                                self.auth_tracker.succeed();
264                                return Some(OKXWsMessage::Authenticated);
265                            }
266
267                            log::error!("WebSocket authentication failed: error={msg}");
268                            self.auth_tracker.fail(msg.clone());
269
270                            let error = OKXWebSocketError {
271                                code,
272                                message: msg,
273                                conn_id: Some(conn_id),
274                                timestamp: nautilus_core::time::get_atomic_clock_realtime()
275                                    .get_time_ns()
276                                    .as_u64(),
277                            };
278                            self.pending_messages.push_back(OKXWsMessage::Error(error));
279                        }
280                        OKXWsFrame::BookData { arg, action, data } => {
281                            return Some(OKXWsMessage::BookData { arg, action, data });
282                        }
283                        OKXWsFrame::RpiBookData { arg, action, data } => {
284                            return Some(OKXWsMessage::RpiBookData { arg, action, data });
285                        }
286                        OKXWsFrame::OrderResponse {
287                            id, op, code, msg, data,
288                        } => {
289                            return Some(OKXWsMessage::OrderResponse {
290                                id, op, code, msg, data,
291                            });
292                        }
293                        OKXWsFrame::Data { arg, data } => {
294                            if let Some(output) = self.route_data_message(arg, data) {
295                                return Some(output);
296                            }
297                        }
298                        OKXWsFrame::Error { arg, code, msg } => {
299                            let arg = arg.or_else(|| subscription_arg_from_error_message(&msg));
300                            if let Some(arg) = arg
301                                && self.handle_subscription_error(&arg, &code, &msg)
302                            {
303                                return Some(OKXWsMessage::SubscriptionFailed {
304                                    channel: arg.channel,
305                                    inst_id: arg.inst_id,
306                                    code,
307                                    msg,
308                                });
309                            }
310
311                            let error = OKXWebSocketError {
312                                code,
313                                message: msg,
314                                conn_id: None,
315                                timestamp: nautilus_core::time::get_atomic_clock_realtime()
316                                    .get_time_ns()
317                                    .as_u64(),
318                            };
319                            return Some(OKXWsMessage::Error(error));
320                        }
321                        OKXWsFrame::Reconnected => {
322                            self.auth_tracker.invalidate();
323                            return Some(OKXWsMessage::Reconnected);
324                        }
325                        OKXWsFrame::Subscription {
326                            event, arg, code, msg,
327                            ..
328                        } => {
329                            let rejected = self
330                                .handle_subscription_ack(&event, &arg, code.as_deref(), msg.as_deref());
331
332                            if rejected {
333                                return Some(OKXWsMessage::SubscriptionFailed {
334                                    channel: arg.channel,
335                                    inst_id: arg.inst_id,
336                                    code: code.unwrap_or_default(),
337                                    msg: msg.unwrap_or_default(),
338                                });
339                            }
340                        }
341                        OKXWsFrame::ChannelConnCount { .. } => {}
342                    }
343                }
344
345                else => {
346                    log::debug!("Handler shutting down: stream ended or command channel closed");
347                    return None;
348                }
349            }
350        }
351    }
352
353    fn route_data_message(&self, arg: OKXWebSocketArg, mut data: Value) -> Option<OKXWsMessage> {
354        let OKXWebSocketArg {
355            channel, inst_id, ..
356        } = arg;
357
358        match channel {
359            OKXWsChannel::Account => Some(OKXWsMessage::Account(data)),
360            OKXWsChannel::Positions => Some(OKXWsMessage::Positions(data)),
361            OKXWsChannel::Orders => {
362                parse_array_items(data, "orders", false).map(OKXWsMessage::Orders)
363            }
364            OKXWsChannel::SprdOrders => {
365                parse_array_items(data, "spread orders", false).map(OKXWsMessage::SpreadOrders)
366            }
367            OKXWsChannel::OrdersAlgo | OKXWsChannel::AlgoAdvance => {
368                parse_array_items(data, "algo orders", false).map(OKXWsMessage::AlgoOrders)
369            }
370            OKXWsChannel::LiquidationWarning => {
371                parse_array_items(data, "liquidation warnings", false)
372                    .map(OKXWsMessage::LiquidationWarnings)
373            }
374            OKXWsChannel::Instruments => {
375                prefer_rpi_response_fields(&mut data);
376                parse_array_items(data, "instruments", true).map(OKXWsMessage::Instruments)
377            }
378            _ => Some(OKXWsMessage::ChannelData {
379                channel,
380                inst_id,
381                data,
382            }),
383        }
384    }
385
386    fn handle_subscription_ack(
387        &self,
388        event: &OKXSubscriptionEvent,
389        arg: &OKXWebSocketArg,
390        code: Option<&str>,
391        msg: Option<&str>,
392    ) -> bool {
393        let topic = topic_from_websocket_arg(arg);
394        let success = code.is_none_or(|c| c == OKX_SUCCESS_CODE);
395
396        match event {
397            OKXSubscriptionEvent::Subscribe => {
398                if success {
399                    self.subscriptions_state.confirm_subscribe(&topic);
400                    false
401                } else {
402                    log::warn!(
403                        "Subscription failed: topic={topic:?}, error={msg:?}, code={code:?}"
404                    );
405                    self.subscriptions_state.mark_failure(&topic);
406                    true
407                }
408            }
409            OKXSubscriptionEvent::Unsubscribe => {
410                if success {
411                    self.subscriptions_state.confirm_unsubscribe(&topic);
412                } else {
413                    log::warn!(
414                        "Unsubscription failed - restoring subscription: \
415                         topic={topic:?}, error={msg:?}, code={code:?}"
416                    );
417                    self.subscriptions_state.confirm_unsubscribe(&topic);
418                    self.subscriptions_state.mark_subscribe(&topic);
419                    self.subscriptions_state.confirm_subscribe(&topic);
420                }
421                false
422            }
423        }
424    }
425
426    fn handle_subscription_error(&self, arg: &OKXWebSocketArg, code: &str, msg: &str) -> bool {
427        let topic = topic_from_websocket_arg(arg);
428        let event = if self
429            .subscriptions_state
430            .pending_unsubscribe_topics()
431            .iter()
432            .any(|pending| pending == &topic)
433        {
434            OKXSubscriptionEvent::Unsubscribe
435        } else if self
436            .subscriptions_state
437            .pending_subscribe_topics()
438            .iter()
439            .any(|pending| pending == &topic)
440        {
441            OKXSubscriptionEvent::Subscribe
442        } else {
443            return false;
444        };
445
446        self.handle_subscription_ack(&event, arg, Some(code), Some(msg))
447    }
448
449    async fn handle_subscribe(&self, args: Vec<OKXSubscriptionArg>) -> anyhow::Result<()> {
450        for arg in &args {
451            log::debug!(
452                "Subscribing to channel: channel={:?}, inst_id={:?}",
453                arg.channel,
454                arg.inst_id
455            );
456        }
457
458        let message = OKXSubscription {
459            op: OKXWsOperation::Subscribe,
460            args,
461        };
462
463        let json_txt = serde_json::to_string(&message)
464            .map_err(|e| anyhow::anyhow!("Failed to serialize subscription: {e}"))?;
465
466        self.send_with_retry(json_txt, Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
467            .await
468            .map_err(|e| anyhow::anyhow!("Failed to send subscription after retries: {e}"))?;
469        Ok(())
470    }
471
472    async fn handle_unsubscribe(&self, args: Vec<OKXSubscriptionArg>) -> anyhow::Result<()> {
473        for arg in &args {
474            log::debug!(
475                "Unsubscribing from channel: channel={:?}, inst_id={:?}",
476                arg.channel,
477                arg.inst_id
478            );
479        }
480
481        let message = OKXSubscription {
482            op: OKXWsOperation::Unsubscribe,
483            args,
484        };
485
486        let json_txt = serde_json::to_string(&message)
487            .map_err(|e| anyhow::anyhow!("Failed to serialize unsubscription: {e}"))?;
488
489        self.send_with_retry(json_txt, Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
490            .await
491            .map_err(|e| anyhow::anyhow!("Failed to send unsubscription after retries: {e}"))?;
492        Ok(())
493    }
494
495    pub(crate) fn parse_raw_message(
496        msg: tokio_tungstenite::tungstenite::Message,
497    ) -> Option<OKXWsFrame> {
498        match msg {
499            tokio_tungstenite::tungstenite::Message::Text(text) => {
500                if text == TEXT_PONG {
501                    log::trace!("Received pong from OKX");
502                    return None;
503                }
504
505                if text == TEXT_PING {
506                    log::trace!("Received ping from OKX (text)");
507                    return Some(OKXWsFrame::Ping);
508                }
509
510                if text == RECONNECTED {
511                    log::debug!("Received WebSocket reconnection signal");
512                    return Some(OKXWsFrame::Reconnected);
513                }
514                log::trace!("Received WebSocket message: {text}");
515
516                match serde_json::from_str(&text) {
517                    Ok(ws_event) => match &ws_event {
518                        OKXWsFrame::Error { code, msg, .. } => {
519                            if should_retry_error_code(code) {
520                                log::warn!("WebSocket error: {code} - {msg}");
521                            } else {
522                                log::error!("WebSocket error: {code} - {msg}");
523                            }
524                            Some(ws_event)
525                        }
526                        OKXWsFrame::Login {
527                            event,
528                            code,
529                            msg,
530                            conn_id,
531                        } => {
532                            if code == OKX_SUCCESS_CODE {
533                                log::debug!("WebSocket authenticated: conn_id={conn_id}");
534                            } else {
535                                log::error!(
536                                    "WebSocket authentication failed: \
537                                     event={event}, code={code}, error={msg}"
538                                );
539                            }
540                            Some(ws_event)
541                        }
542                        OKXWsFrame::Subscription {
543                            event,
544                            arg,
545                            conn_id,
546                            ..
547                        } => {
548                            let channel_str = serde_json::to_string(&arg.channel)
549                                .expect("Invalid OKX websocket channel")
550                                .trim_matches('"')
551                                .to_string();
552                            log::debug!("{event}d: channel={channel_str}, conn_id={conn_id}");
553                            Some(ws_event)
554                        }
555                        OKXWsFrame::ChannelConnCount {
556                            channel,
557                            conn_count,
558                            conn_id,
559                            ..
560                        } => {
561                            let channel_str = serde_json::to_string(channel)
562                                .expect("Invalid OKX websocket channel")
563                                .trim_matches('"')
564                                .to_string();
565                            log::debug!(
566                                "Channel connection status: \
567                                 channel={channel_str}, connections={conn_count}, conn_id={conn_id}",
568                            );
569                            None
570                        }
571                        OKXWsFrame::Ping => {
572                            log::trace!("Ignoring ping event parsed from text payload");
573                            None
574                        }
575                        OKXWsFrame::Data { .. }
576                        | OKXWsFrame::BookData { .. }
577                        | OKXWsFrame::RpiBookData { .. } => Some(ws_event),
578                        OKXWsFrame::OrderResponse {
579                            id, op, code, data, ..
580                        } => {
581                            if code == OKX_SUCCESS_CODE {
582                                log::debug!(
583                                    "Order operation successful: id={id:?}, op={op}, code={code}"
584                                );
585
586                                if let Some(order_data) = data.first() {
587                                    let success_msg = order_data
588                                        .get(OKX_FIELD_SMSG)
589                                        .and_then(|s| s.as_str())
590                                        .unwrap_or("Order operation successful");
591                                    log::debug!("Order success details: {success_msg}");
592                                }
593                            }
594                            Some(ws_event)
595                        }
596                        OKXWsFrame::Reconnected => {
597                            log::warn!("Unexpected Reconnected event from deserialization");
598                            None
599                        }
600                    },
601                    Err(e) => {
602                        log::error!("Failed to parse message: {e}: {text}");
603                        None
604                    }
605                }
606            }
607            Message::Ping(_payload) => {
608                log::trace!("Received binary ping frame from OKX");
609                Some(OKXWsFrame::Ping)
610            }
611            Message::Pong(payload) => {
612                log::trace!("Received pong frame from OKX ({} bytes)", payload.len());
613                None
614            }
615            Message::Binary(msg) => {
616                log::debug!("Raw binary frame ({} bytes)", msg.len());
617                log::trace!("Raw binary: {msg:?}");
618                None
619            }
620            Message::Close(_) => {
621                log::debug!("Received close message");
622                None
623            }
624            msg => {
625                log::warn!("Unexpected message: {msg}");
626                None
627            }
628        }
629    }
630}
631
632fn subscription_arg_from_error_message(msg: &str) -> Option<OKXWebSocketArg> {
633    let descriptor = msg
634        .strip_prefix("Wrong URL or channel:")?
635        .split_whitespace()
636        .next()?;
637    let mut fields = descriptor.split(',');
638    let channel = fields.next()?;
639    let mut arg = Map::new();
640    arg.insert("channel".to_string(), Value::String(channel.to_string()));
641
642    for field in fields {
643        let (key, value) = field.split_once(':')?;
644        if !matches!(key, "instId" | "sprdId" | "instType" | "instFamily") {
645            return None;
646        }
647        arg.insert(key.to_string(), Value::String(value.to_string()));
648    }
649
650    serde_json::from_value(Value::Object(arg)).ok()
651}
652
653/// Returns `true` when an OKX WebSocket order message represents a post-only auto-cancel.
654pub fn is_post_only_auto_cancel(msg: &OKXOrderMsg) -> bool {
655    use crate::common::{consts::OKX_POST_ONLY_CANCEL_SOURCE, enums::OKXOrderStatus};
656
657    if msg.state != OKXOrderStatus::Canceled {
658        return false;
659    }
660
661    let cancel_source_matches = matches!(
662        msg.cancel_source.as_deref(),
663        Some(source) if source == OKX_POST_ONLY_CANCEL_SOURCE
664    );
665
666    let reason_matches = matches!(
667        msg.cancel_source_reason.as_deref(),
668        Some(reason) if reason.contains("POST_ONLY")
669    );
670
671    if !(cancel_source_matches || reason_matches) {
672        return false;
673    }
674
675    msg.acc_fill_sz
676        .as_ref()
677        .is_none_or(|filled| filled == "0" || filled.is_empty())
678}
679
680/// Returns `true` when an RPI order update is canceled without any fill.
681pub fn is_unfilled_rpi_cancel(msg: &OKXOrderMsg) -> bool {
682    msg.ord_type == OKXOrderType::Rpi
683        && msg.state == OKXOrderStatus::Canceled
684        && msg
685            .acc_fill_sz
686            .as_ref()
687            .is_none_or(|filled| filled == "0" || filled.is_empty())
688}
689
690// Per-item deserialization so one malformed entry does not drop the batch.
691fn parse_array_items<T: serde::de::DeserializeOwned>(
692    data: Value,
693    label: &str,
694    warn_on_parse_error: bool,
695) -> Option<Vec<T>> {
696    let Value::Array(items) = data else {
697        if warn_on_parse_error {
698            log::warn!("Expected {label} payload to be a JSON array");
699        } else {
700            log::error!("Expected {label} payload to be a JSON array");
701        }
702        return None;
703    };
704
705    let mut parsed = Vec::with_capacity(items.len());
706    for (idx, item) in items.into_iter().enumerate() {
707        match serde_json::from_value::<T>(item) {
708            Ok(value) => parsed.push(value),
709            Err(e) => {
710                if warn_on_parse_error {
711                    log::warn!("Failed to parse {label} item at index {idx}: {e}");
712                } else {
713                    log::error!("Failed to parse {label} item at index {idx}: {e}");
714                }
715            }
716        }
717    }
718
719    if parsed.is_empty() {
720        None
721    } else {
722        Some(parsed)
723    }
724}
725
726#[inline]
727fn contains_ignore_ascii_case(haystack: &str, needle: &str) -> bool {
728    haystack
729        .as_bytes()
730        .windows(needle.len())
731        .any(|window| window.eq_ignore_ascii_case(needle.as_bytes()))
732}
733
734// Specific phrases rather than bare "connection"/"network", which appear
735// in permanent errors too (e.g. "no active WebSocket client connection").
736const RETRYABLE_CLIENT_ERROR_PHRASES: &[&str] = &[
737    "timeout",
738    "timed out",
739    "connection reset",
740    "connection refused",
741    "connection closed",
742    "connection aborted",
743    "broken pipe",
744    "network unreachable",
745    "network is unreachable",
746    "no route to host",
747];
748
749fn should_retry_okx_error(error: &OKXWsError) -> bool {
750    match error {
751        OKXWsError::OkxError { error_code, .. } => should_retry_error_code(error_code),
752        OKXWsError::TungsteniteError(_)
753        | OKXWsError::SendFailed(_)
754        | OKXWsError::OperationTimeout { .. } => true,
755        OKXWsError::ClientError(msg) => RETRYABLE_CLIENT_ERROR_PHRASES
756            .iter()
757            .any(|phrase| contains_ignore_ascii_case(msg, phrase)),
758        OKXWsError::AuthenticationError(_)
759        | OKXWsError::JsonError(_)
760        | OKXWsError::ParsingError(_)
761        | OKXWsError::NoActiveClient
762        | OKXWsError::HandlerUnavailable(_) => false,
763    }
764}
765
766fn create_okx_retry_error(error: RetryError) -> OKXWsError {
767    match error {
768        RetryError::OperationTimeout { timeout_ms } => OKXWsError::OperationTimeout { timeout_ms },
769        RetryError::InvalidConfiguration { message } => OKXWsError::ClientError(message),
770        RetryError::Canceled => {
771            OKXWsError::SendFailed("Adapter disconnecting or shutting down".to_string())
772        }
773        error @ RetryError::ElapsedBudgetExceeded { .. } => {
774            OKXWsError::SendFailed(error.to_string())
775        }
776    }
777}
778
779#[cfg(test)]
780mod tests {
781    use std::sync::{Arc, atomic::AtomicBool};
782
783    use nautilus_network::websocket::{AuthTracker, SubscriptionState};
784    use rstest::rstest;
785    use serde_json::json;
786
787    use super::*;
788    use crate::common::{
789        consts::OKX_WS_TOPIC_DELIMITER, enums::OKXRpiPermission, testing::load_test_json,
790    };
791
792    fn create_handler() -> OKXWsFeedHandler {
793        let signal = Arc::new(AtomicBool::new(false));
794        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
795        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
796        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
797
798        OKXWsFeedHandler::new(
799            signal,
800            cmd_rx,
801            raw_rx,
802            out_tx,
803            AuthTracker::new(),
804            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
805        )
806    }
807
808    #[rstest]
809    fn test_should_retry_typed_send_and_timeout_errors() {
810        assert!(should_retry_okx_error(&OKXWsError::SendFailed(
811            "connection reset".to_string()
812        )));
813        assert!(should_retry_okx_error(&OKXWsError::OperationTimeout {
814            timeout_ms: 1_000
815        }));
816        assert!(!should_retry_okx_error(&OKXWsError::NoActiveClient));
817        assert!(!should_retry_okx_error(&OKXWsError::HandlerUnavailable(
818            "closed".to_string()
819        )));
820    }
821
822    #[rstest]
823    #[case("Connection reset by peer", true)]
824    #[case("send timeout after 30s", true)]
825    #[case("Connection closed unexpectedly", true)]
826    #[case("Broken pipe", true)]
827    #[case("Network unreachable", true)]
828    #[case("No active WebSocket client connection", false)]
829    #[case("network protocol upgrade required", false)]
830    #[case("invalid frame format", false)]
831    fn test_should_retry_client_error(#[case] msg: &str, #[case] expected: bool) {
832        let err = OKXWsError::ClientError(msg.to_string());
833        assert_eq!(should_retry_okx_error(&err), expected);
834    }
835
836    #[rstest]
837    fn test_subscription_error_restores_failed_unsubscribe() {
838        let handler = create_handler();
839        let arg = OKXWebSocketArg {
840            channel: OKXWsChannel::Books,
841            inst_id: Some(Ustr::from("BTC-USD")),
842            inst_type: None,
843            inst_family: None,
844            bar: None,
845        };
846        let topic = topic_from_websocket_arg(&arg);
847        handler.subscriptions_state.mark_subscribe(&topic);
848        handler.subscriptions_state.confirm_subscribe(&topic);
849        handler.subscriptions_state.mark_unsubscribe(&topic);
850
851        let rejected_subscription =
852            handler.handle_subscription_error(&arg, "60019", "Unsubscription failed");
853
854        assert!(!rejected_subscription);
855        assert_eq!(handler.subscriptions_state.all_topics(), vec![topic]);
856        assert!(
857            handler
858                .subscriptions_state
859                .pending_subscribe_topics()
860                .is_empty()
861        );
862        assert!(
863            handler
864                .subscriptions_state
865                .pending_unsubscribe_topics()
866                .is_empty()
867        );
868    }
869
870    #[rstest]
871    fn test_subscription_arg_from_error_message_matches_mainnet_shape() {
872        let msg = "Wrong URL or channel:books,instId:BTC-USDT-SWAP doesn't exist. Please use the \
873                   correct URL, channel and parameters referring to API document.";
874
875        let arg = subscription_arg_from_error_message(msg).unwrap();
876
877        assert_eq!(arg.channel, OKXWsChannel::Books);
878        assert_eq!(arg.inst_id, Some(Ustr::from("BTC-USDT-SWAP")));
879        assert_eq!(arg.inst_type, None);
880        assert_eq!(arg.inst_family, None);
881        assert_eq!(arg.bar, None);
882    }
883
884    #[rstest]
885    fn test_subscription_error_ignores_non_pending_topic() {
886        let handler = create_handler();
887        let arg = OKXWebSocketArg {
888            channel: OKXWsChannel::Books,
889            inst_id: Some(Ustr::from("BTC-USDT-SWAP")),
890            inst_type: None,
891            inst_family: None,
892            bar: None,
893        };
894
895        let rejected_subscription =
896            handler.handle_subscription_error(&arg, "60018", "Subscription failed");
897
898        assert!(!rejected_subscription);
899        assert!(handler.subscriptions_state.all_topics().is_empty());
900        assert!(
901            handler
902                .subscriptions_state
903                .pending_subscribe_topics()
904                .is_empty()
905        );
906        assert!(
907            handler
908                .subscriptions_state
909                .pending_unsubscribe_topics()
910                .is_empty()
911        );
912    }
913
914    #[derive(serde::Deserialize, Debug, PartialEq)]
915    struct ParseArrayItem {
916        value: i64,
917    }
918
919    #[rstest]
920    fn test_parse_array_items_keeps_good_items_when_one_fails() {
921        let data = json!([
922            {"value": 1},
923            {"value": "not a number"},
924            {"value": 3},
925        ]);
926
927        let parsed: Vec<ParseArrayItem> =
928            parse_array_items(data, "test", false).expect("non-empty");
929        assert_eq!(
930            parsed,
931            vec![ParseArrayItem { value: 1 }, ParseArrayItem { value: 3 }],
932        );
933    }
934
935    #[rstest]
936    fn test_parse_array_items_returns_none_when_payload_not_array() {
937        let data = json!({"not": "an array"});
938        let parsed: Option<Vec<ParseArrayItem>> = parse_array_items(data, "test", false);
939        assert!(parsed.is_none());
940    }
941
942    #[rstest]
943    fn test_parse_array_items_returns_none_when_all_items_fail() {
944        let data = json!([{"value": "bad"}]);
945        let parsed: Option<Vec<ParseArrayItem>> = parse_array_items(data, "test", false);
946        assert!(parsed.is_none());
947    }
948
949    #[rstest]
950    fn test_route_instruments_keeps_valid_items_when_one_item_fails() {
951        let handler = create_handler();
952        let mut frame: Value =
953            serde_json::from_str(&load_test_json("ws_instruments.json")).expect("valid fixture");
954        let data = frame
955            .get_mut("data")
956            .and_then(Value::as_array_mut)
957            .expect("data array");
958        let mut invalid_item = data[0].clone();
959        invalid_item["tickSz"] = json!(7);
960        data.insert(0, invalid_item);
961
962        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
963        let msg = handler
964            .route_data_message(arg, frame["data"].clone())
965            .expect("instruments message");
966
967        match msg {
968            OKXWsMessage::Instruments(instruments) => {
969                assert_eq!(instruments.len(), 1);
970                assert_eq!(instruments[0].inst_id.as_str(), "BTC-USDT-SWAP");
971            }
972            other => panic!("Expected Instruments, was {other:?}"),
973        }
974    }
975
976    #[rstest]
977    fn test_route_instruments_prefers_rpi_over_legacy_alias() {
978        let handler = create_handler();
979        let mut frame: Value =
980            serde_json::from_str(&load_test_json("ws_instruments.json")).expect("valid fixture");
981        let instrument = &mut frame["data"][0];
982        instrument["rpi"] = json!("2");
983        instrument["elp"] = json!("1");
984
985        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
986        let msg = handler
987            .route_data_message(arg, frame["data"].clone())
988            .expect("instruments message");
989
990        match msg {
991            OKXWsMessage::Instruments(instruments) => {
992                assert_eq!(instruments.len(), 1);
993                assert_eq!(instruments[0].rpi, Some(OKXRpiPermission::Permitted));
994            }
995            other => panic!("Expected Instruments, was {other:?}"),
996        }
997    }
998
999    #[rstest]
1000    fn test_route_liquidation_warnings() {
1001        let handler = create_handler();
1002        let frame: Value = serde_json::from_str(&load_test_json("ws_liquidation_warning.json"))
1003            .expect("valid fixture");
1004
1005        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
1006        let msg = handler
1007            .route_data_message(arg, frame["data"].clone())
1008            .expect("liquidation warning message");
1009
1010        match msg {
1011            OKXWsMessage::LiquidationWarnings(warnings) => {
1012                assert_eq!(warnings.len(), 1);
1013                assert_eq!(warnings[0].inst_id.as_str(), "BTC-USDT-SWAP");
1014                assert_eq!(warnings[0].mgn_ratio, "0.62");
1015            }
1016            other => panic!("Expected LiquidationWarnings, was {other:?}"),
1017        }
1018    }
1019}
nautilus_okx/websocket/handler.rs

nautilus_okx/websocket/
handler.rs
